module Data.VCache.Trie
( Trie
, trie_space
, empty, singleton
, null, size
, lookup, lookup'
, prefixKeys
, lookupPrefix, lookupPrefix'
, lookupPrefixNode, lookupPrefixNode'
, deletePrefix
, insert, delete, adjust
, insertList, deleteList
, toList, toListBy, elems, keys
, foldr, foldr', foldrM, foldrWithKey, foldrWithKey', foldrWithKeyM
, foldl, foldl', foldlM, foldlWithKey, foldlWithKey', foldlWithKeyM
, map, mapM, mapWithKey, mapWithKeyM
, toListOnKey
, diff, Diff(..)
, validate
, unsafeTrieAddr
, DerefNode
) where
import Prelude hiding (null, lookup, foldr, foldl, map, mapM, elem)
import Control.Applicative hiding (empty)
import qualified Control.Monad as M
import Control.Exception (assert)
import Data.Word
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.Array.IArray as A
import qualified Data.List as L
import Data.Maybe
import Data.Monoid
import Strict
import Ident
import Data.VCache.Trie.Type
import Database.VCache
type DerefNode a = VRef (Node a) -> Node a
defaultDeref :: DerefNode a
defaultDeref = derefc CacheMode0
empty :: VSpace -> Trie a
empty = Trie Nothing
singleton :: (VCacheable a) => VSpace -> ByteString -> a -> Trie a
singleton vc k a = Trie (Just $! vref vc (singletonNode k a)) vc
singletonNode :: ByteString -> a -> Node a
singletonNode k a = Node (mkChildren []) k (Just a)
mkChildren :: [(Word8, Child a)] -> Children a
mkChildren = A.accumArray (flip const) Nothing (minBound, maxBound)
null :: Trie a -> Bool
null = isNothing . trie_root
size :: Trie a -> Int
size = foldr' (const (1 +)) 0
toList :: Trie a -> [(ByteString, a)]
toList = toListBy (,)
elems :: Trie a -> [a]
elems = toListBy (flip const)
keys :: Trie a -> [ByteString]
keys = toListBy const
toListBy :: (ByteString -> a -> b) -> Trie a -> [b]
toListBy fn = foldrWithKey (\ k v bs -> fn k v : bs) []
prefixKeys :: (VCacheable a) => ByteString -> Trie a -> Trie a
prefixKeys _ t@(Trie Nothing _) = t
prefixKeys prefix (Trie (Just pRoot) vc) =
let root = deref' pRoot in
let p' = prefix `B.append` (trie_prefix root) in
let root' = root { trie_prefix = p' } in
Trie (Just $! vref vc root') vc
lookup' :: DerefNode a -> ByteString -> Trie a -> Maybe a
lookup' d k = _lookup d k . trie_root
lookup :: ByteString -> Trie a -> Maybe a
lookup = lookup' defaultDeref
_lookup :: DerefNode a -> ByteString -> Child a -> Maybe a
_lookup _ _ Nothing = Nothing
_lookup d key (Just c) =
let tn = d c in
let pre = trie_prefix tn in
let s = sharedPrefixLen key pre in
let k = B.length key in
let p = B.length pre in
if (s < p) then Nothing else
assert ((s == p) && (k >= p)) $
if (k == p) then trie_accept tn else
let key' = B.drop (p+1) key in
let c' = (trie_branch tn) A.! (B.index key p) in
_lookup d key' c'
lookupPrefix :: (VCacheable a) => ByteString -> Trie a -> Trie a
lookupPrefix = lookupPrefix' defaultDeref
lookupPrefix' :: (VCacheable a) => DerefNode a -> ByteString -> Trie a -> Trie a
lookupPrefix' d k tr =
let node = lookupPrefixNode' d k tr in
let vc = trie_space tr in
let child = vref vc <$> node in
Trie child vc
lookupPrefixNode :: (VCacheable a) => ByteString -> Trie a -> Maybe (Node a)
lookupPrefixNode = lookupPrefixNode' defaultDeref
lookupPrefixNode' :: (VCacheable a) => DerefNode a -> ByteString -> Trie a -> Maybe (Node a)
lookupPrefixNode' d k = _lookupP d k . trie_root
_lookupP :: (VCacheable a) => DerefNode a -> ByteString -> Child a -> Maybe (Node a)
_lookupP d key c | B.null key = d <$> c
_lookupP _ _ Nothing = Nothing
_lookupP d key (Just c) =
let tn = d c in
let pre = trie_prefix tn in
let s = sharedPrefixLen key pre in
let k = B.length key in
let p = B.length pre in
if (k <= p)
then if (s < k) then Nothing else
assert (s == k) $
let pre' = B.drop k pre in
let tn' = tn { trie_prefix = pre' } in
Just $! tn'
else if (s < p) then Nothing else
assert (s == p) $
let key' = B.drop (p+1) key in
let c' = (trie_branch tn) A.! (B.index key p) in
_lookupP d key' c'
deletePrefix :: (VCacheable a) => ByteString -> Trie a -> Trie a
deletePrefix = deletePrefix' deref'
deletePrefix' :: VCacheable a => DerefNode a -> ByteString -> Trie a -> Trie a
deletePrefix' d p (Trie c vc) = Trie (_deleteP d p c) vc
_deleteP :: (VCacheable a) => DerefNode a -> ByteString -> Child a -> Child a
_deleteP _ _ Nothing = Nothing
_deleteP d key c@(Just pNode) =
let vc = vref_space pNode in
let tn = d pNode in
let pre = trie_prefix tn in
let s = sharedPrefixLen key pre in
let k = B.length key in
let p = B.length pre in
if (k <= p)
then if (s < k) then c else
assert (s == k) $ Nothing
else if (s < p) then c else
assert (s == p) $
let key' = B.drop (p+1) key in
let idx = B.index key p in
let tgt = (trie_branch tn) A.! idx in
let tgt' = _deleteP d key' tgt in
if (tgt == tgt') then c else
let bDel = isJust tgt && isNothing tgt' in
let branch' = trie_branch tn A.// [(idx, tgt')] in
let tn' = tn { trie_branch = branch' } in
collapseIf vc bDel tn'
insert :: (VCacheable a) => ByteString -> a -> Trie a -> Trie a
insert k a = adjust (const (Just a)) k
insertList :: (VCacheable a) => [(ByteString, a)] -> Trie a -> Trie a
insertList = flip $ L.foldl' ins where
ins t (k,v) = insert k v t
delete :: (VCacheable a) => ByteString -> Trie a -> Trie a
delete k = adjust (const Nothing) k
deleteList :: (VCacheable a) => [ByteString] -> Trie a -> Trie a
deleteList = flip (L.foldl' (flip delete))
adjust :: (VCacheable a) => (Maybe a -> Maybe a) -> ByteString -> Trie a -> Trie a
adjust fn k t = runStrict $ adjustM fn' k t where
fn' a = return (fn a)
adjustM :: (VCacheable a, Monad m) => (Maybe a -> m (Maybe a)) -> ByteString -> Trie a -> m (Trie a)
adjustM fn k0 tr = adjustRoot where
vc = trie_space tr
adjustRoot =
wc k0 (trie_root tr) >>= \ c' ->
return (Trie c' vc)
wc key Nothing = fn Nothing >>= \ r -> case r of
Nothing -> return Nothing
Just a ->
let tn' = singletonNode key a in
let c' = Just $! vref vc tn' in
return $! c'
wc key c@(Just pChild) =
let tn = deref' pChild in
let pre = trie_prefix tn in
let s = sharedPrefixLen key pre in
let p = B.length pre in
let k = B.length key in
if (s < p)
then fn Nothing >>= \ r ->
if isNothing r then return c else
let ixP = B.index pre s in
let tnP = tn { trie_prefix = B.drop (s+1) pre } in
let cP = Just $! vref vc tnP in
if (s == k)
then
let br' = mkChildren [(ixP, cP)] in
let tn' = Node br' key r in
let c' = Just $! vref vc tn' in
return $! c'
else
let ixK = B.index key s in
assert ((s < k) && (ixK /= ixP)) $
let brK = mkChildren [] in
let tnK = Node brK (B.drop (s+1) key) r in
let cK = Just $! vref vc tnK in
let br' = mkChildren [(ixK,cK), (ixP, cP)] in
let tn' = Node br' (B.take s pre) Nothing in
let c' = Just $! vref vc tn' in
return $! c'
else if (s < k)
then
let key' = B.drop (s+1) key in
let ixK = B.index key s in
let cK = (trie_branch tn) A.! ixK in
wc key' cK >>= \ cK' ->
if (cK == cK') then return c else
let bDel = isJust cK && isNothing cK' in
let br' = trie_branch tn A.// [(ixK, cK')] in
let tn' = tn { trie_branch = br' } in
return $! collapseIf vc bDel tn'
else assert (s == k) $
let v = trie_accept tn in
fn v >>= \ v' ->
let bDel = isJust v && isNothing v' in
let tn' = tn { trie_accept = v' } in
return $! collapseIf vc bDel tn'
collapseIf :: (VCacheable a) => VSpace -> Bool -> Node a -> Child a
collapseIf vc bDel tn =
if not bDel then Just $! vref vc tn else
case tryCollapse tn of
Nothing -> Nothing
Just tn' -> Just $! vref vc tn'
tryCollapse :: Node a -> Maybe (Node a)
tryCollapse tn =
if isJust (trie_accept tn) then Just tn else
let lChildren = L.filter (isJust . snd) (A.assocs (trie_branch tn)) in
case lChildren of
[] -> Nothing
[(ix, Just c)] ->
let tnC = deref' c in
let key' = mconcat [trie_prefix tn, B.singleton ix, trie_prefix tnC] in
let tn' = tnC { trie_prefix = key' } in
Just tn'
_ -> Just tn
validate :: Trie a -> Bool
validate = maybe True validRef . trie_root where
validRef = validNode . deref'
validNode tn =
let lChildren = catMaybes (A.elems (trie_branch tn)) in
let bBranch = case lChildren of { (_:_:_) -> True; _ -> False } in
let bAccept = isJust (trie_accept tn) in
let bNodeValid = bAccept || bBranch in
bNodeValid && L.all validRef lChildren
foldrWithKey, foldrWithKey' :: (ByteString -> a -> b -> b) -> b -> Trie a -> b
foldlWithKey, foldlWithKey' :: (b -> ByteString -> a -> b) -> b -> Trie a -> b
foldrWithKey fn b t = runIdent $ foldrWithKeyM (apwf fn) b t
foldrWithKey' fn b t = runStrict $ foldrWithKeyM (apwf fn) b t
foldlWithKey fn b t = runIdent $ foldlWithKeyM (apwf fn) b t
foldlWithKey' fn b t = runStrict $ foldlWithKeyM (apwf fn) b t
apwf :: (Applicative f) => (a -> b -> c -> d) -> (a -> b -> c -> f d)
apwf fn a b c = pure (fn a b c)
foldr, foldr' :: (a -> b -> b) -> b -> Trie a -> b
foldl, foldl' :: (b -> a -> b) -> b -> Trie a -> b
foldrM :: Monad m => (a -> b -> m b) -> b -> Trie a -> m b
foldlM :: Monad m => (b -> a -> m b) -> b -> Trie a -> m b
foldr = foldrWithKey . skip1st
foldr' = foldrWithKey' . skip1st
foldrM = foldrWithKeyM . skip1st
foldl = foldlWithKey . skip2nd
foldl' = foldlWithKey' . skip2nd
foldlM = foldlWithKeyM . skip2nd
skip1st :: (b -> c) -> (a -> b -> c)
skip2nd :: (a -> c) -> (a -> b -> c)
skip1st = const
skip2nd = flip . const
foldrWithKeyM :: (Monad m) => (ByteString -> a -> b -> m b) -> b -> Trie a -> m b
foldrWithKeyM ff = wr where
wr b = wc mempty b . trie_root
wc _ b Nothing = return b
wc p b (Just c) =
let tn = deref' c in
let p' = p <> trie_prefix tn in
wlc p' (trie_branch tn) maxBound b >>=
maybe return (ff p') (trie_accept tn)
wlc p a !k b =
let cc = if (minBound == k) then return else wlc p a (k1) in
let p' = p `B.snoc` k in
wc p' b (a A.! k) >>= cc
foldlWithKeyM :: (Monad m) => (b -> ByteString -> a -> m b) -> b -> Trie a -> m b
foldlWithKeyM ff = wr where
wr b = wc mempty b . trie_root
wc _ b Nothing = return b
wc p b (Just c) =
let tn = deref' c in
let p' = p <> trie_prefix tn in
let cc = wlc p' (trie_branch tn) minBound in
case trie_accept tn of
Nothing -> cc b
Just val -> ff b p' val >>= cc
wlc p a !k b =
let cc = if (maxBound == k) then return else wlc p a (k+1) in
let p' = p `B.snoc` k in
wc p' b (a A.! k) >>= cc
toListOnKey :: ByteString -> Trie a -> ([(ByteString, a)], [(ByteString, a)])
toListOnKey fullKey = nodeOnKey 0 . trie_root where
nodeOnKey _ Nothing = ([],[])
nodeOnKey nKeyBytes (Just c) =
let tn = deref' c in
let key = B.drop nKeyBytes fullKey in
let pre = trie_prefix tn in
let s = sharedPrefixLen key pre in
let pathToNode = B.take nKeyBytes fullKey in
let nodeToLeftOfKey = (nodeL pathToNode tn, []) in
let nodeToRightOfKey = ([], nodeR pathToNode tn) in
if (s == B.length key) then nodeToRightOfKey else
assert (s < B.length key) $
let ck = B.index key s in
if (s == B.length pre) then nodeSplitOnKey (nKeyBytes + s) ck tn else
assert (s < B.length pre) $
let cp = B.index pre s in
if (cp > ck) then nodeToRightOfKey else
assert (cp < ck) $ nodeToLeftOfKey
nodeElem p = maybe [] (return . (,) p) . trie_accept
nodeL pathToNode tn =
let p = pathToNode <> trie_prefix tn in
let onC (i, c) = maybe [] (nodeL (p `B.snoc` i) . deref') c in
let children = fmap onC $ A.assocs (trie_branch tn) in
let elem = nodeElem p tn in
mconcat $ L.reverse (elem:children)
nodeR pathToNode tn =
let p = pathToNode <> trie_prefix tn in
let onChild (i,c) = maybe [] (nodeR (p `B.snoc` i) . deref') c in
let children = fmap onChild $ A.assocs (trie_branch tn) in
let elem = nodeElem p tn in
mconcat (elem:children)
nodeSplitOnKey nKeyBytes keyChar tn =
let p = B.take nKeyBytes fullKey in
let elem = nodeElem p tn in
let onCL (i,c) = maybe [] (nodeL (p `B.snoc` i) . deref') c in
let onCR (i,c) = maybe [] (nodeR (p `B.snoc` i) . deref') c in
let cL = fmap onCL $ L.filter ((< keyChar) . fst) $ A.assocs (trie_branch tn) in
let cR = fmap onCR $ L.filter ((> keyChar) . fst) $ A.assocs (trie_branch tn) in
let eL = mconcat $ L.reverse (elem:cL) in
let eR = mconcat cR in
let (keL,keR) = nodeOnKey (nKeyBytes + 1) (trie_branch tn A.! keyChar) in
(keL ++ eL, keR ++ eR)
data Diff a = InL a | Diff a a | InR a
deriving (Show, Eq)
instance Functor Diff where
fmap f (InL a) = InL (f a)
fmap f (Diff a b) = Diff (f a) (f b)
fmap f (InR b) = InR (f b)
diff :: (Eq a) => Trie a -> Trie a -> [(ByteString, Diff a)]
diff = diffRoot where
diffRoot a b = diffChild mempty (trie_root a) (trie_root b)
diffChild _ Nothing Nothing = mempty
diffChild k (Just a) Nothing = subtree k InL (deref' a)
diffChild k Nothing (Just b) = subtree k InR (deref' b)
diffChild k (Just a) (Just b) =
if (a == b) then mempty else
diffNode k (deref' a) (deref' b)
diffNode k tnA tnB =
let kA = trie_prefix tnA in
let kB = trie_prefix tnB in
let s = sharedPrefixLen kA kB in
if (s == B.length kA)
then if (s == B.length kB)
then diffNodeEQ k tnA tnB
else diffNodeAB k tnA tnB
else if (s == B.length kB)
then diffNodeBA k tnA tnB
else let elemsA = subtree k InL tnA in
let elemsB = subtree k InR tnB in
let ca = B.index (trie_prefix tnA) s in
let cb = B.index (trie_prefix tnB) s in
assert (ca /= cb) $
if (ca < cb) then elemsA <> elemsB
else elemsB <> elemsA
diffNodeAB k tnA tnB =
assert (trie_prefix tnA `isStrictPrefixOf` trie_prefix tnB) $
let s = B.length (trie_prefix tnA) in
let c = B.index (trie_prefix tnB) s in
let (diffLeft, diffRight) = splitTree k InL c tnA in
let diffMiddle = case trie_branch tnA A.! c of
Nothing -> subtree k InR tnB
Just childSplit ->
let keySplit = k <> B.take (s + 1) (trie_prefix tnB) in
let tnB' = tnB { trie_prefix = B.drop (s + 1) (trie_prefix tnB) } in
diffNode keySplit (deref' childSplit) tnB'
in
mconcat [diffLeft, diffMiddle, diffRight]
diffNodeBA k tnA tnB =
assert (trie_prefix tnB `isStrictPrefixOf` trie_prefix tnA) $
let s = B.length (trie_prefix tnB) in
let c = B.index (trie_prefix tnA) s in
let (diffLeft, diffRight) = splitTree k InR c tnB in
let diffMiddle = case trie_branch tnB A.! c of
Nothing -> subtree k InL tnA
Just childSplit ->
let keySplit = k <> B.take (s + 1) (trie_prefix tnA) in
let tnA' = tnA { trie_prefix = B.drop (s + 1) (trie_prefix tnA) } in
diffNode keySplit tnA' (deref' childSplit)
in
mconcat [diffLeft, diffMiddle, diffRight]
diffNodeEQ k tnA tnB =
assert (trie_prefix tnA == trie_prefix tnB) $
let fullKey = k <> trie_prefix tnA in
let diffK = case (trie_accept tnA, trie_accept tnB) of
(Just va, Nothing) -> [(fullKey, InL va)]
(Nothing, Just vb) -> [(fullKey, InR vb)]
(Just va, Just vb) | (va /= vb) -> [(fullKey, Diff va vb)]
_ -> []
in
let diffC (i, cA) = diffChild (fullKey `B.snoc` i) cA (trie_branch tnB A.! i) in
let diffChildren = L.concatMap diffC (A.assocs (trie_branch tnA)) in
diffK <> diffChildren
subtree k fn tn =
let fullKey = k <> trie_prefix tn in
let lK = maybeToList $ fmap ((,) fullKey . fn) (trie_accept tn) in
let onC (i,c) = maybe [] (subtree (fullKey `B.snoc` i) fn . deref') c in
let lC = mconcat $ fmap onC $ A.assocs (trie_branch tn) in
lK <> lC
splitTree k fn c tn =
let fullKey = k <> trie_prefix tn in
let elemK = maybeToList $ fmap ((,) fullKey . fn) $ trie_accept tn in
let onC i = maybe [] (subtree (fullKey `B.snoc` i) fn . deref') (trie_branch tn A.! i) in
let rangeL = if (c == minBound) then [] else enumFromTo minBound (pred c) in
let rangeR = if (c == maxBound) then [] else enumFromTo (succ c) maxBound in
let leftElems = mconcat $ elemK : fmap onC rangeL in
let rightElems = mconcat $ fmap onC rangeR in
(leftElems, rightElems)
isStrictPrefixOf :: ByteString -> ByteString -> Bool
isStrictPrefixOf a b = (B.length a < B.length b) && (a `B.isPrefixOf` b)
map :: (VCacheable b) => (a -> b) -> Trie a -> Trie b
map = mapWithKey . skip1st
mapM :: (VCacheable b, Monad m) => (a -> m b) -> Trie a -> m (Trie b)
mapM = mapWithKeyM . skip1st
mapWithKey :: (VCacheable b) => (ByteString -> a -> b) -> Trie a -> Trie b
mapWithKey fn = runStrict . mapWithKeyM fn' where
fn' a b = pure (fn a b)
mapWithKeyM :: (Monad m, VCacheable b) => (ByteString -> a -> m b) -> Trie a -> m (Trie b)
mapWithKeyM ff = wr where
wr (Trie c vc) =
wc mempty c >>= \ c' ->
return (Trie c' vc)
wc _ Nothing = return Nothing
wc p (Just c) =
let tn = deref' c in
let p' = p <> trie_prefix tn in
mbrun (ff p') (trie_accept tn) >>= \ accept' ->
let lcs = A.assocs (trie_branch tn) in
M.mapM (wlc p') lcs >>= \ lcs' ->
let branch' = A.array (minBound, maxBound) lcs' in
let tn' = Node branch' (trie_prefix tn) accept' in
let c' = vref' (vref_space c) tn' in
return $! (Just $! c')
wlc p (ix, child) =
let p' = p `B.snoc` ix in
wc p' child >>= \ child' ->
return (ix, child')
mbrun :: (Monad m) => (a -> m b) -> Maybe a -> m (Maybe b)
mbrun _ Nothing = return Nothing
mbrun action (Just a) = M.liftM Just (action a)