From aaca19262b73074ec1c60bf4d4dad26edb8164b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 16:04:05 +0100
Subject: [PATCH 1/9] Switch instruction: add test

---
 scripts/test/fuzzing.py                      |  3 ++-
 test/lit/basic/stack_switching_switch_2.wast | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 test/lit/basic/stack_switching_switch_2.wast

diff --git a/scripts/test/fuzzing.py b/scripts/test/fuzzing.py
index 78b78a089d8..d9121f12c85 100644
--- a/scripts/test/fuzzing.py
+++ b/scripts/test/fuzzing.py
@@ -97,7 +97,8 @@
     'stack_switching_suspend.wast',
     'stack_switching_resume.wast',
     'stack_switching_resume_throw.wast',
-    'stack_switching_switch.wast'
+    'stack_switching_switch.wast',
+    'stack_switching_switch.wast_2'
 ]
 
 
diff --git a/test/lit/basic/stack_switching_switch_2.wast b/test/lit/basic/stack_switching_switch_2.wast
new file mode 100644
index 00000000000..64d42abdd4c
--- /dev/null
+++ b/test/lit/basic/stack_switching_switch_2.wast
@@ -0,0 +1,18 @@
+;; RUN: not wasm-opt -all %s -S -o - 2>&1 | filecheck %s
+
+;; CHECK: {{.*}}wasm-validator error{{.*}}function body type must match,{{.*}}
+
+(module
+ (type $function (func (param i64)))
+ (type $cont (cont $function))
+ (type $function_2 (func (param i32 (ref $cont))))
+ (type $cont_2 (cont $function_2))
+ (tag $tag)
+
+ (func $switch (param $c (ref $cont_2)) (result i64)
+  (switch $cont_2 $tag
+   (i32.const 0)
+   (local.get $c)
+  )
+ )
+)

From 4afd78db37077c68b8192162771d5b0dbd687866 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 16:20:01 +0100
Subject: [PATCH 2/9] Fix type of switch instruction

---
 src/wasm/wasm-ir-builder.cpp                 |  7 ++++++-
 src/wasm/wasm.cpp                            |  6 +++++-
 test/lit/basic/stack_switching_switch_2.wast | 15 +++++++++++++--
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 78e622b3cb9..8cd62f561ca 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -2477,11 +2477,16 @@ Result<> IRBuilder::makeStackSwitch(HeapType ct, Name tag) {
   }
   StackSwitch curr(wasm.allocator);
   curr.tag = tag;
-  auto nparams = ct.getContinuation().type.getSignature().params.size();
+  Type params = ct.getContinuation().type.getSignature().params;
+  auto nparams = params.size();
   if (nparams < 1) {
     return Err{"arity mismatch: the continuation argument must have, at least, "
                "unary arity"};
   }
+  if (!params[nparams - 1].isContinuation()) {
+    return Err{"the last argument of the continuation argument should be "
+               "itself a continuation"};
+  }
 
   // The continuation argument of the continuation is synthetic,
   // i.e. it is provided by the runtime.
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index c8a82311c30..52506f4f684 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -1441,8 +1441,12 @@ void StackSwitch::finalize() {
   }
 
   assert(this->cont->type.isContinuation());
-  type =
+  Type params =
     this->cont->type.getHeapType().getContinuation().type.getSignature().params;
+  assert(params.size() > 0);
+  Type cont = params[params.size() - 1];
+  assert(cont.isContinuation());
+  type = cont.getHeapType().getContinuation().type.getSignature().params;
 }
 
 size_t Function::getNumParams() { return getParams().size(); }
diff --git a/test/lit/basic/stack_switching_switch_2.wast b/test/lit/basic/stack_switching_switch_2.wast
index 64d42abdd4c..2e7eb91fb35 100644
--- a/test/lit/basic/stack_switching_switch_2.wast
+++ b/test/lit/basic/stack_switching_switch_2.wast
@@ -1,14 +1,25 @@
-;; RUN: not wasm-opt -all %s -S -o - 2>&1 | filecheck %s
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt -all %s -S -o - | filecheck %s
 
-;; CHECK: {{.*}}wasm-validator error{{.*}}function body type must match,{{.*}}
 
 (module
+ ;; CHECK:      (type $function (func (param i64)))
  (type $function (func (param i64)))
+ ;; CHECK:      (type $cont (cont $function))
  (type $cont (cont $function))
+ ;; CHECK:      (type $function_2 (func (param i32 (ref $cont))))
  (type $function_2 (func (param i32 (ref $cont))))
+ ;; CHECK:      (type $cont_2 (cont $function_2))
  (type $cont_2 (cont $function_2))
+ ;; CHECK:      (tag $tag (type $4))
  (tag $tag)
 
+ ;; CHECK:      (func $switch (type $5) (param $c (ref $cont_2)) (result i64)
+ ;; CHECK-NEXT:  (switch $cont_2 $tag
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:   (local.get $c)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $switch (param $c (ref $cont_2)) (result i64)
   (switch $cont_2 $tag
    (i32.const 0)

From e04705b107ac2ada4764a48bd11c36b40e9e7781 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 15:46:47 +0100
Subject: [PATCH 3/9] Add optimization tests

---
 scripts/test/fuzzing.py                       |   7 +-
 test/lit/passes/O3_stack-switching.wast       |  91 ++++++++
 .../coalesce-locals-stack-switching.wast      | 202 ++++++++++++++++++
 test/lit/passes/dce-stack-switching.wast      |  61 ++++++
 .../passes/precompute-stack-switching.wast    |  51 +++++
 test/lit/passes/vacuum-stack-switching.wast   |  65 ++++++
 6 files changed, 476 insertions(+), 1 deletion(-)
 create mode 100644 test/lit/passes/O3_stack-switching.wast
 create mode 100644 test/lit/passes/coalesce-locals-stack-switching.wast
 create mode 100644 test/lit/passes/dce-stack-switching.wast
 create mode 100644 test/lit/passes/precompute-stack-switching.wast
 create mode 100644 test/lit/passes/vacuum-stack-switching.wast

diff --git a/scripts/test/fuzzing.py b/scripts/test/fuzzing.py
index d9121f12c85..8cbf3ebfdfa 100644
--- a/scripts/test/fuzzing.py
+++ b/scripts/test/fuzzing.py
@@ -98,7 +98,12 @@
     'stack_switching_resume.wast',
     'stack_switching_resume_throw.wast',
     'stack_switching_switch.wast',
-    'stack_switching_switch.wast_2'
+    'stack_switching_switch.wast_2',
+    'O3_stack-switching.wast',
+    'coalesce-locals-stack-switching.wast',
+    'dce-stack-switching.wast',
+    'precompute-stack-switching.wast',
+    'vacuum-stack-switching.wast'
 ]
 
 
diff --git a/test/lit/passes/O3_stack-switching.wast b/test/lit/passes/O3_stack-switching.wast
new file mode 100644
index 00000000000..71c2b065865
--- /dev/null
+++ b/test/lit/passes/O3_stack-switching.wast
@@ -0,0 +1,91 @@
+;; RUN: not wasm-opt -all -O3 %s -S -o - 2>&1 | filecheck %s
+
+;; Fairly comprehensive test case
+
+;; CHECK: unimplemented
+(module
+ (type $function_1 (func (param (ref eq) (ref eq)) (result (ref eq))))
+ (type $closure (sub (struct (field (ref $function_1)))))
+ (type $function_2 (func (param (ref eq) (ref eq) (ref eq)) (result (ref eq))))
+ (type $closure_2 (struct (field (ref $function_2))))
+ (type $handlers (struct (field $value (ref $closure)) (field $exn (ref $closure)) (field $effect (ref $closure_2))))
+ (type $cont (cont $function_1))
+ (type $fiber (struct (field $handlers (ref $handlers)) (field $cont (ref $cont))))
+ (tag $exception (param (ref eq)))
+ (tag $effect (param (ref eq)) (result (ref eq) (ref eq)))
+
+ (func $resume (export "resume") (param $fiber (ref $fiber)) (param $f (ref $closure)) (param $v (ref eq)) (result (ref eq))
+  (local $g (ref $closure_2))
+  (local $res (ref eq))
+  (local $exn (ref eq))
+  (local $resume_res (tuple (ref eq) (ref $cont)))
+  (local.set $exn
+   (block $handle_exception (result (ref eq))
+    (local.set $resume_res
+     (block $handle_effect (result (ref eq) (ref $cont))
+      (local.set $res
+       (try_table (result (ref eq)) (catch $exception $handle_exception)
+        (resume $cont (on $effect $handle_effect)
+         (local.get $f)
+         (local.get $v)
+         (struct.get $fiber $cont
+          (local.get $fiber)
+         )
+        )
+       )
+      )
+      (return_call_ref $function_1
+       (local.get $res)
+       (local.tee $f
+        (struct.get $handlers $value
+         (struct.get $fiber $handlers
+          (local.get $fiber)
+         )
+        )
+       )
+       (struct.get $closure 0
+        (local.get $f)
+       )
+      )
+     )
+    )
+    (return_call_ref $function_2
+     (tuple.extract 2 0
+      (local.get $resume_res)
+     )
+     (struct.new $fiber
+      (struct.get $fiber $handlers
+       (local.get $fiber)
+      )
+      (tuple.extract 2 1
+       (local.get $resume_res)
+      )
+     )
+     (local.tee $g
+      (struct.get $handlers $effect
+       (struct.get $fiber $handlers
+        (local.get $fiber)
+       )
+      )
+     )
+     (struct.get $closure_2 0
+      (local.get $g)
+     )
+    )
+   )
+  )
+  (return_call_ref $function_1
+   (local.get $exn)
+   (local.tee $f
+    (struct.get $handlers $exn
+     (struct.get $fiber $handlers
+      (local.get $fiber)
+     )
+    )
+   )
+   (struct.get $closure 0
+    (local.get $f)
+   )
+  )
+ )
+)
diff --git a/test/lit/passes/coalesce-locals-stack-switching.wast b/test/lit/passes/coalesce-locals-stack-switching.wast
new file mode 100644
index 00000000000..679344cae02
--- /dev/null
+++ b/test/lit/passes/coalesce-locals-stack-switching.wast
@@ -0,0 +1,202 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt -all --coalesce-locals %s -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $function (func))
+ (type $function (func))
+ ;; CHECK:      (type $cont (cont $function))
+ (type $cont (cont $function))
+ ;; CHECK:      (type $function_2 (func (param (ref $cont))))
+ (type $function_2 (func (param (ref $cont))))
+ ;; CHECK:      (type $cont_2 (cont $function_2))
+ (type $cont_2 (cont $function_2))
+
+ ;; CHECK:      (tag $tag (type $function))
+ (tag $tag)
+ ;; CHECK:      (tag $exn (type $5) (param i32))
+ (tag $exn (param i32))
+
+ ;; CHECK:      (func $resume (type $2) (param $0 (ref $cont)) (result i32)
+ ;; CHECK-NEXT:  (local $1 i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $handle_exn (result i32)
+ ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
+ ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:      (block $handle_effect (result (ref $cont))
+ ;; CHECK-NEXT:       (resume $cont (on $tag $handle_effect)
+ ;; CHECK-NEXT:        (local.get $0)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:       (return
+ ;; CHECK-NEXT:        (i32.const 0)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (return_call $resume
+ ;; CHECK-NEXT:      (block (result (ref $cont))
+ ;; CHECK-NEXT:       (unreachable)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return
+ ;; CHECK-NEXT:     (i32.const 0)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT: )
+ (func $resume (param $c (ref $cont)) (result i32)
+  (local $x i32)
+  (local $y (ref $cont))
+  ;; resume can raise an exception, so this local.set should not be
+  ;; optimized away
+  (local.set $x
+   (block $handle_exn (result i32)
+    (try_table (catch $exn $handle_exn)
+     ;; this local.set is reachable, so it should not be optimized away
+     (local.set $y
+      (block $handle_effect (result (ref $cont))
+       (resume $cont (on $tag $handle_effect)
+        (local.get $c)
+       )
+       (return
+        (i32.const 0)
+       )
+      )
+     )
+     (return_call $resume
+      (local.get $y)
+     )
+    )
+    (return
+     (i32.const 0)
+    )
+   )
+  )
+  (local.get $x)
+ )
+
+ ;; CHECK:      (func $resume_throw (type $2) (param $0 (ref $cont)) (result i32)
+ ;; CHECK-NEXT:  (local $1 i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $handle_exn (result i32)
+ ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
+ ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:      (block $handle_effect (result (ref $cont))
+ ;; CHECK-NEXT:       (resume_throw $cont $tag (on $tag $handle_effect)
+ ;; CHECK-NEXT:        (local.get $0)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:       (return
+ ;; CHECK-NEXT:        (i32.const 0)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (return_call $resume
+ ;; CHECK-NEXT:      (block (result (ref $cont))
+ ;; CHECK-NEXT:       (unreachable)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return
+ ;; CHECK-NEXT:     (i32.const 0)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT: )
+ (func $resume_throw (param $c (ref $cont)) (result i32)
+  (local $x i32)
+  (local $y (ref $cont))
+  ;; resume can raise an exception, so this local.set should not be
+  ;; optimized away
+  (local.set $x
+   (block $handle_exn (result i32)
+    (try_table (catch $exn $handle_exn)
+     ;; this local.set is reachable, so it should not be optimized away
+     (local.set $y
+      (block $handle_effect (result (ref $cont))
+       (resume_throw $cont $tag (on $tag $handle_effect)
+        (local.get $c)
+       )
+       (return
+        (i32.const 0)
+       )
+      )
+     )
+     (return_call $resume
+      (local.get $y)
+     )
+    )
+    (return
+     (i32.const 0)
+    )
+   )
+  )
+  (local.get $x)
+ )
+
+ ;; CHECK:      (func $suspend (type $6) (result i32)
+ ;; CHECK-NEXT:  (local $0 i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $handle_exn (result i32)
+ ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
+ ;; CHECK-NEXT:     (suspend $tag)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return
+ ;; CHECK-NEXT:     (i32.const 0)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT: )
+ (func $suspend (result i32)
+  (local $x i32)
+  ;; suspend can raise an exception, so this local.set should not be
+  ;; optimized away
+  (local.set $x
+   (block $handle_exn (result i32)
+    (try_table (catch $exn $handle_exn)
+     (suspend $tag)
+    )
+    (return
+     (i32.const 0)
+    )
+   )
+  )
+  (local.get $x)
+ )
+
+ ;; CHECK:      (func $switch (type $7) (param $0 (ref $cont_2)) (result i32)
+ ;; CHECK-NEXT:  (local $1 i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $handle_exn (result i32)
+ ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
+ ;; CHECK-NEXT:     (switch $cont_2 $tag
+ ;; CHECK-NEXT:      (local.get $0)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return
+ ;; CHECK-NEXT:     (i32.const 0)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT: )
+ (func $switch (param $c (ref $cont_2)) (result i32)
+  (local $x i32)
+  ;; switch can raise an exception, so this local.set should not be
+  ;; optimized away
+  (local.set $x
+   (block $handle_exn (result i32)
+    (try_table (catch $exn $handle_exn)
+     (switch $cont_2 $tag
+      (local.get $c)
+     )
+    )
+    (return
+     (i32.const 0)
+    )
+   )
+  )
+  (local.get $x)
+ )
+)
diff --git a/test/lit/passes/dce-stack-switching.wast b/test/lit/passes/dce-stack-switching.wast
new file mode 100644
index 00000000000..f7444165860
--- /dev/null
+++ b/test/lit/passes/dce-stack-switching.wast
@@ -0,0 +1,61 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt -all --dce %s -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $function (func))
+ (type $function (func))
+ ;; CHECK:      (type $cont (cont $function))
+ (type $cont (cont $function))
+ ;; CHECK:      (tag $tag (type $function))
+ (tag $tag)
+
+ ;; CHECK:      (func $resume (type $2) (param $c (ref $cont))
+ ;; CHECK-NEXT:  (block $exit
+ ;; CHECK-NEXT:   (block $handle_effect
+ ;; CHECK-NEXT:    (resume $cont (on $tag $handle_effect)
+ ;; CHECK-NEXT:     (local.get $c)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (br $exit)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume (param $c (ref $cont))
+  (block $exit
+  (drop
+    ;; The return type of the block should not be dropped since we can
+    ;; jump there when handling $tag.
+    (block $handle_effect (result (ref $cont))
+     (resume $cont (on $tag $handle_effect)
+      (local.get $c)
+     )
+     (br $exit)
+    )
+   )
+  )
+ )
+
+ ;; CHECK:      (func $resume_throw (type $2) (param $c (ref $cont))
+ ;; CHECK-NEXT:  (block $exit
+ ;; CHECK-NEXT:   (block $handle_effect
+ ;; CHECK-NEXT:    (resume_throw $cont $tag (on $tag $handle_effect)
+ ;; CHECK-NEXT:     (local.get $c)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (br $exit)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume_throw (param $c (ref $cont))
+  (block $exit
+   (drop
+    ;; The return type of the block should not be dropped since we can
+    ;; jump there when handling $tag.
+    (block $handle_effect (result (ref $cont))
+     (resume_throw $cont $tag (on $tag $handle_effect)
+      (local.get $c)
+     )
+     (br $exit)
+    )
+   )
+  )
+ )
+)
diff --git a/test/lit/passes/precompute-stack-switching.wast b/test/lit/passes/precompute-stack-switching.wast
new file mode 100644
index 00000000000..3216fe22257
--- /dev/null
+++ b/test/lit/passes/precompute-stack-switching.wast
@@ -0,0 +1,51 @@
+;; RUN: not wasm-opt -all --precompute %s -S -o - 2>&1 | grep -o unimplemented | filecheck %s
+
+;; CHECK: unimplemented
+(module
+ (type $function (func))
+ (type $cont (cont $function))
+ (type $function_2 (func (param i32)))
+ (type $cont_2 (cont $function_2))
+ (tag $tag)
+ (type $function_3 (func (param (ref $cont))))
+ (type $cont_3 (cont $function_3))
+
+ (func $cont_new (param $f (ref $function)) (result (ref $cont))
+  (cont.new $cont (local.get $f))
+ )
+
+ (func $cont_bind (param $c (ref $cont_2)) (result (ref $cont))
+  (cont.bind $cont_2 $cont
+   (i32.const 0)
+   (local.get $c)
+  )
+ )
+
+ (func $suspend
+  (suspend $tag)
+ )
+
+ (func $resume (param $c (ref $cont)) (result (ref $cont))
+  (block $handle_effect (result (ref $cont))
+   (resume $cont (on $tag $handle_effect)
+    (local.get $c)
+   )
+   (local.get $c)
+  )
+ )
+
+ (func $resume_throw (param $c (ref $cont)) (result (ref $cont))
+  (block $handle_effect (result (ref $cont))
+   (resume_throw $cont $tag (on $tag $handle_effect)
+    (local.get $c)
+   )
+   (local.get $c)
+  )
+ )
+
+ (func $switch (param $c (ref $cont_3))
+  (switch $cont_3 $tag
+   (local.get $c)
+  )
+ )
+)
diff --git a/test/lit/passes/vacuum-stack-switching.wast b/test/lit/passes/vacuum-stack-switching.wast
new file mode 100644
index 00000000000..acc45e9fca3
--- /dev/null
+++ b/test/lit/passes/vacuum-stack-switching.wast
@@ -0,0 +1,65 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt -all --vacuum %s -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $function (func))
+ (type $function (func))
+ ;; CHECK:      (type $cont (cont $function))
+ (type $cont (cont $function))
+ ;; CHECK:      (tag $tag (type $function))
+ (tag $tag)
+
+ ;; CHECK:      (func $resume (type $2) (param $c (ref $cont))
+ ;; CHECK-NEXT:  (block $exit
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (block $handle_effect
+ ;; CHECK-NEXT:     (resume $cont (on $tag $handle_effect)
+ ;; CHECK-NEXT:      (local.get $c)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (br $exit)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume (param $c (ref $cont))
+  (block $exit
+  (drop
+    ;; The return type of the block should not be dropped since we can
+    ;; jump there when handling $tag.
+    (block $handle_effect (result (ref $cont))
+     (resume $cont (on $tag $handle_effect)
+      (local.get $c)
+     )
+     (br $exit)
+    )
+   )
+  )
+ )
+
+ ;; CHECK:      (func $resume_throw (type $2) (param $c (ref $cont))
+ ;; CHECK-NEXT:  (block $exit
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (block $handle_effect
+ ;; CHECK-NEXT:     (resume_throw $cont $tag (on $tag $handle_effect)
+ ;; CHECK-NEXT:      (local.get $c)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (br $exit)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $resume_throw (param $c (ref $cont))
+  (block $exit
+   (drop
+    ;; The return type of the block should not be dropped since we can
+    ;; jump there when handling $tag.
+    (block $handle_effect (result (ref $cont))
+     (resume_throw $cont $tag (on $tag $handle_effect)
+      (local.get $c)
+     )
+     (br $exit)
+    )
+   )
+  )
+ )
+)

From ef1250d7203d6172aaf04de2951a9c3acc4dd987 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 16:48:42 +0100
Subject: [PATCH 4/9] Fix ReFinalize

Update the type of the blocks that are targets of the handlers of resume
and resume_throw instructions.
---
 src/ir/ReFinalize.cpp | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/ir/ReFinalize.cpp b/src/ir/ReFinalize.cpp
index 42b13919726..df41546ed14 100644
--- a/src/ir/ReFinalize.cpp
+++ b/src/ir/ReFinalize.cpp
@@ -183,8 +183,18 @@ void ReFinalize::visitStringSliceWTF(StringSliceWTF* curr) { curr->finalize(); }
 void ReFinalize::visitContNew(ContNew* curr) { curr->finalize(); }
 void ReFinalize::visitContBind(ContBind* curr) { curr->finalize(); }
 void ReFinalize::visitSuspend(Suspend* curr) { curr->finalize(getModule()); }
-void ReFinalize::visitResume(Resume* curr) { curr->finalize(); }
-void ReFinalize::visitResumeThrow(ResumeThrow* curr) { curr->finalize(); }
+void ReFinalize::visitResume(Resume* curr) {
+  curr->finalize();
+  for (size_t i = 0; i < curr->handlerBlocks.size(); i++) {
+    updateBreakValueType(curr->handlerBlocks[i], curr->sentTypes[i]);
+  }
+}
+void ReFinalize::visitResumeThrow(ResumeThrow* curr) {
+  curr->finalize();
+  for (size_t i = 0; i < curr->handlerBlocks.size(); i++) {
+    updateBreakValueType(curr->handlerBlocks[i], curr->sentTypes[i]);
+  }
+}
 void ReFinalize::visitStackSwitch(StackSwitch* curr) { curr->finalize(); }
 
 void ReFinalize::visitExport(Export* curr) { WASM_UNREACHABLE("unimp"); }

From 000eebf43a72cd7ce899ce2f824c82f9a5aef175 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 16:36:47 +0100
Subject: [PATCH 5/9] Fix operateOnScopeNameUsesAndSentTypes

We were incorrectly comparing a tag name with block types.
---
 src/ir/branch-utils.h                       |  8 ++++----
 test/lit/passes/dce-stack-switching.wast    | 20 ++++++++++++--------
 test/lit/passes/vacuum-stack-switching.wast |  4 ++--
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index 1771f2e0e88..b76c601be3b 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -83,15 +83,15 @@ void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {
         }
       }
     } else if (auto* r = expr->dynCast<Resume>()) {
-      for (Index i = 0; i < r->handlerTags.size(); i++) {
-        auto dest = r->handlerTags[i];
+      for (Index i = 0; i < r->handlerBlocks.size(); i++) {
+        auto dest = r->handlerBlocks[i];
         if (!dest.isNull() && dest == name) {
           func(name, r->sentTypes[i]);
         }
       }
     } else if (auto* r = expr->dynCast<ResumeThrow>()) {
-      for (Index i = 0; i < r->handlerTags.size(); i++) {
-        auto dest = r->handlerTags[i];
+      for (Index i = 0; i < r->handlerBlocks.size(); i++) {
+        auto dest = r->handlerBlocks[i];
         if (!dest.isNull() && dest == name) {
           func(name, r->sentTypes[i]);
         }
diff --git a/test/lit/passes/dce-stack-switching.wast b/test/lit/passes/dce-stack-switching.wast
index f7444165860..763184a9278 100644
--- a/test/lit/passes/dce-stack-switching.wast
+++ b/test/lit/passes/dce-stack-switching.wast
@@ -11,11 +11,13 @@
 
  ;; CHECK:      (func $resume (type $2) (param $c (ref $cont))
  ;; CHECK-NEXT:  (block $exit
- ;; CHECK-NEXT:   (block $handle_effect
- ;; CHECK-NEXT:    (resume $cont (on $tag $handle_effect)
- ;; CHECK-NEXT:     (local.get $c)
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (block $handle_effect (result (ref $cont))
+ ;; CHECK-NEXT:     (resume $cont (on $tag $handle_effect)
+ ;; CHECK-NEXT:      (local.get $c)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (br $exit)
  ;; CHECK-NEXT:    )
- ;; CHECK-NEXT:    (br $exit)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
@@ -36,11 +38,13 @@
 
  ;; CHECK:      (func $resume_throw (type $2) (param $c (ref $cont))
  ;; CHECK-NEXT:  (block $exit
- ;; CHECK-NEXT:   (block $handle_effect
- ;; CHECK-NEXT:    (resume_throw $cont $tag (on $tag $handle_effect)
- ;; CHECK-NEXT:     (local.get $c)
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (block $handle_effect (result (ref $cont))
+ ;; CHECK-NEXT:     (resume_throw $cont $tag (on $tag $handle_effect)
+ ;; CHECK-NEXT:      (local.get $c)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (br $exit)
  ;; CHECK-NEXT:    )
- ;; CHECK-NEXT:    (br $exit)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
diff --git a/test/lit/passes/vacuum-stack-switching.wast b/test/lit/passes/vacuum-stack-switching.wast
index acc45e9fca3..18250baddaa 100644
--- a/test/lit/passes/vacuum-stack-switching.wast
+++ b/test/lit/passes/vacuum-stack-switching.wast
@@ -12,7 +12,7 @@
  ;; CHECK:      (func $resume (type $2) (param $c (ref $cont))
  ;; CHECK-NEXT:  (block $exit
  ;; CHECK-NEXT:   (drop
- ;; CHECK-NEXT:    (block $handle_effect
+ ;; CHECK-NEXT:    (block $handle_effect (result (ref $cont))
  ;; CHECK-NEXT:     (resume $cont (on $tag $handle_effect)
  ;; CHECK-NEXT:      (local.get $c)
  ;; CHECK-NEXT:     )
@@ -39,7 +39,7 @@
  ;; CHECK:      (func $resume_throw (type $2) (param $c (ref $cont))
  ;; CHECK-NEXT:  (block $exit
  ;; CHECK-NEXT:   (drop
- ;; CHECK-NEXT:    (block $handle_effect
+ ;; CHECK-NEXT:    (block $handle_effect (result (ref $cont))
  ;; CHECK-NEXT:     (resume_throw $cont $tag (on $tag $handle_effect)
  ;; CHECK-NEXT:      (local.get $c)
  ;; CHECK-NEXT:     )

From 4de3787f74a69467bf9f05954eeb769ff146c24b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 16:29:09 +0100
Subject: [PATCH 6/9] Update ConstantExpressionRunner

---
 src/wasm-interpreter.h                        | 26 +++++--
 test/lit/passes/O3_stack-switching.wast       | 75 ++++++++++++++++++-
 .../passes/precompute-stack-switching.wast    | 48 +++++++++++-
 3 files changed, 139 insertions(+), 10 deletions(-)

diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 2a019513d2a..3b633077fac 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -2617,15 +2617,29 @@ class ConstantExpressionRunner : public ExpressionRunner<SubType> {
     }
     return ExpressionRunner<SubType>::visitRefAs(curr);
   }
-  Flow visitContNew(ContNew* curr) { WASM_UNREACHABLE("unimplemented"); }
-  Flow visitContBind(ContBind* curr) { WASM_UNREACHABLE("unimplemented"); }
-  Flow visitSuspend(Suspend* curr) { WASM_UNREACHABLE("unimplemented"); }
-  Flow visitResume(Resume* curr) { WASM_UNREACHABLE("unimplemented"); }
+  Flow visitContNew(ContNew* curr) {
+    NOTE_ENTER("ContNew");
+    return Flow(NONCONSTANT_FLOW);
+  }
+  Flow visitContBind(ContBind* curr) {
+    NOTE_ENTER("ContBind");
+    return Flow(NONCONSTANT_FLOW);
+  }
+  Flow visitSuspend(Suspend* curr) {
+    NOTE_ENTER("Suspend");
+    return Flow(NONCONSTANT_FLOW);
+  }
+  Flow visitResume(Resume* curr) {
+    NOTE_ENTER("Resume");
+    return Flow(NONCONSTANT_FLOW);
+  }
   Flow visitResumeThrow(ResumeThrow* curr) {
-    WASM_UNREACHABLE("unimplemented");
+    NOTE_ENTER("ResumeThrow");
+    return Flow(NONCONSTANT_FLOW);
   }
   Flow visitStackSwitch(StackSwitch* curr) {
-    WASM_UNREACHABLE("unimplemented");
+    NOTE_ENTER("StackSwitch");
+    return Flow(NONCONSTANT_FLOW);
   }
 
   void trap(const char* why) override { throw NonconstantException(); }
diff --git a/test/lit/passes/O3_stack-switching.wast b/test/lit/passes/O3_stack-switching.wast
index 71c2b065865..98bc066ffa2 100644
--- a/test/lit/passes/O3_stack-switching.wast
+++ b/test/lit/passes/O3_stack-switching.wast
@@ -1,19 +1,90 @@
-;; RUN: not wasm-opt -all -O3 %s -S -o - 2>&1 | filecheck %s
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt -all -O3 %s -S -o - | filecheck %s
 
 ;; Fairly comprehensive test case
 
-;; CHECK: unimplemented
 (module
+ ;; CHECK:      (type $function_1 (func (param (ref eq) (ref eq)) (result (ref eq))))
  (type $function_1 (func (param (ref eq) (ref eq)) (result (ref eq))))
+ ;; CHECK:      (type $closure (sub (struct (field (ref $function_1)))))
  (type $closure (sub (struct (field (ref $function_1)))))
+ ;; CHECK:      (type $cont (cont $function_1))
+
+ ;; CHECK:      (type $function_2 (func (param (ref eq) (ref eq) (ref eq)) (result (ref eq))))
  (type $function_2 (func (param (ref eq) (ref eq) (ref eq)) (result (ref eq))))
+ ;; CHECK:      (type $closure_2 (struct (field (ref $function_2))))
  (type $closure_2 (struct (field (ref $function_2))))
+ ;; CHECK:      (type $handlers (struct (field $value (ref $closure)) (field $exn (ref $closure)) (field $effect (ref $closure_2))))
  (type $handlers (struct (field $value (ref $closure)) (field $exn (ref $closure)) (field $effect (ref $closure_2))))
  (type $cont (cont $function_1))
+ ;; CHECK:      (type $fiber (struct (field $handlers (ref $handlers)) (field $cont (ref $cont))))
  (type $fiber (struct (field $handlers (ref $handlers)) (field $cont (ref $cont))))
+ ;; CHECK:      (tag $exception (type $8) (param (ref eq)))
  (tag $exception (param (ref eq)))
+ ;; CHECK:      (tag $effect (type $9) (param (ref eq)) (result (ref eq) (ref eq)))
  (tag $effect (param (ref eq)) (result (ref eq) (ref eq)))
 
+ ;; CHECK:      (func $resume (type $10) (param $0 (ref $fiber)) (param $1 (ref $closure)) (param $2 (ref eq)) (result (ref eq))
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $handle_exception (result (ref eq))
+ ;; CHECK-NEXT:    (tuple.drop 2
+ ;; CHECK-NEXT:     (block $handle_effect (type $3) (result (ref eq) (ref $cont))
+ ;; CHECK-NEXT:      (return_call_ref $function_1
+ ;; CHECK-NEXT:       (try_table (result (ref eq)) (catch $exception $handle_exception)
+ ;; CHECK-NEXT:        (resume $cont (on $effect $handle_effect)
+ ;; CHECK-NEXT:         (local.get $1)
+ ;; CHECK-NEXT:         (local.get $2)
+ ;; CHECK-NEXT:         (struct.get $fiber $cont
+ ;; CHECK-NEXT:          (local.get $0)
+ ;; CHECK-NEXT:         )
+ ;; CHECK-NEXT:        )
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:       (local.tee $1
+ ;; CHECK-NEXT:        (struct.get $handlers $value
+ ;; CHECK-NEXT:         (struct.get $fiber $handlers
+ ;; CHECK-NEXT:          (local.get $0)
+ ;; CHECK-NEXT:         )
+ ;; CHECK-NEXT:        )
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:       (struct.get $closure 0
+ ;; CHECK-NEXT:        (local.get $1)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (drop
+ ;; CHECK-NEXT:     (unreachable)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (block ;; (replaces unreachable CallRef we can't emit)
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (unreachable)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (block ;; (replaces unreachable StructGet we can't emit)
+ ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:      (block ;; (replaces unreachable StructGet we can't emit)
+ ;; CHECK-NEXT:       (drop
+ ;; CHECK-NEXT:        (unreachable)
+ ;; CHECK-NEXT:       )
+ ;; CHECK-NEXT:       (unreachable)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (unreachable)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (block ;; (replaces unreachable StructGet we can't emit)
+ ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:      (unreachable)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (unreachable)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (unreachable)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $resume (export "resume") (param $fiber (ref $fiber)) (param $f (ref $closure)) (param $v (ref eq)) (result (ref eq))
   (local $g (ref $closure_2))
   (local $res (ref eq))
diff --git a/test/lit/passes/precompute-stack-switching.wast b/test/lit/passes/precompute-stack-switching.wast
index 3216fe22257..0677fee0ccb 100644
--- a/test/lit/passes/precompute-stack-switching.wast
+++ b/test/lit/passes/precompute-stack-switching.wast
@@ -1,19 +1,39 @@
-;; RUN: not wasm-opt -all --precompute %s -S -o - 2>&1 | grep -o unimplemented | filecheck %s
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt -all --precompute %s -S -o - | filecheck %s
 
-;; CHECK: unimplemented
 (module
+ ;; CHECK:      (type $function (func))
  (type $function (func))
+ ;; CHECK:      (type $cont (cont $function))
  (type $cont (cont $function))
+ ;; CHECK:      (type $function_2 (func (param i32)))
  (type $function_2 (func (param i32)))
+ ;; CHECK:      (type $cont_2 (cont $function_2))
  (type $cont_2 (cont $function_2))
+ ;; CHECK:      (type $function_3 (func (param (ref $cont))))
+
+ ;; CHECK:      (type $cont_3 (cont $function_3))
+
+ ;; CHECK:      (tag $tag (type $function))
  (tag $tag)
  (type $function_3 (func (param (ref $cont))))
  (type $cont_3 (cont $function_3))
 
+ ;; CHECK:      (func $cont_new (type $7) (param $f (ref $function)) (result (ref $cont))
+ ;; CHECK-NEXT:  (cont.new $cont
+ ;; CHECK-NEXT:   (local.get $f)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $cont_new (param $f (ref $function)) (result (ref $cont))
   (cont.new $cont (local.get $f))
  )
 
+ ;; CHECK:      (func $cont_bind (type $8) (param $c (ref $cont_2)) (result (ref $cont))
+ ;; CHECK-NEXT:  (cont.bind $cont_2 $cont
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:   (local.get $c)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $cont_bind (param $c (ref $cont_2)) (result (ref $cont))
   (cont.bind $cont_2 $cont
    (i32.const 0)
@@ -21,10 +41,21 @@
   )
  )
 
+ ;; CHECK:      (func $suspend (type $function)
+ ;; CHECK-NEXT:  (suspend $tag)
+ ;; CHECK-NEXT: )
  (func $suspend
   (suspend $tag)
  )
 
+ ;; CHECK:      (func $resume (type $2) (param $c (ref $cont)) (result (ref $cont))
+ ;; CHECK-NEXT:  (block $handle_effect (result (ref $cont))
+ ;; CHECK-NEXT:   (resume $cont (on $tag $handle_effect)
+ ;; CHECK-NEXT:    (local.get $c)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (local.get $c)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $resume (param $c (ref $cont)) (result (ref $cont))
   (block $handle_effect (result (ref $cont))
    (resume $cont (on $tag $handle_effect)
@@ -34,6 +65,14 @@
   )
  )
 
+ ;; CHECK:      (func $resume_throw (type $2) (param $c (ref $cont)) (result (ref $cont))
+ ;; CHECK-NEXT:  (block $handle_effect (result (ref $cont))
+ ;; CHECK-NEXT:   (resume_throw $cont $tag (on $tag $handle_effect)
+ ;; CHECK-NEXT:    (local.get $c)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (local.get $c)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $resume_throw (param $c (ref $cont)) (result (ref $cont))
   (block $handle_effect (result (ref $cont))
    (resume_throw $cont $tag (on $tag $handle_effect)
@@ -43,6 +82,11 @@
   )
  )
 
+ ;; CHECK:      (func $switch (type $9) (param $c (ref $cont_3))
+ ;; CHECK-NEXT:  (switch $cont_3 $tag
+ ;; CHECK-NEXT:   (local.get $c)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
  (func $switch (param $c (ref $cont_3))
   (switch $cont_3 $tag
    (local.get $c)

From 50d64186e5ca54f3678ee6278d4d23f8951e6070 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Fri, 7 Feb 2025 17:28:20 +0100
Subject: [PATCH 7/9] Update the computation of the control flow graph

---
 src/cfg/cfg-traversal.h                       | 27 ++++++
 test/lit/passes/O3_stack-switching.wast       | 91 ++++++++++---------
 .../coalesce-locals-stack-switching.wast      | 28 +++---
 3 files changed, 87 insertions(+), 59 deletions(-)

diff --git a/src/cfg/cfg-traversal.h b/src/cfg/cfg-traversal.h
index 64877c58cfc..46ecc4b2a30 100644
--- a/src/cfg/cfg-traversal.h
+++ b/src/cfg/cfg-traversal.h
@@ -444,6 +444,19 @@ struct CFGWalker : public PostWalker<SubType, VisitorType> {
     self->tryStack.pop_back();
   }
 
+  static void doEndResume(SubType* self, Expression** currp) {
+    auto* module = self->getModule();
+    if (!module || module->features.hasExceptionHandling()) {
+      // This resume might throw, so run the code to handle that.
+      doEndThrowingInst(self, currp);
+    }
+    auto handlerBlocks = BranchUtils::getUniqueTargets(*currp);
+    // Add branches to the targets.
+    for (auto target : handlerBlocks) {
+      self->branches[target].push_back(self->currBasicBlock);
+    }
+  }
+
   static bool isReturnCall(Expression* curr) {
     switch (curr->_id) {
       case Expression::Id::CallId:
@@ -521,6 +534,20 @@ struct CFGWalker : public PostWalker<SubType, VisitorType> {
         self->pushTask(SubType::doEndThrow, currp);
         break;
       }
+      case Expression::Id::ResumeId:
+      case Expression::Id::ResumeThrowId: {
+        self->pushTask(SubType::doEndResume, currp);
+        break;
+      }
+      case Expression::Id::SuspendId:
+      case Expression::Id::StackSwitchId: {
+        auto* module = self->getModule();
+        if (!module || module->features.hasExceptionHandling()) {
+          // This might throw, so run the code to handle that.
+          self->pushTask(SubType::doEndCall, currp);
+        }
+        break;
+      }
       default: {
         if (Properties::isBranch(curr)) {
           self->pushTask(SubType::doEndBranch, currp);
diff --git a/test/lit/passes/O3_stack-switching.wast b/test/lit/passes/O3_stack-switching.wast
index 98bc066ffa2..d928f6614b2 100644
--- a/test/lit/passes/O3_stack-switching.wast
+++ b/test/lit/passes/O3_stack-switching.wast
@@ -25,64 +25,69 @@
  (tag $effect (param (ref eq)) (result (ref eq) (ref eq)))
 
  ;; CHECK:      (func $resume (type $10) (param $0 (ref $fiber)) (param $1 (ref $closure)) (param $2 (ref eq)) (result (ref eq))
- ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:  (local $3 (tuple (ref eq) (ref $cont)))
+ ;; CHECK-NEXT:  (local $4 (ref $handlers))
+ ;; CHECK-NEXT:  (local $5 (ref $closure_2))
+ ;; CHECK-NEXT:  (return_call_ref $function_1
  ;; CHECK-NEXT:   (block $handle_exception (result (ref eq))
- ;; CHECK-NEXT:    (tuple.drop 2
- ;; CHECK-NEXT:     (block $handle_effect (type $3) (result (ref eq) (ref $cont))
- ;; CHECK-NEXT:      (return_call_ref $function_1
- ;; CHECK-NEXT:       (try_table (result (ref eq)) (catch $exception $handle_exception)
- ;; CHECK-NEXT:        (resume $cont (on $effect $handle_effect)
- ;; CHECK-NEXT:         (local.get $1)
- ;; CHECK-NEXT:         (local.get $2)
- ;; CHECK-NEXT:         (struct.get $fiber $cont
- ;; CHECK-NEXT:          (local.get $0)
+ ;; CHECK-NEXT:    (return_call_ref $function_2
+ ;; CHECK-NEXT:     (tuple.extract 2 0
+ ;; CHECK-NEXT:      (local.tee $3
+ ;; CHECK-NEXT:       (block $handle_effect (type $7) (result (ref eq) (ref $cont))
+ ;; CHECK-NEXT:        (return_call_ref $function_1
+ ;; CHECK-NEXT:         (try_table (result (ref eq)) (catch $exception $handle_exception)
+ ;; CHECK-NEXT:          (resume $cont (on $effect $handle_effect)
+ ;; CHECK-NEXT:           (local.get $1)
+ ;; CHECK-NEXT:           (local.get $2)
+ ;; CHECK-NEXT:           (struct.get $fiber $cont
+ ;; CHECK-NEXT:            (local.get $0)
+ ;; CHECK-NEXT:           )
+ ;; CHECK-NEXT:          )
  ;; CHECK-NEXT:         )
- ;; CHECK-NEXT:        )
- ;; CHECK-NEXT:       )
- ;; CHECK-NEXT:       (local.tee $1
- ;; CHECK-NEXT:        (struct.get $handlers $value
- ;; CHECK-NEXT:         (struct.get $fiber $handlers
- ;; CHECK-NEXT:          (local.get $0)
+ ;; CHECK-NEXT:         (local.tee $1
+ ;; CHECK-NEXT:          (struct.get $handlers $value
+ ;; CHECK-NEXT:           (struct.get $fiber $handlers
+ ;; CHECK-NEXT:            (local.get $0)
+ ;; CHECK-NEXT:           )
+ ;; CHECK-NEXT:          )
+ ;; CHECK-NEXT:         )
+ ;; CHECK-NEXT:         (struct.get $closure 0
+ ;; CHECK-NEXT:          (local.get $1)
  ;; CHECK-NEXT:         )
  ;; CHECK-NEXT:        )
  ;; CHECK-NEXT:       )
- ;; CHECK-NEXT:       (struct.get $closure 0
- ;; CHECK-NEXT:        (local.get $1)
- ;; CHECK-NEXT:       )
  ;; CHECK-NEXT:      )
  ;; CHECK-NEXT:     )
- ;; CHECK-NEXT:    )
- ;; CHECK-NEXT:    (drop
- ;; CHECK-NEXT:     (unreachable)
- ;; CHECK-NEXT:    )
- ;; CHECK-NEXT:   )
- ;; CHECK-NEXT:  )
- ;; CHECK-NEXT:  (block ;; (replaces unreachable CallRef we can't emit)
- ;; CHECK-NEXT:   (drop
- ;; CHECK-NEXT:    (unreachable)
- ;; CHECK-NEXT:   )
- ;; CHECK-NEXT:   (drop
- ;; CHECK-NEXT:    (block ;; (replaces unreachable StructGet we can't emit)
- ;; CHECK-NEXT:     (drop
- ;; CHECK-NEXT:      (block ;; (replaces unreachable StructGet we can't emit)
- ;; CHECK-NEXT:       (drop
- ;; CHECK-NEXT:        (unreachable)
+ ;; CHECK-NEXT:     (struct.new $fiber
+ ;; CHECK-NEXT:      (local.tee $4
+ ;; CHECK-NEXT:       (struct.get $fiber $handlers
+ ;; CHECK-NEXT:        (local.get $0)
  ;; CHECK-NEXT:       )
- ;; CHECK-NEXT:       (unreachable)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:      (tuple.extract 2 1
+ ;; CHECK-NEXT:       (local.get $3)
  ;; CHECK-NEXT:      )
  ;; CHECK-NEXT:     )
- ;; CHECK-NEXT:     (unreachable)
+ ;; CHECK-NEXT:     (local.tee $5
+ ;; CHECK-NEXT:      (struct.get $handlers $effect
+ ;; CHECK-NEXT:       (local.get $4)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (struct.get $closure_2 0
+ ;; CHECK-NEXT:      (local.get $5)
+ ;; CHECK-NEXT:     )
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
- ;; CHECK-NEXT:   (drop
- ;; CHECK-NEXT:    (block ;; (replaces unreachable StructGet we can't emit)
- ;; CHECK-NEXT:     (drop
- ;; CHECK-NEXT:      (unreachable)
+ ;; CHECK-NEXT:   (local.tee $1
+ ;; CHECK-NEXT:    (struct.get $handlers $exn
+ ;; CHECK-NEXT:     (struct.get $fiber $handlers
+ ;; CHECK-NEXT:      (local.get $0)
  ;; CHECK-NEXT:     )
- ;; CHECK-NEXT:     (unreachable)
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
- ;; CHECK-NEXT:   (unreachable)
+ ;; CHECK-NEXT:   (struct.get $closure 0
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
  (func $resume (export "resume") (param $fiber (ref $fiber)) (param $f (ref $closure)) (param $v (ref eq)) (result (ref eq))
diff --git a/test/lit/passes/coalesce-locals-stack-switching.wast b/test/lit/passes/coalesce-locals-stack-switching.wast
index 679344cae02..c0733c9c087 100644
--- a/test/lit/passes/coalesce-locals-stack-switching.wast
+++ b/test/lit/passes/coalesce-locals-stack-switching.wast
@@ -18,10 +18,10 @@
 
  ;; CHECK:      (func $resume (type $2) (param $0 (ref $cont)) (result i32)
  ;; CHECK-NEXT:  (local $1 i32)
- ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:  (local.set $1
  ;; CHECK-NEXT:   (block $handle_exn (result i32)
  ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
- ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:     (local.set $0
  ;; CHECK-NEXT:      (block $handle_effect (result (ref $cont))
  ;; CHECK-NEXT:       (resume $cont (on $tag $handle_effect)
  ;; CHECK-NEXT:        (local.get $0)
@@ -32,9 +32,7 @@
  ;; CHECK-NEXT:      )
  ;; CHECK-NEXT:     )
  ;; CHECK-NEXT:     (return_call $resume
- ;; CHECK-NEXT:      (block (result (ref $cont))
- ;; CHECK-NEXT:       (unreachable)
- ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:      (local.get $0)
  ;; CHECK-NEXT:     )
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:    (return
@@ -42,7 +40,7 @@
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
- ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT:  (local.get $1)
  ;; CHECK-NEXT: )
  (func $resume (param $c (ref $cont)) (result i32)
   (local $x i32)
@@ -77,10 +75,10 @@
 
  ;; CHECK:      (func $resume_throw (type $2) (param $0 (ref $cont)) (result i32)
  ;; CHECK-NEXT:  (local $1 i32)
- ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:  (local.set $1
  ;; CHECK-NEXT:   (block $handle_exn (result i32)
  ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
- ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:     (local.set $0
  ;; CHECK-NEXT:      (block $handle_effect (result (ref $cont))
  ;; CHECK-NEXT:       (resume_throw $cont $tag (on $tag $handle_effect)
  ;; CHECK-NEXT:        (local.get $0)
@@ -91,9 +89,7 @@
  ;; CHECK-NEXT:      )
  ;; CHECK-NEXT:     )
  ;; CHECK-NEXT:     (return_call $resume
- ;; CHECK-NEXT:      (block (result (ref $cont))
- ;; CHECK-NEXT:       (unreachable)
- ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:      (local.get $0)
  ;; CHECK-NEXT:     )
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:    (return
@@ -101,7 +97,7 @@
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
- ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT:  (local.get $1)
  ;; CHECK-NEXT: )
  (func $resume_throw (param $c (ref $cont)) (result i32)
   (local $x i32)
@@ -136,7 +132,7 @@
 
  ;; CHECK:      (func $suspend (type $6) (result i32)
  ;; CHECK-NEXT:  (local $0 i32)
- ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:  (local.set $0
  ;; CHECK-NEXT:   (block $handle_exn (result i32)
  ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
  ;; CHECK-NEXT:     (suspend $tag)
@@ -146,7 +142,7 @@
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
- ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT:  (local.get $0)
  ;; CHECK-NEXT: )
  (func $suspend (result i32)
   (local $x i32)
@@ -167,7 +163,7 @@
 
  ;; CHECK:      (func $switch (type $7) (param $0 (ref $cont_2)) (result i32)
  ;; CHECK-NEXT:  (local $1 i32)
- ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:  (local.set $1
  ;; CHECK-NEXT:   (block $handle_exn (result i32)
  ;; CHECK-NEXT:    (try_table (catch $exn $handle_exn)
  ;; CHECK-NEXT:     (switch $cont_2 $tag
@@ -179,7 +175,7 @@
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
- ;; CHECK-NEXT:  (i32.const 0)
+ ;; CHECK-NEXT:  (local.get $1)
  ;; CHECK-NEXT: )
  (func $switch (param $c (ref $cont_2)) (result i32)
   (local $x i32)

From 2a3d6c0b578158a34dd438603d3076f5220b1511 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Tue, 11 Feb 2025 13:49:14 +0100
Subject: [PATCH 8/9] Update scripts/test/fuzzing.py

Co-authored-by: Max Graey <maxgraey@gmail.com>
---
 scripts/test/fuzzing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/test/fuzzing.py b/scripts/test/fuzzing.py
index 8cbf3ebfdfa..7978ac02621 100644
--- a/scripts/test/fuzzing.py
+++ b/scripts/test/fuzzing.py
@@ -98,7 +98,7 @@
     'stack_switching_resume.wast',
     'stack_switching_resume_throw.wast',
     'stack_switching_switch.wast',
-    'stack_switching_switch.wast_2',
+    'stack_switching_switch_2.wast',
     'O3_stack-switching.wast',
     'coalesce-locals-stack-switching.wast',
     'dce-stack-switching.wast',

From 4d606bb7382c05b4415fd87696528ecf4a045c44 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Tue, 25 Feb 2025 11:23:07 +0100
Subject: [PATCH 9/9] Update GlobalTypeRewriter and SubTypes

---
 src/ir/subtypes.h        |  3 ++-
 src/ir/type-updating.cpp |  8 ++++++--
 src/ir/type-updating.h   | 13 +++++++++++++
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/ir/subtypes.h b/src/ir/subtypes.h
index c69250043b8..5c654ceb7be 100644
--- a/src/ir/subtypes.h
+++ b/src/ir/subtypes.h
@@ -126,7 +126,8 @@ struct SubTypes {
           basic = HeapTypes::array.getBasic(share);
           break;
         case HeapTypeKind::Cont:
-          WASM_UNREACHABLE("TODO: cont");
+          basic = HeapTypes::cont.getBasic(share);
+          break;
         case HeapTypeKind::Basic:
           WASM_UNREACHABLE("unexpected kind");
       }
diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index b90d8eb8790..2403f01a700 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -150,8 +150,12 @@ GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes(
         typeBuilder[i] = newArray;
         break;
       }
-      case HeapTypeKind::Cont:
-        WASM_UNREACHABLE("TODO: cont");
+      case HeapTypeKind::Cont: {
+        auto newCont = HeapType(typeBuilder[i]).getContinuation();
+        modifyContinuation(type, newCont);
+        typeBuilder[i] = newCont;
+        break;
+      }
       case HeapTypeKind::Basic:
         WASM_UNREACHABLE("unexpected kind");
     }
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index fe1cd2806aa..1ccc33c7c8d 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -382,6 +382,7 @@ class GlobalTypeRewriter {
   // used to define the new type in the TypeBuilder.
   virtual void modifyStruct(HeapType oldType, Struct& struct_) {}
   virtual void modifyArray(HeapType oldType, Array& array) {}
+  virtual void modifyContinuation(HeapType oldType, Continuation& sig) {}
   virtual void modifySignature(HeapType oldType, Signature& sig) {}
 
   // This additional hook is called after modify* and other operations, and
@@ -502,6 +503,14 @@ class TypeMapper : public GlobalTypeRewriter {
     return getTempType(type);
   }
 
+  HeapType getNewHeapType(HeapType type) {
+    auto iter = mapping.find(type);
+    if (iter != mapping.end()) {
+      return iter->second;
+    }
+    return type;
+  }
+
   void modifyStruct(HeapType oldType, Struct& struct_) override {
     auto& oldFields = oldType.getStruct().fields;
     for (Index i = 0; i < oldFields.size(); i++) {
@@ -513,6 +522,10 @@ class TypeMapper : public GlobalTypeRewriter {
   void modifyArray(HeapType oldType, Array& array) override {
     array.element.type = getNewType(oldType.getArray().element.type);
   }
+  void modifyContinuation(HeapType oldType,
+                          Continuation& continuation) override {
+    continuation.type = getNewHeapType(oldType.getContinuation().type);
+  }
   void modifySignature(HeapType oldSignatureType, Signature& sig) override {
     auto getUpdatedTypeList = [&](Type type) {
       std::vector<Type> vec;