From f1148b8b41087dba8441a2c79980441522e3b23b Mon Sep 17 00:00:00 2001
From: Vidar Holen <vidar@vidarholen.net>
Date: Mon, 25 Jul 2022 10:00:50 -0700
Subject: [PATCH] Include postdominators in CFGResult

---
 src/ShellCheck/CFG.hs         | 86 ++++++++++++++++++++++++++++++++---
 src/ShellCheck/CFGAnalysis.hs |  5 +-
 2 files changed, 84 insertions(+), 7 deletions(-)

diff --git a/src/ShellCheck/CFG.hs b/src/ShellCheck/CFG.hs
index 6f6d4f1..ad05e93 100644
--- a/src/ShellCheck/CFG.hs
+++ b/src/ShellCheck/CFG.hs
@@ -54,6 +54,8 @@ import qualified Data.Set as S
 import Control.Monad.RWS.Lazy
 import Data.Graph.Inductive.Graph
 import Data.Graph.Inductive.Query.DFS
+import Data.Graph.Inductive.Basic
+import Data.Graph.Inductive.Query.Dominators
 import Data.Graph.Inductive.PatriciaTree as G
 import Debug.Trace -- STRIP
 
@@ -171,9 +173,11 @@ data CFGResult = CFGResult {
     -- Map from Id to nominal start&end node (i.e. assuming normal execution without exits)
     cfIdToRange :: M.Map Id (Node, Node),
     -- A set of all nodes belonging to an Id, recursively
-    cfIdToNodes :: M.Map Id (S.Set Node)
+    cfIdToNodes :: M.Map Id (S.Set Node),
+    -- A map to nodes that the given node postdominates
+    cfPostDominators :: M.Map Node (S.Set Node)
 }
-  deriving (Show)
+  deriving (Show, Generic, NFData)
 
 buildGraph :: CFGParameters -> Token -> CFGResult
 buildGraph params root =
@@ -183,12 +187,20 @@ buildGraph params root =
 --            renumberTopologically $
                 removeUnnecessaryStructuralNodes
                     base
-    in
-        CFGResult {
+
+        idToRange = M.fromList mapping
+        isRealEdge (from, to, edge) = case edge of CFEFlow -> True; _ -> False
+        onlyRealEdges = filter isRealEdge edges
+        (_, mainExit) = fromJust $ M.lookup (getId root) idToRange
+
+        result = CFGResult {
             cfGraph = mkGraph nodes edges,
-            cfIdToRange = M.fromList mapping,
-            cfIdToNodes = M.fromListWith S.union $ map (\(id, n) -> (id, S.singleton n)) association
+            cfIdToRange = idToRange,
+            cfIdToNodes = M.fromListWith S.union $ map (\(id, n) -> (id, S.singleton n)) association,
+            cfPostDominators = findPostDominators mainExit $ mkGraph nodes onlyRealEdges
         }
+    in
+        deepseq result result
 
 remapGraph :: M.Map Node Node -> CFW -> CFW
 remapGraph remap (nodes, edges, mapping, assoc) =
@@ -1190,5 +1202,67 @@ tokenToParts t =
         -- Check if getLiteralString can handle it, if not it's unknown
         _ -> [maybe CFStringUnknown CFStringLiteral $ getLiteralString t]
 
+
+-- Change all subshell invocations to instead link directly to their contents.
+-- This is used for producing dominator trees.
+inlineSubshells :: CFGraph -> CFGraph
+inlineSubshells graph = relinkedGraph
+  where
+    subshells = ufold find [] graph
+    find (incoming, node, label, outgoing) acc =
+        case label of
+            CFExecuteSubshell _ start end -> (node, label, start, end, incoming, outgoing):acc
+            _ -> acc
+
+    relinkedGraph = foldl' relink graph subshells
+    relink graph (node, label, start, end, incoming, outgoing) =
+        let
+            -- Link CFExecuteSubshell to the CFEntryPoint
+            subshellToStart = (incoming, node, label, [(CFEFlow, start)])
+            -- Link the subshell exit to the
+            endToNexts = (endIncoming, endNode, endLabel, outgoing)
+            (endIncoming, endNode, endLabel, _) = context graph end
+        in
+            subshellToStart & (endToNexts & graph)
+
+findEntryNodes :: CFGraph -> [Node]
+findEntryNodes graph = ufold find [] graph
+  where
+    find (incoming, node, label, _) list =
+        case label of
+            CFEntryPoint {} | null incoming -> node:list
+            _ -> list
+
+findDominators main graph = asSetMap
+  where
+    inlined = inlineSubshells graph
+    entryNodes = main : findEntryNodes graph
+    asLists = concatMap (dom inlined) entryNodes
+    asSetMap = M.fromList $ map (\(node, list) -> (node, S.fromList list)) asLists
+
+findTerminalNodes :: CFGraph -> [Node]
+findTerminalNodes graph = ufold find [] graph
+  where
+    find (_, node, label, _) list =
+        case label of
+            CFUnresolvedExit -> node:list
+            CFApplyEffects effects -> f effects list
+            _ -> list
+
+    f [] list = list
+    f (IdTagged _ (CFDefineFunction _ id start end):rest) list = f rest (end:list)
+    f (_:rest) list = f rest list
+
+findPostDominators :: Node -> CFGraph -> M.Map Node (S.Set Node)
+findPostDominators mainexit graph = asSetMap
+  where
+    inlined = inlineSubshells graph
+    terminals = findTerminalNodes inlined
+    (incoming, _, label, outgoing) = context graph mainexit
+    withExitEdges = (incoming ++ map (\c -> (CFEFlow, c)) terminals, mainexit, label, outgoing) & inlined
+    reversed = grev withExitEdges
+    postDoms = dom reversed mainexit
+    asSetMap = M.fromList $ map (\(node, list) -> (node, S.fromList list)) postDoms
+
 return []
 runTests =  $( [| $(forAllProperties) (quickCheckWithResult (stdArgs { maxSuccess = 1 }) ) |])
diff --git a/src/ShellCheck/CFGAnalysis.hs b/src/ShellCheck/CFGAnalysis.hs
index bb90860..dc0a4b1 100644
--- a/src/ShellCheck/CFGAnalysis.hs
+++ b/src/ShellCheck/CFGAnalysis.hs
@@ -99,6 +99,7 @@ data CFGAnalysis = CFGAnalysis {
     graph :: CFGraph,
     tokenToRange :: M.Map Id (Node, Node),
     tokenToNodes :: M.Map Id (S.Set Node),
+    postDominators :: M.Map Node (S.Set Node),
     nodeToData :: M.Map Node (ProgramState, ProgramState)
 } deriving (Show, Generic, NFData)
 
@@ -1304,7 +1305,8 @@ analyzeControlFlow params t =
             graph = cfGraph cfg,
             tokenToRange = cfIdToRange cfg,
             tokenToNodes = cfIdToNodes cfg,
-            nodeToData = nodeToData
+            nodeToData = nodeToData,
+            postDominators = cfPostDominators cfg
         }
 
 
@@ -1355,5 +1357,6 @@ analyzeStragglers ctx state stragglers = do
         transferFunctionValue ctx def
 
 
+
 return []
 runTests = $quickCheckAll