use std::collections::HashSet;
use fxrank_core::confidence::detection_confidence;
use fxrank_core::effect::{Effect, EffectKind, Tier};
use fxrank_core::score::weight_for_class;
use libcst_native::{
Assert, AssignTargetExpression, Call, Expression, Name, Parameters, Raise, SmallStatement,
Statement, Suite,
};
use super::expr::render_expr;
use super::{EffectSink, walk_own_body};
use crate::functions::{FnBody, FnUnit};
use crate::imports::Imports;
use crate::source::{SpanIndex, anchor_of_subslice};
pub fn detect(
unit: &FnUnit,
imports: &Imports,
module_bindings: &HashSet<String>,
span: &SpanIndex,
) -> Vec<(Effect, bool)> {
let params = collect_param_names(unit.params);
let mut globals: HashSet<String> = HashSet::new();
let mut nonlocals: HashSet<String> = HashSet::new();
let mut locals: HashSet<String> = HashSet::new();
prescan_body(&unit.body, &mut globals, &mut nonlocals, &mut locals);
let is_init = unit.symbol == "__init__";
let mut sink = MutSink {
params: ¶ms,
globals: &globals,
nonlocals: &nonlocals,
locals: &locals,
imports,
module_bindings,
is_init,
span,
effects: Vec::new(),
};
walk_own_body(unit, &mut sink);
sink.effects
}
fn collect_param_names(params: &Parameters) -> HashSet<String> {
let mut out = HashSet::new();
let all = params
.posonly_params
.iter()
.chain(¶ms.params)
.chain(¶ms.kwonly_params);
for p in all {
out.insert(p.name.value.to_owned());
}
if let Some(libcst_native::StarArg::Param(p)) = ¶ms.star_arg {
out.insert(p.name.value.to_owned());
}
if let Some(p) = ¶ms.star_kwarg {
out.insert(p.name.value.to_owned());
}
out
}
fn prescan_body(
body: &FnBody,
globals: &mut HashSet<String>,
nonlocals: &mut HashSet<String>,
locals: &mut HashSet<String>,
) {
match body {
FnBody::Suite(suite) => prescan_suite(suite, globals, nonlocals, locals),
FnBody::Expr(_) => {} }
}
fn prescan_suite(
suite: &Suite,
globals: &mut HashSet<String>,
nonlocals: &mut HashSet<String>,
locals: &mut HashSet<String>,
) {
match suite {
Suite::IndentedBlock(b) => {
for stmt in &b.body {
prescan_stmt(stmt, globals, nonlocals, locals);
}
}
Suite::SimpleStatementSuite(s) => {
for small in &s.body {
prescan_small(small, globals, nonlocals, locals);
}
}
}
}
fn prescan_stmt(
stmt: &Statement,
globals: &mut HashSet<String>,
nonlocals: &mut HashSet<String>,
locals: &mut HashSet<String>,
) {
match stmt {
Statement::Simple(line) => {
for small in &line.body {
prescan_small(small, globals, nonlocals, locals);
}
}
Statement::Compound(c) => prescan_compound(c, globals, nonlocals, locals),
}
}
fn prescan_compound(
compound: &libcst_native::CompoundStatement,
globals: &mut HashSet<String>,
nonlocals: &mut HashSet<String>,
locals: &mut HashSet<String>,
) {
use libcst_native::CompoundStatement;
match compound {
CompoundStatement::FunctionDef(_) | CompoundStatement::ClassDef(_) => {}
CompoundStatement::If(i) => {
prescan_suite(&i.body, globals, nonlocals, locals);
if let Some(orelse) = &i.orelse {
prescan_orelse(orelse, globals, nonlocals, locals);
}
}
CompoundStatement::For(f) => {
crate::imports::collect_target_names(&f.target, locals);
prescan_suite(&f.body, globals, nonlocals, locals);
if let Some(orelse) = &f.orelse {
prescan_suite(&orelse.body, globals, nonlocals, locals);
}
}
CompoundStatement::While(w) => {
prescan_suite(&w.body, globals, nonlocals, locals);
if let Some(orelse) = &w.orelse {
prescan_suite(&orelse.body, globals, nonlocals, locals);
}
}
CompoundStatement::Try(t) => {
prescan_suite(&t.body, globals, nonlocals, locals);
for h in &t.handlers {
if let Some(asname) = &h.name {
crate::imports::collect_target_names(&asname.name, locals);
}
prescan_suite(&h.body, globals, nonlocals, locals);
}
if let Some(orelse) = &t.orelse {
prescan_suite(&orelse.body, globals, nonlocals, locals);
}
if let Some(fin) = &t.finalbody {
prescan_suite(&fin.body, globals, nonlocals, locals);
}
}
CompoundStatement::TryStar(t) => {
prescan_suite(&t.body, globals, nonlocals, locals);
for h in &t.handlers {
if let Some(asname) = &h.name {
crate::imports::collect_target_names(&asname.name, locals);
}
prescan_suite(&h.body, globals, nonlocals, locals);
}
if let Some(orelse) = &t.orelse {
prescan_suite(&orelse.body, globals, nonlocals, locals);
}
if let Some(fin) = &t.finalbody {
prescan_suite(&fin.body, globals, nonlocals, locals);
}
}
CompoundStatement::With(w) => {
for item in &w.items {
if let Some(asname) = &item.asname {
crate::imports::collect_target_names(&asname.name, locals);
}
}
prescan_suite(&w.body, globals, nonlocals, locals);
}
CompoundStatement::Match(m) => {
for case in &m.cases {
prescan_suite(&case.body, globals, nonlocals, locals);
}
}
}
}
fn prescan_orelse(
orelse: &libcst_native::OrElse,
globals: &mut HashSet<String>,
nonlocals: &mut HashSet<String>,
locals: &mut HashSet<String>,
) {
match orelse {
libcst_native::OrElse::Elif(elif) => {
prescan_suite(&elif.body, globals, nonlocals, locals);
if let Some(inner) = &elif.orelse {
prescan_orelse(inner, globals, nonlocals, locals);
}
}
libcst_native::OrElse::Else(e) => {
prescan_suite(&e.body, globals, nonlocals, locals);
}
}
}
fn prescan_small(
small: &SmallStatement,
globals: &mut HashSet<String>,
nonlocals: &mut HashSet<String>,
locals: &mut HashSet<String>,
) {
match small {
SmallStatement::Global(g) => {
for item in &g.names {
globals.insert(item.name.value.to_owned());
}
}
SmallStatement::Nonlocal(n) => {
for item in &n.names {
nonlocals.insert(item.name.value.to_owned());
}
}
SmallStatement::Assign(a) => {
for target in &a.targets {
crate::imports::collect_target_names(&target.target, locals);
}
}
SmallStatement::AnnAssign(a) => {
crate::imports::collect_target_names(&a.target, locals);
}
SmallStatement::AugAssign(a) => {
if let AssignTargetExpression::Name(n) = &a.target {
locals.insert(n.value.to_owned());
}
}
_ => {}
}
}
struct MutSink<'a> {
params: &'a HashSet<String>,
globals: &'a HashSet<String>,
nonlocals: &'a HashSet<String>,
locals: &'a HashSet<String>,
imports: &'a Imports,
module_bindings: &'a HashSet<String>,
is_init: bool,
span: &'a SpanIndex<'a>,
effects: Vec<(Effect, bool)>,
}
impl EffectSink for MutSink<'_> {
fn on_call(&mut self, call: &Call) {
let Expression::Attribute(attr) = call.func.as_ref() else {
return;
};
if !is_mutating_method(attr.attr.value) {
return;
}
let Some(root) = root_name_of_expr(&attr.value) else {
return;
};
let line = name_line_expr(&attr.value, self.span);
let receiver = render_expr(&attr.value).unwrap_or_else(|| root.clone());
let evidence = format!("{receiver}.{}(…)", attr.attr.value);
self.classify_and_push(root, line, evidence);
}
fn on_assert(&mut self, _assert: &Assert) {}
fn on_raise(&mut self, _raise: &Raise) {}
fn on_assign_target(&mut self, target: &AssignTargetExpression, is_aug: bool) {
match target {
AssignTargetExpression::Attribute(attr) => {
if let Expression::Name(n) = attr.value.as_ref()
&& n.value == "self"
{
let line = name_line(n, self.span);
if self.is_init {
self.push(
EffectKind::LocalMutation,
Tier::Heuristic,
line,
"self.x = … (constructor init, contained)".to_string(),
true,
);
} else {
self.push(
EffectKind::ThisMutation,
Tier::Heuristic,
line,
format!("self.{} = … (instance state)", attr.attr.value),
false,
);
}
return;
}
if let Some(root) = root_name_of_expr(&attr.value) {
let line = name_line_expr(&attr.value, self.span);
let evidence = format!("{root}.{} = …", attr.attr.value);
self.classify_and_push(root, line, evidence);
}
}
AssignTargetExpression::Name(n) if is_aug => {
let name = n.value.to_owned();
let line = name_line(n, self.span);
let evidence = format!("{name} += …");
self.classify_and_push(name, line, evidence);
}
AssignTargetExpression::Name(n)
if self.globals.contains(n.value) || self.nonlocals.contains(n.value) =>
{
let name = n.value.to_owned();
let line = name_line(n, self.span);
let evidence = format!("{name} = …");
self.classify_and_push(name, line, evidence);
}
AssignTargetExpression::Name(_) => {}
AssignTargetExpression::Subscript(sub) => {
if let Some(root) = root_name_of_expr(&sub.value) {
let line = name_line_expr(&sub.value, self.span);
let evidence = format!("{root}[…] = …");
self.classify_and_push(root, line, evidence);
}
}
_ => {}
}
}
}
impl MutSink<'_> {
fn classify_and_push(&mut self, root: String, line: usize, evidence: String) {
if root == "self" {
self.push(
EffectKind::ThisMutation,
Tier::Heuristic,
line,
evidence,
false,
);
return;
}
if self.globals.contains(&root) {
self.push(
EffectKind::GlobalMutation,
Tier::Exact,
line,
format!("global {root} ({evidence})"),
false,
);
return;
}
if self.nonlocals.contains(&root) {
self.push(
EffectKind::ThisMutation,
Tier::Exact,
line,
format!("nonlocal {root} ({evidence})"),
false,
);
return;
}
if self.params.contains(&root) {
self.push(
EffectKind::ParamMutation,
Tier::Heuristic,
line,
evidence,
false,
);
return;
}
if self.locals.contains(&root) {
self.push(EffectKind::LocalMutation, Tier::Exact, line, evidence, true);
return;
}
if self.imports.resolve(&root).is_some() {
self.push(
EffectKind::GlobalMutation,
Tier::Heuristic,
line,
format!("{evidence} (imported `{root}`)"),
false,
);
return;
}
if self.module_bindings.contains(&root) {
self.push(
EffectKind::GlobalMutation,
Tier::Heuristic,
line,
format!("{evidence} (module-level `{root}`)"),
false,
);
return;
}
self.push_hidden(line, evidence, "captured-binding");
}
fn push_hidden(&mut self, line: usize, evidence: String, subreason: &str) {
let kind = EffectKind::HiddenMutation;
let tier = Tier::Heuristic;
let class = kind.base_class();
self.effects.push((
Effect {
kind,
class,
discounted_to: None,
weight: weight_for_class(class),
line,
tier,
hidden: true,
evidence,
discount: None,
subreason: Some(subreason.to_owned()),
confidence: detection_confidence(tier, false, false),
},
false,
));
}
fn push(
&mut self,
kind: EffectKind,
tier: Tier,
line: usize,
evidence: String,
contained: bool,
) {
let class = kind.base_class();
self.effects.push((
Effect {
kind,
class,
discounted_to: None,
weight: weight_for_class(class),
line,
tier,
hidden: false,
evidence,
discount: None,
subreason: None,
confidence: detection_confidence(tier, false, false),
},
contained,
));
}
}
fn root_name_of_expr(expr: &Expression) -> Option<String> {
match expr {
Expression::Name(n) => Some(n.value.to_owned()),
Expression::Attribute(a) => root_name_of_expr(&a.value),
Expression::Subscript(s) => root_name_of_expr(&s.value),
Expression::Call(c) => root_name_of_expr(&c.func),
_ => None,
}
}
fn is_mutating_method(name: &str) -> bool {
matches!(
name,
"append"
| "extend"
| "insert"
| "remove"
| "pop"
| "clear"
| "sort"
| "reverse"
| "update"
| "add"
| "discard"
| "setdefault"
)
}
fn name_line_expr(expr: &Expression, span: &SpanIndex) -> usize {
leftmost_name(expr).map(|n| name_line(n, span)).unwrap_or(0)
}
fn leftmost_name<'a>(expr: &'a Expression<'a>) -> Option<&'a Name<'a>> {
match expr {
Expression::Name(n) => Some(n),
Expression::Attribute(a) => leftmost_name(&a.value),
Expression::Subscript(s) => leftmost_name(&s.value),
Expression::Call(c) => leftmost_name(&c.func),
_ => None,
}
}
fn name_line(name: &Name, span: &SpanIndex) -> usize {
span.line_col(anchor_of_subslice(span.src(), name.value)).0
}
#[cfg(test)]
mod tests {
use super::*;
use crate::functions;
use fxrank_core::effect::EffectKind::{self, *};
use std::collections::HashMap;
fn mutation_effects(name: &str) -> HashMap<String, Vec<(EffectKind, bool)>> {
let src = std::fs::read_to_string(format!("tests/fixtures/{name}.py")).unwrap();
let module = libcst_native::parse_module(&src, None).unwrap();
let imports = crate::imports::Imports::build(&module);
let module_bindings = crate::imports::module_bindings(&module);
let span = crate::source::SpanIndex::new(&src);
let anchors = crate::source::lambda_anchors(&src).expect("tokenize must succeed");
let (units, _) = functions::collect(&module, &src, &span, &anchors);
let mut out: HashMap<String, Vec<(EffectKind, bool)>> = HashMap::new();
for unit in &units {
let pairs = detect(unit, &imports, &module_bindings, &span);
out.insert(
unit.symbol.clone(),
pairs.iter().map(|(e, c)| (e.kind, *c)).collect(),
);
}
out
}
fn mutation_evidence(name: &str) -> HashMap<String, Vec<(EffectKind, bool, String)>> {
let src = std::fs::read_to_string(format!("tests/fixtures/{name}.py")).unwrap();
let module = libcst_native::parse_module(&src, None).unwrap();
let imports = crate::imports::Imports::build(&module);
let module_bindings = crate::imports::module_bindings(&module);
let span = crate::source::SpanIndex::new(&src);
let anchors = crate::source::lambda_anchors(&src).expect("tokenize must succeed");
let (units, _) = functions::collect(&module, &src, &span, &anchors);
let mut out: HashMap<String, Vec<(EffectKind, bool, String)>> = HashMap::new();
for unit in &units {
let pairs = detect(unit, &imports, &module_bindings, &span);
out.insert(
unit.symbol.clone(),
pairs
.iter()
.map(|(e, c)| (e.kind, *c, e.evidence.clone()))
.collect(),
);
}
out
}
#[test]
fn classifies_mutation_by_escape() {
let m = mutation_effects("mutation");
assert!(
m["uses_global"].contains(&(GlobalMutation, false)),
"uses_global should have GlobalMutation(contained=false), got: {:?}",
m["uses_global"]
);
assert!(
m["bump"].contains(&(ThisMutation, false)),
"bump should have ThisMutation(contained=false), got: {:?}",
m["bump"]
);
assert!(
m["mutates_param"].contains(&(ParamMutation, false)),
"mutates_param should have ParamMutation(contained=false), got: {:?}",
m["mutates_param"]
);
assert!(
m["builds_local"].contains(&(LocalMutation, true)),
"builds_local should have LocalMutation(contained=true), got: {:?}",
m["builds_local"]
);
assert!(
m["__init__"].contains(&(LocalMutation, true)),
"__init__ should have LocalMutation(contained=true), got: {:?}",
m["__init__"]
);
}
#[test]
fn plain_assign_to_global_nonlocal_names_escapes() {
let m = mutation_effects("mutation");
assert!(
m["plain_global_rebind"].contains(&(GlobalMutation, false)),
"plain `=` to a global name must emit GlobalMutation(false), got: {:?}",
m["plain_global_rebind"]
);
assert!(
m["plain_nonlocal_rebind"].contains(&(ThisMutation, false)),
"plain `=` to a nonlocal name must emit ThisMutation(false), got: {:?}",
m["plain_nonlocal_rebind"]
);
assert!(
m["plain_local_binding"].is_empty(),
"plain `=` to a true local must emit NO mutation, got: {:?}",
m["plain_local_binding"]
);
}
#[test]
fn self_method_and_subscript_mutations_escape_even_in_init() {
let m = mutation_effects("mutation");
assert!(
m["__init__"].contains(&(LocalMutation, true)),
"direct `self.attr = …` in __init__ stays LocalMutation(true), got: {:?}",
m["__init__"]
);
assert!(
m["__init__"].contains(&(ThisMutation, false)),
"`self.items.append(…)` in __init__ must be ThisMutation(false), got: {:?}",
m["__init__"]
);
assert!(
m["store"].contains(&(ThisMutation, false)),
"`self[i] = v` must be ThisMutation(false), got: {:?}",
m["store"]
);
}
#[test]
fn push_hidden_emits_hidden_mutation_with_subreason() {
let params = std::collections::HashSet::new();
let globals = std::collections::HashSet::new();
let nonlocals = std::collections::HashSet::new();
let locals = std::collections::HashSet::new();
let src = "x\n";
let module = libcst_native::parse_module(src, None).unwrap();
let imports = crate::imports::Imports::build(&module);
let span = crate::source::SpanIndex::new(src);
let mut sink = MutSink {
params: ¶ms,
globals: &globals,
nonlocals: &nonlocals,
locals: &locals,
imports: &imports,
module_bindings: &HashSet::new(),
is_init: false,
span: &span,
effects: Vec::new(),
};
sink.push_hidden(1, "outer_acc.append(…)".to_string(), "captured-binding");
assert_eq!(sink.effects.len(), 1);
let (effect, contained) = &sink.effects[0];
assert_eq!(effect.kind, EffectKind::HiddenMutation);
assert_eq!(effect.class, 3);
assert!(effect.hidden, "push_hidden must set hidden:true");
assert_eq!(effect.subreason.as_deref(), Some("captured-binding"));
assert!(!contained, "hidden writes escape — contained=false");
}
#[test]
fn detect_accepts_imports_param() {
let src = "def f(lst):\n lst.append(1)\n";
let module = libcst_native::parse_module(src, None).unwrap();
let imports = crate::imports::Imports::build(&module);
let module_bindings = crate::imports::module_bindings(&module);
let span = crate::source::SpanIndex::new(src);
let anchors = crate::source::lambda_anchors(src).expect("tokenize must succeed");
let (units, _) = functions::collect(&module, src, &span, &anchors);
let f = units.iter().find(|u| u.symbol == "f").unwrap();
let pairs = detect(f, &imports, &module_bindings, &span);
assert!(
pairs.iter().any(|(e, _)| e.kind == ParamMutation),
"lst.append where lst is a param → ParamMutation, got: {:?}",
pairs.iter().map(|(e, _)| e.kind).collect::<Vec<_>>()
);
}
#[test]
fn import_rooted_write_is_global_mutation() {
let m = mutation_effects("mutation");
assert!(
m["mutates_imported_module"].contains(&(GlobalMutation, false)),
"config.settings.append(…) where `config` is imported must be GlobalMutation(false), got: {:?}",
m["mutates_imported_module"]
);
}
#[test]
fn captured_binding_subreason_is_set() {
let src = std::fs::read_to_string("tests/fixtures/mutation.py").unwrap();
let module = libcst_native::parse_module(&src, None).unwrap();
let imports = crate::imports::Imports::build(&module);
let module_bindings = crate::imports::module_bindings(&module);
let span = crate::source::SpanIndex::new(&src);
let anchors = crate::source::lambda_anchors(&src).expect("tokenize must succeed");
let (units, _) = functions::collect(&module, &src, &span, &anchors);
let inner = units.iter().find(|u| u.symbol == "inner").unwrap();
let pairs = detect(inner, &imports, &module_bindings, &span);
let hidden = pairs
.iter()
.find(|(e, _)| e.kind == HiddenMutation)
.map(|(e, _)| e)
.expect("inner must emit a HiddenMutation");
assert_eq!(hidden.class, 3);
assert!(
hidden.hidden,
"captured-binding HiddenMutation must be hidden:true"
);
assert_eq!(hidden.subreason.as_deref(), Some("captured-binding"));
assert!(
pairs.iter().any(|(e, c)| e.kind == HiddenMutation && !*c),
"captured-binding write escapes — contained=false"
);
}
#[test]
fn mutating_method_evidence_uses_full_receiver() {
let m = mutation_evidence("mutation");
let init = &m["__init__"];
let append = init
.iter()
.find(|(k, _, _)| *k == ThisMutation)
.unwrap_or_else(|| panic!("expected a ThisMutation in __init__, got: {init:?}"));
assert!(
append.2.contains("self.items"),
"evidence must name the full receiver `self.items`, got: {:?}",
append.2
);
}
fn detect_src(src: &str, fn_name: &str) -> Vec<(Effect, bool)> {
let module = libcst_native::parse_module(src, None).unwrap();
let imports = crate::imports::Imports::build(&module);
let module_bindings = crate::imports::module_bindings(&module);
let span = crate::source::SpanIndex::new(src);
let anchors = crate::source::lambda_anchors(src).expect("tokenize must succeed");
let (units, _) = functions::collect(&module, src, &span, &anchors);
let unit = units
.iter()
.find(|u| u.symbol == fn_name)
.expect("unit not found");
detect(unit, &imports, &module_bindings, &span)
}
#[test]
fn module_level_content_mutation_is_global() {
let src = "_cache = {}\ndef f():\n _cache['k'] = 1\n";
let pairs = detect_src(src, "f");
assert!(
pairs.iter().any(|(e, c)| e.kind == GlobalMutation && !*c),
"module-level `_cache['k']=1` (no `global`) must be GlobalMutation(false), got: {:?}",
pairs.iter().map(|(e, _)| e.kind).collect::<Vec<_>>()
);
assert!(
!pairs.iter().any(|(e, _)| e.kind == HiddenMutation),
"module-level content mutation must not be hidden.mutation"
);
}
#[test]
fn local_shadowing_module_binding_is_local() {
let src = "_cache = {}\ndef f():\n _cache = {}\n _cache['k'] = 1\n";
let pairs = detect_src(src, "f");
assert!(
pairs.iter().any(|(e, c)| e.kind == LocalMutation && *c),
"shadowing local `_cache` must be LocalMutation(true), got: {:?}",
pairs.iter().map(|(e, _)| e.kind).collect::<Vec<_>>()
);
assert!(
!pairs.iter().any(|(e, _)| e.kind == GlobalMutation),
"shadowing local must not escalate to GlobalMutation"
);
}
#[test]
fn for_target_shadow_stays_local() {
let src = "_cache = {}\ndef f():\n for _cache in []:\n _cache['k'] = 1\n";
let pairs = detect_src(src, "f");
let writes: Vec<_> = pairs
.iter()
.filter(|(e, _)| {
matches!(
e.kind,
LocalMutation | GlobalMutation | HiddenMutation | ThisMutation
)
})
.collect();
assert!(
writes.iter().any(|(e, c)| e.kind == LocalMutation && *c),
"expected LocalMutation(contained=true) for for-target shadow, got: {:?}",
writes.iter().map(|(e, c)| (e.kind, *c)).collect::<Vec<_>>()
);
assert!(
!writes.iter().any(|(e, _)| e.kind == GlobalMutation),
"expected NO GlobalMutation for for-target shadow, got: {:?}",
writes.iter().map(|(e, c)| (e.kind, *c)).collect::<Vec<_>>()
);
}
#[test]
fn local_destructured_shadow_stays_local() {
let src = "_cache = {}\ndef f():\n (_cache,) = ({},)\n _cache['k'] = 1\n";
let pairs = detect_src(src, "f");
assert!(
pairs.iter().any(|(e, c)| e.kind == LocalMutation && *c),
"destructuring-local `_cache` must be LocalMutation(true), got: {:?}",
pairs.iter().map(|(e, _)| e.kind).collect::<Vec<_>>()
);
assert!(
!pairs.iter().any(|(e, _)| e.kind == GlobalMutation),
"destructured local must not escalate to GlobalMutation"
);
}
}