module Data.IntTrie
( IntTrie, identity, apply, modify, overwrite )
where
import Control.Applicative
import Data.Bits
import Data.Function (fix)
data IntTrie a = IntTrie (BitTrie a) a (BitTrie a)
data BitTrie a = BitTrie a (BitTrie a) (BitTrie a)
instance Functor BitTrie where
fmap f ~(BitTrie x l r) = BitTrie (f x) (fmap f l) (fmap f r)
instance Applicative BitTrie where
pure x = fix (\g -> BitTrie x g g)
~(BitTrie f fl fr) <*> ~(BitTrie x xl xr) = BitTrie (f x) (fl <*> xl) (fr <*> xr)
instance Functor IntTrie where
fmap f ~(IntTrie neg z pos) = IntTrie (fmap f neg) (f z) (fmap f pos)
instance Applicative IntTrie where
pure x = IntTrie (pure x) x (pure x)
IntTrie fneg fz fpos <*> IntTrie xneg xz xpos =
IntTrie (fneg <*> xneg) (fz xz) (fpos <*> xpos)
apply :: (Ord b, Bits b) => IntTrie a -> b -> a
apply ~(IntTrie neg z pos) x =
case compare x 0 of
LT -> applyPositive neg (x)
EQ -> z
GT -> applyPositive pos x
applyPositive :: (Bits b) => BitTrie a -> b -> a
applyPositive ~(BitTrie one even odd) x
| x == 1 = one
| testBit x 0 = applyPositive odd (x `shiftR` 1)
| otherwise = applyPositive even (x `shiftR` 1)
identity :: (Bits a) => IntTrie a
identity = IntTrie (fmap negate identityPositive) 0 identityPositive
identityPositive :: (Bits a) => BitTrie a
identityPositive = go
where
go = BitTrie 1 (fmap (`shiftL` 1) go) (fmap (\n -> (n `shiftL` 1) .|. 1) go)
modify :: (Ord b, Bits b) => b -> (a -> a) -> IntTrie a -> IntTrie a
modify x f ~(IntTrie neg z pos) =
case compare x 0 of
LT -> IntTrie (modifyPositive (x) f neg) z pos
EQ -> IntTrie neg (f z) pos
GT -> IntTrie neg z (modifyPositive x f pos)
modifyPositive :: (Bits b) => b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive x f ~(BitTrie one even odd)
| x == 1 = BitTrie (f one) even odd
| testBit x 0 = BitTrie one even (modifyPositive (x `shiftR` 1) f odd)
| otherwise = BitTrie one (modifyPositive (x `shiftR` 1) f even) odd
overwrite :: (Ord b, Bits b) => b -> a -> IntTrie a -> IntTrie a
overwrite i x = modify i (const x)