Skip to content

Commit 25de9b2

Browse files
scottgerringscottgerring
committed
feat: add 'no alloc' to stop functions from allocating on the heap
1 parent 51e9bdd commit 25de9b2

File tree

12 files changed

+549
-3
lines changed

12 files changed

+549
-3
lines changed

‎cargo_pup_lint_config/src/function_lint/builder.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,12 @@ impl<'a> FunctionConstraintBuilder<'a> {
129129
self
130130
}
131131

132+
/// Require that the function does not perform heap allocations
133+
pub fn no_allocation(mut self) -> Self {
134+
self.add_rule_internal(FunctionRule::NoAllocation(self.current_severity));
135+
self
136+
}
137+
132138
/// Create a new MaxLength rule with the current severity
133139
pub fn create_max_length_rule(&self, length: usize) -> FunctionRule {
134140
FunctionRule::MaxLength(length, self.current_severity)

‎cargo_pup_lint_config/src/function_lint/types.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ pub enum FunctionRule {
6262
ResultErrorMustImplementError(Severity),
6363
/// Enforces that a function matching the selector must not exist at all
6464
MustNotExist(Severity),
65+
/// Enforces that a function must not perform heap allocations
66+
NoAllocation(Severity),
6567
}
6668

6769
// Helper methods for FunctionRule

‎cargo_pup_lint_impl/src/lints/function_lint/lint.rs‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@ use rustc_lint::{LateContext, LateLintPass, LintStore};
1111
use rustc_middle::ty::TyKind;
1212
use rustc_session::impl_lint_pass;
1313
use rustc_span::BytePos;
14+
use std::collections::HashMap;
15+
use std::sync::Mutex;
16+
17+
use super::no_allocation::detect_allocation_in_mir;
1418

1519
// Helper: retrieve the concrete Self type of the impl the method belongs to, if any
1620
fn get_self_type<'tcx>(
@@ -26,6 +30,8 @@ pub struct FunctionLint {
2630
name: String,
2731
matches: FunctionMatch,
2832
function_rules: Vec<FunctionRule>,
33+
// Cache for allocation detection to avoid re-analyzing the same functions
34+
allocation_cache: Mutex<HashMap<rustc_hir::def_id::DefId, bool>>,
2935
}
3036

3137
impl FunctionLint {
@@ -36,6 +42,7 @@ impl FunctionLint {
3642
name: f.name.clone(),
3743
matches: f.matches.clone(),
3844
function_rules: f.rules.clone(),
45+
allocation_cache: Mutex::new(HashMap::new()),
3946
})
4047
} else {
4148
panic!("Expected a Function lint configuration")
@@ -245,6 +252,7 @@ impl ArchitectureLintRule for FunctionLint {
245252
name: name.clone(),
246253
matches: matches.clone(),
247254
function_rules: function_rules.clone(),
255+
allocation_cache: Mutex::new(HashMap::new()),
248256
})
249257
});
250258
}
@@ -351,6 +359,28 @@ impl<'tcx> LateLintPass<'tcx> for FunctionLint {
351359
"Remove this function to satisfy the architectural rule",
352360
);
353361
}
362+
FunctionRule::NoAllocation(severity) => {
363+
if ctx.tcx.is_mir_available(fn_def_id) {
364+
let mir = ctx.tcx.optimized_mir(fn_def_id);
365+
366+
if let Some(violation) = detect_allocation_in_mir(
367+
ctx.tcx,
368+
mir,
369+
fn_def_id,
370+
&mut self.allocation_cache.lock().unwrap(),
371+
) {
372+
span_lint_and_help(
373+
ctx,
374+
FUNCTION_LINT::get_by_severity(*severity),
375+
self.name().as_str(),
376+
violation.span,
377+
format!("Function allocates heap memory: {}", violation.reason),
378+
None,
379+
"Remove heap allocations to satisfy the NoAllocation rule",
380+
);
381+
}
382+
}
383+
}
354384
}
355385
}
356386
}
@@ -454,6 +484,28 @@ impl<'tcx> LateLintPass<'tcx> for FunctionLint {
454484
"Remove this function to satisfy the architectural rule",
455485
);
456486
}
487+
FunctionRule::NoAllocation(severity) => {
488+
if ctx.tcx.is_mir_available(fn_def_id) {
489+
let mir = ctx.tcx.optimized_mir(fn_def_id);
490+
491+
if let Some(violation) = detect_allocation_in_mir(
492+
ctx.tcx,
493+
mir,
494+
fn_def_id,
495+
&mut self.allocation_cache.lock().unwrap(),
496+
) {
497+
span_lint_and_help(
498+
ctx,
499+
FUNCTION_LINT::get_by_severity(*severity),
500+
self.name().as_str(),
501+
violation.span,
502+
format!("Function allocates heap memory: {}", violation.reason),
503+
None,
504+
"Remove heap allocations to satisfy the NoAllocation rule",
505+
);
506+
}
507+
}
508+
}
457509
}
458510
}
459511
}

‎cargo_pup_lint_impl/src/lints/function_lint/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2024 Datadog, Inc.
22

33
mod lint;
4+
mod no_allocation;
45

56
pub use lint::FunctionLint;

‎cargo_pup_lint_impl/src/lints/function_lint/no_allocation.rs‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2024 Datadog, Inc.
2+
3+
use rustc_hir::def_id::DefId;
4+
use rustc_middle::mir::{Body, TerminatorKind};
5+
use rustc_middle::ty::TyCtxt;
6+
use rustc_span::Span;
7+
use std::collections::HashMap;
8+
9+
/// Represents a violation of the no-allocation rule
10+
#[derive(Debug)]
11+
pub struct AllocationViolation {
12+
pub span: Span,
13+
pub reason: String,
14+
}
15+
16+
/// Detects heap allocations in a function's MIR
17+
pub fn detect_allocation_in_mir<'tcx>(
18+
tcx: TyCtxt<'tcx>,
19+
mir: &Body<'tcx>,
20+
_fn_def_id: DefId,
21+
cache: &mut HashMap<DefId, bool>,
22+
) -> Option<AllocationViolation> {
23+
// Iterate through basic blocks
24+
for (_bb, bb_data) in mir.basic_blocks.iter_enumerated() {
25+
// Check terminator for calls
26+
if let Some(terminator) = &bb_data.terminator {
27+
if let TerminatorKind::Call { func, .. } = &terminator.kind {
28+
// Extract function DefId using const_fn_def
29+
if let Some((callee_def_id, _generics)) = func.const_fn_def() {
30+
let path = tcx.def_path_str(callee_def_id);
31+
32+
// Check if it's a known allocating function
33+
if is_allocating_function(&path) {
34+
return Some(AllocationViolation {
35+
span: terminator.source_info.span,
36+
reason: format!("calls allocating function: {}", path),
37+
});
38+
}
39+
40+
// Check transitively (with cycle detection)
41+
if should_analyze_transitively(tcx, callee_def_id) {
42+
if function_allocates(tcx, callee_def_id, cache) {
43+
return Some(AllocationViolation {
44+
span: terminator.source_info.span,
45+
reason: format!("calls function that allocates: {}", path),
46+
});
47+
}
48+
}
49+
}
50+
}
51+
}
52+
}
53+
54+
None
55+
}
56+
57+
/// Checks if a function path corresponds to a known allocating function
58+
fn is_allocating_function(path: &str) -> bool {
59+
// Direct allocation functions - these are the low-level allocators
60+
if path.contains("alloc::alloc::") {
61+
if path.contains("::alloc")
62+
|| path.contains("::allocate")
63+
|| path.contains("::exchange_malloc")
64+
|| path.contains("::box_free")
65+
{
66+
return true;
67+
}
68+
}
69+
70+
// Box allocations - check for various Box patterns
71+
if path.contains("::Box::") || path.contains("::Box::<") {
72+
if path.contains("::new") {
73+
return true;
74+
}
75+
}
76+
77+
// Vec allocations and operations that may allocate
78+
if path.contains("::Vec::") || path.contains("::Vec::<") {
79+
if path.contains("::new")
80+
|| path.contains("::with_capacity")
81+
|| path.contains("::push")
82+
|| path.contains("::insert")
83+
|| path.contains("::extend")
84+
|| path.contains("::append")
85+
|| path.contains("::resize")
86+
|| path.contains("::from_elem")
87+
{
88+
return true;
89+
}
90+
}
91+
92+
// String allocations
93+
if path.contains("::String::") {
94+
if path.contains("::new")
95+
|| path.contains("::from")
96+
|| path.contains("::from_utf8")
97+
|| path.contains("::from_utf16")
98+
|| path.contains("::push_str")
99+
|| path.contains("::push")
100+
|| path.contains("::insert")
101+
|| path.contains("::insert_str")
102+
{
103+
return true;
104+
}
105+
}
106+
107+
// Format macro and related
108+
if path.contains("::format") || path.contains("fmt::format") {
109+
return true;
110+
}
111+
112+
// Rc and Arc
113+
if path.contains("::Rc::")
114+
|| path.contains("::Rc::<")
115+
|| path.contains("::Arc::")
116+
|| path.contains("::Arc::<")
117+
{
118+
if path.contains("::new") || path.contains("::clone") {
119+
return true;
120+
}
121+
}
122+
123+
// Collection types - broader matching
124+
if path.contains("HashMap")
125+
|| path.contains("BTreeMap")
126+
|| path.contains("HashSet")
127+
|| path.contains("BTreeSet")
128+
|| path.contains("VecDeque")
129+
|| path.contains("LinkedList")
130+
|| path.contains("BinaryHeap")
131+
{
132+
if path.contains(">::new")
133+
|| path.contains(">::with_capacity")
134+
|| path.contains(">::insert")
135+
|| path.contains(">::push")
136+
{
137+
return true;
138+
}
139+
}
140+
141+
// to_string, to_owned methods - these allocate
142+
if path.contains("::to_string") || path.contains("::to_owned") {
143+
return true;
144+
}
145+
146+
// RawVec - internal vec allocator
147+
if path.contains("RawVec") && (path.contains("::new") || path.contains("::allocate")) {
148+
return true;
149+
}
150+
151+
false
152+
}
153+
154+
/// Determines if we should recursively analyze a function
155+
fn should_analyze_transitively(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
156+
// Only analyze functions in the local crate
157+
// External crates are harder to analyze and may not have MIR available
158+
def_id.krate == rustc_hir::def_id::LOCAL_CRATE && tcx.is_mir_available(def_id)
159+
}
160+
161+
/// Recursively checks if a function allocates, with memoization
162+
fn function_allocates<'tcx>(
163+
tcx: TyCtxt<'tcx>,
164+
def_id: DefId,
165+
cache: &mut HashMap<DefId, bool>,
166+
) -> bool {
167+
// Check cache
168+
if let Some(&result) = cache.get(&def_id) {
169+
return result;
170+
}
171+
172+
// Mark as false initially (cycle detection)
173+
cache.insert(def_id, false);
174+
175+
// Try to get MIR
176+
if !tcx.is_mir_available(def_id) {
177+
// Conservative: assume external functions don't allocate
178+
// This prevents false positives for standard library functions
179+
return false;
180+
}
181+
182+
let mir = tcx.optimized_mir(def_id);
183+
let allocates = detect_allocation_in_mir(tcx, mir, def_id, cache).is_some();
184+
185+
// Update cache with actual result
186+
cache.insert(def_id, allocates);
187+
allocates
188+
}

‎test_app/expected_output‎

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -295,5 +295,68 @@ error: Function 'process_item' is forbidden by lint rule
295295
= help: Remove this function to satisfy the architectural rule
296296
= note: Applied by cargo-pup rule 'async_functions_forbidden'.
297297

298-
warning: `test_app` (bin "test_app") generated 20 warnings
299-
error: could not compile `test_app` (bin "test_app") due to 6 previous errors; 20 warnings emitted
298+
warning: Function allocates heap memory: calls allocating function: std::boxed::Box::<T>::new
299+
--> src/no_allocation.rs:12:5
300+
|
301+
12 | Box::new(42)
302+
| ^^^^^^^^^^^^
303+
|
304+
= help: Remove heap allocations to satisfy the NoAllocation rule
305+
= note: Applied by cargo-pup rule 'no_allocation_check'.
306+
307+
warning: Function allocates heap memory: calls allocating function: std::vec::Vec::<T>::new
308+
--> src/no_allocation.rs:17:5
309+
|
310+
17 | Vec::new()
311+
| ^^^^^^^^^^
312+
|
313+
= help: Remove heap allocations to satisfy the NoAllocation rule
314+
= note: Applied by cargo-pup rule 'no_allocation_check'.
315+
316+
warning: Function allocates heap memory: calls allocating function: std::string::ToString::to_string
317+
--> src/no_allocation.rs:22:5
318+
|
319+
22 | x.to_string()
320+
| ^^^^^^^^^^^^^
321+
|
322+
= help: Remove heap allocations to satisfy the NoAllocation rule
323+
= note: Applied by cargo-pup rule 'no_allocation_check'.
324+
325+
warning: Function allocates heap memory: calls function that allocates: no_allocation::helper_that_allocates
326+
--> src/no_allocation.rs:27:5
327+
|
328+
27 | helper_that_allocates(x)
329+
| ^^^^^^^^^^^^^^^^^^^^^^^^
330+
|
331+
= help: Remove heap allocations to satisfy the NoAllocation rule
332+
= note: Applied by cargo-pup rule 'no_allocation_check'.
333+
334+
warning: Function allocates heap memory: calls allocating function: std::vec::Vec::<T>::new
335+
--> src/no_allocation.rs:32:5
336+
|
337+
32 | Vec::new()
338+
| ^^^^^^^^^^
339+
|
340+
= help: Remove heap allocations to satisfy the NoAllocation rule
341+
= note: Applied by cargo-pup rule 'no_allocation_check'.
342+
343+
warning: Function allocates heap memory: calls function that allocates: no_allocation::deep_helper
344+
--> src/no_allocation.rs:36:5
345+
|
346+
36 | deep_helper()
347+
| ^^^^^^^^^^^^^
348+
|
349+
= help: Remove heap allocations to satisfy the NoAllocation rule
350+
= note: Applied by cargo-pup rule 'no_allocation_check'.
351+
352+
warning: Function allocates heap memory: calls function that allocates: no_allocation::middle_helper
353+
--> src/no_allocation.rs:41:5
354+
|
355+
41 | middle_helper()
356+
| ^^^^^^^^^^^^^^^
357+
|
358+
= help: Remove heap allocations to satisfy the NoAllocation rule
359+
= note: Applied by cargo-pup rule 'no_allocation_check'.
360+
361+
warning: `test_app` (bin "test_app") generated 27 warnings
362+
error: could not compile `test_app` (bin "test_app") due to 6 previous errors; 27 warnings emitted

‎test_app/pup.ron‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,5 +129,12 @@
129129
MustNotExist(Error),
130130
],
131131
)),
132+
Function((
133+
name: "no_allocation_check",
134+
matches: InModule("^test_app::no_allocation$"),
135+
rules: [
136+
NoAllocation(Warn),
137+
],
138+
)),
132139
],
133140
)

0 commit comments

Comments
 (0)