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>>,
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
135pub(crate) mod private {
136    use super::*;
137
138    /// An IR object owned by Context
139    pub trait ArenaObj
140    where
141        Self: Sized,
142    {
143        /// Get the arena that has allocated this object.
144        fn get_arena(ctx: &Context) -> &Arena<Self>;
145        /// Get the arena that has allocated this object.
146        fn get_arena_mut(ctx: &mut Context) -> &mut Arena<Self>;
147        /// Get a Ptr to self.
148        fn get_self_ptr(&self, ctx: &Context) -> Ptr<Self>;
149        /// If this object contains any ArenaObj itself, it must dealloc()
150        /// all of those sub-objects. This is called when self is deallocated.
151        fn dealloc_sub_objects(ptr: Ptr<Self>, ctx: &mut Context);
152
153        /// Allocates object on the arena, given a creator function.
154        fn alloc<T: FnOnce(Ptr<Self>) -> Self>(ctx: &mut Context, f: T) -> Ptr<Self> {
155            let creator = |idx: ArenaIndex| {
156                let t = f(Ptr::<Self> {
157                    idx,
158                    _dummy: PhantomData::<Self>,
159                });
160                RefCell::new(t)
161            };
162            Ptr::<Self> {
163                idx: Self::get_arena_mut(ctx).insert_with_key(creator),
164                _dummy: PhantomData,
165            }
166        }
167
168        /// Deallocates this object from the arena.
169        fn dealloc(ptr: Ptr<Self>, ctx: &mut Context) {
170            Self::dealloc_sub_objects(ptr, ctx);
171            Self::get_arena_mut(ctx).remove(ptr.idx);
172        }
173    }
174}
175
176use private::ArenaObj;
177
178/// Pointer to an IR Object owned by Context.
179pub struct Ptr<T: ArenaObj> {
180    pub(crate) idx: ArenaIndex,
181    pub(crate) _dummy: PhantomData<T>,
182}
183
184impl<T: ArenaObj> core::fmt::Debug for Ptr<T> {
185    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
186        write!(f, "Ptr<{}>[{}]", core::any::type_name::<T>(), self.idx)
187    }
188}
189
190#[derive(Debug, thiserror::Error)]
191#[error("Attempt to dereference a dangling Ptr")]
192pub struct DanglingPtrDerefError;
193
194impl<'a, T: ArenaObj> Ptr<T> {
195    /// Borrow the inner [RefCell] and return a [Ref] to the pointee.
196    /// The borrow is live as long as the returned [Ref] lives.
197    /// Panics on dangling [Ptr] or borrow interior mutability errors.
198    /// Run `cargo` with options `+nightly -Zbuild-std -Zbuild-std-features="debug_refcell"`
199    /// to enable printing the interior mutability violation locations.
200    #[track_caller]
201    pub fn deref(&self, ctx: &'a Context) -> Ref<'a, T> {
202        T::get_arena(ctx)
203            .get(self.idx)
204            .expect("Dangling Ptr deref")
205            .borrow()
206    }
207
208    /// Mutably borrow the inner [RefCell] and return a [RefMut] to the pointee.
209    /// The borrow is live as long as the returned [RefMut] lives.
210    /// Panics on dangling [Ptr] or borrow interior mutability errors.
211    /// Run `cargo` with options `+nightly -Zbuild-std -Zbuild-std-features="debug_refcell"`
212    /// to enable printing the interior mutability violation locations.
213    #[track_caller]
214    pub fn deref_mut(&self, ctx: &'a Context) -> RefMut<'a, T> {
215        T::get_arena(ctx)
216            .get(self.idx)
217            .expect("Dangling Ptr deref_mut")
218            .borrow_mut()
219    }
220
221    /// Try and borrow the inner [RefCell] and return a [Ref] to the pointee.
222    /// The borrow is live as long as the returned [Ref] lives.
223    /// If [Ptr] is dangling or already mutably borrowed, an [Error](crate::result::Error)
224    /// with [DanglingPtrDerefError] or [BorrowError](core::cell::BorrowError) is returned.
225    pub fn try_deref(&self, ctx: &'a Context) -> Result<Ref<'a, T>> {
226        T::get_arena(ctx)
227            .get(self.idx)
228            .ok_or_else(|| arg_error_noloc!(DanglingPtrDerefError))?
229            .try_borrow()
230            .map_err(|err| arg_error_noloc!(err))
231    }
232
233    /// Try and mutably borrow the inner [RefCell] and return a [RefMut] to the pointee.
234    /// The borrow is live as long as the returned [RefMut] lives.
235    /// If [Ptr] is dangling or already borrowed, an [Error](crate::result::Error)
236    /// with [DanglingPtrDerefError] or [BorrowMutError](core::cell::BorrowMutError) is returned.
237    pub fn try_deref_mut(&self, ctx: &'a Context) -> Result<RefMut<'a, T>> {
238        T::get_arena(ctx)
239            .get(self.idx)
240            .ok_or_else(|| arg_error_noloc!(DanglingPtrDerefError))?
241            .try_borrow_mut()
242            .map_err(|err| arg_error_noloc!(err))
243    }
244
245    /// Create a unique (to the arena) name based on the arena index.
246    pub(crate) fn make_name(&self, name_base: &str) -> Identifier {
247        let idx = format!("{}", self.idx);
248        (name_base.to_string() + &idx).try_into().unwrap()
249    }
250}
251
252impl<T: ArenaObj> Clone for Ptr<T> {
253    fn clone(&self) -> Ptr<T> {
254        *self
255    }
256}
257
258impl<T: ArenaObj> Copy for Ptr<T> {}
259
260impl<T: ArenaObj> PartialEq for Ptr<T> {
261    fn eq(&self, other: &Self) -> bool {
262        self.idx == other.idx
263    }
264}
265
266impl<T: ArenaObj> Eq for Ptr<T> {}
267
268impl<T: ArenaObj + 'static> Hash for Ptr<T> {
269    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
270        self.idx.hash(state);
271    }
272}
273
274impl<T: ArenaObj + Printable> Printable for Ptr<T> {
275    fn fmt(
276        &self,
277        ctx: &Context,
278        state: &printable::State,
279        f: &mut core::fmt::Formatter<'_>,
280    ) -> core::fmt::Result {
281        self.deref(ctx).fmt(ctx, state, f)
282    }
283}
284
285impl<T: ArenaObj + Verify> Verify for Ptr<T> {
286    fn verify(&self, ctx: &Context) -> Result<()> {
287        self.deref(ctx).verify(ctx)
288    }
289}
290
291#[doc(hidden)]
292/// Declaration of a static [Identifier] for use as a dictionary key.
293#[derive(Eq, PartialEq, Debug, Clone)]
294pub struct DictKeyId {
295    /// The [Identifier] itself.
296    pub id: Identifier,
297    /// The file where this key was declared.
298    pub file: &'static str,
299    /// The line where this key was declared.
300    pub line: u32,
301    /// The column where this key was declared.
302    pub column: u32,
303}
304
305/// These represent registrations that happen automatically at link time.
306/// Every dialect, op, type, and attribute that are linked into your code
307/// will register themselves using this type.
308pub type ContextRegistration = fn(&mut Context);
309
310#[doc(hidden)]
311/// `pliron` uses dictionaries indexed by static [Identifier]s in many places,
312/// such as [Context::aux_data_map], and the states of [Printable](crate::printable::State)
313/// and [Parsable](crate::parsable::State). To avoid collisions in these [Identifier]s,
314/// we use the [crate::dict_key!] macro to verify that all such keys declared using the macro
315/// are unique. The macro adds the keys to this static slice, which is then verified
316/// when a [Context] is created.
317#[cfg(not(target_family = "wasm"))]
318pub mod statics {
319    use super::*;
320
321    #[::pliron::linkme::distributed_slice]
322    #[linkme(crate = ::pliron::linkme)]
323    pub static DICT_KEY_IDS: [LazyLock<DictKeyId>];
324
325    pub fn get_dict_key_ids() -> impl Iterator<Item = &'static LazyLock<DictKeyId>> {
326        DICT_KEY_IDS.iter()
327    }
328
329    #[::pliron::linkme::distributed_slice]
330    #[linkme(crate = ::pliron::linkme)]
331    pub static CONTEXT_REGISTRATIONS: [ContextRegistration];
332
333    pub fn get_context_registrations() -> impl Iterator<Item = &'static ContextRegistration> {
334        CONTEXT_REGISTRATIONS.iter()
335    }
336}
337
338#[cfg(target_family = "wasm")]
339pub mod statics {
340    use super::*;
341    use crate::InventoryWrapper;
342
343    ::pliron::inventory::collect!(InventoryWrapper<LazyLock<DictKeyId>>);
344
345    pub fn get_dict_key_ids() -> impl Iterator<Item = &'static LazyLock<DictKeyId>> {
346        ::pliron::inventory::iter::<InventoryWrapper<LazyLock<DictKeyId>>>().map(|llw| llw.0)
347    }
348
349    ::pliron::inventory::collect!(InventoryWrapper<ContextRegistration>);
350
351    pub fn get_context_registrations() -> impl Iterator<Item = &'static ContextRegistration> {
352        ::pliron::inventory::iter::<InventoryWrapper<ContextRegistration>>().map(|llw| llw.0)
353    }
354}
355
356pub use statics::*;
357
358#[doc(hidden)]
359pub static DICT_KEYS_VERIFIER: LazyLock<Result<()>> = LazyLock::new(verify_dict_keys);
360
361#[doc(hidden)]
362/// Collect `(owner, __all_verifiers)` entries into an ordered verifier map.
363///
364/// Each owner (op/type/attribute) can contribute verifiers through multiple interfaces.
365/// This helper preserves interface dependency order (as returned by each `__all_verifiers`
366/// function) while deduplicating verifier function pointers.
367pub(crate) fn collect_deduped_interface_verifiers<Id, AllVerifiers, Verifier>(
368    interface_verifiers: impl Iterator<Item = &'static (Id, AllVerifiers)>,
369) -> HMap<Id, Vec<Verifier>>
370where
371    Id: Eq + Hash + Clone + 'static,
372    AllVerifiers: Fn() -> Vec<Verifier> + Clone + 'static,
373    Verifier: Eq + Hash + Clone,
374{
375    let mut grouped = IMap::default();
376    for entry in interface_verifiers {
377        let (id, all_verifiers_for_interface) = entry;
378        grouped
379            .entry(id.clone())
380            .and_modify(|verifiers: &mut Vec<AllVerifiers>| {
381                verifiers.push(all_verifiers_for_interface.clone())
382            })
383            .or_insert(vec![all_verifiers_for_interface.clone()]);
384    }
385
386    // Remove duplicates (best effort as rustc may inline functions, resulting in different pointers).
387    // Relies on `__all_verifiers` returning the super-verifiers followed by self verifier
388    // to ensure that super-interfaces are verified first.
389    grouped
390        .into_iter()
391        .map(|(id, verifiers)| {
392            let mut dedupd_verifiers = Vec::new();
393            let mut seen = HSet::default();
394            for verifier_fn_list in verifiers {
395                for verifier in verifier_fn_list() {
396                    if seen.insert(verifier.clone()) {
397                        dedupd_verifiers.push(verifier);
398                    }
399                }
400            }
401            (id, dedupd_verifiers)
402        })
403        .collect()
404}
405
406#[doc(hidden)]
407/// Verify that all dictionary keys are unique. This is called when a [Context] is created.
408/// If any duplicate keys are found, a panic is raised with the file, line, and column
409/// information of the duplicate keys.
410pub fn verify_dict_keys() -> Result<()> {
411    let mut seen: HMap<Identifier, (&'static str, u32, u32)> = HMap::default();
412    for key in get_dict_key_ids() {
413        if let Some((file, line, column)) = seen.get(&key.id) {
414            return verify_err_noloc!(
415                "Duplicate dictionary key \"{}\" declared in {}:{}:{} and {}:{}:{}",
416                key.id,
417                file,
418                line,
419                column,
420                key.file,
421                key.line,
422                key.column
423            );
424        }
425        seen.insert(key.id.clone(), (key.file, key.line, key.column));
426    }
427    Ok(())
428}
429
430/// A macro to declare a static [Identifier] for use as a dictionary key.
431///
432/// Usage:
433/// ```
434/// # use pliron::dict_key;
435/// dict_key!(MY_KEY, "my_key");
436/// let mut ctx = pliron::context::Context::new();
437/// let aux_data_index = ctx.aux_data.insert(Box::new(42));
438/// ctx.aux_data_map.insert(MY_KEY.clone(), aux_data_index);
439/// assert_eq!(ctx.aux_data[aux_data_index].downcast_ref::<i32>(), Some(&42));
440/// assert_eq!(ctx.aux_data_map[&*MY_KEY], aux_data_index);
441/// ```
442/// Here, `MY_KEY` is the name of the static variable, and `"my_key"` is the
443/// string value of the [Identifier]. The macro will create a static variable
444/// of type [`LazyLock<Identifier>`](LazyLock) with the name `MY_KEY`.
445#[macro_export]
446macro_rules! dict_key {
447    (   $(#[$outer:meta])*
448        $decl:ident, $name:expr
449    ) => {
450        // Create a static variable linked to the DICT_KEY_IDS slice
451        // to ensure that all keys are unique.
452        // The static variable is created in a separate anonymous module.
453        const _: () = {
454            #[cfg_attr(not(target_family = "wasm"),
455                ::pliron::linkme::distributed_slice(::pliron::context::DICT_KEY_IDS), linkme(crate = ::pliron::linkme))]
456            pub static $decl: $crate::std_deps::sync::LazyLock<::pliron::context::DictKeyId> =
457                $crate::std_deps::sync::LazyLock::new(|| ::pliron::context::DictKeyId {
458                    id: $name.try_into().unwrap(),
459                    file: file!(),
460                    line: line!(),
461                    column: column!(),
462                });
463
464            #[cfg(target_family = "wasm")]
465            ::pliron::inventory::submit! {
466                ::pliron::InventoryWrapper(&$decl)
467            }
468        };
469        $(#[$outer])*
470        // Create a static variable with the provided name to access the identifier.
471        pub static $decl: $crate::std_deps::sync::LazyLock<::pliron::identifier::Identifier> =
472            $crate::std_deps::sync::LazyLock::new(|| $name.try_into().unwrap());
473    };
474}
475
476/// A macro to register a [ContextRegistration]. The argument function
477/// will be called with a mutable reference to the [Context] when a [Context] is created.
478/// Use this outside of any function (e.g. in the module scope).
479///
480/// Usage:
481/// ```
482/// use pliron::context_registration;
483/// context_registration!(my_registration_fn);
484/// fn my_registration_fn(_: &mut pliron::context::Context) {}
485/// ```
486/// Here, `my_registration_fn` is a function matching [ContextRegistration].
487#[macro_export]
488macro_rules! context_registration {
489    (   $(#[$outer:meta])*
490        $registration:expr
491    ) => {
492        const _: () = {
493            $(#[$outer])*
494            #[cfg_attr(not(target_family = "wasm"),
495                ::pliron::linkme::distributed_slice(::pliron::context::CONTEXT_REGISTRATIONS), linkme(crate = ::pliron::linkme))]
496            static CONTEXT_REGISTRATION: ::pliron::context::ContextRegistration = $registration;
497
498            #[cfg(target_family = "wasm")]
499            ::pliron::inventory::submit! {
500                ::pliron::InventoryWrapper(&CONTEXT_REGISTRATION)
501            }
502        };
503    };
504}
pliron/context.rs

pliron/
context.rs
