From 690985cb874537e3d149d874279056b95f7019da Mon Sep 17 00:00:00 2001
From: Enrico Lumetti <enrico.lumetti@gmail.com>
Date: Sun, 13 Oct 2024 20:22:53 +0200
Subject: [PATCH] Add CFG construction

---
 eoc.cabal       |  8 +++--
 lib/Bril/CFG.hs | 85 +++++++++++++++++++++++++++++++++++++++++++++++++
 tests/Main.hs   | 10 +++++-
 3 files changed, 99 insertions(+), 4 deletions(-)
 create mode 100644 lib/Bril/CFG.hs

diff --git a/eoc.cabal b/eoc.cabal
index db1eaa2..e835148 100644
--- a/eoc.cabal
+++ b/eoc.cabal
@@ -17,7 +17,7 @@ build-type:          Simple
 extra-source-files:  README.md
 
 library eoc-lib
-    exposed-modules: Eoc, Bril
+    exposed-modules: Eoc, Bril, Bril.CFG
     hs-source-dirs: lib
     build-depends:
         base >=4.13 && <4.18,
@@ -25,8 +25,9 @@ library eoc-lib
         aeson >= 2.2,
         sorted-list ^>= 0.2,
         bytestring ^>= 0.11,
-        scientific ^>= 0.3
-    default-language: Haskell2010
+        scientific ^>= 0.3,
+        containers ^>= 0.6.7
+  default-language: Haskell2010
   default-extensions: DeriveGeneric, OverloadedStrings
 
 executable eoc
@@ -46,6 +47,7 @@ test-suite tests
     build-depends:
         base >=4.13 && <4.18,
         HUnit >=1.6,
+        QuickCheck ^>=2.15,
         eoc-lib,
         text >= 2.0,
         aeson >= 2.2
diff --git a/lib/Bril/CFG.hs b/lib/Bril/CFG.hs
new file mode 100644
index 0000000..837d0af
--- /dev/null
+++ b/lib/Bril/CFG.hs
@@ -0,0 +1,85 @@
+module Bril.CFG (BasicBlock, CFG, createCFG) where
+
+import Bril
+import qualified Data.Map.Strict as Map
+
+data BasicBlock = BasicBlock {
+        instrs :: [Instruction],
+        incoming :: [String]
+    }
+   deriving (Show)
+
+data CFG = CFG {
+        cfgBlocks :: Map.Map String BasicBlock,
+        cfgStart :: String
+    }
+   deriving (Show)
+
+createCFG :: [Instruction] -> CFG
+createCFG instrs =
+    let
+        (initCFG, firstLabel) = initializeCFG $ addFallthroughJmp $ splitBBs $ addInitialLabel instrs
+        cfgBlocks' = foldr processBB (cfgBlocks initCFG) (Map.keys (cfgBlocks initCFG))
+    in
+        CFG cfgBlocks' firstLabel
+    where
+        processBB :: String -> Map.Map String BasicBlock -> Map.Map String BasicBlock
+        processBB label m = addBBJumps label (outgoingCFGLabels m label) m
+        addBBJumps :: String -> [String] -> Map.Map String BasicBlock -> Map.Map String BasicBlock
+        addBBJumps bblockLabel incomingLabels m = Map.adjust (addLabels incomingLabels) bblockLabel m
+        addLabels :: [String] -> BasicBlock -> BasicBlock
+        addLabels labels (BasicBlock instrs incoming) = (BasicBlock instrs (incoming ++ labels))
+
+initializeCFG :: [([Instruction], String)] -> (CFG, String)
+initializeCFG instrs =
+    let mappedBlocks = foldr mapBlock Map.empty instrs -- Map.fromList $ map swap instrs
+        firstLabel = snd . head $ instrs
+    in (CFG mappedBlocks firstLabel, firstLabel)
+    where
+        mapBlock (instrs, label) m = Map.insert label (BasicBlock instrs []) m
+
+addFallthroughJmp :: [([Instruction], String)] -> [([Instruction], String)]
+addFallthroughJmp blocks =
+    let blocks' = map addJumpToNextBB $ zip blocks $ tail $ map snd blocks
+    in blocks' ++ [last blocks]
+    where
+        addJumpToNextBB (block@(instrs, label), nextLabel)
+            -- TODO: handle empty instrs
+            | isJump (last instrs) = block
+            | otherwise = (instrs ++ [EffectInstr Jmp [] [] [nextLabel]], label)
+
+splitBBs :: [Instruction] -> [([Instruction], String)]
+splitBBs instrs = zip (tail basicBlocks) labels
+    where
+        (basicBlocks, labels) = foldr splitBBs' ([[]], []) instrs
+        splitBBs' instr (bblocks, labels)
+                | isLabel instr = ([] : bblocks, (labelName instr) : labels)
+                | isJump instr = ([instr] : (tail bblocks) , labels)
+                | otherwise = ((instr : (head bblocks)) : (tail bblocks), labels)
+
+addInitialLabel :: [Instruction] -> [Instruction]
+addInitialLabel instrs@(firstInstr : instrs')
+    | isLabel firstInstr = instrs
+    | otherwise = (Label "@start") : instrs
+addInitialLabel [] = []
+
+labelName (Label name) = name
+isLabel instr = case instr of
+    (Label _) -> True
+    _ -> False
+    
+isJump instr = case instr of
+    (EffectInstr op _ _ _) -> op == Jmp || op == Br
+    _ -> False
+
+outgoingCFGLabels :: Map.Map String BasicBlock -> String -> [String]
+outgoingCFGLabels m label =
+    case Map.lookup label m of
+        Just bb -> outgoingLabels bb
+        _ -> []
+
+outgoingLabels :: BasicBlock -> [String]
+outgoingLabels (BasicBlock instrs _) =
+    let (EffectInstr lastOp _ _ labels) = last instrs
+    in labels
+
diff --git a/tests/Main.hs b/tests/Main.hs
index a96407a..4307a92 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -1,7 +1,8 @@
 module Main where
 
 import qualified Bril as B
-import Data.Aeson (eitherDecode, decode, toJSON, parseJSON, Result(..))
+import qualified Bril.CFG as BCFG
+import Data.Aeson (eitherDecode, decode, encode, toJSON, parseJSON, Result(..))
 import Data.Aeson.Types (parse)
 import Data.Text
 import Data.Maybe
@@ -10,6 +11,7 @@ import Data.Function ((&))
 import Control.Exception (try)
 
 import Test.HUnit
+import qualified Test.QuickCheck as QC
 import qualified System.Exit as Exit
  
 testTypes :: Test
@@ -92,6 +94,10 @@ testProgramFromFiles = TestCase $ do
         (Right parseResult) -> assertBool "parse add.json" $ isRight parseResult
         _ -> assertFailure "cannot open file"
 
+prop_InvariantBrilSerialization :: B.Program -> Bool
+prop_InvariantBrilSerialization program =
+    (decode (encode program)) == (Just program)
+
 
 tests :: Test
 tests = TestList [
@@ -104,4 +110,6 @@ tests = TestList [
 main :: IO ()
 main = do
     result <- runTestTT tests
+
+    -- QC.quickCheck prop_InvariantBrilSerialization
     if failures result > 0 then Exit.exitFailure else Exit.exitSuccess