wazevo(frontend): BCE should not propagate the memory address to unsealed blocks (#2164)

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

internal/engine/wazevo/frontend/frontend.go

@@ -524,7 +524,13 @@ func (c *Compiler) initializeCurrentBlockKnownBounds() {
 	case 1:
 		pred := currentBlk.Pred(0).ID()
 		for _, kb := range c.getKnownSafeBoundsAtTheEndOfBlocks(pred).View() {
-			c.recordKnownSafeBound(kb.id, kb.bound, kb.absoluteAddr)
+			// Unless the block is sealed, we cannot assume the absolute address is valid:
+			// later we might add another predecessor that has no visibility of that value.
+			addr := ssa.ValueInvalid
+			if currentBlk.Sealed() {
+				addr = kb.absoluteAddr
+			}
+			c.recordKnownSafeBound(kb.id, kb.bound, addr)
 		}
 	default:
 		c.pointers = c.pointers[:0]

internal/engine/wazevo/frontend/frontend_test.go

@@ -2891,6 +2891,105 @@ blk3: () <-- (blk1,blk2)
 	v21:i64 = Iadd v19, v20
 	v22:i32 = Load v21, 0x15
 	Jump blk_ret
+`,
+		},
+		{
+			name: "bounds check in loop with if-else",
+			m: &wasm.Module{
+				TypeSection:     []wasm.FunctionType{{Params: []wasm.ValueType{wasm.ValueTypeI32}}},
+				FunctionSection: []wasm.Index{0},
+				CodeSection: []wasm.Code{{
+					Body: []byte{
+						wasm.OpcodeLocalGet, 0,
+						wasm.OpcodeI32Load, 0x2, 0x10, // staticOffset=0x20
+						wasm.OpcodeDrop,
+
+						wasm.OpcodeLoop, 0x40, // blockSignature:vv.
+
+						/* */ wasm.OpcodeLocalGet, 0,
+						// This address must be always recomputed:
+						// v9 in `Load v9, 0x20` won't propagate from blk0 to blk1,
+						/* */ wasm.OpcodeI32Load, 0x2, 0x20, // staticOffset=0x30
+
+						/* */ wasm.OpcodeIf, 0x40, // blockSignature:vv.
+						/* */ /* */ wasm.OpcodeLoop, 0x40, // blockSignature:vv.
+						/* */ /* */ /* */ wasm.OpcodeLocalGet, 0,
+						/* */ /* */ /* */ wasm.OpcodeI32Load, 0x2, 0x50, // staticOffset=0x50
+						/* */ /* */ /* */ wasm.OpcodeBrIf, 2, // quit to the inner loop
+						/* */ /* */ wasm.OpcodeEnd,
+						/* */ wasm.OpcodeElse,
+						/* */ /* */ wasm.OpcodeLocalGet, 0,
+						/* */ /* */ wasm.OpcodeBrIf, 2, // quit to the outer loop
+						/* */ wasm.OpcodeEnd,
+						/* */ wasm.OpcodeBr, 1, // quit to the loop
+
+						wasm.OpcodeEnd, // end loop
+
+						wasm.OpcodeEnd,
+					},
+				}},
+				MemorySection: &wasm.Memory{Min: 1},
+			},
+			features: api.CoreFeaturesV2,
+			exp: `
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
+	v3:i64 = Iconst_64 0x14
+	v4:i64 = UExtend v2, 32->64
+	v5:i64 = Uload32 module_ctx, 0x10
+	v6:i64 = Iadd v4, v3
+	v7:i32 = Icmp lt_u, v5, v6
+	ExitIfTrue v7, exec_ctx, memory_out_of_bounds
+	v8:i64 = Load module_ctx, 0x8
+	v9:i64 = Iadd v8, v4
+	v10:i32 = Load v9, 0x10
+	Jump blk1
+
+blk1: () <-- (blk0,blk6)
+	v11:i64 = Iconst_64 0x24
+	v12:i64 = UExtend v2, 32->64
+	v13:i64 = Uload32 module_ctx, 0x10
+	v14:i64 = Iadd v12, v11
+	v15:i32 = Icmp lt_u, v13, v14
+	ExitIfTrue v15, exec_ctx, memory_out_of_bounds
+	v16:i64 = Load module_ctx, 0x8
+	v17:i64 = Iadd v16, v12
+	v18:i32 = Load v17, 0x20
+	Brz v18, blk4
+	Jump blk3
+
+blk2: ()
+
+blk3: () <-- (blk1)
+	Jump blk6
+
+blk4: () <-- (blk1)
+	Brnz v2, blk_ret
+	Jump blk9
+
+blk5: () <-- (blk7,blk9)
+	Jump blk_ret
+
+blk6: () <-- (blk3)
+	v19:i64 = Iconst_64 0x54
+	v20:i64 = UExtend v2, 32->64
+	v21:i64 = Uload32 module_ctx, 0x10
+	v22:i64 = Iadd v20, v19
+	v23:i32 = Icmp lt_u, v21, v22
+	ExitIfTrue v23, exec_ctx, memory_out_of_bounds
+	v24:i64 = Load module_ctx, 0x8
+	v25:i64 = Iadd v24, v20
+	v26:i32 = Load v25, 0x50
+	Brnz v26, blk1
+	Jump blk8
+
+blk7: () <-- (blk8)
+	Jump blk5
+
+blk8: () <-- (blk6)
+	Jump blk7
+
+blk9: () <-- (blk4)
+	Jump blk5
 `,
 		},
 	} {
@@ -3136,7 +3235,7 @@ func TestCompiler_finalizeKnownSafeBoundsAtTheEndOoBlock(t *testing.T) {
 }
 
 func TestCompiler_initializeCurrentBlockKnownBounds(t *testing.T) {
-	t.Run("single", func(t *testing.T) {
+	t.Run("single (sealed)", func(t *testing.T) {
 		c := NewFrontendCompiler(&wasm.Module{}, ssa.NewBuilder(), nil, false, false, false)
 		builder := c.ssaBuilder
 		child := builder.AllocateBasicBlock()
@@ -3150,6 +3249,8 @@ func TestCompiler_initializeCurrentBlockKnownBounds(t *testing.T) {
 			c.finalizeKnownSafeBoundsAtTheEndOfBlock(parent.ID())
 		}
 
+		// Seal the child: the absolute addresses will be propagated.
+		builder.Seal(child)
 		builder.SetCurrentBlock(child)
 		c.initializeCurrentBlockKnownBounds()
 		kb := c.getKnownSafeBound(1)
@@ -3165,6 +3266,36 @@ func TestCompiler_initializeCurrentBlockKnownBounds(t *testing.T) {
 		require.Equal(t, uint64(666), kb.bound)
 		require.Equal(t, ssa.Value(54321), kb.absoluteAddr)
 	})
+	t.Run("single (unsealed)", func(t *testing.T) {
+		c := NewFrontendCompiler(&wasm.Module{}, ssa.NewBuilder(), nil, false, false, false)
+		builder := c.ssaBuilder
+		child := builder.AllocateBasicBlock()
+		{
+			parent := builder.AllocateBasicBlock()
+			builder.SetCurrentBlock(parent)
+			c.recordKnownSafeBound(1, 99, 9999)
+			c.recordKnownSafeBound(2, 150, 9999)
+			c.recordKnownSafeBound(5, 666, 54321)
+			builder.AllocateInstruction().AsJump(ssa.ValuesNil, child).Insert(builder)
+			c.finalizeKnownSafeBoundsAtTheEndOfBlock(parent.ID())
+		}
+
+		// Do not seal the child: the absolute addresses will be recomputed.
+		builder.SetCurrentBlock(child)
+		c.initializeCurrentBlockKnownBounds()
+		kb := c.getKnownSafeBound(1)
+		require.True(t, kb.valid())
+		require.Equal(t, uint64(99), kb.bound)
+		require.NotEqual(t, ssa.Value(9999), kb.absoluteAddr)
+		kb = c.getKnownSafeBound(2)
+		require.True(t, kb.valid())
+		require.Equal(t, uint64(150), kb.bound)
+		require.NotEqual(t, ssa.Value(9999), kb.absoluteAddr)
+		kb = c.getKnownSafeBound(5)
+		require.True(t, kb.valid())
+		require.Equal(t, uint64(666), kb.bound)
+		require.NotEqual(t, ssa.Value(54321), kb.absoluteAddr)
+	})
 	t.Run("multiple predecessors", func(t *testing.T) {
 		c := NewFrontendCompiler(&wasm.Module{}, ssa.NewBuilder(), nil, false, false, false)
 		builder := c.ssaBuilder

internal/engine/wazevo/ssa/basic_block.go

@@ -55,6 +55,9 @@ type BasicBlock interface {
 	// Valid is true if this block is still valid even after optimizations.
 	Valid() bool
 
+	// Sealed is true if this block has been sealed.
+	Sealed() bool
+
 	// BeginPredIterator returns the first predecessor of this block.
 	BeginPredIterator() BasicBlock
 
@@ -212,6 +215,11 @@ func (bb *basicBlock) Valid() bool {
 	return !bb.invalid
 }
 
+// Sealed implements BasicBlock.Sealed.
+func (bb *basicBlock) Sealed() bool {
+	return bb.sealed
+}
+
 // InsertInstruction implements BasicBlock.InsertInstruction.
 func (bb *basicBlock) InsertInstruction(next *Instruction) {
 	current := bb.currentInstr