1 files changed, 43 insertions, 11 deletions
diff --git a/DataFlow.hs b/DataFlow.hs
index 78ec018..98cf65d 100644
--- a/DataFlow.hs
+++ b/DataFlow.hs
@@ -1,30 +1,53 @@
-module DataFlow(dataFlowAnalysis, Direction(..)) where
+{-# LANGUAGE ScopedTypeVariables, TupleSections #-}
+module DataFlow(dataFlowAnalysis, Direction(..), module Lattice) where
 
 import Data.Function
 import Data.List
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+import Lattice
 
 
 data Direction = Forwards | Backwards
   deriving (Show, Eq)
 
 dataFlowAnalysis
-        :: Ord id
+        :: forall id state. (Ord id, Lattice state)
         -- | The direction of the dataflow analysis
-        => Direction               
+        => Direction
         -- | The control flow graph of basic blocks
-        -> Map.Map id [id]         
-        -- | The initial values
-        -> (id -> state)           
+        -> Map.Map id [id]
+        -- | The initial nodes in the analysis (i.e. either the initial
+        -- node or the final nodes)
+        -> [id]
+        -- | The initial value for sources in the analysis
+        -> state
         -- | The transfer function; note that the meaning of the states
         -- depends on the direction
-        -> (state -> id -> state)  
-        -- | For each block, the incoming state and the outgoing state
+        -> (state -> id -> state)
+        -- | For each block that is reachable from 'entryNodes', the
+        -- incoming state and the outgoing state (in CFG order, not
+        -- analysis order)
         -> Map.Map id (state, state)
-dataFlowAnalysis Backwards cfg initF transF =
+dataFlowAnalysis Backwards cfg entryNodes initValue transF =
     Map.map (\(a, b) -> (b, a)) $
-        dataFlowAnalysis Forwards (reverseGraph cfg) initF transF
-dataFlowAnalysis Forwards cfg initF transF = undefined
+        dataFlowAnalysis Forwards (reverseGraph cfg) entryNodes initValue transF
+dataFlowAnalysis Forwards cfg entryNodes initValue transF =
+    let initWorklist = [(from, to) | (from, tos) <- Map.assocs cfg, to <- tos]
+        reachable = floodFill entryNodes cfg
+        initStates = Map.fromList
+            [(lab, if lab `elem` entryNodes then initValue else bottom)
+            | lab <- reachable]
+    in Map.mapWithKey (\k s -> (s, transF s k)) $ flowFixpoint initWorklist initStates
+  where
+    flowFixpoint :: [(id, id)] -> Map.Map id state -> Map.Map id state
+    flowFixpoint [] stateMap = stateMap
+    flowFixpoint ((from, to) : worklist) stateMap =
+        let after = transF (stateMap Map.! from) from `join` (stateMap Map.! to)
+        in if after /= stateMap Map.! to
+               then flowFixpoint (map (to,) (cfg Map.! to) ++ worklist)
+                                 (Map.insert to after stateMap)
+               else flowFixpoint worklist stateMap
 
 
 reverseGraph :: Ord a => Map.Map a [a] -> Map.Map a [a]
@@ -36,3 +59,12 @@ reverseGraph graph =
     $ [(to, from)
       | (from, tos) <- Map.assocs graph
       , to <- tos]
+
+floodFill :: Ord a => [a] -> Map.Map a [a] -> [a]
+floodFill startNodes graph = Set.toList $ go $ Set.fromList startNodes
+  where
+    go seen =
+        let front = Set.fromList (concatMap (graph Map.!) seen) Set.\\ seen
+        in if Set.size front > 0
+               then go (Set.union seen front)
+               else seen