LSRA/stack_instruction.py at main · PhilippeGSK/LSRA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import enum
import dataclasses
from ir import *

class StackInstructionKind(enum.Enum):
    LdLocal = enum.auto()
    StLocal = enum.auto()
    Push = enum.auto()
    Pop = enum.auto()
    Add = enum.auto()
    Sub = enum.auto()
    Mul = enum.auto()
    Div = enum.auto()
    Eq = enum.auto()
    Jmp = enum.auto()
    Branch = enum.auto()
    Ret = enum.auto()

@dataclasses.dataclass
class StackInstruction:
    kind: StackInstructionKind
    operands: list

@dataclasses.dataclass
class StackFunction:
    local_vars: int
    instructions: list[StackInstruction]

def import_to_ir(fn: StackFunction) -> Ir:
    blocks = BasicBlockList()
    current_block = blocks.first
    tree_stack: list[Tree] = []

    def fold(block: BasicBlock, kind: TreeKind, n: int, operands: list) -> None:
        nonlocal tree_stack
        if (n > len(tree_stack)):
            raise Exception("Not enough stack operands")
        l = len(tree_stack)
        res = tree_stack[l - n:]
        tree_stack = tree_stack[:l - n]

        new_tree = Tree(kind=kind, subtrees=res, operands=operands, parent=None, block=block)
        for subtree in new_tree.subtrees:
            subtree.parent = new_tree

        tree_stack.append(new_tree)

    last_i_ins = len(fn.instructions) - 1
    i_stmt_start = 0

    for i_ins, ins in enumerate(fn.instructions):
        match ins.kind:
            case StackInstructionKind.LdLocal:
                fold(current_block, TreeKind.LdLocal, 0, ins.operands)

            case StackInstructionKind.StLocal:
                fold(current_block, TreeKind.StLocal, 1, ins.operands)
                current_block.append_tree(i_stmt_start, tree_stack.pop())
                i_stmt_start = i_ins + 1
                if tree_stack != []: raise Exception("Leftover stack operands")

            case StackInstructionKind.Push:
                fold(current_block, TreeKind.Const, 0, ins.operands)

            case StackInstructionKind.Pop:
                fold(current_block, TreeKind.Discard, 1, [])
                current_block.append_tree(i_stmt_start, tree_stack.pop())
                i_stmt_start = i_ins + 1
                if tree_stack != []: raise Exception("Leftover stack operands")

            case StackInstructionKind.Add:
                fold(current_block, TreeKind.BinOp, 2, [Operator.Add])

            case StackInstructionKind.Sub:
                fold(current_block, TreeKind.BinOp, 2, [Operator.Sub])

            case StackInstructionKind.Mul:
                fold(current_block, TreeKind.BinOp, 2, [Operator.Mul])

            case StackInstructionKind.Div:
                fold(current_block, TreeKind.BinOp, 2, [Operator.Div])

            case StackInstructionKind.Eq:
                fold(current_block, TreeKind.BinOp, 2, [Operator.Eq])

            case StackInstructionKind.Jmp:
                target = blocks.get_or_insert_block_at(ins.operands[0])
                fold(current_block, TreeKind.Jmp, 0, [BlockEdge(target=target)])
                current_block.append_tree(i_stmt_start, tree_stack.pop())
                i_stmt_start = i_ins + 1
                if tree_stack != []: raise Exception("Leftover stack operands")
                if (i_ins == last_i_ins): break
                current_block = blocks.get_or_insert_block_at(i_ins + 1)

            case StackInstructionKind.Branch:
                if_target = blocks.get_or_insert_block_at(ins.operands[0])
                else_target = blocks.get_or_insert_block_at(ins.operands[1])
                fold(current_block, TreeKind.Branch, 1, [BlockEdge(target=if_target), BlockEdge(target=else_target)])
                current_block.append_tree(i_stmt_start, tree_stack.pop())
                i_stmt_start = i_ins + 1
                if tree_stack != []: raise Exception("Leftover stack operands")
                if (i_ins == last_i_ins): break
                current_block = blocks.get_or_insert_block_at(i_ins + 1)

            case StackInstructionKind.Ret:
                fold(current_block, TreeKind.Ret, 1, [])
                current_block.append_tree(i_stmt_start, tree_stack.pop())
                i_stmt_start = i_ins + 1
                if (i_ins == last_i_ins): break
                current_block = blocks.get_or_insert_block_at(i_ins + 1)

        if (i_ins == last_i_ins):
            raise Exception("Illegal terminator")

    result = Ir(blocks=blocks, local_vars=fn.local_vars)
    result.recompute_predecessors()
    result.reindex()
    return result