/src/MemoTrie.hs

 1{-# LANGUAGE UnicodeSyntax, MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, UndecidableInstances #-}
 2
 3{- Adapted from Conal Elliott's MemoTree package. ( http://haskell.org/haskellwiki/MemoTrie )
 4
 5We don't just use that package, because it is affected by this bug:
 6
 7  http://hackage.haskell.org/trac/ghc/ticket/2888
 8
 9This implementation works around that bug by using functional dependencies and multi-param type classes instead of type families.
10
11-}
12
13module MemoTrie (Trie(..), memo, PairTrie(..), BoolTrie(..)) where
14
15import qualified Data.Char as Char
16import qualified Data.Bits as Bits
17
18class Trie a t | a → t where
19  trie :: (a → b) → (t b)
20  untrie :: (t b) → (a → b)
21
22memo :: Trie b t ⇒ (b → a) → (b → a)
23memo = untrie . trie
24
25data UnitTrie a = UnitTrie a
26instance Trie () UnitTrie where
27  trie f = UnitTrie (f ())
28  untrie (UnitTrie a) = const a
29
30data BoolTrie a = BoolTrie a a
31instance Trie Bool BoolTrie where
32  trie f = BoolTrie (f False) (f True)
33  untrie (BoolTrie f t) b = if b then t else f
34
35data EitherTrie ta tb x = EitherTrie (ta x) (tb x)
36instance (Trie a ta, Trie b tb) ⇒ Trie (Either a b) (EitherTrie ta tb) where
37  trie f = EitherTrie (trie (f . Left)) (trie (f . Right))
38  untrie (EitherTrie s t) = either (untrie s) (untrie t)
39
40data PairTrie ta tb x = PairTrie (ta (tb x))
41instance (Trie a ta, Trie b tb) ⇒ Trie (a, b) (PairTrie ta tb) where
42  trie f = PairTrie (trie (trie . curry f))
43  untrie (PairTrie t) = uncurry (untrie . untrie t)
44
45data ListTrie tx x = ListTrie x (tx (ListTrie tx x))
46instance Trie x tx ⇒ Trie [x] (ListTrie tx) where
47  trie f = ListTrie (f []) (trie (\y → trie $ f . (y:)))
48  untrie (ListTrie p _) [] = p
49  untrie (ListTrie _ q) (h:t) = untrie (untrie q h) t
50
51bits :: (Num t, Bits.Bits t) ⇒ t → [Bool]
52bits 0 = []
53bits x = Bits.testBit x 0 : bits (Bits.shiftR x 1)
54
55unbit :: Num t ⇒ Bool → t
56unbit False = 0
57unbit True = 1
58
59unbits :: (Num t, Bits.Bits t) ⇒ [Bool] → t
60unbits [] = 0
61unbits (x:xs) = (Bits..|.) (unbit x) (Bits.shiftL (unbits xs) 1)
62
63instance Trie Char (ListTrie BoolTrie) where
64  trie f = trie (f . Char.chr . unbits)
65  untrie t = untrie t . bits . Char.ord
66
67{- Test:
68
69    f :: String → Integer
70    f = memo $ \l → case l of
71      [] → 1
72      (h:t) → (if Char.isAlpha h then 2 else 1) * slow_func t + slow_func t
73
74    out = slow_func $ take 21 $ cycle "a9d55"
75
76Replacing memo with id makes f much slower.
77
78-}