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