Merge pull request #76 from Shopify/ap.allow-multiple-passes-of-trampoline

Allow multiple passes of trampoline

Author: adampetro

trampoline/src/lib.rs

@@ -108,7 +108,7 @@ pub struct TrampolineCodegen {
 
 impl TrampolineCodegen {
     pub fn new(module: Module) -> walrus::Result<Self> {
-        let guest_memory_id = module.get_memory_id()?;
+        let guest_memory_id = Self::guest_memory_id(&module)?;
 
         Ok(Self {
             module,
@@ -136,6 +136,21 @@ impl TrampolineCodegen {
         })
     }
 
+    fn guest_memory_id(module: &Module) -> walrus::Result<MemoryId> {
+        let non_imported_memories = module
+            .memories
+            .iter()
+            .filter(|&memory| memory.import.is_none())
+            .map(|memory| memory.id())
+            .collect::<Vec<_>>();
+
+        if let Some((memory_id, [])) = non_imported_memories.split_first() {
+            Ok(*memory_id)
+        } else {
+            anyhow::bail!("expected exactly one non-imported guest memory")
+        }
+    }
+
     fn emit_memcpy_to_guest(&mut self) -> FunctionId {
         let provider_memory_id = self.provider_memory_id();
 
@@ -470,6 +485,18 @@ mod test {
     use super::{TrampolineCodegen, IMPORTS, PROVIDER_MODULE_NAME};
     use walrus::Module;
 
+    fn trampoline_wat(wat_bytes: &[u8]) -> walrus::Result<String> {
+        let wasm_buf = wat::parse_bytes(wat_bytes)?;
+        trampoline_wasm(&wasm_buf)
+    }
+
+    fn trampoline_wasm(wasm_bytes: &[u8]) -> walrus::Result<String> {
+        let module = Module::from_buffer(wasm_bytes)?;
+        let codegen = TrampolineCodegen::new(module)?;
+        let mut result = codegen.apply()?;
+        wasmprinter::print_bytes(result.emit_wasm())
+    }
+
     #[test]
     fn disassemble_trampoline() {
         // to add a new test case, add a new file to the `test_data` directory and run `cargo test`
@@ -478,10 +505,7 @@ mod test {
         // environment variable as documented at https://docs.rs/insta/latest/insta/index.html#updating-snapshots
         insta::glob!("test_data/*.wat", |path| {
             let input = wat::parse_file(path).unwrap();
-            let module = Module::from_buffer(&input).unwrap();
-            let codegen = TrampolineCodegen::new(module).unwrap();
-            let mut result = codegen.apply().unwrap();
-            let actual = wasmprinter::print_bytes(result.emit_wasm()).unwrap();
+            let actual = trampoline_wat(&input).unwrap();
             insta::assert_snapshot!(actual);
         });
     }
@@ -495,4 +519,30 @@ mod test {
             assert!(module.imports.find(PROVIDER_MODULE_NAME, import).is_some());
         }
     }
+
+    #[test]
+    fn test_consumer_second_pass_is_a_no_op() {
+        let input = include_bytes!("test_data/consumer.wat");
+        let first_wat = trampoline_wat(input).unwrap();
+        let second_wat = trampoline_wat(first_wat.as_bytes()).unwrap();
+
+        assert_eq!(first_wat, second_wat);
+    }
+
+    #[test]
+    fn test_error_for_multiple_guest_memories() {
+        let module = r#"
+        (module
+            (memory 1)
+            (memory 1)
+        )
+        "#;
+        let result = trampoline_wat(module.as_bytes());
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "expected exactly one non-imported guest memory"
+        );
+    }
 }

trampoline/src/snapshots/shopify_function_trampoline__test__disassemble_trampoline@consumer.wat.snap

@@ -11,21 +11,22 @@ input_file: trampoline/src/test_data/consumer.wat
   (type (;4;) (func (param i32 i32) (result i64)))
   (type (;5;) (func (param i32 i32 i32)))
   (type (;6;) (func (param i32 i32 i32) (result i32)))
-  (type (;7;) (func (param i32 i32 i32 i32) (result i32)))
-  (type (;8;) (func (param i32 i64) (result i32)))
-  (type (;9;) (func (param i32 i64 i32) (result i64)))
-  (type (;10;) (func (param i32 i64 i32 i32) (result i64)))
-  (type (;11;) (func (param i32 f64) (result i32)))
+  (type (;7;) (func (param i32 i32 i32 i32)))
+  (type (;8;) (func (param i32 i32 i32 i32) (result i32)))
+  (type (;9;) (func (param i32 i64) (result i32)))
+  (type (;10;) (func (param i32 i64 i32) (result i64)))
+  (type (;11;) (func (param i32 i64 i32 i32) (result i64)))
+  (type (;12;) (func (param i32 f64) (result i32)))
   (import "shopify_function_v1" "_shopify_function_context_new" (func (;0;) (type 0)))
   (import "shopify_function_v1" "_shopify_function_input_get" (func (;1;) (type 2)))
-  (import "shopify_function_v1" "_shopify_function_input_get_interned_obj_prop" (func (;2;) (type 9)))
-  (import "shopify_function_v1" "_shopify_function_input_get_at_index" (func (;3;) (type 9)))
-  (import "shopify_function_v1" "_shopify_function_input_get_obj_key_at_index" (func (;4;) (type 9)))
-  (import "shopify_function_v1" "_shopify_function_input_get_val_len" (func (;5;) (type 8)))
+  (import "shopify_function_v1" "_shopify_function_input_get_interned_obj_prop" (func (;2;) (type 10)))
+  (import "shopify_function_v1" "_shopify_function_input_get_at_index" (func (;3;) (type 10)))
+  (import "shopify_function_v1" "_shopify_function_input_get_obj_key_at_index" (func (;4;) (type 10)))
+  (import "shopify_function_v1" "_shopify_function_input_get_val_len" (func (;5;) (type 9)))
   (import "shopify_function_v1" "_shopify_function_output_new_bool" (func (;6;) (type 3)))
   (import "shopify_function_v1" "_shopify_function_output_new_null" (func (;7;) (type 1)))
   (import "shopify_function_v1" "_shopify_function_output_new_i32" (func (;8;) (type 1)))
-  (import "shopify_function_v1" "_shopify_function_output_new_f64" (func (;9;) (type 11)))
+  (import "shopify_function_v1" "_shopify_function_output_new_f64" (func (;9;) (type 12)))
   (import "shopify_function_v1" "_shopify_function_output_new_object" (func (;10;) (type 3)))
   (import "shopify_function_v1" "_shopify_function_output_finish_object" (func (;11;) (type 1)))
   (import "shopify_function_v1" "_shopify_function_output_new_array" (func (;12;) (type 3)))
@@ -34,8 +35,8 @@ input_file: trampoline/src/test_data/consumer.wat
   (import "shopify_function_v1" "_shopify_function_output_finalize" (func (;15;) (type 1)))
   (import "shopify_function_v1" "_shopify_function_input_get_utf8_str_addr" (func (;16;) (type 3)))
   (import "shopify_function_v1" "memory" (memory (;0;) 1))
-  (import "shopify_function_v1" "_shopify_function_input_get_obj_prop" (func (;17;) (type 10)))
-  (import "shopify_function_v1" "shopify_function_realloc" (func (;18;) (type 7)))
+  (import "shopify_function_v1" "_shopify_function_input_get_obj_prop" (func (;17;) (type 11)))
+  (import "shopify_function_v1" "shopify_function_realloc" (func (;18;) (type 8)))
   (import "shopify_function_v1" "_shopify_function_output_new_utf8_str" (func (;19;) (type 4)))
   (import "shopify_function_v1" "_shopify_function_intern_utf8_str" (func (;20;) (type 4)))
   (memory (;1;) 1)
@@ -70,7 +71,7 @@ input_file: trampoline/src/test_data/consumer.wat
     local.get 2
     call 27
   )
-  (func (;23;) (type 10) (param i32 i64 i32 i32) (result i64)
+  (func (;23;) (type 11) (param i32 i64 i32 i32) (result i64)
     (local i32)
     local.get 3
     call 25
@@ -84,7 +85,7 @@ input_file: trampoline/src/test_data/consumer.wat
     local.get 3
     call 17
   )
-  (func (;24;) (type 7) (param i32 i32 i32 i32) (result i32)
+  (func (;24;) (type 7) (param i32 i32 i32 i32)
     local.get 2
     local.get 0
     local.get 1

trampoline/src/test_data/consumer.wat

@@ -10,7 +10,7 @@
     (import "shopify_function_v1" "shopify_function_input_get_at_index" (func (param i32 i64 i32) (result i64)))
     (import "shopify_function_v1" "shopify_function_input_get_obj_key_at_index" (func (param i32 i64 i32) (result i64)))
     (import "shopify_function_v1" "shopify_function_input_get_val_len" (func (param i32 i64) (result i32)))
-    (import "shopify_function_v1" "shopify_function_input_read_utf8_str" (func (param i32 i32 i32 i32) (result i32)))
+    (import "shopify_function_v1" "shopify_function_input_read_utf8_str" (func (param i32 i32 i32 i32)))
 
     ;; Write.
     (import "shopify_function_v1" "shopify_function_output_new_bool" (func (param i32 i32) (result i32)))