<s> from . . import names <EOL> from . . ndtypes import ClosureT , make_closure_type <EOL> from . . syntax import Expr , UntypedFn , Closure , Var , Return , ClosureElt , TypedFn , FormalArgs <EOL> from . . syntax . helpers import get_types <EOL> def unpack_closure ( closure ) : <EOL> """<STR_LIT>""" <EOL> if closure . __class__ is ClosureT : <EOL> fn , closure_args = closure . fn , closure . arg_types <EOL> elif closure . __class__ is Closure : <EOL> fn = closure . fn <EOL> closure_args = closure . args <EOL> elif closure . type . __class__ is ClosureT : <EOL> fn , arg_types = closure . type . fn , closure . type . arg_types <EOL> closure_args = [ ClosureElt ( closure , i , type = arg_t ) <EOL> for ( i , arg_t ) in enumerate ( arg_types ) ] <EOL> else : <EOL> fn = closure <EOL> closure_args = [ ] <EOL> return fn , closure_args <EOL> def make_typed_closure ( clos , typed_fn ) : <EOL> if clos . __class__ is UntypedFn : <EOL> return typed_fn <EOL> assert isinstance ( clos , Expr ) and clos . type . __class__ is ClosureT <EOL> _ , closure_args = unpack_closure ( clos ) <EOL> if len ( closure_args ) == <NUM_LIT:0> : <EOL> return typed_fn <EOL> else : <EOL> t = make_closure_type ( typed_fn , get_types ( closure_args ) ) <EOL> return Closure ( typed_fn , closure_args , t ) <EOL> def mk_untyped_identity ( ) : <EOL> var_name = names . fresh ( '<STR_LIT:x>' ) <EOL> var_expr = Var ( var_name ) <EOL> fn_name = names . fresh ( '<STR_LIT>' ) <EOL> args_obj = FormalArgs ( ) <EOL> args_obj . add_positional ( var_name ) <EOL> return UntypedFn ( name = fn_name , <EOL> args = args_obj , <EOL> body = [ Return ( var_expr ) ] ) <EOL> untyped_identity_function = mk_untyped_identity ( ) <EOL> def _get_fundef ( fn ) : <EOL> c = fn . __class__ <EOL> if c is UntypedFn or c is TypedFn : <EOL> return fn <EOL> assert c is str , "<STR_LIT>" % ( fn , fn . type ) <EOL> return UntypedFn . registry [ fn ] <EOL> def _get_closure_type ( fn ) : <EOL> c = fn . __class__ <EOL> if c is ClosureT : <EOL> return fn <EOL> elif c is Closure : <EOL> return fn . type <EOL> elif c is Var : <EOL> assert fn . type . __class__ is ClosureT <EOL> return fn . type <EOL> else : <EOL> fundef = _get_fundef ( fn ) <EOL> return make_closure_type ( fundef , [ ] ) </s>
