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 (shift x (1))
| otherwise = applyPositive even (shift x (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 (\n -> shift n 1) go) (fmap (\n -> shift n 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 f odd)
| otherwise = BitTrie one (modifyPositive x f even) odd
overwrite :: (Ord b, Bits b) => b -> a -> IntTrie a -> IntTrie a
overwrite i x = modify i (const x)