From 0768860471ab6cdb4526d999ef3e742b5421e247 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Tue, 9 Feb 2021 08:06:42 -0800
Subject: [PATCH] Fix the codegen of the `TableGrow` intrinsic

It's not safe for us to arbitrarily modify the instruction stream since
wasm isn't guaranteed to be an AST! Instead we resort to a few extra
instructions with locals to achieve what we want here.

Closes #2446
---
 crates/externref-xform/src/lib.rs             | 27 +++++++++---
 .../tests/table-grow-intrinsic.wat            | 13 +++---
 .../externref-xform/tests/tee-before-grow.wat | 43 +++++++++++++++++++
 3 files changed, 72 insertions(+), 11 deletions(-)
 create mode 100644 crates/externref-xform/tests/tee-before-grow.wat

diff --git a/crates/externref-xform/src/lib.rs b/crates/externref-xform/src/lib.rs
index 881bba713eb..563cabfdd18 100644
--- a/crates/externref-xform/src/lib.rs
+++ b/crates/externref-xform/src/lib.rs
@@ -681,11 +681,13 @@ impl Transform<'_> {
                 continue;
             }
             let entry = func.entry_block();
+            let scratch_i32 = module.locals.add(ValType::I32);
             dfs_pre_order_mut(
                 &mut Rewrite {
                     clone_ref: self.clone_ref()?,
                     heap_dealloc: self.heap_dealloc()?,
                     xform: self,
+                    scratch_i32,
                 },
                 func,
                 entry,
@@ -698,6 +700,7 @@ impl Transform<'_> {
             xform: &'a Transform<'b>,
             clone_ref: FunctionId,
             heap_dealloc: FunctionId,
+            scratch_i32: LocalId,
         }
 
         impl VisitorMut for Rewrite<'_, '_> {
@@ -723,16 +726,28 @@ impl Transform<'_> {
                     let ty = ValType::Externref;
                     match intrinsic {
                         Intrinsic::TableGrow => {
-                            // Switch this to a `table.grow` instruction...
+                            // Change something that looks like:
+                            //
+                            //      call $table_grow
+                            //
+                            // into:
+                            //
+                            //      local.set $scratch
+                            //      ref.null extern
+                            //      local.get $scratch
+                            //      table.grow $table
+                            //
+                            // Note that things happen backwards here due to the
+                            // order of insertion.
                             seq.instrs[i].0 = TableGrow {
                                 table: self.xform.table,
                             }
                             .into();
-                            // ... and then insert a `ref.null` before the
-                            // preceding instruction as the value to grow the
-                            // table with.
-                            seq.instrs
-                                .insert(i - 1, (RefNull { ty }.into(), InstrLocId::default()));
+                            let loc = seq.instrs[i].1;
+                            let local = self.scratch_i32;
+                            seq.instrs.insert(i, (LocalGet { local }.into(), loc));
+                            seq.instrs.insert(i, (RefNull { ty }.into(), loc));
+                            seq.instrs.insert(i, (LocalSet { local }.into(), loc));
                         }
                         Intrinsic::TableSetNull => {
                             // Switch this to a `table.set` instruction...
diff --git a/crates/externref-xform/tests/table-grow-intrinsic.wat b/crates/externref-xform/tests/table-grow-intrinsic.wat
index 5886a0b5178..70670c045a6 100644
--- a/crates/externref-xform/tests/table-grow-intrinsic.wat
+++ b/crates/externref-xform/tests/table-grow-intrinsic.wat
@@ -17,6 +17,14 @@
   (type (;0;) (func (result i32)))
   (type (;1;) (func (param i32)))
   (type (;2;) (func (param externref)))
+  (func $foo (type 1) (param i32)
+    (local i32)
+    i32.const 0
+    local.set 1
+    ref.null extern
+    local.get 1
+    table.grow 0
+    drop)
   (func $foo_externref_shim (@name "foo externref shim") (type 2) (param externref)
     (local i32)
     call $alloc
@@ -25,11 +33,6 @@
     table.set 0
     local.get 1
     call $foo)
-  (func $foo (type 1) (param i32)
-    ref.null extern
-    i32.const 0
-    table.grow 0
-    drop)
   (func $alloc (type 0) (result i32)
     i32.const 0)
   (table (;0;) 32 externref)
diff --git a/crates/externref-xform/tests/tee-before-grow.wat b/crates/externref-xform/tests/tee-before-grow.wat
new file mode 100644
index 00000000000..4331c398f39
--- /dev/null
+++ b/crates/externref-xform/tests/tee-before-grow.wat
@@ -0,0 +1,43 @@
+;; @xform export "foo" (externref_owned)
+
+(module
+  (import "__wbindgen_externref_xform__" "__wbindgen_externref_table_grow"
+    (func $grow (param i32) (result i32)))
+  (func $foo (export "foo") (param i32)
+    (local i32)
+    i32.const 0
+    local.tee 0
+    call $grow
+    drop)
+  (func $alloc (export "__externref_table_alloc") (result i32)
+    i32.const 0)
+  (func $dealloc (export "__externref_table_dealloc") (param i32))
+)
+
+(; CHECK-ALL:
+(module
+  (type (;0;) (func (result i32)))
+  (type (;1;) (func (param i32)))
+  (type (;2;) (func (param externref)))
+  (func $foo (type 1) (param i32)
+    (local i32)
+    i32.const 0
+    local.tee 0
+    local.set 1
+    ref.null extern
+    local.get 1
+    table.grow 0
+    drop)
+  (func $foo_externref_shim (@name "foo externref shim") (type 2) (param externref)
+    (local i32)
+    call $alloc
+    local.tee 1
+    local.get 0
+    table.set 0
+    local.get 1
+    call $foo)
+  (func $alloc (type 0) (result i32)
+    i32.const 0)
+  (table (;0;) 32 externref)
+  (export "foo" (func $foo_externref_shim)))
+;)