From e3e5b9942d373ca243bd3e85aade61e134d4a6e2 Mon Sep 17 00:00:00 2001
From: Edoardo Vacchi <evacchi@users.noreply.github.com>
Date: Wed, 25 Oct 2023 17:32:25 +0200
Subject: [PATCH] wazevo: fuzz, select does not return the correct value for
 vectors (#1814)

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>
---
 .../wazevo/backend/isa/arm64/lower_instr.go   | 38 ++++++++++++-------
 .../backend/isa/arm64/lower_instr_test.go     |  9 +++--
 .../fuzzcases/fuzzcases_test.go               |  2 +
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index 733eeb3a..253f0ad8 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -1873,21 +1873,31 @@ func (m *machine) lowerSelect(c, x, y, result ssa.Value) {
 }
 
 func (m *machine) lowerSelectVec(rc, rn, rm, rd operand) {
-	tmp := operandNR(m.compiler.AllocateVReg(ssa.TypeI64))
+	// Declare and insert the conditional branch here jump to label `ifNonZero` below:
+	// but we cannot forward reference the label.
+	cbr := m.allocateInstr()
+	m.insert(cbr)
 
-	// Sets all bits to 1 if rc is not zero.
-	alu := m.allocateInstr()
-	alu.asALU(aluOpSub, tmp, operandNR(xzrVReg), rc, true)
-	m.insert(alu)
+	// If rc is zero, mov rd, rm then jump to end.
+	mov0 := m.allocateInstr()
+	mov0.asFpuMov128(rd.nr(), rm.nr())
+	m.insert(mov0)
 
-	// Then move the bits to the result vector register.
-	dup := m.allocateInstr()
-	dup.asVecDup(rd, tmp, vecArrangement2D)
-	m.insert(dup)
+	// Declared and insert the non-conditional jump to label `end` below:
+	// again, we cannot forward reference the label.
+	br := m.allocateInstr()
+	m.insert(br)
 
-	// Now that `rd` has either all bits one or zero depending on `rc`,
-	// we can use bsl to select between `rn` and `rm`.
-	ins := m.allocateInstr()
-	ins.asVecRRR(vecOpBsl, rd, rn, rm, vecArrangement16B)
-	m.insert(ins)
+	// Create and insert the label, and update `cbr` to the real instruction.
+	ifNonZero := m.insertBrTargetLabel()
+	cbr.asCondBr(registerAsRegNotZeroCond(rc.nr()), ifNonZero, true)
+
+	// If rc is non-zero, set mov rd, rn.
+	mov := m.allocateInstr()
+	mov.asFpuMov128(rd.nr(), rn.nr())
+	m.insert(mov)
+
+	// Create and insert the label, and update `br` to the real instruction.
+	end := m.insertBrTargetLabel()
+	br.asBr(end)
 }
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
index dec0ba74..e0c32629 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
@@ -842,9 +842,12 @@ func TestMachine_lowerSelectVec(t *testing.T) {
 
 	m.lowerSelectVec(c, rn, rm, rd)
 	require.Equal(t, `
-sub x5?, xzr, x1?
-dup v4?.2d, x5?
-bsl v4?.16b, v2?.16b, v3?.16b
+cbnz x1?, L1
+mov v4?.16b, v3?.16b
+b L2
+L1:
+mov v4?.16b, v2?.16b
+L2:
 `, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
 }
 
diff --git a/internal/integration_test/fuzzcases/fuzzcases_test.go b/internal/integration_test/fuzzcases/fuzzcases_test.go
index 434a888c..bbd7c312 100644
--- a/internal/integration_test/fuzzcases/fuzzcases_test.go
+++ b/internal/integration_test/fuzzcases/fuzzcases_test.go
@@ -94,6 +94,8 @@ func Test696(t *testing.T) {
 		}{
 			{fnName: "select", in: 1, exp: [2]uint64{0xffffffffffffffff, 0xeeeeeeeeeeeeeeee}},
 			{fnName: "select", in: 0, exp: [2]uint64{0x1111111111111111, 0x2222222222222222}},
+			{fnName: "select", in: 0xffffff, exp: [2]uint64{0xffffffffffffffff, 0xeeeeeeeeeeeeeeee}},
+			{fnName: "select", in: 0x000000, exp: [2]uint64{0x1111111111111111, 0x2222222222222222}},
 			{fnName: "typed select", in: 1, exp: [2]uint64{0xffffffffffffffff, 0xeeeeeeeeeeeeeeee}},
 			{fnName: "typed select", in: 0, exp: [2]uint64{0x1111111111111111, 0x2222222222222222}},
 		} {