Merge pull request #181 from 0xPolygonMiden/bitwalker/assembler-refactor

chore: handle assembler refactoring changes

Cargo.toml

@@ -1,21 +1,21 @@
 [workspace]
 resolver = "2"
 members = [
-  "codegen/*",
-  "hir",
-  "hir-analysis",
-  "hir-macros",
-  "hir-symbol",
-  "hir-transform",
-  "hir-type",
-  "midenc",
-  "midenc-compile",
-  "midenc-driver",
-  "midenc-session",
-  "tools/*",
-  "frontend-wasm",
-  "tests/rust-apps/*",
-  "tests/integration",
+    "codegen/*",
+    "hir",
+    "hir-analysis",
+    "hir-macros",
+    "hir-symbol",
+    "hir-transform",
+    "hir-type",
+    "midenc",
+    "midenc-compile",
+    "midenc-driver",
+    "midenc-session",
+    "tools/*",
+    "frontend-wasm",
+    "tests/rust-apps/*",
+    "tests/integration",
 ]
 exclude = ["tests/rust-apps-wasm", "cargo-ext/tests/data"]
 
@@ -46,6 +46,7 @@ Inflector = "0.11"
 intrusive-collections = "0.9"
 inventory = "0.3"
 log = "0.4"
+miette = { version = "7.1.0", git = "https://github.com/bitwalker/miette", branch = "no-std" }
 paste = "1.0"
 parking_lot = "0.12"
 parking_lot_core = "0.9"
@@ -59,15 +60,15 @@ smallvec = { version = "1.9", features = [
     "const_new",
 ] }
 smallstr = { version = "0.3", features = ["union"] }
-thiserror = "1.0"
+thiserror = { version = "1.0", git = "https://github.com/bitwalker/thiserror", branch = "no-std" }
 toml = { version = "0.5", features = ["preserve_order"] }
 derive_more = "0.99"
 indexmap = "2.1"
-# 211152c631d16a943aae503466b198b93c61150f is latest (as of Jan 25th) commit in the next branch
-miden-assembly = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "211152c631d16a943aae503466b198b93c61150f" }
-miden-core = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "211152c631d16a943aae503466b198b93c61150f" }
-miden-processor = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "211152c631d16a943aae503466b198b93c61150f" }
-miden-stdlib = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "211152c631d16a943aae503466b198b93c61150f" }
+# 0ac54125 is the latest commit in 'next' that includes the assembler refactoring
+miden-assembly = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "0ac54125d5bd43160566736b2be3f5818fd84f69" }
+miden-core = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "0ac54125d5bd43160566736b2be3f5818fd84f69" }
+miden-processor = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "0ac54125d5bd43160566736b2be3f5818fd84f69" }
+miden-stdlib = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "0ac54125d5bd43160566736b2be3f5818fd84f69" }
 miden-codegen-masm = { path = "codegen/masm" }
 miden-diagnostics = "0.1"
 miden-hir = { path = "hir" }
@@ -84,6 +85,10 @@ midenc-session = { path = "midenc-session" }
 miden-integration-tests = { path = "tests/integration" }
 wat = "1.0.69"
 
+[patch.crates-io]
+thiserror = { git = "https://github.com/bitwalker/thiserror", branch = "no-std" }
+miette = { git = "https://github.com/bitwalker/miette", branch = "no-std" }
+
 [profile.dev]
 lto = false
 # Needed for 'inventory' to work

codegen/masm/Cargo.toml

@@ -17,6 +17,7 @@ intrusive-collections.workspace = true
 inventory.workspace = true
 log.workspace = true
 miden-assembly.workspace = true
+miden-core.workspace = true
 miden-diagnostics.workspace = true
 miden-hir.workspace = true
 miden-hir-analysis.workspace = true

codegen/masm/src/codegen/emit/mod.rs

@@ -87,6 +87,7 @@ pub mod unary;
 
 use core::ops::{Deref, DerefMut};
 
+use miden_assembly::ast::InvokeKind;
 use miden_hir::{self as hir, Immediate, Type};
 
 use super::{Operand, OperandStack};
@@ -199,27 +200,51 @@ impl<'a> OpEmitter<'a> {
         self.stack
     }
 
+    #[inline]
+    fn maybe_register_invoke(&mut self, op: &masm::Op) {
+        match op {
+            Op::Exec(id) => {
+                self.function.register_absolute_invocation_target(InvokeKind::Exec, *id)
+            }
+            Op::Call(id) => {
+                self.function.register_absolute_invocation_target(InvokeKind::Call, *id)
+            }
+            Op::Syscall(id) => {
+                self.function.register_absolute_invocation_target(InvokeKind::SysCall, *id)
+            }
+            _ => (),
+        }
+    }
+
     /// Emit `op` to the current block
     #[inline(always)]
     pub fn emit(&mut self, op: masm::Op) {
+        self.maybe_register_invoke(&op);
         self.current_block().push(op)
     }
 
     /// Emit `n` copies of `op` to the current block
     #[inline(always)]
     pub fn emit_n(&mut self, count: usize, op: masm::Op) {
+        self.maybe_register_invoke(&op);
         self.current_block().push_n(count, op);
     }
 
     /// Emit `ops` to the current block
     #[inline(always)]
     pub fn emit_all(&mut self, ops: &[masm::Op]) {
+        for op in ops {
+            self.maybe_register_invoke(op);
+        }
         self.current_block().extend_from_slice(ops);
     }
 
     /// Emit `n` copies of the sequence `ops` to the current block
     #[inline(always)]
     pub fn emit_repeat(&mut self, count: usize, ops: &[masm::Op]) {
+        for op in ops {
+            self.maybe_register_invoke(op);
+        }
         self.current_block().push_repeat(ops, count);
     }
 
@@ -229,9 +254,14 @@ impl<'a> OpEmitter<'a> {
     where
         F: Fn(usize) -> [Op; N],
     {
+        for op in template(0) {
+            self.maybe_register_invoke(&op);
+        }
+
         let block = self.current_block();
         for n in 0..count {
-            block.extend_from_slice(&template(n));
+            let ops = template(n);
+            block.extend_from_slice(&ops);
         }
     }
 
@@ -952,7 +982,7 @@ mod tests {
         assert_eq!(emitter.stack()[0], Type::I1);
         assert_eq!(emitter.stack()[1], one);
 
-        emitter.assert();
+        emitter.assert(None);
         assert_eq!(emitter.stack_len(), 1);
         assert_eq!(emitter.stack()[0], one);
 
@@ -980,7 +1010,7 @@ mod tests {
         assert_eq!(emitter.stack()[0], Type::I1);
         assert_eq!(emitter.stack()[1], one);
 
-        emitter.assertz();
+        emitter.assertz(None);
         assert_eq!(emitter.stack_len(), 1);
         assert_eq!(emitter.stack()[0], one);
 
@@ -1754,7 +1784,7 @@ mod tests {
         emitter.literal(ten);
         assert_eq!(emitter.stack_len(), 1);
 
-        emitter.assert();
+        emitter.assert(None);
         assert_eq!(emitter.stack_len(), 0);
     }
 
@@ -1770,7 +1800,7 @@ mod tests {
         emitter.literal(ten);
         assert_eq!(emitter.stack_len(), 1);
 
-        emitter.assertz();
+        emitter.assertz(None);
         assert_eq!(emitter.stack_len(), 0);
     }
 

codegen/masm/src/codegen/emit/primop.rs

@@ -7,8 +7,9 @@ impl<'a> OpEmitter<'a> {
     /// Assert that an integer value on the stack has the value 1
     ///
     /// This operation consumes the input value.
-    pub fn assert(&mut self) {
+    pub fn assert(&mut self, code: Option<u32>) {
         let arg = self.stack.pop().expect("operand stack is empty");
+        let code = code.unwrap_or_default();
         match arg.ty() {
             Type::Felt
             | Type::U32
@@ -18,13 +19,18 @@ impl<'a> OpEmitter<'a> {
             | Type::U8
             | Type::I8
             | Type::I1 => {
-                self.emit(Op::Assert);
+                self.emit(Op::AssertWithError(code));
             }
             Type::I128 => {
-                self.emit_all(&[Op::Assertz, Op::Assertz, Op::Assertz, Op::Assert]);
+                self.emit_all(&[
+                    Op::AssertzWithError(code),
+                    Op::AssertzWithError(code),
+                    Op::AssertzWithError(code),
+                    Op::AssertWithError(code),
+                ]);
             }
             Type::U64 | Type::I64 => {
-                self.emit_all(&[Op::Assertz, Op::Assert]);
+                self.emit_all(&[Op::AssertzWithError(code), Op::AssertWithError(code)]);
             }
             ty if !ty.is_integer() => {
                 panic!("invalid argument to assert: expected integer, got {ty}")
@@ -36,8 +42,9 @@ impl<'a> OpEmitter<'a> {
     /// Assert that an integer value on the stack has the value 0
     ///
     /// This operation consumes the input value.
-    pub fn assertz(&mut self) {
+    pub fn assertz(&mut self, code: Option<u32>) {
         let arg = self.stack.pop().expect("operand stack is empty");
+        let code = code.unwrap_or_default();
         match arg.ty() {
             Type::Felt
             | Type::U32
@@ -47,10 +54,10 @@ impl<'a> OpEmitter<'a> {
             | Type::U8
             | Type::I8
             | Type::I1 => {
-                self.emit(Op::Assertz);
+                self.emit(Op::AssertzWithError(code));
             }
             ty @ (Type::I128 | Type::U64 | Type::I64) => {
-                self.emit_n(ty.size_in_bits() / 32, Op::Assertz);
+                self.emit_n(ty.size_in_bits() / 32, Op::AssertzWithError(code));
             }
             ty if !ty.is_integer() => {
                 panic!("invalid argument to assertz: expected integer, got {ty}")

codegen/masm/src/codegen/emit/unary.rs

@@ -1280,7 +1280,7 @@ impl<'a> OpEmitter<'a> {
                 self.emit_all(&[Op::Pow2, Op::U32Assert]);
             }
             Type::I32 => {
-                self.emit_all(&[Op::Exec("intrinsics::i32::pow2".parse().unwrap())]);
+                self.emit(Op::Exec("intrinsics::i32::pow2".parse().unwrap()));
             }
             Type::U8 | Type::U16 => {
                 self.emit_all(&[Op::Pow2, Op::U32Assert]);

codegen/masm/src/codegen/emitter.rs

@@ -627,8 +627,20 @@ impl<'b, 'f: 'b> BlockEmitter<'b, 'f> {
     }
 
     fn emit_primop_imm(&mut self, inst_info: &InstInfo, op: &hir::PrimOpImm) {
+        let args = op.args.as_slice(&self.function.f.dfg.value_lists);
         let mut emitter = self.inst_emitter(inst_info.inst);
         match op.op {
+            hir::Opcode::Assert => {
+                assert_eq!(args.len(), 1);
+                emitter
+                    .assert(Some(op.imm.as_u32().expect("invalid assertion error code immediate")));
+            }
+            hir::Opcode::Assertz => {
+                assert_eq!(args.len(), 1);
+                emitter.assertz(Some(
+                    op.imm.as_u32().expect("invalid assertion error code immediate"),
+                ));
+            }
             hir::Opcode::AssertEq => {
                 emitter.assert_eq_imm(op.imm);
             }
@@ -648,12 +660,12 @@ impl<'b, 'f: 'b> BlockEmitter<'b, 'f> {
             // Pop a value of the given type off the stack and assert it's value is one
             hir::Opcode::Assert => {
                 assert_eq!(args.len(), 1);
-                emitter.assert();
+                emitter.assert(None);
             }
             // Pop a value of the given type off the stack and assert it's value is zero
             hir::Opcode::Assertz => {
                 assert_eq!(args.len(), 1);
-                emitter.assertz();
+                emitter.assertz(None);
             }
             // Pop two values of the given type off the stack and assert equality
             hir::Opcode::AssertEq => {
@@ -756,7 +768,29 @@ impl<'b, 'f: 'b> BlockEmitter<'b, 'f> {
                         *body_blk = mapped_blocks[prev_body_blk];
                         rewrites.push((prev_body_blk, *body_blk));
                     }
-                    Op::LocAddr(_) | Op::LocStore(_) | Op::LocStorew(_) => {
+                    Op::Exec(id) => {
+                        self.function.f_prime.register_absolute_invocation_target(
+                            miden_assembly::ast::InvokeKind::Exec,
+                            id,
+                        );
+                    }
+                    Op::Call(id) => {
+                        self.function.f_prime.register_absolute_invocation_target(
+                            miden_assembly::ast::InvokeKind::Call,
+                            id,
+                        );
+                    }
+                    Op::Syscall(id) => {
+                        self.function.f_prime.register_absolute_invocation_target(
+                            miden_assembly::ast::InvokeKind::SysCall,
+                            id,
+                        );
+                    }
+                    Op::LocAddr(_)
+                    | Op::LocLoad(_)
+                    | Op::LocLoadw(_)
+                    | Op::LocStore(_)
+                    | Op::LocStorew(_) => {
                         unimplemented!(
                             "locals are not currently supported in inline assembly blocks"
                         )

codegen/masm/src/convert.rs

@@ -1,3 +1,4 @@
+use miden_assembly::LibraryPath;
 use miden_hir::{
     self as hir,
     pass::{AnalysisManager, ConversionPass, ConversionResult},
@@ -55,7 +56,7 @@ impl ConversionPass for ConvertHirToMasm<hir::Program> {
         analyses: &mut AnalysisManager,
         session: &Session,
     ) -> ConversionResult<Self::To> {
-        let mut masm_program = Box::new(masm::Program::from(program.as_ref()));
+        let mut masm_program = Box::new(masm::Program::from_hir(&program));
 
         // Remove the set of modules to compile from the program
         let modules = program.modules_mut().take();
@@ -104,9 +105,18 @@ impl ConversionPass for ConvertHirToMasm<hir::Module> {
         analyses: &mut AnalysisManager,
         session: &Session,
     ) -> ConversionResult<Self::To> {
+        use miden_assembly::ast::ModuleKind;
         use miden_hir::ProgramAnalysisKey;
 
-        let mut masm_module = Box::new(masm::Module::new(module.name));
+        let kind = if module.is_kernel() {
+            ModuleKind::Kernel
+        } else {
+            ModuleKind::Library
+        };
+        let name = LibraryPath::new(&module.name).unwrap_or_else(|err| {
+            panic!("invalid module name '{}': {}", module.name.as_str(), err)
+        });
+        let mut masm_module = Box::new(masm::Module::new(name, kind));
 
         // Compute import information for this module
         masm_module.imports = module.imports();