CallGraph.hs (3721B)
1 {-# OPTIONS_HADDOCK prune #-} 2 {-# LANGUAGE OverloadedStrings #-} 3 4 -- | 5 -- Module: Audit.AArch64.CallGraph 6 -- Copyright: (c) 2025 Jared Tobin 7 -- License: MIT 8 -- Maintainer: jared@ppad.tech 9 -- 10 -- Inter-procedural call graph construction for AArch64 assembly. 11 -- Used to find all functions reachable from a given symbol. 12 13 module Audit.AArch64.CallGraph ( 14 -- * Call graph 15 CallGraph 16 , buildCallGraph 17 -- * Queries 18 , allSymbols 19 , reachableSymbols 20 , reachingSymbols 21 , symbolExists 22 ) where 23 24 import Audit.AArch64.CFG (isFunctionLabel) 25 import Audit.AArch64.Runtime (RuntimeConfig) 26 import Audit.AArch64.Types (Instr(..), Line(..)) 27 import Data.Graph 28 (Graph, Vertex, graphFromEdges, reachable, transposeG) 29 import Data.Map.Strict (Map) 30 import qualified Data.Map.Strict as Map 31 import Data.Set (Set) 32 import qualified Data.Set as Set 33 import Data.Text (Text) 34 35 -- | Call graph: maps symbols to the set of symbols they call. 36 -- Includes lookup functions for graph traversal. 37 data CallGraph = CallGraph 38 { cgGraph :: !Graph 39 , cgNodeFromV :: !(Vertex -> ((), Text, [Text])) 40 , cgVertexFromK :: !(Text -> Maybe Vertex) 41 , cgSymbols :: !(Set Text) 42 } 43 44 -- | Build a call graph from parsed assembly lines. 45 -- 46 -- Extracts function symbols and their call targets 47 -- (bl instructions). Does not resolve indirect calls (blr). 48 buildCallGraph :: RuntimeConfig -> [Line] -> CallGraph 49 buildCallGraph rt lns = 50 CallGraph graph nodeFromV vertexFromK allSyms 51 where 52 (graph, nodeFromV, vertexFromK) = 53 graphFromEdges edges 54 55 -- Build map from symbol to its instructions. 56 -- State: (current symbol, accumulated map) 57 symInstrs :: Map Text [Instr] 58 symInstrs = snd $ foldl' step ("<unknown>", Map.empty) lns 59 where 60 step (curSym, acc) ln = 61 let sym = case lineLabel ln of 62 Just lbl 63 | isFunctionLabel rt lbl -> lbl 64 _ -> curSym 65 acc' = Map.insertWith (++) sym [] acc 66 in case lineInstr ln of 67 Nothing -> (sym, acc') 68 Just i -> 69 (sym, Map.insertWith (++) sym [i] acc') 70 71 -- Extract call targets from instructions 72 callTargets :: Text -> [Text] 73 callTargets sym = case Map.lookup sym symInstrs of 74 Nothing -> [] 75 Just instrs -> [target | Bl target <- instrs] 76 77 -- All symbols in the assembly 78 allSyms :: Set Text 79 allSyms = Map.keysSet symInstrs 80 81 -- Build graph edges: (node data, key, [adjacent keys]) 82 edges :: [((), Text, [Text])] 83 edges = 84 [((), sym, callTargets sym) | sym <- Set.toList allSyms] 85 86 -- | Get all symbols reachable from a root symbol 87 -- (including the root). 88 -- 89 -- Returns empty set if the root symbol doesn't exist. 90 reachableSymbols :: Text -> CallGraph -> Set Text 91 reachableSymbols root cg = case cgVertexFromK cg root of 92 Nothing -> Set.empty 93 Just v -> Set.fromList 94 [sym | v' <- reachable (cgGraph cg) v 95 , let (_, sym, _) = cgNodeFromV cg v'] 96 97 -- | Get all symbols that can reach a target symbol 98 -- (including the target). 99 -- 100 -- This is the reverse of 'reachableSymbols': finds all 101 -- potential callers. Returns empty set if the target 102 -- symbol doesn't exist. 103 reachingSymbols :: Text -> CallGraph -> Set Text 104 reachingSymbols target cg = 105 case cgVertexFromK cg target of 106 Nothing -> Set.empty 107 Just v -> Set.fromList 108 [sym | v' <- reachable (transposeG (cgGraph cg)) v 109 , let (_, sym, _) = cgNodeFromV cg v'] 110 111 -- | Get all symbols in the call graph. 112 allSymbols :: CallGraph -> Set Text 113 allSymbols = cgSymbols 114 115 -- | Check if a symbol exists in the call graph. 116 symbolExists :: Text -> CallGraph -> Bool 117 symbolExists sym cg = Set.member sym (cgSymbols cg)