1use std::collections::HashSet;
35
36use fxrank_core::confidence::detection_confidence;
37use fxrank_core::effect::{Effect, EffectKind, Tier};
38use fxrank_core::score::weight_for_class;
39use libcst_native::{
40 Assert, AssignTargetExpression, Call, Expression, Name, Parameters, Raise, SmallStatement,
41 Statement, Suite,
42};
43
44use super::expr::render_expr;
45use super::{EffectSink, walk_own_body};
46use crate::functions::{FnBody, FnUnit};
47use crate::source::{SpanIndex, anchor_of_subslice};
48
49pub fn detect(unit: &FnUnit, span: &SpanIndex) -> Vec<(Effect, bool)> {
57 let params = collect_param_names(unit.params);
59
60 let mut globals: HashSet<String> = HashSet::new();
62 let mut nonlocals: HashSet<String> = HashSet::new();
63 let mut locals: HashSet<String> = HashSet::new();
64 prescan_body(&unit.body, &mut globals, &mut nonlocals, &mut locals);
65
66 let is_init = unit.symbol == "__init__";
68 let mut sink = MutSink {
69 params: ¶ms,
70 globals: &globals,
71 nonlocals: &nonlocals,
72 locals: &locals,
73 is_init,
74 span,
75 effects: Vec::new(),
76 };
77 walk_own_body(unit, &mut sink);
78 sink.effects
79}
80
81fn collect_param_names(params: &Parameters) -> HashSet<String> {
90 let mut out = HashSet::new();
91 let all = params
92 .posonly_params
93 .iter()
94 .chain(¶ms.params)
95 .chain(¶ms.kwonly_params);
96 for p in all {
97 out.insert(p.name.value.to_owned());
98 }
99 if let Some(libcst_native::StarArg::Param(p)) = ¶ms.star_arg {
100 out.insert(p.name.value.to_owned());
101 }
102 if let Some(p) = ¶ms.star_kwarg {
103 out.insert(p.name.value.to_owned());
104 }
105 out
106}
107
108fn prescan_body(
118 body: &FnBody,
119 globals: &mut HashSet<String>,
120 nonlocals: &mut HashSet<String>,
121 locals: &mut HashSet<String>,
122) {
123 match body {
124 FnBody::Suite(suite) => prescan_suite(suite, globals, nonlocals, locals),
125 FnBody::Expr(_) => {} }
127}
128
129fn prescan_suite(
130 suite: &Suite,
131 globals: &mut HashSet<String>,
132 nonlocals: &mut HashSet<String>,
133 locals: &mut HashSet<String>,
134) {
135 match suite {
136 Suite::IndentedBlock(b) => {
137 for stmt in &b.body {
138 prescan_stmt(stmt, globals, nonlocals, locals);
139 }
140 }
141 Suite::SimpleStatementSuite(s) => {
142 for small in &s.body {
143 prescan_small(small, globals, nonlocals, locals);
144 }
145 }
146 }
147}
148
149fn prescan_stmt(
150 stmt: &Statement,
151 globals: &mut HashSet<String>,
152 nonlocals: &mut HashSet<String>,
153 locals: &mut HashSet<String>,
154) {
155 match stmt {
156 Statement::Simple(line) => {
157 for small in &line.body {
158 prescan_small(small, globals, nonlocals, locals);
159 }
160 }
161 Statement::Compound(c) => prescan_compound(c, globals, nonlocals, locals),
162 }
163}
164
165fn prescan_compound(
166 compound: &libcst_native::CompoundStatement,
167 globals: &mut HashSet<String>,
168 nonlocals: &mut HashSet<String>,
169 locals: &mut HashSet<String>,
170) {
171 use libcst_native::CompoundStatement;
172 match compound {
173 CompoundStatement::FunctionDef(_) | CompoundStatement::ClassDef(_) => {}
175 CompoundStatement::If(i) => {
176 prescan_suite(&i.body, globals, nonlocals, locals);
177 if let Some(orelse) = &i.orelse {
178 prescan_orelse(orelse, globals, nonlocals, locals);
179 }
180 }
181 CompoundStatement::For(f) => {
182 prescan_suite(&f.body, globals, nonlocals, locals);
183 if let Some(orelse) = &f.orelse {
184 prescan_suite(&orelse.body, globals, nonlocals, locals);
185 }
186 }
187 CompoundStatement::While(w) => {
188 prescan_suite(&w.body, globals, nonlocals, locals);
189 if let Some(orelse) = &w.orelse {
190 prescan_suite(&orelse.body, globals, nonlocals, locals);
191 }
192 }
193 CompoundStatement::Try(t) => {
194 prescan_suite(&t.body, globals, nonlocals, locals);
195 for h in &t.handlers {
196 prescan_suite(&h.body, globals, nonlocals, locals);
197 }
198 if let Some(orelse) = &t.orelse {
199 prescan_suite(&orelse.body, globals, nonlocals, locals);
200 }
201 if let Some(fin) = &t.finalbody {
202 prescan_suite(&fin.body, globals, nonlocals, locals);
203 }
204 }
205 CompoundStatement::TryStar(t) => {
206 prescan_suite(&t.body, globals, nonlocals, locals);
207 for h in &t.handlers {
208 prescan_suite(&h.body, globals, nonlocals, locals);
209 }
210 if let Some(orelse) = &t.orelse {
211 prescan_suite(&orelse.body, globals, nonlocals, locals);
212 }
213 if let Some(fin) = &t.finalbody {
214 prescan_suite(&fin.body, globals, nonlocals, locals);
215 }
216 }
217 CompoundStatement::With(w) => {
218 prescan_suite(&w.body, globals, nonlocals, locals);
219 }
220 CompoundStatement::Match(m) => {
221 for case in &m.cases {
222 prescan_suite(&case.body, globals, nonlocals, locals);
223 }
224 }
225 }
226}
227
228fn prescan_orelse(
229 orelse: &libcst_native::OrElse,
230 globals: &mut HashSet<String>,
231 nonlocals: &mut HashSet<String>,
232 locals: &mut HashSet<String>,
233) {
234 match orelse {
235 libcst_native::OrElse::Elif(elif) => {
236 prescan_suite(&elif.body, globals, nonlocals, locals);
237 if let Some(inner) = &elif.orelse {
238 prescan_orelse(inner, globals, nonlocals, locals);
239 }
240 }
241 libcst_native::OrElse::Else(e) => {
242 prescan_suite(&e.body, globals, nonlocals, locals);
243 }
244 }
245}
246
247fn prescan_small(
248 small: &SmallStatement,
249 globals: &mut HashSet<String>,
250 nonlocals: &mut HashSet<String>,
251 locals: &mut HashSet<String>,
252) {
253 match small {
254 SmallStatement::Global(g) => {
255 for item in &g.names {
256 globals.insert(item.name.value.to_owned());
257 }
258 }
259 SmallStatement::Nonlocal(n) => {
260 for item in &n.names {
261 nonlocals.insert(item.name.value.to_owned());
262 }
263 }
264 SmallStatement::Assign(a) => {
265 for target in &a.targets {
268 if let AssignTargetExpression::Name(n) = &target.target {
269 locals.insert(n.value.to_owned());
270 }
271 }
272 }
273 SmallStatement::AnnAssign(a) => {
275 if let AssignTargetExpression::Name(n) = &a.target {
276 locals.insert(n.value.to_owned());
277 }
278 }
279 _ => {}
280 }
281}
282
283struct MutSink<'a> {
286 params: &'a HashSet<String>,
287 globals: &'a HashSet<String>,
288 nonlocals: &'a HashSet<String>,
289 locals: &'a HashSet<String>,
292 is_init: bool,
294 span: &'a SpanIndex<'a>,
295 effects: Vec<(Effect, bool)>,
296}
297
298impl EffectSink for MutSink<'_> {
299 fn on_call(&mut self, call: &Call) {
300 let Expression::Attribute(attr) = call.func.as_ref() else {
302 return;
303 };
304 if !is_mutating_method(attr.attr.value) {
305 return;
306 }
307 let Some(root) = root_name_of_expr(&attr.value) else {
309 return;
310 };
311 let line = name_line_expr(&attr.value, self.span);
312 let receiver = render_expr(&attr.value).unwrap_or_else(|| root.clone());
316 let evidence = format!("{receiver}.{}(…)", attr.attr.value);
317 self.classify_and_push(root, line, evidence);
318 }
319
320 fn on_assert(&mut self, _assert: &Assert) {}
321 fn on_raise(&mut self, _raise: &Raise) {}
322
323 fn on_assign_target(&mut self, target: &AssignTargetExpression, is_aug: bool) {
324 match target {
325 AssignTargetExpression::Attribute(attr) => {
327 if let Expression::Name(n) = attr.value.as_ref()
328 && n.value == "self"
329 {
330 let line = name_line(n, self.span);
331 if self.is_init {
332 self.push(
333 EffectKind::LocalMutation,
334 Tier::Heuristic,
335 line,
336 "self.x = … (constructor init, contained)".to_string(),
337 true,
338 );
339 } else {
340 self.push(
341 EffectKind::ThisMutation,
342 Tier::Heuristic,
343 line,
344 format!("self.{} = … (instance state)", attr.attr.value),
345 false,
346 );
347 }
348 return;
349 }
350 if let Some(root) = root_name_of_expr(&attr.value) {
352 let line = name_line_expr(&attr.value, self.span);
353 let evidence = format!("{root}.{} = …", attr.attr.value);
354 self.classify_and_push(root, line, evidence);
355 }
356 }
357 AssignTargetExpression::Name(n) if is_aug => {
364 let name = n.value.to_owned();
365 let line = name_line(n, self.span);
366 let evidence = format!("{name} += …");
367 self.classify_and_push(name, line, evidence);
368 }
369 AssignTargetExpression::Name(n)
372 if self.globals.contains(n.value) || self.nonlocals.contains(n.value) =>
373 {
374 let name = n.value.to_owned();
375 let line = name_line(n, self.span);
376 let evidence = format!("{name} = …");
377 self.classify_and_push(name, line, evidence);
378 }
379 AssignTargetExpression::Name(_) => {}
380 AssignTargetExpression::Subscript(sub) => {
382 if let Some(root) = root_name_of_expr(&sub.value) {
383 let line = name_line_expr(&sub.value, self.span);
384 let evidence = format!("{root}[…] = …");
385 self.classify_and_push(root, line, evidence);
386 }
387 }
388 _ => {}
390 }
391 }
392}
393
394impl MutSink<'_> {
395 fn classify_and_push(&mut self, root: String, line: usize, evidence: String) {
397 if root == "self" {
404 self.push(
405 EffectKind::ThisMutation,
406 Tier::Heuristic,
407 line,
408 evidence,
409 false,
410 );
411 return;
412 }
413
414 if self.globals.contains(&root) {
417 self.push(
418 EffectKind::GlobalMutation,
419 Tier::Exact,
420 line,
421 format!("global {root} ({evidence})"),
422 false,
423 );
424 return;
425 }
426
427 if self.nonlocals.contains(&root) {
430 self.push(
431 EffectKind::ThisMutation,
432 Tier::Exact,
433 line,
434 format!("nonlocal {root} ({evidence})"),
435 false,
436 );
437 return;
438 }
439
440 if self.params.contains(&root) {
442 self.push(
443 EffectKind::ParamMutation,
444 Tier::Heuristic,
445 line,
446 evidence,
447 false,
448 );
449 return;
450 }
451
452 if self.locals.contains(&root) {
454 self.push(EffectKind::LocalMutation, Tier::Exact, line, evidence, true);
455 }
456
457 }
461
462 fn push(
463 &mut self,
464 kind: EffectKind,
465 tier: Tier,
466 line: usize,
467 evidence: String,
468 contained: bool,
469 ) {
470 let class = kind.base_class();
471 self.effects.push((
472 Effect {
473 kind,
474 class,
475 discounted_to: None,
476 weight: weight_for_class(class),
477 line,
478 tier,
479 hidden: false,
480 evidence,
481 discount: None,
482 confidence: detection_confidence(tier, false, false),
483 },
484 contained,
485 ));
486 }
487}
488
489fn root_name_of_expr(expr: &Expression) -> Option<String> {
494 match expr {
495 Expression::Name(n) => Some(n.value.to_owned()),
496 Expression::Attribute(a) => root_name_of_expr(&a.value),
497 Expression::Subscript(s) => root_name_of_expr(&s.value),
498 Expression::Call(c) => root_name_of_expr(&c.func),
499 _ => None,
500 }
501}
502
503fn is_mutating_method(name: &str) -> bool {
505 matches!(
506 name,
507 "append"
508 | "extend"
509 | "insert"
510 | "remove"
511 | "pop"
512 | "clear"
513 | "sort"
514 | "reverse"
515 | "update"
516 | "add"
517 | "discard"
518 | "setdefault"
519 )
520}
521
522fn name_line_expr(expr: &Expression, span: &SpanIndex) -> usize {
524 leftmost_name(expr).map(|n| name_line(n, span)).unwrap_or(0)
525}
526
527fn leftmost_name<'a>(expr: &'a Expression<'a>) -> Option<&'a Name<'a>> {
529 match expr {
530 Expression::Name(n) => Some(n),
531 Expression::Attribute(a) => leftmost_name(&a.value),
532 Expression::Subscript(s) => leftmost_name(&s.value),
533 Expression::Call(c) => leftmost_name(&c.func),
534 _ => None,
535 }
536}
537
538fn name_line(name: &Name, span: &SpanIndex) -> usize {
540 span.line_col(anchor_of_subslice(span.src(), name.value)).0
541}
542
543#[cfg(test)]
546mod tests {
547 use super::*;
548 use crate::functions;
549 use fxrank_core::effect::EffectKind::{self, *};
550 use std::collections::HashMap;
551
552 fn mutation_effects(name: &str) -> HashMap<String, Vec<(EffectKind, bool)>> {
555 let src = std::fs::read_to_string(format!("tests/fixtures/{name}.py")).unwrap();
556 let module = libcst_native::parse_module(&src, None).unwrap();
557 let span = crate::source::SpanIndex::new(&src);
558 let anchors = crate::source::lambda_anchors(&src).expect("tokenize must succeed");
559 let (units, _) = functions::collect(&module, &src, &span, &anchors);
560 let mut out: HashMap<String, Vec<(EffectKind, bool)>> = HashMap::new();
561 for unit in &units {
562 let pairs = detect(unit, &span);
563 out.insert(
564 unit.symbol.clone(),
565 pairs.iter().map(|(e, c)| (e.kind, *c)).collect(),
566 );
567 }
568 out
569 }
570
571 fn mutation_evidence(name: &str) -> HashMap<String, Vec<(EffectKind, bool, String)>> {
573 let src = std::fs::read_to_string(format!("tests/fixtures/{name}.py")).unwrap();
574 let module = libcst_native::parse_module(&src, None).unwrap();
575 let span = crate::source::SpanIndex::new(&src);
576 let anchors = crate::source::lambda_anchors(&src).expect("tokenize must succeed");
577 let (units, _) = functions::collect(&module, &src, &span, &anchors);
578 let mut out: HashMap<String, Vec<(EffectKind, bool, String)>> = HashMap::new();
579 for unit in &units {
580 let pairs = detect(unit, &span);
581 out.insert(
582 unit.symbol.clone(),
583 pairs
584 .iter()
585 .map(|(e, c)| (e.kind, *c, e.evidence.clone()))
586 .collect(),
587 );
588 }
589 out
590 }
591
592 #[test]
593 fn classifies_mutation_by_escape() {
594 let m = mutation_effects("mutation");
595
596 assert!(
598 m["uses_global"].contains(&(GlobalMutation, false)),
599 "uses_global should have GlobalMutation(contained=false), got: {:?}",
600 m["uses_global"]
601 );
602
603 assert!(
605 m["bump"].contains(&(ThisMutation, false)),
606 "bump should have ThisMutation(contained=false), got: {:?}",
607 m["bump"]
608 );
609
610 assert!(
612 m["mutates_param"].contains(&(ParamMutation, false)),
613 "mutates_param should have ParamMutation(contained=false), got: {:?}",
614 m["mutates_param"]
615 );
616
617 assert!(
619 m["builds_local"].contains(&(LocalMutation, true)),
620 "builds_local should have LocalMutation(contained=true), got: {:?}",
621 m["builds_local"]
622 );
623
624 assert!(
626 m["__init__"].contains(&(LocalMutation, true)),
627 "__init__ should have LocalMutation(contained=true), got: {:?}",
628 m["__init__"]
629 );
630 }
631
632 #[test]
637 fn plain_assign_to_global_nonlocal_names_escapes() {
638 let m = mutation_effects("mutation");
639
640 assert!(
642 m["plain_global_rebind"].contains(&(GlobalMutation, false)),
643 "plain `=` to a global name must emit GlobalMutation(false), got: {:?}",
644 m["plain_global_rebind"]
645 );
646
647 assert!(
649 m["plain_nonlocal_rebind"].contains(&(ThisMutation, false)),
650 "plain `=` to a nonlocal name must emit ThisMutation(false), got: {:?}",
651 m["plain_nonlocal_rebind"]
652 );
653
654 assert!(
656 m["plain_local_binding"].is_empty(),
657 "plain `=` to a true local must emit NO mutation, got: {:?}",
658 m["plain_local_binding"]
659 );
660 }
661
662 #[test]
668 fn self_method_and_subscript_mutations_escape_even_in_init() {
669 let m = mutation_effects("mutation");
670
671 assert!(
673 m["__init__"].contains(&(LocalMutation, true)),
674 "direct `self.attr = …` in __init__ stays LocalMutation(true), got: {:?}",
675 m["__init__"]
676 );
677
678 assert!(
681 m["__init__"].contains(&(ThisMutation, false)),
682 "`self.items.append(…)` in __init__ must be ThisMutation(false), got: {:?}",
683 m["__init__"]
684 );
685
686 assert!(
688 m["store"].contains(&(ThisMutation, false)),
689 "`self[i] = v` must be ThisMutation(false), got: {:?}",
690 m["store"]
691 );
692 }
693
694 #[test]
698 fn mutating_method_evidence_uses_full_receiver() {
699 let m = mutation_evidence("mutation");
700 let init = &m["__init__"];
701 let append = init
702 .iter()
703 .find(|(k, _, _)| *k == ThisMutation)
704 .unwrap_or_else(|| panic!("expected a ThisMutation in __init__, got: {init:?}"));
705 assert!(
706 append.2.contains("self.items"),
707 "evidence must name the full receiver `self.items`, got: {:?}",
708 append.2
709 );
710 }
711}