Precompute: Avoid caching heap values of partial precomputations

src/passes/Precompute.cpp

@@ -672,10 +672,15 @@ struct Precompute
         auto** pointerToSelect =
           getChildPointerInImmediateParent(stack, selectIndex, func);
         *pointerToSelect = select->ifTrue;
-        auto ifTrue = precomputeExpression(parent);
+        // When we perform these speculative precomputations, we must not use
+        // the normal heapValues, as we are testing modified versions of
+        // |parent|. Results here must not be cached for later.
+        HeapValues temp;
+        auto ifTrue = precomputeExpression(parent, true, &temp);
+        temp.clear();
         if (isValidPrecomputation(ifTrue)) {
           *pointerToSelect = select->ifFalse;
-          auto ifFalse = precomputeExpression(parent);
+          auto ifFalse = precomputeExpression(parent, true, &temp);
           if (isValidPrecomputation(ifFalse)) {
             // Wonderful, we can precompute here! The select can now contain the
             // computed values in its arms.
@@ -714,12 +719,19 @@ struct Precompute
 
 private:
   // Precompute an expression, returning a flow, which may be a constant
-  // (that we can replace the expression with if replaceExpression is set).
-  Flow precomputeExpression(Expression* curr, bool replaceExpression = true) {
+  // (that we can replace the expression with if replaceExpression is set). When
+  // |usedHeapValues| is provided, we use those values instead of the normal
+  // |heapValues| (that is, we do not use the normal heap value cache).
+  Flow precomputeExpression(Expression* curr,
+                            bool replaceExpression = true,
+                            HeapValues* usedHeapValues = nullptr) {
     Flow flow;
     try {
-      flow = PrecomputingExpressionRunner(
-               getModule(), getValues, heapValues, replaceExpression)
+      flow = PrecomputingExpressionRunner(getModule(),
+                                          getValues,
+                                          usedHeapValues ? *usedHeapValues
+                                                         : heapValues,
+                                          replaceExpression)
                .visit(curr);
     } catch (NonconstantException&) {
       return Flow(NONCONSTANT_FLOW);

test/lit/passes/precompute-propagate-partial.wast

@@ -0,0 +1,55 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --precompute-propagate --optimize-level=2 -all -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $empty (struct))
+  (type $empty (struct))
+
+  ;; CHECK:      (type $i32_any (sub (struct (field i32) (field (ref any)))))
+  (type $i32_any (sub (struct (field i32) (field (ref any)))))
+
+  ;; CHECK:      (func $partial-gc-cache (type $1) (result i32)
+  ;; CHECK-NEXT:  (local $x i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (struct.new $i32_any
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:    (select (result (ref (exact $empty)))
+  ;; CHECK-NEXT:     (struct.new_default $empty)
+  ;; CHECK-NEXT:     (struct.new_default $empty)
+  ;; CHECK-NEXT:     (i32.trunc_f64_s
+  ;; CHECK-NEXT:      (f64.const 562949953421311)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $partial-gc-cache (result i32)
+    (local $x i32)
+    ;; The select here looks promising for partial computation. We will therefore
+    ;; try to precompute variations of the outer struct.new. While doing so we
+    ;; should not confuse later computations, and in particular, not store
+    ;; anything in the heap value cache - the risk is that these partial
+    ;; precomputations will not have the i32.trunc which traps, so they will
+    ;; execute without trapping, and if we cached that result for the outer
+    ;; struct.new, we'd later think that results applies here (if we did, we'd
+    ;; think this does not trap, and precompute it into a nop, removing the
+    ;; trap erroneously).
+    (drop
+      (struct.new $i32_any
+        (i32.const 0)
+        (select (result (ref $empty))
+          (struct.new $empty)
+          (struct.new $empty)
+          (i32.trunc_f64_s
+            (f64.const 562949953421311)
+          )
+        )
+      )
+    )
+    ;; This local causes propagation to work, which causes the extra
+    ;; optimizations that lead to the bug above.
+    (local.get $x)
+  )
+)