-- SPDX-FileCopyrightText: 2025 Sören Tempel <soeren+git@soeren-tempel.net>
-- SPDX-FileCopyrightText: 2026 Reliable System Software, Technische Universität Braunschweig <vss@ibr.cs.tu-bs.de>
--
-- SPDX-License-Identifier: GPL-3.0-only

module Language.QBE.Analysis.CFG
  ( -- * Control Flow Graph
    Label,
    CFG (cfgFunction),
    build,
    identToLabel,
    labelToIdent,
    labelToBlock,
    lookupSuccs,

    -- * Graph Representation
    asGraph,
    nodes,
    edges,
    bounds,

    -- * Dominator Analysis
    asDomGraph,
    startNode,
  )
where

import Data.Graph (Bounds, Graph, buildG)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.IntSet qualified as IntSet
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromJust)
import Data.Tuple (swap)
import Language.QBE.Analysis.Graph qualified as DG
import Language.QBE.Types qualified as QBE

-- | Representation of a node in the t'CFG'.
type Label = IntMap.Key

-- | A representation of the control-flow within a 'QBE.FuncDef'.
data CFG
  = CFG
  { -- | Function for which this CFG was built.
    CFG -> FuncDef
cfgFunction :: QBE.FuncDef,
    CFG -> Label
cfgMaxBound :: Int,
    CFG -> Map BlockIdent Label
cfgLabelMap :: Map QBE.BlockIdent Label,
    CFG -> IntMap BlockIdent
cfgBlockMap :: IntMap QBE.BlockIdent,
    CFG -> IntMap [Label]
cfgSuccessors :: IntMap [Label]
  }

-- | Returns a list of all graph nodes in an unspecified order.
nodes :: CFG -> [Label]
nodes :: CFG -> [Label]
nodes = IntMap BlockIdent -> [Label]
forall a. IntMap a -> [Label]
IntMap.keys (IntMap BlockIdent -> [Label])
-> (CFG -> IntMap BlockIdent) -> CFG -> [Label]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFG -> IntMap BlockIdent
cfgBlockMap

-- | Returns a list of graph edges in an unspecified order.
edges :: CFG -> [(Label, Label)]
edges :: CFG -> [(Label, Label)]
edges CFG
cfg = ([(Label, Label)] -> (Label, [Label]) -> [(Label, Label)])
-> [(Label, Label)] -> [(Label, [Label])] -> [(Label, Label)]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl [(Label, Label)] -> (Label, [Label]) -> [(Label, Label)]
forall {t} {a}. [(t, a)] -> (t, [a]) -> [(t, a)]
go [] ([(Label, [Label])] -> [(Label, Label)])
-> [(Label, [Label])] -> [(Label, Label)]
forall a b. (a -> b) -> a -> b
$ IntMap [Label] -> [(Label, [Label])]
forall a. IntMap a -> [(Label, a)]
IntMap.toList (CFG -> IntMap [Label]
cfgSuccessors CFG
cfg)
  where
    go :: [(t, a)] -> (t, [a]) -> [(t, a)]
go [(t, a)]
acc (t
p, [a]
c) = [(t, a)]
acc [(t, a)] -> [(t, a)] -> [(t, a)]
forall a. [a] -> [a] -> [a]
++ (a -> (t, a)) -> [a] -> [(t, a)]
forall a b. (a -> b) -> [a] -> [b]
map (t
p,) [a]
c

-- | Returns the bounds of the t'CFG'. This is useful, for example, to
-- build a subgraph using 'Data.Graph.buildG'.
bounds :: CFG -> Bounds
bounds :: CFG -> (Label, Label)
bounds CFG
cfg = (Label
0, CFG -> Label
cfgMaxBound CFG
cfg)

-- | Convert a 'QBE.BlockIdent' to a CFG node 'Label'.
--
-- This function is partial, on an invalid 'Label', an error is thrown.
identToLabel :: CFG -> QBE.BlockIdent -> Label
identToLabel :: CFG -> BlockIdent -> Label
identToLabel CFG {cfgLabelMap :: CFG -> Map BlockIdent Label
cfgLabelMap = Map BlockIdent Label
m} BlockIdent
blkId =
  Maybe Label -> Label
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Label -> Label) -> Maybe Label -> Label
forall a b. (a -> b) -> a -> b
$ BlockIdent -> Map BlockIdent Label -> Maybe Label
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup BlockIdent
blkId Map BlockIdent Label
m

-- | Convert a CFG node 'Label' to a 'QBE.BlockIdent'.
--
-- This function is partial, on an invalid 'Label', an error is thrown.
labelToIdent :: CFG -> Label -> QBE.BlockIdent
labelToIdent :: CFG -> Label -> BlockIdent
labelToIdent CFG {cfgBlockMap :: CFG -> IntMap BlockIdent
cfgBlockMap = IntMap BlockIdent
m} Label
label =
  Maybe BlockIdent -> BlockIdent
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe BlockIdent -> BlockIdent) -> Maybe BlockIdent -> BlockIdent
forall a b. (a -> b) -> a -> b
$ Label -> IntMap BlockIdent -> Maybe BlockIdent
forall a. Label -> IntMap a -> Maybe a
IntMap.lookup Label
label IntMap BlockIdent
m

-- | Utility function to convert a node 'Label' to a 'QBE.Block'.
-- Performs two \(O(\log n)\) lookups internally.
--
-- This function is partial, on an invalid 'Label', an error is thrown.
labelToBlock :: CFG -> Label -> QBE.Block
labelToBlock :: CFG -> Label -> Block
labelToBlock CFG
cfg Label
label =
  let blocks :: Map BlockIdent Block
blocks = FuncDef -> Map BlockIdent Block
QBE.fBlock (FuncDef -> Map BlockIdent Block)
-> FuncDef -> Map BlockIdent Block
forall a b. (a -> b) -> a -> b
$ CFG -> FuncDef
cfgFunction CFG
cfg
   in Maybe Block -> Block
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Block -> Block) -> Maybe Block -> Block
forall a b. (a -> b) -> a -> b
$ BlockIdent -> Map BlockIdent Block -> Maybe Block
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (CFG -> Label -> BlockIdent
labelToIdent CFG
cfg Label
label) Map BlockIdent Block
blocks

-- | Mapping of 'Label' to its successors in the CFG, represented as an
-- ordered list of zero, one, or two elements. A list with two elements
-- represents a conditional jump where the left child is the is the true
-- branch and the right child is the false branch. A list wih a single
-- element signifies an unconditional jump. If the given node does not
-- have any successors an empty list is returned.
--
-- This function is partial, on an invalid 'Label', an error is thrown.
lookupSuccs :: CFG -> Label -> [Label]
lookupSuccs :: CFG -> Label -> [Label]
lookupSuccs CFG {cfgSuccessors :: CFG -> IntMap [Label]
cfgSuccessors = IntMap [Label]
succs} Label
label =
  Maybe [Label] -> [Label]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Label] -> [Label]) -> Maybe [Label] -> [Label]
forall a b. (a -> b) -> a -> b
$ Label -> IntMap [Label] -> Maybe [Label]
forall a. Label -> IntMap a -> Maybe a
IntMap.lookup Label
label IntMap [Label]
succs

------------------------------------------------------------------------

identStart :: Label
identStart :: Label
identStart = Label
0

-- | Construct a t'CFG' for a given function.
build :: QBE.FuncDef -> CFG
build :: FuncDef -> CFG
build FuncDef
func =
  CFG
    { cfgMaxBound :: Label
cfgMaxBound = (BlockIdent, Label) -> Label
forall a b. (a, b) -> b
snd ((BlockIdent, Label) -> Label) -> (BlockIdent, Label) -> Label
forall a b. (a -> b) -> a -> b
$ [(BlockIdent, Label)] -> (BlockIdent, Label)
forall a. HasCallStack => [a] -> a
last [(BlockIdent, Label)]
blkIdLabels,
      cfgFunction :: FuncDef
cfgFunction = FuncDef
func,
      cfgLabelMap :: Map BlockIdent Label
cfgLabelMap = Map BlockIdent Label
labelMap,
      cfgBlockMap :: IntMap BlockIdent
cfgBlockMap = [(Label, BlockIdent)] -> IntMap BlockIdent
forall a. [(Label, a)] -> IntMap a
IntMap.fromList ([(Label, BlockIdent)] -> IntMap BlockIdent)
-> [(Label, BlockIdent)] -> IntMap BlockIdent
forall a b. (a -> b) -> a -> b
$ ((BlockIdent, Label) -> (Label, BlockIdent))
-> [(BlockIdent, Label)] -> [(Label, BlockIdent)]
forall a b. (a -> b) -> [a] -> [b]
map (BlockIdent, Label) -> (Label, BlockIdent)
forall a b. (a, b) -> (b, a)
swap [(BlockIdent, Label)]
blkIdLabels,
      cfgSuccessors :: IntMap [Label]
cfgSuccessors = [(Label, [Label])] -> IntMap [Label]
forall a. [(Label, a)] -> IntMap a
IntMap.fromList ([(Label, [Label])] -> IntMap [Label])
-> [(Label, [Label])] -> IntMap [Label]
forall a b. (a -> b) -> a -> b
$ Map BlockIdent Label -> [Block] -> [(Label, [Label])]
build' Map BlockIdent Label
labelMap [Block]
blocks
    }
  where
    labelMap :: Map QBE.BlockIdent Label
    labelMap :: Map BlockIdent Label
labelMap = [(BlockIdent, Label)] -> Map BlockIdent Label
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(BlockIdent, Label)]
blkIdLabels

    blocks :: [QBE.Block]
    blocks :: [Block]
blocks = Map BlockIdent Block -> [Block]
forall k a. Map k a -> [a]
Map.elems (Map BlockIdent Block -> [Block])
-> Map BlockIdent Block -> [Block]
forall a b. (a -> b) -> a -> b
$ FuncDef -> Map BlockIdent Block
QBE.fBlock FuncDef
func

    blkIdLabels :: [(QBE.BlockIdent, Label)]
    blkIdLabels :: [(BlockIdent, Label)]
blkIdLabels = [BlockIdent] -> [Label] -> [(BlockIdent, Label)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Block -> BlockIdent) -> [Block] -> [BlockIdent]
forall a b. (a -> b) -> [a] -> [b]
map Block -> BlockIdent
QBE.label [Block]
blocks) [Label
identStart ..]

build' :: Map QBE.BlockIdent Label -> [QBE.Block] -> [(IntMap.Key, [Label])]
build' :: Map BlockIdent Label -> [Block] -> [(Label, [Label])]
build' Map BlockIdent Label
labelMap = ([(Label, [Label])] -> Block -> [(Label, [Label])])
-> [(Label, [Label])] -> [Block] -> [(Label, [Label])]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl [(Label, [Label])] -> Block -> [(Label, [Label])]
go []
  where
    toLabel :: QBE.BlockIdent -> Label
    toLabel :: BlockIdent -> Label
toLabel BlockIdent
ident = Maybe Label -> Label
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Label -> Label) -> Maybe Label -> Label
forall a b. (a -> b) -> a -> b
$ BlockIdent -> Map BlockIdent Label -> Maybe Label
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup BlockIdent
ident Map BlockIdent Label
labelMap

    go :: [(Label, [Label])] -> Block -> [(Label, [Label])]
go [(Label, [Label])]
acc block :: Block
block@(QBE.Block {label :: Block -> BlockIdent
QBE.label = BlockIdent
ident}) =
      let succs :: [Label]
succs = case Block -> JumpInstr
QBE.term Block
block of
            QBE.Jump BlockIdent
target -> [BlockIdent -> Label
toLabel BlockIdent
target]
            QBE.Jnz Value
_ BlockIdent
i1 BlockIdent
i2 -> [BlockIdent -> Label
toLabel BlockIdent
i1, BlockIdent -> Label
toLabel BlockIdent
i2]
            QBE.Return Maybe Value
_ -> []
            JumpInstr
QBE.Halt -> []
       in (BlockIdent -> Label
toLabel BlockIdent
ident, [Label]
succs) (Label, [Label]) -> [(Label, [Label])] -> [(Label, [Label])]
forall a. a -> [a] -> [a]
: [(Label, [Label])]
acc

------------------------------------------------------------------------

asGraph :: CFG -> Graph
asGraph :: CFG -> Graph
asGraph CFG
cfg = (Label, Label) -> [(Label, Label)] -> Graph
buildG (Label
identStart, CFG -> Label
cfgMaxBound CFG
cfg) ([(Label, Label)] -> Graph) -> [(Label, Label)] -> Graph
forall a b. (a -> b) -> a -> b
$ CFG -> [(Label, Label)]
edges CFG
cfg

asDomGraph :: CFG -> DG.Graph
asDomGraph :: CFG -> Graph
asDomGraph CFG
cfg = ([Label] -> IntSet) -> IntMap [Label] -> Graph
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map [Label] -> IntSet
IntSet.fromList (CFG -> IntMap [Label]
cfgSuccessors CFG
cfg)

-- | Determine the entry node of the t'CFG'. Useful, for example, to
-- generated a 'DG.Rooted' representation for the control-flow graph.
startNode :: CFG -> Label
startNode :: CFG -> Label
startNode cfg :: CFG
cfg@(CFG {cfgFunction :: CFG -> FuncDef
cfgFunction = FuncDef
func}) =
  CFG -> BlockIdent -> Label
identToLabel CFG
cfg (FuncDef -> BlockIdent
QBE.fStart FuncDef
func)