module Database.Algebra.Dag
(
AlgebraDag
, Operator(..)
, nodeMap
, rootNodes
, refCountMap
, mkDag
, emptyDag
, addRootNodes
, parents
, topsort
, hasPath
, reachableNodesFrom
, operator
, insert
, insertNoShare
, replaceChild
, replaceRoot
, collect
) where
import Control.Exception.Base
import Data.Aeson
import qualified Data.Graph.Inductive.Graph as G
import Data.Graph.Inductive.PatriciaTree
import qualified Data.Graph.Inductive.Query.DFS as DFS
import qualified Data.IntMap as IM
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Set as S
import Database.Algebra.Dag.Common
data AlgebraDag a = AlgebraDag
{ nodeMap :: NodeMap a
, opMap :: M.Map a AlgNode
, nextNodeID :: AlgNode
, graph :: UGr
, rootNodes :: [AlgNode]
, refCountMap :: NodeMap Int
}
instance ToJSON a => ToJSON (AlgebraDag a) where
toJSON dag = toJSON (nodeMap dag, rootNodes dag)
instance (FromJSON a, Operator a) => FromJSON (AlgebraDag a) where
parseJSON v = do
(nm, rs) <- parseJSON v
return $ mkDag nm rs
class (Ord a, Show a) => Operator a where
opChildren :: a -> [AlgNode]
replaceOpChild :: a -> AlgNode -> AlgNode -> a
initRefCount :: Operator o => [AlgNode] -> NodeMap o -> NodeMap Int
initRefCount rs nm = L.foldl' incParents (IM.foldr' insertEdge IM.empty nm) (L.nub rs)
where
insertEdge op rm = L.foldl' incParents rm (L.nub $ opChildren op)
incParents rm n = IM.insert n ((IM.findWithDefault 0 n rm) + 1) rm
initOpMap :: Ord o => NodeMap o -> M.Map o AlgNode
initOpMap nm = IM.foldrWithKey (\n o om -> M.insert o n om) M.empty nm
mkDag :: Operator a => NodeMap a -> [AlgNode] -> AlgebraDag a
mkDag m rs = AlgebraDag { nodeMap = mNormalized
, graph = g
, rootNodes = rs
, refCountMap = initRefCount rs mNormalized
, opMap = initOpMap mNormalized
, nextNodeID = 1 + (fst $ IM.findMax mNormalized)
}
where
mNormalized = normalizeMap rs m
g = uncurry G.mkUGraph $ IM.foldrWithKey aux ([], []) mNormalized
aux n op (allNodes, allEdges) = (n : allNodes, es ++ allEdges)
where
es = map (\v -> (n, v)) $ opChildren op
emptyDag :: AlgebraDag a
emptyDag =
AlgebraDag { nodeMap = IM.empty
, opMap = M.empty
, nextNodeID = 1
, graph = G.mkUGraph [] []
, rootNodes = []
, refCountMap = IM.empty
}
addRootNodes :: Operator a => AlgebraDag a -> [AlgNode] -> AlgebraDag a
addRootNodes d rs = assert (all (\n -> IM.member n $ nodeMap d) rs) $
d { rootNodes = rs
, nodeMap = mNormalized
, refCountMap = initRefCount rs mNormalized
, opMap = initOpMap mNormalized
, graph = uncurry G.mkUGraph $ IM.foldrWithKey aux ([], []) mNormalized
}
where
mNormalized = normalizeMap rs (nodeMap d)
aux n op (allNodes, allEdges) = (n : allNodes, es ++ allEdges)
where
es = map (\v -> (n, v)) $ opChildren op
reachable :: Operator a => NodeMap a -> [AlgNode] -> S.Set AlgNode
reachable m rs = L.foldl' traverseDag S.empty rs
where traverseDag :: S.Set AlgNode -> AlgNode -> S.Set AlgNode
traverseDag s n = if S.member n s
then s
else L.foldl' traverseDag (S.insert n s) (opChildren $ lookupOp n)
lookupOp n = case IM.lookup n m of
Just op -> op
Nothing -> error $ "node not present in map: " ++ (show n)
normalizeMap :: Operator a => [AlgNode] -> NodeMap a -> NodeMap a
normalizeMap rs m =
let reachableNodes = reachable m rs
in IM.filterWithKey (\n _ -> S.member n reachableNodes) m
lookupRefCount :: AlgNode -> AlgebraDag a -> Int
lookupRefCount n d =
case IM.lookup n (refCountMap d) of
Just c -> c
Nothing -> error $ "no refcount value for node " ++ (show n)
decrRefCount :: AlgebraDag a -> AlgNode -> AlgebraDag a
decrRefCount d n =
let refCount = lookupRefCount n d
refCount' = assert (refCount /= 0) $ refCount 1
in d { refCountMap = IM.insert n refCount' (refCountMap d) }
delete' :: Operator a => AlgNode -> AlgebraDag a -> AlgebraDag a
delete' n d =
let op = operator n d
g' = G.delNode n $ graph d
m' = IM.delete n $ nodeMap d
rc' = IM.delete n $ refCountMap d
opMap' = case M.lookup op $ opMap d of
Just n' | n == n' -> M.delete op $ opMap d
_ -> opMap d
in d { nodeMap = m', graph = g', refCountMap = rc', opMap = opMap' }
refCountSafe :: AlgNode -> AlgebraDag o -> Maybe Int
refCountSafe n d = IM.lookup n $ refCountMap d
collect :: Operator o => S.Set AlgNode -> AlgebraDag o -> AlgebraDag o
collect collectNodes d = S.foldl' tryCollectNode d collectNodes
where tryCollectNode :: (Show o, Operator o) => AlgebraDag o -> AlgNode -> AlgebraDag o
tryCollectNode di n =
case refCountSafe n di of
Just rc -> if rc == 0
then
let cs = L.nub $ opChildren $ operator n di
d' = delete' n di
in L.foldl' cutEdge d' cs
else di
Nothing -> di
cutEdge :: Operator a => AlgebraDag a -> AlgNode -> AlgebraDag a
cutEdge d edgeTarget =
let d' = decrRefCount d edgeTarget
newRefCount = lookupRefCount edgeTarget d'
in if newRefCount == 0
then let cs = L.nub $ opChildren $ operator edgeTarget d'
d'' = delete' edgeTarget d'
in L.foldl' cutEdge d'' cs
else d'
addRefTo :: AlgebraDag a -> AlgNode -> AlgebraDag a
addRefTo d n =
let refCount = lookupRefCount n d
in d { refCountMap = IM.insert n (refCount + 1) (refCountMap d) }
replaceRoot :: Operator a => AlgebraDag a -> AlgNode -> AlgNode -> AlgebraDag a
replaceRoot d old new =
if old `elem` (rootNodes d)
then let rs' = map doReplace $ rootNodes d
doReplace r = if r == old then new else r
d' = d { rootNodes = rs' }
in
assert (old /= new) $ addRefTo (decrRefCount d' old) new
else d
insert :: Operator a => a -> AlgebraDag a -> (AlgNode, AlgebraDag a)
insert op d =
case M.lookup op $ opMap d of
Just n -> (n, d)
Nothing -> insertNoShare op d
insertNoShare :: Operator a => a -> AlgebraDag a -> (AlgNode, AlgebraDag a)
insertNoShare op d =
let cs = L.nub $ opChildren op
n = nextNodeID d
g' = G.insEdges (map (\c -> (n, c, ())) cs) $ G.insNode (n, ()) $ graph d
m' = IM.insert n op $ nodeMap d
rc' = IM.insert n 0 $ refCountMap d
opMap' = M.insert op n $ opMap d
d' = d { nodeMap = m'
, graph = g'
, refCountMap = rc'
, opMap = opMap'
, nextNodeID = n + 1
}
in (n, L.foldl' addRefTo d' cs)
parents :: AlgNode -> AlgebraDag a -> [AlgNode]
parents n d = G.pre (graph d) n
replaceChild :: Operator a => AlgNode -> AlgNode -> AlgNode -> AlgebraDag a -> AlgebraDag a
replaceChild parent old new d =
let op = operator parent d
in if old `elem` opChildren op && old /= new
then let op' = replaceOpChild op old new
m' = IM.insert parent op' $ nodeMap d
om' = M.insert op' parent $ M.delete op $ opMap d
g' = G.insEdge (parent, new, ()) $ G.delEdge (parent, old) $ graph d
d' = d { nodeMap = m', graph = g', opMap = om' }
d'' = decrRefCount d' old
in
if new `elem` G.suc (graph d) parent
then d''
else addRefTo d'' new
else d
operator :: Operator a => AlgNode -> AlgebraDag a -> a
operator n d =
case IM.lookup n $ nodeMap d of
Just op -> op
Nothing -> error $ "AlgebraDag.operator: lookup failed for " ++ (show n) ++ "\n" ++ (show $ map fst $ IM.toList $ nodeMap d)
topsort :: Operator a => AlgebraDag a -> [AlgNode]
topsort d = DFS.topsort $ graph d
reachableNodesFrom :: AlgNode -> AlgebraDag a -> S.Set AlgNode
reachableNodesFrom n d = S.fromList $ DFS.reachable n $ graph d
hasPath :: AlgNode -> AlgNode -> AlgebraDag a -> Bool
hasPath a b d = b `S.member` (reachableNodesFrom a d)