llvm_ir/
function.rs

1//! Function definition: signature, arguments, basic blocks, and flat instruction pool.
2
3use crate::basic_block::BasicBlock;
4use crate::context::{ArgId, BlockId, InstrId, TypeId, ValueRef};
5use crate::instruction::Instruction;
6use crate::value::{Argument, Linkage};
7use std::collections::HashMap;
8
9/// A function definition or declaration.
10pub struct Function {
11    /// Public API for `name`.
12    pub name: String,
13    /// Function type (FunctionType TypeId).
14    pub ty: TypeId,
15    /// Formal arguments.
16    pub args: Vec<Argument>,
17    /// Basic blocks in program order.
18    pub blocks: Vec<BasicBlock>,
19    /// Flat instruction pool; `InstrId(i)` indexes `instructions[i]`.
20    pub instructions: Vec<Instruction>,
21    /// Optional `!dbg !N` attachment for each instruction id.
22    pub instr_dbg_locs: HashMap<InstrId, u32>,
23    /// Arbitrary metadata attachments per instruction, e.g. `!dbg !12`, `!tbaa !7`.
24    pub instr_metadata: HashMap<InstrId, Vec<(String, String)>>,
25    /// Maps result name → InstrId.
26    pub value_names: HashMap<String, InstrId>,
27    /// Maps argument name → ArgId.
28    pub arg_names: HashMap<String, ArgId>,
29    /// True if this is a declaration (no body).
30    pub is_declaration: bool,
31    /// Public API for `linkage`.
32    pub linkage: Linkage,
33    /// Counter for generating unique names.
34    next_name_id: u32,
35}
36
37impl Function {
38    /// Public API for `new`.
39    pub fn new(name: impl Into<String>, ty: TypeId, args: Vec<Argument>, linkage: Linkage) -> Self {
40        let mut f = Function {
41            name: name.into(),
42            ty,
43            args: Vec::new(),
44            blocks: Vec::new(),
45            instructions: Vec::new(),
46            instr_dbg_locs: HashMap::new(),
47            instr_metadata: HashMap::new(),
48            value_names: HashMap::new(),
49            arg_names: HashMap::new(),
50            is_declaration: false,
51            linkage,
52            next_name_id: 0,
53        };
54        for arg in args {
55            let idx = ArgId(f.args.len() as u32);
56            if !arg.name.is_empty() {
57                f.arg_names.insert(arg.name.clone(), idx);
58            }
59            f.args.push(arg);
60        }
61        f
62    }
63
64    /// Public API for `new_declaration`.
65    pub fn new_declaration(
66        name: impl Into<String>,
67        ty: TypeId,
68        args: Vec<Argument>,
69        linkage: Linkage,
70    ) -> Self {
71        let mut f = Self::new(name, ty, args, linkage);
72        f.is_declaration = true;
73        f
74    }
75
76    // -----------------------------------------------------------------------
77    // Block management
78    // -----------------------------------------------------------------------
79
80    /// Add a new basic block and return its `BlockId`.
81    pub fn add_block(&mut self, bb: BasicBlock) -> BlockId {
82        let id = BlockId(self.blocks.len() as u32);
83        self.blocks.push(bb);
84        id
85    }
86
87    /// Public API for `block`.
88    pub fn block(&self, id: BlockId) -> &BasicBlock {
89        &self.blocks[id.0 as usize]
90    }
91
92    /// Public API for `block_mut`.
93    pub fn block_mut(&mut self, id: BlockId) -> &mut BasicBlock {
94        &mut self.blocks[id.0 as usize]
95    }
96
97    /// Public API for `num_blocks`.
98    pub fn num_blocks(&self) -> usize {
99        self.blocks.len()
100    }
101
102    // -----------------------------------------------------------------------
103    // Instruction pool
104    // -----------------------------------------------------------------------
105
106    /// Allocate an instruction in the flat pool, register its name if any,
107    /// and return the `InstrId`.
108    pub fn alloc_instr(&mut self, mut instr: Instruction) -> InstrId {
109        // Auto-number unnamed value-producing instructions.
110        if instr.name.as_deref() == Some("") {
111            let name = self.fresh_name();
112            instr.name = Some(name);
113        }
114        let id = InstrId(self.instructions.len() as u32);
115        if let Some(ref n) = instr.name {
116            if !n.is_empty() {
117                self.value_names.insert(n.clone(), id);
118            }
119        }
120        self.instructions.push(instr);
121        id
122    }
123
124    /// Public API for `instr`.
125    pub fn instr(&self, id: InstrId) -> &Instruction {
126        &self.instructions[id.0 as usize]
127    }
128
129    /// Public API for `instr_mut`.
130    pub fn instr_mut(&mut self, id: InstrId) -> &mut Instruction {
131        &mut self.instructions[id.0 as usize]
132    }
133
134    /// Public API for `num_instrs`.
135    pub fn num_instrs(&self) -> usize {
136        self.instructions.len()
137    }
138
139    /// Public API for `set_instr_dbg_loc`.
140    pub fn set_instr_dbg_loc(&mut self, id: InstrId, loc_id: u32) {
141        self.instr_dbg_locs.insert(id, loc_id);
142    }
143
144    /// Public API for `instr_dbg_loc`.
145    pub fn instr_dbg_loc(&self, id: InstrId) -> Option<u32> {
146        self.instr_dbg_locs.get(&id).copied()
147    }
148
149    /// Public API for `add_instr_metadata`.
150    pub fn add_instr_metadata(&mut self, id: InstrId, key: impl Into<String>, value: impl Into<String>) {
151        self.instr_metadata
152            .entry(id)
153            .or_default()
154            .push((key.into(), value.into()));
155    }
156
157    /// Public API for `instr_metadata`.
158    pub fn instr_metadata(&self, id: InstrId) -> Option<&[(String, String)]> {
159        self.instr_metadata.get(&id).map(Vec::as_slice)
160    }
161
162    // -----------------------------------------------------------------------
163    // Arguments
164    // -----------------------------------------------------------------------
165
166    /// Public API for `arg`.
167    pub fn arg(&self, id: ArgId) -> &Argument {
168        &self.args[id.0 as usize]
169    }
170
171    /// Public API for `num_args`.
172    pub fn num_args(&self) -> usize {
173        self.args.len()
174    }
175
176    // -----------------------------------------------------------------------
177    // Name lookups
178    // -----------------------------------------------------------------------
179
180    /// Public API for `lookup_value`.
181    pub fn lookup_value(&self, name: &str) -> Option<ValueRef> {
182        if let Some(&iid) = self.value_names.get(name) {
183            return Some(ValueRef::Instruction(iid));
184        }
185        if let Some(&aid) = self.arg_names.get(name) {
186            return Some(ValueRef::Argument(aid));
187        }
188        None
189    }
190
191    /// Public API for `lookup_block`.
192    pub fn lookup_block(&self, name: &str) -> Option<BlockId> {
193        self.blocks
194            .iter()
195            .enumerate()
196            .find(|(_, bb)| bb.name == name)
197            .map(|(i, _)| BlockId(i as u32))
198    }
199
200    // -----------------------------------------------------------------------
201    // Type of SSA values
202    // -----------------------------------------------------------------------
203
204    /// Public API for `type_of_value`.
205    pub fn type_of_value(&self, vref: ValueRef) -> Option<TypeId> {
206        match vref {
207            ValueRef::Instruction(id) => Some(self.instructions[id.0 as usize].ty),
208            ValueRef::Argument(id) => Some(self.args[id.0 as usize].ty),
209            ValueRef::Constant(_) | ValueRef::Global(_) => None, // caller must consult Context/Module
210        }
211    }
212
213    // -----------------------------------------------------------------------
214    // Name generation
215    // -----------------------------------------------------------------------
216
217    /// Produce a unique name like `"1"`, `"2"`, … for unnamed SSA values.
218    pub fn fresh_name(&mut self) -> String {
219        let n = self.next_name_id;
220        self.next_name_id += 1;
221        // Prefix with "v" so the printed name is "%v0", "%v1", etc.  Plain
222        // integer names ("%0", "%1") collide with LLVM's implicit slot
223        // numbering and are rejected by llvm-as / clang with:
224        //   "instruction expected to be numbered '%N' or greater"
225        format!("v{n}")
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use crate::context::Context;
233
234    #[test]
235    fn function_fresh_names() {
236        let mut ctx = Context::new();
237        let fn_ty = ctx.mk_fn_type(ctx.void_ty, vec![], false);
238        let mut f = Function::new("test", fn_ty, vec![], Linkage::External);
239        assert_eq!(f.fresh_name(), "v0");
240        assert_eq!(f.fresh_name(), "v1");
241        assert_eq!(f.fresh_name(), "v2");
242    }
243
244    #[test]
245    fn function_add_block() {
246        let mut ctx = Context::new();
247        let fn_ty = ctx.mk_fn_type(ctx.void_ty, vec![], false);
248        let mut f = Function::new("test", fn_ty, vec![], Linkage::External);
249        let bb = BasicBlock::new("entry");
250        let bid = f.add_block(bb);
251        assert_eq!(bid, BlockId(0));
252        assert_eq!(f.block(bid).name, "entry");
253    }
254}
llvm_ir/function.rs

llvm_ir/
function.rs
