{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MagicHash #-}

module Data.Trie.Binary.Magma where

import Data.Bits
import Data.Coerce
import GHC.Exts
import GHC.TypeLits

{- ORMOLU_DISABLE -}
-- * Examples
{- $examples
=== __ABC308G Minimum Xor Pair Query__

> data A = Done !Word | Partial !Word
>
> -- non-associative
> instance Magma A where
>     magma = m
>       where
>         m (Done x) (Done y) = Done (min x y)
>         m (Done x) Partial{} = Done x
>         m Partial{} (Done y) = Done y
>         m (Partial x) (Partial y) = Done (xor x y)
>     {-# INLINE magma #-}
>
> type B = Int
>
> tip :: Word -> B -> A
> tip w b
>     | b >= 2 = Done 0
>     | otherwise = Partial w
> {-# INLINE tip #-}
>
> ins :: Maybe B -> Maybe B
> ins = Just . maybe 1 (+ 1)
> {-# INLINE ins #-}
>
> del :: Maybe B -> Maybe B
> del ma = case ma of
>     Just a | a > 1 -> Just (a - 1)
>     _ -> Nothing
> {-# INLINE del #-}
>
> insert :: (KnownNat h) => Word -> Trie h A B -> Trie h A B
> insert = alterWith tip ins
> {-# INLINE insert #-}
>
> delete :: (KnownNat h) => Word -> Trie h A B -> Trie h A B
> delete = alterWith tip del
> {-# INLINE delete #-}
>
> magmaAll :: Trie h A B -> Maybe A
> magmaAll = magmaAllWith tip
> {-# INLINE magmaAll #-}
-}
{- ORMOLU_ENABLE -}

-- * Binary Trie

data Trie (h :: Nat) a b
  = Bin !a !(Trie h a b) !(Trie h a b)
  | Tip !Word !b
  | Nil
  deriving (Int -> Trie h a b -> ShowS
[Trie h a b] -> ShowS
Trie h a b -> String
(Int -> Trie h a b -> ShowS)
-> (Trie h a b -> String)
-> ([Trie h a b] -> ShowS)
-> Show (Trie h a b)
forall (h :: Nat) a b.
(Show a, Show b) =>
Int -> Trie h a b -> ShowS
forall (h :: Nat) a b. (Show a, Show b) => [Trie h a b] -> ShowS
forall (h :: Nat) a b. (Show a, Show b) => Trie h a b -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (h :: Nat) a b.
(Show a, Show b) =>
Int -> Trie h a b -> ShowS
showsPrec :: Int -> Trie h a b -> ShowS
$cshow :: forall (h :: Nat) a b. (Show a, Show b) => Trie h a b -> String
show :: Trie h a b -> String
$cshowList :: forall (h :: Nat) a b. (Show a, Show b) => [Trie h a b] -> ShowS
showList :: [Trie h a b] -> ShowS
Show)

magmaAllWith ::
  -- | Tip to Magma
  (Word -> b -> a) ->
  Trie h a b ->
  Maybe a
magmaAllWith :: forall b a (h :: Nat). (Word -> b -> a) -> Trie h a b -> Maybe a
magmaAllWith Word -> b -> a
_ (Bin a
a Trie h a b
_ Trie h a b
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
a
magmaAllWith Word -> b -> a
u (Tip Word
w b
b) = a -> Maybe a
forall a. a -> Maybe a
Just (Word -> b -> a
u Word
w b
b)
magmaAllWith Word -> b -> a
_ Trie h a b
Nil = Maybe a
forall a. Maybe a
Nothing

binWith ::
  (Magma a) =>
  -- | Tip to Magma
  (Word -> b -> a) ->
  -- | left child
  Trie h a b ->
  -- | right child
  Trie h a b ->
  Trie h a b
binWith :: forall a b (h :: Nat).
Magma a =>
(Word -> b -> a) -> Trie h a b -> Trie h a b -> Trie h a b
binWith Word -> b -> a
_ t0 :: Trie h a b
t0@(Bin a
a0 Trie h a b
_ Trie h a b
_) t1 :: Trie h a b
t1@(Bin a
a1 Trie h a b
_ Trie h a b
_) = a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin (a -> a -> a
forall a. Magma a => a -> a -> a
magma a
a0 a
a1) Trie h a b
t0 Trie h a b
t1
binWith Word -> b -> a
u t0 :: Trie h a b
t0@(Bin a
a0 Trie h a b
_ Trie h a b
_) t1 :: Trie h a b
t1@(Tip Word
w1 b
b1) = a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin (a -> a -> a
forall a. Magma a => a -> a -> a
magma a
a0 (Word -> b -> a
u Word
w1 b
b1)) Trie h a b
t0 Trie h a b
t1
binWith Word -> b -> a
_ t0 :: Trie h a b
t0@(Bin a
a0 Trie h a b
_ Trie h a b
_) t1 :: Trie h a b
t1@Trie h a b
Nil = a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin a
a0 Trie h a b
t0 Trie h a b
t1
binWith Word -> b -> a
u t0 :: Trie h a b
t0@(Tip Word
w0 b
b0) t1 :: Trie h a b
t1@(Bin a
a1 Trie h a b
_ Trie h a b
_) = a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin (a -> a -> a
forall a. Magma a => a -> a -> a
magma (Word -> b -> a
u Word
w0 b
b0) a
a1) Trie h a b
t0 Trie h a b
t1
binWith Word -> b -> a
u t0 :: Trie h a b
t0@(Tip Word
w0 b
b0) t1 :: Trie h a b
t1@(Tip Word
w1 b
b1) = a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin (a -> a -> a
forall a. Magma a => a -> a -> a
magma (Word -> b -> a
u Word
w0 b
b0) (Word -> b -> a
u Word
w1 b
b1)) Trie h a b
t0 Trie h a b
t1
binWith Word -> b -> a
_ t0 :: Trie h a b
t0@Tip{} Trie h a b
Nil = Trie h a b
t0
binWith Word -> b -> a
_ t0 :: Trie h a b
t0@Trie h a b
Nil t1 :: Trie h a b
t1@(Bin a
a1 Trie h a b
_ Trie h a b
_) = a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin a
a1 Trie h a b
t0 Trie h a b
t1
binWith Word -> b -> a
_ Trie h a b
Nil t1 :: Trie h a b
t1@Tip{} = Trie h a b
t1
binWith Word -> b -> a
_ Trie h a b
Nil Trie h a b
Nil = Trie h a b
forall (h :: Nat) a b. Trie h a b
Nil

alterWith ::
  forall h a b.
  (Magma a, KnownNat h) =>
  -- | Tip to Magma
  (Word -> b -> a) ->
  -- | insert/delete
  (Maybe b -> Maybe b) ->
  Word ->
  Trie h a b ->
  Trie h a b
alterWith :: forall (h :: Nat) a b.
(Magma a, KnownNat h) =>
(Word -> b -> a)
-> (Maybe b -> Maybe b) -> Word -> Trie h a b -> Trie h a b
alterWith Word -> b -> a
u Maybe b -> Maybe b
f Word
v = Word -> Trie h a b -> Trie h a b
forall {h :: Nat}. Word -> Trie h a b -> Trie h a b
go (Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL Word
1 (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' @h Proxy# h
forall {k} (a :: k). Proxy# a
proxy#) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
  where
    go :: Word -> Trie h a b -> Trie h a b
go !Word
flg Trie h a b
trie = case Trie h a b
trie of
      Bin a
_ Trie h a b
t0 Trie h a b
t1
        | Word
v Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
flg Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0 -> (Word -> b -> a) -> Trie h a b -> Trie h a b -> Trie h a b
forall a b (h :: Nat).
Magma a =>
(Word -> b -> a) -> Trie h a b -> Trie h a b -> Trie h a b
binWith Word -> b -> a
u (Word -> Trie h a b -> Trie h a b
go (Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR Word
flg Int
1) Trie h a b
t0) Trie h a b
t1
        | Bool
otherwise -> (Word -> b -> a) -> Trie h a b -> Trie h a b -> Trie h a b
forall a b (h :: Nat).
Magma a =>
(Word -> b -> a) -> Trie h a b -> Trie h a b -> Trie h a b
binWith Word -> b -> a
u Trie h a b
t0 (Word -> Trie h a b -> Trie h a b
go (Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR Word
flg Int
1) Trie h a b
t1)
      Tip Word
w b
a
        | Word
v Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
w -> Trie h a b -> (b -> Trie h a b) -> Maybe b -> Trie h a b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Trie h a b
forall (h :: Nat) a b. Trie h a b
Nil (Word -> b -> Trie h a b
forall (h :: Nat) a b. Word -> b -> Trie h a b
Tip Word
v) (Maybe b -> Trie h a b) -> Maybe b -> Trie h a b
forall a b. (a -> b) -> a -> b
$ Maybe b -> Maybe b
f (b -> Maybe b
forall a. a -> Maybe a
Just b
a)
        | Word
w Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
flg Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0 -> Word -> Trie h a b -> Trie h a b
go Word
flg (Trie h a b -> Trie h a b) -> Trie h a b -> Trie h a b
forall a b. (a -> b) -> a -> b
$ a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin (Word -> b -> a
u Word
w b
a) (Word -> b -> Trie h a b
forall (h :: Nat) a b. Word -> b -> Trie h a b
Tip Word
w b
a) Trie h a b
forall (h :: Nat) a b. Trie h a b
Nil
        | Bool
otherwise -> Word -> Trie h a b -> Trie h a b
go Word
flg (Trie h a b -> Trie h a b) -> Trie h a b -> Trie h a b
forall a b. (a -> b) -> a -> b
$ a -> Trie h a b -> Trie h a b -> Trie h a b
forall (h :: Nat) a b. a -> Trie h a b -> Trie h a b -> Trie h a b
Bin (Word -> b -> a
u Word
w b
a) Trie h a b
forall (h :: Nat) a b. Trie h a b
Nil (Word -> b -> Trie h a b
forall (h :: Nat) a b. Word -> b -> Trie h a b
Tip Word
w b
a)
      Trie h a b
Nil -> Trie h a b -> (b -> Trie h a b) -> Maybe b -> Trie h a b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Trie h a b
forall (h :: Nat) a b. Trie h a b
Nil (Word -> b -> Trie h a b
forall (h :: Nat) a b. Word -> b -> Trie h a b
Tip Word
v) (Maybe b -> Trie h a b) -> Maybe b -> Trie h a b
forall a b. (a -> b) -> a -> b
$ Maybe b -> Maybe b
f Maybe b
forall a. Maybe a
Nothing
{-# INLINE alterWith #-}

-- * Magma
class Magma a where
  magma :: a -> a -> a

newtype WrappedSemigroup m = WrapSemigroup m

instance (Semigroup m) => Magma (WrappedSemigroup m) where
  magma :: WrappedSemigroup m -> WrappedSemigroup m -> WrappedSemigroup m
magma = (m -> m -> m)
-> WrappedSemigroup m -> WrappedSemigroup m -> WrappedSemigroup m
forall a b. Coercible a b => a -> b
coerce (m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) :: m -> m -> m)
  {-# INLINE magma #-}