[Wasm GC] Add Wasm GC support to InstrumentMemory (#3976)

This adds calls to imports around every struct load and store, to
note their values, and also to arrays (where it also notes the
index).

This has been very useful in debugging LowerGC (lowering of Wasm
GC to wasm MVP).

3 files changed, 373 insertions, 2 deletions

diff --git a/scripts/wasm2js.js b/scripts/wasm2js.js
index 869734c6f..a72ccfbfb 100644
--- a/scripts/wasm2js.js
+++ b/scripts/wasm2js.js
@@ -163,7 +163,7 @@ var asmLibraryArg = {
     return low;
   },
   load_val_f32: function(loc, value) {
-    console.log('loaload_val_i32d_ptr ' + [loc, value]);
+    console.log('load_val_f32 ' + [loc, value]);
     return value;
   },
   load_val_f64: function(loc, value) {
@@ -184,13 +184,87 @@ var asmLibraryArg = {
     return low;
   },
   store_val_f32: function(loc, value) {
-    console.log('loastore_val_i32d_ptr ' + [loc, value]);
+    console.log('store_val_f32 ' + [loc, value]);
     return value;
   },
   store_val_f64: function(loc, value) {
     console.log('store_val_f64 ' + [loc, value]);
     return value;
   },
+
+  struct_get_val_i32: function(loc, value) {
+    console.log('struct_get_val_i32 ' + [loc, value]);
+    return value;
+  },
+  struct_get_val_i64: function(loc, value) {
+    console.log('struct_get_val_i64 ' + [loc, value]);
+    return value;
+  },
+  struct_get_val_f32: function(loc, value) {
+    console.log('struct_get_val_f32 ' + [loc, value]);
+    return value;
+  },
+  struct_get_val_f64: function(loc, value) {
+    console.log('struct_get_val_f64 ' + [loc, value]);
+    return value;
+  },
+  struct_set_val_i32: function(loc, value) {
+    console.log('struct_set_val_i32 ' + [loc, value]);
+    return value;
+  },
+  struct_set_val_i64: function(loc, value) {
+    console.log('struct_set_val_i64 ' + [loc, value]);
+    return value;
+  },
+  struct_set_val_f32: function(loc, value) {
+    console.log('struct_set_val_f32 ' + [loc, value]);
+    return value;
+  },
+  struct_set_val_f64: function(loc, value) {
+    console.log('struct_set_val_f64 ' + [loc, value]);
+    return value;
+  },
+
+  array_get_val_i32: function(loc, value) {
+    console.log('array_get_val_i32 ' + [loc, value]);
+    return value;
+  },
+  array_get_val_i64: function(loc, value) {
+    console.log('array_get_val_i64 ' + [loc, value]);
+    return value;
+  },
+  array_get_val_f32: function(loc, value) {
+    console.log('array_get_val_f32 ' + [loc, value]);
+    return value;
+  },
+  array_get_val_f64: function(loc, value) {
+    console.log('array_get_val_f64 ' + [loc, value]);
+    return value;
+  },
+  array_set_val_i32: function(loc, value) {
+    console.log('array_set_val_i32 ' + [loc, value]);
+    return value;
+  },
+  array_set_val_i64: function(loc, value) {
+    console.log('array_set_val_i64 ' + [loc, value]);
+    return value;
+  },
+  array_set_val_f32: function(loc, value) {
+    console.log('array_set_val_f32 ' + [loc, value]);
+    return value;
+  },
+  array_set_val_f64: function(loc, value) {
+    console.log('array_set_val_f64 ' + [loc, value]);
+    return value;
+  },
+  array_get_index: function(loc, value) {
+    console.log('array_get_index ' + [loc, value]);
+    return value;
+  },
+  array_set_index: function(loc, value) {
+    console.log('array_set_index ' + [loc, value]);
+    return value;
+  },
 };
 
 var wasmMemory = new WebAssembly.Memory({ initial: 1 });
diff --git a/src/passes/InstrumentMemory.cpp b/src/passes/InstrumentMemory.cpp
index 38a23a658..073c78822 100644
--- a/src/passes/InstrumentMemory.cpp
+++ b/src/passes/InstrumentMemory.cpp
@@ -51,6 +51,11 @@
 //     (i32.const 4)
 //    )
 //   )
+//
+// GC struct and array operations are similarly instrumented, but without their
+// pointers (which are references), and we only log MVP wasm types (i.e., not
+// references or rtts).
+//
 
 #include "asmjs/shared-constants.h"
 #include "shared-constants.h"
@@ -70,6 +75,24 @@ static Name store_val_i32("store_val_i32");
 static Name store_val_i64("store_val_i64");
 static Name store_val_f32("store_val_f32");
 static Name store_val_f64("store_val_f64");
+static Name struct_get_val_i32("struct_get_val_i32");
+static Name struct_get_val_i64("struct_get_val_i64");
+static Name struct_get_val_f32("struct_get_val_f32");
+static Name struct_get_val_f64("struct_get_val_f64");
+static Name struct_set_val_i32("struct_set_val_i32");
+static Name struct_set_val_i64("struct_set_val_i64");
+static Name struct_set_val_f32("struct_set_val_f32");
+static Name struct_set_val_f64("struct_set_val_f64");
+static Name array_get_val_i32("array_get_val_i32");
+static Name array_get_val_i64("array_get_val_i64");
+static Name array_get_val_f32("array_get_val_f32");
+static Name array_get_val_f64("array_get_val_f64");
+static Name array_set_val_i32("array_set_val_i32");
+static Name array_set_val_i64("array_set_val_i64");
+static Name array_set_val_f32("array_set_val_f32");
+static Name array_set_val_f64("array_set_val_f64");
+static Name array_get_index("array_get_index");
+static Name array_set_index("array_set_index");
 
 // TODO: Add support for atomicRMW/cmpxchg
 
@@ -138,20 +161,132 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
       target, {builder.makeConst(int32_t(id)), curr->value}, curr->value->type);
   }
 
+  void visitStructGet(StructGet* curr) {
+    Builder builder(*getModule());
+    Name target;
+    if (curr->type == Type::i32) {
+      target = struct_get_val_i32;
+    } else if (curr->type == Type::i64) {
+      target = struct_get_val_i64;
+    } else if (curr->type == Type::f32) {
+      target = struct_get_val_f32;
+    } else if (curr->type == Type::f64) {
+      target = struct_get_val_f64;
+    } else {
+      return; // TODO: other types, unreachable, etc.
+    }
+    replaceCurrent(builder.makeCall(
+      target, {builder.makeConst(int32_t(id++)), curr}, curr->type));
+  }
+
+  void visitStructSet(StructSet* curr) {
+    Builder builder(*getModule());
+    Name target;
+    if (curr->value->type == Type::i32) {
+      target = struct_set_val_i32;
+    } else if (curr->value->type == Type::i64) {
+      target = struct_set_val_i64;
+    } else if (curr->value->type == Type::f32) {
+      target = struct_set_val_f32;
+    } else if (curr->value->type == Type::f64) {
+      target = struct_set_val_f64;
+    } else {
+      return; // TODO: other types, unreachable, etc.
+    }
+    curr->value =
+      builder.makeCall(target,
+                       {builder.makeConst(int32_t(id++)), curr->value},
+                       curr->value->type);
+  }
+
+  void visitArrayGet(ArrayGet* curr) {
+    Builder builder(*getModule());
+    curr->index =
+      builder.makeCall(array_get_index,
+                       {builder.makeConst(int32_t(id++)), curr->index},
+                       Type::i32);
+    Name target;
+    if (curr->type == Type::i32) {
+      target = array_get_val_i32;
+    } else if (curr->type == Type::i64) {
+      target = array_get_val_i64;
+    } else if (curr->type == Type::f32) {
+      target = array_get_val_f32;
+    } else if (curr->type == Type::f64) {
+      target = array_get_val_f64;
+    } else {
+      return; // TODO: other types, unreachable, etc.
+    }
+    replaceCurrent(builder.makeCall(
+      target, {builder.makeConst(int32_t(id++)), curr}, curr->type));
+  }
+
+  void visitArraySet(ArraySet* curr) {
+    Builder builder(*getModule());
+    curr->index =
+      builder.makeCall(array_set_index,
+                       {builder.makeConst(int32_t(id++)), curr->index},
+                       Type::i32);
+    Name target;
+    if (curr->value->type == Type::i32) {
+      target = array_set_val_i32;
+    } else if (curr->value->type == Type::i64) {
+      target = array_set_val_i64;
+    } else if (curr->value->type == Type::f32) {
+      target = array_set_val_f32;
+    } else if (curr->value->type == Type::f64) {
+      target = array_set_val_f64;
+    } else {
+      return; // TODO: other types, unreachable, etc.
+    }
+    curr->value =
+      builder.makeCall(target,
+                       {builder.makeConst(int32_t(id++)), curr->value},
+                       curr->value->type);
+  }
+
   void visitModule(Module* curr) {
     auto indexType = curr->memory.indexType;
+
+    // Load.
     addImport(
       curr, load_ptr, {Type::i32, Type::i32, indexType, indexType}, indexType);
     addImport(curr, load_val_i32, {Type::i32, Type::i32}, Type::i32);
     addImport(curr, load_val_i64, {Type::i32, Type::i64}, Type::i64);
     addImport(curr, load_val_f32, {Type::i32, Type::f32}, Type::f32);
     addImport(curr, load_val_f64, {Type::i32, Type::f64}, Type::f64);
+
+    // Store.
     addImport(
       curr, store_ptr, {Type::i32, Type::i32, indexType, indexType}, indexType);
     addImport(curr, store_val_i32, {Type::i32, Type::i32}, Type::i32);
     addImport(curr, store_val_i64, {Type::i32, Type::i64}, Type::i64);
     addImport(curr, store_val_f32, {Type::i32, Type::f32}, Type::f32);
     addImport(curr, store_val_f64, {Type::i32, Type::f64}, Type::f64);
+
+    if (curr->features.hasGC()) {
+      // Struct get/set.
+      addImport(curr, struct_get_val_i32, {Type::i32, Type::i32}, Type::i32);
+      addImport(curr, struct_get_val_i64, {Type::i32, Type::i64}, Type::i64);
+      addImport(curr, struct_get_val_f32, {Type::i32, Type::f32}, Type::f32);
+      addImport(curr, struct_get_val_f64, {Type::i32, Type::f64}, Type::f64);
+      addImport(curr, struct_set_val_i32, {Type::i32, Type::i32}, Type::i32);
+      addImport(curr, struct_set_val_i64, {Type::i32, Type::i64}, Type::i64);
+      addImport(curr, struct_set_val_f32, {Type::i32, Type::f32}, Type::f32);
+      addImport(curr, struct_set_val_f64, {Type::i32, Type::f64}, Type::f64);
+
+      // Array get/set.
+      addImport(curr, array_get_val_i32, {Type::i32, Type::i32}, Type::i32);
+      addImport(curr, array_get_val_i64, {Type::i32, Type::i64}, Type::i64);
+      addImport(curr, array_get_val_f32, {Type::i32, Type::f32}, Type::f32);
+      addImport(curr, array_get_val_f64, {Type::i32, Type::f64}, Type::f64);
+      addImport(curr, array_set_val_i32, {Type::i32, Type::i32}, Type::i32);
+      addImport(curr, array_set_val_i64, {Type::i32, Type::i64}, Type::i64);
+      addImport(curr, array_set_val_f32, {Type::i32, Type::f32}, Type::f32);
+      addImport(curr, array_set_val_f64, {Type::i32, Type::f64}, Type::f64);
+      addImport(curr, array_get_index, {Type::i32, Type::i32}, Type::i32);
+      addImport(curr, array_set_index, {Type::i32, Type::i32}, Type::i32);
+    }
   }
 
 private:
diff --git a/test/lit/passes/instrument-memory-gc.wast b/test/lit/passes/instrument-memory-gc.wast
new file mode 100644
index 000000000..6e4192d16
--- /dev/null
+++ b/test/lit/passes/instrument-memory-gc.wast
@@ -0,0 +1,162 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: foreach %s %t wasm-opt --instrument-memory -all -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
+
+  ;; CHECK:      (type $i32_i64_=>_i64 (func (param i32 i64) (result i64)))
+
+  ;; CHECK:      (type $i32_f32_=>_f32 (func (param i32 f32) (result f32)))
+
+  ;; CHECK:      (type $i32_f64_=>_f64 (func (param i32 f64) (result f64)))
+
+  ;; CHECK:      (type $struct (struct (field (mut i32)) (field f32) (field $named f64)))
+  (type $struct (struct
+    (field (mut i32))
+    (field f32)
+    (field $named f64)
+  ))
+  ;; CHECK:      (type $i32_i32_i32_i32_=>_i32 (func (param i32 i32 i32 i32) (result i32)))
+
+  ;; CHECK:      (type $ref|$struct|_=>_none (func (param (ref $struct))))
+
+  ;; CHECK:      (type $ref|$array|_=>_none (func (param (ref $array))))
+
+  ;; CHECK:      (type $array (array (mut f64)))
+  (type $array (array (mut f64)))
+
+  ;; CHECK:      (import "env" "load_ptr" (func $load_ptr (param i32 i32 i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "load_val_i32" (func $load_val_i32 (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "load_val_i64" (func $load_val_i64 (param i32 i64) (result i64)))
+
+  ;; CHECK:      (import "env" "load_val_f32" (func $load_val_f32 (param i32 f32) (result f32)))
+
+  ;; CHECK:      (import "env" "load_val_f64" (func $load_val_f64 (param i32 f64) (result f64)))
+
+  ;; CHECK:      (import "env" "store_ptr" (func $store_ptr (param i32 i32 i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "store_val_i32" (func $store_val_i32 (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "store_val_i64" (func $store_val_i64 (param i32 i64) (result i64)))
+
+  ;; CHECK:      (import "env" "store_val_f32" (func $store_val_f32 (param i32 f32) (result f32)))
+
+  ;; CHECK:      (import "env" "store_val_f64" (func $store_val_f64 (param i32 f64) (result f64)))
+
+  ;; CHECK:      (import "env" "struct_get_val_i32" (func $struct_get_val_i32 (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "struct_get_val_i64" (func $struct_get_val_i64 (param i32 i64) (result i64)))
+
+  ;; CHECK:      (import "env" "struct_get_val_f32" (func $struct_get_val_f32 (param i32 f32) (result f32)))
+
+  ;; CHECK:      (import "env" "struct_get_val_f64" (func $struct_get_val_f64 (param i32 f64) (result f64)))
+
+  ;; CHECK:      (import "env" "struct_set_val_i32" (func $struct_set_val_i32 (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "struct_set_val_i64" (func $struct_set_val_i64 (param i32 i64) (result i64)))
+
+  ;; CHECK:      (import "env" "struct_set_val_f32" (func $struct_set_val_f32 (param i32 f32) (result f32)))
+
+  ;; CHECK:      (import "env" "struct_set_val_f64" (func $struct_set_val_f64 (param i32 f64) (result f64)))
+
+  ;; CHECK:      (import "env" "array_get_val_i32" (func $array_get_val_i32 (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "array_get_val_i64" (func $array_get_val_i64 (param i32 i64) (result i64)))
+
+  ;; CHECK:      (import "env" "array_get_val_f32" (func $array_get_val_f32 (param i32 f32) (result f32)))
+
+  ;; CHECK:      (import "env" "array_get_val_f64" (func $array_get_val_f64 (param i32 f64) (result f64)))
+
+  ;; CHECK:      (import "env" "array_set_val_i32" (func $array_set_val_i32 (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "array_set_val_i64" (func $array_set_val_i64 (param i32 i64) (result i64)))
+
+  ;; CHECK:      (import "env" "array_set_val_f32" (func $array_set_val_f32 (param i32 f32) (result f32)))
+
+  ;; CHECK:      (import "env" "array_set_val_f64" (func $array_set_val_f64 (param i32 f64) (result f64)))
+
+  ;; CHECK:      (import "env" "array_get_index" (func $array_get_index (param i32 i32) (result i32)))
+
+  ;; CHECK:      (import "env" "array_set_index" (func $array_set_index (param i32 i32) (result i32)))
+
+  ;; CHECK:      (func $structs (param $x (ref $struct))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $struct_get_val_i32
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:    (struct.get $struct 0
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $struct_get_val_f32
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:    (struct.get $struct 1
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $struct_get_val_f64
+  ;; CHECK-NEXT:    (i32.const 2)
+  ;; CHECK-NEXT:    (struct.get $struct $named
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.set $struct 0
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (call $struct_set_val_i32
+  ;; CHECK-NEXT:    (i32.const 3)
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $structs (param $x (ref $struct))
+    (drop
+      (struct.get $struct 0 (local.get $x))
+    )
+    (drop
+      (struct.get $struct 1 (local.get $x))
+    )
+    (drop
+      (struct.get $struct 2 (local.get $x))
+    )
+    (struct.set $struct 0 (local.get $x) (i32.const 42))
+  )
+
+  ;; CHECK:      (func $arrays (param $x (ref $array))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $array_get_val_f64
+  ;; CHECK-NEXT:    (i32.const 5)
+  ;; CHECK-NEXT:    (array.get $array
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (call $array_get_index
+  ;; CHECK-NEXT:      (i32.const 4)
+  ;; CHECK-NEXT:      (i32.const 10)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (array.set $array
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (call $array_set_index
+  ;; CHECK-NEXT:    (i32.const 6)
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $array_set_val_f64
+  ;; CHECK-NEXT:    (i32.const 7)
+  ;; CHECK-NEXT:    (f64.const 3.14159)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $arrays (param $x (ref $array))
+    (drop
+      (array.get $array (local.get $x) (i32.const 10))
+    )
+    (array.set $array (local.get $x) (i32.const 42) (f64.const 3.14159))
+  )
+)
+