modelverse/kernel/modelverse_jit/bytecode_to_cfg.py

"""Converts bytecode IR to CFG IR."""

import modelverse_jit.bytecode_ir as bytecode_ir
import modelverse_jit.cfg_ir as cfg_ir
import modelverse_jit.runtime as jit_runtime

class AnalysisState(object):
    """State that is common to the bytecode->CFG transformation of a function."""
    def __init__(self):
        self.counter = cfg_ir.SharedCounter()
        self.analyzed_instructions = set()
        self.loop_instructions = {}
        self.entry_point = cfg_ir.BasicBlock(self.counter)
        self.current_block = self.entry_point
        self.root_node = None
        self.__write_prolog()

    def __write_prolog(self):
        # Write a prolog in CFG IR.
        # We want to create the following definition:
        #
        #     !entry_point():
        #         $jit_locals = alloc-root-node
        #
        # We also want to store '$jit_locals' in an attribute, so we can
        # use it to shield locals from the GC.
        self.root_node = self.current_block.append_definition(cfg_ir.AllocateRootNode())

    def analyze(self, instruction):
        """Analyzes the given instruction as a basic block."""
        if instruction in self.analyzed_instructions:
            raise jit_runtime.JitCompilationFailedException(
                'Cannot jit non-tree instruction graph.')

        self.analyzed_instructions.add(instruction)

        # Find an analyzer.
        instruction_type = type(instruction)
        if instruction_type in self.instruction_analyzers:
            # Analyze the instruction.
            result = self.instruction_analyzers[instruction_type](self, instruction)

            # Check if the instruction has a 'next' instruction. If so, analyze it!
            if instruction.next_instruction is not None:
                next_result = self.analyze(instruction.next_instruction)
                if next_result.value.has_value() or (not result.value.has_value()):
                    result = next_result

            return result
        else:
            raise jit_runtime.JitCompilationFailedException(
                "Unknown instruction type: '%s'" % type(instruction))

    def emit_select(self, create_condition, create_if_body, create_else_body):
        """Emits a 'select' instruction."""
        # Create blocks that look like this:
        #
        # !current_block(...):
        #     ...
        #     $cond = <condition>
        #     select $cond, !if_block(), !else_block()
        #
        # !if_block():
        #     $if_result = <if-body>
        #     jump !phi_block($if_result)
        #
        # !else_block():
        #     $else_result = <else-body>
        #     jump !phi_block($else_result)
        #
        # !phi_block($result = block-parameter):
        #     ...
        #
        if_block = cfg_ir.BasicBlock(self.counter)
        else_block = cfg_ir.BasicBlock(self.counter)
        phi_block = cfg_ir.BasicBlock(self.counter)
        param_def = phi_block.append_parameter(cfg_ir.BlockParameter())

        condition = create_condition()
        self.current_block.flow = cfg_ir.SelectFlow(
            condition, cfg_ir.Branch(if_block), cfg_ir.Branch(else_block))

        self.current_block = if_block
        if_result = create_if_body()
        self.current_block.flow = cfg_ir.create_jump(phi_block, [if_result])

        self.current_block = else_block
        else_result = create_else_body()
        self.current_block.flow = cfg_ir.create_jump(phi_block, [else_result])

        self.current_block = phi_block
        return param_def

    def analyze_if(self, instruction):
        """Analyzes an 'if' instruction."""
        return self.emit_select(
            lambda: self.analyze(instruction.condition),
            lambda: self.analyze(instruction.if_clause),
            lambda:
            self.current_block.append_definition(cfg_ir.Literal(None))
            if instruction.else_clause is None
            else self.analyze(instruction.else_clause))

    def analyze_while(self, instruction):
        """Analyzes a 'while' instruction."""
        # Create blocks that look like this:
        #
        # !current_block(...):
        #     ...
        #     jump !loop_condition()
        #
        # !loop_condition():
        #     $condition = <condition>
        #     select $condition, !loop_body(), !loop_exit()
        #
        # !loop_body():
        #     $result = <body>
        #     jump !loop_condition()
        #
        # !loop_exit():
        #     $nothing = literal None
        #     ...
        loop_condition_block = cfg_ir.BasicBlock(self.counter)
        loop_body_block = cfg_ir.BasicBlock(self.counter)
        loop_exit_block = cfg_ir.BasicBlock(self.counter)
        self.loop_instructions[instruction] = (loop_condition_block, loop_exit_block)

        self.current_block.flow = cfg_ir.create_jump(loop_condition_block)

        self.current_block = loop_condition_block
        condition_value = self.analyze(instruction.condition)
        self.current_block.flow = cfg_ir.SelectFlow(
            condition_value, cfg_ir.Branch(loop_body_block), cfg_ir.Branch(loop_exit_block))

        self.current_block = loop_body_block
        self.analyze(instruction.body)
        self.current_block.flow = cfg_ir.create_jump(loop_condition_block)

        self.current_block = loop_exit_block
        return loop_exit_block.append_definition(cfg_ir.Literal(None))

    def analyze_return(self, instruction):
        """Analyzes a 'return' instruction."""
        if instruction.value is None:
            return_value = self.current_block.append_definition(cfg_ir.Literal(None))
        else:
            return_value = self.analyze(instruction.value)
        # Don't forget to deallocate the root node.
        self.current_block.append_definition(cfg_ir.DeallocateRootNode(self.root_node))
        self.current_block.flow = cfg_ir.ReturnFlow(return_value)
        self.current_block = cfg_ir.BasicBlock(self.counter)
        return self.current_block.append_definition(cfg_ir.Literal(None))

    def analyze_constant(self, instruction):
        """Analyzes a 'constant' instruction."""
        return self.current_block.append_definition(cfg_ir.Literal(instruction.constant_id))

    def analyze_resolve(self, instruction):
        """Analyzes a 'resolve' instruction."""
        return self.emit_select(
            lambda:
            self.current_block.append_definition(
                cfg_ir.CheckLocalExists(instruction.variable)),
            lambda:
            self.current_block.append_definition(
                cfg_ir.ResolveLocal(instruction.variable)),
            lambda:
            self.current_block.append_definition(
                cfg_ir.ResolveGlobal(instruction.variable)))

    def analyze_declare(self, instruction):
        """Analyzes a 'declare' instruction."""
        return self.current_block.append_definition(
            cfg_ir.DeclareLocal(instruction.variable, self.root_node))

    def analyze_global(self, instruction):
        """Analyzes a 'global' instruction."""
        return self.current_block.append_definition(cfg_ir.DeclareGlobal(instruction.variable))

    def analyze_assign(self, instruction):
        """Analyzes an 'assign' instruction."""
        pointer_result = self.analyze(instruction.pointer)
        value_result = self.analyze(instruction.value)
        return self.current_block.append_definition(
            cfg_ir.StoreAtPointer(pointer_result, value_result))

    def analyze_access(self, instruction):
        """Analyzes an 'access' instruction."""
        pointer_result = self.analyze(instruction.pointer)
        return self.current_block.append_definition(cfg_ir.LoadPointer(pointer_result))

    def analyze_output(self, instruction):
        """Analyzes an 'output' instruction."""
        value_result = self.analyze(instruction.value)
        return self.current_block.append_definition(cfg_ir.Output(value_result))

    def analyze_input(self, _):
        """Analyzes an 'input' instruction."""
        return self.current_block.append_definition(cfg_ir.Input())

    def analyze_break(self, instruction):
        """Analyzes a 'break' instruction."""
        if instruction.loop not in self.loop_instructions:
            raise jit_runtime.JitCompilationFailedException(
                "'break' instruction targets a 'while' loop that has not been defined yet.")

        _, exit_block = self.loop_instructions[instruction.loop]
        self.current_block.flow = cfg_ir.create_jump(exit_block)
        self.current_block = cfg_ir.BasicBlock(self.counter)
        return self.current_block.append_definition(cfg_ir.Literal(None))

    def analyze_continue(self, instruction):
        """Analyzes a 'continue' instruction."""
        if instruction.loop not in self.loop_instructions:
            raise jit_runtime.JitCompilationFailedException(
                "'continue' instruction targets a 'while' loop that has not been defined yet.")

        cond_block, _ = self.loop_instructions[instruction.loop]
        self.current_block.flow = cfg_ir.create_jump(cond_block)
        self.current_block = cfg_ir.BasicBlock(self.counter)
        return self.current_block.append_definition(cfg_ir.Literal(None))

    def analyze_call(self, instruction):
        """Analyzes the given 'call' instruction."""
        target = self.analyze(instruction.target)
        arg_list = []
        for key, arg_instruction in instruction.argument_list:
            arg_list.append((key, self.analyze(arg_instruction)))

        return self.current_block.append_definition(cfg_ir.JitFunctionCall(target, arg_list))

    instruction_analyzers = {
        bytecode_ir.SelectInstruction : analyze_if,
        bytecode_ir.WhileInstruction : analyze_while,
        bytecode_ir.ReturnInstruction : analyze_return,
        bytecode_ir.ConstantInstruction : analyze_constant,
        bytecode_ir.ResolveInstruction : analyze_resolve,
        bytecode_ir.DeclareInstruction : analyze_declare,
        bytecode_ir.GlobalInstruction : analyze_global,
        bytecode_ir.AssignInstruction : analyze_assign,
        bytecode_ir.AccessInstruction : analyze_access,
        bytecode_ir.OutputInstruction : analyze_output,
        bytecode_ir.InputInstruction : analyze_input,
        bytecode_ir.CallInstruction : analyze_call,
        bytecode_ir.BreakInstruction : analyze_break,
        bytecode_ir.ContinueInstruction : analyze_continue
    }