Repository: tetratelabs/wazero
Pull request #1770 (2023-10-10 01:20:48)

wazevo(arm64): fixes 32bit shifted operand with const amount
detected by the fuzzer

#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
index 246a4a88..2e1fa4a0 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
@@ -263,10 +263,10 @@ func (m *machine) getOperand_SR_NR(def *backend.SSAValueDefinition, mode extMode
 		amountDef := m.compiler.ValueDefinition(amountVal)
 		if amountDef.IsFromInstr() && amountDef.Instr.Constant() {
 			// If that is the case, we can use the shifted register operand (SR).
-			c := amountDef.Instr.ConstantVal() & 63 // Clears the unnecessary bits.
+			c := byte(amountDef.Instr.ConstantVal()) & (targetVal.Type().Bits() - 1) // Clears the unnecessary bits.
 			m.compiler.MarkLowered(def.Instr)
 			m.compiler.MarkLowered(amountDef.Instr)
-			return operandSR(targetVReg, byte(c), shiftOpLSL)
+			return operandSR(targetVReg, c, shiftOpLSL)
 		}
 	}
 	return m.getOperand_NR(def, mode)


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
index 03306820..4d25ca5a 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
@@ -199,6 +199,49 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 			setup: ishlWithConstAmount,
 			exp:   operandSR(regalloc.VReg(1234), 14, shiftOpLSL),
 		},
+		{
+			name: "ishl with const amount with i32",
+			setup: func(
+				ctx *mockCompiler, builder ssa.Builder, m *machine,
+			) (def *backend.SSAValueDefinition, mode extMode, verify func(t *testing.T)) {
+				blk := builder.CurrentBlock()
+				// (p1+p2) << amount
+				p1 := blk.AddParam(builder, ssa.TypeI32)
+				p2 := blk.AddParam(builder, ssa.TypeI32)
+				add := builder.AllocateInstruction()
+				add.AsIadd(p1, p2)
+				builder.InsertInstruction(add)
+				addResult := add.Return()
+
+				amount := builder.AllocateInstruction()
+				amount.AsIconst32(45) // should be taken modulo by 31.
+				builder.InsertInstruction(amount)
+
+				amountVal := amount.Return()
+
+				ishl := builder.AllocateInstruction()
+				ishl.AsIshl(addResult, amountVal)
+				builder.InsertInstruction(ishl)
+
+				ctx.definitions[p1] = &backend.SSAValueDefinition{BlkParamVReg: regalloc.VReg(1), BlockParamValue: p1}
+				ctx.definitions[p2] = &backend.SSAValueDefinition{BlkParamVReg: regalloc.VReg(2), BlockParamValue: p2}
+				ctx.definitions[addResult] = &backend.SSAValueDefinition{Instr: add, N: 0}
+				ctx.definitions[amountVal] = &backend.SSAValueDefinition{Instr: amount, N: 0}
+
+				ctx.vRegMap[addResult] = regalloc.VReg(1234)
+				ctx.vRegMap[ishl.Return()] = regalloc.VReg(10)
+				def = &backend.SSAValueDefinition{Instr: ishl, N: 0}
+				mode = extModeNone
+				verify = func(t *testing.T) {
+					_, ok := ctx.lowered[ishl]
+					require.True(t, ok)
+					_, ok = ctx.lowered[amount]
+					require.True(t, ok)
+				}
+				return
+			},
+			exp: operandSR(regalloc.VReg(1234), 13, shiftOpLSL),
+		},
 		{
 			name: "ishl with const amount with const shift target",
 			setup: func(