winch: Rework `br_table` jumps (#7628)

This commit reworks the `br_table` logic so that it correctly handles
all the jumps involved to each of the targets.

Even though it is safe to use the default branch for type information,
it is not safe to use it to derive the base stack pointer and base value
stack length. This change ensures that each target offset is taken into
account to balance the value stack prior to each jump.

tests/misc_testsuite/winch/br_table.wast

@@ -10,6 +10,25 @@
   (func (export "type-i64")
     (block (drop (i64.ctz (br_table 0 0 (i32.const 0)))))
   )
+  (func (export "type-f32")
+    (block (drop (f32.neg (br_table 0 0 (i32.const 0)))))
+  )
+  (func (export "type-f64")
+    (block (drop (f64.neg (br_table 0 0 (i32.const 0)))))
+  )
+
+  (func (export "type-i32-value") (result i32)
+    (block (result i32) (i32.ctz (br_table 0 0 (i32.const 1) (i32.const 0))))
+  )
+  (func (export "type-i64-value") (result i64)
+    (block (result i64) (i64.ctz (br_table 0 0 (i64.const 2) (i32.const 0))))
+  )
+  (func (export "type-f32-value") (result f32)
+    (block (result f32) (f32.neg (br_table 0 0 (f32.const 3) (i32.const 0))))
+  )
+  (func (export "type-f64-value") (result f64)
+    (block (result f64) (f64.neg (br_table 0 0 (f64.const 4) (i32.const 0))))
+  )
 
   (func (export "empty") (param i32) (result i32)
     (block (br_table 0 (local.get 0)) (return (i32.const 21)))
@@ -1123,6 +1142,12 @@
 
 (assert_return (invoke "type-i32"))
 (assert_return (invoke "type-i64"))
+(assert_return (invoke "type-f32"))
+(assert_return (invoke "type-f64"))
+(assert_return (invoke "type-i32-value") (i32.const 1))
+(assert_return (invoke "type-i64-value") (i64.const 2))
+(assert_return (invoke "type-f32-value") (f32.const 3))
+(assert_return (invoke "type-f64-value") (f64.const 4))
 
 (assert_return (invoke "empty" (i32.const 0)) (i32.const 22))
 (assert_return (invoke "empty" (i32.const 1)) (i32.const 22))

tests/misc_testsuite/winch/misc.wast

@@ -0,0 +1,14 @@
+;; Additional run tests for Winch not covered in the official spec test suite.
+
+(module
+  (func (export "br-table-ensure-sp") (result i32)
+    (block (result i32)
+       (i32.const 0)
+    )
+    (i32.const 0)
+    (i32.const 0)
+    (br_table 0)
+  )
+)
+
+(assert_return (invoke "br-table-ensure-sp") (i32.const 0))

winch/codegen/src/visitor.rs

@@ -1246,6 +1246,11 @@ where
         };
 
         self.masm.jmp_table(&labels, index.into(), tmp);
+        // Save the original stack pointer offset; we will reset the stack
+        // pointer to this offset after jumping to each of the targets. Each
+        // jump might adjust the stack according to the base offset of the
+        // target.
+        let current_sp = self.masm.sp_offset();
 
         for (t, l) in targets
             .targets()
@@ -1255,14 +1260,26 @@ where
         {
             let control_index = control_index(t.unwrap(), self.control_frames.len());
             let frame = &mut self.control_frames[control_index];
+            // Reset the stack pointer to its original offset. This is needed
+            // because each jump will potentially adjust the stack pointer
+            // according to the base offset of the target.
+            self.masm.reset_stack_pointer(current_sp);
 
             // NB: We don't perform any result handling as it was
             // already taken care of above before jumping to the
             // jump table.
             self.masm.bind(*l);
+            // Ensure that the stack pointer is correctly positioned before
+            // jumping to the jump table code.
+            let (_, offset) = frame.base_stack_len_and_sp();
+            self.masm.ensure_sp_for_jump(offset);
             self.masm.jmp(*frame.label());
             frame.set_as_target();
         }
+        // Finally reset the stack pointer to the original location.
+        // The reachability analysis, will ensure it's correctly located
+        // once reachability is restored.
+        self.masm.reset_stack_pointer(current_sp);
         self.context.reachable = false;
         self.context.free_reg(index.reg);
         self.context.free_reg(tmp);

winch/filetests/filetests/x64/br_table/ensure_sp_state.wat

@@ -0,0 +1,34 @@
+;;! target = "x86_64"
+
+(module
+  (func (export "") (result i32)
+    block (result i32)
+       i32.const 0
+    end
+    i32.const 0
+    i32.const 0
+    br_table 0
+  )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 b800000000           	mov	eax, 0
+;;   11:	 4883ec04             	sub	rsp, 4
+;;   15:	 890424               	mov	dword ptr [rsp], eax
+;;   18:	 b900000000           	mov	ecx, 0
+;;   1d:	 b800000000           	mov	eax, 0
+;;   22:	 ba00000000           	mov	edx, 0
+;;   27:	 39ca                 	cmp	edx, ecx
+;;   29:	 0f42ca               	cmovb	ecx, edx
+;;   2c:	 4c8d1d0a000000       	lea	r11, [rip + 0xa]
+;;   33:	 4963148b             	movsxd	rdx, dword ptr [r11 + rcx*4]
+;;   37:	 4901d3               	add	r11, rdx
+;;   3a:	 41ffe3               	jmp	r11
+;;   3d:	 0400                 	add	al, 0
+;;   3f:	 0000                 	add	byte ptr [rax], al
+;;   41:	 4883c404             	add	rsp, 4
+;;   45:	 4883c408             	add	rsp, 8
+;;   49:	 5d                   	pop	rbp
+;;   4a:	 c3                   	ret	

winch/filetests/filetests/x64/br_table/stack_handling.wat

@@ -26,13 +26,16 @@
 ;;   33:	 49630c83             	movsxd	rcx, dword ptr [r11 + rax*4]
 ;;   37:	 4901cb               	add	r11, rcx
 ;;   3a:	 41ffe3               	jmp	r11
-;;   3d:	 0c00                 	or	al, 0
+;;   3d:	 1a00                 	sbb	al, byte ptr [rax]
 ;;   3f:	 0000                 	add	byte ptr [rax], al
-;;   41:	 1000                 	adc	byte ptr [rax], al
+;;   41:	 1100                 	adc	dword ptr [rax], eax
 ;;   43:	 0000                 	add	byte ptr [rax], al
-;;   45:	 0c00                 	or	al, 0
+;;   45:	 1a00                 	sbb	al, byte ptr [rax]
 ;;   47:	 0000                 	add	byte ptr [rax], al
-;;   49:	 4883c404             	add	rsp, 4
-;;   4d:	 4883c410             	add	rsp, 0x10
-;;   51:	 5d                   	pop	rbp
-;;   52:	 c3                   	ret	
+;;   49:	 e909000000           	jmp	0x57
+;;   4e:	 4883c404             	add	rsp, 4
+;;   52:	 e904000000           	jmp	0x5b
+;;   57:	 4883c404             	add	rsp, 4
+;;   5b:	 4883c410             	add	rsp, 0x10
+;;   5f:	 5d                   	pop	rbp
+;;   60:	 c3                   	ret