qute

 1-- SPDX-FileCopyrightText: 2026 Sören Tempel <soeren+git@soeren-tempel.net>
 2--
 3-- SPDX-License-Identifier: GPL-3.0-only
 4
 5-- | This module implements a "linker" for QBE source files. It takes multiple
 6-- 'Language.QBE.Program' values and merges them by resolving conflicting
 7-- non-exporting definitions by assising them a new (unique) name.
 8module Language.QBE.Linker (link) where
 9
10import Control.Monad (foldM)
11import Control.Monad.State (State, evalState, get, gets, modify)
12import Data.Map (Map)
13import Data.Map qualified as Map
14import Language.QBE (Definition (DefData, DefFile, DefFunc, DefType), Program, localDefs, typeDefs)
15import Language.QBE.Types qualified as QBE
16import Language.QBE.Uses qualified as USE
17
18data Env
19  = Env
20  { envGlobals :: Map QBE.GlobalIdent Int,
21    envTypes :: Map QBE.UserIdent Int
22  }
23
24mkEnv :: Env
25mkEnv = Env Map.empty Map.empty
26
27------------------------------------------------------------------------
28
29-- TODO: Some sort of type class for Idents?
30rename :: (Ord a) => (a -> Int -> a) -> Map a Int -> a -> a
31rename inc varOcc i = maybe i (inc' i) $ Map.lookup i varOcc
32  where
33    inc' ident n = if n >= 1 then inc ident n else ident
34
35renameGlobal :: Map QBE.GlobalIdent Int -> QBE.GlobalIdent -> QBE.GlobalIdent
36renameGlobal = rename (\(QBE.GlobalIdent s) n -> QBE.GlobalIdent $ s ++ ".QG" ++ show n)
37
38renameType :: Map QBE.UserIdent Int -> QBE.UserIdent -> QBE.UserIdent
39renameType = rename (\(QBE.UserIdent s) n -> QBE.UserIdent $ s ++ ".QU" ++ show n)
40
41renameDef :: Env -> Definition -> Definition
42renameDef Env {envGlobals = varOcc, envTypes = tyOcc} (DefFunc funcDef) =
43  DefFunc $ USE.mapGlobals (USE.mapType funcDef (renameType tyOcc)) (renameGlobal varOcc)
44renameDef Env {envGlobals = varOcc} (DefData dataDef) =
45  DefData $ USE.mapGlobals dataDef (renameGlobal varOcc)
46renameDef Env {envTypes = tyOcc} (DefType typeDef) =
47  DefType $ USE.mapType typeDef (renameType tyOcc)
48renameDef _ def@(DefFile _) = def
49
50renameDefs :: Env -> Program -> Program
51renameDefs env = map (renameDef env)
52
53renameProg :: Program -> State Env Program
54renameProg prog = gets (`renameDefs` prog)
55
56------------------------------------------------------------------------
57
58count :: (Ord a) => Map a Int -> [a] -> Map a Int
59count knownVars defs =
60  let inc = maybe 0 (+ 1)
61   in foldl (flip $ upsert inc) knownVars defs
62  where
63    -- TODO: Remove this once we upgrade to containers >= 0.8.1.
64    -- See: <https://github.com/haskell/containers/issues/809>.
65    upsert f = Map.alter (Just . f)
66
67uniqProg :: Program -> State Env Program
68uniqProg prog = do
69  Env {envGlobals = knownVars, envTypes = knownTypes} <- get
70  let varOcc = count knownVars (localDefs prog)
71      tyOcc = count knownTypes (typeDefs prog)
72  modify (\e -> e {envGlobals = varOcc, envTypes = tyOcc})
73  renameProg prog
74
75-- | Merge multiple QBE sources files into one, renaming non-exported definitions.
76link :: [Program] -> Program
77link progs = evalState (foldM go [] progs) mkEnv
78  where
79    go acc x = (acc ++) <$> uniqProg x