module Trie where

import Text.PrettyPrint as Doc hiding ((<>))

data Trie = Node Membership [(Char, Trie)]
    deriving (Eq, Show)

data Membership = Member | NotMember
    deriving (Eq, Show)

-- If you prefer to convert Membership to Bool:
isMember :: Membership -> Bool
isMember Member = True
isMember NotMember = False

-- Some examples

albertTrie :: Trie
albertTrie = Node Member [('a', a), ('p', p), ('t', t)]
  where
    a = Node NotMember [('c', ac)]
    ac = Node NotMember [('e', ace)]
    ace = Node Member []
    p = Node NotMember [('i', pi)]
    pi = Node Member [('t', pit)]
    pit = Node Member []
    t = Node NotMember [('o', to)]
    to = Node NotMember [('n', ton), ('p', top)]
    top = Node Member []
    ton = Node Member []

-- bertieTrie is like albertTrie but the empty string is not a member.
bertieTrie :: Trie
bertieTrie = Node NotMember cs
  where
    Node _ cs = albertTrie

-- Nice formatting of a trie.

printTrie :: Trie -> IO ()
printTrie = putStrLn . prettyTrie

prettyTrie :: Trie -> String
prettyTrie = render . prettyRoot
  where
    prettyRoot (Node mval children) =
        text "\"\"" <> prettyVal mval
        $+$
        prettyChildren children

    prettyNonRoot (c, Node mval children) =
        char c <> prettyVal mval
        $+$
        prettyChildren children

    prettyVal NotMember = Doc.empty
    prettyVal Member = text (" member")

    prettyChildren = nest 2 . vcat . map prettyNonRoot