1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
//! Trivial copy propagation pass.
//!
//! This uses def-use analysis to remove values that have exactly one def and one use, which must
//! be an assignment.
//!
//! To give an example, we look for patterns that look like:
//!
//!     DEST = SRC
//!     ...
//!     USE(DEST)
//!
//! where `DEST` and `SRC` are both locals of some form. We replace that with:
//!
//!     NOP
//!     ...
//!     USE(SRC)
//!
//! The assignment `DEST = SRC` must be (a) the only mutation of `DEST` and (b) the only
//! (non-mutating) use of `SRC`. These restrictions are conservative and may be relaxed in the
//! future.

use crate::transform::{MirPass, MirSource};
use crate::util::def_use::DefUseAnalysis;
use rustc_middle::mir::visit::MutVisitor;
use rustc_middle::mir::{
    Body, Constant, Local, LocalKind, Location, Operand, Place, Rvalue, StatementKind,
};
use rustc_middle::ty::TyCtxt;

pub struct CopyPropagation;

impl<'tcx> MirPass<'tcx> for CopyPropagation {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, _source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        // We only run when the MIR optimization level is > 1.
        // This avoids a slow pass, and messing up debug info.
        if tcx.sess.opts.debugging_opts.mir_opt_level <= 1 {
            return;
        }

        let mut def_use_analysis = DefUseAnalysis::new(body);
        loop {
            def_use_analysis.analyze(body);

            if eliminate_self_assignments(body, &def_use_analysis) {
                def_use_analysis.analyze(body);
            }

            let mut changed = false;
            for dest_local in body.local_decls.indices() {
                debug!("considering destination local: {:?}", dest_local);

                let action;
                let location;
                {
                    // The destination must have exactly one def.
                    let dest_use_info = def_use_analysis.local_info(dest_local);
                    let dest_def_count = dest_use_info.def_count_not_including_drop();
                    if dest_def_count == 0 {
                        debug!("  Can't copy-propagate local: dest {:?} undefined", dest_local);
                        continue;
                    }
                    if dest_def_count > 1 {
                        debug!(
                            "  Can't copy-propagate local: dest {:?} defined {} times",
                            dest_local,
                            dest_use_info.def_count()
                        );
                        continue;
                    }
                    if dest_use_info.use_count() == 0 {
                        debug!("  Can't copy-propagate local: dest {:?} unused", dest_local);
                        continue;
                    }
                    // Conservatively gives up if the dest is an argument,
                    // because there may be uses of the original argument value.
                    // Also gives up on the return place, as we cannot propagate into its implicit
                    // use by `return`.
                    if matches!(
                        body.local_kind(dest_local),
                        LocalKind::Arg | LocalKind::ReturnPointer
                    ) {
                        debug!("  Can't copy-propagate local: dest {:?} (argument)", dest_local);
                        continue;
                    }
                    let dest_place_def = dest_use_info.defs_not_including_drop().next().unwrap();
                    location = dest_place_def.location;

                    let basic_block = &body[location.block];
                    let statement_index = location.statement_index;
                    let statement = match basic_block.statements.get(statement_index) {
                        Some(statement) => statement,
                        None => {
                            debug!("  Can't copy-propagate local: used in terminator");
                            continue;
                        }
                    };

                    // That use of the source must be an assignment.
                    match &statement.kind {
                        StatementKind::Assign(box (place, Rvalue::Use(operand))) => {
                            if let Some(local) = place.as_local() {
                                if local == dest_local {
                                    let maybe_action = match operand {
                                        Operand::Copy(src_place) | Operand::Move(src_place) => {
                                            Action::local_copy(&body, &def_use_analysis, *src_place)
                                        }
                                        Operand::Constant(ref src_constant) => {
                                            Action::constant(src_constant)
                                        }
                                    };
                                    match maybe_action {
                                        Some(this_action) => action = this_action,
                                        None => continue,
                                    }
                                } else {
                                    debug!(
                                        "  Can't copy-propagate local: source use is not an \
                                    assignment"
                                    );
                                    continue;
                                }
                            } else {
                                debug!(
                                    "  Can't copy-propagate local: source use is not an \
                                    assignment"
                                );
                                continue;
                            }
                        }
                        _ => {
                            debug!(
                                "  Can't copy-propagate local: source use is not an \
                                    assignment"
                            );
                            continue;
                        }
                    }
                }

                changed =
                    action.perform(body, &def_use_analysis, dest_local, location, tcx) || changed;
                // FIXME(pcwalton): Update the use-def chains to delete the instructions instead of
                // regenerating the chains.
                break;
            }
            if !changed {
                break;
            }
        }
    }
}

fn eliminate_self_assignments(body: &mut Body<'_>, def_use_analysis: &DefUseAnalysis) -> bool {
    let mut changed = false;

    for dest_local in body.local_decls.indices() {
        let dest_use_info = def_use_analysis.local_info(dest_local);

        for def in dest_use_info.defs_not_including_drop() {
            let location = def.location;
            if let Some(stmt) = body[location.block].statements.get(location.statement_index) {
                match &stmt.kind {
                    StatementKind::Assign(box (
                        place,
                        Rvalue::Use(Operand::Copy(src_place) | Operand::Move(src_place)),
                    )) => {
                        if let (Some(local), Some(src_local)) =
                            (place.as_local(), src_place.as_local())
                        {
                            if local == dest_local && dest_local == src_local {
                            } else {
                                continue;
                            }
                        } else {
                            continue;
                        }
                    }
                    _ => {
                        continue;
                    }
                }
            } else {
                continue;
            }
            debug!("deleting a self-assignment for {:?}", dest_local);
            body.make_statement_nop(location);
            changed = true;
        }
    }

    changed
}

enum Action<'tcx> {
    PropagateLocalCopy(Local),
    PropagateConstant(Constant<'tcx>),
}

impl<'tcx> Action<'tcx> {
    fn local_copy(
        body: &Body<'tcx>,
        def_use_analysis: &DefUseAnalysis,
        src_place: Place<'tcx>,
    ) -> Option<Action<'tcx>> {
        // The source must be a local.
        let src_local = if let Some(local) = src_place.as_local() {
            local
        } else {
            debug!("  Can't copy-propagate local: source is not a local");
            return None;
        };

        // We're trying to copy propagate a local.
        // There must be exactly one use of the source used in a statement (not in a terminator).
        let src_use_info = def_use_analysis.local_info(src_local);
        let src_use_count = src_use_info.use_count();
        if src_use_count == 0 {
            debug!("  Can't copy-propagate local: no uses");
            return None;
        }
        if src_use_count != 1 {
            debug!("  Can't copy-propagate local: {} uses", src_use_info.use_count());
            return None;
        }

        // Verify that the source doesn't change in between. This is done conservatively for now,
        // by ensuring that the source has exactly one mutation. The goal is to prevent things
        // like:
        //
        //     DEST = SRC;
        //     SRC = X;
        //     USE(DEST);
        //
        // From being misoptimized into:
        //
        //     SRC = X;
        //     USE(SRC);
        let src_def_count = src_use_info.def_count_not_including_drop();
        // allow function arguments to be propagated
        let is_arg = body.local_kind(src_local) == LocalKind::Arg;
        if (is_arg && src_def_count != 0) || (!is_arg && src_def_count != 1) {
            debug!(
                "  Can't copy-propagate local: {} defs of src{}",
                src_def_count,
                if is_arg { " (argument)" } else { "" },
            );
            return None;
        }

        Some(Action::PropagateLocalCopy(src_local))
    }

    fn constant(src_constant: &Constant<'tcx>) -> Option<Action<'tcx>> {
        Some(Action::PropagateConstant(*src_constant))
    }

    fn perform(
        self,
        body: &mut Body<'tcx>,
        def_use_analysis: &DefUseAnalysis,
        dest_local: Local,
        location: Location,
        tcx: TyCtxt<'tcx>,
    ) -> bool {
        match self {
            Action::PropagateLocalCopy(src_local) => {
                // Eliminate the destination and the assignment.
                //
                // First, remove all markers.
                //
                // FIXME(pcwalton): Don't do this. Merge live ranges instead.
                debug!("  Replacing all uses of {:?} with {:?} (local)", dest_local, src_local);
                for place_use in &def_use_analysis.local_info(dest_local).defs_and_uses {
                    if place_use.context.is_storage_marker() {
                        body.make_statement_nop(place_use.location)
                    }
                }
                for place_use in &def_use_analysis.local_info(src_local).defs_and_uses {
                    if place_use.context.is_storage_marker() {
                        body.make_statement_nop(place_use.location)
                    }
                }

                // Replace all uses of the destination local with the source local.
                def_use_analysis.replace_all_defs_and_uses_with(dest_local, body, src_local, tcx);

                // Finally, zap the now-useless assignment instruction.
                debug!("  Deleting assignment");
                body.make_statement_nop(location);

                true
            }
            Action::PropagateConstant(src_constant) => {
                // First, remove all markers.
                //
                // FIXME(pcwalton): Don't do this. Merge live ranges instead.
                debug!(
                    "  Replacing all uses of {:?} with {:?} (constant)",
                    dest_local, src_constant
                );
                let dest_local_info = def_use_analysis.local_info(dest_local);
                for place_use in &dest_local_info.defs_and_uses {
                    if place_use.context.is_storage_marker() {
                        body.make_statement_nop(place_use.location)
                    }
                }

                // Replace all uses of the destination local with the constant.
                let mut visitor = ConstantPropagationVisitor::new(dest_local, src_constant, tcx);
                for dest_place_use in &dest_local_info.defs_and_uses {
                    visitor.visit_location(body, dest_place_use.location)
                }

                // Zap the assignment instruction if we eliminated all the uses. We won't have been
                // able to do that if the destination was used in a projection, because projections
                // must have places on their LHS.
                let use_count = dest_local_info.use_count();
                if visitor.uses_replaced == use_count {
                    debug!(
                        "  {} of {} use(s) replaced; deleting assignment",
                        visitor.uses_replaced, use_count
                    );
                    body.make_statement_nop(location);
                    true
                } else if visitor.uses_replaced == 0 {
                    debug!("  No uses replaced; not deleting assignment");
                    false
                } else {
                    debug!(
                        "  {} of {} use(s) replaced; not deleting assignment",
                        visitor.uses_replaced, use_count
                    );
                    true
                }
            }
        }
    }
}

struct ConstantPropagationVisitor<'tcx> {
    dest_local: Local,
    constant: Constant<'tcx>,
    tcx: TyCtxt<'tcx>,
    uses_replaced: usize,
}

impl<'tcx> ConstantPropagationVisitor<'tcx> {
    fn new(
        dest_local: Local,
        constant: Constant<'tcx>,
        tcx: TyCtxt<'tcx>,
    ) -> ConstantPropagationVisitor<'tcx> {
        ConstantPropagationVisitor { dest_local, constant, tcx, uses_replaced: 0 }
    }
}

impl<'tcx> MutVisitor<'tcx> for ConstantPropagationVisitor<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
        self.super_operand(operand, location);

        match operand {
            Operand::Copy(place) | Operand::Move(place) => {
                if let Some(local) = place.as_local() {
                    if local == self.dest_local {
                    } else {
                        return;
                    }
                } else {
                    return;
                }
            }
            _ => return,
        }

        *operand = Operand::Constant(box self.constant);
        self.uses_replaced += 1
    }
}