    {-# LANGUAGE CPP #-}
    {-# LANGUAGE DeriveGeneric #-}
    {-# LANGUAGE TemplateHaskell #-}
    {-# OPTIONS_HADDOCK show-extensions #-}

    -- |
    -- Module      :  Data.Trie
    -- License     :  GPL-2
    -- Maintainer  :  fuuzetsu@fuuzetsu.co.uk
    -- Stability   :  experimental
    -- Portability :  portable
    --
    -- An implementation of a trie over a words. Properties:
    --
    -- @
    -- 'fromList' . 'toList' ≡ 'id'
    -- 'toList' . 'fromString' ≡ (:[])
    -- 'sort' . 'nub' . 'toList' . 'fromList' ≡ 'sort' . 'nub'
    -- @

    module Data.Trie ( empty, insert, fromString, fromList
                     , toList, lookupPrefix, forcedNext, Trie
                     , possibleSuffixes, certainSuffix
                     ) where

    import           Control.Monad
    import           Data.Binary
    import qualified Data.Map as Map
    import           GHC.Generics (Generic)

    data Trie = Trie Bool (Map.Map Char Trie) deriving (Int -> Trie -> ShowS
    [Trie] -> ShowS
    Trie -> String
    (Int -> Trie -> ShowS)
    -> (Trie -> String) -> ([Trie] -> ShowS) -> Show Trie
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Trie -> ShowS
    showsPrec :: Int -> Trie -> ShowS
    $cshow :: Trie -> String
    show :: Trie -> String
    $cshowList :: [Trie] -> ShowS
    showList :: [Trie] -> ShowS
    Show, Trie -> Trie -> Bool
    (Trie -> Trie -> Bool) -> (Trie -> Trie -> Bool) -> Eq Trie
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Trie -> Trie -> Bool
    == :: Trie -> Trie -> Bool
    $c/= :: Trie -> Trie -> Bool
    /= :: Trie -> Trie -> Bool
    Eq, (forall x. Trie -> Rep Trie x)
    -> (forall x. Rep Trie x -> Trie) -> Generic Trie
    forall x. Rep Trie x -> Trie
    forall x. Trie -> Rep Trie x
    forall a.
    (forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
    $cfrom :: forall x. Trie -> Rep Trie x
    from :: forall x. Trie -> Rep Trie x
    $cto :: forall x. Rep Trie x -> Trie
    to :: forall x. Rep Trie x -> Trie
    Generic)

    -- | A blank Trie
    empty :: Trie
    empty :: Trie
    empty = Bool -> Map Char Trie -> Trie
    Trie Bool
    False Map Char Trie
    forall k a. Map k a
    Map.empty

    -- | Insert a new string into the trie.
    insert :: String -> Trie -> Trie
    insert :: String -> Trie -> Trie
    insert []     (Trie Bool
    _ Map Char Trie
    m) = Bool -> Map Char Trie -> Trie
    Trie Bool
    True Map Char Trie
    m
    insert (Char
    x:String
    xs) (Trie Bool
    b Map Char Trie
    m) =
      Bool -> Map Char Trie -> Trie
    Trie Bool
    b (Map Char Trie -> Trie) -> Map Char Trie -> Trie
    forall a b. (a -> b) -> a -> b
    $ (Maybe Trie -> Maybe Trie)
    -> Char -> Map Char Trie -> Map Char Trie
    forall k a.
    Ord k =>
    (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
    Map.alter (Trie -> Maybe Trie
    forall a. a -> Maybe a
    Just (Trie -> Maybe Trie)
    -> (Maybe Trie -> Trie) -> Maybe Trie -> Maybe Trie
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Trie -> (Trie -> Trie) -> Maybe Trie -> Trie
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe (String -> Trie
    fromString String
    xs) (String -> Trie -> Trie
    insert String
    xs)) Char
    x Map Char Trie
    m

    fromString :: String -> Trie
    fromString :: String -> Trie
    fromString =
      (Char -> Trie -> Trie) -> Trie -> String -> Trie
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr (\Char
    x Trie
    xs -> Bool -> Map Char Trie -> Trie
    Trie Bool
    False (Char -> Trie -> Map Char Trie
    forall k a. k -> a -> Map k a
    Map.singleton Char
    x Trie
    xs)) (Bool -> Map Char Trie -> Trie
    Trie Bool
    True Map Char Trie
    forall k a. Map k a
    Map.empty)

    -- | Take a list of String and compress it into a Trie
    fromList :: [String] -> Trie
    fromList :: [String] -> Trie
    fromList = (String -> Trie -> Trie) -> Trie -> [String] -> Trie
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr String -> Trie -> Trie
    insert Trie
    empty

    -- | Take a trie and expand it into the strings that it represents
    toList :: Trie -> [String]
    toList :: Trie -> [String]
    toList (Trie Bool
    b Map Char Trie
    m) =
        if Bool
    b then String
    ""String -> [String] -> [String]
    forall a. a -> [a] -> [a]
    :[String]
    expand
        else [String]
    expand
        where expand :: [String]
    expand = [ Char
    charChar -> ShowS
    forall a. a -> [a] -> [a]
    :String
    word | (Char
    char, Trie
    trie) <- Map Char Trie -> [(Char, Trie)]
    forall k a. Map k a -> [(k, a)]
    Map.toList Map Char Trie
    m,
                                     String
    word <- Trie -> [String]
    toList Trie
    trie ]

    -- | Takes a trie and a prefix and returns the sub-trie that
    -- of words with that prefix
    lookupPrefix :: (MonadPlus m) => String -> Trie -> m Trie
    lookupPrefix :: forall (m :: * -> *). MonadPlus m => String -> Trie -> m Trie
    lookupPrefix [] Trie
    trie = Trie -> m Trie
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Trie
    trie
    lookupPrefix (Char
    x:String
    xs) (Trie Bool
    _ Map Char Trie
    m) = Maybe Trie -> m Trie
    forall (m :: * -> *) a. MonadPlus m => Maybe a -> m a
    liftMaybe (Char -> Map Char Trie -> Maybe Trie
    forall k a. Ord k => k -> Map k a -> Maybe a
    Map.lookup Char
    x Map Char Trie
    m) m Trie -> (Trie -> m Trie) -> m Trie
    forall a b. m a -> (a -> m b) -> m b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= String -> Trie -> m Trie
    forall (m :: * -> *). MonadPlus m => String -> Trie -> m Trie
    lookupPrefix String
    xs

    liftMaybe :: MonadPlus m => Maybe a -> m a
    liftMaybe :: forall (m :: * -> *) a. MonadPlus m => Maybe a -> m a
    liftMaybe Maybe a
    Nothing = m a
    forall a. m a
    forall (m :: * -> *) a. MonadPlus m => m a
    mzero
    liftMaybe (Just a
    x) = a -> m a
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    x

    -- | Finds the longest certain path down the trie starting at a the root
    -- Invariant Assumption: All paths have at least one 'true' node below them
    forcedNext :: Trie -> String
    forcedNext :: Trie -> String
    forcedNext (Trie Bool
    _ Map Char Trie
    m) =
        if [(Char, Trie)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(Char, Trie)]
    ls Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    1 then
            let (Char
    char, Trie
    trie) = [(Char, Trie)] -> (Char, Trie)
    forall a. HasCallStack => [a] -> a
    head [(Char, Trie)]
    ls in
            Char
    charChar -> ShowS
    forall a. a -> [a] -> [a]
    :Trie -> String
    forcedNext Trie
    trie
        else []
        where ls :: [(Char, Trie)]
    ls = Map Char Trie -> [(Char, Trie)]
    forall k a. Map k a -> [(k, a)]
    Map.toList Map Char Trie
    m

    -- | Helper function, finds all the suffixes of a given prefix
    possibleSuffixes :: String -> Trie -> [String]
    possibleSuffixes :: String -> Trie -> [String]
    possibleSuffixes String
    prefix Trie
    fulltrie =
        String -> Trie -> [Trie]
    forall (m :: * -> *). MonadPlus m => String -> Trie -> m Trie
    lookupPrefix String
    prefix Trie
    fulltrie [Trie] -> (Trie -> [String]) -> [String]
    forall a b. [a] -> (a -> [b]) -> [b]
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Trie -> [String]
    toList

    -- | Helper function, finds the longest certain path down the trie
    -- starting at a given word
    certainSuffix :: String -> Trie -> String
    certainSuffix :: String -> Trie -> String
    certainSuffix String
    prefix Trie
    fulltrie =
        String -> Trie -> [Trie]
    forall (m :: * -> *). MonadPlus m => String -> Trie -> m Trie
    lookupPrefix String
    prefix Trie
    fulltrie [Trie] -> (Trie -> String) -> String
    forall a b. [a] -> (a -> [b]) -> [b]
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Trie -> String
    forcedNext

    instance Binary Trie