graduate_code/solutions/language.py

"""G02: typed, pure, finite expressions. No eval or host-language calls.
AST: int, bool, variable str, ('+',a,b), ('*',a,b), ('<',a,b),
('if',condition,then,otherwise), ('let',name,value,body).
"""


def infer(expr, env=None, depth=0):
    env = {} if env is None else env
    if depth > 48:
        raise ValueError("expression depth exceeds 48")
    if type(expr) is bool:
        return "bool"
    if type(expr) is int:
        return "int"
    if type(expr) is str:
        if expr not in env:
            raise ValueError("unbound name: " + expr)
        return env[expr]
    if type(expr) is not tuple or not expr:
        raise ValueError("malformed expression")
    tag = expr[0]
    if tag in ("+", "*", "<") and len(expr) == 3:
        if any(infer(x, env, depth + 1) != "int" for x in expr[1:]):
            raise TypeError("arithmetic/comparison requires integers")
        return "bool" if tag == "<" else "int"
    if tag == "if" and len(expr) == 4:
        if infer(expr[1], env, depth + 1) != "bool":
            raise TypeError("condition must be Boolean")
        left = infer(expr[2], env, depth + 1)
        right = infer(expr[3], env, depth + 1)
        if left != right:
            raise TypeError("branches must have identical types")
        return left
    if tag == "let" and len(expr) == 4 and type(expr[1]) is str:
        typ = infer(expr[2], env, depth + 1)
        child = dict(env)
        child[expr[1]] = typ
        return infer(expr[3], child, depth + 1)
    raise ValueError("unknown node or wrong arity")


def evaluate(expr, env=None):
    values = {} if env is None else dict(env)
    if any(type(x) not in (int, bool) for x in values.values()):
        raise TypeError("environment values must be integers or Booleans")
    infer(expr, {k: "bool" if type(v) is bool else "int"
                 for k, v in values.items()})

    def go(e, scope):
        if type(e) in (int, bool):
            return e
        if type(e) is str:
            return scope[e]
        tag = e[0]
        if tag == "let":
            value = go(e[2], scope)  # value uses the outer scope
            return go(e[3], {**scope, e[1]: value})
        if tag == "if":
            return go(e[2] if go(e[1], scope) else e[3], scope)
        a, b = go(e[1], scope), go(e[2], scope)
        if tag == "+":
            return a + b
        if tag == "*":
            return a * b
        return a < b

    return go(expr, values)


if __name__ == "__main__":
    e = ("let", "x", 3, ("if", ("<", "x", 4), ("*", "x", 2), 0))
    print(infer(e), evaluate(e))