Auto merge of #142531 - ohadravid:better-storage-calls-copy-prop, r=<… · rust-lang/rust@ef7d206 (original) (raw)

``

1

`+

use std::borrow::Cow;

`

``

2

+

1

3

`use rustc_index::IndexSlice;

`

2

4

`use rustc_index::bit_set::DenseBitSet;

`

3

5

`use rustc_middle::mir::visit::*;

`

4

6

`use rustc_middle::mir::*;

`

5

7

`use rustc_middle::ty::TyCtxt;

`

``

8

`+

use rustc_mir_dataflow::impls::{MaybeStorageDead, always_storage_live_locals};

`

``

9

`+

use rustc_mir_dataflow::{Analysis, ResultsCursor};

`

6

10

`use tracing::{debug, instrument};

`

7

11

``

8

12

`use crate::ssa::SsaLocals;

`

`@@ -16,7 +20,7 @@ use crate::ssa::SsaLocals;

`

16

20

`/// _d = move? _c

`

17

21

`/// where each of the locals is only assigned once.

`

18

22

`///

`

19

``

`` -

/// We want to replace all those locals by _a, either copied or moved.

``

``

23

`` +

/// We want to replace all those locals by _a (the "head"), either copied or moved.

``

20

24

`pub(super) struct CopyProp;

`

21

25

``

22

26

`impl<'tcx> crate::MirPass<'tcx> for CopyProp {

`

`@@ -34,15 +38,40 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {

`

34

38

`let fully_moved = fully_moved_locals(&ssa, body);

`

35

39

`debug!(?fully_moved);

`

36

40

``

37

``

`-

let mut storage_to_remove = DenseBitSet::new_empty(fully_moved.domain_size());

`

``

41

`+

let mut head_storage_to_check = DenseBitSet::new_empty(fully_moved.domain_size());

`

``

42

+

38

43

`for (local, &head) in ssa.copy_classes().iter_enumerated() {

`

39

44

`if local != head {

`

40

``

`-

storage_to_remove.insert(head);

`

``

45

`+

// We need to determine if we can keep the head's storage statements (which enables better optimizations).

`

``

46

`` +

// For every local's usage location, if the head is in maybe_storage_dead, we have to remove the storage statements for it.

``

``

47

`+

head_storage_to_check.insert(head);

`

41

48

`}

`

42

49

`}

`

43

50

``

44

51

`let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);

`

45

52

``

``

53

`+

let storage_to_remove = if any_replacement {

`

``

54

`+

let always_live_locals = &always_storage_live_locals(body);

`

``

55

+

``

56

`+

let maybe_storage_dead = MaybeStorageDead::new(Cow::Borrowed(always_live_locals))

`

``

57

`+

.iterate_to_fixpoint(tcx, body, None)

`

``

58

`+

.into_results_cursor(body);

`

``

59

+

``

60

`+

let mut storage_checker = StorageChecker {

`

``

61

`+

copy_classes: ssa.copy_classes(),

`

``

62

`+

maybe_storage_dead,

`

``

63

`+

head_storage_to_check,

`

``

64

`+

storage_to_remove: DenseBitSet::new_empty(fully_moved.domain_size()),

`

``

65

`+

};

`

``

66

+

``

67

`+

storage_checker.visit_body(body);

`

``

68

+

``

69

`+

storage_checker.storage_to_remove

`

``

70

`+

} else {

`

``

71

`+

// Will be empty anyway.

`

``

72

`+

head_storage_to_check

`

``

73

`+

};

`

``

74

+

46

75

`Replacer {

`

47

76

` tcx,

`

48

77

`copy_classes: ssa.copy_classes(),

`

`@@ -119,6 +148,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {

`

119

148

`if self.borrowed_locals.contains(*local) {

`

120

149

`return;

`

121

150

`}

`

``

151

+

122

152

`match ctxt {

`

123

153

`// Do not modify the local in storage statements.

`

124

154

`PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}

`

`@@ -172,3 +202,29 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {

`

172

202

`}

`

173

203

`}

`

174

204

`}

`

``

205

+

``

206

`+

struct StorageChecker<'a, 'tcx> {

`

``

207

`+

storage_to_remove: DenseBitSet,

`

``

208

`+

head_storage_to_check: DenseBitSet,

`

``

209

`+

maybe_storage_dead: ResultsCursor<'a, 'tcx, MaybeStorageDead<'a>>,

`

``

210

`+

copy_classes: &'a IndexSlice<Local, Local>,

`

``

211

`+

}

`

``

212

+

``

213

`+

impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {

`

``

214

`+

fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {

`

``

215

`+

let head = self.copy_classes[local];

`

``

216

+

``

217

`+

if context.is_use() && self.head_storage_to_check.contains(head) {

`

``

218

`+

self.maybe_storage_dead.seek_after_primary_effect(location);

`

``

219

`+

if self.maybe_storage_dead.get().contains(head) {

`

``

220

`+

debug!(

`

``

221

`+

?location,

`

``

222

`+

?local,

`

``

223

`+

?head,

`

``

224

`+

"found use of local with head at a location in which head is maybe dead, marking head for storage removal"

`

``

225

`+

);

`

``

226

`+

self.storage_to_remove.insert(head);

`

``

227

`+

}

`

``

228

`+

}

`

``

229

`+

}

`

``

230

`+

}

`