{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
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)
-- TODO: difference
-- | Non-Empty Rose Tree with explicit emptyness
data PseudoTrie t a = More !(t, Maybe a) !(NonEmpty (PseudoTrie t a))
| Rest !(NonEmpty t) a
| Nil
deriving (Show, Eq, Functor)
-- | Overwriting instance
instance (Eq t) => Monoid (PseudoTrie t a) where
mempty = Nil
mappend = merge
-- | Depth first
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
-- | Provides a form of deletion by setting a path to @Nothing@, but doesn't
-- cleanup like @prune@
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 -- disjoint
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
-- | Overwrite the LHS point-wise with the RHS's contents
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
-- | Simple test on the heads of two tries
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
-- | The meet of two @PseudoTrie@s
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
-- implicit root
| 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
-- difference :: Eq t =>
-- PseudoTrie t a
-- -> PseudoTrie t a
-- -> PseudoTrie t a
-- | Needless intermediary elements are turned into shortcuts, @Nil@'s in
-- subtrees are also removed.
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