Merge pull request #1123 from 0xPolygonMiden/andrew-add-get-module-to…

…-library

Add `get_module_ast(path)` method to the `Library` trait

assembly/src/library/mod.rs

@@ -41,6 +41,11 @@ pub trait Library {
 
     /// Returns the dependency libraries of this library.
     fn dependencies(&self) -> &[LibraryNamespace];
+
+    /// Returns the AST of the module stored at the provided path.
+    fn get_module_ast(&self, path: &LibraryPath) -> Option<&ModuleAst> {
+        self.modules().find(|&module| module.path == *path).map(|module| &module.ast)
+    }
 }
 
 impl<T> Library for &T
@@ -66,6 +71,10 @@ where
     fn dependencies(&self) -> &[LibraryNamespace] {
         T::dependencies(self)
     }
+
+    fn get_module_ast(&self, path: &LibraryPath) -> Option<&ModuleAst> {
+        T::get_module_ast(self, path)
+    }
 }
 
 // MODULE

assembly/src/library/tests.rs

@@ -1,4 +1,4 @@
-use super::{LibraryNamespace, LibraryPath, MaslLibrary, Module, ModuleAst, Version};
+use super::{Library, LibraryNamespace, LibraryPath, MaslLibrary, Module, ModuleAst, Version};
 use vm_core::utils::{Deserializable, Serializable, SliceReader};
 
 #[test]
@@ -56,3 +56,37 @@ fn masl_locations_serialization() {
     bundle.clear_locations();
     assert_eq!(bundle, deserialized);
 }
+
+#[test]
+fn get_module_by_path() {
+    // declare foo module
+    let foo_source = r#"
+        export.foo
+            add
+        end
+    "#;
+    let path = LibraryPath::new("test::foo").unwrap();
+    let ast = ModuleAst::parse(foo_source).unwrap();
+    let foo = Module::new(path, ast);
+
+    let modules = [foo].to_vec();
+
+    // create the bundle with locations
+    let namespace = LibraryNamespace::new("test").unwrap();
+    let version = Version::MIN;
+    let locations = true;
+    let bundle =
+        MaslLibrary::new(namespace, version, locations, modules.clone(), Vec::new()).unwrap();
+
+    // get AST associated with "test::foo" path
+    let foo_ast = bundle.get_module_ast(&LibraryPath::new("test::foo").unwrap()).unwrap();
+    let foo_ast_str = format!("{foo_ast}");
+    let foo_expected = "export.foo.0
+    add
+end
+
+";
+    assert_eq!(foo_ast_str, foo_expected);
+
+    assert!(bundle.get_module_ast(&LibraryPath::new("test::bar").unwrap()).is_none());
+}

stdlib/src/lib.rs

@@ -1,6 +1,9 @@
 #![no_std]
 
-use assembly::{utils::Deserializable, Library, LibraryNamespace, MaslLibrary, Version};
+use assembly::{
+    ast::ModuleAst, utils::Deserializable, Library, LibraryNamespace, LibraryPath, MaslLibrary,
+    Version,
+};
 
 // STANDARD LIBRARY
 // ================================================================================================
@@ -40,6 +43,10 @@ impl Library for StdLibrary {
     fn dependencies(&self) -> &[assembly::LibraryNamespace] {
         self.0.dependencies()
     }
+
+    fn get_module_ast(&self, path: &LibraryPath) -> Option<&ModuleAst> {
+        self.0.get_module_ast(path)
+    }
 }
 
 #[test]