From c64eb8b129f86d46a54ab841f4088b94d4007812 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Wed, 12 Nov 2025 12:59:54 +0100 Subject: [PATCH 1/2] feat: add 'no alloc' to stop functions from allocating on the heap --- .../src/function_lint/builder.rs | 6 + .../src/function_lint/types.rs | 2 + .../src/lints/function_lint/lint.rs | 52 +++++ .../src/lints/function_lint/mod.rs | 1 + .../src/lints/function_lint/no_allocation.rs | 180 ++++++++++++++++++ test_app/expected_output | 67 ++++++- test_app/pup.ron | 7 + test_app/src/main.rs | 1 + test_app/src/no_allocation.rs | 42 ++++ tests/ui/function_lint/no_allocation.rs | 74 +++++++ tests/ui/function_lint/no_allocation.stderr | 93 +++++++++ tests/ui/function_lint/pup.ron | 19 +- 12 files changed, 541 insertions(+), 3 deletions(-) create mode 100644 cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs create mode 100644 test_app/src/no_allocation.rs create mode 100644 tests/ui/function_lint/no_allocation.rs create mode 100644 tests/ui/function_lint/no_allocation.stderr diff --git a/cargo_pup_lint_config/src/function_lint/builder.rs b/cargo_pup_lint_config/src/function_lint/builder.rs index c8c10ee..c3f5d27 100644 --- a/cargo_pup_lint_config/src/function_lint/builder.rs +++ b/cargo_pup_lint_config/src/function_lint/builder.rs @@ -129,6 +129,12 @@ impl<'a> FunctionConstraintBuilder<'a> { self } + /// Require that the function does not perform heap allocations + pub fn no_allocation(mut self) -> Self { + self.add_rule_internal(FunctionRule::NoAllocation(self.current_severity)); + self + } + /// Create a new MaxLength rule with the current severity pub fn create_max_length_rule(&self, length: usize) -> FunctionRule { FunctionRule::MaxLength(length, self.current_severity) diff --git a/cargo_pup_lint_config/src/function_lint/types.rs b/cargo_pup_lint_config/src/function_lint/types.rs index dd867a7..de47032 100644 --- a/cargo_pup_lint_config/src/function_lint/types.rs +++ b/cargo_pup_lint_config/src/function_lint/types.rs @@ -62,6 +62,8 @@ pub enum FunctionRule { ResultErrorMustImplementError(Severity), /// Enforces that a function matching the selector must not exist at all MustNotExist(Severity), + /// Enforces that a function must not perform heap allocations + NoAllocation(Severity), } // Helper methods for FunctionRule diff --git a/cargo_pup_lint_impl/src/lints/function_lint/lint.rs b/cargo_pup_lint_impl/src/lints/function_lint/lint.rs index 70e90e9..cb33f84 100644 --- a/cargo_pup_lint_impl/src/lints/function_lint/lint.rs +++ b/cargo_pup_lint_impl/src/lints/function_lint/lint.rs @@ -11,6 +11,10 @@ use rustc_lint::{LateContext, LateLintPass, LintStore}; use rustc_middle::ty::TyKind; use rustc_session::impl_lint_pass; use rustc_span::BytePos; +use std::collections::HashMap; +use std::sync::Mutex; + +use super::no_allocation::detect_allocation_in_mir; // Helper: retrieve the concrete Self type of the impl the method belongs to, if any fn get_self_type<'tcx>( @@ -26,6 +30,8 @@ pub struct FunctionLint { name: String, matches: FunctionMatch, function_rules: Vec, + // Cache for allocation detection to avoid re-analyzing the same functions + allocation_cache: Mutex>, } impl FunctionLint { @@ -36,6 +42,7 @@ impl FunctionLint { name: f.name.clone(), matches: f.matches.clone(), function_rules: f.rules.clone(), + allocation_cache: Mutex::new(HashMap::new()), }) } else { panic!("Expected a Function lint configuration") @@ -245,6 +252,7 @@ impl ArchitectureLintRule for FunctionLint { name: name.clone(), matches: matches.clone(), function_rules: function_rules.clone(), + allocation_cache: Mutex::new(HashMap::new()), }) }); } @@ -351,6 +359,28 @@ impl<'tcx> LateLintPass<'tcx> for FunctionLint { "Remove this function to satisfy the architectural rule", ); } + FunctionRule::NoAllocation(severity) => { + if ctx.tcx.is_mir_available(fn_def_id) { + let mir = ctx.tcx.optimized_mir(fn_def_id); + + if let Some(violation) = detect_allocation_in_mir( + ctx.tcx, + mir, + fn_def_id, + &mut self.allocation_cache.lock().unwrap(), + ) { + span_lint_and_help( + ctx, + FUNCTION_LINT::get_by_severity(*severity), + self.name().as_str(), + violation.span, + format!("Function allocates heap memory: {}", violation.reason), + None, + "Remove heap allocations to satisfy the NoAllocation rule", + ); + } + } + } } } } @@ -454,6 +484,28 @@ impl<'tcx> LateLintPass<'tcx> for FunctionLint { "Remove this function to satisfy the architectural rule", ); } + FunctionRule::NoAllocation(severity) => { + if ctx.tcx.is_mir_available(fn_def_id) { + let mir = ctx.tcx.optimized_mir(fn_def_id); + + if let Some(violation) = detect_allocation_in_mir( + ctx.tcx, + mir, + fn_def_id, + &mut self.allocation_cache.lock().unwrap(), + ) { + span_lint_and_help( + ctx, + FUNCTION_LINT::get_by_severity(*severity), + self.name().as_str(), + violation.span, + format!("Function allocates heap memory: {}", violation.reason), + None, + "Remove heap allocations to satisfy the NoAllocation rule", + ); + } + } + } } } } diff --git a/cargo_pup_lint_impl/src/lints/function_lint/mod.rs b/cargo_pup_lint_impl/src/lints/function_lint/mod.rs index 0ce0633..645d944 100644 --- a/cargo_pup_lint_impl/src/lints/function_lint/mod.rs +++ b/cargo_pup_lint_impl/src/lints/function_lint/mod.rs @@ -1,5 +1,6 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2024 Datadog, Inc. mod lint; +mod no_allocation; pub use lint::FunctionLint; diff --git a/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs b/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs new file mode 100644 index 0000000..f0c8282 --- /dev/null +++ b/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs @@ -0,0 +1,180 @@ +// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2024 Datadog, Inc. + +use rustc_hir::def_id::DefId; +use rustc_middle::mir::{Body, TerminatorKind}; +use rustc_middle::ty::TyCtxt; +use rustc_span::Span; +use std::collections::HashMap; + +/// Represents a violation of the no-allocation rule +#[derive(Debug)] +pub struct AllocationViolation { + pub span: Span, + pub reason: String, +} + +/// Detects heap allocations in a function's MIR +pub fn detect_allocation_in_mir<'tcx>( + tcx: TyCtxt<'tcx>, + mir: &Body<'tcx>, + _fn_def_id: DefId, + cache: &mut HashMap, +) -> Option { + // Iterate through basic blocks + for (_bb, bb_data) in mir.basic_blocks.iter_enumerated() { + // Check terminator for calls + if let Some(terminator) = &bb_data.terminator + && let TerminatorKind::Call { func, .. } = &terminator.kind + { + // Extract function DefId using const_fn_def + if let Some((callee_def_id, _generics)) = func.const_fn_def() { + let path = tcx.def_path_str(callee_def_id); + + // Check if it's a known allocating function + if is_allocating_function(&path) { + return Some(AllocationViolation { + span: terminator.source_info.span, + reason: format!("calls allocating function: {path}"), + }); + } + + // Check transitively (with cycle detection) + if should_analyze_transitively(tcx, callee_def_id) + && function_allocates(tcx, callee_def_id, cache) + { + return Some(AllocationViolation { + span: terminator.source_info.span, + reason: format!("calls function that allocates: {path}"), + }); + } + } + } + } + + None +} + +/// Checks if a function path corresponds to a known allocating function +fn is_allocating_function(path: &str) -> bool { + // Direct allocation functions - these are the low-level allocators + if path.contains("alloc::alloc::") + && (path.contains("::alloc") + || path.contains("::allocate") + || path.contains("::exchange_malloc") + || path.contains("::box_free")) + { + return true; + } + + // Box allocations - check for various Box patterns + if (path.contains("::Box::") || path.contains("::Box::<")) && path.contains("::new") { + return true; + } + + // Vec allocations and operations that may allocate + if (path.contains("::Vec::") || path.contains("::Vec::<")) + && (path.contains("::new") + || path.contains("::with_capacity") + || path.contains("::push") + || path.contains("::insert") + || path.contains("::extend") + || path.contains("::append") + || path.contains("::resize") + || path.contains("::from_elem")) + { + return true; + } + + // String allocations + if path.contains("::String::") + && (path.contains("::new") + || path.contains("::from") + || path.contains("::from_utf8") + || path.contains("::from_utf16") + || path.contains("::push_str") + || path.contains("::push") + || path.contains("::insert") + || path.contains("::insert_str")) + { + return true; + } + + // Format macro and related + if path.contains("::format") || path.contains("fmt::format") { + return true; + } + + // Rc and Arc + if (path.contains("::Rc::") + || path.contains("::Rc::<") + || path.contains("::Arc::") + || path.contains("::Arc::<")) + && (path.contains("::new") || path.contains("::clone")) + { + return true; + } + + // Collection types - broader matching + if (path.contains("HashMap") + || path.contains("BTreeMap") + || path.contains("HashSet") + || path.contains("BTreeSet") + || path.contains("VecDeque") + || path.contains("LinkedList") + || path.contains("BinaryHeap")) + && (path.contains(">::new") + || path.contains(">::with_capacity") + || path.contains(">::insert") + || path.contains(">::push")) + { + return true; + } + + // to_string, to_owned methods - these allocate + if path.contains("::to_string") || path.contains("::to_owned") { + return true; + } + + // RawVec - internal vec allocator + if path.contains("RawVec") && (path.contains("::new") || path.contains("::allocate")) { + return true; + } + + false +} + +/// Determines if we should recursively analyze a function +fn should_analyze_transitively(tcx: TyCtxt<'_>, def_id: DefId) -> bool { + // Only analyze functions in the local crate + // External crates are harder to analyze and may not have MIR available + def_id.krate == rustc_hir::def_id::LOCAL_CRATE && tcx.is_mir_available(def_id) +} + +/// Recursively checks if a function allocates, with memoization +fn function_allocates<'tcx>( + tcx: TyCtxt<'tcx>, + def_id: DefId, + cache: &mut HashMap, +) -> bool { + // Check cache + if let Some(&result) = cache.get(&def_id) { + return result; + } + + // Mark as false initially (cycle detection) + cache.insert(def_id, false); + + // Try to get MIR + if !tcx.is_mir_available(def_id) { + // Conservative: assume external functions don't allocate + // This prevents false positives for standard library functions + return false; + } + + let mir = tcx.optimized_mir(def_id); + let allocates = detect_allocation_in_mir(tcx, mir, def_id, cache).is_some(); + + // Update cache with actual result + cache.insert(def_id, allocates); + allocates +} diff --git a/test_app/expected_output b/test_app/expected_output index b4f34bc..fe6daa2 100644 --- a/test_app/expected_output +++ b/test_app/expected_output @@ -295,5 +295,68 @@ error: Function 'process_item' is forbidden by lint rule = help: Remove this function to satisfy the architectural rule = note: Applied by cargo-pup rule 'async_functions_forbidden'. -warning: `test_app` (bin "test_app") generated 20 warnings -error: could not compile `test_app` (bin "test_app") due to 6 previous errors; 20 warnings emitted +warning: Function allocates heap memory: calls allocating function: std::boxed::Box::::new + --> src/no_allocation.rs:12:5 + | +12 | Box::new(42) + | ^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: Function allocates heap memory: calls allocating function: std::vec::Vec::::new + --> src/no_allocation.rs:17:5 + | +17 | Vec::new() + | ^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: Function allocates heap memory: calls allocating function: std::string::ToString::to_string + --> src/no_allocation.rs:22:5 + | +22 | x.to_string() + | ^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: Function allocates heap memory: calls function that allocates: no_allocation::helper_that_allocates + --> src/no_allocation.rs:27:5 + | +27 | helper_that_allocates(x) + | ^^^^^^^^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: Function allocates heap memory: calls allocating function: std::vec::Vec::::new + --> src/no_allocation.rs:32:5 + | +32 | Vec::new() + | ^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: Function allocates heap memory: calls function that allocates: no_allocation::deep_helper + --> src/no_allocation.rs:36:5 + | +36 | deep_helper() + | ^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: Function allocates heap memory: calls function that allocates: no_allocation::middle_helper + --> src/no_allocation.rs:41:5 + | +41 | middle_helper() + | ^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_check'. + +warning: `test_app` (bin "test_app") generated 27 warnings +error: could not compile `test_app` (bin "test_app") due to 6 previous errors; 27 warnings emitted diff --git a/test_app/pup.ron b/test_app/pup.ron index e03fc29..861a48e 100644 --- a/test_app/pup.ron +++ b/test_app/pup.ron @@ -129,5 +129,12 @@ MustNotExist(Error), ], )), + Function(( + name: "no_allocation_check", + matches: InModule("^test_app::no_allocation$"), + rules: [ + NoAllocation(Warn), + ], + )), ], ) \ No newline at end of file diff --git a/test_app/src/main.rs b/test_app/src/main.rs index 9f72930..af1ea5f 100644 --- a/test_app/src/main.rs +++ b/test_app/src/main.rs @@ -14,6 +14,7 @@ mod result_error; mod trait_impl; mod builder_style; mod async_functions; +mod no_allocation; fn main() { println!("Hello, world!"); diff --git a/test_app/src/no_allocation.rs b/test_app/src/no_allocation.rs new file mode 100644 index 0000000..c6745ab --- /dev/null +++ b/test_app/src/no_allocation.rs @@ -0,0 +1,42 @@ +// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2024 Datadog, Inc. + +// Module to test NoAllocation lint + +// This function should pass - no allocations +pub fn pure_math(x: i32, y: i32) -> i32 { + x + y +} + +// This function should fail - allocates a Box +pub fn allocates_box() -> Box { + Box::new(42) +} + +// This function should fail - allocates a Vec +pub fn allocates_vec() -> Vec { + Vec::new() +} + +// Helper function that allocates - this will be flagged +fn helper_that_allocates(x: i32) -> String { + x.to_string() +} + +// This function should fail - calls a function that allocates (transitive) +pub fn calls_allocating_helper(x: i32) -> String { + helper_that_allocates(x) +} + +// Chain of calls to demonstrate transitive detection +fn deep_helper() -> Vec { + Vec::new() +} + +fn middle_helper() -> Vec { + deep_helper() +} + +// This should fail - indirectly allocates through multiple levels +pub fn deeply_nested_allocation() -> Vec { + middle_helper() +} diff --git a/tests/ui/function_lint/no_allocation.rs b/tests/ui/function_lint/no_allocation.rs new file mode 100644 index 0000000..67bcdd1 --- /dev/null +++ b/tests/ui/function_lint/no_allocation.rs @@ -0,0 +1,74 @@ +//@compile-flags: --crate-name=test_crate + +// Test functions that should pass (no allocations) + +fn pure_computation(x: i32, y: i32) -> i32 { + x + y +} + +fn stack_only(x: i32) -> [i32; 4] { + [x, x + 1, x + 2, x + 3] +} + +fn uses_references(x: &str) -> usize { + x.len() +} + +// Test functions that should fail (allocate) + +fn allocates_box() -> Box { + Box::new(42) //~ ERROR: Function allocates heap memory +} + +fn allocates_vec_new() { + let _v = Vec::::new(); //~ ERROR: Function allocates heap memory +} + +fn allocates_string_new() { + let _s = String::new(); //~ ERROR: Function allocates heap memory +} + +fn allocates_to_string(x: i32) -> String { + x.to_string() //~ ERROR: Function allocates heap memory +} + +fn allocates_rc() -> std::rc::Rc { + std::rc::Rc::new(42) //~ ERROR: Function allocates heap memory +} + +fn allocates_arc() -> std::sync::Arc { + std::sync::Arc::new(42) //~ ERROR: Function allocates heap memory +} + +fn allocates_hashmap() { + let _map = std::collections::HashMap::::new(); //~ ERROR: Function allocates heap memory +} + +// Transitive allocation tests + +fn helper_allocates() -> Vec { + let v = Vec::new(); //~ ERROR: Function allocates heap memory + v +} + +fn calls_allocating_function() -> usize { + let v = helper_allocates(); //~ ERROR: Function allocates heap memory + v.len() +} + +// Method tests +struct MyStruct { + value: i32, +} + +impl MyStruct { + fn no_alloc_method(&self) -> i32 { + self.value * 2 + } + + fn alloc_method(&self) -> String { + self.value.to_string() //~ ERROR: Function allocates heap memory + } +} + +fn main() {} diff --git a/tests/ui/function_lint/no_allocation.stderr b/tests/ui/function_lint/no_allocation.stderr new file mode 100644 index 0000000..bacb2d1 --- /dev/null +++ b/tests/ui/function_lint/no_allocation.stderr @@ -0,0 +1,93 @@ +error: Function allocates heap memory: calls allocating function: std::boxed::Box::::new + --> tests/ui/function_lint/no_allocation.rs:20:5 + | +LL | Box::new(42) + | ^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + = note: `#[deny(function_lint)]` on by default + +error: Function allocates heap memory: calls allocating function: std::vec::Vec::::new + --> tests/ui/function_lint/no_allocation.rs:24:14 + | +LL | let _v = Vec::::new(); + | ^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::string::String::new + --> tests/ui/function_lint/no_allocation.rs:28:14 + | +LL | let _s = String::new(); + | ^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::string::ToString::to_string + --> tests/ui/function_lint/no_allocation.rs:32:5 + | +LL | x.to_string() + | ^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::rc::Rc::::new + --> tests/ui/function_lint/no_allocation.rs:36:5 + | +LL | std::rc::Rc::new(42) + | ^^^^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::sync::Arc::::new + --> tests/ui/function_lint/no_allocation.rs:40:5 + | +LL | std::sync::Arc::new(42) + | ^^^^^^^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::collections::HashMap::::new + --> tests/ui/function_lint/no_allocation.rs:44:16 + | +LL | let _map = std::collections::HashMap::::new(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::vec::Vec::::new + --> tests/ui/function_lint/no_allocation.rs:50:13 + | +LL | let v = Vec::new(); + | ^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls function that allocates: helper_allocates + --> tests/ui/function_lint/no_allocation.rs:55:13 + | +LL | let v = helper_allocates(); + | ^^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: Function allocates heap memory: calls allocating function: std::string::ToString::to_string + --> tests/ui/function_lint/no_allocation.rs:70:9 + | +LL | self.value.to_string() + | ^^^^^^^^^^^^^^^^^^^^^^ + | + = help: Remove heap allocations to satisfy the NoAllocation rule + = note: Applied by cargo-pup rule 'no_allocation_test'. + +error: aborting due to 10 previous errors + diff --git a/tests/ui/function_lint/pup.ron b/tests/ui/function_lint/pup.ron index 6c64d32..e9f4c45 100644 --- a/tests/ui/function_lint/pup.ron +++ b/tests/ui/function_lint/pup.ron @@ -339,7 +339,7 @@ // ====================================================================== // SECTION: IsAsync Tests (for is_async.rs) // ====================================================================== - + // Test that async functions are detected and can be linted Function( ( @@ -351,6 +351,23 @@ ) ] ) + ), + + // ====================================================================== + // SECTION: NoAllocation Tests (for no_allocation.rs) + // ====================================================================== + + // Test that functions with specific names do not allocate + Function( + ( + name: "no_allocation_test", + matches: InModule("test_crate"), + rules: [ + NoAllocation( + Error, + ) + ] + ) ) ] ) \ No newline at end of file From 15ebdbadb32d4cb1011fb4673d84074737886426 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Wed, 12 Nov 2025 17:29:07 +0100 Subject: [PATCH 2/2] feat: make 'no alloc' work on closure bodies too --- .../src/function_lint/builder.rs | 3 ++ .../src/lints/function_lint/no_allocation.rs | 46 ++++++++++++++++++- 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/cargo_pup_lint_config/src/function_lint/builder.rs b/cargo_pup_lint_config/src/function_lint/builder.rs index c3f5d27..5f75b9f 100644 --- a/cargo_pup_lint_config/src/function_lint/builder.rs +++ b/cargo_pup_lint_config/src/function_lint/builder.rs @@ -130,6 +130,9 @@ impl<'a> FunctionConstraintBuilder<'a> { } /// Require that the function does not perform heap allocations + /// This is a best-efforts lint! There are limits to the extent to which we can + /// scan _everything_ used in function bodies to find allocations, but in most + /// cases it should work well. pub fn no_allocation(mut self) -> Self { self.add_rule_internal(FunctionRule::NoAllocation(self.current_severity)); self diff --git a/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs b/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs index f0c8282..5d0554d 100644 --- a/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs +++ b/cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs @@ -24,12 +24,56 @@ pub fn detect_allocation_in_mir<'tcx>( for (_bb, bb_data) in mir.basic_blocks.iter_enumerated() { // Check terminator for calls if let Some(terminator) = &bb_data.terminator - && let TerminatorKind::Call { func, .. } = &terminator.kind + && let TerminatorKind::Call { func, args, .. } = &terminator.kind { // Extract function DefId using const_fn_def if let Some((callee_def_id, _generics)) = func.const_fn_def() { let path = tcx.def_path_str(callee_def_id); + // Check arguments for closures + for arg in args.iter() { + use rustc_middle::mir::Operand; + + // Try to extract closure DefId from the operand + let closure_def_id = match &arg.node { + Operand::Constant(constant) => { + // Check if this is a closure type + let ty = constant.const_.ty(); + if let rustc_middle::ty::TyKind::Closure(def_id, _) = ty.kind() { + Some(*def_id) + } else if let rustc_middle::ty::TyKind::FnDef(def_id, _) = ty.kind() { + Some(*def_id) + } else { + None + } + } + Operand::Move(place) | Operand::Copy(place) => { + // For Move/Copy operands, check the type of the place + let ty = place.ty(mir, tcx).ty; + if let rustc_middle::ty::TyKind::Closure(def_id, _) = ty.kind() { + Some(*def_id) + } else if let rustc_middle::ty::TyKind::FnDef(def_id, _) = ty.kind() { + Some(*def_id) + } else { + None + } + } + }; + + if let Some(closure_def_id) = closure_def_id { + // Analyze the closure if it's local + if closure_def_id.krate == rustc_hir::def_id::LOCAL_CRATE + && tcx.is_mir_available(closure_def_id) + && function_allocates(tcx, closure_def_id, cache) + { + return Some(AllocationViolation { + span: terminator.source_info.span, + reason: format!("passes allocating closure to {path}"), + }); + } + } + } + // Check if it's a known allocating function if is_allocating_function(&path) { return Some(AllocationViolation {