Switch Yield from a Statement to an Expr

src/bytecode_vm/compiler/mod.rs

@@ -434,7 +434,6 @@ impl Compiler {
 
     fn compile_expr(&mut self, expr: &Expr) -> Result<Bytecode, CompilerError> {
         match expr {
-            Expr::NoOp => Ok(vec![]),
             Expr::None => {
                 let index = self.get_or_set_constant_index(Constant::None);
                 Ok(vec![Opcode::LoadConst(index)])

src/parser/mod.rs

@@ -599,6 +599,21 @@ impl Parser {
                     right: Box::new(right),
                 })
             }
+            Token::Yield => {
+                self.consume(Token::Yield)?;
+
+                if self.current_token == Token::From {
+                    self.consume(Token::From)?;
+                    let expr = self.parse_simple_expr()?;
+                    Ok(Expr::YieldFrom(Box::new(expr)))
+                // The [`Token::RParen`] can be found on generator lambdas.
+                } else if self.end_of_statement() || self.current_token == Token::RParen {
+                    Ok(Expr::Yield(None))
+                } else {
+                    let expr = self.parse_simple_expr()?;
+                    Ok(Expr::Yield(Some(Box::new(expr))))
+                }
+            }
             Token::Not => {
                 self.consume(Token::Not)?;
                 let right = self.parse_term()?;
@@ -1306,20 +1321,6 @@ impl Parser {
                 self.consume(Token::Continue)?;
                 Ok(Statement::Continue)
             }
-            Token::Yield => {
-                self.consume(Token::Yield)?;
-
-                if self.current_token == Token::From {
-                    self.consume(Token::From)?;
-                    let expr = self.parse_simple_expr()?;
-                    Ok(Statement::YieldFrom(expr))
-                } else if self.end_of_statement() {
-                    Ok(Statement::Yield(None))
-                } else {
-                    let expr = self.parse_simple_expr()?;
-                    Ok(Statement::Yield(Some(expr)))
-                }
-            }
             Token::Nonlocal => {
                 self.consume(Token::Nonlocal)?;
                 Ok(Statement::Nonlocal(self.parse_identifier()?))
@@ -1419,17 +1420,16 @@ impl Parser {
 
         let expr = if self.current_token == Token::LParen {
             self.consume(Token::LParen)?;
-            self.consume(Token::Yield)?;
+            let expr = self.parse_simple_expr()?;
             self.consume(Token::RParen)?;
-            Expr::NoOp
+            expr
         } else {
             self.parse_simple_expr()?
         };
 
         Ok(Expr::Lambda {
             args: Box::new(args),
-            expr: Box::new(expr.clone()),
-            is_generator: matches!(expr, Expr::NoOp),
+            expr: Box::new(expr),
         })
     }
 
@@ -1538,7 +1538,7 @@ impl Parser {
             match self.current_token {
                 Token::Comma => {
                     self.consume(Token::Comma)?;
-                    items.insert(item.clone(), Expr::NoOp);
+                    items.insert(item.clone(), Expr::None);
 
                     // Handle trailing comma
                     if self.current_token == Token::RBrace {
@@ -1548,7 +1548,7 @@ impl Parser {
                 }
                 Token::RBrace => {
                     self.consume(Token::RBrace)?;
-                    items.insert(item.clone(), Expr::NoOp);
+                    items.insert(item.clone(), Expr::None);
                     return Ok(Expr::Set(items.keys().cloned().collect()));
                 }
                 Token::For => {
@@ -2151,7 +2151,6 @@ lambda: 4
                 kwargs_var: None,
             }),
             expr: Box::new(Expr::Integer(4)),
-            is_generator: false,
         });
 
         match parse(&mut parser) {
@@ -2174,7 +2173,6 @@ lambda index: 4
                 kwargs_var: None,
             }),
             expr: Box::new(Expr::Integer(4)),
-            is_generator: false,
         });
 
         match parse(&mut parser) {
@@ -2203,7 +2201,6 @@ lambda index, val: 4
                 kwargs_var: None,
             }),
             expr: Box::new(Expr::Integer(4)),
-            is_generator: false,
         });
 
         match parse(&mut parser) {
@@ -2222,8 +2219,7 @@ lambda: (yield)
                 args_var: None,
                 kwargs_var: None,
             }),
-            expr: Box::new(Expr::NoOp),
-            is_generator: true,
+            expr: Box::new(Expr::Yield(None)),
         });
 
         match parse(&mut parser) {
@@ -2250,8 +2246,7 @@ lambda: (yield)
                     args_var: None,
                     kwargs_var: None,
                 }),
-                expr: Box::new(Expr::NoOp),
-                is_generator: true,
+                expr: Box::new(Expr::Yield(None)),
             })),
         });
 
@@ -3894,7 +3889,9 @@ def countdown(n):
                         },
                         body: Block {
                             statements: vec![
-                                Statement::Yield(Some(Expr::Variable("n".to_string()))),
+                                Statement::Expression(Expr::Yield(Some(Box::new(Expr::Variable(
+                                    "n".to_string(),
+                                ))))),
                                 Statement::Assignment {
                                     left: Expr::Variable("n".to_string()),
                                     right: Expr::BinaryOperation {
@@ -3920,7 +3917,8 @@ yield from a
 "#;
         let mut parser = init(input);
 
-        let expected_ast = Statement::YieldFrom(Expr::Variable("a".into()));
+        let expected_ast =
+            Statement::Expression(Expr::YieldFrom(Box::new(Expr::Variable("a".into()))));
 
         match parse(&mut parser) {
             Err(e) => panic!("Parser error: {:?}", e),

src/parser/static_analysis.rs

@@ -14,11 +14,17 @@ pub struct YieldDetector {
     pub found_yield: bool,
 }
 
+impl YieldDetector {
+    pub fn new() -> Self {
+        Self { found_yield: false }
+    }
+}
+
 impl Visitor for YieldDetector {
     fn visit_block(&mut self, _program: &Block) {}
 
     fn visit_statement(&mut self, statement: &Statement) {
-        if matches!(statement, Statement::Yield(_)) {
+        if matches!(statement, Statement::Expression(Expr::Yield(_))) {
             self.found_yield = true;
         }
     }
@@ -103,17 +109,13 @@ impl FunctionAnalysisVisitor {
     // variable reads, attribute accesses on objects.
     fn check_for_accessed_vars(&mut self, statement: &Statement) {
         match statement {
-            Statement::Expression(expr) => match expr {
-                Expr::FunctionCall { args, .. } => {
-                    for arg in args.args.iter() {
-                        if let Some(name) = arg.as_variable() {
-                            self.accessed_vars.push(name.into());
-                        }
+            Statement::Expression(Expr::FunctionCall { args, .. }) => {
+                for arg in args.args.iter() {
+                    if let Some(name) = arg.as_variable() {
+                        self.accessed_vars.push(name.into());
                     }
                 }
-                Expr::NoOp => {}
-                _ => {}
-            },
+            }
             Statement::Nonlocal(var) => {
                 self.accessed_vars.push(var.into());
             }

src/parser/types.rs

@@ -239,7 +239,6 @@ pub struct ForClause {
 
 #[derive(Debug, PartialEq, Clone)]
 pub enum Expr {
-    NoOp,
     None,
     NotImplemented,
     Ellipsis,
@@ -254,6 +253,8 @@ pub enum Expr {
     Dict(HashMap<Expr, Expr>),
     Tuple(Vec<Expr>),
     FString(Vec<FStringPart>),
+    Yield(Option<Box<Expr>>),
+    YieldFrom(Box<Expr>),
     BinaryOperation {
         left: Box<Expr>,
         op: BinOp,
@@ -324,7 +325,6 @@ pub enum Expr {
     Lambda {
         args: Box<ParsedArgDefinitions>,
         expr: Box<Expr>,
-        is_generator: bool,
     },
     TypeNode(TypeNode),
 }
@@ -419,6 +419,10 @@ impl Block {
         Self { statements }
     }
 
+    pub fn from_expr(expr: Expr) -> Self {
+        Self::new(vec![Statement::Expression(expr)])
+    }
+
     pub fn accept<V: Visitor>(&self, visitor: &mut V) {
         visitor.visit_block(self);
         for statement in &self.statements {
@@ -575,8 +579,6 @@ pub enum Statement {
         body: Block,
     },
     Return(Vec<Expr>),
-    Yield(Option<Expr>),
-    YieldFrom(Expr),
     Nonlocal(Variable),
     Global(Variable),
     IfElse {

src/treewalk/interpreter.rs

@@ -714,12 +714,8 @@ impl Interpreter {
         &self,
         arguments: &ParsedArgDefinitions,
         expr: &Expr,
-        is_generator: &bool,
     ) -> Result<ExprResult, InterpreterError> {
-        let block = match is_generator {
-            false => Block::new(vec![Statement::Expression(expr.clone())]),
-            true => Block::new(vec![Statement::Yield(None)]),
-        };
+        let block = Block::from_expr(expr.clone());
 
         let function = Container::new(Function::new_lambda(
             self.state.clone(),
@@ -1265,7 +1261,6 @@ impl Interpreter {
 
     pub fn evaluate_expr(&self, expr: &Expr) -> Result<ExprResult, InterpreterError> {
         match expr {
-            Expr::NoOp => unreachable!(),
             Expr::None => Ok(ExprResult::None),
             Expr::Ellipsis => Ok(ExprResult::Ellipsis),
             Expr::NotImplemented => Ok(ExprResult::NotImplemented),
@@ -1325,12 +1320,10 @@ impl Interpreter {
                 self.evaluate_slice_operation(object, params)
             }
             Expr::FString(parts) => self.evaluate_f_string(parts),
-            Expr::Lambda {
-                args,
-                expr,
-                is_generator,
-            } => self.evaluate_lambda(args, expr, is_generator),
+            Expr::Lambda { args, expr } => self.evaluate_lambda(args, expr),
             Expr::TypeNode(type_node) => self.evaluate_type_node(type_node),
+            // This is unreachable because it should be handled inside `GeneratorExecutor`.
+            Expr::Yield(_) | Expr::YieldFrom(_) => unreachable!(),
         }
     }
 
@@ -1345,8 +1338,6 @@ impl Interpreter {
         let result = match stmt {
             // These are handled above
             Statement::Expression(_) | Statement::Return(_) => unreachable!(),
-            // This is unreachable because it should be handled inside `GeneratorExecutor`.
-            Statement::Yield(_) | Statement::YieldFrom(_) => unreachable!(),
             Statement::Pass => Ok(()),
             Statement::Break => Err(InterpreterError::EncounteredBreak),
             Statement::Continue => Err(InterpreterError::EncounteredContinue),
@@ -1924,7 +1915,7 @@ y = _f.__type_params__
                         body: Block { ref statements, .. },
                         ..
                     } if statements.len() == 1 &&
-                        matches!(statements[0], Statement::Yield(None))
+                        matches!(statements[0], Statement::Expression(Expr::Yield(None)))
                 ));
                 assert!(matches!(
                     interpreter.state.read("e").unwrap(),

src/treewalk/types/function.rs

@@ -107,7 +107,7 @@ impl Function {
     }
 
     pub fn is_generator(&self) -> bool {
-        let mut detector = YieldDetector { found_yield: false };
+        let mut detector = YieldDetector::new();
         self.body.accept(&mut detector);
         detector.found_yield
     }

src/treewalk/types/generator.rs

@@ -45,7 +45,7 @@ impl Generator {
         if let Some((first_clause, remaining_clauses)) = clauses.split_first() {
             let loop_body = if remaining_clauses.is_empty() {
                 // Base case: Yield the body
-                Block::new(vec![Statement::Yield(Some(body.clone()))])
+                Block::from_expr(Expr::Yield(Some(Box::new(body.clone()))))
             } else {
                 // Recursive case: Build nested loop for the remaining clauses
                 Self::build_nested_loops(body, remaining_clauses)
@@ -129,9 +129,11 @@ impl Container<Generator> {
     ) -> Result<Option<ExprResult>, InterpreterError> {
         if !control_flow {
             match statement {
-                Statement::Yield(None) => Ok(None),
-                Statement::Yield(Some(expr)) => Ok(Some(interpreter.evaluate_expr(&expr)?)),
-                Statement::YieldFrom(_) => unimplemented!(),
+                Statement::Expression(Expr::Yield(None)) => Ok(None),
+                Statement::Expression(Expr::Yield(Some(expr))) => {
+                    Ok(Some(interpreter.evaluate_expr(&expr)?))
+                }
+                Statement::Expression(Expr::YieldFrom(_)) => unimplemented!(),
                 _ => {
                     // would we ever need to support a return statement here?
                     let _ = interpreter.evaluate_statement(&statement)?;