module Data.Trie.Pseudo where
import Control.Applicative
import Control.Arrow (second)
import Control.Monad (replicateM)
import Data.Foldable hiding (all)
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty (..), fromList, toList)
import qualified Data.List.NonEmpty as NE
import Data.Maybe (fromMaybe)
import Data.Monoid
import qualified Data.Semigroup as S
import Prelude hiding (foldl, foldr, foldr1, lookup,
map)
data PseudoTrie t a = More !(t, Maybe a) !(NonEmpty (PseudoTrie t a))
| Rest !(NonEmpty t) a
| Nil
deriving (Show, Eq, Functor)
instance (Eq t) => Monoid (PseudoTrie t a) where
mempty = Nil
mappend = merge
instance Foldable (PseudoTrie t) where
foldr _ acc Nil = acc
foldr f acc (Rest _ x) = f x acc
foldr f acc (More (t, Nothing) xs) = foldr go acc xs
where
go z bcc = foldr f bcc z
foldr f acc (More (t, Just x) xs) = foldr go (f x acc) xs
where
go z bcc = foldr f bcc z
beginsWith :: (Eq t) => PseudoTrie t a -> t -> Bool
beginsWith Nil _ = False
beginsWith (Rest (t:|_) _) p = t == p
beginsWith (More (t,_) _) p = t == p
assign :: (Eq t) => NonEmpty t -> Maybe a -> PseudoTrie t a -> PseudoTrie t a
assign ts (Just x) Nil = Rest ts x
assign _ Nothing Nil = Nil
assign tss@(t:|ts) mx ys@(Rest pss@(p:|ps) y)
| tss == pss = case mx of
(Just x) -> Rest pss x
Nothing -> Nil
| t == p = case (ts,ps) of
([], p':_) -> More (t,mx) $ Rest (NE.fromList ps) y :| []
(t':_, []) -> case mx of
Just x -> More (p,Just y) $ Rest (NE.fromList ts) x :| []
Nothing -> ys
(t':_,p':_) -> if t' == p'
then More (t,Nothing) $
assign (NE.fromList ts) mx (Rest (NE.fromList ps) y) :| []
else case mx of
Nothing -> ys
Just x -> More (t,Nothing) $ NE.fromList $
[ Rest (NE.fromList ps) y
, Rest (NE.fromList ts) x
]
| otherwise = ys
assign (t:|ts) mx y@(More (p,my) ys)
| t == p = case ts of
[] -> More (p,mx) ys
_ -> More (p,my) $ fmap (assign (NE.fromList ts) mx) ys
| otherwise = y
merge :: (Eq t) => PseudoTrie t a -> PseudoTrie t a -> PseudoTrie t a
merge Nil y = y
merge x Nil = x
merge xx@(Rest tss@(t:|ts) x) (Rest pss@(p:|ps) y)
| tss == pss = Rest pss y
| t == p = case (ts,ps) of
([],p':ps') -> More (t,Just x) $ Rest (NE.fromList ps) y :| []
(t':ts',[]) -> More (t,Just y) $ Rest (NE.fromList ts) x :| []
(_,_) -> More (t,Nothing) $
merge (Rest (NE.fromList ts) x)
(Rest (NE.fromList ps) y) :| []
| otherwise = xx
merge xx@(More (t,mx) xs) (More (p,my) ys)
| t == p = More (p,my) $ NE.fromList $
foldr go [] $ NE.toList xs ++ NE.toList ys
| otherwise = xx
where
go q [] = [q]
go q (z:zs) | areDisjoint q z = q : z : zs
| otherwise = merge q z : zs
merge xx@(More (t,mx) xs) (Rest pss@(p:|ps) y)
| t == p = case ps of
[] -> More (t,Just y) xs
_ -> More (t,mx) $
fmap (flip merge $ Rest (NE.fromList ps) y) xs
| otherwise = xx
merge xx@(Rest tss@(t:|ts) x) (More (p,my) ys)
| t == p = case ts of
[] -> More (p,Just x) ys
_ -> More (p,my) $
fmap (merge $ Rest (NE.fromList ts) x) ys
| otherwise = xx
add :: (Eq t) => NonEmpty t -> PseudoTrie t a -> PseudoTrie t a -> PseudoTrie t a
add ts input container =
let ts' = NE.toList ts in
merge container $ mkMores ts' input
where
mkMores :: (Eq t) => [t] -> PseudoTrie t a -> PseudoTrie t a
mkMores [] trie = trie
mkMores (t:ts) trie = More (t,Nothing) $
mkMores ts trie :| []
toAssocs :: PseudoTrie t a -> [(NonEmpty t, a)]
toAssocs = go [] []
where
go :: [t] -> [(NonEmpty t, a)] -> PseudoTrie t a -> [(NonEmpty t, a)]
go depth acc Nil = acc
go depth acc (Rest ts x) = (NE.fromList $ depth ++ NE.toList ts, x) : acc
go depth acc (More (t, Nothing) xs) =
foldr (flip $ go $ depth ++ [t]) acc $ NE.toList xs
go depth acc (More (t, Just x) xs) =
(NE.fromList $ depth ++ [t], x) :
(foldr $ flip $ go $ depth ++ [t]) acc (NE.toList xs)
fromAssocs :: (Eq t) => [(NonEmpty t, a)] -> PseudoTrie t a
fromAssocs = foldr (uncurry assign) Nil . fmap (second Just)
lookup :: (Eq t) => NonEmpty t -> PseudoTrie t a -> Maybe a
lookup _ Nil = Nothing
lookup tss (Rest pss a)
| tss == pss = Just a
| otherwise = Nothing
lookup tss@(t:|ts) (More (p,mx) xs)
| t == p = case ts of
[] -> mx
(t':_) -> find (hasNextTag t') xs >>= lookup (fromList ts)
| otherwise = Nothing
where
hasNextTag :: (Eq t) => t -> PseudoTrie t a -> Bool
hasNextTag t Nil = False
hasNextTag t (More (p,_) _) = t == p
hasNextTag t (Rest (p:|_) _) = t == p
lookupNearestParent :: (Eq t) => [t] -> PseudoTrie t a -> Maybe a
lookupNearestParent [] (More (_,mx) _) = mx
lookupNearestParent [] (Rest _ x) = Just x
lookupNearestParent _ Nil = Nothing
lookupNearestParent tss@(t:ts) (Rest pss@(p:|ps) a)
| (NE.fromList tss) == pss || (ps == [] && ts /= []) = Just a
| t == p = case ts of
[] -> if ps == [] then Just a
else Nothing
_ -> lookupNearestParent ts $ Rest (NE.fromList ps) a
| otherwise = Nothing
lookupNearestParent (t:ts) (More (p,mx) xs)
| t == p = case ts of
[] -> mx
(t':_) -> find (hasNextTag t') xs >>= lookupNearestParent ts
where
hasNextTag :: (Eq t) => t -> PseudoTrie t a -> Bool
hasNextTag t Nil = False
hasNextTag t (More (p,_) _) = t == p
hasNextTag t (Rest (p:|_) _) = t == p
areDisjoint :: (Eq t) => PseudoTrie t a -> PseudoTrie t a -> Bool
areDisjoint (More (t,_) _) (More (p,_) _)
| t == p = False
| otherwise = True
areDisjoint (Rest (t:|_) _) (Rest (p:|_) _)
| t == p = False
| otherwise = True
areDisjoint _ _ = True
intersectionWith :: (Eq t) =>
(a -> b -> c)
-> PseudoTrie t a
-> PseudoTrie t b
-> PseudoTrie t c
intersectionWith _ _ Nil = Nil
intersectionWith _ Nil _ = Nil
intersectionWith f (Rest tss@(t:|ts) x) (Rest pss@(p:|ps) y)
| tss == pss = Rest pss $ f x y
| otherwise = Nil
intersectionWith f (More (t,mx) xs) (More (p,my) ys)
| t == p = case [intersectionWith f x' y' | x' <- NE.toList xs, y' <- NE.toList ys] of
[] -> case f <$> mx <*> my of
Nothing -> Nil
Just c -> Rest (p :| []) c
zs -> More (p,f <$> mx <*> my) $ NE.fromList zs
| otherwise = Nil
intersectionWith f (More (t,mx) xs) (Rest pss@(p:|ps) y)
| t == p = case ps of
[] -> case f <$> mx <*> Just y of
Nothing -> Nil
Just c -> Rest (p :| []) c
_ -> More (p,Nothing) $ fmap (flip (intersectionWith f) $ Rest (fromList ps) y) xs
| otherwise = Nil
intersectionWith f (Rest tss@(t:|ts) x) (More (p,my) ys)
| t == p = case ts of
[] -> case f <$> Just x <*> my of
Nothing -> Nil
Just c -> Rest (t :| []) c
_ -> More (t,Nothing) $ fmap (intersectionWith f $ Rest (fromList ts) x) ys
| otherwise = Nil
prune :: PseudoTrie t a -> PseudoTrie t a
prune = go
where
go Nil = Nil
go xx@(Rest ts x) = xx
go (More (t,Nothing) xs) =
case cleaned xs of
[Nil] -> Nil
[Rest ts x] -> Rest (t:|NE.toList ts) x
xs' -> More (t,Nothing) $ NE.fromList xs'
go (More (t,Just x) xs) =
case cleaned xs of
[Nil] -> Rest (t:|[]) x
xs' -> More (t,Just x) $ NE.fromList xs'
cleaned xs = removeNils (NE.toList $ fmap go xs)
removeNils xs = case removeNils' xs of
[] -> [Nil]
ys -> ys
where
removeNils' [] = []
removeNils' (Nil:xs) = removeNils' xs
removeNils' (x:xs) = x : removeNils' xs