miden_assembly_syntax/library/
module.rs

1use alloc::{sync::Arc, vec::Vec};
2use core::ops::Index;
3
4use midenc_hir_type::FunctionType;
5
6use crate::{
7    Path, Word,
8    ast::{self, AttributeSet, ConstantValue, Ident, ItemIndex, ProcedureName},
9};
10
11// MODULE INFO
12// ================================================================================================
13
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct ModuleInfo {
16    path: Arc<Path>,
17    items: Vec<ItemInfo>,
18}
19
20impl ModuleInfo {
21    /// Returns a new [`ModuleInfo`] instantiated library path.
22    pub fn new(path: Arc<Path>) -> Self {
23        Self { path, items: Vec::new() }
24    }
25
26    /// Adds a procedure to the module.
27    pub fn add_procedure(
28        &mut self,
29        name: ProcedureName,
30        digest: Word,
31        signature: Option<Arc<FunctionType>>,
32        attributes: AttributeSet,
33    ) {
34        self.items
35            .push(ItemInfo::Procedure(ProcedureInfo { name, digest, signature, attributes }));
36    }
37
38    /// Adds a constant to the module.
39    pub fn add_constant(&mut self, name: Ident, value: ConstantValue) {
40        self.items.push(ItemInfo::Constant(ConstantInfo { name, value }));
41    }
42
43    /// Adds a type declaration to the module.
44    pub fn add_type(&mut self, name: Ident, ty: ast::types::Type) {
45        self.items.push(ItemInfo::Type(TypeInfo { name, ty }));
46    }
47
48    /// Returns the module's library path.
49    pub fn path(&self) -> &Path {
50        &self.path
51    }
52
53    /// Returns the number of procedures in the module.
54    pub fn num_procedures(&self) -> usize {
55        self.items.iter().filter(|item| matches!(item, ItemInfo::Procedure(_))).count()
56    }
57
58    /// Returns the [`ItemInfo`] of the item at the provided index, if any.
59    pub fn get_item_by_index(&self, index: ItemIndex) -> Option<&ItemInfo> {
60        self.items.get(index.as_usize())
61    }
62
63    /// Returns the [ItemIndex] of an item by its local name
64    pub fn get_item_index_by_name(&self, name: &str) -> Option<ItemIndex> {
65        self.items.iter().enumerate().find_map(|(idx, info)| {
66            if info.name().as_str() == name {
67                Some(ItemIndex::new(idx))
68            } else {
69                None
70            }
71        })
72    }
73
74    /// Returns the digest of the procedure with the provided name, if any.
75    pub fn get_procedure_digest_by_name(&self, name: &str) -> Option<Word> {
76        self.items.iter().find_map(|info| match info {
77            ItemInfo::Procedure(proc) if proc.name.as_str() == name => Some(proc.digest),
78            _ => None,
79        })
80    }
81
82    /// Returns an iterator over the items in the module with their corresponding item index in the
83    /// module.
84    pub fn items(&self) -> impl ExactSizeIterator<Item = (ItemIndex, &ItemInfo)> {
85        self.items.iter().enumerate().map(|(idx, item)| (ItemIndex::new(idx), item))
86    }
87
88    /// Returns an iterator over the procedure infos in the module with their corresponding
89    /// item index in the module.
90    pub fn procedures(&self) -> impl Iterator<Item = (ItemIndex, &ProcedureInfo)> {
91        self.items.iter().enumerate().filter_map(|(idx, item)| match item {
92            ItemInfo::Procedure(proc) => Some((ItemIndex::new(idx), proc)),
93            _ => None,
94        })
95    }
96
97    /// Returns an iterator over the MAST roots of procedures defined in this module.
98    pub fn procedure_digests(&self) -> impl Iterator<Item = Word> + '_ {
99        self.items.iter().filter_map(|item| match item {
100            ItemInfo::Procedure(proc) => Some(proc.digest),
101            _ => None,
102        })
103    }
104
105    /// Access the constants associated with this module
106    pub fn constants(&self) -> impl Iterator<Item = (ItemIndex, &ConstantInfo)> {
107        self.items.iter().enumerate().filter_map(|(idx, item)| match item {
108            ItemInfo::Constant(info) => Some((ItemIndex::new(idx), info)),
109            _ => None,
110        })
111    }
112
113    /// Access the type declarations associated with this module
114    pub fn types(&self) -> impl Iterator<Item = (ItemIndex, &TypeInfo)> {
115        self.items.iter().enumerate().filter_map(|(idx, item)| match item {
116            ItemInfo::Type(info) => Some((ItemIndex::new(idx), info)),
117            _ => None,
118        })
119    }
120}
121
122impl Index<ItemIndex> for ModuleInfo {
123    type Output = ItemInfo;
124
125    fn index(&self, index: ItemIndex) -> &Self::Output {
126        &self.items[index.as_usize()]
127    }
128}
129
130/// Stores information about an item
131#[derive(Debug, Clone, PartialEq, Eq)]
132pub enum ItemInfo {
133    Procedure(ProcedureInfo),
134    Constant(ConstantInfo),
135    Type(TypeInfo),
136}
137
138impl ItemInfo {
139    pub fn name(&self) -> &Ident {
140        match self {
141            Self::Procedure(info) => info.name.as_ref(),
142            Self::Constant(info) => &info.name,
143            Self::Type(info) => &info.name,
144        }
145    }
146
147    pub fn attributes(&self) -> Option<&AttributeSet> {
148        match self {
149            Self::Procedure(info) => Some(&info.attributes),
150            Self::Constant(_) | Self::Type(_) => None,
151        }
152    }
153
154    pub fn unwrap_procedure(&self) -> &ProcedureInfo {
155        match self {
156            Self::Procedure(info) => info,
157            Self::Constant(_) | Self::Type(_) => panic!("expected item to be a procedure"),
158        }
159    }
160}
161
162/// Stores the name and digest of a procedure.
163#[derive(Debug, Clone, PartialEq, Eq)]
164pub struct ProcedureInfo {
165    pub name: ProcedureName,
166    pub digest: Word,
167    pub signature: Option<Arc<FunctionType>>,
168    pub attributes: AttributeSet,
169}
170
171/// Stores the name and value of a constant
172#[derive(Debug, Clone, PartialEq, Eq)]
173pub struct ConstantInfo {
174    pub name: Ident,
175    pub value: ConstantValue,
176}
177
178/// Stores the name and concrete type of a type declaration
179#[derive(Debug, Clone, PartialEq, Eq)]
180pub struct TypeInfo {
181    pub name: Ident,
182    pub ty: ast::types::Type,
183}