Skip to main content

llvm_ir/
context.rs

1//! Context: interning tables for IR types and constants, plus all newtype index types.
2
3use crate::types::{FloatKind, FunctionType, StructType, TypeData};
4use crate::value::ConstantData;
5use std::collections::HashMap;
6
7// ---------------------------------------------------------------------------
8// Newtype index types — all u32, Copy
9// ---------------------------------------------------------------------------
10
11/// Public API for `TypeId`.
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
13pub struct TypeId(pub u32);
14
15/// Public API for `FunctionId`.
16#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
17pub struct FunctionId(pub u32);
18
19/// Public API for `BlockId`.
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
21pub struct BlockId(pub u32);
22
23/// Public API for `InstrId`.
24#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
25pub struct InstrId(pub u32);
26
27/// Public API for `ArgId`.
28#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
29pub struct ArgId(pub u32);
30
31/// Public API for `ConstId`.
32#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
33pub struct ConstId(pub u32);
34
35/// Public API for `GlobalId`.
36#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
37pub struct GlobalId(pub u32);
38
39// ---------------------------------------------------------------------------
40// Universal SSA value reference
41// ---------------------------------------------------------------------------
42
43/// A `Copy` reference to any SSA value.
44#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
45pub enum ValueRef {
46    /// `Instruction` variant.
47    Instruction(InstrId),
48    /// `Argument` variant.
49    Argument(ArgId),
50    /// `Constant` variant.
51    Constant(ConstId),
52    /// `Global` variant.
53    Global(GlobalId),
54}
55
56// ---------------------------------------------------------------------------
57// Constant deduplication key (scalars only)
58// ---------------------------------------------------------------------------
59
60/// Public API for `ConstantKey`.
61#[derive(Clone, Debug, PartialEq, Eq, Hash)]
62pub enum ConstantKey {
63    /// `Int` variant.
64    Int(TypeId, u64),
65    /// `Float` variant.
66    Float(TypeId, u64), // raw bits
67    /// `Null` variant.
68    Null(TypeId),
69    /// `Undef` variant.
70    Undef(TypeId),
71    /// `Poison` variant.
72    Poison(TypeId),
73    /// `ZeroInitializer` variant.
74    ZeroInitializer(TypeId),
75}
76
77// ---------------------------------------------------------------------------
78// Context
79// ---------------------------------------------------------------------------
80
81/// Public API for `Context`.
82pub struct Context {
83    // `types` field.
84    types: Vec<TypeData>,
85    /// Structural type interning (anonymous types).
86    type_map: HashMap<TypeData, TypeId>,
87    /// Named struct lookup by name.
88    named_struct_map: HashMap<String, TypeId>,
89    /// Constant pool.
90    pub constants: Vec<ConstantData>,
91    // `const_map` field.
92    const_map: HashMap<ConstantKey, ConstId>,
93
94    // Pre-interned singletons.
95    /// Public API for `void_ty`.
96    pub void_ty: TypeId,
97    /// Public API for `i1_ty`.
98    pub i1_ty: TypeId,
99    /// Public API for `i8_ty`.
100    pub i8_ty: TypeId,
101    /// Public API for `i16_ty`.
102    pub i16_ty: TypeId,
103    /// Public API for `i32_ty`.
104    pub i32_ty: TypeId,
105    /// Public API for `i64_ty`.
106    pub i64_ty: TypeId,
107    /// Public API for `f32_ty`.
108    pub f32_ty: TypeId,
109    /// Public API for `f64_ty`.
110    pub f64_ty: TypeId,
111    /// Public API for `ptr_ty`.
112    pub ptr_ty: TypeId,
113    /// Public API for `label_ty`.
114    pub label_ty: TypeId,
115}
116
117impl Context {
118    /// Public API for `new`.
119    pub fn new() -> Self {
120        let mut ctx = Context {
121            // `types` field.
122            types: Vec::new(),
123            // `type_map` field.
124            type_map: HashMap::new(),
125            // `named_struct_map` field.
126            named_struct_map: HashMap::new(),
127            // `constants` field.
128            constants: Vec::new(),
129            // `const_map` field.
130            const_map: HashMap::new(),
131            // `void_ty` field.
132            void_ty: TypeId(0),
133            // `i1_ty` field.
134            i1_ty: TypeId(0),
135            // `i8_ty` field.
136            i8_ty: TypeId(0),
137            // `i16_ty` field.
138            i16_ty: TypeId(0),
139            // `i32_ty` field.
140            i32_ty: TypeId(0),
141            // `i64_ty` field.
142            i64_ty: TypeId(0),
143            // `f32_ty` field.
144            f32_ty: TypeId(0),
145            // `f64_ty` field.
146            f64_ty: TypeId(0),
147            // `ptr_ty` field.
148            ptr_ty: TypeId(0),
149            // `label_ty` field.
150            label_ty: TypeId(0),
151        };
152        ctx.void_ty = ctx.intern_anon(TypeData::Void);
153        ctx.i1_ty = ctx.intern_anon(TypeData::Integer(1));
154        ctx.i8_ty = ctx.intern_anon(TypeData::Integer(8));
155        ctx.i16_ty = ctx.intern_anon(TypeData::Integer(16));
156        ctx.i32_ty = ctx.intern_anon(TypeData::Integer(32));
157        ctx.i64_ty = ctx.intern_anon(TypeData::Integer(64));
158        ctx.f32_ty = ctx.intern_anon(TypeData::Float(FloatKind::Single));
159        ctx.f64_ty = ctx.intern_anon(TypeData::Float(FloatKind::Double));
160        ctx.ptr_ty = ctx.intern_anon(TypeData::Pointer);
161        ctx.label_ty = ctx.intern_anon(TypeData::Label);
162        ctx
163    }
164
165    /// Intern a non-named-struct type by structural equality.
166    fn intern_anon(&mut self, td: TypeData) -> TypeId {
167        if let Some(&id) = self.type_map.get(&td) {
168            return id;
169        }
170        let id = TypeId(self.types.len() as u32);
171        self.type_map.insert(td.clone(), id);
172        self.types.push(td);
173        id
174    }
175
176    // -----------------------------------------------------------------------
177    // Type constructors
178    // -----------------------------------------------------------------------
179
180    /// Public API for `mk_int`.
181    pub fn mk_int(&mut self, bits: u32) -> TypeId {
182        self.intern_anon(TypeData::Integer(bits))
183    }
184
185    /// Public API for `mk_float`.
186    pub fn mk_float(&mut self, kind: FloatKind) -> TypeId {
187        self.intern_anon(TypeData::Float(kind))
188    }
189
190    /// Public API for `mk_ptr`.
191    pub fn mk_ptr(&mut self) -> TypeId {
192        self.ptr_ty
193    }
194
195    /// Public API for `mk_label`.
196    pub fn mk_label(&mut self) -> TypeId {
197        self.label_ty
198    }
199
200    /// Public API for `mk_metadata`.
201    pub fn mk_metadata(&mut self) -> TypeId {
202        self.intern_anon(TypeData::Metadata)
203    }
204
205    /// Public API for `mk_array`.
206    pub fn mk_array(&mut self, element: TypeId, len: u64) -> TypeId {
207        self.intern_anon(TypeData::Array { element, len })
208    }
209
210    /// Public API for `mk_vector`.
211    pub fn mk_vector(&mut self, element: TypeId, len: u32, scalable: bool) -> TypeId {
212        self.intern_anon(TypeData::Vector {
213            element,
214            len,
215            scalable,
216        })
217    }
218
219    /// Public API for `mk_fn_type`.
220    pub fn mk_fn_type(&mut self, ret: TypeId, params: Vec<TypeId>, variadic: bool) -> TypeId {
221        self.intern_anon(TypeData::Function(FunctionType {
222            ret,
223            params,
224            variadic,
225        }))
226    }
227
228    /// Public API for `mk_struct_anon`.
229    pub fn mk_struct_anon(&mut self, fields: Vec<TypeId>, packed: bool) -> TypeId {
230        self.intern_anon(TypeData::Struct(StructType {
231            name: None,
232            fields,
233            packed,
234        }))
235    }
236
237    /// Create or look up a named struct. If the name is new, an opaque (empty-body)
238    /// struct is allocated. Call `define_struct_body` to fill in fields later.
239    pub fn mk_struct_named(&mut self, name: String) -> TypeId {
240        if let Some(&id) = self.named_struct_map.get(&name) {
241            return id;
242        }
243        let id = TypeId(self.types.len() as u32);
244        self.types.push(TypeData::Struct(StructType {
245            name: Some(name.clone()),
246            fields: Vec::new(),
247            packed: false,
248        }));
249        self.named_struct_map.insert(name, id);
250        id
251    }
252
253    /// Fill in the body of a previously-created named struct.
254    pub fn define_struct_body(&mut self, id: TypeId, fields: Vec<TypeId>, packed: bool) {
255        if let TypeData::Struct(st) = &mut self.types[id.0 as usize] {
256            st.fields = fields;
257            st.packed = packed;
258        }
259    }
260
261    /// Look up a named struct by name.
262    pub fn get_named_struct(&self, name: &str) -> Option<TypeId> {
263        self.named_struct_map.get(name).copied()
264    }
265
266    // -----------------------------------------------------------------------
267    // Type accessors
268    // -----------------------------------------------------------------------
269
270    /// Public API for `get_type`.
271    pub fn get_type(&self, id: TypeId) -> &TypeData {
272        &self.types[id.0 as usize]
273    }
274
275    /// Public API for `get_type_mut`.
276    pub fn get_type_mut(&mut self, id: TypeId) -> &mut TypeData {
277        &mut self.types[id.0 as usize]
278    }
279
280    /// Total number of interned types.
281    pub fn num_types(&self) -> usize {
282        self.types.len()
283    }
284
285    /// Iterate over all (TypeId, TypeData) pairs.
286    pub fn types(&self) -> impl Iterator<Item = (TypeId, &TypeData)> {
287        self.types
288            .iter()
289            .enumerate()
290            .map(|(i, td)| (TypeId(i as u32), td))
291    }
292
293    // -----------------------------------------------------------------------
294    // Constant constructors
295    // -----------------------------------------------------------------------
296
297    /// Public API for `const_int`.
298    pub fn const_int(&mut self, ty: TypeId, val: u64) -> ConstId {
299        let key = ConstantKey::Int(ty, val);
300        if let Some(&id) = self.const_map.get(&key) {
301            return id;
302        }
303        let id = ConstId(self.constants.len() as u32);
304        self.constants.push(ConstantData::Int { ty, val });
305        self.const_map.insert(key, id);
306        id
307    }
308
309    /// Store a float constant as raw bits (f32 bits in low 32 for Single,
310    /// full u64 bits for Double / other).
311    pub fn const_float(&mut self, ty: TypeId, bits: u64) -> ConstId {
312        let key = ConstantKey::Float(ty, bits);
313        if let Some(&id) = self.const_map.get(&key) {
314            return id;
315        }
316        let id = ConstId(self.constants.len() as u32);
317        self.constants.push(ConstantData::Float { ty, bits });
318        self.const_map.insert(key, id);
319        id
320    }
321
322    /// Public API for `const_null`.
323    pub fn const_null(&mut self, ty: TypeId) -> ConstId {
324        let key = ConstantKey::Null(ty);
325        if let Some(&id) = self.const_map.get(&key) {
326            return id;
327        }
328        let id = ConstId(self.constants.len() as u32);
329        self.constants.push(ConstantData::Null(ty));
330        self.const_map.insert(key, id);
331        id
332    }
333
334    /// Public API for `const_undef`.
335    pub fn const_undef(&mut self, ty: TypeId) -> ConstId {
336        let key = ConstantKey::Undef(ty);
337        if let Some(&id) = self.const_map.get(&key) {
338            return id;
339        }
340        let id = ConstId(self.constants.len() as u32);
341        self.constants.push(ConstantData::Undef(ty));
342        self.const_map.insert(key, id);
343        id
344    }
345
346    /// Public API for `const_poison`.
347    pub fn const_poison(&mut self, ty: TypeId) -> ConstId {
348        let key = ConstantKey::Poison(ty);
349        if let Some(&id) = self.const_map.get(&key) {
350            return id;
351        }
352        let id = ConstId(self.constants.len() as u32);
353        self.constants.push(ConstantData::Poison(ty));
354        self.const_map.insert(key, id);
355        id
356    }
357
358    /// Public API for `const_zero`.
359    pub fn const_zero(&mut self, ty: TypeId) -> ConstId {
360        let key = ConstantKey::ZeroInitializer(ty);
361        if let Some(&id) = self.const_map.get(&key) {
362            return id;
363        }
364        let id = ConstId(self.constants.len() as u32);
365        self.constants.push(ConstantData::ZeroInitializer(ty));
366        self.const_map.insert(key, id);
367        id
368    }
369
370    /// Push a complex (non-scalar) constant without deduplication.
371    pub fn push_const(&mut self, c: ConstantData) -> ConstId {
372        let id = ConstId(self.constants.len() as u32);
373        self.constants.push(c);
374        id
375    }
376
377    // -----------------------------------------------------------------------
378    // Constant accessors
379    // -----------------------------------------------------------------------
380
381    /// Public API for `get_const`.
382    pub fn get_const(&self, id: ConstId) -> &ConstantData {
383        &self.constants[id.0 as usize]
384    }
385
386    /// Public API for `type_of_const`.
387    pub fn type_of_const(&self, id: ConstId) -> TypeId {
388        match &self.constants[id.0 as usize] {
389            ConstantData::Int { ty, .. } => *ty,
390            ConstantData::IntWide { ty, .. } => *ty,
391            ConstantData::Float { ty, .. } => *ty,
392            ConstantData::Null(ty) => *ty,
393            ConstantData::Undef(ty) => *ty,
394            ConstantData::Poison(ty) => *ty,
395            ConstantData::ZeroInitializer(ty) => *ty,
396            ConstantData::Array { ty, .. } => *ty,
397            ConstantData::Struct { ty, .. } => *ty,
398            ConstantData::Vector { ty, .. } => *ty,
399            ConstantData::GlobalRef { ty, .. } => *ty, // name field ignored here
400        }
401    }
402}
403
404impl Default for Context {
405    fn default() -> Self {
406        Self::new()
407    }
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    #[test]
415    fn singleton_types() {
416        let ctx = Context::new();
417        // Verify pre-interned singletons are distinct
418        assert_ne!(ctx.void_ty, ctx.i32_ty);
419        assert_ne!(ctx.i32_ty, ctx.i64_ty);
420        assert_ne!(ctx.f32_ty, ctx.f64_ty);
421        assert_ne!(ctx.ptr_ty, ctx.i32_ty);
422    }
423
424    #[test]
425    fn type_interning() {
426        let mut ctx = Context::new();
427        let a = ctx.mk_int(32);
428        let b = ctx.mk_int(32);
429        assert_eq!(a, b);
430        let c = ctx.mk_int(64);
431        assert_ne!(a, c);
432        assert_eq!(a, ctx.i32_ty);
433        assert_eq!(c, ctx.i64_ty);
434    }
435
436    #[test]
437    fn named_struct() {
438        let mut ctx = Context::new();
439        let id1 = ctx.mk_struct_named("Foo".to_string());
440        let id2 = ctx.mk_struct_named("Foo".to_string());
441        assert_eq!(id1, id2);
442        let id3 = ctx.mk_struct_named("Bar".to_string());
443        assert_ne!(id1, id3);
444    }
445
446    #[test]
447    fn const_int_dedup() {
448        let mut ctx = Context::new();
449        let c1 = ctx.const_int(ctx.i32_ty, 42);
450        let c2 = ctx.const_int(ctx.i32_ty, 42);
451        assert_eq!(c1, c2);
452        let c3 = ctx.const_int(ctx.i32_ty, 0);
453        assert_ne!(c1, c3);
454    }
455}