CFG.hs (10780B)
1 {-# OPTIONS_HADDOCK prune #-} 2 {-# LANGUAGE DeriveAnyClass #-} 3 {-# LANGUAGE DeriveGeneric #-} 4 {-# LANGUAGE OverloadedStrings #-} 5 6 -- | 7 -- Module: Audit.AArch64.CFG 8 -- Copyright: (c) 2025 Jared Tobin 9 -- License: MIT 10 -- Maintainer: jared@ppad.tech 11 -- 12 -- Control flow graph construction for AArch64 assembly. 13 14 module Audit.AArch64.CFG ( 15 BasicBlock(..) 16 , CFG(..) 17 , buildCFG 18 , blockLabels 19 , blockSuccessors 20 , cfgBlockCount 21 , indexBlock 22 , blockFunction 23 -- * Function partitioning 24 , isFunctionLabel 25 , functionBlocks 26 , functionLabels 27 , callTargets 28 , buildCallGraph 29 -- * Cached analysis structures (deprecated) 30 , buildFunctionBlocksMap 31 , buildCallerMap 32 ) where 33 34 import Audit.AArch64.Runtime (RuntimeConfig(..)) 35 import Audit.AArch64.Types 36 import Control.DeepSeq (NFData) 37 import Data.Map.Strict (Map) 38 import Data.Maybe (listToMaybe) 39 import qualified Data.Map.Strict as Map 40 import Data.Primitive.Array (Array) 41 import qualified Data.Primitive.Array as A 42 import Data.Set (Set) 43 import qualified Data.Set as Set 44 import Data.Text (Text) 45 import qualified Data.Text as T 46 import GHC.Generics (Generic) 47 48 -- | A basic block: a sequence of instructions with a single entry point. 49 data BasicBlock = BasicBlock 50 { bbLabel :: !(Maybe Text) -- ^ Optional label at block start 51 , bbLines :: ![Line] -- ^ Instructions in the block 52 , bbSuccs :: ![Text] -- ^ Successor block labels 53 , bbLastInstr :: !(Maybe Instr) -- ^ Last instruction (cached) 54 , bbSuccIdxs :: ![Int] -- ^ Successor block indices (cached) 55 , bbHasFallthrough :: !Bool -- ^ Block falls through to next 56 } deriving (Eq, Show, Generic, NFData) 57 58 -- | Control flow graph. 59 data CFG = CFG 60 { cfgBlocks :: !(Array BasicBlock) -- ^ All basic blocks (O(1) index) 61 , cfgLabelMap :: !(Map Text Int) -- ^ Label -> block index 62 , cfgEntry :: !Int -- ^ Entry block index 63 , cfgFuncBlocks :: !(Map Text [Int]) -- ^ Function label -> block indices 64 } deriving (Eq, Show, Generic, NFData) 65 66 -- | Get all labels defined in the CFG. 67 blockLabels :: CFG -> Set Text 68 blockLabels cfg = Set.fromList 69 [ lbl | bb <- toList (cfgBlocks cfg), Just lbl <- [bbLabel bb] ] 70 where 71 toList arr = [A.indexArray arr i | i <- [0 .. A.sizeofArray arr - 1]] 72 73 -- | Get the number of blocks in the CFG. 74 cfgBlockCount :: CFG -> Int 75 cfgBlockCount cfg = A.sizeofArray (cfgBlocks cfg) 76 77 -- | Index into the CFG's block array. 78 indexBlock :: CFG -> Int -> BasicBlock 79 indexBlock cfg i = A.indexArray (cfgBlocks cfg) i 80 81 -- | Find the enclosing function label for a block index. 82 -- Returns Nothing if the block doesn't belong to any function. 83 blockFunction :: CFG -> Int -> Maybe Text 84 blockFunction cfg idx = 85 listToMaybe [ func | (func, idxs) <- Map.toList (cfgFuncBlocks cfg) 86 , idx `elem` idxs ] 87 88 -- | Build a CFG from parsed assembly lines. 89 buildCFG :: RuntimeConfig -> [Line] -> CFG 90 buildCFG rt lns = cfg 91 where 92 -- Split into basic blocks at labels and control flow instructions 93 rawBlocks = buildBlocks lns 94 nBlocks = length rawBlocks 95 96 -- Build label map first 97 labelMap = Map.fromList 98 [ (lbl, idx) 99 | (idx, bb) <- zip [0..] rawBlocks 100 , Just lbl <- [bbLabel bb] 101 ] 102 103 -- Compute successor indices and fallthrough for each block 104 annotatedBlocks = 105 [ bb { bbSuccIdxs = succIdxs, bbHasFallthrough = fallthrough } 106 | (idx, bb) <- zip [0..] rawBlocks 107 , let jumpTargets = [ i | lbl <- bbSuccs bb 108 , Just i <- [Map.lookup lbl labelMap] ] 109 fallthrough = hasFallthroughInstr (bbLastInstr bb) 110 fallthroughEdge = if fallthrough && idx + 1 < nBlocks 111 then [idx + 1] 112 else [] 113 succIdxs = jumpTargets ++ fallthroughEdge 114 ] 115 116 -- Build the block array 117 blocks = A.arrayFromList annotatedBlocks 118 119 -- Build function blocks map in a single pass 120 funcBlocksMap = buildFuncBlocksOnce rt annotatedBlocks 121 122 cfg = CFG 123 { cfgBlocks = blocks 124 , cfgLabelMap = labelMap 125 , cfgEntry = 0 126 , cfgFuncBlocks = funcBlocksMap 127 } 128 129 -- | Check if an instruction falls through to the next block. 130 hasFallthroughInstr :: Maybe Instr -> Bool 131 hasFallthroughInstr Nothing = True 132 hasFallthroughInstr (Just instr) = case instr of 133 B _ -> False -- Unconditional branch 134 Br _ -> False -- Indirect branch 135 Ret _ -> False -- Return 136 _ -> True -- Conditional branches, calls fall through 137 138 -- | Build function blocks map in a single pass over blocks. 139 buildFuncBlocksOnce :: RuntimeConfig -> [BasicBlock] 140 -> Map Text [Int] 141 buildFuncBlocksOnce rt bbs = 142 finalize $ foldl' step (Nothing, Map.empty) indexed 143 where 144 nBlocks = length bbs 145 indexed = zip [0..] bbs 146 147 step (mCur, acc) (idx, bb) = 148 case bbLabel bb of 149 Just lbl | isFunctionLabel rt lbl -> 150 let acc' = closeCurrent idx mCur acc 151 in (Just (lbl, idx), acc') 152 _ -> (mCur, acc) 153 154 closeCurrent _ Nothing acc = acc 155 closeCurrent endIdx (Just (lbl, start)) acc = 156 Map.insert lbl [start .. endIdx - 1] acc 157 158 finalize (mCur, acc) = closeCurrent nBlocks mCur acc 159 160 -- | Build basic blocks from lines. 161 buildBlocks :: [Line] -> [BasicBlock] 162 buildBlocks [] = [] 163 buildBlocks lns = go [] Nothing lns 164 where 165 go acc mLabel [] = 166 if null acc 167 then [] 168 else [finishBlock mLabel (reverse acc)] 169 170 go acc mLabel (l:ls) = 171 case (lineLabel l, lineInstr l) of 172 -- New label starts a new block 173 (Just lbl, _) -> 174 let current = if null acc then [] else [finishBlock mLabel (reverse acc)] 175 in current ++ go [l] (Just lbl) ls 176 177 -- Control flow instruction ends a block 178 (Nothing, Just instr) | isTerminator instr -> 179 let block = finishBlock mLabel (reverse (l:acc)) 180 in block : go [] Nothing ls 181 182 -- Regular instruction 183 _ -> go (l:acc) mLabel ls 184 185 finishBlock mLabel lns' = 186 let lastI = lastInstr lns' 187 in BasicBlock 188 { bbLabel = mLabel 189 , bbLines = lns' 190 , bbSuccs = successorLabels lastI 191 , bbLastInstr = lastI 192 , bbSuccIdxs = [] -- Filled in by buildCFG 193 , bbHasFallthrough = True -- Filled in by buildCFG 194 } 195 196 lastInstr [] = Nothing 197 lastInstr xs = lineInstr (last xs) 198 199 -- | Check if an instruction is a control flow terminator. 200 isTerminator :: Instr -> Bool 201 isTerminator instr = case instr of 202 B _ -> True 203 BCond _ _ -> True 204 Bl _ -> True 205 Blr _ -> True 206 Br _ -> True 207 Ret _ -> True 208 Cbz _ _ -> True 209 Cbnz _ _ -> True 210 Tbz _ _ _ -> True 211 Tbnz _ _ _ -> True 212 _ -> False 213 214 -- | Extract successor labels from a terminating instruction. 215 -- Note: Bl/Blr are NOT included as successors; calls are interprocedural 216 -- boundaries. Fallthrough after call is handled by hasFallthrough. 217 -- B to function labels (tail calls) ARE included - they create CFG edges 218 -- for intra-file dataflow propagation. 219 successorLabels :: Maybe Instr -> [Text] 220 successorLabels Nothing = [] 221 successorLabels (Just instr) = case instr of 222 B lbl -> [lbl] -- Includes both local branches and tail calls 223 BCond _ lbl -> [lbl] -- Plus fallthrough 224 Bl _ -> [] -- Call; fallthrough handled separately 225 Blr _ -> [] -- Indirect call 226 Br _ -> [] -- Indirect branch 227 Ret _ -> [] -- Return 228 Cbz _ lbl -> [lbl] -- Plus fallthrough 229 Cbnz _ lbl -> [lbl] -- Plus fallthrough 230 Tbz _ _ lbl -> [lbl] -- Plus fallthrough 231 Tbnz _ _ lbl -> [lbl] -- Plus fallthrough 232 _ -> [] 233 234 -- | Get successor block indices for a given block. 235 -- Returns cached successor indices (precomputed during CFG construction). 236 blockSuccessors :: CFG -> Int -> [Int] 237 blockSuccessors cfg idx 238 | idx < 0 || idx >= cfgBlockCount cfg = [] 239 | otherwise = bbSuccIdxs (indexBlock cfg idx) 240 241 -- | Check if a label is a function entry (not a local label). 242 -- LLVM local label checks are base; runtime-specific checks 243 -- are delegated to 'rtIsLocalLabel'. 244 isFunctionLabel :: RuntimeConfig -> Text -> Bool 245 isFunctionLabel rt lbl 246 | T.isPrefixOf "LBB" lbl = False -- LLVM basic block 247 | T.isPrefixOf "Lloh" lbl = False -- LLVM linker hint 248 | T.isPrefixOf "LCPI" lbl = False -- LLVM constant pool 249 | T.isPrefixOf "ltmp" lbl = False -- LLVM temporary 250 | T.isPrefixOf "l_" lbl = False -- Local label 251 | rtIsLocalLabel rt lbl = False -- Runtime-specific 252 | otherwise = True -- Likely a function 253 254 -- | Get all function entry labels in the CFG. 255 -- Returns keys from the cached function blocks map. 256 functionLabels :: CFG -> [Text] 257 functionLabels cfg = Map.keys (cfgFuncBlocks cfg) 258 259 -- | Get block indices belonging to a function. 260 -- Returns cached block indices (precomputed during CFG construction). 261 functionBlocks :: CFG -> Text -> [Int] 262 functionBlocks cfg funcLabel = 263 Map.findWithDefault [] funcLabel (cfgFuncBlocks cfg) 264 265 -- | Extract call targets from a block's instructions. 266 -- Includes both bl (call with link) and b to function labels (tail call). 267 callTargets :: RuntimeConfig -> BasicBlock -> [Text] 268 callTargets rt bb = 269 [ target | l <- bbLines bb 270 , Just instr <- [lineInstr l] 271 , target <- getCallTarget instr ] 272 where 273 getCallTarget (Bl target) = [target] 274 getCallTarget (B target) 275 | isFunctionLabel rt target = [target] -- Tail call 276 getCallTarget _ = [] 277 278 -- | Build call graph: maps each function to its callees. 279 -- Uses cached function blocks map for efficient lookup. 280 buildCallGraph :: RuntimeConfig -> CFG -> Map Text [Text] 281 buildCallGraph rt cfg = Map.fromList 282 [ (func, callees) 283 | (func, indices) <- Map.toList (cfgFuncBlocks cfg) 284 , let callees = concatMap 285 (callTargets rt . indexBlock cfg) indices 286 ] 287 288 -- | Build map of function labels to their block index ranges. 289 -- Deprecated: use cfgFuncBlocks directly; this returns the cached map. 290 buildFunctionBlocksMap :: CFG -> Map Text [Int] 291 buildFunctionBlocksMap = cfgFuncBlocks 292 293 -- | Build caller map: maps each function to its callers. 294 -- Takes precomputed function blocks map to avoid rescanning. 295 buildCallerMap :: RuntimeConfig -> CFG -> Map Text [Int] 296 -> Map Text [Text] 297 buildCallerMap rt cfg funcBlocksMap = 298 Map.fromListWith (++) 299 [ (callee, [caller]) 300 | (caller, callees) <- Map.toList callGraph 301 , callee <- callees 302 ] 303 where 304 callGraph = Map.fromList 305 [ (func, concatMap 306 (callTargets rt . indexBlock cfg) idxs) 307 | (func, idxs) <- Map.toList funcBlocksMap 308 ]