{-# LANGUAGE CPP #-} {-# LANGUAGE StrictData #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE DeriveTraversable #-} {-# LANGUAGE DeriveLift #-} module Text.Collate.Trie ( Trie , empty , insert , alter , unfoldTrie , matchLongestPrefix ) where import qualified Data.IntMap as M import Data.Bifunctor (first) import Data.Binary (Binary(..)) import Language.Haskell.TH.Syntax (Lift(..)) import Instances.TH.Lift () import Data.Maybe (fromMaybe) #if MIN_VERSION_base(4,11,0) #else import Data.Semigroup (Semigroup(..)) #endif -- import Debug.Trace data Trie a = Trie (Maybe a) (Maybe (M.IntMap (Trie a))) deriving (Show, Eq, Ord, Lift, Functor, Foldable, Traversable) instance Semigroup (Trie a) where trie1 <> trie2 = foldr (uncurry insert) trie1 (unfoldTrie trie2) instance Monoid (Trie a) where mempty = Trie Nothing Nothing mappend = (<>) instance Binary a => Binary (Trie a) where put (Trie mbv mbm) = put (mbv, mbm) get = do (mbv,mbm) <- get return $ Trie mbv mbm empty :: Trie a empty = Trie Nothing Nothing unfoldTrie :: Trie a -> [([Int], a)] unfoldTrie = map (first reverse) . go [] where go xs (Trie (Just v) (Just m)) = (xs, v) : concatMap (gopair xs) (M.toList m) go xs (Trie (Just v) Nothing) = [(xs, v)] go xs (Trie Nothing (Just m)) = concatMap (gopair xs) (M.toList m) go _ (Trie Nothing Nothing) = [] gopair xs (i, trie) = go (i:xs) trie insert :: [Int] -> a -> Trie a -> Trie a insert [] x (Trie _ mbm) = Trie (Just x) mbm insert (c:cs) x (Trie mbv (Just m)) = case M.lookup c m of Nothing -> Trie mbv (Just (M.insert c (insert cs x empty) m)) Just trie -> Trie mbv (Just (M.insert c (insert cs x trie) m)) insert (c:cs) x (Trie mbv Nothing) = Trie mbv (Just (M.insert c (insert cs x empty) mempty)) alter :: (Maybe a -> Maybe a) -> [Int] -> Trie a -> Trie a alter f [] (Trie mbv mbm) = Trie (f mbv) mbm alter f (c:cs) (Trie mbv (Just m)) = Trie mbv (Just (M.insert c (alter f cs $ fromMaybe empty $ M.lookup c m) m)) alter f (c:cs) (Trie mbv Nothing) = Trie mbv (Just (M.insert c (alter f cs empty) mempty)) matchLongestPrefix :: Trie a -> [Int] -> Maybe (a, [Int], Trie a) matchLongestPrefix = go Nothing where go _ (Trie (Just x) mbm) [] = Just (x, [], Trie Nothing mbm) go best (Trie Nothing _) [] = best go best (Trie mbv mbm) (c:cs) = case mbm >>= M.lookup c of Nothing -> case mbv of Nothing -> best Just x -> Just (x, c:cs, Trie Nothing mbm) Just trie -> go (case mbv of Nothing -> best Just x -> Just (x, c:cs, trie)) trie cs