|
| 1 | +//! Rewrite final-use copies before return into moves. |
| 2 | +//! |
| 3 | +//! # The problem |
| 4 | +//! |
| 5 | +//! MIR building represents reads of values whose type is `Copy` using `Operand::Copy`, including |
| 6 | +//! when such a local is returned. If that local's address has ever been observed, then the local's |
| 7 | +//! allocation is semantically valid until its `StorageDead` or function exit. This keeps the local |
| 8 | +//! live across `_0 = copy local`, so its live range overlaps with the return place and |
| 9 | +//! `MoveElimination` cannot unify the source local with `_0`. |
| 10 | +//! |
| 11 | +//! # The solution |
| 12 | +//! |
| 13 | +//! At function return, all local allocations are about to become invalid anyway. After borrowck, |
| 14 | +//! this pass can therefore turn a final-use `Copy` into a `Move`, as long as shortening the source |
| 15 | +//! local's live range has no observable effect before the return happens. |
| 16 | +//! Concretely, between the transformed copy (now a move) and the return, there may only be writes |
| 17 | +//! to unborrowed locals, storage markers, nops, and gotos. |
| 18 | +//! |
| 19 | +//! # The algorithm |
| 20 | +//! |
| 21 | +//! Start from every `Return` terminator, with `_0` treated as used by the return. Then scan |
| 22 | +//! predecessor blocks backward through `Goto` edges, forming a return-tail tree. The scan maintains |
| 23 | +//! `used_after`, the set of locals accessed later on that path. |
| 24 | +//! |
| 25 | +//! A `Copy` operand is rewritten to a `Move` when its base local is not in `used_after`. Then any |
| 26 | +//! locals touched by that operand, including index-projection locals, are added to `used_after` |
| 27 | +//! before the backward scan continues. |
| 28 | +//! |
| 29 | +//! The scan stops when accessing an indirect place because that may access any borrowed local, |
| 30 | +//! which would make the pass unable to prove any useful final uses. It also stops at writes to |
| 31 | +//! borrowed locals, because those can create a new address-observed allocation range whose overlap |
| 32 | +//! with an earlier borrowed local must be preserved. |
| 33 | +
|
| 34 | +use rustc_index::bit_set::DenseBitSet; |
| 35 | +use rustc_middle::mir::*; |
| 36 | +use rustc_middle::ty::TyCtxt; |
| 37 | +use rustc_mir_dataflow::impls::borrowed_locals; |
| 38 | + |
| 39 | +pub(super) struct TailCopyToMove; |
| 40 | + |
| 41 | +impl<'tcx> crate::MirPass<'tcx> for TailCopyToMove { |
| 42 | + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
| 43 | + sess.mir_opt_level() >= 2 |
| 44 | + } |
| 45 | + |
| 46 | + #[tracing::instrument(level = "trace", skip(self, _tcx, body))] |
| 47 | + fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 48 | + let borrowed = borrowed_locals(body); |
| 49 | + let predecessors = body.basic_blocks.predecessors().clone(); |
| 50 | + let mut stack = Vec::new(); |
| 51 | + |
| 52 | + // A return terminator implicitly uses the return place. Walking backward through |
| 53 | + // assignments records the locals accessed later on this path. |
| 54 | + for (bb, data) in body.basic_blocks.iter_enumerated() { |
| 55 | + if matches!(data.terminator().kind, TerminatorKind::Return) { |
| 56 | + let mut used_after = DenseBitSet::new_empty(body.local_decls.len()); |
| 57 | + used_after.insert(RETURN_PLACE); |
| 58 | + stack.push(TailState { block: bb, used_after }); |
| 59 | + } |
| 60 | + } |
| 61 | + |
| 62 | + while let Some(mut state) = stack.pop() { |
| 63 | + loop { |
| 64 | + // `scan_block` rewrites final-use copies in this block and updates `used_after` to |
| 65 | + // the locals whose allocation is accessed after the block starts. If the block is not |
| 66 | + // pure tail code, this path is done. |
| 67 | + if !scan_block(body, state.block, &mut state.used_after, &borrowed) { |
| 68 | + break; |
| 69 | + } |
| 70 | + |
| 71 | + // Continue through predecessor blocks only when the predecessor's terminator is a |
| 72 | + // plain `Goto` to this block. Other terminators are control-flow or effect |
| 73 | + // boundaries. Follow one predecessor in-place and push only sibling paths, so |
| 74 | + // straight-line return tails don't clone the bitset at every block. |
| 75 | + let mut next = None; |
| 76 | + for pred in predecessors[state.block].iter().copied() { |
| 77 | + let terminator = body.basic_blocks[pred].terminator(); |
| 78 | + if let TerminatorKind::Goto { target } = terminator.kind { |
| 79 | + debug_assert_eq!(target, state.block); |
| 80 | + if next.is_none() { |
| 81 | + next = Some(pred); |
| 82 | + } else { |
| 83 | + stack.push(TailState { |
| 84 | + block: pred, |
| 85 | + used_after: state.used_after.clone(), |
| 86 | + }); |
| 87 | + } |
| 88 | + } |
| 89 | + } |
| 90 | + |
| 91 | + let Some(next) = next else { |
| 92 | + break; |
| 93 | + }; |
| 94 | + state.block = next; |
| 95 | + } |
| 96 | + } |
| 97 | + } |
| 98 | + |
| 99 | + fn is_required(&self) -> bool { |
| 100 | + false |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +struct TailState { |
| 105 | + block: BasicBlock, |
| 106 | + used_after: DenseBitSet<Local>, |
| 107 | +} |
| 108 | + |
| 109 | +/// Scan a block backward while the return-tail invariant still holds. |
| 110 | +/// |
| 111 | +/// The invariant is that a whole-local `Copy` can be changed to a `Move` only if this path has no |
| 112 | +/// later access to that local's allocation before returning, and no later operation whose observable |
| 113 | +/// behavior could depend on ending an address-observed local's allocation early. `used_after` |
| 114 | +/// tracks those later local-allocation accesses. |
| 115 | +fn scan_block<'tcx>( |
| 116 | + body: &mut Body<'tcx>, |
| 117 | + block: BasicBlock, |
| 118 | + used_after: &mut DenseBitSet<Local>, |
| 119 | + borrowed: &DenseBitSet<Local>, |
| 120 | +) -> bool { |
| 121 | + for statement in body.basic_blocks.as_mut_preserves_cfg()[block].statements.iter_mut().rev() { |
| 122 | + match &mut statement.kind { |
| 123 | + // Under the local lifetime semantics from RFC 3943, `StorageLive` does not allocate, |
| 124 | + // and `StorageDead` has no effect if the local was already freed by a move. These |
| 125 | + // markers therefore do not affect whether a copy can be treated as a final use. |
| 126 | + StatementKind::StorageLive(_) | StatementKind::StorageDead(_) | StatementKind::Nop => {} |
| 127 | + StatementKind::Assign(assign) => { |
| 128 | + let dest = assign.0; |
| 129 | + |
| 130 | + // Accessing an indirect place may touch any borrowed local, so continuing would |
| 131 | + // require treating all borrowed locals as used after this point. |
| 132 | + if dest.is_indirect_first_projection() { |
| 133 | + return false; |
| 134 | + } |
| 135 | + |
| 136 | + // Writing to a borrowed local can start a new allocation range. Shortening an |
| 137 | + // earlier borrowed local could remove an overlap with that new range. |
| 138 | + if borrowed.contains(dest.local) { |
| 139 | + return false; |
| 140 | + } |
| 141 | + |
| 142 | + // A destination write accesses the base local, and evaluating the destination may |
| 143 | + // also access projection locals, such as an index. |
| 144 | + record_place_locals(dest, used_after); |
| 145 | + |
| 146 | + // This pass only models `Use` and `Aggregate` rvalues whose operands are direct. |
| 147 | + // Other rvalues are outside the conservative return-tail shape handled here. |
| 148 | + if !process_rvalue(&mut assign.1, used_after) { |
| 149 | + return false; |
| 150 | + } |
| 151 | + } |
| 152 | + StatementKind::SetDiscriminant { place, .. } => { |
| 153 | + let dest = **place; |
| 154 | + |
| 155 | + // Accessing an indirect place may touch any borrowed local, so continuing would |
| 156 | + // require treating all borrowed locals as used after this point. |
| 157 | + if dest.is_indirect_first_projection() { |
| 158 | + return false; |
| 159 | + } |
| 160 | + |
| 161 | + // Writing to a borrowed local can start a new allocation range. Shortening an |
| 162 | + // earlier borrowed local could remove an overlap with that new range. |
| 163 | + if borrowed.contains(dest.local) { |
| 164 | + return false; |
| 165 | + } |
| 166 | + |
| 167 | + // `SetDiscriminant` has a validity invariant on the rest of the place, so treat |
| 168 | + // the base local as accessed along with any projection locals. |
| 169 | + record_place_locals(dest, used_after); |
| 170 | + } |
| 171 | + _ => { |
| 172 | + // Anything else may perform effects or evaluate places in ways this pass does not |
| 173 | + // model, so it is not part of the pure return tail. |
| 174 | + return false; |
| 175 | + } |
| 176 | + } |
| 177 | + } |
| 178 | + |
| 179 | + true |
| 180 | +} |
| 181 | + |
| 182 | +/// Records all locals used in a place, including `Index` projections in `used_after`. |
| 183 | +fn record_place_locals<'tcx>(place: Place<'tcx>, used_after: &mut DenseBitSet<Local>) { |
| 184 | + for local in place.as_ref().accessed_locals() { |
| 185 | + used_after.insert(local); |
| 186 | + } |
| 187 | +} |
| 188 | + |
| 189 | +/// Process the RHS of an assignment in a pure return tail. |
| 190 | +fn process_rvalue<'tcx>(rvalue: &mut Rvalue<'tcx>, used_after: &mut DenseBitSet<Local>) -> bool { |
| 191 | + match rvalue { |
| 192 | + Rvalue::Use(operand, _) => process_operand(operand, used_after), |
| 193 | + Rvalue::Aggregate(_, operands) => { |
| 194 | + // Operands are evaluated left-to-right. We scan them right-to-left so `used_after` |
| 195 | + // includes uses later in the same statement. If an operand accesses an indirect place, |
| 196 | + // only earlier operands and earlier statements are outside the pure tail. |
| 197 | + for operand in operands.iter_mut().rev() { |
| 198 | + if !process_operand(operand, used_after) { |
| 199 | + return false; |
| 200 | + } |
| 201 | + } |
| 202 | + true |
| 203 | + } |
| 204 | + _ => { |
| 205 | + // This pass doesn't model other rvalues, so they are not part of the pure return tail. |
| 206 | + return false; |
| 207 | + } |
| 208 | + } |
| 209 | +} |
| 210 | + |
| 211 | +/// Process one operand in an rvalue. |
| 212 | +fn process_operand<'tcx>(operand: &mut Operand<'tcx>, used_after: &mut DenseBitSet<Local>) -> bool { |
| 213 | + let place = match operand { |
| 214 | + Operand::Copy(place) | Operand::Move(place) if place.is_indirect_first_projection() => { |
| 215 | + // Accessing an indirect place may touch any borrowed local. Continuing would |
| 216 | + // require treating all borrowed locals as used after this point, which would prevent |
| 217 | + // the useful copy-to-move rewrites this pass is looking for. |
| 218 | + return false; |
| 219 | + } |
| 220 | + Operand::Copy(place) => { |
| 221 | + let place = *place; |
| 222 | + // No later operation in the scanned tail accesses this local's allocation, so this |
| 223 | + // copy is a final use on the current return path and can be represented as a move. |
| 224 | + if !used_after.contains(place.local) { |
| 225 | + *operand = Operand::Move(place); |
| 226 | + } |
| 227 | + Some(place) |
| 228 | + } |
| 229 | + Operand::Move(place) => Some(*place), |
| 230 | + Operand::Constant(_) | Operand::RuntimeChecks(_) => None, |
| 231 | + }; |
| 232 | + |
| 233 | + if let Some(place) = place { |
| 234 | + // Reading an operand place accesses its base local, and evaluating its projections may |
| 235 | + // access additional locals, such as the index local in `place[index]`. |
| 236 | + record_place_locals(place, used_after); |
| 237 | + } |
| 238 | + |
| 239 | + true |
| 240 | +} |
0 commit comments