DAG.hs

module DAG where

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet

data DAG a = DAG { dagValues :: IntMap a, dagEdges :: IntMap IntSet }

type NodeIndex = Int
type NodeIndexes = IntSet

nodeChildren :: DAG -> NodeIndex -> NodeIndexes
nodeChildren DAG{dagEdges = edges} v = case edges `IntMap.lookup` v of
  Nothing -> IntSet.empty
  Just relatives -> IntSet.map (\o -> o + v) $ IntSet.filter (> 0) relatives
  
nodeParents :: DAG -> NodeIndex -> NodeIndexes
nodeParents DAG{dagEdges = edges} v = case edges `IntMap.lookup` v of
  Nothing -> IntSet.empty
  Just relatives -> IntSet.map (\o -> o + v) $ IntSet.filter (< 0) relatives