Add insert_after to Builder class (#327)

swernli · web-flow · commit 4bd90b8d607b · 2025-04-04T23:18:23.000Z
This adds the ability to set the insertion point for a builder to after the given instruction.
diff --git a/pyqir/pyqir/_native.pyi b/pyqir/pyqir/_native.pyi
@@ -151,6 +151,14 @@ class Builder:
         """
         ...
 
+    def insert_after(self, instr: Instruction) -> None:
+        """
+        Tells the builder to insert subsequent instructions after the given instruction.
+
+        :param inst: The instruction to insert after.
+        """
+        ...
+
     def instr(self, instr: Instruction) -> None:
         """
         Inserts an instruction into the current block.
diff --git a/pyqir/src/builder.rs b/pyqir/src/builder.rs
@@ -77,6 +77,31 @@ impl Builder {
         Ok(())
     }
 
+    /// Tells this builder to insert subsequent instructions after the given instruction.
+    ///
+    /// :param Value instr: The instruction to insert after.
+    /// :rtype: None
+    #[pyo3(text_signature = "(instr)")]
+    fn insert_after(&mut self, py: Python, instr: &Value) -> PyResult<()> {
+        let owner = Owner::merge(py, [&self.owner, instr.owner()])?;
+        if *owner.context(py).borrow(py) != *self.owner.context(py).borrow(py) {
+            Err(PyValueError::new_err(
+                "Instruction is not from the same context as builder.",
+            ))?;
+        }
+
+        unsafe {
+            let next_instr = LLVMGetNextInstruction(instr.cast().as_ptr());
+            if next_instr.is_null() {
+                let block = LLVMGetInstructionParent(instr.cast().as_ptr());
+                LLVMPositionBuilderAtEnd(self.cast().as_ptr(), block);
+            } else {
+                LLVMPositionBuilderBefore(self.cast().as_ptr(), next_instr);
+            }
+        }
+        Ok(())
+    }
+
     /// Inserts the given instruction.
     fn instr(&mut self, py: Python, instr: &Value) -> PyResult<()> {
         let owner = Owner::merge(py, [&self.owner, instr.owner()])?;
diff --git a/pyqir/tests/test_transform.py b/pyqir/tests/test_transform.py
@@ -28,13 +28,14 @@ def test_reordering_instrs() -> None:
     x_instr = module.entry_block.instructions[0]
     y_instr = module.entry_block.instructions[1]
     z_instr = module.entry_block.instructions[2]
-    builder.insert_before(x_instr)
-    y_instr.remove()
-    builder.instr(y_instr)
     builder.insert_before(y_instr)
     z_instr.remove()
     builder.instr(z_instr)
-    builder.insert_at_end(module.entry_block)
+    builder.insert_after(y_instr)
+    x_instr.remove()
+    builder.instr(x_instr)
+    # Inserting after the last instruction will be equivalent to `insert_at_end` on the block.
+    builder.insert_after(x_instr)
     after_ir = module.ir()
     file = os.path.join(os.path.dirname(__file__), "resources/test_reordering_after.ll")
     expected = Path(file).read_text()
