<s> from . . import names <EOL> from . . analysis import contains_calls , can_inline <EOL> from . . syntax import ( If , Assign , While , Return , ForLoop , Var , TypedFn , Const , ExprStmt ) <EOL> from subst import subst_stmt_list <EOL> from transform import Transform <EOL> def replace_returns ( stmts , output_var ) : <EOL> """<STR_LIT>""" <EOL> new_stmts = [ ] <EOL> for stmt in stmts : <EOL> c = stmt . __class__ <EOL> if c is Return : <EOL> if output_var : <EOL> new_stmts . append ( Assign ( output_var , stmt . value ) ) <EOL> continue <EOL> if c is If : <EOL> stmt . true = replace_returns ( stmt . true , output_var ) <EOL> stmt . false = replace_returns ( stmt . false , output_var ) <EOL> elif c in ( ForLoop , While ) : <EOL> stmt . body = replace_returns ( stmt . body , output_var ) <EOL> new_stmts . append ( stmt ) <EOL> return new_stmts <EOL> def replace_return_with_var ( body , type_env , return_type ) : <EOL> result_name = names . fresh ( "<STR_LIT:result>" ) <EOL> type_env [ result_name ] = return_type <EOL> result_var = Var ( result_name , type = return_type ) <EOL> new_body = replace_returns ( body , result_var ) <EOL> return result_var , new_body <EOL> def do_inline ( src_fundef , args , dest_type_env , dest_block , result_var = None ) : <EOL> arg_names = src_fundef . arg_names <EOL> n_expected = len ( arg_names ) <EOL> n_given = len ( args ) <EOL> arg_str = "<STR_LIT:U+002C>" . join ( src_fundef . arg_names ) <EOL> assert n_expected == n_given , "<STR_LIT>" % ( src_fundef . name , n_expected , arg_str , n_given , args ) <EOL> rename_dict = { } <EOL> def rename_var ( old_name ) : <EOL> t = src_fundef . type_env [ old_name ] <EOL> new_name = names . refresh ( old_name ) <EOL> new_var = Var ( new_name , type = t ) <EOL> rename_dict [ old_name ] = new_var <EOL> dest_type_env [ new_name ] = t <EOL> return new_var <EOL> for ( old_formal , actual ) in zip ( arg_names , args ) : <EOL> if actual . __class__ in ( Var , Const ) : <EOL> rename_dict [ old_formal ] = actual <EOL> else : <EOL> new_var = rename_var ( old_formal ) <EOL> dest_block . append ( Assign ( new_var , actual ) ) <EOL> for old_name in src_fundef . type_env . iterkeys ( ) : <EOL> if old_name not in arg_names : <EOL> rename_var ( old_name ) <EOL> new_body = subst_stmt_list ( src_fundef . body , rename_dict ) <EOL> if result_var is None : <EOL> result_var , new_body = replace_return_with_var ( new_body , dest_type_env , src_fundef . return_type ) <EOL> else : <EOL> new_body = replace_returns ( new_body , result_var ) <EOL> dest_block . extend ( new_body ) <EOL> return result_var <EOL> class Inliner ( Transform ) : <EOL> def __init__ ( self ) : <EOL> Transform . __init__ ( self ) <EOL> self . count = <NUM_LIT:0> <EOL> def transform_TypedFn ( self , expr ) : <EOL> return Inliner ( ) . apply ( expr ) <EOL> def transform_ExprStmt ( self , stmt ) : <EOL> return Transform . transform_ExprStmt ( self , stmt ) <EOL> def transform_Call ( self , expr ) : <EOL> fn = self . transform_expr ( expr . fn ) <EOL> target = self . get_fn ( fn ) <EOL> if target . __class__ is TypedFn : <EOL> closure_args = self . closure_elts ( fn ) <EOL> if not can_inline ( target ) : <EOL> return expr <EOL> self . count += <NUM_LIT:1> <EOL> curr_block = self . blocks . current ( ) <EOL> combined_args = tuple ( closure_args ) + tuple ( expr . args ) <EOL> result_var = do_inline ( target , combined_args , <EOL> self . type_env , curr_block ) <EOL> return result_var <EOL> else : <EOL> return expr <EOL> def apply ( self , fn ) : <EOL> if contains_calls ( fn ) : <EOL> return Transform . apply ( self , fn ) <EOL> else : <EOL> return fn </s>
