1use rustc_middle::mir::{
2 Body, Local, Operand, Rvalue, StatementKind,
3 TerminatorKind,
4};
5use rustc_middle::ty::TyCtxt;
6use rustc_span::DUMMY_SP;
7
8use crate::compat::FxHashMap;
9
10pub(crate) fn follow_parents(parents: &FxHashMap<Local, Local>, start: Local) -> Local {
11 let mut current = start;
12 let mut seen = std::collections::HashSet::new();
13 while seen.insert(current) {
14 let Some(next) = parents.get(¤t) else {
15 break;
16 };
17 current = *next;
18 }
19 current
20}
21
22pub(crate) fn resolve_through_casts<'tcx>(body: &Body<'tcx>, local: Local) -> Local {
23 let mut current = local;
24 let mut seen = std::collections::HashSet::new();
25 while seen.insert(current) {
26 let found = body.basic_blocks.iter().any(|data| {
27 data.statements.iter().any(|stmt| {
28 let StatementKind::Assign(assign) = &stmt.kind else {
29 return false;
30 };
31 let (target, rvalue) = assign.as_ref();
32 if target.local != current || !target.projection.is_empty() {
33 return false;
34 }
35 if let Rvalue::Cast(_, operand, _) = rvalue {
36 #[allow(unreachable_patterns)]
37 match operand {
38 Operand::Copy(p) | Operand::Move(p) if p.projection.is_empty() => {
39 current = p.local;
40 return true;
41 }
42 _ => {}
43 }
44 }
45 false
46 })
47 });
48 if !found {
49 break;
50 }
51 }
52 current
53}
54
55fn scalar_constant(operand: &Operand<'_>) -> Option<u128> {
56 let constant = match operand {
57 Operand::Constant(c) => c,
58 _ => return None,
59 };
60 constant.const_.try_to_scalar_int().map(|s| s.to_uint(s.size()))
61}
62
63pub(crate) fn collect_all_const_bytes_worklist<'tcx>(
64 tcx: TyCtxt<'tcx>,
65 body: &Body<'tcx>,
66 root: Local,
67) -> Vec<Vec<u8>> {
68 let mut results: Vec<Vec<u8>> = Vec::new();
69 let mut worklist: Vec<Local> = vec![root];
70 let mut visited: std::collections::HashSet<Local> = std::collections::HashSet::new();
71
72 while let Some(local) = worklist.pop() {
73 if !visited.insert(local) {
74 continue;
75 }
76
77 for data in body.basic_blocks.iter() {
78 for statement in &data.statements {
79 let StatementKind::Assign(assign) = &statement.kind else {
80 continue;
81 };
82 let (target, rvalue) = assign.as_ref();
83 if target.local != local || !target.projection.is_empty() {
84 continue;
85 }
86
87 if let Rvalue::Ref(_, _, place) = rvalue {
88 if let Some(bytes) = const_bytes_for_local(tcx, body, place.local) {
89 results.push(bytes);
90 }
91 continue;
92 }
93
94 if let Rvalue::Use(operand, ..) = rvalue {
95 #[allow(unreachable_patterns)]
96 match operand {
97 Operand::Copy(p) | Operand::Move(p) => {
98 worklist.push(p.local);
99 if let Some(bytes) = const_bytes_for_local(tcx, body, p.local) {
100 results.push(bytes);
101 }
102 continue;
103 }
104 Operand::Constant(_) => {}
105 _ => continue,
106 }
107 }
108
109 let constant = match rvalue {
110 Rvalue::Use(Operand::Constant(constant), ..)
111 | Rvalue::Cast(_, Operand::Constant(constant), _) => constant,
112 _ => continue,
113 };
114 let Ok(value) = constant.const_.eval(
115 tcx,
116 rustc_middle::ty::TypingEnv::fully_monomorphized(),
117 DUMMY_SP,
118 ) else {
119 continue;
120 };
121 if let Some(bytes) = crate::helpers::mir_utils::const_value_bytes(tcx, value, 0) {
122 results.push(bytes);
123 }
124 }
125 }
126
127 for data in body.basic_blocks.iter() {
128 if let Some(terminator) = &data.terminator {
129 if let TerminatorKind::Call { destination, func, args, .. } = &terminator.kind {
130 let dlocal = destination.local;
131 if dlocal != local {
132 continue;
133 }
134 if !destination.projection.is_empty() {
135 continue;
136 }
137 let name = crate::helpers::mir_utils::call_name(tcx, func);
138 if name.contains("as_ptr") || name.contains("::as_") {
139 for arg in args {
140 if let Some(bytes) = const_bytes_from_operand(tcx, body, &arg.node) {
141 results.push(bytes);
142 }
143 }
144 }
145 if name.contains("::add") {
146 if let Some(offset) = args.get(1).and_then(|a| scalar_constant(&a.node)) {
147 if let Some(base) = args.first() {
148 if let Some(bytes) = const_bytes_from_operand(tcx, body, &base.node) {
149 let start = offset as usize;
150 if start < bytes.len() {
151 results.push(bytes[start..].to_vec());
152 }
153 }
154 }
155 }
156 }
157 if name.contains("box_assume_init_into_vec_unsafe") {
158 if let Some(box_op) = args.first() {
159 if let Operand::Copy(p) | Operand::Move(p) = &box_op.node {
160 if p.projection.is_empty() {
161 worklist.push(p.local);
162 }
163 }
164 }
165 }
166 }
167 }
168 }
169 }
170
171 {
172 let mut agg_roots = std::collections::HashSet::new();
173 let mut seen = std::collections::HashSet::new();
174 let mut work = vec![root];
175 while let Some(local) = work.pop() {
176 if !seen.insert(local) {
177 continue;
178 }
179 for data in body.basic_blocks.iter() {
180 for statement in &data.statements {
181 let StatementKind::Assign(assign) = &statement.kind else {
182 continue;
183 };
184 let (target, rvalue) = assign.as_ref();
185 if target.local != local || !target.projection.is_empty() {
186 continue;
187 }
188 if let Rvalue::Use(Operand::Copy(p) | Operand::Move(p), ..) = rvalue {
189 work.push(p.local);
190 }
191 if let Rvalue::Cast(_, Operand::Copy(p) | Operand::Move(p), _) = rvalue {
192 if p.projection.is_empty() {
193 work.push(p.local);
194 }
195 }
196 }
197 }
198 for data in body.basic_blocks.iter() {
199 if let Some(terminator) = &data.terminator {
200 if let TerminatorKind::Call { destination, func, args, .. } = &terminator.kind {
201 if destination.local == local
202 && destination.projection.is_empty()
203 {
204 let name = crate::helpers::mir_utils::call_name(tcx, func);
205 if name.contains("box_assume_init_into_vec_unsafe") {
206 if let Some(box_op) = args.first() {
207 if let Operand::Copy(p) | Operand::Move(p) = &box_op.node {
208 if p.projection.is_empty() {
209 work.push(p.local);
210 }
211 }
212 }
213 }
214 }
215 }
216 }
217 }
218 agg_roots.insert(local);
219 }
220
221 for data in body.basic_blocks.iter() {
222 for statement in &data.statements {
223 let StatementKind::Assign(assign) = &statement.kind else {
224 continue;
225 };
226 let (_, rvalue) = assign.as_ref();
227 let Rvalue::Aggregate(_, operands) = rvalue else {
228 continue;
229 };
230 if operands.len() < 2 {
231 continue;
232 }
233 let last_op = operands.iter().last().unwrap();
234 if !is_constant_zero_u8(last_op) {
235 continue;
236 }
237 let mut all_nonzero = true;
238 for op in operands.iter().take(operands.len() - 1) {
239 if !aggregate_op_is_nonzero(tcx, body, op) {
240 all_nonzero = false;
241 break;
242 }
243 }
244 if all_nonzero {
245 let len = operands.len();
246 let mut bytes = Vec::with_capacity(len);
247 for _ in 0..len - 1 {
248 bytes.push(b'x');
249 }
250 bytes.push(0);
251 results.push(bytes);
252 }
253 }
254 }
255 }
256
257 for data in body.basic_blocks.iter() {
258 if let Some(terminator) = &data.terminator {
259 if let TerminatorKind::Call { func, args, .. } = &terminator.kind {
260 let name = crate::helpers::mir_utils::call_name(tcx, func);
261 if name.contains("as_ptr") || name.contains("::as_") {
262 for arg in args {
263 if let Some(bytes) = operand_const_bytes(tcx, &arg.node) {
264 results.push(bytes);
265 } else if let Operand::Copy(p) | Operand::Move(p) = &arg.node {
266 if p.projection.is_empty() {
267 if let Some(bytes) = const_bytes_for_local(tcx, body, p.local) {
268 results.push(bytes);
269 }
270 }
271 }
272 }
273 }
274 }
275 }
276 }
277
278 results
279}
280
281fn const_bytes_from_operand<'tcx>(
282 tcx: TyCtxt<'tcx>,
283 body: &Body<'tcx>,
284 operand: &Operand<'tcx>,
285) -> Option<Vec<u8>> {
286 if let Some(bytes) = operand_const_bytes(tcx, operand) {
287 return Some(bytes);
288 }
289 match operand {
290 Operand::Copy(p) | Operand::Move(p) if p.projection.is_empty() => {
291 if let Some(bytes) = const_bytes_for_local(tcx, body, p.local) {
292 return Some(bytes);
293 }
294 const_bytes_from_call_dest(tcx, body, p.local)
295 }
296 _ => None,
297 }
298}
299
300fn const_bytes_from_call_dest<'tcx>(
301 tcx: TyCtxt<'tcx>,
302 body: &Body<'tcx>,
303 local: Local,
304) -> Option<Vec<u8>> {
305 for data in body.basic_blocks.iter() {
306 if let Some(terminator) = &data.terminator {
307 if let TerminatorKind::Call { destination, func, args, .. } = &terminator.kind {
308 if destination.local != local || !destination.projection.is_empty() {
309 continue;
310 }
311 let name = crate::helpers::mir_utils::call_name(tcx, func);
312 if name.contains("as_ptr") || name.contains("::as_") {
313 for arg in args {
314 if let Some(bytes) = const_bytes_from_operand(tcx, body, &arg.node) {
315 return Some(bytes);
316 }
317 }
318 }
319 }
320 }
321 }
322 None
323}
324
325pub(crate) fn const_bytes_for_local<'tcx>(
326 tcx: TyCtxt<'tcx>,
327 body: &Body<'tcx>,
328 root: Local,
329) -> Option<Vec<u8>> {
330 for data in body.basic_blocks.iter() {
331 for statement in &data.statements {
332 let StatementKind::Assign(assign) = &statement.kind else {
333 continue;
334 };
335 let (target, rvalue) = assign.as_ref();
336 if target.local != root || !target.projection.is_empty() {
337 continue;
338 }
339 if let Rvalue::Ref(_, _, place) = rvalue {
340 let deref_local = place.local;
341 if let Some(bytes) = const_bytes_for_local(tcx, body, deref_local) {
342 return Some(bytes);
343 }
344 continue;
345 }
346 if let Rvalue::Use(operand, ..) = rvalue {
347 #[allow(unreachable_patterns)]
348 match operand {
349 Operand::Copy(p) | Operand::Move(p) => {
350 if let Some(bytes) = const_bytes_for_local(tcx, body, p.local) {
351 return Some(bytes);
352 }
353 if let Some(bytes) = const_bytes_from_call_dest(tcx, body, p.local) {
354 return Some(bytes);
355 }
356 continue;
357 }
358 Operand::Constant(_) => {}
359 _ => continue,
360 }
361 }
362 if let Rvalue::Cast(_, operand, _) = rvalue {
363 if let Operand::Copy(p) | Operand::Move(p) = operand {
364 if p.projection.is_empty() {
365 if let Some(bytes) = const_bytes_for_local(tcx, body, p.local) {
366 return Some(bytes);
367 }
368 }
369 }
370 continue;
371 }
372 let constant = match rvalue {
373 Rvalue::Use(Operand::Constant(constant), ..)
374 | Rvalue::Cast(_, Operand::Constant(constant), _) => constant,
375 _ => continue,
376 };
377 let value = constant
378 .const_
379 .eval(
380 tcx,
381 rustc_middle::ty::TypingEnv::fully_monomorphized(),
382 DUMMY_SP,
383 )
384 .ok()?;
385 return crate::helpers::mir_utils::const_value_bytes(tcx, value, 0);
386 }
387 }
388 None
389}
390
391fn aggregate_op_is_nonzero<'tcx>(
392 tcx: TyCtxt<'tcx>,
393 body: &Body<'tcx>,
394 operand: &Operand<'tcx>,
395) -> bool {
396 if is_constant_zero_u8(operand) {
397 return false;
398 }
399 if operand_const_bytes(tcx, operand).is_some() {
400 return true;
401 }
402 match operand {
403 Operand::Copy(p) | Operand::Move(p) if p.projection.is_empty() => {
404 for data in body.basic_blocks.iter() {
405 if let Some(terminator) = &data.terminator {
406 if let TerminatorKind::Call { destination, func, .. } = &terminator.kind {
407 if destination.local == p.local && destination.projection.is_empty() {
408 return fn_always_returns_nonzero(tcx, func);
409 }
410 }
411 }
412 }
413 false
414 }
415 Operand::Constant(c) => {
416 c.const_
417 .try_to_scalar_int()
418 .map_or(false, |s| s.to_uint(s.size()) != 0)
419 }
420 _ => false,
421 }
422}
423
424fn operand_const_bytes<'tcx>(tcx: TyCtxt<'tcx>, operand: &Operand<'tcx>) -> Option<Vec<u8>> {
425 let constant = match operand {
426 Operand::Constant(c) => c,
427 _ => return None,
428 };
429 let value = constant
430 .const_
431 .eval(
432 tcx,
433 rustc_middle::ty::TypingEnv::fully_monomorphized(),
434 DUMMY_SP,
435 )
436 .ok()?;
437 crate::helpers::mir_utils::const_value_bytes(tcx, value, 0)
438}
439
440fn is_constant_zero_u8(operand: &Operand<'_>) -> bool {
441 let constant = match operand {
442 Operand::Constant(c) => c,
443 _ => return false,
444 };
445 constant
446 .const_
447 .try_to_scalar_int()
448 .map_or(false, |s| s.to_uint(s.size()) == 0)
449}
450
451fn fn_always_returns_nonzero<'tcx>(
452 tcx: TyCtxt<'tcx>,
453 func: &Operand<'tcx>,
454) -> bool {
455 let Some(fn_def_id) = crate::helpers::mir_utils::dep_callee_def_id(func) else { return false };
456 let callee_body = tcx.optimized_mir(fn_def_id);
457
458 let mut has_return = false;
459 for bb_data in callee_body.basic_blocks.iter() {
460 if let Some(terminator) = &bb_data.terminator {
461 if matches!(terminator.kind, TerminatorKind::Return) {
462 has_return = true;
463 }
464 }
465 for stmt in &bb_data.statements {
466 let StatementKind::Assign(assign) = &stmt.kind else { continue };
467 let (target, rvalue) = assign.as_ref();
468 if target.local != Local::from_usize(0) || !target.projection.is_empty() {
469 continue;
470 }
471 if !rvalue_is_nonzero(tcx, rvalue, callee_body) {
472 return false;
473 }
474 }
475 }
476
477 has_return
478}
479
480fn rvalue_is_nonzero<'tcx>(_tcx: TyCtxt<'tcx>, rvalue: &Rvalue<'tcx>, _body: &Body<'tcx>) -> bool {
481 match rvalue {
482 Rvalue::Use(Operand::Constant(c), ..) => {
483 c.const_
484 .try_to_scalar_int()
485 .map_or(false, |s| s.to_uint(s.size()) != 0)
486 }
487 Rvalue::Use(Operand::Copy(_), ..) | Rvalue::Use(Operand::Move(_), ..) => true,
488 _ => false,
489 }
490}
491
492pub(crate) fn body_parents<'tcx>(
493 tcx: TyCtxt<'tcx>,
494 body: &Body<'tcx>,
495) -> FxHashMap<Local, Local> {
496 let mut parents: FxHashMap<Local, Local> = Default::default();
497 for data in body.basic_blocks.iter() {
498 for statement in &data.statements {
499 let StatementKind::Assign(assign) = &statement.kind else {
500 continue;
501 };
502 let (target, rvalue) = assign.as_ref();
503 let source = match rvalue {
504 Rvalue::Use(Operand::Copy(place) | Operand::Move(place), ..)
505 | Rvalue::Cast(_, Operand::Copy(place) | Operand::Move(place), _)
506 | Rvalue::Ref(_, _, place)
507 | Rvalue::RawPtr(_, place)
508 | Rvalue::CopyForDeref(place) => Some(place.local),
509 _ => None,
510 };
511 if let Some(source) = source {
512 parents.entry(target.local).or_insert(source);
513 }
514 }
515 let Some(terminator) = &data.terminator else {
516 continue;
517 };
518 let TerminatorKind::Call {
519 func,
520 args,
521 destination,
522 ..
523 } = &terminator.kind
524 else {
525 continue;
526 };
527 let name = crate::helpers::mir_utils::call_name(tcx, func);
528 if !crate::helpers::api_classify::is_as_ptr(&name) {
529 continue;
530 }
531 let Some(source) = args.first().and_then(|arg| match &arg.node {
532 Operand::Copy(place) | Operand::Move(place) => Some(place.local),
533 _ => None,
534 }) else {
535 continue;
536 };
537 parents.entry(destination.local).or_insert(source);
538 }
539 parents
540}