module Data.Queue.TrieQueue where
import Data.Queue.Class
import Data.Queue.QueueHelpers
import Control.Arrow((***))
import Control.Monad(liftM2)
import Data.Function
import Data.Maybe
import Data.List (sortBy, groupBy)
import Data.Monoid
import Data.Ord
import Data.Sequence (Seq, viewl, ViewL(..), (><), (|>), (<|))
import qualified Data.Sequence as Seq
import Data.Map (Map, findMin, minViewWithKey, fromDistinctAscList)
import qualified Data.Map as Map
import qualified Data.Foldable as Fold (toList)
import GHC.Exts
type Label e = Seq e
data Trie e = Leaf (Label e) !Int
| Edge (Label e) !Int (Map e (Trie e))
deriving (Eq)
type MTrie e = Maybe (Trie e)
newtype TrieQueue e = TQ (HeapQ (Trie e)) deriving (Eq, Show, Monoid)
mkTQ :: Int -> Trie e -> TrieQueue e
mkTQ n t = TQ (HQ n (Just t))
mkLab :: [e] -> Label e
mkLab = Seq.fromList
instance Show e => Show (Trie e) where
show (Leaf xs xn) = "(" ++ show xn ++ "x" ++ show (Fold.toList xs) ++ ")"
show (Edge xs xn m) = "==" ++ show xn ++ "x" ++ show (Fold.toList xs) ++ "==>" ++ show m
instance Ord e => Monoid (Trie e) where
mappend = merger
instance Ord e => IQueue (TrieQueue e) where
type QueueKey (TrieQueue e) = [e]
empty = mempty
merge = mappend
mergeAll = mconcat
singleton xs = mkTQ 1 $ Leaf (mkLab xs) 1
extract (TQ (HQ n t)) = fmap ((Fold.toList *** (TQ . HQ (n1))) . extractMin') t
size (TQ (HQ n _)) = n
catTrie :: Ord e => Label e -> Trie e -> Trie e
catTrie xs (Leaf ys n) = Leaf (xs >< ys) n
catTrie xs (Edge ys n m) = compactTrie (Edge (xs >< ys) n m)
compactTrie :: Ord e => Trie e -> Trie e
compactTrie = fromJust . compactTrie'
compactTrie' :: Ord e => Trie e -> MTrie e
compactTrie' t@(Edge xs 0 m)
| Map.null m = Nothing
| Map.size m == 1, (y, yT) <- findMin m
= Just $ catTrie (xs |> y) yT
| otherwise = Just t
compactTrie' (Edge xs n m)
| Map.null m = Just $ Leaf xs n
compactTrie' (Leaf _ 0) = Nothing
compactTrie' t = Just t
extractMin' :: Ord e => Trie e -> (Label e, MTrie e)
extractMin' (Leaf xs n) = (xs, compactTrie' (Leaf xs (n1)))
extractMin' (Edge xs (n+1) m) = (xs, compactTrie' (Edge xs n m))
extractMin' (Edge xs 0 m)
| Just ((y, yT), m') <- minViewWithKey m,
(ys, yT') <- extractMin' yT
= (xs >< (y <| ys), maybe (compactTrie' (Edge xs 0 m')) (\ yT' -> Just $ Edge xs 0 $ Map.insert y yT' m') yT')
extractMin' _ = error "Internal failure to note empty queue"
type TailMaker e = Label e -> Trie e
merge' :: Ord e => Label e -> Label e -> TailMaker e -> TailMaker e -> (e -> TailMaker e) -> (e -> TailMaker e) -> Trie e -> Trie e
merge' xs0 ys0 xTail yTail xCons yCons xy = merge'' 0 xs0 ys0 where
merge'' n xs ys = case (viewl xs, viewl ys) of
(x :< xs1, y :< ys1) -> let pfx = Seq.take n xs0; xT = xTail xs1; yT = yTail ys1; in case x `compare` y of
LT -> Edge pfx 0 $ fromDistinctAscList [(x, xT), (y, yT)]
EQ -> merge'' (n+1) xs1 ys1
GT -> Edge pfx 0 $ fromDistinctAscList [(y, yT), (x, xT)]
(x :< xs1, EmptyL) -> yCons x xs1
(EmptyL, y :< ys1) -> xCons y ys1
(EmptyL, EmptyL) -> xy
merger :: Ord e => Trie e -> Trie e -> Trie e
Leaf xs0 xn `merger` Leaf ys0 yn =
merge' xs0 ys0 (flip Leaf xn) (flip Leaf yn) (edger xs0 xn yn) (edger ys0 yn xn) (Leaf xs0 (xn + yn))
where edger xs xn yn y ys = Edge xs xn $ Map.singleton y (Leaf ys yn)
Leaf xs0 xn `merger` Edge ys0 yn yM =
merge' xs0 ys0 (flip Leaf xn) (\ ys -> Edge ys yn yM) (\ y ys -> Edge xs0 xn $ Map.singleton y (Edge ys yn yM))
(\ x xs -> Edge ys0 yn $ Map.insertWith merger x (Leaf xs xn) yM) (Edge xs0 (xn + yn) yM)
x@Edge{} `merger` y@Leaf{} = merger y x
Edge xs0 xn xM `merger` Edge ys0 yn yM
= merge' xs0 ys0 (edger xn xM) (edger yn yM) (cons xs0 xn yn yM) (cons ys0 yn xn xM)
(Edge xs0 (xn + yn) $ Map.unionWith merger xM yM)
where edger n m l = Edge l n m
cons xs xn yn yM y = Edge xs xn . Map.singleton y . edger yn yM
trieFromList :: Ord e => [[e]] -> MTrie e
trieFromList = extractCommon mempty
where groupHeads = groupBy ((==) `on` listToMaybe) . sortBy (comparing listToMaybe)
extractCommon pfx xs = case groupHeads xs of
[] -> Nothing
[empties@([]:_)] -> Just $ Leaf pfx (length empties)
(empties@([]:_):xss) -> Just $ Edge pfx (length empties) (fromGroups xss)
[(y:ys):yss] -> extractCommon (pfx |> y) (ys:map tail yss)
xss -> Just $ Edge pfx 0 (fromGroups xss)
fromGroups xss = fromDistinctAscList [(y, fromJust $ trieFromList $ ys : map tail yss) | ((y:ys):yss) <- xss]