[Wasm GC] Heap2Local: Handle branches (#3881)

If we branch to a block, and there are no other branches or a final
value on the block either, then there is no mixing, and we may be able
to optimize the allocation. Before this PR, all branches stopped us.

To do this, add some helpers in BranchUtils.

The main flow logic in Heap2Local used to stop when we reached a
child for the second time. With branches, however, a child can flow both to
its immediate parent, and to branch targets, and so the proper thing to look
at is when we reach a parent for the second time (which would definitely
indicate mixing).

Tests are added for the new functionality. Note that some existing tests
already covered some things we should not optimize, and so no tests
were needed for them. The existing ones are: $get-through-block,
$branch-to-block.

1 files changed, 221 insertions, 7 deletions

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 29c25a8df..3736e4990 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -239,11 +239,13 @@
 
   ;; CHECK:      (func $ignore-unreachable
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (block
-  ;; CHECK-NEXT:    (struct.new_with_rtt $struct.A
-  ;; CHECK-NEXT:     (i32.const 2)
-  ;; CHECK-NEXT:     (unreachable)
-  ;; CHECK-NEXT:     (rtt.canon $struct.A)
+  ;; CHECK-NEXT:   (block ;; (replaces something unreachable we can't emit)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (struct.new_with_rtt $struct.A
+  ;; CHECK-NEXT:      (i32.const 2)
+  ;; CHECK-NEXT:      (unreachable)
+  ;; CHECK-NEXT:      (rtt.canon $struct.A)
+  ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
@@ -729,7 +731,7 @@
     )
   )
 
-  ;; CHECK:      (func $branch-value (result f64)
+  ;; CHECK:      (func $block-value (result f64)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
@@ -758,7 +760,7 @@
   ;; CHECK-NEXT:   (local.get $2)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  (func $branch-value (result f64)
+  (func $block-value (result f64)
     (local $ref (ref null $struct.A))
     (local.set $ref
       (struct.new_default_with_rtt $struct.A
@@ -1530,6 +1532,218 @@
     )
   )
 
+  ;; CHECK:      (func $branch-to-block-no-fallthrough (result f64)
+  ;; CHECK-NEXT:  (local $0 (ref null $struct.A))
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result (ref $struct.A))
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.new_default_with_rtt $struct.A
+  ;; CHECK-NEXT:     (rtt.canon $struct.A)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block (result f64)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (block $block (result (ref null $struct.A))
+  ;; CHECK-NEXT:     (drop
+  ;; CHECK-NEXT:      (br_if $block
+  ;; CHECK-NEXT:       (local.get $0)
+  ;; CHECK-NEXT:       (i32.const 0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (return
+  ;; CHECK-NEXT:      (f64.const 2.1828)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.get $2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $branch-to-block-no-fallthrough (result f64)
+    (local $0 (ref null $struct.A))
+    (local.set $0
+      (struct.new_default_with_rtt $struct.A
+        (rtt.canon $struct.A)
+      )
+    )
+    (struct.get $struct.A 1
+      (block $block (result (ref null $struct.A))
+        (drop
+          ;; A branch to the block of our allocation. In this case there is no
+          ;; other value reaching the block, and so our branch is the sole value
+          ;; which means there is no mixing, and we can optimize this.
+          (br_if $block
+            (local.get $0)
+            (i32.const 0)
+          )
+        )
+        (return (f64.const 2.1828))
+      )
+    )
+  )
+
+  ;; CHECK:      (func $two-branches (result f64)
+  ;; CHECK-NEXT:  (local $0 (ref null $struct.A))
+  ;; CHECK-NEXT:  (local.set $0
+  ;; CHECK-NEXT:   (struct.new_default_with_rtt $struct.A
+  ;; CHECK-NEXT:    (rtt.canon $struct.A)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct.A 1
+  ;; CHECK-NEXT:   (block $block (result (ref null $struct.A))
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (br_if $block
+  ;; CHECK-NEXT:      (local.get $0)
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (br_if $block
+  ;; CHECK-NEXT:      (ref.null $struct.A)
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (f64.const 2.1828)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $two-branches (result f64)
+    (local $0 (ref null $struct.A))
+    (local.set $0
+      (struct.new_default_with_rtt $struct.A
+        (rtt.canon $struct.A)
+      )
+    )
+    (struct.get $struct.A 1
+      (block $block (result (ref null $struct.A))
+        (drop
+          ;; A branch to the block of our allocation.
+          (br_if $block
+            (local.get $0)
+            (i32.const 0)
+          )
+        )
+        (drop
+          ;; Another branch, causing mixing that prevents optimizations.
+          (br_if $block
+            (ref.null $struct.A)
+            (i32.const 0)
+          )
+        )
+        (return (f64.const 2.1828))
+      )
+    )
+  )
+
+  ;; CHECK:      (func $two-branches-b (result f64)
+  ;; CHECK-NEXT:  (local $0 (ref null $struct.A))
+  ;; CHECK-NEXT:  (local.set $0
+  ;; CHECK-NEXT:   (struct.new_default_with_rtt $struct.A
+  ;; CHECK-NEXT:    (rtt.canon $struct.A)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct.A 1
+  ;; CHECK-NEXT:   (block $block (result (ref null $struct.A))
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (br_if $block
+  ;; CHECK-NEXT:      (local.get $0)
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (br_if $block
+  ;; CHECK-NEXT:      (local.get $0)
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (f64.const 2.1828)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $two-branches-b (result f64)
+    (local $0 (ref null $struct.A))
+    (local.set $0
+      (struct.new_default_with_rtt $struct.A
+        (rtt.canon $struct.A)
+      )
+    )
+    (struct.get $struct.A 1
+      (block $block (result (ref null $struct.A))
+        (drop
+          (br_if $block
+            (local.get $0)
+            (i32.const 0)
+          )
+        )
+        (drop
+          ;; As in $two-branches, but the value here is our allocation, the same
+          ;; as in the first branch above us. We do not yet optimize such merges
+          ;; of our allocation, but we could in the future.
+          (br_if $block
+            (local.get $0)
+            (i32.const 0)
+          )
+        )
+        (return (f64.const 2.1828))
+      )
+    )
+  )
+
+  ;; CHECK:      (func $br_if_flow (result f64)
+  ;; CHECK-NEXT:  (local $0 (ref null $struct.A))
+  ;; CHECK-NEXT:  (local.set $0
+  ;; CHECK-NEXT:   (struct.new_default_with_rtt $struct.A
+  ;; CHECK-NEXT:    (rtt.canon $struct.A)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct.A 1
+  ;; CHECK-NEXT:   (block $block (result (ref null $struct.A))
+  ;; CHECK-NEXT:    (call $send-ref
+  ;; CHECK-NEXT:     (br_if $block
+  ;; CHECK-NEXT:      (local.get $0)
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (f64.const 2.1828)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $br_if_flow (result f64)
+    (local $0 (ref null $struct.A))
+    (local.set $0
+      (struct.new_default_with_rtt $struct.A
+        (rtt.canon $struct.A)
+      )
+    )
+    (struct.get $struct.A 1
+      (block $block (result (ref null $struct.A))
+        ;; If it were not for the call here then we would be able to optimize
+        ;; the allocation in this function. (The branch with the allocation is
+        ;; ok, but the br_if also flows the value into a call, that escapes it.)
+        (call $send-ref
+          (br_if $block
+            (local.get $0)
+            (i32.const 0)
+          )
+        )
+        (return (f64.const 2.1828))
+      )
+    )
+  )
+
   ;; CHECK:      (func $ref-as-non-null
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)