Merge pull request #13 from golemcloud/vigoo/collect-memory-info

Ability to enumerate all linear memories

Author: vigoo

src/analysis/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::component::*;
+use crate::core::Mem;
 use crate::AstCustomization;
 use mappable_rc::Mrc;
 use std::cell::RefCell;
@@ -233,6 +234,24 @@ impl<Ast: AstCustomization + 'static> AnalysisContext<Ast> {
         Ok(result)
     }
 
+    /// Gets all the memories (not just the exported ones) from all modules within the WASM component
+    pub fn get_all_memories(&self) -> AnalysisResult<Vec<Mem>> {
+        let mut result = Vec::new();
+
+        let mut component_stack = vec![self.get_component()];
+        while let Some(component) = component_stack.pop() {
+            for module in component.modules() {
+                for mem in module.mems() {
+                    result.push((*mem).clone());
+                }
+            }
+            for inner_component in component.components() {
+                component_stack.push(inner_component.clone());
+            }
+        }
+        Ok(result)
+    }
+
     pub fn warnings(&self) -> Vec<AnalysisWarning> {
         self.warnings.borrow().clone()
     }

src/core/mod.rs

@@ -371,8 +371,8 @@ impl Section<CoreIndexSpace, CoreSectionType> for FuncType {
 ///
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Limits {
-    min: u64,
-    max: Option<u64>,
+    pub min: u64,
+    pub max: Option<u64>,
 }
 
 /// Memory types classify linear memories and their size range.

tests/mems.rs

@@ -0,0 +1,83 @@
+use golem_wasm_ast::analysis::AnalysisContext;
+use golem_wasm_ast::component::Component;
+use golem_wasm_ast::core::{Limits, Mem, MemType};
+use golem_wasm_ast::IgnoreAllButMetadata;
+
+#[test]
+fn mems_shopping_cart_component() {
+    let source_bytes = include_bytes!("../wasm/shopping-cart.wasm");
+    let component: Component<IgnoreAllButMetadata> = Component::from_bytes(source_bytes).unwrap();
+
+    let state = AnalysisContext::new(component);
+    let mems = state.get_all_memories().unwrap();
+
+    pretty_assertions::assert_eq!(
+        mems,
+        vec![Mem {
+            mem_type: MemType {
+                limits: Limits { min: 17, max: None }
+            }
+        }]
+    )
+}
+
+#[test]
+fn mems_shopping_cart_resource_component() {
+    let source_bytes = include_bytes!("../wasm/shopping-cart-resource.wasm");
+    let component: Component<IgnoreAllButMetadata> = Component::from_bytes(source_bytes).unwrap();
+
+    let state = AnalysisContext::new(component);
+    let mems = state.get_all_memories().unwrap();
+
+    pretty_assertions::assert_eq!(
+        mems,
+        vec![Mem {
+            mem_type: MemType {
+                limits: Limits { min: 17, max: None }
+            }
+        }]
+    )
+}
+
+#[test]
+fn mems_file_service_component() {
+    let source_bytes = include_bytes!("../wasm/file-service.wasm");
+    let component: Component<IgnoreAllButMetadata> = Component::from_bytes(source_bytes).unwrap();
+
+    let state = AnalysisContext::new(component);
+    let mems = state.get_all_memories().unwrap();
+
+    pretty_assertions::assert_eq!(
+        mems,
+        vec![Mem {
+            mem_type: MemType {
+                limits: Limits { min: 17, max: None }
+            }
+        }]
+    )
+}
+
+#[test]
+fn mems_auction_registry_composed_component() {
+    let source_bytes = include_bytes!("../wasm/auction_registry_composed.wasm");
+    let component: Component<IgnoreAllButMetadata> = Component::from_bytes(source_bytes).unwrap();
+
+    let state = AnalysisContext::new(component);
+    let mems = state.get_all_memories().unwrap();
+
+    pretty_assertions::assert_eq!(
+        mems,
+        vec![
+            Mem {
+                mem_type: MemType {
+                    limits: Limits { min: 17, max: None }
+                }
+            },
+            Mem {
+                mem_type: MemType {
+                    limits: Limits { min: 17, max: None }
+                }
+            },
+        ]
+    )
+}