Fix bug in transform example

examples/transform_replace_add.py

@@ -1,4 +1,9 @@
-"""Parse LLVM IR and replace every integer add with an equivalent sub.
+"""Parse LLVM IR and replace plain integer adds with equivalent subs.
+
+The rewrite preserves both operands by converting ``x + y`` into
+``x - (0 - y)``. Adds with ``nuw``/``nsw`` flags are intentionally left alone
+because those flags carry poison semantics that this small example does not try
+to reproduce.
 
 This demonstrates the current transformation APIs:
 
@@ -19,9 +24,9 @@ import llvm
 
 
 INPUT_IR = """
-define i32 @add_one(i32 %x) {
+define i32 @add_values(i32 %x, i32 %y) {
 entry:
-  %sum = add i32 %x, 1
+  %sum = add i32 %x, %y
   ret i32 %sum
 }
 """
@@ -38,13 +43,19 @@ def transform(ir_text: str) -> str:
                     inst
                     for bb in func.basic_blocks
                     for inst in bb.instructions
-                    if inst.opcode == llvm.Opcode.Add
+                    if (
+                        inst.opcode == llvm.Opcode.Add
+                        and not inst.nsw
+                        and not inst.nuw
+                    )
                 ]
 
                 for inst in additions:
                     lhs, rhs = inst.operands
+                    base_name = inst.name or "add"
                     with inst.create_builder() as builder:
-                        replacement = builder.sub(lhs, rhs.type.constant(-1), inst.name + ".repl")
+                        neg_rhs = builder.neg(rhs, base_name + ".rhs.neg")
+                        replacement = builder.sub(lhs, neg_rhs, base_name + ".repl")
                     inst.replace_all_uses_with(replacement)
                     inst.erase_from_parent()
 

tests/test_examples.py

@@ -52,7 +52,39 @@ def test_readme_quick_start_snippet() -> None:
 
 def test_transform_replace_add_example() -> None:
     output = run_example("examples/transform_replace_add.py")
-    assert "define i32 @add_one(i32 %x)" in output
-    assert "%sum.repl = sub i32 %x, -1" in output
+    assert "define i32 @add_values(i32 %x, i32 %y)" in output
+    assert "%sum.rhs.neg = sub i32 0, %y" in output
+    assert "%sum.repl = sub i32 %x, %sum.rhs.neg" in output
     assert "ret i32 %sum.repl" in output
     assert " add i32 " not in output
+
+
+def test_transform_replace_add_preserves_rhs_operand() -> None:
+    from examples.transform_replace_add import transform
+
+    output = transform(
+        """
+        define i32 @f(i32 %x, i32 %y) {
+        entry:
+          %sum = add i32 %x, %y
+          ret i32 %sum
+        }
+        """
+    )
+    assert "%sum.rhs.neg = sub i32 0, %y" in output
+    assert "%sum.repl = sub i32 %x, %sum.rhs.neg" in output
+
+
+def test_transform_replace_add_leaves_no_wrap_adds_alone() -> None:
+    from examples.transform_replace_add import transform
+
+    output = transform(
+        """
+        define i32 @f(i32 %x, i32 %y) {
+        entry:
+          %sum = add nsw i32 %x, %y
+          ret i32 %sum
+        }
+        """
+    )
+    assert "%sum = add nsw i32 %x, %y" in output