0 | module Libraries.Data.StringTrie
 1 |
 2 | import Data.These
 3 | import Libraries.Data.StringMap
 4 |
 5 | %default total
 6 |
 7 | -- prefix tree specialised to use `String`s as keys
 8 |
 9 | public export
10 | record StringTrie a where
11 |   constructor MkStringTrie
12 |   node : These a (StringMap (StringTrie a))
13 |
14 | public export
15 | Show a => Show (StringTrie a) where
16 |   show (MkStringTrie node) = assert_total $ show node
17 |
18 | public export
19 | Functor StringTrie where
20 |   map f (MkStringTrie node) = MkStringTrie $ assert_total $ bimap f (map (map f)) node
21 |
22 | public export
23 | empty : StringTrie a
24 | empty = MkStringTrie $ That empty
25 |
26 | public export
27 | singleton : List String -> a -> StringTrie a
28 | singleton []      v = MkStringTrie $ This v
29 | singleton (k::ks) v = MkStringTrie $ That $ singleton k (singleton ks v)
30 |
31 | -- insert using supplied function to resolve clashes
32 | public export
33 | insertWith : List String -> (Maybe a -> a) -> StringTrie a -> StringTrie a
34 | insertWith []      f (MkStringTrie nd) =
35 |   MkStringTrie $ these (This . f . Just) (Both (f Nothing)) (Both . f . Just) nd
36 | insertWith (k::ks) f (MkStringTrie nd) =
37 |   MkStringTrie $ these (\x => Both x (singleton k end)) (That . rec) (\x => Both x . rec) nd
38 |   where
39 |   end : StringTrie a
40 |   end = singleton ks (f Nothing)
41 |   rec : StringMap (StringTrie a) -> StringMap (StringTrie a)
42 |   rec sm = maybe (insert k end sm) (\tm => insert k (insertWith ks f tm) sm) (lookup k sm)
43 |
44 | public export
45 | insert : List String -> a -> StringTrie a -> StringTrie a
46 | insert ks v = insertWith ks (const v)
47 |
48 | -- fold the trie in a depth-first fashion performing monadic actions on values, then keys
49 | -- note that for `Both` the action on node values will be performed first because of `bitraverse` implementation
50 | public export
51 | foldWithKeysM : (Monad m, Monoid b) => (List String -> m b) -> (List String -> a -> m b) -> StringTrie a -> m b
52 | foldWithKeysM {a} {m} {b} fk fv = go []
53 |   where
54 |   go : List String -> StringTrie a -> m b
55 |   go ks (MkStringTrie nd) =
56 |     bifold <$> bitraverse
57 |                 (fv ks)
58 |                 (\sm => foldlM
59 |                           (\x, (k, vs) => do let ks' = ks++[k]
60 |                                              y <- assert_total $ go ks' vs
61 |                                              z <- fk ks'
62 |                                              pure $ x <+> y <+> z)
63 |                           neutral
64 |                           (StringMap.toList sm))
65 |                 nd
66 |