depyler_analysis/
generator_state.rs

1//! Generator State Analysis
2//!
3//! Analyzes generator functions to determine:
4//! - Which local variables need to be preserved across yields
5//! - Yield points and control flow
6//! - State machine structure
7
8use depyler_hir::hir::{HirExpr, HirFunction, HirStmt, Type};
9use std::collections::HashSet;
10
11/// Information about a generator's state requirements
12#[derive(Debug, Clone)]
13pub struct GeneratorStateInfo {
14    /// Local variables that need to be in the state struct
15    pub state_variables: Vec<StateVariable>,
16    /// Parameters that are used after yield points
17    pub captured_params: Vec<String>,
18    /// Number of yield points in the function
19    pub yield_count: usize,
20    /// Whether the generator has loops
21    pub has_loops: bool,
22}
23
24#[derive(Debug, Clone)]
25pub struct StateVariable {
26    pub name: String,
27    pub ty: Type,
28}
29
30impl GeneratorStateInfo {
31    /// Analyze a generator function to determine state requirements
32    pub fn analyze(func: &HirFunction) -> Self {
33        let mut analyzer = StateAnalyzer {
34            state_variables: Vec::new(),
35            captured_params: HashSet::new(),
36            yield_count: 0,
37            has_loops: false,
38            declared_vars: HashSet::new(),
39        };
40
41        analyzer.analyze_statements(&func.body);
42
43        // Convert param names to Vec
44        let captured_params: Vec<String> = func
45            .params
46            .iter()
47            .filter(|p| analyzer.captured_params.contains(&p.name))
48            .map(|p| p.name.clone())
49            .collect();
50
51        GeneratorStateInfo {
52            state_variables: analyzer.state_variables,
53            captured_params,
54            yield_count: analyzer.yield_count,
55            has_loops: analyzer.has_loops,
56        }
57    }
58}
59
60struct StateAnalyzer {
61    state_variables: Vec<StateVariable>,
62    captured_params: HashSet<String>,
63    yield_count: usize,
64    has_loops: bool,
65    declared_vars: HashSet<String>,
66}
67
68impl StateAnalyzer {
69    fn analyze_statements(&mut self, stmts: &[HirStmt]) {
70        for stmt in stmts {
71            self.analyze_statement(stmt);
72        }
73    }
74
75    fn analyze_statement(&mut self, stmt: &HirStmt) {
76        match stmt {
77            HirStmt::Assign {
78                target,
79                value,
80                type_annotation,
81            } => {
82                self.analyze_assign(target, value, type_annotation);
83            }
84            HirStmt::For { iter, body, .. } => {
85                self.analyze_for_loop(iter, body);
86            }
87            HirStmt::While { condition, body } => {
88                self.analyze_while_loop(condition, body);
89            }
90            HirStmt::If {
91                condition,
92                then_body,
93                else_body,
94            } => {
95                self.analyze_if_stmt(condition, then_body, else_body);
96            }
97            HirStmt::Expr(expr) | HirStmt::Return(Some(expr)) => {
98                self.analyze_expression(expr);
99            }
100            // DEPYLER-0561: Analyze with statements for generator state capture
101            HirStmt::With { context, body, .. } => {
102                self.analyze_expression(context);
103                self.analyze_statements(body);
104            }
105            HirStmt::Try {
106                body,
107                handlers,
108                orelse,
109                finalbody,
110            } => {
111                self.analyze_statements(body);
112                for handler in handlers {
113                    self.analyze_statements(&handler.body);
114                }
115                if let Some(else_stmts) = orelse {
116                    self.analyze_statements(else_stmts);
117                }
118                if let Some(final_stmts) = finalbody {
119                    self.analyze_statements(final_stmts);
120                }
121            }
122            _ => {}
123        }
124    }
125
126    /// DEPYLER-0258: Infer type from value expression when no annotation provided
127    /// Complexity: 8 (within ≤10 target)
128    fn infer_type_from_expression(expr: &HirExpr) -> Type {
129        match expr {
130            HirExpr::Literal(lit) => match lit {
131                depyler_hir::hir::Literal::Int(_) => Type::Int,
132                depyler_hir::hir::Literal::Float(_) => Type::Float,
133                depyler_hir::hir::Literal::String(_) => Type::String,
134                depyler_hir::hir::Literal::Bytes(_) => Type::Custom("bytes".to_string()),
135                depyler_hir::hir::Literal::Bool(_) => Type::Bool,
136                depyler_hir::hir::Literal::None => Type::None,
137            },
138            HirExpr::List(items) => {
139                // Infer element type from first item
140                let elem_type = items
141                    .first()
142                    .map(Self::infer_type_from_expression)
143                    .unwrap_or(Type::Unknown);
144                Type::List(Box::new(elem_type))
145            }
146            HirExpr::Dict(_) => Type::Dict(Box::new(Type::String), Box::new(Type::Unknown)),
147            HirExpr::Set(_) => Type::Set(Box::new(Type::Unknown)),
148            // For complex expressions, default to Unknown
149            _ => Type::Unknown,
150        }
151    }
152
153    /// DEPYLER-0494: Analyze assignment target (handles both Symbol and Tuple)
154    /// Complexity: 7 (within ≤10 target)
155    fn analyze_assign(
156        &mut self,
157        target: &depyler_hir::hir::AssignTarget,
158        value: &HirExpr,
159        type_annotation: &Option<Type>,
160    ) {
161        match target {
162            depyler_hir::hir::AssignTarget::Symbol(name) => {
163                let name_str = name.as_str();
164                if !self.declared_vars.contains(name_str) {
165                    self.declared_vars.insert(name_str.to_string());
166                    // DEPYLER-0258 FIX: Infer type from value expression if no annotation
167                    let ty = type_annotation
168                        .clone()
169                        .unwrap_or_else(|| Self::infer_type_from_expression(value));
170                    self.state_variables.push(StateVariable {
171                        name: name_str.to_string(),
172                        ty,
173                    });
174                }
175            }
176            // DEPYLER-0494 FIX: Handle tuple unpacking (a, b) = (0, 1)
177            depyler_hir::hir::AssignTarget::Tuple(targets) => {
178                // Infer element types from tuple value if possible
179                let element_types = if let HirExpr::Tuple(values) = value {
180                    // Parallel tuple: (a, b) = (val1, val2)
181                    values
182                        .iter()
183                        .map(Self::infer_type_from_expression)
184                        .collect::<Vec<_>>()
185                } else {
186                    // Can't infer from value, use Unknown
187                    vec![]
188                };
189
190                for (idx, target_elem) in targets.iter().enumerate() {
191                    if let depyler_hir::hir::AssignTarget::Symbol(name) = target_elem {
192                        let name_str = name.as_str();
193                        if !self.declared_vars.contains(name_str) {
194                            self.declared_vars.insert(name_str.to_string());
195                            // Try to get type from parallel value or use Unknown
196                            let ty = element_types.get(idx).cloned().unwrap_or(Type::Unknown);
197                            self.state_variables.push(StateVariable {
198                                name: name_str.to_string(),
199                                ty,
200                            });
201                        }
202                    }
203                }
204            }
205            _ => {} // Index and Attribute assignments don't declare new variables
206        }
207        self.analyze_expression(value);
208    }
209
210    fn analyze_for_loop(&mut self, iter: &HirExpr, body: &[HirStmt]) {
211        self.has_loops = true;
212        self.analyze_expression(iter);
213        self.analyze_statements(body);
214    }
215
216    fn analyze_while_loop(&mut self, condition: &HirExpr, body: &[HirStmt]) {
217        self.has_loops = true;
218        self.analyze_expression(condition);
219        self.analyze_statements(body);
220    }
221
222    fn analyze_if_stmt(
223        &mut self,
224        condition: &HirExpr,
225        then_body: &[HirStmt],
226        else_body: &Option<Vec<HirStmt>>,
227    ) {
228        self.analyze_expression(condition);
229        self.analyze_statements(then_body);
230        if let Some(else_stmts) = else_body {
231            self.analyze_statements(else_stmts);
232        }
233    }
234
235    fn analyze_expression(&mut self, expr: &HirExpr) {
236        match expr {
237            HirExpr::Yield { value } => self.analyze_yield(value),
238            HirExpr::Var(name) => self.analyze_variable(name),
239            HirExpr::Binary { left, right, .. } => self.analyze_binary(left, right),
240            HirExpr::Unary { operand, .. } => self.analyze_expression(operand),
241            HirExpr::Call { args, .. } | HirExpr::List(args) | HirExpr::Tuple(args) => {
242                self.analyze_expressions(args);
243            }
244            HirExpr::Index { base, index } => self.analyze_binary(base, index),
245            HirExpr::MethodCall { object, args, .. } => {
246                self.analyze_expression(object);
247                self.analyze_expressions(args);
248            }
249            _ => {}
250        }
251    }
252
253    fn analyze_yield(&mut self, value: &Option<Box<HirExpr>>) {
254        self.yield_count += 1;
255        if let Some(v) = value {
256            self.analyze_expression(v);
257        }
258    }
259
260    fn analyze_variable(&mut self, name: &str) {
261        let name_str = name;
262        if !self.declared_vars.contains(name_str) {
263            self.captured_params.insert(name_str.to_string());
264        }
265    }
266
267    fn analyze_binary(&mut self, left: &HirExpr, right: &HirExpr) {
268        self.analyze_expression(left);
269        self.analyze_expression(right);
270    }
271
272    fn analyze_expressions(&mut self, exprs: &[HirExpr]) {
273        for expr in exprs {
274            self.analyze_expression(expr);
275        }
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use depyler_hir::hir::{BinOp, ExceptHandler, FunctionProperties, HirParam, Literal, UnaryOp};
283    use depyler_annotations::TranspilationAnnotations;
284    use smallvec::smallvec;
285
286    // ============================================
287    // Helper function for creating test functions
288    // ============================================
289
290    fn make_func(name: &str, params: Vec<HirParam>, body: Vec<HirStmt>) -> HirFunction {
291        HirFunction {
292            name: name.to_string(),
293            params: params.into(),
294            ret_type: Type::Int,
295            body,
296            properties: FunctionProperties::default(),
297            annotations: TranspilationAnnotations::default(),
298            docstring: None,
299        }
300    }
301
302    fn make_param(name: &str, ty: Type) -> HirParam {
303        HirParam::new(name.to_string(), ty)
304    }
305
306    fn make_assign(name: &str, value: HirExpr, ty: Option<Type>) -> HirStmt {
307        HirStmt::Assign {
308            target: depyler_hir::hir::AssignTarget::Symbol(name.to_string()),
309            value,
310            type_annotation: ty,
311        }
312    }
313
314    fn make_yield(value: Option<HirExpr>) -> HirExpr {
315        HirExpr::Yield {
316            value: value.map(Box::new),
317        }
318    }
319
320    // ============================================
321    // GeneratorStateInfo struct tests
322    // ============================================
323
324    #[test]
325    fn test_generator_state_info_clone() {
326        let info = GeneratorStateInfo {
327            state_variables: vec![StateVariable {
328                name: "x".to_string(),
329                ty: Type::Int,
330            }],
331            captured_params: vec!["n".to_string()],
332            yield_count: 2,
333            has_loops: true,
334        };
335        let cloned = info.clone();
336        assert_eq!(cloned.state_variables.len(), 1);
337        assert_eq!(cloned.yield_count, 2);
338        assert!(cloned.has_loops);
339    }
340
341    #[test]
342    fn test_generator_state_info_debug() {
343        let info = GeneratorStateInfo {
344            state_variables: vec![],
345            captured_params: vec![],
346            yield_count: 0,
347            has_loops: false,
348        };
349        let debug_str = format!("{:?}", info);
350        assert!(debug_str.contains("GeneratorStateInfo"));
351    }
352
353    // ============================================
354    // StateVariable struct tests
355    // ============================================
356
357    #[test]
358    fn test_state_variable_clone() {
359        let var = StateVariable {
360            name: "counter".to_string(),
361            ty: Type::Float,
362        };
363        let cloned = var.clone();
364        assert_eq!(cloned.name, "counter");
365        assert!(matches!(cloned.ty, Type::Float));
366    }
367
368    #[test]
369    fn test_state_variable_debug() {
370        let var = StateVariable {
371            name: "x".to_string(),
372            ty: Type::Bool,
373        };
374        let debug_str = format!("{:?}", var);
375        assert!(debug_str.contains("StateVariable"));
376        assert!(debug_str.contains("x"));
377    }
378
379    // ============================================
380    // Empty and minimal function tests
381    // ============================================
382
383    #[test]
384    fn test_analyze_empty_function() {
385        let func = make_func("empty", vec![], vec![]);
386        let info = GeneratorStateInfo::analyze(&func);
387
388        assert_eq!(info.state_variables.len(), 0);
389        assert_eq!(info.captured_params.len(), 0);
390        assert_eq!(info.yield_count, 0);
391        assert!(!info.has_loops);
392    }
393
394    #[test]
395    fn test_analyze_function_with_only_yield() {
396        let func = make_func(
397            "simple_yield",
398            vec![],
399            vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
400                Literal::Int(42),
401            ))))],
402        );
403        let info = GeneratorStateInfo::analyze(&func);
404
405        assert_eq!(info.yield_count, 1);
406        assert!(!info.has_loops);
407    }
408
409    #[test]
410    fn test_analyze_multiple_yields() {
411        let func = make_func(
412            "multi_yield",
413            vec![],
414            vec![
415                HirStmt::Expr(make_yield(Some(HirExpr::Literal(Literal::Int(1))))),
416                HirStmt::Expr(make_yield(Some(HirExpr::Literal(Literal::Int(2))))),
417                HirStmt::Expr(make_yield(Some(HirExpr::Literal(Literal::Int(3))))),
418            ],
419        );
420        let info = GeneratorStateInfo::analyze(&func);
421
422        assert_eq!(info.yield_count, 3);
423    }
424
425    #[test]
426    fn test_analyze_yield_none() {
427        let func = make_func("yield_none", vec![], vec![HirStmt::Expr(make_yield(None))]);
428        let info = GeneratorStateInfo::analyze(&func);
429
430        assert_eq!(info.yield_count, 1);
431    }
432
433    // ============================================
434    // Type inference tests
435    // ============================================
436
437    #[test]
438    fn test_infer_type_int_literal() {
439        let func = make_func(
440            "int_infer",
441            vec![],
442            vec![make_assign("x", HirExpr::Literal(Literal::Int(42)), None)],
443        );
444        let info = GeneratorStateInfo::analyze(&func);
445
446        assert_eq!(info.state_variables.len(), 1);
447        assert_eq!(info.state_variables[0].name, "x");
448        assert!(matches!(info.state_variables[0].ty, Type::Int));
449    }
450
451    #[test]
452    fn test_infer_type_float_literal() {
453        let func = make_func(
454            "float_infer",
455            vec![],
456            vec![make_assign(
457                "y",
458                HirExpr::Literal(Literal::Float(3.15)),
459                None,
460            )],
461        );
462        let info = GeneratorStateInfo::analyze(&func);
463
464        assert!(matches!(info.state_variables[0].ty, Type::Float));
465    }
466
467    #[test]
468    fn test_infer_type_string_literal() {
469        let func = make_func(
470            "string_infer",
471            vec![],
472            vec![make_assign(
473                "s",
474                HirExpr::Literal(Literal::String("hello".to_string())),
475                None,
476            )],
477        );
478        let info = GeneratorStateInfo::analyze(&func);
479
480        assert!(matches!(info.state_variables[0].ty, Type::String));
481    }
482
483    #[test]
484    fn test_infer_type_bool_literal() {
485        let func = make_func(
486            "bool_infer",
487            vec![],
488            vec![make_assign(
489                "b",
490                HirExpr::Literal(Literal::Bool(true)),
491                None,
492            )],
493        );
494        let info = GeneratorStateInfo::analyze(&func);
495
496        assert!(matches!(info.state_variables[0].ty, Type::Bool));
497    }
498
499    #[test]
500    fn test_infer_type_none_literal() {
501        let func = make_func(
502            "none_infer",
503            vec![],
504            vec![make_assign("n", HirExpr::Literal(Literal::None), None)],
505        );
506        let info = GeneratorStateInfo::analyze(&func);
507
508        assert!(matches!(info.state_variables[0].ty, Type::None));
509    }
510
511    #[test]
512    fn test_infer_type_bytes_literal() {
513        let func = make_func(
514            "bytes_infer",
515            vec![],
516            vec![make_assign(
517                "data",
518                HirExpr::Literal(Literal::Bytes(vec![1, 2, 3])),
519                None,
520            )],
521        );
522        let info = GeneratorStateInfo::analyze(&func);
523
524        assert!(matches!(info.state_variables[0].ty, Type::Custom(ref s) if s == "bytes"));
525    }
526
527    #[test]
528    fn test_infer_type_list_with_elements() {
529        let func = make_func(
530            "list_infer",
531            vec![],
532            vec![make_assign(
533                "items",
534                HirExpr::List(vec![
535                    HirExpr::Literal(Literal::Int(1)),
536                    HirExpr::Literal(Literal::Int(2)),
537                ]),
538                None,
539            )],
540        );
541        let info = GeneratorStateInfo::analyze(&func);
542
543        if let Type::List(inner) = &info.state_variables[0].ty {
544            assert!(matches!(**inner, Type::Int));
545        } else {
546            panic!("Expected List type");
547        }
548    }
549
550    #[test]
551    fn test_infer_type_empty_list() {
552        let func = make_func(
553            "empty_list",
554            vec![],
555            vec![make_assign("items", HirExpr::List(vec![]), None)],
556        );
557        let info = GeneratorStateInfo::analyze(&func);
558
559        if let Type::List(inner) = &info.state_variables[0].ty {
560            assert!(matches!(**inner, Type::Unknown));
561        } else {
562            panic!("Expected List type");
563        }
564    }
565
566    #[test]
567    fn test_infer_type_dict() {
568        let func = make_func(
569            "dict_infer",
570            vec![],
571            vec![make_assign("d", HirExpr::Dict(vec![]), None)],
572        );
573        let info = GeneratorStateInfo::analyze(&func);
574
575        assert!(matches!(info.state_variables[0].ty, Type::Dict(_, _)));
576    }
577
578    #[test]
579    fn test_infer_type_set() {
580        let func = make_func(
581            "set_infer",
582            vec![],
583            vec![make_assign(
584                "s",
585                HirExpr::Set(vec![HirExpr::Literal(Literal::Int(1))]),
586                None,
587            )],
588        );
589        let info = GeneratorStateInfo::analyze(&func);
590
591        assert!(matches!(info.state_variables[0].ty, Type::Set(_)));
592    }
593
594    #[test]
595    fn test_infer_type_complex_expression() {
596        let func = make_func(
597            "complex",
598            vec![],
599            vec![make_assign(
600                "z",
601                HirExpr::Binary {
602                    left: Box::new(HirExpr::Var("a".to_string())),
603                    op: BinOp::Add,
604                    right: Box::new(HirExpr::Var("b".to_string())),
605                },
606                None,
607            )],
608        );
609        let info = GeneratorStateInfo::analyze(&func);
610
611        // Complex expressions default to Unknown
612        assert!(matches!(info.state_variables[0].ty, Type::Unknown));
613    }
614
615    #[test]
616    fn test_explicit_type_annotation_overrides_inference() {
617        let func = make_func(
618            "explicit",
619            vec![],
620            vec![make_assign(
621                "x",
622                HirExpr::Literal(Literal::Int(42)),
623                Some(Type::Float), // Explicit annotation
624            )],
625        );
626        let info = GeneratorStateInfo::analyze(&func);
627
628        // Explicit annotation takes precedence
629        assert!(matches!(info.state_variables[0].ty, Type::Float));
630    }
631
632    // ============================================
633    // Loop detection tests
634    // ============================================
635
636    #[test]
637    fn test_for_loop_detection() {
638        let func = make_func(
639            "for_gen",
640            vec![],
641            vec![HirStmt::For {
642                target: depyler_hir::hir::AssignTarget::Symbol("i".to_string()),
643                iter: HirExpr::List(vec![]),
644                body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Var(
645                    "i".to_string(),
646                ))))],
647            }],
648        );
649        let info = GeneratorStateInfo::analyze(&func);
650
651        assert!(info.has_loops);
652        assert_eq!(info.yield_count, 1);
653    }
654
655    #[test]
656    fn test_while_loop_detection() {
657        let func = make_func(
658            "while_gen",
659            vec![],
660            vec![HirStmt::While {
661                condition: HirExpr::Literal(Literal::Bool(true)),
662                body: vec![HirStmt::Expr(make_yield(None))],
663            }],
664        );
665        let info = GeneratorStateInfo::analyze(&func);
666
667        assert!(info.has_loops);
668    }
669
670    #[test]
671    fn test_nested_loops() {
672        let func = make_func(
673            "nested",
674            vec![],
675            vec![HirStmt::For {
676                target: depyler_hir::hir::AssignTarget::Symbol("i".to_string()),
677                iter: HirExpr::List(vec![]),
678                body: vec![HirStmt::For {
679                    target: depyler_hir::hir::AssignTarget::Symbol("j".to_string()),
680                    iter: HirExpr::List(vec![]),
681                    body: vec![HirStmt::Expr(make_yield(None))],
682                }],
683            }],
684        );
685        let info = GeneratorStateInfo::analyze(&func);
686
687        assert!(info.has_loops);
688    }
689
690    // ============================================
691    // Parameter capture tests
692    // ============================================
693
694    #[test]
695    fn test_capture_parameter_used_in_expression() {
696        let func = make_func(
697            "param_capture",
698            vec![make_param("n", Type::Int)],
699            vec![HirStmt::Expr(make_yield(Some(HirExpr::Var(
700                "n".to_string(),
701            ))))],
702        );
703        let info = GeneratorStateInfo::analyze(&func);
704
705        assert!(info.captured_params.contains(&"n".to_string()));
706    }
707
708    #[test]
709    fn test_unused_parameter_not_captured() {
710        let func = make_func(
711            "unused_param",
712            vec![make_param("unused", Type::Int)],
713            vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
714                Literal::Int(42),
715            ))))],
716        );
717        let info = GeneratorStateInfo::analyze(&func);
718
719        assert!(info.captured_params.is_empty());
720    }
721
722    #[test]
723    fn test_multiple_params_partial_capture() {
724        let func = make_func(
725            "partial",
726            vec![
727                make_param("a", Type::Int),
728                make_param("b", Type::Int),
729                make_param("c", Type::Int),
730            ],
731            vec![HirStmt::Expr(make_yield(Some(HirExpr::Binary {
732                left: Box::new(HirExpr::Var("a".to_string())),
733                op: BinOp::Add,
734                right: Box::new(HirExpr::Var("c".to_string())),
735            })))],
736        );
737        let info = GeneratorStateInfo::analyze(&func);
738
739        assert!(info.captured_params.contains(&"a".to_string()));
740        assert!(info.captured_params.contains(&"c".to_string()));
741        assert!(!info.captured_params.contains(&"b".to_string()));
742    }
743
744    // ============================================
745    // State variable tracking tests
746    // ============================================
747
748    #[test]
749    fn test_variable_declared_before_use_not_param() {
750        let func = make_func(
751            "declared",
752            vec![],
753            vec![
754                make_assign("x", HirExpr::Literal(Literal::Int(0)), Some(Type::Int)),
755                HirStmt::Expr(make_yield(Some(HirExpr::Var("x".to_string())))),
756            ],
757        );
758        let info = GeneratorStateInfo::analyze(&func);
759
760        assert!(info.captured_params.is_empty());
761        assert_eq!(info.state_variables.len(), 1);
762        assert_eq!(info.state_variables[0].name, "x");
763    }
764
765    #[test]
766    fn test_variable_reassignment_not_duplicated() {
767        let func = make_func(
768            "reassign",
769            vec![],
770            vec![
771                make_assign("x", HirExpr::Literal(Literal::Int(0)), Some(Type::Int)),
772                make_assign("x", HirExpr::Literal(Literal::Int(1)), None),
773                make_assign("x", HirExpr::Literal(Literal::Int(2)), None),
774            ],
775        );
776        let info = GeneratorStateInfo::analyze(&func);
777
778        // Only one state variable entry despite multiple assignments
779        assert_eq!(info.state_variables.len(), 1);
780    }
781
782    // ============================================
783    // Tuple unpacking tests (DEPYLER-0494)
784    // ============================================
785
786    #[test]
787    fn test_tuple_unpacking_simple() {
788        let func = make_func(
789            "tuple_unpack",
790            vec![],
791            vec![HirStmt::Assign {
792                target: depyler_hir::hir::AssignTarget::Tuple(vec![
793                    depyler_hir::hir::AssignTarget::Symbol("a".to_string()),
794                    depyler_hir::hir::AssignTarget::Symbol("b".to_string()),
795                ]),
796                value: HirExpr::Tuple(vec![
797                    HirExpr::Literal(Literal::Int(1)),
798                    HirExpr::Literal(Literal::Int(2)),
799                ]),
800                type_annotation: None,
801            }],
802        );
803        let info = GeneratorStateInfo::analyze(&func);
804
805        assert_eq!(info.state_variables.len(), 2);
806        let names: Vec<&str> = info
807            .state_variables
808            .iter()
809            .map(|v| v.name.as_str())
810            .collect();
811        assert!(names.contains(&"a"));
812        assert!(names.contains(&"b"));
813    }
814
815    #[test]
816    fn test_tuple_unpacking_type_inference() {
817        let func = make_func(
818            "tuple_types",
819            vec![],
820            vec![HirStmt::Assign {
821                target: depyler_hir::hir::AssignTarget::Tuple(vec![
822                    depyler_hir::hir::AssignTarget::Symbol("x".to_string()),
823                    depyler_hir::hir::AssignTarget::Symbol("y".to_string()),
824                ]),
825                value: HirExpr::Tuple(vec![
826                    HirExpr::Literal(Literal::Int(42)),
827                    HirExpr::Literal(Literal::String("hello".to_string())),
828                ]),
829                type_annotation: None,
830            }],
831        );
832        let info = GeneratorStateInfo::analyze(&func);
833
834        let x_var = info.state_variables.iter().find(|v| v.name == "x").unwrap();
835        let y_var = info.state_variables.iter().find(|v| v.name == "y").unwrap();
836        assert!(matches!(x_var.ty, Type::Int));
837        assert!(matches!(y_var.ty, Type::String));
838    }
839
840    #[test]
841    fn test_tuple_unpacking_non_tuple_value() {
842        let func = make_func(
843            "non_tuple",
844            vec![],
845            vec![HirStmt::Assign {
846                target: depyler_hir::hir::AssignTarget::Tuple(vec![
847                    depyler_hir::hir::AssignTarget::Symbol("a".to_string()),
848                    depyler_hir::hir::AssignTarget::Symbol("b".to_string()),
849                ]),
850                value: HirExpr::Var("some_tuple".to_string()), // Not a literal tuple
851                type_annotation: None,
852            }],
853        );
854        let info = GeneratorStateInfo::analyze(&func);
855
856        // Can't infer types, should be Unknown
857        for var in &info.state_variables {
858            assert!(matches!(var.ty, Type::Unknown));
859        }
860    }
861
862    // ============================================
863    // If statement analysis tests
864    // ============================================
865
866    #[test]
867    fn test_if_statement_then_only() {
868        let func = make_func(
869            "if_then",
870            vec![make_param("cond", Type::Bool)],
871            vec![HirStmt::If {
872                condition: HirExpr::Var("cond".to_string()),
873                then_body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
874                    Literal::Int(1),
875                ))))],
876                else_body: None,
877            }],
878        );
879        let info = GeneratorStateInfo::analyze(&func);
880
881        assert_eq!(info.yield_count, 1);
882        assert!(info.captured_params.contains(&"cond".to_string()));
883    }
884
885    #[test]
886    fn test_if_statement_with_else() {
887        let func = make_func(
888            "if_else",
889            vec![],
890            vec![HirStmt::If {
891                condition: HirExpr::Literal(Literal::Bool(true)),
892                then_body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
893                    Literal::Int(1),
894                ))))],
895                else_body: Some(vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
896                    Literal::Int(2),
897                ))))]),
898            }],
899        );
900        let info = GeneratorStateInfo::analyze(&func);
901
902        assert_eq!(info.yield_count, 2);
903    }
904
905    // ============================================
906    // Try/Except statement tests
907    // ============================================
908
909    #[test]
910    fn test_try_except_body() {
911        let func = make_func(
912            "try_except",
913            vec![],
914            vec![HirStmt::Try {
915                body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
916                    Literal::Int(1),
917                ))))],
918                handlers: vec![ExceptHandler {
919                    exception_type: None,
920                    name: None,
921                    body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
922                        Literal::Int(2),
923                    ))))],
924                }],
925                orelse: None,
926                finalbody: None,
927            }],
928        );
929        let info = GeneratorStateInfo::analyze(&func);
930
931        assert_eq!(info.yield_count, 2);
932    }
933
934    #[test]
935    fn test_try_with_finally() {
936        let func = make_func(
937            "try_finally",
938            vec![],
939            vec![HirStmt::Try {
940                body: vec![HirStmt::Expr(make_yield(None))],
941                handlers: vec![],
942                orelse: None,
943                finalbody: Some(vec![make_assign(
944                    "cleanup",
945                    HirExpr::Literal(Literal::Bool(true)),
946                    None,
947                )]),
948            }],
949        );
950        let info = GeneratorStateInfo::analyze(&func);
951
952        assert_eq!(info.yield_count, 1);
953        assert!(info.state_variables.iter().any(|v| v.name == "cleanup"));
954    }
955
956    #[test]
957    fn test_try_with_orelse() {
958        let func = make_func(
959            "try_else",
960            vec![],
961            vec![HirStmt::Try {
962                body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
963                    Literal::Int(1),
964                ))))],
965                handlers: vec![],
966                orelse: Some(vec![HirStmt::Expr(make_yield(Some(HirExpr::Literal(
967                    Literal::Int(2),
968                ))))]),
969                finalbody: None,
970            }],
971        );
972        let info = GeneratorStateInfo::analyze(&func);
973
974        assert_eq!(info.yield_count, 2);
975    }
976
977    // ============================================
978    // With statement tests (DEPYLER-0561)
979    // ============================================
980
981    #[test]
982    fn test_with_statement() {
983        let func = make_func(
984            "with_gen",
985            vec![make_param(
986                "ctx",
987                Type::Custom("ContextManager".to_string()),
988            )],
989            vec![HirStmt::With {
990                context: HirExpr::Var("ctx".to_string()),
991                target: Some("f".to_string()),
992                body: vec![HirStmt::Expr(make_yield(Some(HirExpr::Var(
993                    "f".to_string(),
994                ))))],
995                is_async: false,
996            }],
997        );
998        let info = GeneratorStateInfo::analyze(&func);
999
1000        assert_eq!(info.yield_count, 1);
1001        // 'ctx' is a param used in the context expression, so it's captured
1002        assert!(info.captured_params.contains(&"ctx".to_string()));
1003    }
1004
1005    // ============================================
1006    // Expression analysis tests
1007    // ============================================
1008
1009    #[test]
1010    fn test_binary_expression_captures_variables() {
1011        let func = make_func(
1012            "binary",
1013            vec![make_param("a", Type::Int), make_param("b", Type::Int)],
1014            vec![HirStmt::Expr(make_yield(Some(HirExpr::Binary {
1015                left: Box::new(HirExpr::Var("a".to_string())),
1016                op: BinOp::Add,
1017                right: Box::new(HirExpr::Var("b".to_string())),
1018            })))],
1019        );
1020        let info = GeneratorStateInfo::analyze(&func);
1021
1022        assert!(info.captured_params.contains(&"a".to_string()));
1023        assert!(info.captured_params.contains(&"b".to_string()));
1024    }
1025
1026    #[test]
1027    fn test_unary_expression() {
1028        let func = make_func(
1029            "unary",
1030            vec![make_param("x", Type::Int)],
1031            vec![HirStmt::Expr(make_yield(Some(HirExpr::Unary {
1032                op: UnaryOp::Neg,
1033                operand: Box::new(HirExpr::Var("x".to_string())),
1034            })))],
1035        );
1036        let info = GeneratorStateInfo::analyze(&func);
1037
1038        assert!(info.captured_params.contains(&"x".to_string()));
1039    }
1040
1041    #[test]
1042    fn test_call_expression() {
1043        let func = make_func(
1044            "call",
1045            vec![make_param("arg", Type::Int)],
1046            vec![HirStmt::Expr(make_yield(Some(HirExpr::Call {
1047                func: "some_func".to_string(),
1048                args: vec![HirExpr::Var("arg".to_string())],
1049                kwargs: vec![],
1050            })))],
1051        );
1052        let info = GeneratorStateInfo::analyze(&func);
1053
1054        assert!(info.captured_params.contains(&"arg".to_string()));
1055    }
1056
1057    #[test]
1058    fn test_method_call_expression() {
1059        let func = make_func(
1060            "method_call",
1061            vec![make_param("obj", Type::Custom("MyClass".to_string()))],
1062            vec![HirStmt::Expr(make_yield(Some(HirExpr::MethodCall {
1063                object: Box::new(HirExpr::Var("obj".to_string())),
1064                method: "method".to_string(),
1065                args: vec![HirExpr::Literal(Literal::Int(42))],
1066                kwargs: vec![],
1067            })))],
1068        );
1069        let info = GeneratorStateInfo::analyze(&func);
1070
1071        assert!(info.captured_params.contains(&"obj".to_string()));
1072    }
1073
1074    #[test]
1075    fn test_index_expression() {
1076        let func = make_func(
1077            "index",
1078            vec![make_param("arr", Type::List(Box::new(Type::Int)))],
1079            vec![HirStmt::Expr(make_yield(Some(HirExpr::Index {
1080                base: Box::new(HirExpr::Var("arr".to_string())),
1081                index: Box::new(HirExpr::Literal(Literal::Int(0))),
1082            })))],
1083        );
1084        let info = GeneratorStateInfo::analyze(&func);
1085
1086        assert!(info.captured_params.contains(&"arr".to_string()));
1087    }
1088
1089    #[test]
1090    fn test_list_expression_analyzes_elements() {
1091        let func = make_func(
1092            "list_expr",
1093            vec![make_param("x", Type::Int)],
1094            vec![HirStmt::Expr(make_yield(Some(HirExpr::List(vec![
1095                HirExpr::Var("x".to_string()),
1096                HirExpr::Literal(Literal::Int(42)),
1097            ]))))],
1098        );
1099        let info = GeneratorStateInfo::analyze(&func);
1100
1101        assert!(info.captured_params.contains(&"x".to_string()));
1102    }
1103
1104    #[test]
1105    fn test_tuple_expression_analyzes_elements() {
1106        let func = make_func(
1107            "tuple_expr",
1108            vec![make_param("a", Type::Int), make_param("b", Type::Int)],
1109            vec![HirStmt::Expr(make_yield(Some(HirExpr::Tuple(vec![
1110                HirExpr::Var("a".to_string()),
1111                HirExpr::Var("b".to_string()),
1112            ]))))],
1113        );
1114        let info = GeneratorStateInfo::analyze(&func);
1115
1116        assert!(info.captured_params.contains(&"a".to_string()));
1117        assert!(info.captured_params.contains(&"b".to_string()));
1118    }
1119
1120    // ============================================
1121    // Return statement tests
1122    // ============================================
1123
1124    #[test]
1125    fn test_return_with_expression() {
1126        let func = make_func(
1127            "return_expr",
1128            vec![make_param("x", Type::Int)],
1129            vec![
1130                HirStmt::Expr(make_yield(Some(HirExpr::Var("x".to_string())))),
1131                HirStmt::Return(Some(HirExpr::Var("x".to_string()))),
1132            ],
1133        );
1134        let info = GeneratorStateInfo::analyze(&func);
1135
1136        assert!(info.captured_params.contains(&"x".to_string()));
1137    }
1138
1139    #[test]
1140    fn test_return_none() {
1141        let func = make_func("return_none", vec![], vec![HirStmt::Return(None)]);
1142        let info = GeneratorStateInfo::analyze(&func);
1143
1144        // Return(None) should not crash
1145        assert_eq!(info.yield_count, 0);
1146    }
1147
1148    // ============================================
1149    // Other statement types tests
1150    // ============================================
1151
1152    #[test]
1153    fn test_pass_statement() {
1154        let func = make_func("pass_func", vec![], vec![HirStmt::Pass]);
1155        let info = GeneratorStateInfo::analyze(&func);
1156
1157        assert_eq!(info.state_variables.len(), 0);
1158        assert_eq!(info.yield_count, 0);
1159    }
1160
1161    #[test]
1162    fn test_break_statement() {
1163        let func = make_func(
1164            "break_func",
1165            vec![],
1166            vec![HirStmt::While {
1167                condition: HirExpr::Literal(Literal::Bool(true)),
1168                body: vec![HirStmt::Break { label: None }],
1169            }],
1170        );
1171        let info = GeneratorStateInfo::analyze(&func);
1172
1173        assert!(info.has_loops);
1174    }
1175
1176    #[test]
1177    fn test_continue_statement() {
1178        let func = make_func(
1179            "continue_func",
1180            vec![],
1181            vec![HirStmt::While {
1182                condition: HirExpr::Literal(Literal::Bool(true)),
1183                body: vec![HirStmt::Continue { label: None }],
1184            }],
1185        );
1186        let info = GeneratorStateInfo::analyze(&func);
1187
1188        assert!(info.has_loops);
1189    }
1190
1191    // ============================================
1192    // Complex scenario tests
1193    // ============================================
1194
1195    #[test]
1196    fn test_fibonacci_generator() {
1197        // def fib(n): a, b = 0, 1; while a < n: yield a; a, b = b, a + b
1198        let func = make_func(
1199            "fib",
1200            vec![make_param("n", Type::Int)],
1201            vec![
1202                HirStmt::Assign {
1203                    target: depyler_hir::hir::AssignTarget::Tuple(vec![
1204                        depyler_hir::hir::AssignTarget::Symbol("a".to_string()),
1205                        depyler_hir::hir::AssignTarget::Symbol("b".to_string()),
1206                    ]),
1207                    value: HirExpr::Tuple(vec![
1208                        HirExpr::Literal(Literal::Int(0)),
1209                        HirExpr::Literal(Literal::Int(1)),
1210                    ]),
1211                    type_annotation: None,
1212                },
1213                HirStmt::While {
1214                    condition: HirExpr::Binary {
1215                        left: Box::new(HirExpr::Var("a".to_string())),
1216                        op: BinOp::Lt,
1217                        right: Box::new(HirExpr::Var("n".to_string())),
1218                    },
1219                    body: vec![
1220                        HirStmt::Expr(make_yield(Some(HirExpr::Var("a".to_string())))),
1221                        HirStmt::Assign {
1222                            target: depyler_hir::hir::AssignTarget::Tuple(vec![
1223                                depyler_hir::hir::AssignTarget::Symbol("a".to_string()),
1224                                depyler_hir::hir::AssignTarget::Symbol("b".to_string()),
1225                            ]),
1226                            value: HirExpr::Tuple(vec![
1227                                HirExpr::Var("b".to_string()),
1228                                HirExpr::Binary {
1229                                    left: Box::new(HirExpr::Var("a".to_string())),
1230                                    op: BinOp::Add,
1231                                    right: Box::new(HirExpr::Var("b".to_string())),
1232                                },
1233                            ]),
1234                            type_annotation: None,
1235                        },
1236                    ],
1237                },
1238            ],
1239        );
1240        let info = GeneratorStateInfo::analyze(&func);
1241
1242        assert!(info.has_loops);
1243        assert_eq!(info.yield_count, 1);
1244        assert!(info.captured_params.contains(&"n".to_string()));
1245        assert_eq!(info.state_variables.len(), 2); // a and b
1246    }
1247
1248    #[test]
1249    fn test_range_like_generator() {
1250        // def range_gen(start, stop, step): i = start; while i < stop: yield i; i += step
1251        let func = make_func(
1252            "range_gen",
1253            vec![
1254                make_param("start", Type::Int),
1255                make_param("stop", Type::Int),
1256                make_param("step", Type::Int),
1257            ],
1258            vec![
1259                make_assign("i", HirExpr::Var("start".to_string()), Some(Type::Int)),
1260                HirStmt::While {
1261                    condition: HirExpr::Binary {
1262                        left: Box::new(HirExpr::Var("i".to_string())),
1263                        op: BinOp::Lt,
1264                        right: Box::new(HirExpr::Var("stop".to_string())),
1265                    },
1266                    body: vec![
1267                        HirStmt::Expr(make_yield(Some(HirExpr::Var("i".to_string())))),
1268                        make_assign(
1269                            "i",
1270                            HirExpr::Binary {
1271                                left: Box::new(HirExpr::Var("i".to_string())),
1272                                op: BinOp::Add,
1273                                right: Box::new(HirExpr::Var("step".to_string())),
1274                            },
1275                            None,
1276                        ),
1277                    ],
1278                },
1279            ],
1280        );
1281        let info = GeneratorStateInfo::analyze(&func);
1282
1283        assert!(info.has_loops);
1284        assert_eq!(info.yield_count, 1);
1285        // start is used to initialize i, which is a local var
1286        // stop and step are used after yield points
1287        assert!(info.captured_params.contains(&"start".to_string()));
1288        assert!(info.captured_params.contains(&"stop".to_string()));
1289        assert!(info.captured_params.contains(&"step".to_string()));
1290        assert_eq!(info.state_variables.len(), 1); // i
1291    }
1292
1293    #[test]
1294    fn test_simple_counter_analysis() {
1295        // def counter(n): current = 0; while current < n: yield current; current += 1
1296        let func = HirFunction {
1297            name: "counter".to_string(),
1298            params: smallvec![HirParam::new("n".to_string(), Type::Int)],
1299            ret_type: Type::Int,
1300            body: vec![
1301                HirStmt::Assign {
1302                    target: depyler_hir::hir::AssignTarget::Symbol("current".to_string()),
1303                    value: HirExpr::Literal(depyler_hir::hir::Literal::Int(0)),
1304                    type_annotation: Some(Type::Int),
1305                },
1306                HirStmt::While {
1307                    condition: HirExpr::Binary {
1308                        left: Box::new(HirExpr::Var("current".to_string())),
1309                        op: depyler_hir::hir::BinOp::Lt,
1310                        right: Box::new(HirExpr::Var("n".to_string())),
1311                    },
1312                    body: vec![
1313                        HirStmt::Expr(HirExpr::Yield {
1314                            value: Some(Box::new(HirExpr::Var("current".to_string()))),
1315                        }),
1316                        HirStmt::Assign {
1317                            target: depyler_hir::hir::AssignTarget::Symbol("current".to_string()),
1318                            value: HirExpr::Binary {
1319                                left: Box::new(HirExpr::Var("current".to_string())),
1320                                op: depyler_hir::hir::BinOp::Add,
1321                                right: Box::new(HirExpr::Literal(depyler_hir::hir::Literal::Int(1))),
1322                            },
1323                            type_annotation: Some(Type::Int),
1324                        },
1325                    ],
1326                },
1327            ],
1328            properties: FunctionProperties::default(),
1329            annotations: TranspilationAnnotations::default(),
1330            docstring: None,
1331        };
1332
1333        let state_info = GeneratorStateInfo::analyze(&func);
1334
1335        assert_eq!(state_info.yield_count, 1, "Should find 1 yield");
1336        assert!(state_info.has_loops, "Should detect loop");
1337        assert_eq!(
1338            state_info.state_variables.len(),
1339            1,
1340            "Should find 'current' variable"
1341        );
1342        assert_eq!(state_info.state_variables[0].name, "current");
1343        assert!(
1344            state_info.captured_params.contains(&"n".to_string()),
1345            "Should capture parameter 'n'"
1346        );
1347    }
1348
1349    #[test]
1350    #[allow(non_snake_case)]
1351    fn test_depyler_0258_type_inference_from_literal_values() {
1352        // BUG #1: DynamicType inference should infer from value expressions
1353        // Current: i = 0 (no type annotation) → Type::Unknown
1354        // Expected: i = 0 (no type annotation) → Type::Int
1355
1356        let func = HirFunction {
1357            name: "count_up".to_string(),
1358            params: smallvec![HirParam::new("n".to_string(), Type::Int)],
1359            ret_type: Type::Int,
1360            body: vec![
1361                // Assignment WITHOUT type annotation - should infer from literal
1362                HirStmt::Assign {
1363                    target: depyler_hir::hir::AssignTarget::Symbol("i".to_string()),
1364                    value: HirExpr::Literal(depyler_hir::hir::Literal::Int(0)),
1365                    type_annotation: None, // ← No annotation, must infer!
1366                },
1367                HirStmt::While {
1368                    condition: HirExpr::Binary {
1369                        left: Box::new(HirExpr::Var("i".to_string())),
1370                        op: depyler_hir::hir::BinOp::Lt,
1371                        right: Box::new(HirExpr::Var("n".to_string())),
1372                    },
1373                    body: vec![
1374                        HirStmt::Expr(HirExpr::Yield {
1375                            value: Some(Box::new(HirExpr::Var("i".to_string()))),
1376                        }),
1377                        HirStmt::Assign {
1378                            target: depyler_hir::hir::AssignTarget::Symbol("i".to_string()),
1379                            value: HirExpr::Binary {
1380                                left: Box::new(HirExpr::Var("i".to_string())),
1381                                op: depyler_hir::hir::BinOp::Add,
1382                                right: Box::new(HirExpr::Literal(depyler_hir::hir::Literal::Int(1))),
1383                            },
1384                            type_annotation: None, // ← Reassignment, type already known
1385                        },
1386                    ],
1387                },
1388            ],
1389            properties: FunctionProperties::default(),
1390            annotations: TranspilationAnnotations::default(),
1391            docstring: None,
1392        };
1393
1394        let state_info = GeneratorStateInfo::analyze(&func);
1395
1396        // Assert: Should find state variable 'i'
1397        assert_eq!(
1398            state_info.state_variables.len(),
1399            1,
1400            "Should find 'i' variable"
1401        );
1402        assert_eq!(state_info.state_variables[0].name, "i");
1403
1404        // Assert: Type should be inferred as Int from literal value
1405        // This WILL FAIL (RED phase) because current code uses Type::Unknown
1406        assert_eq!(
1407            state_info.state_variables[0].ty,
1408            Type::Int,
1409            "DEPYLER-0258: Should infer Type::Int from literal value, not Type::Unknown"
1410        );
1411    }
1412}
depyler_analysis/generator_state.rs

depyler_analysis/
generator_state.rs
