module Tokstyle.Analysis.Dataflow
    ( Dataflow (..)
    , solve
    , solveBackward
    ) where

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Sequence   (Seq (..), (|>))
import qualified Data.Sequence   as Seq
import           Data.Set        (Set)
import qualified Data.Set        as Set

-- | A generic dataflow problem definition.
data Dataflow node edge state = Dataflow
    { Dataflow node edge state -> node -> state -> state
transfer     :: node -> state -> state               -- ^ How a node changes the state
    , Dataflow node edge state -> node -> edge -> state -> state
edgeTransfer :: node -> edge -> state -> state       -- ^ How an edge changes the state
    , Dataflow node edge state -> state -> state -> state
merge        :: state -> state -> state              -- ^ How to combine states from two paths
    , Dataflow node edge state -> state
initial      :: state                                -- ^ Starting state for Entry node (forward) or Exit node (backward)
    }

-- | Solves the forward dataflow problem using a worklist algorithm.
solve :: (Ord node, Eq state)
      => node                        -- ^ The entry node
      -> Map node [(edge, node)]     -- ^ The CFG
      -> Dataflow node edge state
      -> Map node state
solve :: node
-> Map node [(edge, node)]
-> Dataflow node edge state
-> Map node state
solve node
entry Map node [(edge, node)]
cfg Dataflow node edge state
problem =
    let startStates :: Map node state
startStates = node -> state -> Map node state
forall k a. k -> a -> Map k a
Map.singleton node
entry (Dataflow node edge state -> state
forall node edge state. Dataflow node edge state -> state
initial Dataflow node edge state
problem)
        worklist :: Seq node
worklist = node -> Seq node
forall a. a -> Seq a
Seq.singleton node
entry
        inWorklist :: Set node
inWorklist = node -> Set node
forall a. a -> Set a
Set.singleton node
entry
    in Seq node -> Set node -> Map node state -> Map node state
loop Seq node
worklist Set node
inWorklist Map node state
startStates
  where
    loop :: Seq node -> Set node -> Map node state -> Map node state
loop Seq node
Empty Set node
_ Map node state
states = Map node state
states
    loop (node
u :<| Seq node
ws) Set node
inWS Map node state
states =
        let currentIn :: state
currentIn = state -> node -> Map node state -> state
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault (Dataflow node edge state -> state
forall node edge state. Dataflow node edge state -> state
initial Dataflow node edge state
problem) node
u Map node state
states
            nodeOut :: state
nodeOut = Dataflow node edge state -> node -> state -> state
forall node edge state.
Dataflow node edge state -> node -> state -> state
transfer Dataflow node edge state
problem node
u state
currentIn
            successors :: [(edge, node)]
successors = [(edge, node)] -> node -> Map node [(edge, node)] -> [(edge, node)]
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] node
u Map node [(edge, node)]
cfg
            (Seq node
nextWS, Set node
nextInWS, Map node state
nextStates) = ((Seq node, Set node, Map node state)
 -> (edge, node) -> (Seq node, Set node, Map node state))
-> (Seq node, Set node, Map node state)
-> [(edge, node)]
-> (Seq node, Set node, Map node state)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (node
-> state
-> (Seq node, Set node, Map node state)
-> (edge, node)
-> (Seq node, Set node, Map node state)
forall k.
Ord k =>
node
-> state
-> (Seq k, Set k, Map k state)
-> (edge, k)
-> (Seq k, Set k, Map k state)
update node
u state
nodeOut) (Seq node
ws, node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.delete node
u Set node
inWS, Map node state
states) [(edge, node)]
successors
        in Seq node -> Set node -> Map node state -> Map node state
loop Seq node
nextWS Set node
nextInWS Map node state
nextStates

    update :: node
-> state
-> (Seq k, Set k, Map k state)
-> (edge, k)
-> (Seq k, Set k, Map k state)
update node
u state
nodeOut (Seq k
ws, Set k
inWS, Map k state
states) (edge
e, k
v) =
        let edgeOut :: state
edgeOut = Dataflow node edge state -> node -> edge -> state -> state
forall node edge state.
Dataflow node edge state -> node -> edge -> state -> state
edgeTransfer Dataflow node edge state
problem node
u edge
e state
nodeOut
            oldIn :: Maybe state
oldIn = k -> Map k state -> Maybe state
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
v Map k state
states
            newIn :: state
newIn = case Maybe state
oldIn of
                Just state
s  -> Dataflow node edge state -> state -> state -> state
forall node edge state.
Dataflow node edge state -> state -> state -> state
merge Dataflow node edge state
problem state
s state
edgeOut
                Maybe state
Nothing -> state
edgeOut
        in if state -> Maybe state
forall a. a -> Maybe a
Just state
newIn Maybe state -> Maybe state -> Bool
forall a. Eq a => a -> a -> Bool
/= Maybe state
oldIn
           then (if k
v k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set k
inWS then Seq k
ws else Seq k
ws Seq k -> k -> Seq k
forall a. Seq a -> a -> Seq a
|> k
v, k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert k
v Set k
inWS, k -> state -> Map k state -> Map k state
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
v state
newIn Map k state
states)
           else (Seq k
ws, Set k
inWS, Map k state
states)

-- | Solves the backward dataflow problem using a worklist algorithm.
solveBackward :: (Ord node, Eq state)
              => [node]                      -- ^ Exit nodes
              -> Map node [(edge, node)]     -- ^ The CFG
              -> Dataflow node edge state
              -> Map node state
solveBackward :: [node]
-> Map node [(edge, node)]
-> Dataflow node edge state
-> Map node state
solveBackward [node]
exits Map node [(edge, node)]
cfg Dataflow node edge state
problem =
    let startStates :: Map node state
startStates = [(node, state)] -> Map node state
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (node
exit, Dataflow node edge state -> state
forall node edge state. Dataflow node edge state -> state
initial Dataflow node edge state
problem) | node
exit <- [node]
exits ]
        worklist :: Seq node
worklist = [node] -> Seq node
forall a. [a] -> Seq a
Seq.fromList [node]
exits
        inWorklist :: Set node
inWorklist = [node] -> Set node
forall a. Ord a => [a] -> Set a
Set.fromList [node]
exits
    in Seq node -> Set node -> Map node state -> Map node state
loop Seq node
worklist Set node
inWorklist Map node state
startStates
  where
    revCfg :: Map node [(edge, node)]
revCfg = Map node [(edge, node)] -> Map node [(edge, node)]
forall node edge.
Ord node =>
Map node [(edge, node)] -> Map node [(edge, node)]
reverseCfg Map node [(edge, node)]
cfg

    loop :: Seq node -> Set node -> Map node state -> Map node state
loop Seq node
Empty Set node
_ Map node state
states = Map node state
states
    loop (node
u :<| Seq node
ws) Set node
inWS Map node state
states =
        let currentOut :: state
currentOut = state -> node -> Map node state -> state
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault (Dataflow node edge state -> state
forall node edge state. Dataflow node edge state -> state
initial Dataflow node edge state
problem) node
u Map node state
states
            nodeIn :: state
nodeIn = Dataflow node edge state -> node -> state -> state
forall node edge state.
Dataflow node edge state -> node -> state -> state
transfer Dataflow node edge state
problem node
u state
currentOut
            predecessors :: [(edge, node)]
predecessors = [(edge, node)] -> node -> Map node [(edge, node)] -> [(edge, node)]
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] node
u Map node [(edge, node)]
revCfg
            (Seq node
nextWS, Set node
nextInWS, Map node state
nextStates) = ((Seq node, Set node, Map node state)
 -> (edge, node) -> (Seq node, Set node, Map node state))
-> (Seq node, Set node, Map node state)
-> [(edge, node)]
-> (Seq node, Set node, Map node state)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (node
-> state
-> (Seq node, Set node, Map node state)
-> (edge, node)
-> (Seq node, Set node, Map node state)
forall k.
Ord k =>
node
-> state
-> (Seq k, Set k, Map k state)
-> (edge, k)
-> (Seq k, Set k, Map k state)
update node
u state
nodeIn) (Seq node
ws, node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.delete node
u Set node
inWS, Map node state
states) [(edge, node)]
predecessors
        in Seq node -> Set node -> Map node state -> Map node state
loop Seq node
nextWS Set node
nextInWS Map node state
nextStates

    update :: node
-> state
-> (Seq k, Set k, Map k state)
-> (edge, k)
-> (Seq k, Set k, Map k state)
update node
u state
nodeIn (Seq k
ws, Set k
inWS, Map k state
states) (edge
e, k
v) =
        let edgeIn :: state
edgeIn = Dataflow node edge state -> node -> edge -> state -> state
forall node edge state.
Dataflow node edge state -> node -> edge -> state -> state
edgeTransfer Dataflow node edge state
problem node
u edge
e state
nodeIn
            oldOut :: Maybe state
oldOut = k -> Map k state -> Maybe state
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
v Map k state
states
            newOut :: state
newOut = case Maybe state
oldOut of
                Just state
s  -> Dataflow node edge state -> state -> state -> state
forall node edge state.
Dataflow node edge state -> state -> state -> state
merge Dataflow node edge state
problem state
s state
edgeIn
                Maybe state
Nothing -> state
edgeIn
        in if state -> Maybe state
forall a. a -> Maybe a
Just state
newOut Maybe state -> Maybe state -> Bool
forall a. Eq a => a -> a -> Bool
/= Maybe state
oldOut
           then (if k
v k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set k
inWS then Seq k
ws else Seq k
ws Seq k -> k -> Seq k
forall a. Seq a -> a -> Seq a
|> k
v, k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert k
v Set k
inWS, k -> state -> Map k state -> Map k state
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
v state
newOut Map k state
states)
           else (Seq k
ws, Set k
inWS, Map k state
states)

    reverseCfg :: Ord node => Map node [(edge, node)] -> Map node [(edge, node)]
    reverseCfg :: Map node [(edge, node)] -> Map node [(edge, node)]
reverseCfg Map node [(edge, node)]
c = ([(edge, node)] -> [(edge, node)] -> [(edge, node)])
-> [(node, [(edge, node)])] -> Map node [(edge, node)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [(edge, node)] -> [(edge, node)] -> [(edge, node)]
forall a. [a] -> [a] -> [a]
(++)
        [ (node
v, [(edge
e, node
u)]) | (node
u, [(edge, node)]
edges) <- Map node [(edge, node)] -> [(node, [(edge, node)])]
forall k a. Map k a -> [(k, a)]
Map.toList Map node [(edge, node)]
c, (edge
e, node
v) <- [(edge, node)]
edges ]