1use alloc::boxed::Box;
2use alloc::string::ToString;
3use alloc::sync::Arc;
4
5use miden_core::mast::MastNodeExt;
6use miden_mast_package::Package;
7use miden_mast_package::debug_info::PackageDebugInfo;
8use miden_processor::LoadedMastForest;
9use thiserror::Error;
10
11use crate::assembly::Path;
12use crate::package::{loaded_mast_forest, package_debug_info};
13use crate::utils::create_external_node_forest;
14use crate::utils::serde::{
15 ByteReader,
16 ByteWriter,
17 Deserializable,
18 DeserializationError,
19 Serializable,
20};
21use crate::vm::AdviceMap;
22use crate::{MastForest, MastNodeId, Word};
23
24#[derive(Debug, Error)]
29pub enum MastForestScriptError {
30 #[error("entrypoint node {0} is not in the provided MAST forest")]
31 EntrypointNotInForest(MastNodeId),
32 #[error("package does not contain a procedure with '@{0}' attribute")]
33 NoProcedureWithAttribute(Box<str>),
34 #[error("package contains multiple procedures with '@{0}' attribute")]
35 MultipleProceduresWithAttribute(Box<str>),
36 #[error("procedure at path '{0}' not found in package")]
37 ProcedureNotFound(Box<str>),
38 #[error("procedure at path '{0}' does not have the specified attribute")]
39 ProcedureMissingAttribute(Box<str>),
40 #[error("expected a library package, but the provided package is an executable")]
41 ExecutablePackage,
42}
43
44#[derive(Debug, Clone)]
54pub(crate) struct MastForestScript {
55 mast: Arc<MastForest>,
56 entrypoint: MastNodeId,
57 package_debug_info: Option<Arc<PackageDebugInfo>>,
58}
59
60impl MastForestScript {
61 pub fn from_parts(
69 mast: Arc<MastForest>,
70 entrypoint: MastNodeId,
71 ) -> Result<Self, MastForestScriptError> {
72 if mast.get_node_by_id(entrypoint).is_none() {
73 return Err(MastForestScriptError::EntrypointNotInForest(entrypoint));
74 }
75 Ok(Self {
76 mast,
77 entrypoint,
78 package_debug_info: None,
79 })
80 }
81
82 pub(crate) fn from_package(
88 package: &Package,
89 attribute: &str,
90 ) -> Result<Self, MastForestScriptError> {
91 if package.is_program() {
92 return Err(MastForestScriptError::ExecutablePackage);
93 }
94
95 let mut entrypoint = None;
96
97 for export in package.manifest.exports() {
98 if let Some(proc_export) = export.as_procedure()
99 && proc_export.attributes.has(attribute)
100 {
101 if entrypoint.is_some() {
102 return Err(MastForestScriptError::MultipleProceduresWithAttribute(
103 attribute.into(),
104 ));
105 }
106 entrypoint = Some(proc_export.node.ok_or_else(|| {
107 MastForestScriptError::NoProcedureWithAttribute(attribute.into())
108 })?);
109 }
110 }
111
112 let entrypoint = entrypoint
113 .ok_or_else(|| MastForestScriptError::NoProcedureWithAttribute(attribute.into()))?;
114
115 Ok(Self::from_parts(package.mast_forest().clone(), entrypoint)?
116 .with_package_debug_info(package))
117 }
118
119 pub(crate) fn from_package_reference(
128 package: &Package,
129 path: &Path,
130 attribute: &str,
131 ) -> Result<Self, MastForestScriptError> {
132 let export = package
133 .manifest
134 .exports()
135 .find(|e| e.path().as_ref() == path)
136 .ok_or_else(|| MastForestScriptError::ProcedureNotFound(path.to_string().into()))?;
137
138 let proc_export = export
139 .as_procedure()
140 .ok_or_else(|| MastForestScriptError::ProcedureNotFound(path.to_string().into()))?;
141
142 if !proc_export.attributes.has(attribute) {
143 return Err(MastForestScriptError::ProcedureMissingAttribute(path.to_string().into()));
144 }
145
146 let digest = proc_export.digest;
147
148 let (mast, entrypoint) = create_external_node_forest(digest);
149
150 Ok(Self::from_parts(Arc::new(mast), entrypoint)?.with_package_debug_info(package))
151 }
152
153 pub fn mast(&self) -> Arc<MastForest> {
158 self.mast.clone()
159 }
160
161 pub fn loaded_mast_forest(&self) -> LoadedMastForest {
163 loaded_mast_forest(self.mast.clone(), self.package_debug_info.clone())
164 }
165
166 pub fn digest(&self) -> Word {
168 self.mast[self.entrypoint].digest()
169 }
170
171 pub fn entrypoint(&self) -> MastNodeId {
173 self.entrypoint
174 }
175
176 pub fn clear_debug_info(&mut self) {
178 self.package_debug_info = None;
179 }
180
181 pub fn with_package_debug_info(mut self, package: &Package) -> Self {
184 self.package_debug_info = package_debug_info(package);
185 self
186 }
187
188 pub fn with_advice_map(mut self, advice_map: AdviceMap) -> Self {
194 if advice_map.is_empty() {
195 return self;
196 }
197
198 let mast = (*self.mast).clone().with_advice_map(advice_map);
199 self.mast = Arc::new(mast);
200 self
201 }
202}
203
204impl PartialEq for MastForestScript {
205 fn eq(&self, other: &Self) -> bool {
206 self.mast == other.mast && self.entrypoint == other.entrypoint
207 }
208}
209
210impl Eq for MastForestScript {}
211
212impl Serializable for MastForestScript {
216 fn write_into<W: ByteWriter>(&self, target: &mut W) {
217 self.mast.write_into(target);
218 target.write_u32(u32::from(self.entrypoint));
219 }
220
221 fn get_size_hint(&self) -> usize {
222 let mast_size = self.mast.to_bytes().len();
226 let u32_size = 0u32.get_size_hint();
227
228 mast_size + u32_size
229 }
230}
231
232impl Deserializable for MastForestScript {
233 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
234 let mast = MastForest::read_from(source)?;
235 let entrypoint = MastNodeId::from_u32_safe(source.read_u32()?, &mast)?;
236
237 Self::from_parts(Arc::new(mast), entrypoint)
238 .map_err(|e| DeserializationError::InvalidValue(e.to_string()))
239 }
240}