Skip to main content

midenc_codegen_masm/
artifact.rs

1use alloc::sync::Arc;
2use core::fmt;
3
4use miden_assembly::{Path, ast::InvocationTarget};
5use miden_core::Word;
6use miden_mast_package::Package;
7use midenc_hir::{constants::ConstantData, dialects::builtin, interner::Symbol};
8use midenc_session::{
9    Emit, OutputMode, OutputType, Session, Writer,
10    diagnostics::{IntoDiagnostic, Report, SourceSpan, Span, WrapErr},
11};
12
13use crate::{TraceEvent, lower::NativePtr, masm};
14
15mod project_support;
16
17pub struct MasmComponent {
18    pub id: Option<builtin::ComponentId>,
19    /// The path of the root module for this component
20    ///
21    /// All components must have a canonical root module, even if empty
22    pub root: Arc<Path>,
23    /// The symbol name of the component initializer function
24    ///
25    /// This function is responsible for initializing global variables and writing data segments
26    /// into memory at program startup, and at cross-context call boundaries (in callee prologue).
27    pub init: Option<masm::InvocationTarget>,
28    /// The symbol name of the program entrypoint, if this component is executable.
29    ///
30    /// If unset, it indicates that the component is a library, even if it could be made executable.
31    pub entrypoint: Option<masm::InvocationTarget>,
32    /// The kernel library to link against
33    pub kernel: Option<masm::KernelLibrary>,
34    /// The rodata segments of this component keyed by the offset of the segment
35    pub rodata: Vec<Rodata>,
36    /// The address of the start of the global heap
37    pub heap_base: u32,
38    /// The address of the `__stack_pointer` global, if such a global has been defined
39    pub stack_pointer: Option<u32>,
40    /// The set of modules in this component
41    pub modules: Vec<Arc<masm::Module>>,
42}
43
44impl Emit for MasmComponent {
45    fn name(&self) -> Option<Symbol> {
46        None
47    }
48
49    fn output_type(&self, _mode: OutputMode) -> OutputType {
50        OutputType::Masm
51    }
52
53    fn write_to<W: Writer>(
54        &self,
55        mut writer: W,
56        mode: OutputMode,
57        _session: &Session,
58    ) -> anyhow::Result<()> {
59        if mode != OutputMode::Text {
60            anyhow::bail!("masm emission does not support binary mode");
61        }
62        writer.write_fmt(core::format_args!("{self}"))?;
63        Ok(())
64    }
65}
66
67/// Represents a read-only data segment, combined with its content digest
68#[derive(Clone, PartialEq, Eq)]
69pub struct Rodata {
70    /// The component to which this read-only data segment belongs
71    pub component: builtin::ComponentId,
72    /// The content digest computed for `data`
73    pub digest: Word,
74    /// The address at which the data for this segment begins
75    pub start: NativePtr,
76    /// The raw binary data for this segment
77    pub data: Arc<ConstantData>,
78}
79impl fmt::Debug for Rodata {
80    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
81        f.debug_struct("Rodata")
82            .field("digest", &format_args!("{}", &self.digest))
83            .field("start", &self.start)
84            .field_with("data", |f| {
85                f.debug_struct("ConstantData")
86                    .field("len", &self.data.len())
87                    .finish_non_exhaustive()
88            })
89            .finish()
90    }
91}
92impl Rodata {
93    pub fn size_in_bytes(&self) -> usize {
94        self.data.len()
95    }
96
97    pub fn size_in_felts(&self) -> usize {
98        self.data.len().div_ceil(4)
99    }
100
101    pub fn size_in_words(&self) -> usize {
102        self.size_in_felts().div_ceil(4)
103    }
104
105    /// Attempt to convert this rodata object to its equivalent representation in felts
106    ///
107    /// See [Self::bytes_to_elements] for more details.
108    pub fn to_elements(&self) -> Vec<miden_processor::Felt> {
109        Self::bytes_to_elements(self.data.as_slice())
110    }
111
112    /// Attempt to convert the given bytes to their equivalent representation in felts
113    ///
114    /// The resulting felts will be in padded out to the nearest number of words, i.e. if the data
115    /// only takes up 3 felts worth of bytes, then the resulting `Vec` will contain 4 felts, so that
116    /// the total size is a valid number of words.
117    pub fn bytes_to_elements(bytes: &[u8]) -> Vec<miden_processor::Felt> {
118        use miden_processor::Felt;
119
120        let mut felts = Vec::with_capacity(bytes.len() / 4);
121        let mut iter = bytes.iter().copied().array_chunks::<4>();
122        felts.extend(
123            iter.by_ref().map(|chunk| Felt::new_unchecked(u32::from_le_bytes(chunk) as u64)),
124        );
125        let remainder = iter.into_remainder();
126        if remainder.len() > 0 {
127            let mut chunk = [0u8; 4];
128            for (i, byte) in remainder.enumerate() {
129                chunk[i] = byte;
130            }
131            felts.push(Felt::new_unchecked(u32::from_le_bytes(chunk) as u64));
132        }
133
134        let size_in_felts = bytes.len().div_ceil(4);
135        let size_in_words = size_in_felts.div_ceil(4);
136        let padding = (size_in_words * 4).abs_diff(felts.len());
137        felts.resize(felts.len() + padding, Felt::ZERO);
138        debug_assert_eq!(felts.len() % 4, 0, "expected to be a valid number of words");
139        felts
140    }
141}
142
143inventory::submit! {
144    midenc_session::CompileFlag::new("test_harness")
145        .long("test-harness")
146        .action(midenc_session::FlagAction::SetTrue)
147        .help("If present, causes the code generator to emit extra code for the VM test harness")
148        .help_heading("Testing")
149}
150
151impl fmt::Display for MasmComponent {
152    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
153        use crate::intrinsics::INTRINSICS_MODULE_NAMES;
154
155        for module in self.modules.iter() {
156            // Skip printing the standard library modules and intrinsics modules to focus on the
157            // user-defined modules and avoid the
158            // stack overflow error when printing large programs
159            // https://github.com/0xMiden/miden-formatting/issues/4
160            let module_name = module.path().as_str();
161            let module_name_trimmed = module_name.trim_start_matches("::");
162            if INTRINSICS_MODULE_NAMES.contains(&module_name) {
163                continue;
164            }
165            if module.is_in_namespace(Path::new("std"))
166                || module_name_trimmed.starts_with("miden::core")
167                || module_name_trimmed.starts_with("miden::protocol")
168            {
169                continue;
170            } else {
171                writeln!(f, "# mod {}\n", &module_name)?;
172                writeln!(f, "{module}")?;
173            }
174        }
175        Ok(())
176    }
177}
178
179impl MasmComponent {
180    /// Assemble this component into a Miden package.
181    pub fn assemble(
182        &self,
183        account_component_metadata_bytes: Option<&[u8]>,
184        session: &Session,
185    ) -> Result<Arc<Package>, Report> {
186        project_support::assemble(self, account_component_metadata_bytes, session)
187    }
188
189    /// Assemble this component into a Miden package using a pre-populated package registry.
190    pub fn assemble_with_registry(
191        &self,
192        account_component_metadata_bytes: Option<&[u8]>,
193        session: &Session,
194        registry: &mut midenc_session::registry::HybridPackageRegistry,
195    ) -> Result<Arc<Package>, Report> {
196        project_support::assemble_with_registry(
197            self,
198            account_component_metadata_bytes,
199            session,
200            registry,
201        )
202    }
203
204    /// Generate an executable module which when run expects the raw data segment data to be
205    /// provided on the advice stack in the same order as initialization, and the operands of
206    /// the entrypoint function on the operand stack.
207    fn generate_main(
208        &self,
209        entrypoint: &InvocationTarget,
210        emit_test_harness: bool,
211        source_manager: Arc<dyn midenc_session::SourceManager + Send + Sync>,
212    ) -> Result<Box<masm::Module>, Report> {
213        use masm::{Instruction as Inst, Op};
214
215        let mut exe = Box::new(masm::Module::new_executable());
216        let span = SourceSpan::default();
217        let mut invoked = Vec::new();
218        let body = {
219            let mut block = masm::Block::new(span, Vec::with_capacity(64));
220            // Invoke component initializer, if present
221            if let Some(init) = self.init.as_ref() {
222                invoked.push(masm::Invoke::new(masm::InvokeKind::Exec, init.clone()));
223                block.push(Op::Inst(Span::new(span, Inst::Exec(init.clone()))));
224            }
225
226            // Initialize test harness, if requested
227            if emit_test_harness {
228                self.emit_test_harness(&mut block);
229            }
230
231            // Invoke the program entrypoint
232            block.push(Op::Inst(Span::new(
233                span,
234                Inst::Trace(TraceEvent::FrameStart.as_u32().into()),
235            )));
236            invoked.push(masm::Invoke::new(masm::InvokeKind::Exec, entrypoint.clone()));
237            block.push(Op::Inst(Span::new(span, Inst::Exec(entrypoint.clone()))));
238            block
239                .push(Op::Inst(Span::new(span, Inst::Trace(TraceEvent::FrameEnd.as_u32().into()))));
240
241            // Truncate the stack to 16 elements on exit
242            let truncate_stack = {
243                let name = masm::ProcedureName::new("truncate_stack").unwrap();
244                let module = masm::LibraryPath::new("::miden::core::sys").unwrap();
245                let qualified = masm::QualifiedProcedureName::new(module.as_path(), name);
246                InvocationTarget::Path(Span::new(span, qualified.into_inner()))
247            };
248            invoked.push(masm::Invoke::new(masm::InvokeKind::Exec, truncate_stack.clone()));
249            block.push(Op::Inst(Span::new(span, Inst::Exec(truncate_stack))));
250            block
251        };
252        let mut start = masm::Procedure::new(
253            span,
254            masm::Visibility::Public,
255            masm::ProcedureName::main(),
256            0,
257            body,
258        );
259        start.extend_invoked(invoked);
260        exe.define_procedure(start, source_manager)
261            .into_diagnostic()
262            .wrap_err("failed to define executable `main` procedure")?;
263        Ok(exe)
264    }
265
266    fn emit_test_harness(&self, block: &mut masm::Block) {
267        use masm::{Instruction as Inst, IntValue, Op, PushValue};
268        use miden_core::Felt;
269
270        let span = SourceSpan::default();
271
272        let pipe_words_to_memory = {
273            let name = masm::ProcedureName::new("pipe_words_to_memory").unwrap();
274            let module = masm::LibraryPath::new("::miden::core::mem").unwrap();
275            let qualified = masm::QualifiedProcedureName::new(module.as_path(), name);
276            InvocationTarget::Path(Span::new(span, qualified.into_inner()))
277        };
278
279        // Step 1: Get the number of initializers to run
280        // => [inits] on operand stack
281        block.push(Op::Inst(Span::new(span, Inst::AdvPush)));
282
283        // Step 2: Evaluate the initial state of the loop condition `inits > 0`
284        // => [inits, inits]
285        block.push(Op::Inst(Span::new(span, Inst::Dup0)));
286        // => [inits > 0, inits]
287        block.push(Op::Inst(Span::new(span, Inst::Push(PushValue::Int(IntValue::U8(0)).into()))));
288        block.push(Op::Inst(Span::new(span, Inst::Gt)));
289
290        // Step 3: Loop until `inits == 0`
291        let mut loop_body = Vec::with_capacity(16);
292
293        // State of operand stack on entry to `loop_body`: [inits]
294        // State of advice stack on entry to `loop_body`: [dest_ptr, num_words, ...]
295        //
296        // Step 3a: Compute next value of `inits`, i.e. `inits'`
297        // => [inits - 1]
298        loop_body.push(Op::Inst(Span::new(span, Inst::SubImm(Felt::ONE.into()))));
299
300        // Step 3b: Copy initializer data to memory
301        // => [num_words, dest_ptr, inits']
302        loop_body.push(Op::Inst(Span::new(span, Inst::AdvPush)));
303        loop_body.push(Op::Inst(Span::new(span, Inst::AdvPush)));
304        // => [C, B, A, dest_ptr, inits'] on operand stack
305        loop_body
306            .push(Op::Inst(Span::new(span, Inst::Trace(TraceEvent::FrameStart.as_u32().into()))));
307        loop_body.push(Op::Inst(Span::new(span, Inst::Exec(pipe_words_to_memory))));
308        loop_body
309            .push(Op::Inst(Span::new(span, Inst::Trace(TraceEvent::FrameEnd.as_u32().into()))));
310        // Drop C, B, A
311        loop_body.push(Op::Inst(Span::new(span, Inst::DropW)));
312        loop_body.push(Op::Inst(Span::new(span, Inst::DropW)));
313        loop_body.push(Op::Inst(Span::new(span, Inst::DropW)));
314        // => [inits']
315        loop_body.push(Op::Inst(Span::new(span, Inst::Drop)));
316
317        // Step 3c: Evaluate loop condition `inits' > 0`
318        // => [inits', inits']
319        loop_body.push(Op::Inst(Span::new(span, Inst::Dup0)));
320        // => [inits' > 0, inits']
321        loop_body
322            .push(Op::Inst(Span::new(span, Inst::Push(PushValue::Int(IntValue::U8(0)).into()))));
323        loop_body.push(Op::Inst(Span::new(span, Inst::Gt)));
324
325        // Step 4: Enter (or skip) loop
326        block.push(Op::While {
327            span,
328            body: masm::Block::new(span, loop_body),
329        });
330
331        // Step 5: Drop `inits` after loop is evaluated
332        block.push(Op::Inst(Span::new(span, Inst::Drop)));
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use proptest::prelude::*;
339
340    use super::*;
341
342    fn validate_bytes_to_elements(bytes: &[u8]) {
343        let result = Rodata::bytes_to_elements(bytes);
344
345        // Each felt represents 4 bytes
346        let expected_felts = bytes.len().div_ceil(4);
347        // Felts should be padded to a multiple of 4 (1 word = 4 felts)
348        let expected_total_felts = expected_felts.div_ceil(4) * 4;
349
350        assert_eq!(
351            result.len(),
352            expected_total_felts,
353            "For {} bytes, expected {} felts (padded from {} felts), but got {}",
354            bytes.len(),
355            expected_total_felts,
356            expected_felts,
357            result.len()
358        );
359
360        // Verify padding is zeros
361        for (i, felt) in result.iter().enumerate().skip(expected_felts) {
362            assert_eq!(*felt, miden_processor::Felt::ZERO, "Padding at index {i} should be zero");
363        }
364    }
365
366    #[test]
367    fn test_bytes_to_elements_edge_cases() {
368        validate_bytes_to_elements(&[]);
369        validate_bytes_to_elements(&[1]);
370        validate_bytes_to_elements(&[0u8; 4]);
371        validate_bytes_to_elements(&[0u8; 15]);
372        validate_bytes_to_elements(&[0u8; 16]);
373        validate_bytes_to_elements(&[0u8; 17]);
374        validate_bytes_to_elements(&[0u8; 31]);
375        validate_bytes_to_elements(&[0u8; 32]);
376        validate_bytes_to_elements(&[0u8; 33]);
377        validate_bytes_to_elements(&[0u8; 64]);
378    }
379
380    proptest! {
381        #![proptest_config(ProptestConfig::with_cases(1000))]
382        #[test]
383        fn proptest_bytes_to_elements(bytes in prop::collection::vec(any::<u8>(), 0..=1000)) {
384            validate_bytes_to_elements(&bytes);
385        }
386
387        #[test]
388        fn proptest_bytes_to_elements_word_boundaries(size_factor in 0u32..=100) {
389            // Test specifically around word boundaries
390            // Test sizes around multiples of 16 (since 1 word = 4 felts = 16 bytes)
391            let base_size = size_factor * 16;
392            for offset in -2i32..=2 {
393                let size = (base_size as i32 + offset).max(0) as usize;
394                let bytes = vec![0u8; size];
395                validate_bytes_to_elements(&bytes);
396            }
397        }
398    }
399}