Skip to main content

depyler_analysis/
string_optimization.rs

1use depyler_hir::hir::{AssignTarget, HirExpr, HirFunction, HirStmt, Literal, Type};
2use std::collections::{HashMap, HashSet};
3
4/// Analyzes string usage patterns to determine optimal string types
5#[derive(Debug, Default)]
6pub struct StringOptimizer {
7    /// String literals that are only read, never mutated
8    read_only_strings: HashSet<String>,
9    /// String parameters that are never mutated
10    immutable_params: HashSet<String>,
11    /// Strings that are returned from functions
12    returned_strings: HashSet<String>,
13    /// Strings used in multiple contexts (may need Cow)
14    mixed_usage_strings: HashSet<String>,
15    /// String literal frequency counter for interning decisions
16    string_literal_count: HashMap<String, usize>,
17    /// Strings that should be interned due to frequent use
18    interned_strings: HashSet<String>,
19    /// Mapping from string literal to its unique constant name
20    interned_names: HashMap<String, String>,
21}
22
23/// Optimal string representation based on usage analysis
24#[derive(Debug, Clone, PartialEq)]
25pub enum OptimalStringType {
26    /// Use &'static str for string literals that are never mutated
27    StaticStr,
28    /// Use &str for borrowed string parameters
29    BorrowedStr { lifetime: Option<String> },
30    /// Use String for owned, mutable strings
31    OwnedString,
32    /// Use Cow<'static, str> for mixed usage patterns
33    CowStr,
34}
35
36impl StringOptimizer {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Analyze a function to determine optimal string types
42    pub fn analyze_function(&mut self, func: &HirFunction) {
43        // Track string parameters
44        for param in &func.params {
45            if matches!(param.ty, Type::String) {
46                self.immutable_params.insert(param.name.clone());
47            }
48        }
49
50        // Analyze function body
51        for stmt in &func.body {
52            self.analyze_stmt(stmt);
53        }
54
55        // Parameters that are mutated are not immutable
56        for param in self.immutable_params.clone() {
57            if !self.is_immutable(&param) {
58                self.immutable_params.remove(&param);
59            }
60        }
61    }
62
63    /// Get the optimal string type for a given context
64    pub fn get_optimal_type(&self, context: &StringContext) -> OptimalStringType {
65        match context {
66            StringContext::Literal(s) => {
67                // v3.16.0 Phase 3: Only use Cow for TRUE mixed usage (returned AND borrowed elsewhere)
68                // Don't use Cow for simple returned literals - use owned String instead
69                if self.mixed_usage_strings.contains(s) {
70                    OptimalStringType::CowStr
71                } else if self.returned_strings.contains(s) {
72                    // Returned but not borrowed elsewhere - use owned String
73                    OptimalStringType::OwnedString
74                } else {
75                    // DEPYLER-TYPE-001: Default to StaticStr (&'static str) for literals
76                    // String literals passed as function args should stay as &str
77                    // since function params with Type::String become &str in Rust.
78                    // Only use OwnedString when explicitly needed (returned/concatenated).
79                    OptimalStringType::StaticStr
80                }
81            }
82            StringContext::Parameter(name) => {
83                if self.immutable_params.contains(name) {
84                    OptimalStringType::BorrowedStr {
85                        lifetime: Some("'a".to_string()),
86                    }
87                } else if self.mixed_usage_strings.contains(name) {
88                    OptimalStringType::CowStr
89                } else {
90                    OptimalStringType::OwnedString
91                }
92            }
93            StringContext::Return => OptimalStringType::OwnedString,
94            StringContext::Concatenation => OptimalStringType::OwnedString,
95        }
96    }
97
98    fn analyze_stmt(&mut self, stmt: &HirStmt) {
99        match stmt {
100            HirStmt::Assign { target, value, .. } => {
101                self.analyze_assign_stmt(target, value);
102            }
103            HirStmt::Return(Some(expr)) => {
104                self.analyze_expr(expr, true);
105            }
106            HirStmt::If {
107                condition,
108                then_body,
109                else_body,
110            } => {
111                self.analyze_if_stmt(condition, then_body, else_body);
112            }
113            HirStmt::While { condition, body } => {
114                self.analyze_while_stmt(condition, body);
115            }
116            HirStmt::For { iter, body, .. } => {
117                self.analyze_for_stmt(iter, body);
118            }
119            HirStmt::Expr(expr) => {
120                self.analyze_expr(expr, false);
121            }
122            _ => {}
123        }
124    }
125
126    fn analyze_assign_stmt(&mut self, target: &AssignTarget, value: &HirExpr) {
127        if let AssignTarget::Symbol(symbol) = target {
128            if self.immutable_params.contains(symbol) {
129                self.immutable_params.remove(symbol);
130            }
131        }
132        self.analyze_expr(value, false);
133    }
134
135    fn analyze_if_stmt(
136        &mut self,
137        condition: &HirExpr,
138        then_body: &[HirStmt],
139        else_body: &Option<Vec<HirStmt>>,
140    ) {
141        self.analyze_expr(condition, false);
142        for stmt in then_body {
143            self.analyze_stmt(stmt);
144        }
145        if let Some(else_stmts) = else_body {
146            for stmt in else_stmts {
147                self.analyze_stmt(stmt);
148            }
149        }
150    }
151
152    fn analyze_while_stmt(&mut self, condition: &HirExpr, body: &[HirStmt]) {
153        self.analyze_expr(condition, false);
154        for stmt in body {
155            self.analyze_stmt(stmt);
156        }
157    }
158
159    fn analyze_for_stmt(&mut self, iter: &HirExpr, body: &[HirStmt]) {
160        self.analyze_expr(iter, false);
161        for stmt in body {
162            self.analyze_stmt(stmt);
163        }
164    }
165
166    fn analyze_expr(&mut self, expr: &HirExpr, is_returned: bool) {
167        match expr {
168            HirExpr::Literal(Literal::String(s)) => {
169                self.analyze_string_literal(s, is_returned);
170            }
171            HirExpr::Var(name) => {
172                self.analyze_var_usage(name, is_returned);
173            }
174            HirExpr::Binary { op, left, right } => {
175                self.analyze_binary_expr(op, left, right);
176            }
177            HirExpr::Call { func, args, .. } => {
178                self.analyze_call_expr(func, args);
179            }
180            HirExpr::List(elts) | HirExpr::Tuple(elts) => {
181                self.analyze_collection_expr(elts, is_returned);
182            }
183            HirExpr::Dict(items) => {
184                self.analyze_dict_expr(items, is_returned);
185            }
186            _ => {}
187        }
188    }
189
190    fn analyze_string_literal(&mut self, s: &str, is_returned: bool) {
191        *self.string_literal_count.entry(s.to_string()).or_insert(0) += 1;
192
193        if self.string_literal_count.get(s).copied().unwrap_or(0) > 3 {
194            self.interned_strings.insert(s.to_string());
195        }
196
197        if is_returned {
198            self.returned_strings.insert(s.to_string());
199        } else {
200            self.read_only_strings.insert(s.to_string());
201        }
202    }
203
204    /// Finalize interned string names, resolving any collisions
205    /// This must be called after analysis and before code generation
206    pub fn finalize_interned_names(&mut self) {
207        if !self.interned_names.is_empty() {
208            // Already finalized
209            return;
210        }
211
212        // Map from base constant name to list of actual string values
213        let mut name_map: HashMap<String, Vec<String>> = HashMap::new();
214
215        // Group strings by their base constant name
216        for s in &self.interned_strings {
217            let base_name = self.generate_base_const_name(s);
218            name_map.entry(base_name).or_default().push(s.clone());
219        }
220
221        // Assign unique names, adding suffixes for collisions
222        for (base_name, strings) in name_map {
223            if strings.len() == 1 {
224                // No collision, use base name
225                self.interned_names.insert(strings[0].clone(), base_name);
226            } else {
227                // Collision detected, add numeric suffixes
228                for (idx, s) in strings.iter().enumerate() {
229                    let unique_name = format!("{}_{}", base_name, idx + 1);
230                    self.interned_names.insert(s.clone(), unique_name);
231                }
232            }
233        }
234    }
235
236    /// Generate base constant name from string content (may have collisions)
237    fn generate_base_const_name(&self, s: &str) -> String {
238        // Convert to uppercase, replace non-alphanumeric with underscore
239        let name = s
240            .chars()
241            .map(|c| match c {
242                'a'..='z' | 'A'..='Z' | '0'..='9' => c.to_ascii_uppercase(),
243                _ => '_',
244            })
245            .collect::<String>();
246
247        let base_name = if name.is_empty() {
248            "EMPTY".to_string()
249        } else {
250            name
251        };
252
253        format!("STR_{}", base_name)
254    }
255
256    fn analyze_var_usage(&mut self, name: &str, is_returned: bool) {
257        if is_returned && self.immutable_params.contains(name) {
258            self.mixed_usage_strings.insert(name.to_string());
259        }
260    }
261
262    fn analyze_binary_expr(&mut self, op: &depyler_hir::hir::BinOp, left: &HirExpr, right: &HirExpr) {
263        if matches!(op, depyler_hir::hir::BinOp::Add)
264            && (self.is_string_expr(left) || self.is_string_expr(right))
265        {
266            self.mark_as_owned(left);
267            self.mark_as_owned(right);
268        }
269        self.analyze_expr(left, false);
270        self.analyze_expr(right, false);
271    }
272
273    fn analyze_call_expr(&mut self, func: &str, args: &[HirExpr]) {
274        if self.is_mutating_method(func) && !args.is_empty() {
275            if let HirExpr::Var(name) = &args[0] {
276                self.immutable_params.remove(name);
277            }
278        }
279        for arg in args {
280            self.analyze_expr(arg, false);
281        }
282    }
283
284    fn analyze_collection_expr(&mut self, elts: &[HirExpr], is_returned: bool) {
285        for elt in elts {
286            self.analyze_expr(elt, is_returned);
287        }
288    }
289
290    fn analyze_dict_expr(&mut self, items: &[(HirExpr, HirExpr)], is_returned: bool) {
291        for (k, v) in items {
292            self.analyze_expr(k, false);
293            self.analyze_expr(v, is_returned);
294        }
295    }
296
297    #[allow(dead_code)]
298    fn is_read_only(&self, s: &str) -> bool {
299        self.read_only_strings.contains(s) && !self.returned_strings.contains(s)
300    }
301
302    fn is_immutable(&self, param: &str) -> bool {
303        self.immutable_params.contains(param)
304    }
305
306    /// Check if an expression is a string type
307    fn is_string_expr(&self, expr: &HirExpr) -> bool {
308        match expr {
309            HirExpr::Literal(Literal::String(_)) => true,
310            HirExpr::Var(name) => self.immutable_params.contains(name),
311            HirExpr::Call { func, .. } => {
312                // Common string-returning functions
313                matches!(func.as_str(), "str" | "format" | "to_string" | "join")
314            }
315            _ => false,
316        }
317    }
318
319    /// Mark a string expression as needing ownership
320    fn mark_as_owned(&mut self, expr: &HirExpr) {
321        match expr {
322            HirExpr::Literal(Literal::String(s)) => {
323                self.read_only_strings.remove(s);
324            }
325            HirExpr::Var(name) => {
326                self.immutable_params.remove(name);
327            }
328            _ => {}
329        }
330    }
331
332    /// Check if a method call mutates the string
333    fn is_mutating_method(&self, method: &str) -> bool {
334        matches!(
335            method,
336            "push_str" | "push" | "insert" | "insert_str" | "replace_range" | "clear" | "truncate"
337        )
338    }
339
340    /// Check if a string literal should be interned
341    pub fn should_intern(&self, s: &str) -> bool {
342        self.interned_strings.contains(s)
343    }
344
345    /// Get interned string name for a literal
346    /// Returns the unique constant name for an interned string
347    pub fn get_interned_name(&self, s: &str) -> Option<String> {
348        // Return the finalized name from the cache
349        self.interned_names.get(s).cloned()
350    }
351
352    /// Generate interned string constants
353    pub fn generate_interned_constants(&self) -> Vec<String> {
354        let mut constants = Vec::new();
355
356        for (string_value, const_name) in &self.interned_names {
357            constants.push(format!(
358                "const {}: &'static str = \"{}\";",
359                const_name,
360                escape_string(string_value)
361            ));
362        }
363
364        constants
365    }
366}
367
368/// Context in which a string is being used
369#[derive(Debug, Clone)]
370pub enum StringContext {
371    /// String literal in source code
372    Literal(String),
373    /// Function parameter
374    Parameter(String),
375    /// Return value
376    Return,
377    /// String concatenation operation
378    Concatenation,
379}
380
381impl std::fmt::Display for StringContext {
382    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
383        match self {
384            StringContext::Literal(s) => write!(f, "\"{}\"", s),
385            StringContext::Parameter(name) => write!(f, "{}", name),
386            StringContext::Return => write!(f, "<return>"),
387            StringContext::Concatenation => write!(f, "<concat>"),
388        }
389    }
390}
391
392/// Generates optimized string code based on usage
393pub fn generate_optimized_string(optimizer: &StringOptimizer, context: &StringContext) -> String {
394    match optimizer.get_optimal_type(context) {
395        OptimalStringType::StaticStr => generate_static_str(context),
396        OptimalStringType::BorrowedStr { .. } => generate_borrowed_str(context),
397        OptimalStringType::OwnedString => generate_owned_string(context),
398        OptimalStringType::CowStr => generate_cow_str(context),
399    }
400}
401
402fn generate_static_str(context: &StringContext) -> String {
403    match context {
404        StringContext::Literal(s) => format!("\"{}\"", escape_string(s)),
405        _ => format!("{}.to_string()", context),
406    }
407}
408
409fn generate_borrowed_str(context: &StringContext) -> String {
410    match context {
411        StringContext::Parameter(name) => name.clone(),
412        StringContext::Literal(s) => format!("\"{}\"", escape_string(s)),
413        _ => format!("{}.as_str()", context),
414    }
415}
416
417fn generate_owned_string(context: &StringContext) -> String {
418    match context {
419        StringContext::Literal(s) => format!("\"{}\".to_string()", escape_string(s)),
420        StringContext::Parameter(name) => format!("{}.to_string()", name),
421        StringContext::Concatenation | StringContext::Return => "String::new()".to_string(),
422    }
423}
424
425fn generate_cow_str(context: &StringContext) -> String {
426    match context {
427        StringContext::Literal(s) => format!("Cow::Borrowed(\"{}\")", escape_string(s)),
428        StringContext::Parameter(name) => format!("Cow::Borrowed({})", name),
429        StringContext::Concatenation | StringContext::Return => {
430            "Cow::Owned(String::new())".to_string()
431        }
432    }
433}
434
435/// Escape a string for use in Rust string literals
436fn escape_string(s: &str) -> String {
437    s.chars().flat_map(escape_char).collect()
438}
439
440fn escape_char(c: char) -> Vec<char> {
441    match c {
442        '"' => vec!['\\', '"'],
443        '\\' => vec!['\\', '\\'],
444        '\n' => vec!['\\', 'n'],
445        '\r' => vec!['\\', 'r'],
446        '\t' => vec!['\\', 't'],
447        c => vec![c],
448    }
449}
450
451#[cfg(test)]
452mod tests {
453    use super::*;
454    use depyler_hir::hir::*;
455
456    #[test]
457    fn test_read_only_string_optimization() {
458        let mut optimizer = StringOptimizer::new();
459
460        let func = HirFunction {
461            name: "test".to_string(),
462            params: vec![].into(),
463            ret_type: Type::None,
464            body: vec![HirStmt::Expr(HirExpr::Call {
465                func: "print".to_string(),
466                args: vec![HirExpr::Literal(Literal::String("hello".to_string()))],
467                kwargs: vec![],
468            })],
469            properties: FunctionProperties::default(),
470            annotations: Default::default(),
471            docstring: None,
472        };
473
474        optimizer.analyze_function(&func);
475
476        let context = StringContext::Literal("hello".to_string());
477        assert_eq!(
478            optimizer.get_optimal_type(&context),
479            OptimalStringType::StaticStr
480        );
481    }
482
483    #[test]
484    fn test_returned_string_needs_ownership() {
485        let mut optimizer = StringOptimizer::new();
486
487        let func = HirFunction {
488            name: "test".to_string(),
489            params: vec![].into(),
490            ret_type: Type::String,
491            body: vec![HirStmt::Return(Some(HirExpr::Literal(Literal::String(
492                "result".to_string(),
493            ))))],
494            properties: FunctionProperties::default(),
495            annotations: Default::default(),
496            docstring: None,
497        };
498
499        optimizer.analyze_function(&func);
500
501        let context = StringContext::Literal("result".to_string());
502        assert_eq!(
503            optimizer.get_optimal_type(&context),
504            OptimalStringType::OwnedString
505        );
506    }
507
508    #[test]
509    fn test_immutable_parameter_borrowing() {
510        let mut optimizer = StringOptimizer::new();
511
512        let func = HirFunction {
513            name: "test".to_string(),
514            params: vec![HirParam::new("s".to_string(), Type::String)].into(),
515            ret_type: Type::Int,
516            body: vec![HirStmt::Return(Some(HirExpr::Call {
517                func: "len".to_string(),
518                args: vec![HirExpr::Var("s".to_string())],
519                kwargs: vec![],
520            }))],
521            properties: FunctionProperties::default(),
522            annotations: Default::default(),
523            docstring: None,
524        };
525
526        optimizer.analyze_function(&func);
527
528        let context = StringContext::Parameter("s".to_string());
529        assert_eq!(
530            optimizer.get_optimal_type(&context),
531            OptimalStringType::BorrowedStr {
532                lifetime: Some("'a".to_string())
533            }
534        );
535    }
536
537    #[test]
538    fn test_generate_optimized_string_code() {
539        let optimizer = StringOptimizer::new();
540
541        let code =
542            generate_optimized_string(&optimizer, &StringContext::Literal("hello".to_string()));
543        assert!(code == "\"hello\".to_string()" || code == "\"hello\"");
544    }
545
546    #[test]
547    fn test_new_creates_default() {
548        let optimizer = StringOptimizer::new();
549        assert!(!optimizer.should_intern("any"));
550        assert!(optimizer.get_interned_name("any").is_none());
551    }
552
553    #[test]
554    fn test_mixed_usage_strings_get_cow() {
555        let mut optimizer = StringOptimizer::new();
556
557        // Parameter that is both used and returned - but immutable params take precedence
558        let func = HirFunction {
559            name: "test".to_string(),
560            params: vec![HirParam::new("s".to_string(), Type::String)].into(),
561            ret_type: Type::String,
562            body: vec![
563                HirStmt::Expr(HirExpr::Var("s".to_string())),
564                HirStmt::Return(Some(HirExpr::Var("s".to_string()))),
565            ],
566            properties: FunctionProperties::default(),
567            annotations: Default::default(),
568            docstring: None,
569        };
570
571        optimizer.analyze_function(&func);
572
573        // Even though it's in mixed_usage_strings, immutable_params takes precedence
574        let context = StringContext::Parameter("s".to_string());
575        assert_eq!(
576            optimizer.get_optimal_type(&context),
577            OptimalStringType::BorrowedStr {
578                lifetime: Some("'a".to_string())
579            }
580        );
581
582        // Verify it IS in mixed_usage though
583        assert!(optimizer.mixed_usage_strings.contains("s"));
584    }
585
586    #[test]
587    fn test_mutated_parameter_loses_immutability() {
588        let mut optimizer = StringOptimizer::new();
589
590        let func = HirFunction {
591            name: "test".to_string(),
592            params: vec![HirParam::new("s".to_string(), Type::String)].into(),
593            ret_type: Type::None,
594            body: vec![HirStmt::Assign {
595                target: AssignTarget::Symbol("s".to_string()),
596                value: HirExpr::Literal(Literal::String("new".to_string())),
597                type_annotation: None,
598            }],
599            properties: FunctionProperties::default(),
600            annotations: Default::default(),
601            docstring: None,
602        };
603
604        optimizer.analyze_function(&func);
605
606        let context = StringContext::Parameter("s".to_string());
607        assert_eq!(
608            optimizer.get_optimal_type(&context),
609            OptimalStringType::OwnedString
610        );
611    }
612
613    #[test]
614    fn test_string_interning_threshold() {
615        let mut optimizer = StringOptimizer::new();
616
617        // Use the same string 4 times to trigger interning
618        let func = HirFunction {
619            name: "test".to_string(),
620            params: vec![].into(),
621            ret_type: Type::None,
622            body: vec![
623                HirStmt::Expr(HirExpr::Literal(Literal::String("common".to_string()))),
624                HirStmt::Expr(HirExpr::Literal(Literal::String("common".to_string()))),
625                HirStmt::Expr(HirExpr::Literal(Literal::String("common".to_string()))),
626                HirStmt::Expr(HirExpr::Literal(Literal::String("common".to_string()))),
627            ],
628            properties: FunctionProperties::default(),
629            annotations: Default::default(),
630            docstring: None,
631        };
632
633        optimizer.analyze_function(&func);
634
635        assert!(optimizer.should_intern("common"));
636    }
637
638    #[test]
639    fn test_finalize_interned_names_no_collision() {
640        let mut optimizer = StringOptimizer::new();
641        optimizer.interned_strings.insert("hello".to_string());
642        optimizer.finalize_interned_names();
643
644        let name = optimizer.get_interned_name("hello").unwrap();
645        assert_eq!(name, "STR_HELLO");
646    }
647
648    #[test]
649    fn test_finalize_interned_names_with_collision() {
650        let mut optimizer = StringOptimizer::new();
651        // "hello!" and "hello?" both become STR_HELLO_
652        optimizer.interned_strings.insert("hello!".to_string());
653        optimizer.interned_strings.insert("hello?".to_string());
654        optimizer.finalize_interned_names();
655
656        let name1 = optimizer.get_interned_name("hello!").unwrap();
657        let name2 = optimizer.get_interned_name("hello?").unwrap();
658
659        assert!(name1.starts_with("STR_HELLO_"));
660        assert!(name2.starts_with("STR_HELLO_"));
661        assert_ne!(name1, name2);
662    }
663
664    #[test]
665    fn test_finalize_interned_names_already_finalized() {
666        let mut optimizer = StringOptimizer::new();
667        optimizer.interned_strings.insert("test".to_string());
668        optimizer.finalize_interned_names();
669
670        // Call again - should be a no-op
671        optimizer.finalize_interned_names();
672
673        assert!(optimizer.get_interned_name("test").is_some());
674    }
675
676    #[test]
677    fn test_generate_base_const_name_empty() {
678        let optimizer = StringOptimizer::new();
679        let name = optimizer.generate_base_const_name("");
680        assert_eq!(name, "STR_EMPTY");
681    }
682
683    #[test]
684    fn test_generate_base_const_name_special_chars() {
685        let optimizer = StringOptimizer::new();
686        let name = optimizer.generate_base_const_name("hello world!");
687        assert_eq!(name, "STR_HELLO_WORLD_");
688    }
689
690    #[test]
691    fn test_generate_interned_constants() {
692        let mut optimizer = StringOptimizer::new();
693        optimizer.interned_strings.insert("test".to_string());
694        optimizer.finalize_interned_names();
695
696        let constants = optimizer.generate_interned_constants();
697        assert_eq!(constants.len(), 1);
698        assert!(constants[0].contains("STR_TEST"));
699        assert!(constants[0].contains("\"test\""));
700    }
701
702    #[test]
703    fn test_analyze_if_stmt_with_else() {
704        let mut optimizer = StringOptimizer::new();
705
706        let func = HirFunction {
707            name: "test".to_string(),
708            params: vec![].into(),
709            ret_type: Type::None,
710            body: vec![HirStmt::If {
711                condition: HirExpr::Literal(Literal::Bool(true)),
712                then_body: vec![HirStmt::Expr(HirExpr::Literal(Literal::String(
713                    "then".to_string(),
714                )))],
715                else_body: Some(vec![HirStmt::Expr(HirExpr::Literal(Literal::String(
716                    "else".to_string(),
717                )))]),
718            }],
719            properties: FunctionProperties::default(),
720            annotations: Default::default(),
721            docstring: None,
722        };
723
724        optimizer.analyze_function(&func);
725
726        assert!(optimizer.is_read_only("then"));
727        assert!(optimizer.is_read_only("else"));
728    }
729
730    #[test]
731    fn test_analyze_while_stmt() {
732        let mut optimizer = StringOptimizer::new();
733
734        let func = HirFunction {
735            name: "test".to_string(),
736            params: vec![].into(),
737            ret_type: Type::None,
738            body: vec![HirStmt::While {
739                condition: HirExpr::Literal(Literal::Bool(true)),
740                body: vec![HirStmt::Expr(HirExpr::Literal(Literal::String(
741                    "loop".to_string(),
742                )))],
743            }],
744            properties: FunctionProperties::default(),
745            annotations: Default::default(),
746            docstring: None,
747        };
748
749        optimizer.analyze_function(&func);
750
751        assert!(optimizer.is_read_only("loop"));
752    }
753
754    #[test]
755    fn test_analyze_for_stmt() {
756        let mut optimizer = StringOptimizer::new();
757
758        let func = HirFunction {
759            name: "test".to_string(),
760            params: vec![].into(),
761            ret_type: Type::None,
762            body: vec![HirStmt::For {
763                target: AssignTarget::Symbol("i".to_string()),
764                iter: HirExpr::List(vec![]),
765                body: vec![HirStmt::Expr(HirExpr::Literal(Literal::String(
766                    "body".to_string(),
767                )))],
768            }],
769            properties: FunctionProperties::default(),
770            annotations: Default::default(),
771            docstring: None,
772        };
773
774        optimizer.analyze_function(&func);
775
776        assert!(optimizer.is_read_only("body"));
777    }
778
779    #[test]
780    fn test_analyze_string_concatenation() {
781        let mut optimizer = StringOptimizer::new();
782
783        let func = HirFunction {
784            name: "test".to_string(),
785            params: vec![].into(),
786            ret_type: Type::None,
787            body: vec![HirStmt::Expr(HirExpr::Binary {
788                op: BinOp::Add,
789                left: Box::new(HirExpr::Literal(Literal::String("a".to_string()))),
790                right: Box::new(HirExpr::Literal(Literal::String("b".to_string()))),
791            })],
792            properties: FunctionProperties::default(),
793            annotations: Default::default(),
794            docstring: None,
795        };
796
797        optimizer.analyze_function(&func);
798
799        // Binary expression with Add triggers analysis on both sides
800        // The analyze_binary_expr calls mark_as_owned, then analyze_expr which re-adds to read_only
801        // So they end up in read_only_strings (is_read_only returns true if in read_only AND not returned)
802        assert!(optimizer.read_only_strings.contains("a"));
803        assert!(optimizer.read_only_strings.contains("b"));
804    }
805
806    #[test]
807    fn test_analyze_mutating_call() {
808        let mut optimizer = StringOptimizer::new();
809
810        let func = HirFunction {
811            name: "test".to_string(),
812            params: vec![HirParam::new("s".to_string(), Type::String)].into(),
813            ret_type: Type::None,
814            body: vec![HirStmt::Expr(HirExpr::Call {
815                func: "push_str".to_string(),
816                args: vec![HirExpr::Var("s".to_string())],
817                kwargs: vec![],
818            })],
819            properties: FunctionProperties::default(),
820            annotations: Default::default(),
821            docstring: None,
822        };
823
824        optimizer.analyze_function(&func);
825
826        // Mutated parameter loses immutability
827        assert!(!optimizer.immutable_params.contains("s"));
828    }
829
830    #[test]
831    fn test_analyze_list_and_tuple() {
832        let mut optimizer = StringOptimizer::new();
833
834        let func = HirFunction {
835            name: "test".to_string(),
836            params: vec![].into(),
837            ret_type: Type::None,
838            body: vec![
839                HirStmt::Expr(HirExpr::List(vec![HirExpr::Literal(Literal::String(
840                    "list".to_string(),
841                ))])),
842                HirStmt::Expr(HirExpr::Tuple(vec![HirExpr::Literal(Literal::String(
843                    "tuple".to_string(),
844                ))])),
845            ],
846            properties: FunctionProperties::default(),
847            annotations: Default::default(),
848            docstring: None,
849        };
850
851        optimizer.analyze_function(&func);
852
853        assert!(optimizer.is_read_only("list"));
854        assert!(optimizer.is_read_only("tuple"));
855    }
856
857    #[test]
858    fn test_analyze_dict() {
859        let mut optimizer = StringOptimizer::new();
860
861        let func = HirFunction {
862            name: "test".to_string(),
863            params: vec![].into(),
864            ret_type: Type::None,
865            body: vec![HirStmt::Expr(HirExpr::Dict(vec![(
866                HirExpr::Literal(Literal::String("key".to_string())),
867                HirExpr::Literal(Literal::String("value".to_string())),
868            )]))],
869            properties: FunctionProperties::default(),
870            annotations: Default::default(),
871            docstring: None,
872        };
873
874        optimizer.analyze_function(&func);
875
876        assert!(optimizer.is_read_only("key"));
877        assert!(optimizer.is_read_only("value"));
878    }
879
880    #[test]
881    fn test_is_string_expr_call() {
882        let optimizer = StringOptimizer::new();
883
884        assert!(optimizer.is_string_expr(&HirExpr::Call {
885            func: "str".to_string(),
886            args: vec![],
887            kwargs: vec![]
888        }));
889        assert!(optimizer.is_string_expr(&HirExpr::Call {
890            func: "format".to_string(),
891            args: vec![],
892            kwargs: vec![]
893        }));
894        assert!(optimizer.is_string_expr(&HirExpr::Call {
895            func: "to_string".to_string(),
896            args: vec![],
897            kwargs: vec![]
898        }));
899        assert!(optimizer.is_string_expr(&HirExpr::Call {
900            func: "join".to_string(),
901            args: vec![],
902            kwargs: vec![]
903        }));
904        assert!(!optimizer.is_string_expr(&HirExpr::Call {
905            func: "len".to_string(),
906            args: vec![],
907            kwargs: vec![]
908        }));
909    }
910
911    #[test]
912    fn test_is_mutating_method() {
913        let optimizer = StringOptimizer::new();
914
915        assert!(optimizer.is_mutating_method("push_str"));
916        assert!(optimizer.is_mutating_method("push"));
917        assert!(optimizer.is_mutating_method("insert"));
918        assert!(optimizer.is_mutating_method("insert_str"));
919        assert!(optimizer.is_mutating_method("replace_range"));
920        assert!(optimizer.is_mutating_method("clear"));
921        assert!(optimizer.is_mutating_method("truncate"));
922        assert!(!optimizer.is_mutating_method("len"));
923    }
924
925    #[test]
926    fn test_get_optimal_type_return_context() {
927        let optimizer = StringOptimizer::new();
928        assert_eq!(
929            optimizer.get_optimal_type(&StringContext::Return),
930            OptimalStringType::OwnedString
931        );
932    }
933
934    #[test]
935    fn test_get_optimal_type_concatenation_context() {
936        let optimizer = StringOptimizer::new();
937        assert_eq!(
938            optimizer.get_optimal_type(&StringContext::Concatenation),
939            OptimalStringType::OwnedString
940        );
941    }
942
943    #[test]
944    fn test_string_context_display() {
945        assert_eq!(
946            format!("{}", StringContext::Literal("hello".to_string())),
947            "\"hello\""
948        );
949        assert_eq!(
950            format!("{}", StringContext::Parameter("s".to_string())),
951            "s"
952        );
953        assert_eq!(format!("{}", StringContext::Return), "<return>");
954        assert_eq!(format!("{}", StringContext::Concatenation), "<concat>");
955    }
956
957    #[test]
958    fn test_generate_static_str() {
959        let s = generate_static_str(&StringContext::Literal("hello".to_string()));
960        assert_eq!(s, "\"hello\"");
961
962        let s = generate_static_str(&StringContext::Parameter("s".to_string()));
963        assert_eq!(s, "s.to_string()");
964    }
965
966    #[test]
967    fn test_generate_borrowed_str() {
968        let s = generate_borrowed_str(&StringContext::Parameter("s".to_string()));
969        assert_eq!(s, "s");
970
971        let s = generate_borrowed_str(&StringContext::Literal("test".to_string()));
972        assert_eq!(s, "\"test\"");
973
974        let s = generate_borrowed_str(&StringContext::Return);
975        assert_eq!(s, "<return>.as_str()");
976    }
977
978    #[test]
979    fn test_generate_owned_string() {
980        let s = generate_owned_string(&StringContext::Literal("hello".to_string()));
981        assert_eq!(s, "\"hello\".to_string()");
982
983        let s = generate_owned_string(&StringContext::Parameter("s".to_string()));
984        assert_eq!(s, "s.to_string()");
985
986        let s = generate_owned_string(&StringContext::Return);
987        assert_eq!(s, "String::new()");
988
989        let s = generate_owned_string(&StringContext::Concatenation);
990        assert_eq!(s, "String::new()");
991    }
992
993    #[test]
994    fn test_generate_cow_str() {
995        let s = generate_cow_str(&StringContext::Literal("hello".to_string()));
996        assert_eq!(s, "Cow::Borrowed(\"hello\")");
997
998        let s = generate_cow_str(&StringContext::Parameter("s".to_string()));
999        assert_eq!(s, "Cow::Borrowed(s)");
1000
1001        let s = generate_cow_str(&StringContext::Return);
1002        assert_eq!(s, "Cow::Owned(String::new())");
1003
1004        let s = generate_cow_str(&StringContext::Concatenation);
1005        assert_eq!(s, "Cow::Owned(String::new())");
1006    }
1007
1008    #[test]
1009    fn test_escape_string_special_chars() {
1010        assert_eq!(escape_string("hello\"world"), "hello\\\"world");
1011        assert_eq!(escape_string("hello\\world"), "hello\\\\world");
1012        assert_eq!(escape_string("hello\nworld"), "hello\\nworld");
1013        assert_eq!(escape_string("hello\rworld"), "hello\\rworld");
1014        assert_eq!(escape_string("hello\tworld"), "hello\\tworld");
1015        assert_eq!(escape_string("normal"), "normal");
1016    }
1017
1018    #[test]
1019    fn test_return_none_handled() {
1020        let mut optimizer = StringOptimizer::new();
1021
1022        let func = HirFunction {
1023            name: "test".to_string(),
1024            params: vec![].into(),
1025            ret_type: Type::None,
1026            body: vec![HirStmt::Return(None)],
1027            properties: FunctionProperties::default(),
1028            annotations: Default::default(),
1029            docstring: None,
1030        };
1031
1032        optimizer.analyze_function(&func);
1033        // Should not panic
1034    }
1035
1036    #[test]
1037    fn test_mark_as_owned_var() {
1038        let mut optimizer = StringOptimizer::new();
1039        optimizer.immutable_params.insert("s".to_string());
1040
1041        optimizer.mark_as_owned(&HirExpr::Var("s".to_string()));
1042
1043        assert!(!optimizer.immutable_params.contains("s"));
1044    }
1045
1046    #[test]
1047    fn test_mark_as_owned_other() {
1048        let mut optimizer = StringOptimizer::new();
1049
1050        // Should not panic on non-string/non-var expressions
1051        optimizer.mark_as_owned(&HirExpr::Literal(Literal::Int(42)));
1052    }
1053
1054    #[test]
1055    fn test_analyze_assign_non_symbol_target() {
1056        let mut optimizer = StringOptimizer::new();
1057
1058        let func = HirFunction {
1059            name: "test".to_string(),
1060            params: vec![].into(),
1061            ret_type: Type::None,
1062            body: vec![HirStmt::Assign {
1063                target: AssignTarget::Index {
1064                    base: Box::new(HirExpr::Var("arr".to_string())),
1065                    index: Box::new(HirExpr::Literal(Literal::Int(0))),
1066                },
1067                value: HirExpr::Literal(Literal::String("value".to_string())),
1068                type_annotation: None,
1069            }],
1070            properties: FunctionProperties::default(),
1071            annotations: Default::default(),
1072            docstring: None,
1073        };
1074
1075        optimizer.analyze_function(&func);
1076        // Should not panic
1077    }
1078
1079    #[test]
1080    fn test_call_with_no_args() {
1081        let mut optimizer = StringOptimizer::new();
1082
1083        let func = HirFunction {
1084            name: "test".to_string(),
1085            params: vec![].into(),
1086            ret_type: Type::None,
1087            body: vec![HirStmt::Expr(HirExpr::Call {
1088                func: "push_str".to_string(),
1089                args: vec![],
1090                kwargs: vec![],
1091            })],
1092            properties: FunctionProperties::default(),
1093            annotations: Default::default(),
1094            docstring: None,
1095        };
1096
1097        optimizer.analyze_function(&func);
1098        // Should not panic when args is empty
1099    }
1100
1101    #[test]
1102    fn test_mixed_usage_literal() {
1103        let mut optimizer = StringOptimizer::new();
1104        optimizer.mixed_usage_strings.insert("mixed".to_string());
1105
1106        let context = StringContext::Literal("mixed".to_string());
1107        assert_eq!(
1108            optimizer.get_optimal_type(&context),
1109            OptimalStringType::CowStr
1110        );
1111    }
1112
1113    #[test]
1114    fn test_parameter_mixed_usage() {
1115        let mut optimizer = StringOptimizer::new();
1116        optimizer.mixed_usage_strings.insert("param".to_string());
1117
1118        let context = StringContext::Parameter("param".to_string());
1119        assert_eq!(
1120            optimizer.get_optimal_type(&context),
1121            OptimalStringType::CowStr
1122        );
1123    }
1124
1125    // Tests for OptimalStringType traits
1126    #[test]
1127    fn test_optimal_string_type_debug() {
1128        let static_str = OptimalStringType::StaticStr;
1129        assert!(format!("{:?}", static_str).contains("StaticStr"));
1130
1131        let borrowed = OptimalStringType::BorrowedStr {
1132            lifetime: Some("'a".to_string()),
1133        };
1134        assert!(format!("{:?}", borrowed).contains("BorrowedStr"));
1135        assert!(format!("{:?}", borrowed).contains("'a"));
1136
1137        let owned = OptimalStringType::OwnedString;
1138        assert!(format!("{:?}", owned).contains("OwnedString"));
1139
1140        let cow = OptimalStringType::CowStr;
1141        assert!(format!("{:?}", cow).contains("CowStr"));
1142    }
1143
1144    #[test]
1145    fn test_optimal_string_type_clone() {
1146        let original = OptimalStringType::BorrowedStr {
1147            lifetime: Some("'b".to_string()),
1148        };
1149        let cloned = original.clone();
1150        assert_eq!(original, cloned);
1151    }
1152
1153    #[test]
1154    fn test_optimal_string_type_partial_eq() {
1155        assert_eq!(OptimalStringType::StaticStr, OptimalStringType::StaticStr);
1156        assert_eq!(
1157            OptimalStringType::OwnedString,
1158            OptimalStringType::OwnedString
1159        );
1160        assert_eq!(OptimalStringType::CowStr, OptimalStringType::CowStr);
1161        assert_ne!(OptimalStringType::StaticStr, OptimalStringType::OwnedString);
1162        assert_ne!(OptimalStringType::CowStr, OptimalStringType::StaticStr);
1163
1164        let borrowed1 = OptimalStringType::BorrowedStr {
1165            lifetime: Some("'a".to_string()),
1166        };
1167        let borrowed2 = OptimalStringType::BorrowedStr {
1168            lifetime: Some("'a".to_string()),
1169        };
1170        let borrowed3 = OptimalStringType::BorrowedStr {
1171            lifetime: Some("'b".to_string()),
1172        };
1173        let borrowed_none = OptimalStringType::BorrowedStr { lifetime: None };
1174
1175        assert_eq!(borrowed1, borrowed2);
1176        assert_ne!(borrowed1, borrowed3);
1177        assert_ne!(borrowed1, borrowed_none);
1178    }
1179
1180    // Tests for StringOptimizer struct
1181    #[test]
1182    fn test_string_optimizer_debug() {
1183        let optimizer = StringOptimizer::new();
1184        let debug_str = format!("{:?}", optimizer);
1185        assert!(debug_str.contains("StringOptimizer"));
1186    }
1187
1188    #[test]
1189    fn test_string_optimizer_default() {
1190        let optimizer = StringOptimizer::default();
1191        assert!(optimizer.read_only_strings.is_empty());
1192        assert!(optimizer.immutable_params.is_empty());
1193        assert!(optimizer.returned_strings.is_empty());
1194        assert!(optimizer.mixed_usage_strings.is_empty());
1195        assert!(optimizer.string_literal_count.is_empty());
1196        assert!(optimizer.interned_strings.is_empty());
1197        assert!(optimizer.interned_names.is_empty());
1198    }
1199
1200    // Tests for escape_char
1201    #[test]
1202    fn test_escape_char_backslash() {
1203        let result = escape_char('\\');
1204        assert_eq!(result, vec!['\\', '\\']);
1205    }
1206
1207    #[test]
1208    fn test_escape_char_quote() {
1209        let result = escape_char('"');
1210        assert_eq!(result, vec!['\\', '"']);
1211    }
1212
1213    #[test]
1214    fn test_escape_char_newline() {
1215        let result = escape_char('\n');
1216        assert_eq!(result, vec!['\\', 'n']);
1217    }
1218
1219    #[test]
1220    fn test_escape_char_carriage_return() {
1221        let result = escape_char('\r');
1222        assert_eq!(result, vec!['\\', 'r']);
1223    }
1224
1225    #[test]
1226    fn test_escape_char_tab() {
1227        let result = escape_char('\t');
1228        assert_eq!(result, vec!['\\', 't']);
1229    }
1230
1231    #[test]
1232    fn test_escape_char_normal() {
1233        let result = escape_char('a');
1234        assert_eq!(result, vec!['a']);
1235
1236        let result = escape_char('Z');
1237        assert_eq!(result, vec!['Z']);
1238
1239        let result = escape_char('5');
1240        assert_eq!(result, vec!['5']);
1241    }
1242
1243    // Tests for StringContext variants
1244    #[test]
1245    fn test_string_context_literal() {
1246        let ctx = StringContext::Literal("hello".to_string());
1247        if let StringContext::Literal(s) = ctx {
1248            assert_eq!(s, "hello");
1249        } else {
1250            panic!("Expected Literal");
1251        }
1252    }
1253
1254    #[test]
1255    fn test_string_context_parameter() {
1256        let ctx = StringContext::Parameter("name".to_string());
1257        if let StringContext::Parameter(s) = ctx {
1258            assert_eq!(s, "name");
1259        } else {
1260            panic!("Expected Parameter");
1261        }
1262    }
1263
1264    #[test]
1265    fn test_string_context_return() {
1266        let ctx = StringContext::Return;
1267        assert!(matches!(ctx, StringContext::Return));
1268    }
1269
1270    #[test]
1271    fn test_string_context_concatenation() {
1272        let ctx = StringContext::Concatenation;
1273        assert!(matches!(ctx, StringContext::Concatenation));
1274    }
1275
1276    // Tests for is_read_only and is_immutable
1277    #[test]
1278    fn test_is_read_only_true() {
1279        let mut optimizer = StringOptimizer::new();
1280        optimizer.read_only_strings.insert("readonly".to_string());
1281        assert!(optimizer.is_read_only("readonly"));
1282    }
1283
1284    #[test]
1285    fn test_is_read_only_false() {
1286        let optimizer = StringOptimizer::new();
1287        assert!(!optimizer.is_read_only("nonexistent"));
1288    }
1289
1290    #[test]
1291    fn test_is_immutable_param_true() {
1292        let mut optimizer = StringOptimizer::new();
1293        optimizer.immutable_params.insert("param".to_string());
1294        assert!(optimizer.is_immutable("param"));
1295    }
1296
1297    #[test]
1298    fn test_is_immutable_param_false() {
1299        let optimizer = StringOptimizer::new();
1300        assert!(!optimizer.is_immutable("nonexistent")); // Returns false if not in immutable_params set
1301    }
1302}