1use depyler_hir::hir::{AssignTarget, HirExpr, HirFunction, HirStmt, Literal, Type};
2use std::collections::{HashMap, HashSet};
3
4#[derive(Debug, Default)]
6pub struct StringOptimizer {
7 read_only_strings: HashSet<String>,
9 immutable_params: HashSet<String>,
11 returned_strings: HashSet<String>,
13 mixed_usage_strings: HashSet<String>,
15 string_literal_count: HashMap<String, usize>,
17 interned_strings: HashSet<String>,
19 interned_names: HashMap<String, String>,
21}
22
23#[derive(Debug, Clone, PartialEq)]
25pub enum OptimalStringType {
26 StaticStr,
28 BorrowedStr { lifetime: Option<String> },
30 OwnedString,
32 CowStr,
34}
35
36impl StringOptimizer {
37 pub fn new() -> Self {
38 Self::default()
39 }
40
41 pub fn analyze_function(&mut self, func: &HirFunction) {
43 for param in &func.params {
45 if matches!(param.ty, Type::String) {
46 self.immutable_params.insert(param.name.clone());
47 }
48 }
49
50 for stmt in &func.body {
52 self.analyze_stmt(stmt);
53 }
54
55 for param in self.immutable_params.clone() {
57 if !self.is_immutable(¶m) {
58 self.immutable_params.remove(¶m);
59 }
60 }
61 }
62
63 pub fn get_optimal_type(&self, context: &StringContext) -> OptimalStringType {
65 match context {
66 StringContext::Literal(s) => {
67 if self.mixed_usage_strings.contains(s) {
70 OptimalStringType::CowStr
71 } else if self.returned_strings.contains(s) {
72 OptimalStringType::OwnedString
74 } else {
75 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 pub fn finalize_interned_names(&mut self) {
207 if !self.interned_names.is_empty() {
208 return;
210 }
211
212 let mut name_map: HashMap<String, Vec<String>> = HashMap::new();
214
215 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 for (base_name, strings) in name_map {
223 if strings.len() == 1 {
224 self.interned_names.insert(strings[0].clone(), base_name);
226 } else {
227 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 fn generate_base_const_name(&self, s: &str) -> String {
238 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 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 matches!(func.as_str(), "str" | "format" | "to_string" | "join")
314 }
315 _ => false,
316 }
317 }
318
319 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 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 pub fn should_intern(&self, s: &str) -> bool {
342 self.interned_strings.contains(s)
343 }
344
345 pub fn get_interned_name(&self, s: &str) -> Option<String> {
348 self.interned_names.get(s).cloned()
350 }
351
352 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#[derive(Debug, Clone)]
370pub enum StringContext {
371 Literal(String),
373 Parameter(String),
375 Return,
377 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
392pub 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
435fn 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 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 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 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 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 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 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 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 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 }
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 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 }
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 }
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 #[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 #[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 #[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 #[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 #[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")); }
1302}