wazevo: passes multiple return spectests (#1701)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/compiler.go

@@ -237,6 +237,7 @@ func (c *compiler) assignVirtualRegisters() {
 		// Assigns each value to a virtual register produced by instructions.
 		for cur := blk.Root(); cur != nil; cur = cur.Next() {
 			r, rs := cur.Returns()
+			var N int
 			if r.Valid() {
 				id := r.ID()
 				ssaTyp := r.Type()
@@ -248,18 +249,20 @@ func (c *compiler) assignVirtualRegisters() {
 					RefCount: refCounts[id],
 				}
 				c.ssaTypeOfVRegID[vReg.ID()] = ssaTyp
+				N++
 			}
-			for i, r := range rs {
+			for _, r := range rs {
 				id := r.ID()
 				ssaTyp := r.Type()
 				vReg := c.AllocateVReg(regalloc.RegTypeOf(ssaTyp))
 				c.ssaValueToVRegs[id] = vReg
 				c.ssaValueDefinitions[id] = SSAValueDefinition{
 					Instr:    cur,
-					N:        i,
+					N:        N,
 					RefCount: refCounts[id],
 				}
 				c.ssaTypeOfVRegID[vReg.ID()] = ssaTyp
+				N++
 			}
 		}
 	}

internal/engine/wazevo/e2e_test.go

@@ -128,12 +128,25 @@ func TestSpectestV2(t *testing.T) {
 	for _, tc := range []struct {
 		name string
 	}{
-		{"conversions"}, // includes non-trapping conversions.
+		{"block"},
+		{"br"},
+		{"call"},
+		{"call_indirect"},
+		{"conversions"},
+		{"if"},
+		{"loop"},
 		{"simd_const"},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			spectest.RunCase(t, v2.Testcases, tc.name, context.Background(), config,
-				-1, 0, math.MaxInt)
+			t.Run("normal", func(t *testing.T) {
+				spectest.RunCase(t, v2.Testcases, tc.name, context.Background(), config,
+					-1, 0, math.MaxInt)
+			})
+			t.Run("reg high pressure", func(t *testing.T) {
+				ctx := wazevoapi.EnableHighRegisterPressure(context.Background())
+				spectest.RunCase(t, v2.Testcases, tc.name, ctx, config,
+					-1, 0, math.MaxInt)
+			})
 		})
 	}
 }

internal/engine/wazevo/frontend/lower.go

@@ -1032,7 +1032,7 @@ func (c *Compiler) lowerCurrentOpcode() {
 
 		var args []ssa.Value
 		if len(bt.Params) > 0 {
-			args = cloneValuesList(state.values[len(state.values)-1-len(bt.Params):])
+			args = cloneValuesList(state.values[len(state.values)-len(bt.Params):])
 		}
 
 		// Insert the conditional jump to the Else block.
@@ -1117,7 +1117,7 @@ func (c *Compiler) lowerCurrentOpcode() {
 			// If this is the end of Then block, we have to emit the empty Else block.
 			elseBlk := ctrl.blk
 			builder.SetCurrentBlock(elseBlk)
-			c.insertJumpToBlock(nil, followingBlk)
+			c.insertJumpToBlock(ctrl.clonedArgs, followingBlk)
 		}
 
 		builder.Seal(ctrl.followingBlock)