Merge pull request #450 from maticnetwork/refactor-stdlib-in-assembler

Provide stdlib via module provider to assembler

assembly/Cargo.toml

@@ -27,7 +27,6 @@ std = ["vm-core/std"]
 crypto = { package = "winter-crypto", version = "0.4", default-features = false }
 num_enum = "0.5.7"
 vm-core = { package = "miden-core", path = "../core", version = "0.3", default-features = false }
-vm-stdlib = { package = "miden-stdlib", path = "../stdlib", version = "0.2", default-features = false }
 
 [dev-dependencies]
-criterion = "0.3"
+criterion = "0.3"

assembly/src/lib.rs

@@ -9,10 +9,10 @@ use vm_core::{
     utils::{
         collections::{BTreeMap, Vec},
         string::{String, ToString},
+        Box,
     },
-    CodeBlockTable, Kernel, Library, Program,
+    CodeBlockTable, Kernel, Program,
 };
-use vm_stdlib::StdLibrary;
 
 mod context;
 use context::AssemblyContext;
@@ -43,6 +43,28 @@ const MODULE_PATH_DELIM: &str = "::";
 type ProcMap = BTreeMap<String, Procedure>;
 type ModuleMap = BTreeMap<String, ProcMap>;
 
+// MODULE PROVIDER
+// ================================================================================================
+
+/// The module provider is now a simplified version of a module cache. It is expected to evolve to
+/// a general solution for the module lookup
+///
+/// TODO compute a procedure index deterministically from a module path & procedure label. The
+/// initial expected layout for the index is `[u8; 20]`. The module provider should map procedures
+/// from these indexes, and this is expected to result in code simplification + optimization for
+/// the compiler.
+pub trait ModuleProvider {
+    /// Fetch source contents provided a module path
+    fn get_source(&self, path: &str) -> Option<&str>;
+}
+
+// A default provider that won't resolve modules
+impl ModuleProvider for () {
+    fn get_source(&self, _path: &str) -> Option<&str> {
+        None
+    }
+}
+
 // ASSEMBLER
 // ================================================================================================
 
@@ -54,7 +76,7 @@ type ModuleMap = BTreeMap<String, ProcMap>;
 /// - Via the `new()` constructor. In this case, the kernel is assumed to be empty, and the
 ///   programs compiled using such assembler cannot contain `syscall` instructions.
 pub struct Assembler {
-    stdlib: StdLibrary,
+    module_provider: Box<dyn ModuleProvider>,
     parsed_modules: ModuleMap,
     kernel_procs: ProcMap,
     kernel: Kernel,
@@ -65,28 +87,37 @@ impl Assembler {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
     /// Returns a new instance of [Assembler] instantiated with empty module map.
-    ///
-    /// Debug related decorators are added to span blocks when debug mode is enabled.
-    pub fn new(in_debug_mode: bool) -> Self {
+    pub fn new() -> Self {
         Self {
-            stdlib: StdLibrary::default(),
+            module_provider: Box::new(()),
             parsed_modules: BTreeMap::default(),
             kernel_procs: BTreeMap::default(),
             kernel: Kernel::default(),
-            in_debug_mode,
+            in_debug_mode: false,
         }
     }
 
-    /// Returns a new instance of [Assembler] instantiated with the specified kernel.
-    ///
     /// Debug related decorators are added to span blocks when debug mode is enabled.
+    pub fn with_debug_mode(mut self, in_debug_mode: bool) -> Self {
+        self.in_debug_mode = in_debug_mode;
+        self
+    }
+
+    /// Create a new assembler with a given module provider
+    pub fn with_module_provider<P>(mut self, provider: P) -> Self
+    where
+        P: ModuleProvider + 'static,
+    {
+        self.module_provider = Box::new(provider);
+        self
+    }
+
+    /// Returns a new instance of [Assembler] instantiated with the specified kernel.
     ///
     /// # Errors
     /// Returns an error if compiling kernel source results in an error.
-    pub fn with_kernel(kernel_source: &str, in_debug_mode: bool) -> Result<Self, AssemblyError> {
-        let mut assembler = Self::new(in_debug_mode);
-        assembler.set_kernel(kernel_source)?;
-        Ok(assembler)
+    pub fn with_kernel(mut self, kernel_source: &str) -> Result<Self, AssemblyError> {
+        self.set_kernel(kernel_source).map(|_| self)
     }
 
     // PUBLIC ACCESSORS
@@ -192,11 +223,12 @@ impl Assembler {
                     // and attempt to parse it; if the parsing is successful, this will also add
                     // the parsed module to `self.parsed_modules`
                     if !self.parsed_modules.contains_key(module_path) {
-                        let module_source =
-                            self.stdlib.get_module_source(module_path).map_err(|_| {
-                                AssemblyError::missing_import_source(token, module_path)
+                        self.module_provider
+                            .get_source(module_path)
+                            .ok_or_else(|| AssemblyError::missing_import_source(token, module_path))
+                            .and_then(|module_source| {
+                                self.parse_module(module_source, module_path, dep_chain, cb_table)
                             })?;
-                        self.parse_module(module_source, module_path, dep_chain, cb_table)?;
                     }
 
                     // get procedures from the module at the specified path; we are guaranteed to
@@ -267,6 +299,8 @@ impl Assembler {
 
         // insert exported procedures into `self.parsed_procedures`
         // TODO: figure out how to do this using interior mutability
+        // When the module provider maps index to procedures, it might be implemented with a
+        // send/sync friendly approach (maybe std::sync?).
         unsafe {
             let path = path.to_string();
             let mutable_self = &mut *(self as *const _ as *mut Assembler);
@@ -330,9 +364,8 @@ impl Assembler {
 }
 
 impl Default for Assembler {
-    /// Returns a new instance of [Assembler] instantiated with empty module map in non-debug mode.
     fn default() -> Self {
-        Self::new(false)
+        Self::new()
     }
 }
 

assembly/src/tests.rs

@@ -1,3 +1,5 @@
+use crate::ModuleProvider;
+
 // SIMPLE PROGRAMS
 // ================================================================================================
 #[test]
@@ -177,14 +179,38 @@ fn program_with_exported_procedure() {
 
 #[test]
 fn program_with_one_import() {
-    let assembler = super::Assembler::default();
-    let source = "\
-        use.std::math::u256
-        begin \
-            push.4 push.3 \
-            exec.u256::iszero_unsafe \
-        end";
-    let program = assembler.compile(source).unwrap();
+    const MODULE: &str = "dummy::math::u256";
+
+    #[derive(Default)]
+    struct DummyProvider;
+
+    impl ModuleProvider for DummyProvider {
+        fn get_source(&self, path: &str) -> Option<&str> {
+            (path == MODULE).then_some(
+                r#"
+                export.iszero_unsafe
+                    eq.0
+                    repeat.7
+                        swap
+                        eq.0
+                        and
+                    end
+                end"#,
+            )
+        }
+    }
+
+    let assembler = super::Assembler::new().with_module_provider(DummyProvider::default());
+    let source = format!(
+        r#"
+        use.{}
+        begin
+            push.4 push.3
+            exec.u256::iszero_unsafe
+        end"#,
+        MODULE
+    );
+    let program = assembler.compile(&source).unwrap();
     let expected = "\
         begin \
             span \

miden/Cargo.toml

@@ -42,6 +42,7 @@ prover = { package = "miden-prover", path = "../prover", version = "0.3", defaul
 serde = {version = "1.0.117", optional = true }
 serde_derive = {version = "1.0.117", optional = true }
 serde_json = {version = "1.0.59", optional = true }
+stdlib = { package = "miden-stdlib", path = "../stdlib", version = "0.2", default-features = false }
 structopt = { version = "0.3", default-features = false, optional = true }
 verifier = { package = "miden-verifier", path = "../verifier", version = "0.3", default-features = false }
 winter-utils = { package = "winter-utils", version = "0.4", optional = true }

miden/src/cli/data.rs

@@ -1,8 +1,9 @@
-use assembly::Assembler;
+use miden::Assembler;
 use prover::StarkProof;
 use serde_derive::{Deserialize, Serialize};
 use std::path::{Path, PathBuf};
 use std::{fs, io::Write, time::Instant};
+use stdlib::StdLibrary;
 use vm_core::ProgramOutputs;
 use vm_core::{chiplets::hasher::Digest, Program, ProgramInputs};
 use winter_utils::{Deserializable, SliceReader};
@@ -179,7 +180,8 @@ impl ProgramFile {
         let now = Instant::now();
 
         // compile program
-        let program = Assembler::default()
+        let program = Assembler::new()
+            .with_module_provider(StdLibrary::default())
             .compile(&program_file)
             .map_err(|err| format!("Failed to compile program - {}", err))?;
 

miden/src/examples/fibonacci.rs

@@ -1,5 +1,6 @@
 use super::Example;
 use miden::{Assembler, Program, ProgramInputs};
+use stdlib::StdLibrary;
 use vm_core::{Felt, FieldElement, StarkField};
 
 // EXAMPLE BUILDER
@@ -40,8 +41,10 @@ fn generate_fibonacci_program(n: usize) -> Program {
         n - 1
     );
 
-    let assembler = Assembler::default();
-    assembler.compile(&program).unwrap()
+    Assembler::new()
+        .with_module_provider(StdLibrary::default())
+        .compile(&program)
+        .unwrap()
 }
 
 /// Computes the `n`-th term of Fibonacci sequence

miden/src/tools/mod.rs

@@ -1,8 +1,10 @@
 use super::cli::InputFile;
-use assembly::{Assembler, AssemblyError};
+use assembly::AssemblyError;
 use core::fmt;
+use miden::Assembler;
 use processor::{AsmOpInfo, ExecutionError};
 use std::path::PathBuf;
+use stdlib::StdLibrary;
 use structopt::StructOpt;
 use vm_core::{utils::collections::Vec, Operation, ProgramInputs};
 
@@ -139,8 +141,9 @@ impl fmt::Display for ProgramInfo {
 
 /// Returns program analysis of a given program.
 pub fn analyze(program: &str, inputs: ProgramInputs) -> Result<ProgramInfo, ProgramError> {
-    let assembler = Assembler::new(true);
-    let program = assembler
+    let program = Assembler::new()
+        .with_debug_mode(true)
+        .with_module_provider(StdLibrary::default())
         .compile(program)
         .map_err(ProgramError::AssemblyError)?;
     let vm_state_iterator = processor::execute_iter(&program, &inputs);

miden/tests/integration/helpers/mod.rs

@@ -1,6 +1,7 @@
 pub use miden::{ProofOptions, StarkProof};
 use processor::{ExecutionError, ExecutionTrace, Process, VmStateIterator};
 use proptest::prelude::*;
+use stdlib::StdLibrary;
 pub use vm_core::{
     stack::STACK_TOP_SIZE, Felt, FieldElement, Program, ProgramInputs, ProgramOutputs,
 };
@@ -137,14 +138,18 @@ impl Test {
 
     /// Compiles a test's source and returns the resulting Program.
     pub fn compile(&self) -> Program {
-        let assembler = match self.kernel {
-            Some(ref kernel) => assembly::Assembler::with_kernel(kernel, self.in_debug_mode)
+        let assembler = assembly::Assembler::new()
+            .with_debug_mode(self.in_debug_mode)
+            .with_module_provider(StdLibrary::default());
+
+        match self.kernel.as_ref() {
+            Some(kernel) => assembler
+                .with_kernel(kernel)
                 .expect("kernel compilation failed"),
-            None => assembly::Assembler::new(self.in_debug_mode),
-        };
-        assembler
-            .compile(&self.source)
-            .expect("Failed to compile test source.")
+            None => assembler,
+        }
+        .compile(&self.source)
+        .expect("Failed to compile test source.")
     }
 
     /// Compiles the test's source to a Program and executes it with the tests inputs. Returns a

stdlib/Cargo.toml

@@ -17,3 +17,4 @@ doctest = false
 
 [dependencies]
 vm-core = { package = "miden-core", default-features = false, path = "../core", version = "0.3" }
+vm-assembly = { package = "miden-assembly", default-features = false, path = "../assembly", version = "0.3" }

stdlib/src/lib.rs

@@ -1,5 +1,6 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
+use vm_assembly::ModuleProvider;
 use vm_core::{
     errors::LibraryError,
     utils::{collections::BTreeMap, string::ToString},
@@ -27,6 +28,12 @@ pub struct StdLibrary {
     modules: ModuleMap,
 }
 
+impl ModuleProvider for StdLibrary {
+    fn get_source(&self, path: &str) -> Option<&str> {
+        self.get_module_source(path).ok()
+    }
+}
+
 impl Library for StdLibrary {
     /// Returns root namespace of the standard library, which is always "std".
     fn root_ns(&self) -> &str {