{-# 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
showList :: [Trie] -> ShowS
$cshowList :: [Trie] -> ShowS
show :: Trie -> String
$cshow :: Trie -> String
showsPrec :: Int -> Trie -> ShowS
$cshowsPrec :: Int -> Trie -> ShowS
Show, Trie -> Trie -> Bool
(Trie -> Trie -> Bool) -> (Trie -> Trie -> Bool) -> Eq Trie
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Trie -> Trie -> Bool
$c/= :: Trie -> Trie -> Bool
== :: Trie -> Trie -> Bool
$c== :: 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
$cto :: forall x. Rep Trie x -> Trie
$cfrom :: forall x. Trie -> Rep Trie x
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 (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 (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 (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 (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 (m :: * -> *) a. MonadPlus m => m a
mzero
liftMaybe (Just a
x) = 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 (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. [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 (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 (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Trie -> String
forcedNext

instance Binary Trie