1use depyler_hir::hir::{AssignTarget, HirExpr, HirStmt, Type};
6use std::collections::HashSet;
7
8pub fn uses_hashmap(ty: &Type) -> bool {
10 match ty {
11 Type::Dict(_, _) => true,
12 Type::List(inner) | Type::Optional(inner) => uses_hashmap(inner),
13 Type::Tuple(types) => types.iter().any(uses_hashmap),
14 Type::Function { params, ret } => params.iter().any(uses_hashmap) || uses_hashmap(ret),
15 _ => false,
16 }
17}
18
19pub fn stmt_uses_hashmap(stmt: &HirStmt) -> bool {
21 match stmt {
22 HirStmt::Assign { value, .. } => expr_uses_hashmap(value),
23 HirStmt::Return(Some(expr)) => expr_uses_hashmap(expr),
24 HirStmt::If {
25 condition,
26 then_body,
27 else_body,
28 } => {
29 expr_uses_hashmap(condition)
30 || body_uses_hashmap(then_body)
31 || else_body
32 .as_ref()
33 .is_some_and(|body| body_uses_hashmap(body))
34 }
35 HirStmt::While { condition, body } => {
36 expr_uses_hashmap(condition) || body_uses_hashmap(body)
37 }
38 HirStmt::For { iter, body, .. } => expr_uses_hashmap(iter) || body_uses_hashmap(body),
39 HirStmt::Expr(expr) => expr_uses_hashmap(expr),
40 _ => false,
41 }
42}
43
44pub fn body_uses_hashmap(body: &[HirStmt]) -> bool {
46 body.iter().any(stmt_uses_hashmap)
47}
48
49pub fn expr_uses_hashmap(expr: &HirExpr) -> bool {
51 match expr {
52 HirExpr::Dict(_) => true,
53 HirExpr::Binary { left, right, .. } => expr_uses_hashmap(left) || expr_uses_hashmap(right),
54 HirExpr::Unary { operand, .. } => expr_uses_hashmap(operand),
55 HirExpr::Call { args, .. } => args.iter().any(expr_uses_hashmap),
56 HirExpr::Index { base, index } => expr_uses_hashmap(base) || expr_uses_hashmap(index),
57 HirExpr::List(items) | HirExpr::Tuple(items) => items.iter().any(expr_uses_hashmap),
58 _ => false,
59 }
60}
61
62#[derive(Debug, Default)]
64pub struct ScopeTracker {
65 declared_vars: Vec<HashSet<String>>,
66}
67
68impl ScopeTracker {
69 pub fn new() -> Self {
70 Self {
71 declared_vars: vec![HashSet::new()],
72 }
73 }
74
75 pub fn enter_scope(&mut self) {
76 self.declared_vars.push(HashSet::new());
77 }
78
79 pub fn exit_scope(&mut self) {
80 self.declared_vars.pop();
81 }
82
83 pub fn declare(&mut self, name: &str) {
84 if let Some(scope) = self.declared_vars.last_mut() {
85 scope.insert(name.to_string());
86 }
87 }
88
89 pub fn is_declared(&self, name: &str) -> bool {
90 self.declared_vars.iter().any(|scope| scope.contains(name))
91 }
92
93 pub fn is_declared_in_current_scope(&self, name: &str) -> bool {
94 self.declared_vars
95 .last()
96 .is_some_and(|scope| scope.contains(name))
97 }
98
99 pub fn depth(&self) -> usize {
100 self.declared_vars.len()
101 }
102}
103
104pub fn extract_var_name(target: &AssignTarget) -> Option<String> {
106 match target {
107 AssignTarget::Symbol(name) => Some(name.clone()),
108 AssignTarget::Attribute { value: _, attr: _ } => None,
109 AssignTarget::Index { base: _, index: _ } => None,
110 AssignTarget::Tuple(targets) => {
111 targets.first().and_then(extract_var_name)
113 }
114 }
115}
116
117pub fn is_simple_literal(expr: &HirExpr) -> bool {
119 matches!(expr, HirExpr::Literal(_))
120}
121
122pub fn is_constant_expr(expr: &HirExpr) -> bool {
124 match expr {
125 HirExpr::Literal(_) => true,
126 HirExpr::List(items) | HirExpr::Tuple(items) => items.iter().all(is_constant_expr),
127 HirExpr::Unary { operand, .. } => is_constant_expr(operand),
128 HirExpr::Binary { left, right, .. } => is_constant_expr(left) && is_constant_expr(right),
129 _ => false,
130 }
131}
132
133pub fn expr_complexity(expr: &HirExpr) -> usize {
135 match expr {
136 HirExpr::Literal(_) => 1,
137 HirExpr::Var(_) => 1,
138 HirExpr::List(items) | HirExpr::Tuple(items) => {
139 1 + items.iter().map(expr_complexity).sum::<usize>()
140 }
141 HirExpr::Dict(pairs) => {
142 1 + pairs
143 .iter()
144 .map(|(k, v)| expr_complexity(k) + expr_complexity(v))
145 .sum::<usize>()
146 }
147 HirExpr::Binary { left, right, .. } => 1 + expr_complexity(left) + expr_complexity(right),
148 HirExpr::Unary { operand, .. } => 1 + expr_complexity(operand),
149 HirExpr::Call { args, .. } => 2 + args.iter().map(expr_complexity).sum::<usize>(),
150 HirExpr::Index { base, index } => 1 + expr_complexity(base) + expr_complexity(index),
151 HirExpr::Attribute { value, .. } => 1 + expr_complexity(value),
152 HirExpr::MethodCall { object, args, .. } => {
153 2 + expr_complexity(object) + args.iter().map(expr_complexity).sum::<usize>()
154 }
155 HirExpr::Lambda { body, .. } => 3 + expr_complexity(body),
156 HirExpr::IfExpr { test, body, orelse } => {
157 2 + expr_complexity(test) + expr_complexity(body) + expr_complexity(orelse)
158 }
159 HirExpr::ListComp {
160 element,
161 generators,
162 ..
163 } => 3 + expr_complexity(element) + generators.len(),
164 _ => 1,
165 }
166}
167
168pub fn needs_boxing(ty: &Type) -> bool {
170 match ty {
171 Type::Custom(name) => {
172 name.starts_with("Box<") || name.contains("Rc<") || name.contains("Arc<")
173 }
174 Type::List(inner) => needs_boxing(inner),
175 Type::Optional(inner) => needs_boxing(inner),
176 Type::Tuple(types) => types.iter().any(needs_boxing),
177 _ => false,
178 }
179}
180
181pub fn is_reference_type(ty: &Type) -> bool {
183 match ty {
184 Type::String | Type::List(_) | Type::Dict(_, _) | Type::Set(_) => true,
185 Type::Custom(name) => {
186 name.starts_with("Vec<")
187 || name.starts_with("HashMap<")
188 || name.starts_with("String")
189 || name.starts_with("&")
190 }
191 _ => false,
192 }
193}
194
195pub fn is_primitive_type(ty: &Type) -> bool {
197 matches!(ty, Type::Int | Type::Float | Type::Bool)
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203 use depyler_hir::hir::{BinOp, Literal};
204
205 fn int_expr(n: i64) -> HirExpr {
207 HirExpr::Literal(Literal::Int(n))
208 }
209
210 fn float_expr(f: f64) -> HirExpr {
211 HirExpr::Literal(Literal::Float(f))
212 }
213
214 fn bool_expr(b: bool) -> HirExpr {
215 HirExpr::Literal(Literal::Bool(b))
216 }
217
218 fn string_expr(s: &str) -> HirExpr {
219 HirExpr::Literal(Literal::String(s.to_string()))
220 }
221
222 fn none_expr() -> HirExpr {
223 HirExpr::Literal(Literal::None)
224 }
225
226 fn var_expr(name: &str) -> HirExpr {
227 HirExpr::Var(name.to_string())
228 }
229
230 #[test]
231 fn test_uses_hashmap_dict() {
232 assert!(uses_hashmap(&Type::Dict(
233 Box::new(Type::String),
234 Box::new(Type::Int)
235 )));
236 }
237
238 #[test]
239 fn test_uses_hashmap_primitives() {
240 assert!(!uses_hashmap(&Type::Int));
241 assert!(!uses_hashmap(&Type::Float));
242 assert!(!uses_hashmap(&Type::Bool));
243 assert!(!uses_hashmap(&Type::String));
244 }
245
246 #[test]
247 fn test_uses_hashmap_nested() {
248 let nested = Type::List(Box::new(Type::Dict(
249 Box::new(Type::String),
250 Box::new(Type::Int),
251 )));
252 assert!(uses_hashmap(&nested));
253 }
254
255 #[test]
256 fn test_uses_hashmap_optional() {
257 let opt_dict = Type::Optional(Box::new(Type::Dict(
258 Box::new(Type::String),
259 Box::new(Type::Int),
260 )));
261 assert!(uses_hashmap(&opt_dict));
262
263 let opt_int = Type::Optional(Box::new(Type::Int));
264 assert!(!uses_hashmap(&opt_int));
265 }
266
267 #[test]
268 fn test_uses_hashmap_tuple() {
269 let tuple_with_dict = Type::Tuple(vec![
270 Type::Int,
271 Type::Dict(Box::new(Type::String), Box::new(Type::Int)),
272 ]);
273 assert!(uses_hashmap(&tuple_with_dict));
274
275 let tuple_no_dict = Type::Tuple(vec![Type::Int, Type::String]);
276 assert!(!uses_hashmap(&tuple_no_dict));
277 }
278
279 #[test]
280 fn test_uses_hashmap_function() {
281 let func_with_dict_param = Type::Function {
282 params: vec![Type::Dict(Box::new(Type::String), Box::new(Type::Int))],
283 ret: Box::new(Type::Int),
284 };
285 assert!(uses_hashmap(&func_with_dict_param));
286
287 let func_with_dict_ret = Type::Function {
288 params: vec![Type::Int],
289 ret: Box::new(Type::Dict(Box::new(Type::String), Box::new(Type::Int))),
290 };
291 assert!(uses_hashmap(&func_with_dict_ret));
292 }
293
294 #[test]
295 fn test_expr_uses_hashmap_dict() {
296 let dict_expr = HirExpr::Dict(vec![]);
297 assert!(expr_uses_hashmap(&dict_expr));
298 }
299
300 #[test]
301 fn test_expr_uses_hashmap_primitives() {
302 assert!(!expr_uses_hashmap(&int_expr(42)));
303 assert!(!expr_uses_hashmap(&float_expr(3.15)));
304 assert!(!expr_uses_hashmap(&bool_expr(true)));
305 assert!(!expr_uses_hashmap(&string_expr("test")));
306 }
307
308 #[test]
309 fn test_expr_uses_hashmap_nested() {
310 let nested = HirExpr::List(vec![HirExpr::Dict(vec![])]);
311 assert!(expr_uses_hashmap(&nested));
312 }
313
314 #[test]
315 fn test_expr_uses_hashmap_binary() {
316 let binary_with_dict = HirExpr::Binary {
317 left: Box::new(HirExpr::Dict(vec![])),
318 op: BinOp::Add,
319 right: Box::new(int_expr(1)),
320 };
321 assert!(expr_uses_hashmap(&binary_with_dict));
322
323 let binary_no_dict = HirExpr::Binary {
324 left: Box::new(int_expr(1)),
325 op: BinOp::Add,
326 right: Box::new(int_expr(2)),
327 };
328 assert!(!expr_uses_hashmap(&binary_no_dict));
329 }
330
331 #[test]
332 fn test_scope_tracker_new() {
333 let tracker = ScopeTracker::new();
334 assert_eq!(tracker.depth(), 1);
335 }
336
337 #[test]
338 fn test_scope_tracker_declare() {
339 let mut tracker = ScopeTracker::new();
340 tracker.declare("x");
341 assert!(tracker.is_declared("x"));
342 assert!(!tracker.is_declared("y"));
343 }
344
345 #[test]
346 fn test_scope_tracker_nested_scopes() {
347 let mut tracker = ScopeTracker::new();
348 tracker.declare("outer");
349 tracker.enter_scope();
350 tracker.declare("inner");
351
352 assert!(tracker.is_declared("outer"));
353 assert!(tracker.is_declared("inner"));
354 assert!(tracker.is_declared_in_current_scope("inner"));
355 assert!(!tracker.is_declared_in_current_scope("outer"));
356
357 tracker.exit_scope();
358 assert!(tracker.is_declared("outer"));
359 assert!(!tracker.is_declared("inner"));
360 }
361
362 #[test]
363 fn test_scope_tracker_depth() {
364 let mut tracker = ScopeTracker::new();
365 assert_eq!(tracker.depth(), 1);
366 tracker.enter_scope();
367 assert_eq!(tracker.depth(), 2);
368 tracker.enter_scope();
369 assert_eq!(tracker.depth(), 3);
370 tracker.exit_scope();
371 assert_eq!(tracker.depth(), 2);
372 }
373
374 #[test]
375 fn test_extract_var_name_symbol() {
376 let target = AssignTarget::Symbol("x".to_string());
377 assert_eq!(extract_var_name(&target), Some("x".to_string()));
378 }
379
380 #[test]
381 fn test_extract_var_name_attribute() {
382 let target = AssignTarget::Attribute {
383 value: Box::new(var_expr("obj")),
384 attr: "field".to_string(),
385 };
386 assert_eq!(extract_var_name(&target), None);
387 }
388
389 #[test]
390 fn test_extract_var_name_index() {
391 let target = AssignTarget::Index {
392 base: Box::new(var_expr("arr")),
393 index: Box::new(int_expr(0)),
394 };
395 assert_eq!(extract_var_name(&target), None);
396 }
397
398 #[test]
399 fn test_extract_var_name_tuple() {
400 let target = AssignTarget::Tuple(vec![
401 AssignTarget::Symbol("a".to_string()),
402 AssignTarget::Symbol("b".to_string()),
403 ]);
404 assert_eq!(extract_var_name(&target), Some("a".to_string()));
405 }
406
407 #[test]
408 fn test_is_simple_literal() {
409 assert!(is_simple_literal(&int_expr(42)));
410 assert!(is_simple_literal(&float_expr(3.15)));
411 assert!(is_simple_literal(&bool_expr(true)));
412 assert!(is_simple_literal(&string_expr("test")));
413 assert!(is_simple_literal(&none_expr()));
414 assert!(!is_simple_literal(&var_expr("x")));
415 assert!(!is_simple_literal(&HirExpr::List(vec![])));
416 }
417
418 #[test]
419 fn test_is_constant_expr_literals() {
420 assert!(is_constant_expr(&int_expr(42)));
421 assert!(is_constant_expr(&float_expr(3.15)));
422 assert!(is_constant_expr(&bool_expr(true)));
423 assert!(is_constant_expr(&string_expr("test")));
424 assert!(is_constant_expr(&none_expr()));
425 }
426
427 #[test]
428 fn test_is_constant_expr_list() {
429 let const_list = HirExpr::List(vec![int_expr(1), int_expr(2)]);
430 assert!(is_constant_expr(&const_list));
431
432 let non_const_list = HirExpr::List(vec![var_expr("x")]);
433 assert!(!is_constant_expr(&non_const_list));
434 }
435
436 #[test]
437 fn test_is_constant_expr_unary() {
438 use depyler_hir::hir::UnaryOp;
439 let const_unary = HirExpr::Unary {
440 op: UnaryOp::Neg,
441 operand: Box::new(int_expr(42)),
442 };
443 assert!(is_constant_expr(&const_unary));
444 }
445
446 #[test]
447 fn test_is_constant_expr_binary() {
448 let const_binary = HirExpr::Binary {
449 left: Box::new(int_expr(1)),
450 op: BinOp::Add,
451 right: Box::new(int_expr(2)),
452 };
453 assert!(is_constant_expr(&const_binary));
454
455 let non_const_binary = HirExpr::Binary {
456 left: Box::new(var_expr("x")),
457 op: BinOp::Add,
458 right: Box::new(int_expr(2)),
459 };
460 assert!(!is_constant_expr(&non_const_binary));
461 }
462
463 #[test]
464 fn test_expr_complexity_literals() {
465 assert_eq!(expr_complexity(&int_expr(42)), 1);
466 assert_eq!(expr_complexity(&float_expr(3.15)), 1);
467 assert_eq!(expr_complexity(&bool_expr(true)), 1);
468 assert_eq!(expr_complexity(&string_expr("test")), 1);
469 assert_eq!(expr_complexity(&none_expr()), 1);
470 }
471
472 #[test]
473 fn test_expr_complexity_variable() {
474 assert_eq!(expr_complexity(&var_expr("x")), 1);
475 }
476
477 #[test]
478 fn test_expr_complexity_list() {
479 let list = HirExpr::List(vec![int_expr(1), int_expr(2), int_expr(3)]);
480 assert_eq!(expr_complexity(&list), 4); }
482
483 #[test]
484 fn test_expr_complexity_binary() {
485 let binary = HirExpr::Binary {
486 left: Box::new(int_expr(1)),
487 op: BinOp::Add,
488 right: Box::new(int_expr(2)),
489 };
490 assert_eq!(expr_complexity(&binary), 3); }
492
493 #[test]
494 fn test_needs_boxing() {
495 assert!(needs_boxing(&Type::Custom("Box<Node>".to_string())));
496 assert!(needs_boxing(&Type::Custom("Rc<Node>".to_string())));
497 assert!(needs_boxing(&Type::Custom("Arc<Node>".to_string())));
498 assert!(!needs_boxing(&Type::Int));
499 assert!(!needs_boxing(&Type::String));
500 }
501
502 #[test]
503 fn test_needs_boxing_nested() {
504 let nested = Type::List(Box::new(Type::Custom("Box<Node>".to_string())));
505 assert!(needs_boxing(&nested));
506 }
507
508 #[test]
509 fn test_is_reference_type() {
510 assert!(is_reference_type(&Type::String));
511 assert!(is_reference_type(&Type::List(Box::new(Type::Int))));
512 assert!(is_reference_type(&Type::Dict(
513 Box::new(Type::String),
514 Box::new(Type::Int)
515 )));
516 assert!(is_reference_type(&Type::Set(Box::new(Type::Int))));
517 assert!(is_reference_type(&Type::Custom("Vec<i32>".to_string())));
518 assert!(is_reference_type(&Type::Custom(
519 "HashMap<String, i32>".to_string()
520 )));
521 assert!(!is_reference_type(&Type::Int));
522 assert!(!is_reference_type(&Type::Float));
523 }
524
525 #[test]
526 fn test_is_primitive_type() {
527 assert!(is_primitive_type(&Type::Int));
528 assert!(is_primitive_type(&Type::Float));
529 assert!(is_primitive_type(&Type::Bool));
530 assert!(!is_primitive_type(&Type::String));
531 assert!(!is_primitive_type(&Type::List(Box::new(Type::Int))));
532 }
533
534 #[test]
535 fn test_stmt_uses_hashmap_assign() {
536 let stmt = HirStmt::Assign {
537 target: AssignTarget::Symbol("x".to_string()),
538 value: HirExpr::Dict(vec![]),
539 type_annotation: None,
540 };
541 assert!(stmt_uses_hashmap(&stmt));
542 }
543
544 #[test]
545 fn test_stmt_uses_hashmap_return() {
546 let stmt = HirStmt::Return(Some(HirExpr::Dict(vec![])));
547 assert!(stmt_uses_hashmap(&stmt));
548
549 let stmt_no_dict = HirStmt::Return(Some(int_expr(42)));
550 assert!(!stmt_uses_hashmap(&stmt_no_dict));
551 }
552
553 #[test]
554 fn test_stmt_uses_hashmap_expr() {
555 let stmt = HirStmt::Expr(HirExpr::Dict(vec![]));
556 assert!(stmt_uses_hashmap(&stmt));
557 }
558
559 #[test]
560 fn test_body_uses_hashmap() {
561 let body_with_dict = vec![
562 HirStmt::Expr(int_expr(1)),
563 HirStmt::Expr(HirExpr::Dict(vec![])),
564 ];
565 assert!(body_uses_hashmap(&body_with_dict));
566
567 let body_no_dict = vec![HirStmt::Expr(int_expr(1)), HirStmt::Expr(int_expr(2))];
568 assert!(!body_uses_hashmap(&body_no_dict));
569 }
570}