StackIR local2stack: Make sure we do not break non-nullable validation (#5919)

local2stack removes a pair of

local.set 0
local.get 0

when that set is not used anywhere else: whatever value is put into the local,
we can just leave it on the stack to replace the get. However, we only handled
actual uses of the set which we checked using LocalGraph. There may be code
that does not actually use the set local, but needs that set purely for validation
reasons:

local.set 0
local.get 0
block
  local.set 0
end
local.get

That last get reads the value set in the block, so the first set is not used by it.
But for validation purposes, the inner set stops helping at the block end, so
we do need that initial set.

To fix this, check for gets that need our set to validate before removing any.

Fixes #5917

2 files changed, 954 insertions, 1 deletions

diff --git a/src/passes/StackIR.cpp b/src/passes/StackIR.cpp
index bc91d0703..3d5426a17 100644
--- a/src/passes/StackIR.cpp
+++ b/src/passes/StackIR.cpp
@@ -216,7 +216,10 @@ private:
                 auto& sets = localGraph.getSetses[get];
                 if (sets.size() == 1 && *sets.begin() == set) {
                   auto& setInfluences = localGraph.setInfluences[set];
-                  if (setInfluences.size() == 1) {
+                  // If this has the proper value of 1, also do the potentially-
+                  // expensive check of whether we can remove this pair at all.
+                  if (setInfluences.size() == 1 &&
+                      canRemoveSetGetPair(index, i)) {
                     assert(*setInfluences.begin() == get);
                     // Do it! The set and the get can go away, the proper
                     // value is on the stack.
@@ -354,6 +357,120 @@ private:
     // Otherwise, for basic instructions, just count the expression children.
     return ChildIterator(inst->origin).children.size();
   }
+
+  // Given a pair of a local.set and local.get, see if we can remove them
+  // without breaking validation. Specifically, we must keep sets of non-
+  // nullable locals that dominate a get until the end of the block, such as
+  // here:
+  //
+  //  local.set 0    ;; Can we remove
+  //  local.get 0    ;; this pair?
+  //  if
+  //    local.set 0
+  //  else
+  //    local.set 0
+  //  end
+  //  local.get 0    ;; This get poses a problem.
+  //
+  // Logically the 2nd&3rd sets ensure a value is applied to the local before we
+  // read it, but the validation rules only track each set until the end of its
+  // scope, so the 1st set (before the if, in the pair) is necessary.
+  //
+  // The logic below is related to LocalStructuralDominance, but sharing code
+  // with it is difficult as this uses StackIR and not BinaryenIR, and it checks
+  // a particular set/get pair.
+  //
+  // We are given the indexes of the set and get instructions in |insts|.
+  bool canRemoveSetGetPair(Index setIndex, Index getIndex) {
+    // The set must be before the get.
+    assert(setIndex < getIndex);
+
+    auto* set = insts[setIndex]->origin->cast<LocalSet>();
+    if (func->isParam(set->index) ||
+        !func->getLocalType(set->index).isNonNullable()) {
+      // This local cannot pose a problem for validation (params are always
+      // initialized, and nullable locals may be uninitialized).
+      return true;
+    }
+
+    // Track the depth (in block/if/loop/etc. scopes) relative to our starting
+    // point. Anything less deep than that is not interesting, as we can only
+    // help things at our depth or deeper to validate.
+    Index currDepth = 0;
+
+    // Look for a different get than the one in getIndex (since that one is
+    // being removed) which would stop validating without the set. While doing
+    // so, note other sets that ensure validation even if our set is removed. We
+    // track those in this stack of booleans, one for each scope, which is true
+    // if another sets covers us and ours is not needed.
+    //
+    // We begin in the current scope and with no other set covering us.
+    std::vector<bool> coverStack = {false};
+
+    // Track the total number of covers as well, for quick checking below.
+    Index covers = 0;
+
+    // TODO: We could look before us as well, but then we might end up scanning
+    //       much of the function every time.
+    for (Index i = setIndex + 1; i < insts.size(); i++) {
+      auto* inst = insts[i];
+      if (!inst) {
+        continue;
+      }
+      if (isControlFlowBegin(inst)) {
+        // A new scope begins.
+        currDepth++;
+        coverStack.push_back(false);
+      } else if (isControlFlowEnd(inst)) {
+        if (currDepth == 0) {
+          // Less deep than the start, so we found no problem.
+          return true;
+        }
+        currDepth--;
+
+        if (coverStack.back()) {
+          // A cover existed in the scope which ended.
+          covers--;
+        }
+        coverStack.pop_back();
+      } else if (isControlFlowBarrier(inst)) {
+        // A barrier, like the else in an if-else, not only ends a scope but
+        // opens a new one.
+        if (currDepth == 0) {
+          // Another scope with the same depth begins, but ours ended, so stop.
+          return true;
+        }
+
+        if (coverStack.back()) {
+          // A cover existed in the scope which ended.
+          covers--;
+        }
+        coverStack.back() = false;
+      } else if (auto* otherSet = inst->origin->dynCast<LocalSet>()) {
+        // We are covered in this scope henceforth.
+        if (otherSet->index == set->index) {
+          if (!coverStack.back()) {
+            covers++;
+            if (currDepth == 0) {
+              // We have a cover at depth 0, so everything from here on out
+              // will be covered.
+              return true;
+            }
+            coverStack.back() = true;
+          }
+        }
+      } else if (auto* otherGet = inst->origin->dynCast<LocalGet>()) {
+        if (otherGet->index == set->index && i != getIndex && !covers) {
+          // We found a get that might be a problem: it uses the same index, but
+          // is not the get we were told about, and no other set covers us.
+          return false;
+        }
+      }
+    }
+
+    // No problem.
+    return true;
+  }
 };
 
 struct OptimizeStackIR : public WalkerPass<PostWalker<OptimizeStackIR>> {
diff --git a/test/lit/passes/stack-ir-non-nullable.wast b/test/lit/passes/stack-ir-non-nullable.wast
new file mode 100644
index 000000000..ac34b5107
--- /dev/null
+++ b/test/lit/passes/stack-ir-non-nullable.wast
@@ -0,0 +1,836 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s --generate-stack-ir --optimize-stack-ir --shrink-level=1 \
+;; RUN:   -all --print-stack-ir | filecheck %s
+
+;; Shrink level is set to 1 to enable local2stack in StackIR opts.
+
+(module
+ ;; CHECK:      (func $if (type $1) (param $param (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT: )
+ (func $if (param $param (ref eq)) (result (ref eq))
+  (local $temp (ref eq))
+  ;; Copy the param into $temp. $temp is then set in both arms of the if, so
+  ;; it is set before the get at the end of the function, but we still need to
+  ;; keep this set for validation purposes. Specifically, there is a set of
+  ;; $temp followed by a get of it in the if condition, which local2stack could
+  ;; remove in principle, if not for that final get at the function end.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.get $temp)
+ )
+
+ ;; CHECK:      (func $if-no-last-get (type $1) (param $param (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT: )
+ (func $if-no-last-get (param $param (ref eq)) (result (ref eq))
+  ;; As the original, but now there is no final get, so we can remove the set-
+  ;; get pair of $temp before the if.
+  (local $temp (ref eq))
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.get $param) ;; this changed from $temp to $param
+ )
+
+ ;; CHECK:      (func $if-extra-set (type $1) (param $param (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT: )
+ (func $if-extra-set (param $param (ref eq)) (result (ref eq))
+  ;; As the original, but now there is an extra set before the final get, so
+  ;; we can optimize - the extra set ensures validation.
+  (local $temp (ref eq))
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.set $temp    ;; This set is new.
+   (local.get $param)
+  )
+  (local.get $temp)
+ )
+
+ ;; CHECK:      (func $if-wrong-extra-set (type $1) (param $param (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $param
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT: )
+ (func $if-wrong-extra-set (param $param (ref eq)) (result (ref eq))
+  ;; As the last testcase, but the extra set's index is wrong, so we cannot
+  ;; optimize.
+  (local $temp (ref eq))
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.set $param    ;; This set now writes to $param.
+   (local.get $param)
+  )
+  (local.get $temp)
+ )
+
+ ;; CHECK:      (func $if-wrong-extra-get (type $1) (param $param (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT: )
+ (func $if-wrong-extra-get (param $param (ref eq)) (result (ref eq))
+  ;; As the last testcase, but now it is the get that has the wrong index to
+  ;; stop us, so we can optimize.
+  (local $temp (ref eq))
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.set $temp
+   (local.get $param)
+  )
+  (local.get $param) ;; This get does not affect optimizing the pair before the
+                     ;; if, because it is of another local.
+ )
+
+ ;; CHECK:      (func $if-param (type $2) (param $param (ref eq)) (param $temp (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT: )
+ (func $if-param (param $param (ref eq)) (param $temp (ref eq)) (result (ref eq))
+  ;; As the original testcase, but now $temp is a param. Validation is no
+  ;; longer an issue, so we can optimize away the pair.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.get $temp)
+ )
+
+ ;; CHECK:      (func $if-nullable (type $3) (param $param (ref eq)) (result eqref)
+ ;; CHECK-NEXT:  (local $temp eqref)
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   ref.i31
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT: )
+ (func $if-nullable (param $param (ref eq)) (result (ref null eq))
+  (local $temp (ref null eq)) ;; this changed
+  ;; As the original testcase, but now $temp is a nullable. Validation is no
+  ;; longer an issue, so we can optimize away the pair.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 1)
+    )
+   )
+   (local.set $temp
+    (i31.new
+     (i32.const 2)
+    )
+   )
+  )
+  (local.get $temp)
+ )
+
+ ;; CHECK:      (func $if-non-ref (type $4) (param $param i32) (result i32)
+ ;; CHECK-NEXT:  (local $temp i32)
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  i32.eqz
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   i32.const 1
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   i32.const 2
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT: )
+ (func $if-non-ref (param $param i32) (result i32)
+  (local $temp i32)
+  ;; As the original testcase, but now $temp is not a ref. Validation is no
+  ;; longer an issue, so we can optimize away the pair.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (i32.eqz
+    (local.get $temp)
+    (i32.const 0)
+   )
+   (local.set $temp
+    (i32.const 1)
+   )
+   (local.set $temp
+    (i32.const 2)
+   )
+  )
+  (local.get $temp)
+ )
+
+ ;; CHECK:      (func $nesting (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; The if arms contain optimization opportunities, even though there are 2
+  ;; gets in each one, because this top set helps them all validate. Atm we do
+  ;; not look backwards, however, so we fail to optimize here.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (i32.const 0)
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+  )
+  (drop
+   (local.get $temp)
+  )
+ )
+
+ ;; CHECK:      (func $nesting-left (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT: )
+ (func $nesting-left (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; As $nesting, but now the left arm has one get, and can be optimized.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (i32.const 0)
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    ;; A get was removed here.
+   )
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+  )
+ )
+
+ ;; CHECK:      (func $nesting-right (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:   local.get $temp
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT: )
+ (func $nesting-right (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; As above, but now we can optimize the right arm.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (i32.const 0)
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    ;; A get was removed here.
+   )
+  )
+ )
+
+ ;; CHECK:      (func $nesting-both (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT: )
+ (func $nesting-both (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; As above, but now we can optimize both arms.
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (i32.const 0)
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    ;; A get was removed here.
+   )
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+    ;; A get was removed here.
+   )
+  )
+ )
+
+ ;; CHECK:      (func $nesting-both-less (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT: )
+ (func $nesting-both-less (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; As above, but without the initial set-get at the top. We can still
+  ;; optimize both arms.
+  (if
+   (i32.const 0)
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+  )
+ )
+
+ ;; CHECK:      (func $nesting-both-after (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   drop
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting-both-after (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; As above, but now there is a set-get at the end of the function. The get
+  ;; there should not confuse us; we can still optimize both arms, and after
+  ;; them as well.
+  (if
+   (i32.const 0)
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+   (block
+    (local.set $temp
+     (local.get $param)
+    )
+    (drop
+     (local.get $temp)
+    )
+   )
+  )
+  (local.set $temp
+   (local.get $param)
+  )
+  (drop
+   (local.get $temp)
+  )
+ )
+
+ ;; CHECK:      (func $nesting-irrelevant (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting-irrelevant (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; The block in the middle here adds a scope, but it does not prevent us from
+  ;; optimizing.
+  (local.set $temp
+   (local.get $param)
+  )
+  (block $block
+  )
+  (drop
+   (local.get $temp)
+  )
+ )
+
+ ;; CHECK:      (func $nesting-relevant (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting-relevant (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; As above, but now there is a get in that scope, which is a problem.
+  (local.set $temp
+   (local.get $param)
+  )
+  (block $block
+   (drop
+    (local.get $temp)
+   )
+  )
+  (drop
+   (local.get $temp)
+  )
+ )
+
+ ;; CHECK:      (func $nesting-after (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting-after (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; A set-get pair with another after it in a block. We can optimize both.
+  (local.set $temp
+   (local.get $param)
+  )
+  (drop
+   (local.get $temp)
+  )
+  (block $block
+   (local.set $temp
+    (local.get $param)
+   )
+   (drop
+    (local.get $temp)
+   )
+  )
+ )
+
+ ;; CHECK:      (func $nesting-reverse (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting-reverse (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; The reverse of the last case, now the block is first. We can optimize
+  ;; both pairs.
+  (block $block
+   (local.set $temp
+    (local.get $param)
+   )
+   (drop
+    (local.get $temp)
+   )
+  )
+  (local.set $temp
+   (local.get $param)
+  )
+  (drop
+   (local.get $temp)
+  )
+ )
+
+ ;; CHECK:      (func $nesting-covered-but-ended (type $0) (param $param (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  drop
+ ;; CHECK-NEXT: )
+ (func $nesting-covered-but-ended (param $param (ref eq))
+  (local $temp (ref eq))
+  ;; We cannot optimize this first pair, because while we see another set after
+  ;; us, its scope ends, so our set must help the very final get validate.
+  (local.set $temp
+   (local.get $param)
+  )
+  (drop
+   (local.get $temp)
+  )
+  (block $block
+   ;; This pair we can almost optimize, but the get after the block reads from
+   ;; it, so we don't.
+   (local.set $temp
+    (local.get $param)
+   )
+   (drop
+    (local.get $temp)
+   )
+  )
+  (drop
+   (local.get $temp)
+  )
+ )
+
+ ;; CHECK:      (func $two-covers (type $1) (param $param (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (local $temp (ref eq))
+ ;; CHECK-NEXT:  local.get $param
+ ;; CHECK-NEXT:  local.set $temp
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT:  i32.const 0
+ ;; CHECK-NEXT:  ref.i31
+ ;; CHECK-NEXT:  ref.eq
+ ;; CHECK-NEXT:  if
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   local.tee $temp
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  else
+ ;; CHECK-NEXT:   local.get $param
+ ;; CHECK-NEXT:   local.set $temp
+ ;; CHECK-NEXT:  end
+ ;; CHECK-NEXT:  local.get $temp
+ ;; CHECK-NEXT: )
+ (func $two-covers (param $param (ref eq)) (result (ref eq))
+  (local $temp (ref eq))
+  (local.set $temp
+   (local.get $param)
+  )
+  (if
+   (ref.eq
+    (local.get $temp)
+    (i31.new
+     (i32.const 0)
+    )
+   )
+   ;; In this if arm we write to $temp twice. That shouldn't confuse us; there's
+   ;; still a use after the if, and we should not remove the set-get pair before
+   ;; the if.
+   (local.set $temp
+    (local.tee $temp
+     (local.get $param)
+    )
+   )
+   (local.set $temp
+    (local.get $param)
+   )
+  )
+  (local.get $temp)
+ )
+)