{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_HADDOCK show-extensions #-}
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)
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 :: String -> Trie -> Trie
insert :: String -> Trie -> Trie
insert [] (Trie _ m :: Map Char Trie
m) = Bool -> Map Char Trie -> Trie
Trie Bool
True Map Char Trie
m
insert (x :: Char
x:xs :: String
xs) (Trie b :: Bool
b m :: 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 (\x :: Char
x xs :: 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)
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
toList :: Trie -> [String]
toList :: Trie -> [String]
toList (Trie b :: Bool
b m :: Map Char Trie
m) =
if Bool
b then ""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
char, trie :: 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 ]
lookupPrefix :: (MonadPlus m) => String -> Trie -> m Trie
lookupPrefix :: String -> Trie -> m Trie
lookupPrefix [] trie :: Trie
trie = Trie -> m Trie
forall (m :: * -> *) a. Monad m => a -> m a
return Trie
trie
lookupPrefix (x :: Char
x:xs :: String
xs) (Trie _ m :: 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 :: Maybe a -> m a
liftMaybe Nothing = m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
liftMaybe (Just x :: a
x) = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
forcedNext :: Trie -> String
forcedNext :: Trie -> String
forcedNext (Trie _ m :: 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
== 1 then
let (char :: Char
char, trie :: 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
possibleSuffixes :: String -> Trie -> [String]
possibleSuffixes :: String -> Trie -> [String]
possibleSuffixes prefix :: String
prefix fulltrie :: 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
certainSuffix :: String -> Trie -> String
certainSuffix :: String -> Trie -> String
certainSuffix prefix :: String
prefix fulltrie :: 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