compiler/compiler/ir.py

from __future__ import annotations

import abc
from abc import abstractmethod

from .errors import OverrideMandatoryError
from .logger import Logger
from .source import SourceLocation, IRLocation as Location
from .typechecking import typecheck

logger = Logger(__name__)


class IRItem:
    def __init__(self, location: SourceLocation):
        self.location = location

    @abstractmethod
    def codegen(self) -> str:
        raise OverrideMandatoryError()

    @abstractmethod
    def destination(self) -> IRValue:
        raise OverrideMandatoryError()

    def __repr__(self):
        return f"{self.codegen()}"


class IRAction(IRItem, abc.ABC):
    def __init__(self, location: SourceLocation, reads: list[IRValue], writes: list[IRAssignable]):
        super().__init__(location)

        # Must be changed updated in a second pass, when every item has been generated
        self.ir_location: Location = Location(-1, [])
        self.reads = reads
        self.writes = writes

        for read in reads:
            read.readers += [self]

        for write in writes:
            write.writers += [self]


class IRNop(IRAction):
    def __init__(self, location: SourceLocation):
        super().__init__(location, reads=[], writes=[])

    def codegen(self) -> str:
        return "NOP"

    def destination(self) -> IRValue:
        return IRVoid(self.location)


class IRValue(IRItem, abc.ABC):
    values: list[IRValue] = []

    def __init__(self, location: SourceLocation):
        super().__init__(location)
        self.readers: list[IRAction] = []
        IRValue.values += [self]

    def __del__(self):
        IRValue.values.remove(self)

    def destination(self) -> IRValue:
        return self

    def __str__(self):
        return self.codegen()

    def __repr__(self) -> str:
        return f"{self.codegen()} (readers: {repr(self.readers)})"


class IRMove(IRAction):

    @typecheck
    def __init__(self, location: SourceLocation, dest: IRAssignable, source: IRValue):
        super().__init__(location, reads=[source], writes=[dest])
        self.dest = dest
        self.source = source

    def destination(self) -> IRValue:
        return self.dest

    def codegen(self) -> str:
        return f"MOVE {self.source} -> {self.dest}"


class IRCall(IRAction):
    def __init__(self, location: SourceLocation, dest: IRAssignable, function: IRVariable, arguments: list[IRValue]):
        super().__init__(location, reads=arguments, writes=[dest])
        self.dest = dest
        self.function = function
        self.arguments = arguments

    def codegen(self) -> str:
        return f"CALL {self.function} -> {self.dest} : {', '.join(str(arg) for arg in self.arguments)}"

    def destination(self) -> IRValue:
        return self.dest


class IRImmediate(IRValue):
    @typecheck
    def __init__(self, location: SourceLocation, value: int | float | str):
        super().__init__(location)
        self.value = value

    def codegen(self):
        return f"{self.value}"


class IRVoid(IRValue):
    @typecheck
    def __init__(self, location: SourceLocation):
        super().__init__(location)

    def codegen(self) -> str:
        return "<VOID>"


class IRAssignable(IRValue, metaclass=abc.ABCMeta):
    def __init__(self, location: SourceLocation):
        super().__init__(location)
        self.writers: list[IRAction] = []

    def __str__(self):
        return self.codegen()

    def __repr__(self) -> str:
        return f"{self.codegen()} (readers: {repr([reader for reader in self.readers])}, " + \
            f"writers: {repr([writer for writer in self.writers])})"


class IRRegister(IRAssignable):
    register_id: int = 0
    registers: dict[int, IRRegister] = {}

    def __init__(self, location: SourceLocation):
        super().__init__(location)
        self.id = IRRegister.register_id
        IRRegister.register_id += 1
        IRRegister.registers[self.id] = self

    def __hash__(self) -> int:
        return self.id

    def remove(self):
        del IRRegister.registers[self.id]

    def codegen(self):
        return f"%r{self.id}"


class IRVariable(IRAssignable):
    def __init__(self, location: SourceLocation, fq_identifier: str):
        super().__init__(location)
        self.fq_identifier = fq_identifier

    def codegen(self) -> str:
        return f"@{self.fq_identifier}"


class IRAdd(IRAction):

    @typecheck
    def __init__(self, location: SourceLocation, dest: IRAssignable, *values: IRValue):
        super().__init__(location, reads=list(values), writes=[dest])
        assert all(isinstance(v, IRValue) for v in values)

        self.values = values
        self.dest = dest

    def destination(self) -> IRValue:
        return self.dest

    def codegen(self) -> str:
        values = [str(value) for value in self.values]
        return f"ADD {', '.join(values)} -> {self.dest}"


class IRMul(IRAction):

    @typecheck
    def __init__(self, location: SourceLocation, dest: IRAssignable, *values: IRValue):
        super().__init__(location, reads=list(values), writes=[dest])
        assert all(isinstance(v, IRValue) for v in values)

        self.values = values
        self.dest = dest

    def destination(self) -> IRValue:
        return self.dest

    def codegen(self) -> str:
        values = [str(value) for value in self.values]
        return f"MUL {', '.join(values)} -> {self.dest}"


class IRNegation(IRAction):

    @typecheck
    def __init__(self, location: SourceLocation, dest: IRAssignable, source: IRValue):
        super().__init__(location, reads=[source], writes=[dest])

        self.source = source
        self.dest = dest

    def destination(self) -> IRValue:
        return self.dest

    def codegen(self) -> str:
        return f"NEG {self.source} -> {self.dest}"


class IRInvert(IRAction):

    @typecheck
    def __init__(self, location: SourceLocation, dest: IRAssignable, source: IRValue):
        super().__init__(location, reads=[source], writes=[dest])

        self.source = source
        self.dest = dest

    def destination(self) -> IRValue:
        return self.dest

    def codegen(self) -> str:
        return f"INVERT {self.source} -> {self.dest}"