[GC] Read and lower Let instructions (#3485)

For now we don't support non-nullability, and can therefore lower a let into
simpler things. That is,

(let $x = ...
  ;;
)

=>

(block
  $x = ...
  ;;
)
This lets us handle wasm binaries with let, so that we can optimize them
(with the current downside of losing non-nullability).

This is still not trivial to do, sadly, because the indexing of lets is somewhat
odd in the binary. A let modifies the indexes of other things declared before it,
which means that index "0" means different things at different times. And this
is trickier for us because we add more locals as needed for tuples and stacky
code. So this PR makes us track the absolute local indexes from which each
let started to allocate its locals.

The binary testcase was created from this wat using wasp:

(module
  (type $vector (array (field (mut f64))))
  (func $main
    (local $x i32)
    (local $y i32)
    (drop (local.get $x)) ;; 0 is the index appearing in the binary
    ;; first let
    (array.new_with_rtt $vector
      (f64.const 3.14159)
      (i32.const 1)
      (rtt.canon $vector)
    )
    (let (local $v (ref $vector))
      (drop (local.get $v)) ;; 0
      (drop (local.get $x)) ;; 1
      ;; another one, nested
      (array.new_with_rtt $vector
        (f64.const 1234)
        (i32.const 2)
        (rtt.canon $vector)
      )
      (let (local $w (ref $vector))
        (drop (local.get $v)) ;; 1
        (drop (local.get $w)) ;; 0
        (drop (local.get $x)) ;; 2
      )
    )
    ;; another one, later
    (array.new_with_rtt $vector
      (f64.const 2.1828)
      (i32.const 3)
      (rtt.canon $vector)
    )
    (let (local $v (ref $vector))
      (drop (local.get $v)) ;; 0
      (drop (local.get $x)) ;; 1
    )
    (drop (local.get $x)) ;; 0
  )
)

4 files changed, 177 insertions, 12 deletions

diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index ea18fdf67..bda793216 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1004,6 +1004,7 @@ enum ASTNodes {
 
   CallRef = 0x14,
   RetCallRef = 0x15,
+  Let = 0x17,
 
   // gc opcodes
 
@@ -1348,6 +1349,7 @@ public:
   void requireFunctionContext(const char* error);
 
   void readFunctions();
+  void readVars();
 
   std::map<Export*, Index> exportIndices;
   std::vector<Export*> exportOrder;
@@ -1368,6 +1370,29 @@ public:
 
   std::vector<Expression*> expressionStack;
 
+  // Each let block in the binary adds new locals to the bottom of the index
+  // space. That is, all previously-existing indexes are bumped to higher
+  // indexes. getAbsoluteLocalIndex does this computation.
+  // Note that we must track not just the number of locals added in each let,
+  // but also the absolute index from which they were allocated, as binaryen
+  // will add new locals as it goes for things like stacky code and tuples (so
+  // there isn't a simple way to get to the absolute index from a relative one).
+  // Hence each entry here is a pair of the number of items, and the absolute
+  // index they begin at.
+  struct LetData {
+    // How many items are defined in this let.
+    Index num;
+    // The absolute index from which they are allocated from. That is, if num is
+    // 5 and absoluteStart is 10, then we use indexes 10-14.
+    Index absoluteStart;
+  };
+  std::vector<LetData> letStack;
+
+  // Given a relative index of a local (the one used in the wasm binary), get
+  // the absolute one which takes into account lets, and is the one used in
+  // Binaryen IR.
+  Index getAbsoluteLocalIndex(Index index);
+
   // Control flow structure parsing: these have not just the normal binary
   // data for an instruction, but also some bytes later on like "end" or "else".
   // We must be aware of the connection between those things, for debug info.
@@ -1522,6 +1547,8 @@ public:
   void visitRethrow(Rethrow* curr);
   void visitBrOnExn(BrOnExn* curr);
   void visitCallRef(CallRef* curr);
+  // Let is lowered into a block.
+  void visitLet(Block* curr);
 
   void throwError(std::string text);
 
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 9b4794997..bf2e4afef 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -1833,15 +1833,9 @@ void WasmBinaryBuilder::readFunctions() {
     readNextDebugLocation();
 
     BYN_TRACE("reading " << i << std::endl);
-    size_t numLocalTypes = getU32LEB();
-    for (size_t t = 0; t < numLocalTypes; t++) {
-      auto num = getU32LEB();
-      auto type = getConcreteType();
-      while (num > 0) {
-        func->vars.push_back(type);
-        num--;
-      }
-    }
+
+    readVars();
+
     std::swap(func->prologLocation, debugLocation);
     {
       // process the function body
@@ -1854,6 +1848,7 @@ void WasmBinaryBuilder::readFunctions() {
       assert(breakStack.empty());
       assert(expressionStack.empty());
       assert(controlFlowStack.empty());
+      assert(letStack.empty());
       assert(depth == 0);
       func->body = getBlockOrSingleton(func->sig.results);
       assert(depth == 0);
@@ -1863,6 +1858,7 @@ void WasmBinaryBuilder::readFunctions() {
         throwError("stack not empty on function exit");
       }
       assert(controlFlowStack.empty());
+      assert(letStack.empty());
       if (pos != endOfFunction) {
         throwError("binary offset at function exit not at expected location");
       }
@@ -1875,6 +1871,18 @@ void WasmBinaryBuilder::readFunctions() {
   BYN_TRACE(" end function bodies\n");
 }
 
+void WasmBinaryBuilder::readVars() {
+  size_t numLocalTypes = getU32LEB();
+  for (size_t t = 0; t < numLocalTypes; t++) {
+    auto num = getU32LEB();
+    auto type = getConcreteType();
+    while (num > 0) {
+      currFunction->vars.push_back(type);
+      num--;
+    }
+  }
+}
+
 void WasmBinaryBuilder::readExports() {
   BYN_TRACE("== readExports\n");
   size_t num = getU32LEB();
@@ -2886,6 +2894,10 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
       visitCallRef(call);
       break;
     }
+    case BinaryConsts::Let: {
+      visitLet((curr = allocator.alloc<Block>())->cast<Block>());
+      break;
+    }
     case BinaryConsts::AtomicPrefix: {
       code = static_cast<uint8_t>(getU32LEB());
       if (maybeVisitLoad(curr, code, /*isAtomic=*/true)) {
@@ -3056,6 +3068,36 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
   return BinaryConsts::ASTNodes(code);
 }
 
+Index WasmBinaryBuilder::getAbsoluteLocalIndex(Index index) {
+  // Wasm binaries put each let at the bottom of the index space, which may be
+  // good for binary size as often the uses of the let variables are close to
+  // the let itself. However, in Binaryen IR we just have a simple flat index
+  // space of absolute values, which we add to as we parse, and we depend on
+  // later optimizations to reorder locals for size.
+  //
+  // For example, if we have $x, then we add a let with $y, the binary would map
+  // 0 => y, 1 => x, while in Binaryen IR $x always stays at 0, and $y is added
+  // at 1.
+  //
+  // Compute the relative index in the let we were added. We start by looking at
+  // the last let added, and if we belong to it, we are already relative to it.
+  // We will continue relativizing as we go down, til we find our let.
+  int64_t relative = index;
+  for (auto i = int64_t(letStack.size()) - 1; i >= 0; i--) {
+    auto& info = letStack[i];
+    int64_t currNum = info.num;
+    // There were |currNum| let items added in this let. Check if we were one of
+    // them.
+    if (relative < currNum) {
+      return info.absoluteStart + relative;
+    }
+    relative -= currNum;
+  }
+  // We were not a let, but a normal var from the beginning. In that case, after
+  // we subtracted the let items, we have the proper absolute index.
+  return relative;
+}
+
 void WasmBinaryBuilder::startControlFlow(Expression* curr) {
   if (DWARF && currFunction) {
     controlFlowStack.push_back(curr);
@@ -3322,9 +3364,8 @@ void WasmBinaryBuilder::visitCallIndirect(CallIndirect* curr) {
 
 void WasmBinaryBuilder::visitLocalGet(LocalGet* curr) {
   BYN_TRACE("zz node: LocalGet " << pos << std::endl);
-  ;
   requireFunctionContext("local.get");
-  curr->index = getU32LEB();
+  curr->index = getAbsoluteLocalIndex(getU32LEB());
   if (curr->index >= currFunction->getNumLocals()) {
     throwError("bad local.get index");
   }
@@ -3335,7 +3376,7 @@ void WasmBinaryBuilder::visitLocalGet(LocalGet* curr) {
 void WasmBinaryBuilder::visitLocalSet(LocalSet* curr, uint8_t code) {
   BYN_TRACE("zz node: Set|LocalTee\n");
   requireFunctionContext("local.set outside of function");
-  curr->index = getU32LEB();
+  curr->index = getAbsoluteLocalIndex(getU32LEB());
   if (curr->index >= currFunction->getNumLocals()) {
     throwError("bad local.set index");
   }
@@ -5669,6 +5710,31 @@ void WasmBinaryBuilder::visitCallRef(CallRef* curr) {
   curr->finalize(sig.results);
 }
 
+void WasmBinaryBuilder::visitLet(Block* curr) {
+  // A let is lowered into a block that contains the value, and we allocate
+  // locals as needed, which works as we remove non-nullability.
+
+  startControlFlow(curr);
+  // Get the output type.
+  curr->type = getType();
+  // Get the new local types. First, get the absolute index from which we will
+  // start to allocate them.
+  Index absoluteStart = currFunction->vars.size();
+  readVars();
+  Index numNewVars = currFunction->vars.size() - absoluteStart;
+  // Assign the values into locals.
+  Builder builder(wasm);
+  for (Index i = 0; i < numNewVars; i++) {
+    auto* value = popNonVoidExpression();
+    curr->list.push_back(builder.makeLocalSet(absoluteStart + i, value));
+  }
+  // Read the body, with adjusted local indexes.
+  letStack.emplace_back(LetData{numNewVars, absoluteStart});
+  curr->list.push_back(getBlockOrSingleton(curr->type));
+  letStack.pop_back();
+  curr->finalize(curr->type);
+}
+
 bool WasmBinaryBuilder::maybeVisitI31New(Expression*& out, uint32_t code) {
   if (code != BinaryConsts::I31New) {
     return false;
diff --git a/test/let.wasm b/test/let.wasm
new file mode 100644
index 000000000..8a24b7d35
--- /dev/null
+++ b/test/let.wasm
diff --git a/test/let.wasm.fromBinary b/test/let.wasm.fromBinary
new file mode 100644
index 000000000..1de54a96d
--- /dev/null
+++ b/test/let.wasm.fromBinary
@@ -0,0 +1,72 @@
+(module
+ (type $[mut:f64] (array (mut f64)))
+ (type $none_=>_none (func))
+ (func $0
+  (local $0 i32)
+  (local $1 i32)
+  (local $2 (ref null $[mut:f64]))
+  (local $3 (ref null $[mut:f64]))
+  (local $4 (ref null $[mut:f64]))
+  (drop
+   (local.get $0)
+  )
+  (block
+   (local.set $2
+    (array.new_with_rtt $[mut:f64]
+     (rtt.canon $[mut:f64])
+     (i32.const 1)
+     (f64.const 3.14159)
+    )
+   )
+   (block
+    (drop
+     (local.get $2)
+    )
+    (drop
+     (local.get $0)
+    )
+    (block
+     (local.set $3
+      (array.new_with_rtt $[mut:f64]
+       (rtt.canon $[mut:f64])
+       (i32.const 2)
+       (f64.const 1234)
+      )
+     )
+     (block
+      (drop
+       (local.get $2)
+      )
+      (drop
+       (local.get $3)
+      )
+      (drop
+       (local.get $0)
+      )
+     )
+    )
+   )
+  )
+  (block
+   (local.set $4
+    (array.new_with_rtt $[mut:f64]
+     (rtt.canon $[mut:f64])
+     (i32.const 3)
+     (f64.const 2.1828)
+    )
+   )
+   (block
+    (drop
+     (local.get $4)
+    )
+    (drop
+     (local.get $0)
+    )
+   )
+  )
+  (drop
+   (local.get $0)
+  )
+ )
+)
+