lift_core/
context.rs

1use crate::attributes::Attributes;
2use crate::blocks::{BlockData, BlockKey};
3use crate::dialect::{register_builtin_dialects, DialectRegistry};
4use crate::functions::FunctionData;
5use crate::interning::{StringId, StringInterner, TypeInterner};
6use crate::location::Location;
7use crate::module::ModuleData;
8use crate::operations::{OpKey, OperationData};
9use crate::regions::{RegionData, RegionKey};
10use crate::types::{
11    CoreType, DataType, Dimension, MemoryLayout, QubitTypeInfo, TensorTypeInfo, TypeData, TypeId,
12};
13use crate::values::{DefSite, ValueData, ValueKey};
14use slotmap::SlotMap;
15
16#[derive(Debug)]
17pub struct Context {
18    pub values: SlotMap<ValueKey, ValueData>,
19    pub ops: SlotMap<OpKey, OperationData>,
20    pub blocks: SlotMap<BlockKey, BlockData>,
21    pub regions: SlotMap<RegionKey, RegionData>,
22    pub strings: StringInterner,
23    pub types: TypeInterner,
24    pub modules: Vec<ModuleData>,
25    pub dialects: DialectRegistry,
26}
27
28impl Context {
29    pub fn new() -> Self {
30        let mut ctx = Self {
31            values: SlotMap::with_key(),
32            ops: SlotMap::with_key(),
33            blocks: SlotMap::with_key(),
34            regions: SlotMap::with_key(),
35            strings: StringInterner::new(),
36            types: TypeInterner::new(),
37            modules: Vec::new(),
38            dialects: DialectRegistry::new(),
39        };
40        register_builtin_dialects(&mut ctx.dialects);
41        ctx
42    }
43
44    // ── String interning ──
45
46    pub fn intern_string(&mut self, s: &str) -> StringId {
47        self.strings.intern(s)
48    }
49
50    pub fn resolve_string(&self, id: StringId) -> &str {
51        self.strings.resolve(id)
52    }
53
54    // ── Type interning ──
55
56    pub fn intern_type(&mut self, ty: CoreType) -> TypeId {
57        self.types.intern(ty)
58    }
59
60    pub fn resolve_type(&self, id: TypeId) -> &CoreType {
61        self.types.resolve(id)
62    }
63
64    // ── Convenience type constructors ──
65
66    pub fn make_integer_type(&mut self, bits: u32, signed: bool) -> TypeId {
67        self.intern_type(CoreType::Integer { bits, signed })
68    }
69
70    pub fn make_float_type(&mut self, bits: u32) -> TypeId {
71        self.intern_type(CoreType::Float { bits })
72    }
73
74    pub fn make_bool_type(&mut self) -> TypeId {
75        self.intern_type(CoreType::Boolean)
76    }
77
78    pub fn make_void_type(&mut self) -> TypeId {
79        self.intern_type(CoreType::Void)
80    }
81
82    pub fn make_index_type(&mut self) -> TypeId {
83        self.intern_type(CoreType::Index)
84    }
85
86    pub fn make_tensor_type(
87        &mut self,
88        shape: Vec<Dimension>,
89        dtype: DataType,
90        layout: MemoryLayout,
91    ) -> TypeId {
92        let dialect = self.intern_string("tensor");
93        let name = self.intern_string("tensor");
94        self.intern_type(CoreType::Opaque {
95            dialect,
96            name,
97            data: TypeData::Tensor(TensorTypeInfo {
98                shape,
99                dtype,
100                layout,
101            }),
102        })
103    }
104
105    pub fn make_qubit_type(&mut self) -> TypeId {
106        let dialect = self.intern_string("quantum");
107        let name = self.intern_string("qubit");
108        self.intern_type(CoreType::Opaque {
109            dialect,
110            name,
111            data: TypeData::Qubit(QubitTypeInfo::Logical),
112        })
113    }
114
115    pub fn make_physical_qubit_type(
116        &mut self,
117        id: usize,
118        t1: f64,
119        t2: f64,
120        freq: f64,
121        fidelity: f64,
122    ) -> TypeId {
123        let dialect = self.intern_string("quantum");
124        let name = self.intern_string("physical_qubit");
125        use crate::types::OrderedFloat;
126        self.intern_type(CoreType::Opaque {
127            dialect,
128            name,
129            data: TypeData::Qubit(QubitTypeInfo::Physical {
130                id,
131                t1_us: OrderedFloat::from_f64(t1),
132                t2_us: OrderedFloat::from_f64(t2),
133                freq_ghz: OrderedFloat::from_f64(freq),
134                fidelity: OrderedFloat::from_f64(fidelity),
135            }),
136        })
137    }
138
139    pub fn make_bit_type(&mut self) -> TypeId {
140        let dialect = self.intern_string("quantum");
141        let name = self.intern_string("bit");
142        self.intern_type(CoreType::Opaque {
143            dialect,
144            name,
145            data: TypeData::ClassicalBit,
146        })
147    }
148
149    pub fn make_hamiltonian_type(&mut self, num_qubits: usize) -> TypeId {
150        let dialect = self.intern_string("quantum");
151        let name = self.intern_string("hamiltonian");
152        self.intern_type(CoreType::Opaque {
153            dialect,
154            name,
155            data: TypeData::Hamiltonian { num_qubits },
156        })
157    }
158
159    pub fn make_function_type(&mut self, params: Vec<TypeId>, returns: Vec<TypeId>) -> TypeId {
160        self.intern_type(CoreType::Function { params, returns })
161    }
162
163    pub fn make_tuple_type(&mut self, elems: Vec<TypeId>) -> TypeId {
164        self.intern_type(CoreType::Tuple(elems))
165    }
166
167    // ── Value creation ──
168
169    pub fn create_value(&mut self, ty: TypeId, name: Option<StringId>, def: DefSite) -> ValueKey {
170        self.values.insert(ValueData { ty, name, def })
171    }
172
173    pub fn get_value(&self, key: ValueKey) -> Option<&ValueData> {
174        self.values.get(key)
175    }
176
177    pub fn value_type(&self, key: ValueKey) -> Option<TypeId> {
178        self.values.get(key).map(|v| v.ty)
179    }
180
181    // ── Operation creation ──
182
183    pub fn create_op(
184        &mut self,
185        name: &str,
186        dialect: &str,
187        inputs: Vec<ValueKey>,
188        result_types: Vec<TypeId>,
189        attrs: Attributes,
190        location: Location,
191    ) -> (OpKey, Vec<ValueKey>) {
192        let name_id = self.intern_string(name);
193        let dialect_id = self.intern_string(dialect);
194
195        let op_key = self.ops.insert(OperationData {
196            name: name_id,
197            dialect: dialect_id,
198            inputs,
199            results: Vec::new(),
200            attrs,
201            regions: Vec::new(),
202            location,
203            parent_block: None,
204        });
205
206        let mut result_keys = Vec::with_capacity(result_types.len());
207        for (i, ty) in result_types.iter().enumerate() {
208            let val_key = self.create_value(
209                *ty,
210                None,
211                DefSite::OpResult {
212                    op: op_key,
213                    result_index: i as u32,
214                },
215            );
216            result_keys.push(val_key);
217        }
218
219        self.ops[op_key].results = result_keys.clone();
220        (op_key, result_keys)
221    }
222
223    pub fn get_op(&self, key: OpKey) -> Option<&OperationData> {
224        self.ops.get(key)
225    }
226
227    pub fn get_op_mut(&mut self, key: OpKey) -> Option<&mut OperationData> {
228        self.ops.get_mut(key)
229    }
230
231    pub fn op_name(&self, key: OpKey) -> &str {
232        let op = &self.ops[key];
233        self.strings.resolve(op.name)
234    }
235
236    pub fn op_dialect(&self, key: OpKey) -> &str {
237        let op = &self.ops[key];
238        self.strings.resolve(op.dialect)
239    }
240
241    // ── Block creation ──
242
243    pub fn create_block(&mut self) -> BlockKey {
244        self.blocks.insert(BlockData::new())
245    }
246
247    pub fn create_block_arg(&mut self, block: BlockKey, ty: TypeId) -> ValueKey {
248        let arg_index = self.blocks[block].args.len() as u32;
249        let val_key = self.create_value(ty, None, DefSite::BlockArg { block, arg_index });
250        self.blocks[block].args.push(val_key);
251        val_key
252    }
253
254    pub fn add_op_to_block(&mut self, block: BlockKey, op: OpKey) {
255        self.blocks[block].ops.push(op);
256        self.ops[op].parent_block = Some(block);
257    }
258
259    /// Inserts an operation immediately before `before` in its block.
260    /// If `before` has no parent block, falls back to appending at the end.
261    pub fn insert_op_before(&mut self, before: OpKey, op: OpKey) {
262        let parent = match self.ops.get(before) {
263            Some(b) => b.parent_block,
264            None => None,
265        };
266        let Some(block) = parent else {
267            // No parent block: append to every block is wrong; just detach.
268            return;
269        };
270        if let Some(ops) = self.blocks.get_mut(block) {
271            if let Some(pos) = ops.ops.iter().position(|&k| k == before) {
272                ops.ops.insert(pos, op);
273            } else {
274                ops.ops.push(op);
275            }
276        }
277        self.ops[op].parent_block = Some(block);
278    }
279
280    pub fn get_block(&self, key: BlockKey) -> Option<&BlockData> {
281        self.blocks.get(key)
282    }
283
284    // ── Region creation ──
285
286    pub fn create_region(&mut self) -> RegionKey {
287        self.regions.insert(RegionData::new())
288    }
289
290    pub fn add_block_to_region(&mut self, region: RegionKey, block: BlockKey) {
291        if self.regions[region].entry_block.is_none() {
292            self.regions[region].entry_block = Some(block);
293        }
294        self.regions[region].blocks.push(block);
295        self.blocks[block].parent_region = Some(region);
296    }
297
298    pub fn attach_region_to_op(&mut self, op: OpKey, region: RegionKey) {
299        self.ops[op].regions.push(region);
300        self.regions[region].parent_op = Some(op);
301    }
302
303    pub fn get_region(&self, key: RegionKey) -> Option<&RegionData> {
304        self.regions.get(key)
305    }
306
307    // ── Module creation ──
308
309    pub fn create_module(&mut self, name: &str) -> usize {
310        let name_id = self.intern_string(name);
311        let module = ModuleData::new(name_id);
312        self.modules.push(module);
313        self.modules.len() - 1
314    }
315
316    pub fn get_module(&self, index: usize) -> Option<&ModuleData> {
317        self.modules.get(index)
318    }
319
320    pub fn get_module_mut(&mut self, index: usize) -> Option<&mut ModuleData> {
321        self.modules.get_mut(index)
322    }
323
324    pub fn add_function_to_module(&mut self, module_index: usize, func: FunctionData) {
325        if let Some(module) = self.modules.get_mut(module_index) {
326            module.add_function(func);
327        }
328    }
329
330    // ── Snapshot/restore for pass rollback ──
331
332    pub fn snapshot(&self) -> ContextSnapshot {
333        ContextSnapshot {
334            num_values: self.values.len(),
335            num_ops: self.ops.len(),
336            num_blocks: self.blocks.len(),
337            num_regions: self.regions.len(),
338        }
339    }
340
341    // ── Type queries ──
342
343    pub fn is_qubit_type(&self, ty: TypeId) -> bool {
344        matches!(
345            self.resolve_type(ty),
346            CoreType::Opaque {
347                data: TypeData::Qubit(_),
348                ..
349            }
350        )
351    }
352
353    pub fn is_tensor_type(&self, ty: TypeId) -> bool {
354        matches!(
355            self.resolve_type(ty),
356            CoreType::Opaque {
357                data: TypeData::Tensor(_),
358                ..
359            }
360        )
361    }
362
363    pub fn is_bit_type(&self, ty: TypeId) -> bool {
364        matches!(
365            self.resolve_type(ty),
366            CoreType::Opaque {
367                data: TypeData::ClassicalBit,
368                ..
369            }
370        )
371    }
372
373    pub fn get_tensor_info(&self, ty: TypeId) -> Option<&TensorTypeInfo> {
374        match self.resolve_type(ty) {
375            CoreType::Opaque {
376                data: TypeData::Tensor(info),
377                ..
378            } => Some(info),
379            _ => None,
380        }
381    }
382}
383
384impl Default for Context {
385    fn default() -> Self {
386        Self::new()
387    }
388}
389
390#[derive(Debug, Clone)]
391pub struct ContextSnapshot {
392    pub num_values: usize,
393    pub num_ops: usize,
394    pub num_blocks: usize,
395    pub num_regions: usize,
396}
lift_core/context.rs

lift_core/
context.rs
