Videre

linedance-app/venv/lib/python3.12/site-packages/llvmlite/ir/transforms.py

@@ -0,0 +1,64 @@
+from llvmlite.ir import CallInstr
+
+
+class Visitor(object):
+    def visit(self, module):
+        self._module = module
+        for func in module.functions:
+            self.visit_Function(func)
+
+    def visit_Function(self, func):
+        self._function = func
+        for bb in func.blocks:
+            self.visit_BasicBlock(bb)
+
+    def visit_BasicBlock(self, bb):
+        self._basic_block = bb
+        for instr in bb.instructions:
+            self.visit_Instruction(instr)
+
+    def visit_Instruction(self, instr):
+        raise NotImplementedError
+
+    @property
+    def module(self):
+        return self._module
+
+    @property
+    def function(self):
+        return self._function
+
+    @property
+    def basic_block(self):
+        return self._basic_block
+
+
+class CallVisitor(Visitor):
+    def visit_Instruction(self, instr):
+        if isinstance(instr, CallInstr):
+            self.visit_Call(instr)
+
+    def visit_Call(self, instr):
+        raise NotImplementedError
+
+
+class ReplaceCalls(CallVisitor):
+    def __init__(self, orig, repl):
+        super(ReplaceCalls, self).__init__()
+        self.orig = orig
+        self.repl = repl
+        self.calls = []
+
+    def visit_Call(self, instr):
+        if instr.callee == self.orig:
+            instr.replace_callee(self.repl)
+            self.calls.append(instr)
+
+
+def replace_all_calls(mod, orig, repl):
+    """Replace all calls to `orig` to `repl` in module `mod`.
+    Returns the references to the returned calls
+    """
+    rc = ReplaceCalls(orig, repl)
+    rc.visit(mod)
+    return rc.calls