module Data.MemoTrie
( HasTrie(..)
, memo, memo2, memo3, mup
) 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 () where
data () :->: a = UnitTrie a
trie f = UnitTrie (f ())
untrie (UnitTrie x) () = x
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 a, HasTrie b) => HasTrie (Either a b) where
data (Either a b) :->: x = EitherTrie (a :->: x) (b :->: x)
untrie (EitherTrie f g) = either (untrie f) (untrie g)
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))
trie f = PairTrie $ trie $ \a -> trie $ \b -> f (a,b)
untrie (PairTrie t) = uncurry (untrie . untrie t)
trip :: ((a,b),c) -> (a,b,c)
trip ((a,b),c) = (a,b,c)
detrip :: (a,b,c) -> ((a,b),c)
detrip (a,b,c) = ((a,b),c)
instance (HasTrie a, HasTrie b, HasTrie c) => HasTrie (a,b,c) where
data (a,b,c) :->: x = TripleTrie (((a,b),c) :->: x)
trie f = TripleTrie (trie (f . trip))
untrie (TripleTrie t) = untrie t . detrip
list :: Either () (x,[x]) -> [x]
list = either (const []) (uncurry (:))
delist :: [x] -> Either () (x,[x])
delist [] = Left ()
delist (x:xs) = Right (x,xs)
instance HasTrie x => HasTrie [x] where
data [x] :->: a = ListTrie (Either () (x,[x]) :->: a)
trie f = ListTrie (trie (f . list))
untrie (ListTrie t) = untrie t . delist
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
instance HasTrie Word where
data Word :->: a = WordTrie ([Bool] :->: a)
trie f = WordTrie (trie (f . unbits))
untrie (WordTrie t) = untrie t . bits
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, Monoid b) => Monoid (a :->: b) where
mempty = trie mempty
s `mappend` t = trie (untrie s `mappend` untrie t)
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 => Monad ((:->:) a) where
return a = trie (return a)
u >>= k = trie (untrie u >>= untrie . k)