ir.py

from dataclasses import dataclass, field
from enum import Enum, auto
from itertools import groupby, combinations
from pprint import pprint
from typing import Optional, Any, List, Set, Dict, Tuple

DEBUG = True


class StatementType(Enum):
    """
    Statement type
    """

    ASSIGN = auto()
    FUNCTION_CALL = auto()
    PRIMITIVE = auto()
    UNREACHABLE = auto()
    RETURN = auto()
    GOTO = auto()


class ValueType(Enum):
    """
    Value type
    """

    CONST = auto()
    LOCATION = auto()
    BORROW = auto()
    DEREF = auto()
    CALL = auto()
    UNWRAP = auto()


@dataclass
class Definition:
    location: int = None
    bb_index: int = None
    stmt_index: int = None

    def __eq__(self, other):
        return self.__hash__() == other.__hash__()

    def __hash__(self):
        return hash((self.location, self.bb_index, self.stmt_index))

    def __repr__(self):
        return f"Def(loc={self.location}, bb{self.bb_index}, s{self.stmt_index})"


class Mode(Enum):
    """
    Mode of ownership transfer (move or copy)
    """

    NONE = auto()
    MOVE = auto()
    NOT_MOVE = auto()
    COPY = auto()


@dataclass(kw_only=True)
class Statement:
    """
    Statement occurring in MIR with concomitant data relevant for building CFG+BB IR
    """

    stmt_type: Optional[StatementType] = None
    lhs_location: int = None
    mutable: Optional[bool] = False
    mode: Optional[Mode] = Mode.NONE
    value_type: Optional[ValueType] = None
    rhs_location: Optional[int] = None
    rhs_value: Optional[Any] = None
    pred: Set[Tuple[Optional[int], int]] = field(default_factory=set)
    succ: Set[Tuple[Optional[int], int]] = field(default_factory=set)

    # def_in[n] = U OUT[p] for each p in pred[n]
    def_in: Set[Definition] = field(default_factory=set)
    # def_out[i] = GEN[n] U (IN[n] - KILL[n])
    def_out: Set[Definition] = field(default_factory=set)

    # IN[n] = use[n] U (OUT[n] - def[n])
    live_in: Set = field(default_factory=set)
    # OUT[n] = U IN[s] for each s in succ[n]
    live_out: Set = field(default_factory=set)
    # defs[i] = {x | x is used in Statement() i}
    defs: Set[int] = field(default_factory=set)
    # uses[i] = {x | x is defined in Statement() i}
    uses: Set[int] = field(default_factory=set)

    def borrow(self) -> bool:
        """
        Check if statement borrows
        :return: True if statement borrows
        """
        return self.value_type == ValueType.BORROW or (self.value_type == ValueType.UNWRAP and self.rhs_location)

    def gen_defs(self):
        """
        All assignment MIR statements define a location. (_1 = const 42_u32;)
        :return: the set of locations that are generated by this statement.
        """
        return {self.lhs_location}

    def gen_uses(self):
        """
        MIR statements may use locations. (_n = _1;) here _1 is used.
        :return: the set of locations that are used by this statement.
        """
        if self.value_type == ValueType.UNWRAP:
            return {self.rhs_location}
        elif self.rhs_location:
            return {self.rhs_location}
        else:
            return set()

    def __repr__(self):
        return (
            f"\tStatement(\n"
            f"\t\tstmt_type={self.stmt_type},\n"
            f"\t\tlhs_location={self.lhs_location}, \n"
            f"\t\trhs_location={self.rhs_location}, \n"
            f"\t\trhs_value={self.rhs_value}, \n"
            f"\t\trhs_type={self.value_type}, \n"
            f"\t\tmutability={self.mutable}\n"
            f"\t\tgen={self.defs}\n"
            f"\t\tkill={self.uses}\n"
            f"\t)"
        )


@dataclass(kw_only=True)
class FunctionArg:
    """
    Function argument, mode (move, etc), location, constant, type of
    """

    mode: Mode = Mode.NONE
    location: [int | None] = None
    # constant (value, type)
    constant: Optional[Tuple[Any, str]] = None
    type: Optional[str] = None

    def __repr__(self):
        return (
            f"\tFunctionArg(\n"
            f"\t\tmode={self.mode},\n"
            f"\t\tlocation={self.location}, \n"
            f"\t\tconstant={self.constant}, \n"
            f"\t\ttype={self.type}\n"
        )


@dataclass(kw_only=True)
class FunctionStatement(Statement):
    """
    Function statement occurring in MIR with concomitant data relevant for building CFG+BB IR
    """

    function_method: Optional[str] = None
    function_type: Optional[str] = None
    function_args: Optional[List[FunctionArg]] = field(default_factory=list)
    bb_goto: Optional[int] = None

    def reborrow(self, borrows: Set['Borrow']):
        for use in self.uses:
            for borrow in borrows:
                if "get" in self.function_method and any("HashMap" in t for t in self.function_type):
                    if use == borrow.borrower:
                        return use, borrow
        return None, None

    def gen_uses(self):
        """
        Function statements may use locations in FunctionArgs (e.g. switchInt(move _1) -> ...), here _1 is used.
        :return: the set of locations that are used by this statement.
        """
        use_set = set()
        for arg in self.function_args:
            if arg.location is not None:
                if arg.mode in [Mode.MOVE]:
                    if arg.location:
                        use_set.add(arg.location)
        return use_set

    def __repr__(self):
        return (
            f"\tFunctionStatement(\n"
            f"\t\tfunction_method={self.function_method},\n"
            f"\t\tfunction_type={self.function_type}, \n"
            f"\t\tfunction_args={self.function_args}, \n"
            f"\t\tbb_goto={self.bb_goto}, \n"
        )


@dataclass(kw_only=True)
class PrimitiveFunctionStatement(Statement):
    """
    Primitive statement occurring in MIR with concomitant data relevant for building CFG+BB IR
    """

    primitive_type: str = None
    primitive_args: Optional[List[FunctionArg]] = field(default_factory=list)
    bb_goto: Optional[List[int]] = None

    def gen_uses(self):
        """
        Function statements may use locations in FunctionArgs (e.g. switchInt(move _1) -> ...), here _1 is used.
        :return: the set of locations that are used by this statement.
        """
        use_set = set()
        # add super uses
        use_set.update(super().gen_uses())
        for arg in self.primitive_args:

            # TODO: fix this, happens somewhere in valueargs, where arg is nested for some reason
            # TODO: HACK: if arg is list, unwrap to one FunctionArg
            if isinstance(arg, list):
                arg = arg[0]

            if arg.location is not None:
                if arg.mode in [Mode.MOVE] or self.primitive_type == "drop":
                    use_set.add(arg.location)
        if self.primitive_type == "return":
            use_set.add(0)
        return use_set

    # repr
    def __repr__(self):
        return (
            f"\tPrimitiveFunctionStatement(\n"
            f"\t\tprimitive_type={self.primitive_type},\n"
            f"\t\tprimitive_args=\n{self.primitive_args}, \n"
            f"\t\tbb_goto={self.bb_goto}\n"
        )


@dataclass
class BasicBlock:
    name: int = None
    stmts: List[Statement] = field(default_factory=list)
    succ: Set[int] = field(default_factory=set)
    pred: Set[int] = field(default_factory=set)
    def_in: Set = field(default_factory=set)
    def_out: Set = field(default_factory=set)
    live_in: Set = field(default_factory=set)
    live_out: Set = field(default_factory=set)

    def __repr__(self):
        return (
            f"\tBasicBlock(id={self.name}, \n"
            f"\tstmts=\n{self.stmts}, \n"
            f"\tsucc={self.succ}, \n"
            f"\tpred={self.pred}, \n"
            f"\tlivein={self.def_in}, \n"
            f"\tliveout={self.def_out}, \n"
        )

    def add_statements(self, stmts: List[Statement]):
        self.stmts.extend(stmts)


@dataclass()
class Borrow:
    """
    Borrow is a dataclass that represents a borrow of a location.
    Borrower is the location holding the reference to a borrowee
    """

    borrower: int = None
    borrowee: int = None
    mutable: bool = False
    bb_index: int = None
    stmt_index: int = None

    def __hash__(self):
        return hash((self.borrower, self.borrowee, self.mutable, self.bb_index, self.stmt_index))

    def __eq__(self, other):
        return self.__hash__() == other.__hash__()

    def __repr__(self):
        return (
            f"Borrow(er={self.borrower}, ee={self.borrowee}, mut={self.mutable}, bb{self.bb_index}, s{self.stmt_index})"
        )

    def overlaps(self, cfg: 'CFG', other_borrow: 'Borrow') -> tuple[bool, Optional[int], Optional[int]]:
        """
        Returns true if the borrower of this borrow is live in the other borrow
        """
        a = self.borrower
        b = other_borrow.borrower
        # enumerate with index all stmts in cfg
        for bb_index, bb in enumerate(cfg.bbs):
            for stmt_index, stmt in enumerate(bb.stmts):
                # check if both borrows are live at same program point
                if a in stmt.live_in and b in stmt.live_in:
                    # then problem
                    return True, bb_index, stmt_index
        return False, None, None


class CFG:
    """
    Control flow graph data structure using BasicBlock as nodes, and edges by pred and succ on BasicBlock/linkedlist
    """

    # list of BB nodes
    bbs: List[BasicBlock] = []
    # list of BB edges
    edges: List[Tuple[int, int]] = []
    # entry and exit nodes, referring to bb ids
    entry: int = None
    exit: int = None
    # list of types of locations in the CFG k:v -> location:type
    _types: Dict[int, str] = {}

    def find_and_set_entry_exit(self):
        """
        Find entry and exit nodes in the CFG
        """
        for bb in self.bbs:
            if not bb.pred:
                self.entry = bb.name
            if not bb.succ:
                self.exit = bb.name

    def add_edge(self, pred: int, succ: List[int]):
        # if only one succ, add it
        if isinstance(succ, int):
            # if pred > succ:
            #    pred, succ = succ, pred
            self.edges.append((pred, succ))
        else:
            # for each succ, add them to set of bb succs
            for s in succ:
                self.edges.append((pred, s))

    def fill_in_bb_pred_succ(self):
        # for each bb, add pred and succ according to self.edges
        for e in self.edges:
            pred, succ = e
            for bb in self.bbs:
                if bb.name == pred:
                    bb.succ.add(succ)
                if bb.name == succ:
                    bb.pred.add(pred)

    # noinspection PyUnresolvedReferences
    def finalise_cfg(self):
        # set bb succ, pred, entry, exit
        self.fill_in_bb_pred_succ()
        self.find_and_set_entry_exit()

        # compute succ and pred for each stmt in each bb for the combination of bools
        # first, last, is_entry, is_exit
        for bb in self.bbs:
            for i, stmt in enumerate(bb.stmts):
                # use tuples (bb_index, stmt_index), bb_index None if within same bb
                # bb_index points to succ/pred bb_index if traversing basic blocks

                last = i == len(bb.stmts) - 1
                first = i == 0

                # last stmt in bb
                if not first and last:
                    stmt.pred.add((None, i - 1))
                    for succ in bb.succ:
                        stmt.succ.add((succ, 0))

                # first stmt in bb
                elif first and not last:
                    stmt.succ.add((None, i + 1))
                    for pred in bb.pred:
                        stmt.pred.add((pred, len(self.bbs[pred].stmts) - 1))

                # middle stmt in bb
                elif not first and not last:
                    assert i > 0
                    stmt.succ.add((None, i + 1))
                    stmt.pred.add((None, i - 1))

                elif first and last:
                    for pred in bb.pred:
                        stmt.pred.add((pred, len(self.bbs[pred].stmts) - 1))
                    for succ in bb.succ:
                        stmt.succ.add((succ, 0))

                else:
                    assert False, f"Impossible state during {self.__class__.__name__} finalise_cfg"

                # if entry or exit, reset pred/succ
                if stmt == self.entry:
                    stmt.pred = None

                elif stmt == self.exit:
                    stmt.succ = None

                # run gens and kills on each stmt
                stmt.defs = stmt.gen_defs()
                stmt.uses = stmt.gen_uses()

        # assert that all succ and preds in CFG are of the type Set[Tuple[Optional[int], int]]
        for bb in self.bbs:
            for stmt in bb.stmts:
                assert all(isinstance(s, tuple) for s in stmt.succ)
                assert all(isinstance(p, tuple) for p in stmt.pred)

    def index_of(self, elem) -> int:
        # if elem is a bb, return index of bb in bbs
        if isinstance(elem, BasicBlock):
            return self.bbs.index(elem)
        # if elem is a stmt, return index of bb in bbs
        elif isinstance(elem, Statement):
            for bb in self.bbs:
                if elem in bb.stmts:
                    return self.bbs.index(bb)

    def add_bb(self, node: BasicBlock):
        self.bbs.append(node)
        # sort bbs by name (int)
        self.bbs.sort(key=lambda x: x.name)

    def __repr__(self):
        return f"CFG(bbs={self.bbs}, edges={self.edges}, entry={self.entry}, exit={self.exit})"

    def pprint(self):
        for n in self.bbs:
            print(f"BB {n.name}:\n\tSucc: {n.succ}\n\tPred: {n.pred}\n\tStmts (num: {len(n.stmts)}: {n.stmts}")
        print(f"CFG edges: {self.edges}")
        # print succs and preds for each bb
        for bb in self.bbs:
            print(f"BB {bb.name}:\n\tPred: {bb.pred}\n\tSucc: {bb.succ}")
        print(f"CFG entry: {self.entry}, exit: {self.exit}")

    def compute_reaching_definitions(self):
        """
        Compute reaching definitions. Use tuple of (bb_index, and stmt_index) as definitions points.
        Generates and kills functions are defined for the Statement class.
        Use the stmt.live_in and stmt.live_out as the IN and OUT sets for each statement.
        """
        # iterate until no change for IN and OUT sets
        # IN[n] = U OUT[p] for each p in pred[n]
        # OUT[n] = GEN[n] U (IN[n] - KILL[n])
        while True:
            change = False
            for b_i, bb in enumerate(self.bbs):
                for s_i, stmt in enumerate(bb.stmts):
                    # remember old def_in and def_out
                    old_def_in = stmt.def_in.copy()
                    old_def_out = stmt.def_out.copy()

                    # compute IN[n] = U OUT[p] for each p in pred[n]
                    for pred in stmt.pred:
                        # if pred is within same bb
                        if pred[0] is None:
                            stmt.def_in.update(bb.stmts[pred[1]].def_out)
                        # if pred is in a different bb
                        else:
                            stmt.def_in.update(self.bbs[pred[0]].def_out)

                    # assert all IN are of type Set[Definition]
                    assert isinstance(stmt.def_in, Set)
                    assert all(isinstance(d, Definition) for d in stmt.def_in)

                    # compute OUT[n] = DEFS[n] U (IN[n] - USES[n])
                    in_sans_uses = set(filter(lambda x: x.location not in stmt.uses, stmt.def_in))
                    # for each location definition, create Definition object with current indicies
                    defs = {Definition(s, b_i, s_i) for s in stmt.defs}
                    # union of defs and in_sans_uses (defs | in_sans_uses), filtering None locations
                    stmt.def_out = set(filter(lambda x: x.location is not None, defs | in_sans_uses))

                    # update def_in for bb if first statement in bb
                    if bb.stmts.index(stmt) == 0:
                        bb.def_in = stmt.def_in.copy()

                    # update def_out for bb if last statement in bb
                    if stmt == bb.stmts[-1]:
                        bb.def_out = stmt.def_out.copy()

                    # check if change
                    if stmt.def_in != old_def_in or stmt.def_out != old_def_out:
                        change = True
            if not change:
                break

    def compute_liveness(self):
        """
        Compute liveness. Use tuple of (bb_index, and stmt_index) as definitions points.
        Generates and kills functions are defined for the Statement class.
        Use the stmt.live_in and stmt.live_out as the IN and OUT sets for each statement.
        """
        # iterate until no change for IN and OUT sets.
        # OUT[n] = U IN[s] for each s in succ[n]
        # IN[n] = use[n] U (OUT[n] - def[n])
        while True:
            change = False
            for b_i, bb in reversed(list(enumerate(self.bbs))):
                for s_i, stmt in reversed(list(enumerate(bb.stmts))):
                    # remember old live_in and live_out
                    old_live_in = stmt.live_in.copy()
                    old_live_out = stmt.live_out.copy()

                    # compute OUT[n] = U IN[s] for each s in succ[n]
                    for succ in stmt.succ:
                        # if succ is within same bb
                        if succ[0] is None:
                            stmt.live_out.update(bb.stmts[succ[1]].live_in)
                        # if succ is in a different bb
                        else:
                            stmt.live_out.update(self.bbs[succ[0]].live_in)

                    # compute IN[n] = use[n] U (OUT[n] - def[n])
                    stmt.live_in = stmt.uses.union(stmt.live_out.difference(stmt.defs))

                    # update live_in for bb if first statement in bb
                    if bb.stmts.index(stmt) == 0:
                        bb.live_in = stmt.live_in.copy()

                    # update live_out for bb if last statement in bb
                    if stmt == bb.stmts[-1]:
                        bb.live_out = stmt.live_out.copy()

                    # check if change
                    if stmt.live_in != old_live_in or stmt.live_out != old_live_out:
                        change = True

            if not change:
                break

    def compute_borrows(self) -> list[Borrow]:
        # get all statements which borrow() is true, with their bb_index and stmt_index and borrow info
        borrows = set(
            Borrow(stmt.lhs_location, stmt.rhs_location, stmt.mutable, b_i, s_i)
            for b_i, bb in enumerate(self.bbs)
            for s_i, stmt in enumerate(bb.stmts)
            if stmt.borrow()
        )

        # add 'reborrows' from stmt semantics, e.g. 'reborrow = Hashmap::get(move ref)'
        function_borrows = set()
        for b_i, bb in enumerate(self.bbs):
            for s_i, stmt in enumerate(bb.stmts):
                # if statement which may reborrow
                if isinstance(stmt, FunctionStatement):
                    # get possible reborrows
                    use, borrow = stmt.reborrow(borrows)
                    # add to set of borrows if actually reborrows, borrowee is the original borrow's borrowee
                    if use and borrow:
                        function_borrows.add(Borrow(stmt.lhs_location, borrow.borrowee, borrow.mutable, b_i, s_i))

        # union newly found borrows, sorting by borrowee, then borrower - casting to list, to keep order for groupby
        return list(sorted(borrows | function_borrows, key=lambda x: (x.borrowee, x.borrower)))

    def borrow_check(self, borrows: list[Borrow]) -> bool:
        """
        Check if borrows are still valid according to Rust borrowing rules, per
        https://doc.rust-lang.org/book/ch04-02-references-and-borrowing.html:
        - At any given time, you can have either one mutable reference or any number of immutable references. (crux)
        - References must always be valid (moot, since previous compiler phases do this)
        """
        # for each set of borrows (set of Borrow-objects) where borrowee (location) is identical
        for borrowee, borrow_set in groupby(borrows, lambda x: x.borrowee):
            # pedantic: avoid consumption by advancing groupby-iterator
            # (https://stackoverflow.com/questions/34644010/iterator-produced-by-itertools-groupby-is-consumed-unexpectedly)
            borrow_set = set(borrow_set)

            # debug print state, with newline and tabbed borrow-set elements
            print(f"Bcking borrowee {borrowee} with borrows: ")
            for borrow in borrow_set:
                print(f"\t{borrow}")

            # if one or more mutable borrows, must check that mut-borrow is not live at same time as any immut-borrows
            if len(list(filter(lambda x: x.mutable, borrow_set))) >= 1:

                # for each combination of borrow for a given borrowee, check if they overlap
                for borrow1, borrow2 in combinations(borrow_set, 2):
                    # type annotation
                    borrow1: Borrow
                    borrow2: Borrow
                    print(f"checking \t{borrow1} " f"\nand \t\t{borrow2}")

                    # if both are immutable, we don't care about overlap
                    if not borrow1.mutable and not borrow2.mutable:
                        print("both immutable, don't care about overlap\n")
                        continue

                    # if overlap with mutable borrow, return false
                    overlap, lap_b_i, lap_s_i = borrow1.overlaps(self, borrow2)
                    if overlap:
                        print(f"BCK ERROR: found overlap at b_i: {lap_b_i}, s_i: {lap_s_i}\n")
                        return False
                    else:
                        print(f"no overlap\n")
            else:
                # no mutable borrows, so no need to check for overlap
                print(f"no mutable borrows for: {borrowee}, have {len(borrow_set)} immutable borrows, all good.")
        return True