TrieQueue.hs

{-# LANGUAGE PatternGuards, TypeFamilies, GeneralizedNewtypeDeriving #-}

-- | An experimental trie-based priority queue for lists.
module Data.Queue.TrieQueue (TrieQueue) where

import GHC.Exts
import Data.Queue.Class
import Control.Arrow(first)
import Control.Monad(liftM2)
import Data.Function
import Data.Maybe
import Data.List (sortBy, groupBy)
import Data.Word
import Data.Monoid
import Data.Ord
import Data.Sequence (Seq, viewl, ViewL(..), (><), (|>), (<|))
import qualified Data.Sequence as Seq
import Data.Map (Map, findMin)
import qualified Data.Map as Map
import qualified Data.Foldable as Fold (toList)
import Data.Queue.QueueHelpers

type Label e = Seq e
data Trie e = Empty | Leaf (Label e) {-# UNPACK #-} !Int | Edge (Label e) {-# UNPACK #-} !Int (Map e (Trie e)) deriving (Eq)
newtype TrieQueue e = TQ (MonoidQ (Trie e)) deriving (Eq, Show, Monoid)

instance Ord e => Monoid (Trie e) where
	mempty = Empty
	mappend = merger

instance Ord e => IQueue (TrieQueue e) where
	type QueueKey (TrieQueue e) = [e]
	empty = mempty
	merge = mappend
	mergeAll = mconcat

	singleton xs = TQ (HQ 1 (Leaf (Seq.fromList xs) 1))
	fromList = TQ . liftM2 HQ length trieFromList
	
	extract (TQ (HQ n t)) = fmap (\ (x, t) -> (Fold.toList x, TQ (HQ (n-1) t))) (extractMin' t)
	size (TQ (HQ n _)) = n

instance Show e => Show (Trie e) where
	show Empty = "Empty"
	show (Leaf xs xn) = "(" ++ show xn ++ "x" ++ show xs ++ ")"
	show (Edge xs xn m) = "==" ++ show xn ++ "x" ++ show (Fold.toList xs) ++ "==>" ++ show m

catTrie :: Ord e => Label e -> Trie e -> Trie e
catTrie _ Empty = Empty
catTrie xs (Leaf ys n) = Leaf (xs `mappend` ys) n
catTrie xs (Edge ys n m) = compactTrie (Edge (xs `mappend` ys) n m)

compactTrie :: Ord e => Trie e -> Trie e
compactTrie (Edge xs n m)
	| Map.null m		= compactTrie (Leaf xs n)
	| Map.size m == 1, (y, yT) <- findMin m
		= catTrie (xs |> y) yT
compactTrie (Leaf _ 0) = Empty
compactTrie t = t

extractMin' :: Ord e => Trie e -> Maybe (Label e, Trie e)
extractMin' (Leaf xs (n+1)) = Just (xs, compactTrie (Leaf xs n))
extractMin' (Edge xs (n+1) m) = Just (xs, compactTrie (Edge xs n m))
extractMin' (Edge xs 0 m) 
	| Just ((y, yT), m') <- Map.minViewWithKey m,
		Just (ys, yT') <- extractMin' yT
	= Just (xs >< (y <| ys), case yT' of
		Empty	-> compactTrie (Edge xs 0 m')
		_		-> Edge xs 0 (Map.insert y yT' m'))
extractMin' _ = Nothing

--extractMin :: Ord e => Trie e -> Maybe ([e], Trie e)
--extractMin = fmap (first toList) . extractMin'

type TailMaker e = Label e -> Trie e

{-# INLINE merge' #-}
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 =
	let merge'' n xs ys = case (viewl xs, viewl ys) of
		(x :< xs1, y :< ys1) -> let xT = xTail xs1; yT = yTail ys1; in case x `compare` y of
			LT	-> Edge (Seq.take n xs0) 0 $ Map.fromDistinctAscList [(x, xT), (y, yT)]
			EQ	-> merge'' (n+1) xs1 ys1
			GT	-> Edge (Seq.take n xs0) 0 $ Map.fromDistinctAscList [(y, yT), (x, xT)]
		(EmptyL, y :< ys1)
			-> xCons y ys1
		(x :< xs1, EmptyL)
			-> yCons x xs1
		(EmptyL, EmptyL)
			-> xy
		in merge'' 0 xs0 ys0
	
merger :: Ord e => Trie e -> Trie e -> Trie e	
Leaf xs0 xn `merger` Leaf ys0 yn =
	inline 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 =
	inline 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
	= inline 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
Empty `merger` trie = trie
trie `merger` Empty = trie

trieFromList :: Ord e => [[e]] -> Trie e
trieFromList = extractCommon mempty
	where	extractCommon pfx xs = case groupBy ((==) `on` listToMaybe) (sortBy (comparing listToMaybe) xs) of
			[]			-> Empty
			(empties@([]:_):xs')	-> Edge pfx (fromIntegral $ length empties) (fromGroups xs')
			[(y:ys):yss]		-> extractCommon (pfx |> y) (ys:map tail yss)
			xs'			-> Edge pfx 0 (fromGroups xs')
		fromGroups xss = Map.fromDistinctAscList [(y, trieFromList $ ys : map tail yss) | ((y:ys):yss) <- xss]