interpreterv4.py

from copy import deepcopy

from brewparse import parse_program
from element import Element
from intbase import InterpreterBase, ErrorType


class Variable:
    def __init__(self, element):
        self.element = element


class Object:
    def __init__(self):
        self.members = {"proto": Element("nil")}

    def get_member(self, interpreter, member_name):
        if member_name == "proto":
            prototype = self.members["proto"]
            if prototype.elem_type == "nil":
                interpreter.error(
                    ErrorType.NAME_ERROR, f"Prototype has not been defined"
                )
            return prototype
        elif member_name in self.members:
            return self.members[member_name]
        else:
            prototype = self.members["proto"]
            if prototype.elem_type != "nil":
                return prototype.get("val").get_member(interpreter, member_name)
            else:
                interpreter.error(
                    ErrorType.NAME_ERROR, f"Member {member_name} not in object"
                )

    def assign_member(self, interpreter, member_name, value):
        if member_name == "proto" and value.elem_type not in {"object", "nil"}:
            interpreter.error(ErrorType.TYPE_ERROR, f"Prototype must be an object")
        self.members[member_name] = value


class Return(Exception):
    def __init__(self, value):
        super().__init__()
        self.value = value


class Interpreter(InterpreterBase):
    builtin_functions = {"print", "inputi", "inputs"}

    def __init__(self, console_output=True, inp=None, trace_output=False):
        super().__init__(console_output, inp)
        self.trace_output = trace_output

        self.function_defs = {}
        self.variables = {}
        self.scopes = []

    def run(self, program):
        program_node = parse_program(program)
        functions = program_node.get("functions")
        for function in functions:
            name = function.get("name")
            params = function.get("args")
            self.function_defs.setdefault(name, {})
            self.function_defs[name][len(params)] = function

        main_function = Element("fcall", name="main", args=[])
        self.run_function(main_function)

    def run_statement(self, statement):
        match statement.elem_type:
            case "if":
                self.run_if(statement)
            case "while":
                self.run_while(statement)
            case "return":
                self.run_return(statement)
            case "fcall" | "mcall":
                self.run_function(statement)
            case "=":
                self.run_assignment(statement)

    def create_scope(self):
        self.scopes.append(set())

    def delete_scope(self):
        for variable_name in self.scopes[-1]:
            self.pop_variable(variable_name)
        self.scopes.pop()

    def push_variable(self, variable_name, variable):
        self.scopes[-1].add(variable_name)
        self.variables.setdefault(variable_name, [])
        self.variables[variable_name].append(variable)

    def pop_variable(self, variable_name):
        self.variables[variable_name].pop()
        if not self.variables[variable_name]:
            del self.variables[variable_name]

    def run_if(self, if_block):
        statements = if_block.get("statements")
        else_statements = if_block.get("else_statements")
        self.create_scope()
        try:
            condition = self.to_bool(if_block.get("condition"))
            if condition.get("val"):
                for statement in statements:
                    self.run_statement(statement)
            elif else_statements:
                for statement in else_statements:
                    self.run_statement(statement)
        except TypeError:
            self.error(
                ErrorType.TYPE_ERROR, "If condition does not evaluate to a boolean"
            )
        finally:
            self.delete_scope()

    def run_while(self, while_block):
        statements = while_block.get("statements")
        self.create_scope()
        try:
            while True:
                condition = self.to_bool(while_block.get("condition"))
                if condition.get("val"):
                    for statement in statements:
                        self.run_statement(statement)
                else:
                    break
        except TypeError:
            self.error(
                ErrorType.TYPE_ERROR, "While condition does not evaluate to a boolean"
            )
        finally:
            self.delete_scope()

    def run_return(self, statement):
        expression = statement.get("expression")
        if expression:
            raise Return(deepcopy(self.evaluate_expression(expression)))
        else:
            raise Return(Element("nil"))

    def run_function(self, function):
        name = function.get("name")
        args = function.get("args")
        if function.elem_type == "fcall":
            if name in self.function_defs and len(args) in self.function_defs[name]:
                function_def = self.function_defs[name][len(args)]
                params = function_def.get("args")
            elif name in self.variables:
                function_def = self.variables[name][-1].element

                if function_def.elem_type not in {"func", "closure"}:
                    self.error(
                        ErrorType.TYPE_ERROR,
                        f"Variable {name} does not hold a function",
                    )

                params = function_def.get("args")
                if len(params) != len(args):
                    self.error(
                        ErrorType.TYPE_ERROR,
                        f"{name} takes {len(params)} parameters: {len(args)} arguments given",
                    )
            elif name in self.builtin_functions:
                return self.run_builtin(function)
            else:
                self.error(
                    ErrorType.NAME_ERROR,
                    f"No {name}() function found that takes {len(args)} parameters",
                )
        elif function.elem_type == "mcall":
            object_name = function.get("objref")
            if object_name in self.variables:
                object_value = self.variables[object_name][-1].element
            else:
                self.error(
                    ErrorType.NAME_ERROR,
                    f"Variable {object_name} has not been defined",
                )

            if object_value.elem_type != "object":
                self.error(
                    ErrorType.TYPE_ERROR,
                    f"{object_name} is not an object",
                )

            object_object = object_value.get("val")
            function_def = object_object.get_member(self, name)
            if function_def.elem_type not in {"func", "closure"}:
                self.error(
                    ErrorType.TYPE_ERROR,
                    f"{object_name}.{name} does not hold a function",
                )

            params = function_def.get("args")
            if len(params) != len(args):
                self.error(
                    ErrorType.NAME_ERROR,
                    f"{name} takes {len(params)} parameters: {len(args)} arguments given",
                )

        statements = function_def.get("statements")
        param_names = {param.get("name") for param in params}

        # Get arguments before shadowing
        arg_variables = []
        for param, arg in zip(params, args):
            param_name = param.get("name")
            arg_name = arg.get("name")
            if param.elem_type == "refarg" and arg_name in self.variables:
                arg_variables.append(self.variables[arg_name][-1])
            elif param.elem_type == "refarg" and arg_name in self.function_defs:
                arg_variables.append(Variable(self.evaluate_expression(arg)))
            else:
                arg_variables.append(Variable(deepcopy(self.evaluate_expression(arg))))

        self.create_scope()
        if function.elem_type == "mcall":
            self.push_variable("this", self.variables[object_name][-1])

        if function_def.elem_type == "closure":
            captures = function_def.get("captures")
            exposed_captures = dict(
                filter(lambda capture: capture[0] not in param_names, captures.items())
            )
            for capture_name, capture_variable in exposed_captures.items():
                self.push_variable(capture_name, capture_variable)

        for param, arg_variable in zip(params, arg_variables):
            param_name = param.get("name")
            self.push_variable(param_name, arg_variable)

        try:
            for statement in statements:
                self.run_statement(statement)
            return Element("nil")
        except Return as ret:
            return ret.value
        finally:
            if function_def.elem_type == "closure":
                for capture_name in exposed_captures:
                    captures[capture_name] = self.variables[capture_name][-1]

            self.delete_scope()

    def run_builtin(self, function):
        name = function.get("name")
        args = function.get("args")
        match name:
            case "print":
                self.output("".join(self.fmt(arg) for arg in args))
                return Element("nil")
            case "inputi":
                if len(args) == 1:
                    self.output(self.evaluate_expression(args[0]).get("val"))
                elif len(args) > 1:
                    self.error(
                        ErrorType.NAME_ERROR,
                        "No inputi() function found that takes > 1 parameter",
                    )
                return Element("int", val=int(self.get_input()))
            case "inputs":
                if len(args) == 1:
                    self.output(self.evaluate_expression(args[0]).get("val"))
                elif len(args) > 1:
                    self.error(
                        ErrorType.NAME_ERROR,
                        "No inputs() function found that takes > 1 parameter",
                    )
                return Element("string", val=self.get_input())

    def fmt(self, expression):
        element = self.evaluate_expression(expression)
        match element.elem_type:
            case "int" | "string":
                return str(element.get("val"))
            case "bool":
                return str(element.get("val")).lower()

    def to_bool(self, expression):
        element = self.evaluate_expression(expression)
        match element.elem_type:
            case "bool":
                return element
            case "int":
                return Element("bool", val=bool(element.get("val")))
            case _:
                raise TypeError

    def to_int(self, expression):
        element = self.evaluate_expression(expression)
        match element.elem_type:
            case "int":
                return element
            case "bool":
                return Element("int", val=int(element.get("val")))
            case _:
                raise TypeError

    def run_assignment(self, assignment):
        name = assignment.get("name")
        expression = assignment.get("expression")
        value = self.evaluate_expression(expression)
        if "." in name:
            object_name, member_name = name.split(".")
            if object_name in self.variables:
                object_value = self.variables[object_name][-1].element
            else:
                self.error(
                    ErrorType.NAME_ERROR,
                    f"Variable {object_name} has not been defined",
                )

            if object_value.elem_type != "object":
                self.error(
                    ErrorType.TYPE_ERROR,
                    f"{object_name} is not an object",
                )

            object_object = object_value.get("val")
            object_object.assign_member(self, member_name, value)
        else:
            if name not in self.variables:
                self.push_variable(name, Variable(value))
            else:
                self.variables[name][-1].element = value

    def evaluate_expression(self, expression):
        match expression.elem_type:
            case "neg" | "!":
                return self.evaluate_unary_operation(expression)
            case "+" | "-" | "*" | "/" | "==" | "!=" | "<" | ">" | "<=" | ">=" | "||" | "&&":
                return self.evaluate_binary_operation(expression)
            case "fcall" | "mcall":
                return self.run_function(expression)
            case "lambda":
                return self.evaluate_lambda(expression)
            case "@":
                return Element("object", val=Object())
            case "var":  # Includes named functions and bound lambdas
                return self.evaluate_variable(expression)
            case "int" | "string" | "bool" | "nil" | "func" | "closure" | "object":
                return expression

    def evaluate_lambda(self, lambda_def):
        captures = {}
        for name, stack in self.variables.items():
            variable = stack[-1]
            value = variable.element
            if value.elem_type in {"closure", "object"}:
                captures[name] = variable
            else:
                captures[name] = deepcopy(variable)

        return Element(
            "closure",
            args=lambda_def.get("args"),
            statements=lambda_def.get("statements"),
            captures=captures,
        )

    def evaluate_variable(self, variable):
        name = variable.get("name")
        if "." in name:
            object_name, member_name = name.split(".")
            if object_name in self.variables:
                object_value = self.variables[object_name][-1].element
            else:
                self.error(
                    ErrorType.NAME_ERROR,
                    f"Variable {object_name} has not been defined",
                )

            if object_value.elem_type != "object":
                self.error(
                    ErrorType.TYPE_ERROR,
                    f"{object_name} is not an object",
                )

            object_object = object_value.get("val")
            return object_object.get_member(self, member_name)
        else:
            if name in self.variables:
                return self.variables[name][-1].element
            elif name in self.function_defs:
                if len(self.function_defs[name]) != 1:
                    self.error(ErrorType.NAME_ERROR, f"{name}() function is ambiguous")
                return next(iter(self.function_defs[name].values()))
            else:
                self.error(
                    ErrorType.NAME_ERROR,
                    f"Variable {name} has not been defined",
                )

    def evaluate_unary_operation(self, operation):
        op1 = self.evaluate_expression(operation.get("op1"))
        try:
            match operation.elem_type:
                case "neg":
                    if op1.elem_type == "int":
                        return Element("int", val=-op1.get("val"))
                    else:
                        raise TypeError
                case "!":
                    op1 = self.to_bool(op1)
                    return Element("bool", val=not op1.get("val"))
        except TypeError:
            self.error(
                ErrorType.TYPE_ERROR,
                f"Incompatible type for operation {operation.elem_type}: {op1.elem_type}",
            )

    def evaluate_binary_operation(self, operation):
        # Strict evaluation
        op1 = self.evaluate_expression(operation.get("op1"))
        op2 = self.evaluate_expression(operation.get("op2"))
        try:
            match operation.elem_type:
                case "+":
                    if op1.elem_type == op2.elem_type == "string":
                        return Element("string", val=op1.get("val") + op2.get("val"))
                    op1 = self.to_int(op1)
                    op2 = self.to_int(op2)
                    return Element("int", val=op1.get("val") + op2.get("val"))
                case "-":
                    op1 = self.to_int(op1)
                    op2 = self.to_int(op2)
                    return Element("int", val=op1.get("val") - op2.get("val"))
                case "*":
                    op1 = self.to_int(op1)
                    op2 = self.to_int(op2)
                    return Element("int", val=op1.get("val") * op2.get("val"))
                case "/":
                    op1 = self.to_int(op1)
                    op2 = self.to_int(op2)
                    return Element("int", val=op1.get("val") // op2.get("val"))
                case "==":
                    if (op1.elem_type, op2.elem_type) in {
                        ("int", "bool"),
                        ("bool", "int"),
                    }:
                        op1 = self.to_bool(op1)
                        op2 = self.to_bool(op2)
                        return self.evaluate_binary_operation(
                            Element("==", op1=op1, op2=op2)
                        )
                    elif op1.elem_type != op2.elem_type:
                        return Element("bool", val=False)
                    elif op1.elem_type in {"func", "closure", "object"}:
                        return Element("bool", val=op1 is op2)
                    else:
                        return Element("bool", val=op1.get("val") == op2.get("val"))
                case "!=":
                    if (op1.elem_type, op2.elem_type) in {
                        ("int", "bool"),
                        ("bool", "int"),
                    }:
                        op1 = self.to_bool(op1)
                        op2 = self.to_bool(op2)
                        return self.evaluate_binary_operation(
                            Element("!=", op1=op1, op2=op2)
                        )
                    elif op1.elem_type != op2.elem_type:
                        return Element("bool", val=True)
                    elif op1.elem_type in {"func", "closure", "object"}:
                        return Element("bool", val=op1 is not op2)
                    else:
                        return Element("bool", val=op1.get("val") != op2.get("val"))
                case "<":
                    if op1.elem_type == op2.elem_type == "int":
                        return Element("bool", val=op1.get("val") < op2.get("val"))
                    else:
                        raise TypeError
                case ">":
                    if op1.elem_type == op2.elem_type == "int":
                        return Element("bool", val=op1.get("val") > op2.get("val"))
                    else:
                        raise TypeError
                case "<=":
                    if op1.elem_type == op2.elem_type == "int":
                        return Element("bool", val=op1.get("val") <= op2.get("val"))
                    else:
                        raise TypeError
                case ">=":
                    if op1.elem_type == op2.elem_type == "int":
                        return Element("bool", val=op1.get("val") >= op2.get("val"))
                    else:
                        raise TypeError
                case "||":
                    op1 = self.to_bool(op1)
                    op2 = self.to_bool(op2)
                    return Element("bool", val=op1.get("val") or op2.get("val"))
                case "&&":
                    op1 = self.to_bool(op1)
                    op2 = self.to_bool(op2)
                    return Element("bool", val=op1.get("val") and op2.get("val"))
        except TypeError:
            self.error(
                ErrorType.TYPE_ERROR,
                f"Incompatible types for operation {operation.elem_type}: {op1.elem_type} and {op2.elem_type}",
            )