module Data.MemoTrie
( HasTrie(..)
, memo, memo2, memo3, mup
, trieBits, untrieBits
) where
import Data.Bits
import Data.Word
import Control.Applicative
import Data.Monoid
class HasTrie a where
data (:->:) a :: * -> *
trie :: (a -> b) -> (a :->: b)
untrie :: (a :->: b) -> (a -> b)
memo :: HasTrie t => (t -> a) -> (t -> a)
memo = untrie . trie
memo2 :: (HasTrie s,HasTrie t) => (s -> t -> a) -> (s -> t -> a)
memo3 :: (HasTrie r,HasTrie s,HasTrie t) => (r -> s -> t -> a) -> (r -> s -> t -> a)
mup :: HasTrie t => (b -> c) -> (t -> b) -> (t -> c)
mup mem f = memo (mem . f)
memo2 = mup memo
memo3 = mup memo2
instance HasTrie Bool where
data Bool :->: a = BoolTrie a a
trie f = BoolTrie (f False) (f True)
untrie (BoolTrie f _) False = f
untrie (BoolTrie _ t) True = t
instance HasTrie () where
data () :->: a = UnitTrie a
trie f = UnitTrie (f ())
untrie (UnitTrie x) () = x
instance (HasTrie a, HasTrie b) => HasTrie (Either a b) where
data (Either a b) :->: x = EitherTrie (a :->: x) (b :->: x)
untrie (EitherTrie f _) (Left x) = untrie f x
untrie (EitherTrie _ g) (Right y) = untrie g y
trie f = EitherTrie (trie (f . Left)) (trie (f . Right))
instance (HasTrie a, HasTrie b) => HasTrie (a,b) where
data (a,b) :->: x = PairTrie (a :->: (b :->: x))
untrie (PairTrie f) (a,b) = untrie (untrie f a) b
trie f = PairTrie $ trie $ \a -> trie $ \b -> f (a,b)
instance (HasTrie a, HasTrie b, HasTrie c) => HasTrie (a,b, c) where
data (a,b,c) :->: x = TripleTrie (a :->: (b :->: (c :->: x)))
untrie (TripleTrie f) (a,b,c) = untrie (untrie (untrie f a) b) c
trie f = TripleTrie $
trie $ \a -> trie $ \b -> trie $ \ c -> f (a,b,c)
instance HasTrie x => HasTrie [x] where
data [x] :->: a = ListTrie a (x :->: ([x] :->: a))
trie f = ListTrie (f []) $ trie (\x -> trie (f . (x:)))
untrie (ListTrie n _) [] = n
untrie (ListTrie _ t) (x:xs) = untrie (untrie t x) xs
bits :: Bits t => t -> [Bool]
bits 0 = []
bits x = testBit x 0 : bits (shiftR x 1)
unbit :: Num t => Bool -> t
unbit False = 0
unbit True = 1
unbits :: Bits t => [Bool] -> t
unbits [] = 0
unbits (x:xs) = unbit x .|. shiftL (unbits xs) 1
trieBits :: Bits t => (t -> a) -> ([Bool] :->: a)
trieBits f = trie (f . unbits)
untrieBits :: Bits t => ([Bool] :->: a) -> (t -> a)
untrieBits t x = untrie t (bits x)
instance HasTrie Word where
data Word :->: a = WordTrie ([Bool] :->: a)
untrie (WordTrie t) = untrieBits t
trie = WordTrie . trieBits
instance HasTrie Int where
data Int :->: a = IntTrie (Word :->: a)
untrie (IntTrie t) n = untrie t (fromIntegral n)
trie f = IntTrie (trie (f . fromIntegral . toInteger))
instance HasTrie a => Functor ((:->:) a) where
fmap f t = trie (fmap f (untrie t))
instance HasTrie a => Applicative ((:->:) a) where
pure b = trie (pure b)
tf <*> tx = trie (untrie tf <*> untrie tx)
instance (HasTrie a, Monoid b) => Monoid (a :->: b) where
mempty = trie mempty
s `mappend` t = trie (untrie s `mappend` untrie t)