type_check_Llambda.py

import ast
from ast import *
from type_check_Lfun import TypeCheckLfun
from utils import *
import typing

# This type checker uses bidirectional type checking to work-around
# the lack of type annotations in Python's lambdas.

class TypeCheckLlambda(TypeCheckLfun):

 # Example:
 # given: Callable[[tuple[bottom,tuple[int]], int], int]
 # output: tuple[Callable[[tuple[],int], int]]:
 def closure_type(self, fun_ty):
 match fun_ty:
 case FunctionType(params, rt):
 return TupleType([FunctionType([TupleType([])] + params[1:], rt)])
 case _:
 raise Exception('erase_closure_type expected function type, not '
 + str(fun_ty))
 
 def type_check_exp(self, e, env):
 match e:
 case Name(id):
 e.has_type = env[id]
 return env[id]
 case FunRef(id, arity):
 e.has_type = env[id]
 return e.has_type
 case UncheckedCast(exp, ty):
 self.type_check_exp(exp, env)
 return ty
 case Closure(arity, es):
 ts = [self.type_check_exp(e, env) for e in es]
 # closure values with different shapes must be given compatible types
 # e.has_type and e.erased_type is used in expose_allocation
 e.has_type = TupleType(ts)
 function_type = ts[0]
 e.erased_type = self.closure_type(function_type)
 return e.erased_type
 case Lambda(params, body):
 raise Exception('cannot synthesize a type for lambda: ' + str(e))
 case AllocateClosure(length, typ, arity):
 return typ
 case Call(Name('arity'), [func]):
 func_t = self.type_check_exp(func, env)
 match func_t:
 case FunctionType(params_t, return_t):
 return IntType()
 case TupleType(elts_t): # after closure conversion
 return IntType()
 case _:
 raise Exception('type_check_exp: in arity, unexpected ' + \
 repr(func_t))
 case Call(Name(f), args) if f in builtin_functions:
 return super().type_check_exp(e, env) 
 case Call(func, args):
 func_t = self.type_check_exp(func, env)
 match func_t:
 case FunctionType(params_t, return_t):
 for (arg, param_t) in zip(args, params_t):
 self.check_exp(arg, param_t, env)
 return return_t
 case _:
 raise Exception('type_check_exp: unexpected ' + \
 repr(func_t) + '\nin call\n' + str(e))
 case Uninitialized(ty):
 return ty
 case _:
 return super().type_check_exp(e, env)
 
 def check_exp(self, e, ty, env):
 match e:
 case Lambda(params, body):
 e.has_type = ty
 if isinstance(params, ast.arguments):
 new_params = [a.arg for a in params.args]
 e.args = new_params
 else:
 new_params = params
 match ty:
 case FunctionType(params_t, return_t):
 new_env = {x:t for (x,t) in env.items()}
 for (p,t) in zip(new_params, params_t):
 new_env[p] = t
 self.check_exp(body, return_t, new_env)
 case Bottom():
 pass
 case _:
 raise Exception('lambda does not have type ' + str(ty))
 case _:
 t = self.type_check_exp(e, env)
 self.check_type_equal(t, ty, e)

 # Use check_stmts in contexts where there is an expected return type,
 # such as inside the body of a function.
 def check_stmts(self, ss, return_ty, env):
 if len(ss) == 0:
 return
 #trace('*** check_stmts ' + repr(ss[0]) + '\n')
 match ss[0]:
 case FunctionDef(name, params, body, dl, returns, comment):
 #trace('*** tc_check ' + name)
 new_env = {x: t for (x,t) in env.items()}
 if isinstance(params, ast.arguments):
 new_params = [(p.arg, self.parse_type_annot(p.annotation)) for p in params.args]
 ss[0].args = new_params
 new_returns = self.parse_type_annot(returns)
 ss[0].returns = new_returns
 else:
 new_params = params
 new_returns = returns
 for (x,t) in new_params:
 new_env[x] = t
 rt = self.check_stmts(body, new_returns, new_env)
 self.check_stmts(ss[1:], return_ty, env)
 case Return(value):
 #trace('** tc_check return ' + repr(value))
 self.check_exp(value, return_ty, env)
 case Assign([v], value) if isinstance(v, Name):
 if v.id in env:
 self.check_exp(value, env[v.id], env)
 else:
 env[v.id] = self.type_check_exp(value, env)
 v.has_type = env[v.id]
 self.check_stmts(ss[1:], return_ty, env)
 case Assign([Subscript(tup, Constant(index), Store())], value):
 tup_t = self.type_check_exp(tup, env)
 match tup_t:
 case TupleType(ts):
 self.check_exp(value, ts[index], env)
 case Bottom():
 pass
 case ListType(elt_type):
 self.check_exp(value, elt_type, env)
 case _:
 raise Exception('check_stmts: expected a tuple, not ' \
 + repr(tup_t))
 self.check_stmts(ss[1:], return_ty, env)
 case AnnAssign(v, type_annot, value, simple) if isinstance(v, Name):
 ty_annot = self.parse_type_annot(type_annot)
 ss[0].annotation = ty_annot
 if v.id in env:
 self.check_type_equal(env[v.id], ty_annot)
 else:
 env[v.id] = ty_annot
 v.has_type = env[v.id]
 self.check_exp(value, ty_annot, env)
 self.check_stmts(ss[1:], return_ty, env)
 case _:
 self.type_check_stmts(ss, env)

 def type_check_stmts(self, ss, env):
 if len(ss) == 0:
 return
 match ss[0]:
 case Assign([v], value) if isinstance(v, Name):
 t = self.type_check_exp(value, env)
 if v.id in env:
 self.check_type_equal(env[v.id], t, value)
 else:
 env[v.id] = t
 v.has_type = env[v.id]
 return self.type_check_stmts(ss[1:], env)
 case Pass():
 return self.type_check_stmts(ss[1:], env)
 case _:
 return super().type_check_stmts(ss, env)
 
 def type_check(self, p):
 #trace('*** type check Llambda')
 match p:
 case Module(body):
 env = {}
 for s in body:
 match s:
 case FunctionDef(name, params, bod, dl, returns, comment):
 if isinstance(params, ast.arguments):
 params_t = [self.parse_type_annot(p.annotation) \
 for p in params.args]
 returns_t = self.parse_type_annot(returns)
 else:
 params_t = [t for (x,t) in params]
 returns_t = returns
 env[name] = FunctionType(params_t, returns_t)
 self.check_stmts(body, IntType(), env)
 case _:
 raise Exception('type_check: unexpected ' + repr(p))