@@ -17,6 +17,7 @@ struct Compiler {
1717 source_path : Option < String > ,
1818 current_source_location : ast:: Location ,
1919 in_loop : bool ,
20+ in_function_def : bool ,
2021}
2122
2223/// Compile a given sourcecode into a bytecode object.
@@ -75,6 +76,7 @@ impl Compiler {
7576 source_path : None ,
7677 current_source_location : ast:: Location :: default ( ) ,
7778 in_loop : false ,
79+ in_function_def : false ,
7880 }
7981 }
8082
@@ -232,9 +234,10 @@ impl Compiler {
232234
233235 self . compile_test ( test, None , Some ( else_label) , EvalContext :: Statement ) ?;
234236
237+ let was_in_loop = self . in_loop ;
235238 self . in_loop = true ;
236239 self . compile_statements ( body) ?;
237- self . in_loop = false ;
240+ self . in_loop = was_in_loop ;
238241 self . emit ( Instruction :: Jump {
239242 target : start_label,
240243 } ) ;
@@ -295,10 +298,10 @@ impl Compiler {
295298 // Start of loop iteration, set targets:
296299 self . compile_store ( target) ?;
297300
298- // Body of loop:
301+ let was_in_loop = self . in_loop ;
299302 self . in_loop = true ;
300303 self . compile_statements ( body) ?;
301- self . in_loop = false ;
304+ self . in_loop = was_in_loop ;
302305
303306 self . emit ( Instruction :: Jump {
304307 target : start_label,
@@ -431,6 +434,11 @@ impl Compiler {
431434 decorator_list,
432435 } => {
433436 // Create bytecode for this function:
437+ // remember to restore self.in_loop to the original after the function is compiled
438+ let was_in_loop = self . in_loop ;
439+ let was_in_function_def = self . in_function_def ;
440+ self . in_loop = false ;
441+ self . in_function_def = true ;
434442 let flags = self . enter_function ( name, args) ?;
435443 self . compile_statements ( body) ?;
436444
@@ -458,6 +466,8 @@ impl Compiler {
458466 self . emit ( Instruction :: StoreName {
459467 name : name. to_string ( ) ,
460468 } ) ;
469+ self . in_loop = was_in_loop;
470+ self . in_function_def = was_in_function_def;
461471 }
462472 ast:: Statement :: ClassDef {
463473 name,
@@ -466,6 +476,8 @@ impl Compiler {
466476 keywords,
467477 decorator_list,
468478 } => {
479+ let was_in_loop = self . in_loop ;
480+ self . in_loop = false ;
469481 self . prepare_decorators ( decorator_list) ?;
470482 self . emit ( Instruction :: LoadBuildClass ) ;
471483 let line_number = self . get_source_line_number ( ) ;
@@ -546,6 +558,7 @@ impl Compiler {
546558 self . emit ( Instruction :: StoreName {
547559 name : name. to_string ( ) ,
548560 } ) ;
561+ self . in_loop = was_in_loop;
549562 }
550563 ast:: Statement :: Assert { test, msg } => {
551564 // TODO: if some flag, ignore all assert statements!
@@ -584,6 +597,9 @@ impl Compiler {
584597 self . emit ( Instruction :: Continue ) ;
585598 }
586599 ast:: Statement :: Return { value } => {
600+ if !self . in_function_def {
601+ return Err ( CompileError :: InvalidReturn ) ;
602+ }
587603 match value {
588604 Some ( e) => {
589605 let size = e. len ( ) ;
@@ -663,7 +679,6 @@ impl Compiler {
663679 name : & str ,
664680 args : & ast:: Parameters ,
665681 ) -> Result < bytecode:: FunctionOpArg , CompileError > {
666- self . in_loop = false ;
667682 let have_kwargs = !args. defaults . is_empty ( ) ;
668683 if have_kwargs {
669684 // Construct a tuple:
@@ -971,6 +986,9 @@ impl Compiler {
971986 self . emit ( Instruction :: BuildSlice { size } ) ;
972987 }
973988 ast:: Expression :: Yield { value } => {
989+ if !self . in_function_def {
990+ return Err ( CompileError :: InvalidYield ) ;
991+ }
974992 self . mark_generator ( ) ;
975993 match value {
976994 Some ( expression) => self . compile_expression ( expression) ?,
@@ -1021,6 +1039,7 @@ impl Compiler {
10211039 }
10221040 ast:: Expression :: Lambda { args, body } => {
10231041 let name = "<lambda>" . to_string ( ) ;
1042+ // no need to worry about the self.loop_depth because there are no loops in lambda expressions
10241043 let flags = self . enter_function ( & name, args) ?;
10251044 self . compile_expression ( body) ?;
10261045 self . emit ( Instruction :: ReturnValue ) ;
@@ -1369,10 +1388,11 @@ impl Compiler {
13691388
13701389 // Low level helper functions:
13711390 fn emit ( & mut self , instruction : Instruction ) {
1372- self . current_code_object ( ) . instructions . push ( instruction) ;
1373- // TODO: insert source filename
13741391 let location = self . current_source_location . clone ( ) ;
1375- self . current_code_object ( ) . locations . push ( location) ;
1392+ let mut cur_code_obj = self . current_code_object ( ) ;
1393+ cur_code_obj. instructions . push ( instruction) ;
1394+ cur_code_obj. locations . push ( location) ;
1395+ // TODO: insert source filename
13761396 }
13771397
13781398 fn current_code_object ( & mut self ) -> & mut CodeObject {
@@ -1413,6 +1433,7 @@ mod tests {
14131433 use crate :: bytecode:: Constant :: * ;
14141434 use crate :: bytecode:: Instruction :: * ;
14151435 use rustpython_parser:: parser;
1436+
14161437 fn compile_exec ( source : & str ) -> CodeObject {
14171438 let mut compiler = Compiler :: new ( ) ;
14181439 compiler. source_path = Some ( "source_path" . to_string ( ) ) ;
0 commit comments