module Data.Trie.Pred.Base where
import Prelude hiding (lookup)
import Data.Trie.Pred.Base.Step (PredStep (..), Pred (..))
import Data.Trie.Class (Trie (..))
import qualified Data.Trie.HashMap as HT
import qualified Data.HashMap.Lazy as HM
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Data (Typeable)
import Data.Monoid (First (..), Last (..), (<>))
import Data.Maybe (fromMaybe)
import Data.Hashable (Hashable)
import qualified Data.HashMap.Strict as HMS
import Test.QuickCheck (Arbitrary (..))
import Control.DeepSeq (NFData (..))
data PredTrie k a = PredTrie
{ predLits :: !(HT.HashMapStep PredTrie k a)
, predPreds :: PredStep k PredTrie k a
} deriving (Show, Functor, Typeable)
instance (NFData k, NFData a) => NFData (PredTrie k a) where
rnf (PredTrie hs ps) = rnf hs `seq` rnf ps
instance ( Arbitrary k
, Arbitrary a
, Eq k
, Hashable k
) => Arbitrary (PredTrie k a) where
arbitrary = flip PredTrie mempty <$> arbitrary
instance ( Hashable k
, Eq k
) => Trie NonEmpty k PredTrie where
lookup ts (PredTrie ls ps) =
getFirst (First (lookup ts ls) <> First (lookup ts ps))
delete ts (PredTrie ls ps) = PredTrie (delete ts ls) (delete ts ps)
insert ts x (PredTrie ls ps) = PredTrie (HT.insert ts x ls) ps
instance ( Hashable k
, Eq k
) => Monoid (PredTrie k a) where
mempty = PredTrie mempty mempty
mappend (PredTrie ls1 ps1) (PredTrie ls2 ps2) =
PredTrie (ls1 <> ls2) (ps1 <> ps2)
emptyPT :: PredTrie k a
emptyPT = PredTrie HT.empty (PredStep HMS.empty)
matchPT :: ( Hashable k
, Eq k
) => NonEmpty k -> PredTrie k a -> Maybe (NonEmpty k, a, [k])
matchPT (t:|ts) (PredTrie ls (PredStep ps)) = getFirst $
First (goLit ls) <> foldMap (First . goPred) ps
where
goLit (HT.HashMapStep xs) = do
(HT.HashMapChildren mx mxs) <- HM.lookup t xs
let mFoundHere = (t:|[],,[]) <$> mx
if null ts
then mFoundHere
else getFirst $ First (do (pre,y,suff) <- matchPT (NE.fromList ts) =<< mxs
pure (NE.cons t pre, y, suff))
<> First mFoundHere
goPred (Pred p mx xs) = do
r <- p t
let mFoundHere = do x <- ($ r) <$> mx
pure (t:|[], x, [])
if null ts
then mFoundHere
else getFirst $ First (do (pre,y,suff) <- matchPT (NE.fromList ts) xs
pure (NE.cons t pre, y r, suff))
<> First mFoundHere
matchesPT :: ( Hashable k
, Eq k
) => NonEmpty k -> PredTrie k a -> [(NonEmpty k, a, [k])]
matchesPT (t:|ts) (PredTrie ls (PredStep ps)) =
fromMaybe [] (getFirst (First (goLit ls) <> foldMap (First . goPred) ps))
where
goLit (HT.HashMapStep xs) = do
(HT.HashMapChildren mx mxs) <- HM.lookup t xs
let mFoundHere = do x <- mx
pure [(t:|[],x,ts)]
prependAncestry (pre,x,suff) = (NE.cons t pre,x,suff)
if null ts
then mFoundHere
else do foundHere <- mFoundHere
let rs = fromMaybe [] (matchesPT (NE.fromList ts) <$> mxs)
pure (foundHere ++ (prependAncestry <$> rs))
goPred (Pred p mx xs) = do
r <- p t
let mFoundHere = do x <- ($ r) <$> mx
pure [(t:|[],x,ts)]
prependAncestryAndApply (pre,x,suff) = (NE.cons t pre,x r,suff)
if null ts
then mFoundHere
else do foundHere <- mFoundHere
let rs = matchesPT (NE.fromList ts) xs
pure (foundHere ++ (prependAncestryAndApply <$> rs))
data RootedPredTrie k a = RootedPredTrie
{ rootedBase :: !(Maybe a)
, rootedSub :: !(PredTrie k a)
} deriving (Show, Functor, Typeable)
instance ( Hashable k
, Eq k
) => Trie [] k RootedPredTrie where
lookup [] (RootedPredTrie mx _) = mx
lookup ts (RootedPredTrie _ xs) = lookup (NE.fromList ts) xs
delete [] (RootedPredTrie _ xs) = RootedPredTrie Nothing xs
delete ts (RootedPredTrie mx xs) = RootedPredTrie mx (delete (NE.fromList ts) xs)
insert [] x (RootedPredTrie _ xs) = RootedPredTrie (Just x) xs
insert ts x (RootedPredTrie mx xs) = RootedPredTrie mx (insert (NE.fromList ts) x xs)
instance ( Hashable k
, Eq k
) => Monoid (RootedPredTrie k a) where
mempty = emptyRPT
mappend (RootedPredTrie mx xs) (RootedPredTrie my ys) = RootedPredTrie
(getLast (Last mx <> Last my)) (xs <> ys)
emptyRPT :: RootedPredTrie k a
emptyRPT = RootedPredTrie Nothing emptyPT
matchRPT :: ( Hashable k
, Eq k
) => [k] -> RootedPredTrie k a -> Maybe ([k], a, [k])
matchRPT [] (RootedPredTrie mx _) = ([],,[]) <$> mx
matchRPT ts (RootedPredTrie mx xs) = getFirst $
First mFoundThere <> First (([],,[]) <$> mx)
where
mFoundThere = do (pre,x,suff) <- matchPT (NE.fromList ts) xs
pure (NE.toList pre,x,suff)
matchesRPT :: ( Hashable k
, Eq k
) => [k] -> RootedPredTrie k a -> [([k], a, [k])]
matchesRPT [] (RootedPredTrie mx _) = fromMaybe [] $ (\x -> [([],x,[])]) <$> mx
matchesRPT ts (RootedPredTrie mx xs) =
foundHere ++ fmap allowRoot (matchesPT (NE.fromList ts) xs)
where
foundHere = fromMaybe [] $ do x <- mx
pure [([],x,[])]
allowRoot (pre,x,suff) = (NE.toList pre,x,suff)