pliron/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) The pliron contributors
3
4//! [Context] and [Ptr] together provide memory management for `pliron`.
5
6use core::{
7    any::Any,
8    cell::{Cell, Ref, RefCell, RefMut},
9    fmt::Display,
10    hash::Hash,
11    marker::PhantomData,
12};
13
14use crate::{
15    arg_error_noloc,
16    basic_block::BasicBlock,
17    common_traits::Verify,
18    dialect::{Dialect, DialectName},
19    identifier::Identifier,
20    operation::Operation,
21    printable::{self, Printable},
22    region::Region,
23    result::Result,
24    std_deps::sync::LazyLock,
25    storage_uniquer::UniqueStore,
26    r#type::TypeObj,
27    uniqued_any::UniquedAny,
28    utils::table::{HMap, HSet, IMap},
29    verify_err_noloc,
30};
31use alloc::{boxed::Box, format, string::ToString, vec, vec::Vec};
32use slotmap::{SlotMap, new_key_type};
33
34new_key_type! {
35    /// The index type for the [SlotMap] used to store IR objects.
36    pub struct ArenaIndex;
37}
38
39new_key_type! {
40    /// The index type for the [SlotMap] used to store auxiliary data.
41    pub struct AuxDataIndex;
42}
43
44impl Display for ArenaIndex {
45    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
46        write!(f, "{:?}", self.0)
47    }
48}
49
50/// An arena allocation pool for IR objects.
51pub type Arena<T> = SlotMap<ArenaIndex, RefCell<T>>;
52
53/// A context stores all IR data of this compilation session.
54pub struct Context {
55    /// A unique number for each `Value` in the context.
56    /// Serves as a generation counter against accessing invalid `Value`s.
57    pub(crate) value_counter: Cell<u64>,
58    /// A unique number for each `Use` in the context.
59    /// Serves as a generation counter against accessing invalid `Use`s.
60    pub(crate) use_counter: Cell<u64>,
61    /// Allocation pool for [Operation]s.
62    pub(crate) operations: Arena<Operation>,
63    /// Allocation pool for [BasicBlock]s.
64    pub(crate) basic_blocks: Arena<BasicBlock>,
65    /// Allocation pool for [Region]s.
66    pub(crate) regions: Arena<Region>,
67    /// Registered [Dialect]s.
68    pub(crate) dialects: HMap<DialectName, Dialect>,
69    /// Storage for uniqued [TypeObj]s.
70    pub(crate) type_store: UniqueStore<TypeObj>,
71    /// Storage for other uniqued objects.
72    pub(crate) uniqued_any_store: UniqueStore<UniquedAny>,
73    /// Arbitrary data storage. Use [Self::aux_data_map] for dictionary access.
74    pub aux_data: SlotMap<AuxDataIndex, Box<dyn Any + Send>>,
75    /// A dictionary with keys mapping to an index in [Self::aux_data].
76    pub aux_data_map: HMap<Identifier, AuxDataIndex>,
77}
78
79impl Context {
80    pub fn new() -> Context {
81        Self::default()
82    }
83
84    /// Is the IR in this context empty?
85    /// An IR is considered empty if it has no operations, basic blocks, or regions.
86    /// This does not check for types, dialects, ops, or aux_data stored in the context.
87    pub fn is_ir_empty(&self) -> bool {
88        self.operations.is_empty() && self.basic_blocks.is_empty() && self.regions.is_empty()
89    }
90
91    /// Get a unique number for a new value.
92    pub(crate) fn get_new_value_uid(&self) -> u64 {
93        let uid = self.value_counter.get();
94        self.value_counter.set(uid + 1);
95        uid
96    }
97
98    /// Get a unique number for a new use.
99    pub(crate) fn get_new_use_uid(&self) -> u64 {
100        let uid = self.use_counter.get();
101        self.use_counter.set(uid + 1);
102        uid
103    }
104}
105
106impl Default for Context {
107    fn default() -> Self {
108        let mut ctx = Context {
109            value_counter: Cell::new(0),
110            use_counter: Cell::new(0),
111            operations: Arena::default(),
112            basic_blocks: Arena::default(),
113            regions: Arena::default(),
114            dialects: HMap::default(),
115            type_store: UniqueStore::default(),
116            uniqued_any_store: UniqueStore::default(),
117            aux_data: SlotMap::with_key(),
118            aux_data_map: HMap::default(),
119        };
120
121        // Verify that all dictionary keys are unique.
122        if let Err(err) = &*DICT_KEYS_VERIFIER {
123            panic!("{}", err.err);
124        }
125
126        // Run all context registrations
127        for registration in get_context_registrations() {
128            registration(&mut ctx);
129        }
130
131        ctx
132    }
133}
134
135// Static assertion to ensure that [Context] is [Send].
136const _: () = {
137    const fn assert_send<T: Send>() {}
138    assert_send::<Context>();
139};
140
141pub(crate) mod private {
142    use super::*;
143
144    /// An IR object owned by Context
145    pub trait ArenaObj
146    where
147        Self: Sized,
148    {
149        /// Get the arena that has allocated this object.
150        fn get_arena(ctx: &Context) -> &Arena<Self>;
151        /// Get the arena that has allocated this object.
152        fn get_arena_mut(ctx: &mut Context) -> &mut Arena<Self>;
153        /// Get a Ptr to self.
154        fn get_self_ptr(&self, ctx: &Context) -> Ptr<Self>;
155        /// If this object contains any ArenaObj itself, it must dealloc()
156        /// all of those sub-objects. This is called when self is deallocated.
157        fn dealloc_sub_objects(ptr: Ptr<Self>, ctx: &mut Context);
158
159        /// Allocates object on the arena, given a creator function.
160        fn alloc<T: FnOnce(Ptr<Self>) -> Self>(ctx: &mut Context, f: T) -> Ptr<Self> {
161            let creator = |idx: ArenaIndex| {
162                let t = f(Ptr::<Self> {
163                    idx,
164                    _dummy: PhantomData::<Self>,
165                });
166                RefCell::new(t)
167            };
168            Ptr::<Self> {
169                idx: Self::get_arena_mut(ctx).insert_with_key(creator),
170                _dummy: PhantomData,
171            }
172        }
173
174        /// Deallocates this object from the arena.
175        fn dealloc(ptr: Ptr<Self>, ctx: &mut Context) {
176            Self::dealloc_sub_objects(ptr, ctx);
177            Self::get_arena_mut(ctx).remove(ptr.idx);
178        }
179    }
180}
181
182use private::ArenaObj;
183
184/// Pointer to an IR Object owned by Context.
185pub struct Ptr<T: ArenaObj> {
186    pub(crate) idx: ArenaIndex,
187    pub(crate) _dummy: PhantomData<T>,
188}
189
190impl<T: ArenaObj> core::fmt::Debug for Ptr<T> {
191    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
192        write!(f, "Ptr<{}>[{}]", core::any::type_name::<T>(), self.idx)
193    }
194}
195
196#[derive(Debug, thiserror::Error)]
197#[error("Attempt to dereference a dangling Ptr")]
198pub struct DanglingPtrDerefError;
199
200impl<'a, T: ArenaObj> Ptr<T> {
201    /// Borrow the inner [RefCell] and return a [Ref] to the pointee.
202    /// The borrow is live as long as the returned [Ref] lives.
203    /// Panics on dangling [Ptr] or borrow interior mutability errors.
204    /// Run `cargo` with options `+nightly -Zbuild-std -Zbuild-std-features="debug_refcell"`
205    /// to enable printing the interior mutability violation locations.
206    #[track_caller]
207    pub fn deref(&self, ctx: &'a Context) -> Ref<'a, T> {
208        T::get_arena(ctx)
209            .get(self.idx)
210            .expect("Dangling Ptr deref")
211            .borrow()
212    }
213
214    /// Mutably borrow the inner [RefCell] and return a [RefMut] to the pointee.
215    /// The borrow is live as long as the returned [RefMut] lives.
216    /// Panics on dangling [Ptr] or borrow interior mutability errors.
217    /// Run `cargo` with options `+nightly -Zbuild-std -Zbuild-std-features="debug_refcell"`
218    /// to enable printing the interior mutability violation locations.
219    #[track_caller]
220    pub fn deref_mut(&self, ctx: &'a Context) -> RefMut<'a, T> {
221        T::get_arena(ctx)
222            .get(self.idx)
223            .expect("Dangling Ptr deref_mut")
224            .borrow_mut()
225    }
226
227    /// Try and borrow the inner [RefCell] and return a [Ref] to the pointee.
228    /// The borrow is live as long as the returned [Ref] lives.
229    /// If [Ptr] is dangling or already mutably borrowed, an [Error](crate::result::Error)
230    /// with [DanglingPtrDerefError] or [BorrowError](core::cell::BorrowError) is returned.
231    pub fn try_deref(&self, ctx: &'a Context) -> Result<Ref<'a, T>> {
232        T::get_arena(ctx)
233            .get(self.idx)
234            .ok_or_else(|| arg_error_noloc!(DanglingPtrDerefError))?
235            .try_borrow()
236            .map_err(|err| arg_error_noloc!(err))
237    }
238
239    /// Try and mutably borrow the inner [RefCell] and return a [RefMut] to the pointee.
240    /// The borrow is live as long as the returned [RefMut] lives.
241    /// If [Ptr] is dangling or already borrowed, an [Error](crate::result::Error)
242    /// with [DanglingPtrDerefError] or [BorrowMutError](core::cell::BorrowMutError) is returned.
243    pub fn try_deref_mut(&self, ctx: &'a Context) -> Result<RefMut<'a, T>> {
244        T::get_arena(ctx)
245            .get(self.idx)
246            .ok_or_else(|| arg_error_noloc!(DanglingPtrDerefError))?
247            .try_borrow_mut()
248            .map_err(|err| arg_error_noloc!(err))
249    }
250
251    /// Create a unique (to the arena) name based on the arena index.
252    pub(crate) fn make_name(&self, name_base: &str) -> Identifier {
253        let idx = format!("{}", self.idx);
254        (name_base.to_string() + &idx).try_into().unwrap()
255    }
256}
257
258impl<T: ArenaObj> Clone for Ptr<T> {
259    fn clone(&self) -> Ptr<T> {
260        *self
261    }
262}
263
264impl<T: ArenaObj> Copy for Ptr<T> {}
265
266impl<T: ArenaObj> PartialEq for Ptr<T> {
267    fn eq(&self, other: &Self) -> bool {
268        self.idx == other.idx
269    }
270}
271
272impl<T: ArenaObj> Eq for Ptr<T> {}
273
274impl<T: ArenaObj + 'static> Hash for Ptr<T> {
275    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
276        self.idx.hash(state);
277    }
278}
279
280impl<T: ArenaObj + Printable> Printable for Ptr<T> {
281    fn fmt(
282        &self,
283        ctx: &Context,
284        state: &printable::State,
285        f: &mut core::fmt::Formatter<'_>,
286    ) -> core::fmt::Result {
287        self.deref(ctx).fmt(ctx, state, f)
288    }
289}
290
291impl<T: ArenaObj + Verify> Verify for Ptr<T> {
292    fn verify(&self, ctx: &Context) -> Result<()> {
293        self.deref(ctx).verify(ctx)
294    }
295}
296
297#[doc(hidden)]
298/// Declaration of a static [Identifier] for use as a dictionary key.
299#[derive(Eq, PartialEq, Debug, Clone)]
300pub struct DictKeyId {
301    /// The [Identifier] itself.
302    pub id: Identifier,
303    /// The file where this key was declared.
304    pub file: &'static str,
305    /// The line where this key was declared.
306    pub line: u32,
307    /// The column where this key was declared.
308    pub column: u32,
309}
310
311/// These represent registrations that happen automatically at link time.
312/// Every dialect, op, type, and attribute that are linked into your code
313/// will register themselves using this type.
314pub type ContextRegistration = fn(&mut Context);
315
316#[doc(hidden)]
317/// `pliron` uses dictionaries indexed by static [Identifier]s in many places,
318/// such as [Context::aux_data_map], and the states of [Printable](crate::printable::State)
319/// and [Parsable](crate::parsable::State). To avoid collisions in these [Identifier]s,
320/// we use the [crate::dict_key!] macro to verify that all such keys declared using the macro
321/// are unique. The macro adds the keys to this static slice, which is then verified
322/// when a [Context] is created.
323#[cfg(not(target_family = "wasm"))]
324pub mod statics {
325    use super::*;
326
327    #[::pliron::linkme::distributed_slice]
328    #[linkme(crate = ::pliron::linkme)]
329    pub static DICT_KEY_IDS: [DictKeyId];
330
331    pub fn get_dict_key_ids() -> impl Iterator<Item = &'static DictKeyId> {
332        DICT_KEY_IDS.iter()
333    }
334
335    #[::pliron::linkme::distributed_slice]
336    #[linkme(crate = ::pliron::linkme)]
337    pub static CONTEXT_REGISTRATIONS: [ContextRegistration];
338
339    pub fn get_context_registrations() -> impl Iterator<Item = &'static ContextRegistration> {
340        CONTEXT_REGISTRATIONS.iter()
341    }
342}
343
344#[cfg(target_family = "wasm")]
345pub mod statics {
346    use super::*;
347    use crate::InventoryWrapper;
348
349    ::pliron::inventory::collect!(InventoryWrapper<DictKeyId>);
350
351    pub fn get_dict_key_ids() -> impl Iterator<Item = &'static DictKeyId> {
352        ::pliron::inventory::iter::<InventoryWrapper<DictKeyId>>().map(|llw| llw.0)
353    }
354
355    ::pliron::inventory::collect!(InventoryWrapper<ContextRegistration>);
356
357    pub fn get_context_registrations() -> impl Iterator<Item = &'static ContextRegistration> {
358        ::pliron::inventory::iter::<InventoryWrapper<ContextRegistration>>().map(|llw| llw.0)
359    }
360}
361
362pub use statics::*;
363
364#[doc(hidden)]
365pub static DICT_KEYS_VERIFIER: LazyLock<Result<()>> = LazyLock::new(verify_dict_keys);
366
367#[doc(hidden)]
368/// Collect `(owner, __all_verifiers)` entries into an ordered verifier map.
369///
370/// Each owner (op/type/attribute) can contribute verifiers through multiple interfaces.
371/// This helper preserves interface dependency order (as returned by each `__all_verifiers`
372/// function) while deduplicating verifier function pointers.
373pub(crate) fn collect_deduped_interface_verifiers<Id, AllVerifiers, Verifier>(
374    interface_verifiers: impl Iterator<Item = &'static &'static [(Id, AllVerifiers)]>,
375) -> HMap<Id, Vec<Verifier>>
376where
377    Id: Eq + Hash + Clone + 'static,
378    AllVerifiers: Fn() -> Vec<Verifier> + Clone + 'static,
379    Verifier: Eq + Hash + Clone,
380{
381    let mut grouped = IMap::default();
382    for entry in interface_verifiers.flat_map(|it| it.iter()) {
383        let (id, all_verifiers_for_interface) = entry;
384        grouped
385            .entry(id.clone())
386            .and_modify(|verifiers: &mut Vec<AllVerifiers>| {
387                verifiers.push(all_verifiers_for_interface.clone())
388            })
389            .or_insert(vec![all_verifiers_for_interface.clone()]);
390    }
391
392    // Remove duplicates (best effort as rustc may inline functions, resulting in different pointers).
393    // Relies on `__all_verifiers` returning the super-verifiers followed by self verifier
394    // to ensure that super-interfaces are verified first.
395    grouped
396        .into_iter()
397        .map(|(id, verifiers)| {
398            let mut dedupd_verifiers = Vec::new();
399            let mut seen = HSet::default();
400            for verifier_fn_list in verifiers {
401                for verifier in verifier_fn_list() {
402                    if seen.insert(verifier.clone()) {
403                        dedupd_verifiers.push(verifier);
404                    }
405                }
406            }
407            (id, dedupd_verifiers)
408        })
409        .collect()
410}
411
412#[doc(hidden)]
413/// Verify that all dictionary keys are unique. This is called when a [Context] is created.
414/// If any duplicate keys are found, a panic is raised with the file, line, and column
415/// information of the duplicate keys.
416pub fn verify_dict_keys() -> Result<()> {
417    let mut seen: HMap<Identifier, (&'static str, u32, u32)> = HMap::default();
418    for key in get_dict_key_ids() {
419        if let Some((file, line, column)) = seen.get(&key.id) {
420            return verify_err_noloc!(
421                "Duplicate dictionary key \"{}\" declared in {}:{}:{} and {}:{}:{}",
422                key.id,
423                file,
424                line,
425                column,
426                key.file,
427                key.line,
428                key.column
429            );
430        }
431        seen.insert(key.id.clone(), (key.file, key.line, key.column));
432    }
433    Ok(())
434}
435
436/// A macro to declare a static [Identifier] for use as a dictionary key.
437///
438/// Usage:
439/// ```
440/// # use pliron::dict_key;
441/// dict_key!(MY_KEY, "my_key");
442/// let mut ctx = pliron::context::Context::new();
443/// let aux_data_index = ctx.aux_data.insert(Box::new(42));
444/// ctx.aux_data_map.insert(MY_KEY.clone(), aux_data_index);
445/// assert_eq!(ctx.aux_data[aux_data_index].downcast_ref::<i32>(), Some(&42));
446/// assert_eq!(ctx.aux_data_map[&MY_KEY], aux_data_index);
447/// ```
448/// Here, `MY_KEY` is the name of the static variable, and `"my_key"` is the
449/// string value of the [Identifier]. The macro will create a static variable
450/// of type [Identifier] with the name `MY_KEY`.
451#[macro_export]
452macro_rules! dict_key {
453    (   $(#[$outer:meta])*
454        $decl:ident, $name:literal
455    ) => {
456        // Create a static variable linked to the DICT_KEY_IDS slice
457        // to ensure that all keys are unique.
458        // The static variable is created in a separate anonymous module.
459        const _: () = {
460            #[cfg_attr(not(target_family = "wasm"),
461                ::pliron::linkme::distributed_slice(::pliron::context::DICT_KEY_IDS), linkme(crate = ::pliron::linkme))]
462            pub static $decl: ::pliron::context::DictKeyId = ::pliron::context::DictKeyId {
463                id: $crate::ident!($name),
464                file: file!(),
465                line: line!(),
466                column: column!(),
467            };
468
469            #[cfg(target_family = "wasm")]
470            ::pliron::inventory::submit! {
471                ::pliron::InventoryWrapper(&$decl)
472            }
473        };
474        $(#[$outer])*
475        // Create a static variable with the provided name to access the identifier.
476        pub static $decl: ::pliron::identifier::Identifier = $crate::ident!($name);
477    };
478}
479
480/// A macro to register a [ContextRegistration]. The argument function
481/// will be called with a mutable reference to the [Context] when a [Context] is created.
482/// Use this outside of any function (e.g. in the module scope).
483///
484/// Usage:
485/// ```
486/// use pliron::context_registration;
487/// context_registration!(my_registration_fn);
488/// fn my_registration_fn(_: &mut pliron::context::Context) {}
489/// ```
490/// Here, `my_registration_fn` is a function matching [ContextRegistration].
491#[macro_export]
492macro_rules! context_registration {
493    (   $(#[$outer:meta])*
494        $registration:expr
495    ) => {
496        const _: () = {
497            $(#[$outer])*
498            #[cfg_attr(not(target_family = "wasm"),
499                ::pliron::linkme::distributed_slice(::pliron::context::CONTEXT_REGISTRATIONS), linkme(crate = ::pliron::linkme))]
500            static CONTEXT_REGISTRATION: ::pliron::context::ContextRegistration = $registration;
501
502            #[cfg(target_family = "wasm")]
503            ::pliron::inventory::submit! {
504                ::pliron::InventoryWrapper(&CONTEXT_REGISTRATION)
505            }
506        };
507    };
508}
pliron/context.rs

pliron/
context.rs
