fix[venom]: fix branch eliminator cases in sccp

tests/unit/compiler/venom/test_sccp.py

@@ -1,5 +1,8 @@
+import pytest
+
+from vyper.exceptions import StaticAssertionException
 from vyper.venom.analysis.analysis import IRAnalysesCache
-from vyper.venom.basicblock import IRBasicBlock, IRLabel, IRVariable
+from vyper.venom.basicblock import IRBasicBlock, IRLabel, IRLiteral, IRVariable
 from vyper.venom.context import IRContext
 from vyper.venom.passes.make_ssa import MakeSSA
 from vyper.venom.passes.sccp import SCCP
@@ -28,6 +31,73 @@ def test_simple_case():
  assert sccp.lattice[IRVariable("%4")].value == 96
 
 
+def test_branch_eliminator_simple():
+ ctx = IRContext()
+ fn = ctx.create_function("_global")
+
+ bb1 = fn.get_basic_block()
+
+ br1 = IRBasicBlock(IRLabel("then"), fn)
+ br1.append_instruction("stop")
+ br2 = IRBasicBlock(IRLabel("else"), fn)
+ br2.append_instruction("jmp", IRLabel("foo"))
+
+ fn.append_basic_block(br1)
+ fn.append_basic_block(br2)
+
+ bb1.append_instruction("jnz", IRLiteral(1), br1.label, br2.label)
+
+ bb2 = IRBasicBlock(IRLabel("foo"), fn)
+ bb2.append_instruction("jnz", IRLiteral(0), br1.label, br2.label)
+ fn.append_basic_block(bb2)
+
+ ac = IRAnalysesCache(fn)
+ MakeSSA(ac, fn).run_pass()
+ sccp = SCCP(ac, fn)
+ sccp.run_pass()
+
+ assert bb1.instructions[-1].opcode == "jmp"
+ assert bb1.instructions[-1].operands == [br1.label]
+ assert bb2.instructions[-1].opcode == "jmp"
+ assert bb2.instructions[-1].operands == [br2.label]
+
+
+def test_assert_elimination():
+ ctx = IRContext()
+ fn = ctx.create_function("_global")
+
+ bb = fn.get_basic_block()
+
+ bb.append_instruction("assert", IRLiteral(1))
+ bb.append_instruction("assert_unreachable", IRLiteral(1))
+ bb.append_instruction("stop")
+
+ ac = IRAnalysesCache(fn)
+ MakeSSA(ac, fn).run_pass()
+ sccp = SCCP(ac, fn)
+ sccp.run_pass()
+
+ for inst in bb.instructions[:-1]:
+ assert inst.opcode == "nop"
+
+
+@pytest.mark.parametrize("asserter", ("assert", "assert_unreachable"))
+def test_assert_false(asserter):
+ ctx = IRContext()
+ fn = ctx.create_function("_global")
+
+ bb = fn.get_basic_block()
+
+ bb.append_instruction(asserter, IRLiteral(0))
+ bb.append_instruction("stop")
+
+ ac = IRAnalysesCache(fn)
+ MakeSSA(ac, fn).run_pass()
+ sccp = SCCP(ac, fn)
+ with pytest.raises(StaticAssertionException):
+ sccp.run_pass()
+
+
 def test_cont_jump_case():
  ctx = IRContext()
  fn = ctx.create_function("_global")

vyper/venom/passes/sccp/sccp.py

@@ -153,6 +153,7 @@ def _visit_phi(self, inst: IRInstruction):
  in_vars.append(self.lattice[var])
  value = reduce(_meet, in_vars, LatticeEnum.TOP) # type: ignore
  assert inst.output in self.lattice, "Got undefined var for phi"
+
  if value != self.lattice[inst.output]:
  self.lattice[inst.output] = value
  self._add_ssa_work_items(inst)
@@ -169,7 +170,8 @@ def _visit_expr(self, inst: IRInstruction):
  target = self.fn.get_basic_block(inst.operands[0].value)
  self.work_list.append(FlowWorkItem(inst.parent, target))
  elif opcode == "jnz":
- lat = self.lattice[inst.operands[0]]
+ lat = _get_lattice(self.lattice, inst.operands[0])
+
  assert lat != LatticeEnum.TOP, f"Got undefined var at jmp at {inst.parent}"
  if lat == LatticeEnum.BOTTOM:
  for out_bb in inst.parent.cfg_out:
@@ -276,7 +278,8 @@ def _replace_constants(self, inst: IRInstruction, lattice: Lattice):
  case of jumps and asserts as needed.
  """
  if inst.opcode == "jnz":
- lat = lattice[inst.operands[0]]
+ lat = _get_lattice(lattice, inst.operands[0])
+
  if isinstance(lat, IRLiteral):
  if lat.value == 0:
  target = inst.operands[2]
@@ -285,8 +288,10 @@ def _replace_constants(self, inst: IRInstruction, lattice: Lattice):
  inst.opcode = "jmp"
  inst.operands = [target]
  self.cfg_dirty = True
- elif inst.opcode == "assert":
- lat = lattice[inst.operands[0]]
+
+ elif inst.opcode in ("assert", "assert_unreachable"):
+ lat = _get_lattice(lattice, inst.operands[0])
+
  if isinstance(lat, IRLiteral):
  if lat.value > 0:
  inst.opcode = "nop"
@@ -334,3 +339,10 @@ def _meet(x: LatticeItem, y: LatticeItem) -> LatticeItem:
  if y == LatticeEnum.TOP or x == y:
  return x
  return LatticeEnum.BOTTOM
+
+
+def _get_lattice(lattice: Lattice, x: IROperand) -> LatticeItem:
+ if isinstance(x, IRLiteral):
+ return x
+
+ return lattice[x]