[MoveOnlyAddressChecker] Fix repr for reinits.

The address checker records instructions that reinit fields in its
reinitInsts map.  Previously, that map mapped from an instruction to a
range of fields of the type.  But an instruction can use multiple
discontiguous fields of a single value.  (Indeed an attempt to add a
second range that was reinit'd by an already reinit'ing
instruction--even if it were overlapping or adjacent--would have no
effect and the map wouldn't be updated.)  Here, this is fixed by fixing
the representation and updating the storage whenver a new range is seen
to be reinit'd by the instruction.  As in
#66728 , a SmallBitVector is the
representation chosen.

rdar://111356251

Author: nate-chandler

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -626,7 +626,7 @@ struct UseState {
 
   /// memInstMustReinitialize insts. Contains both insts like copy_addr/store
   /// [assign] that are reinits that we will convert to inits and true reinits.
-  llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4> reinitInsts;
+  llvm::SmallMapVector<SILInstruction *, SmallBitVector, 4> reinitInsts;
 
   /// The set of drop_deinits of this mark_must_check
   SmallSetVector<SILInstruction *, 2> dropDeinitInsts;
@@ -671,6 +671,16 @@ struct UseState {
     range.setBits(bits);
   }
 
+  void recordReinitUse(SILInstruction *inst, TypeTreeLeafTypeRange range) {
+    auto iter = reinitInsts.find(inst);
+    if (iter == reinitInsts.end()) {
+      iter =
+          reinitInsts.insert({inst, SmallBitVector(getNumSubelements())}).first;
+    }
+    auto &bits = iter->second;
+    range.setBits(bits);
+  }
+
   /// Returns true if this is a terminator instruction that although it doesn't
   /// use our inout argument directly is used by the pass to ensure that we
   /// reinit said argument if we consumed it in the body of the function.
@@ -778,6 +788,11 @@ struct UseState {
     range.setBits(consumingBits);
   }
 
+  void recordConsumingBlock(SILBasicBlock *block, SmallBitVector &bits) {
+    auto &consumingBits = getOrCreateConsumingBlock(block);
+    consumingBits |= bits;
+  }
+
   void
   initializeLiveness(FieldSensitiveMultiDefPrunedLiveRange &prunedLiveness);
 
@@ -842,7 +857,7 @@ struct UseState {
     if (!isReinitToInitConvertibleInst(inst)) {
       auto iter = reinitInsts.find(inst);
       if (iter != reinitInsts.end()) {
-        if (span.setIntersection(iter->second))
+        if (span.intersects(iter->second))
           return true;
       }
     }
@@ -867,7 +882,7 @@ struct UseState {
     if (isReinitToInitConvertibleInst(inst)) {
       auto iter = reinitInsts.find(inst);
       if (iter != reinitInsts.end()) {
-        if (span.setIntersection(iter->second))
+        if (span.intersects(iter->second))
           return true;
       }
     }
@@ -1744,11 +1759,10 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
 
   if (::memInstMustReinitialize(op)) {
     LLVM_DEBUG(llvm::dbgs() << "Found reinit: " << *user);
-    assert(!useState.reinitInsts.count(user));
     auto leafRange = TypeTreeLeafTypeRange::get(op->get(), getRootAddress());
     if (!leafRange)
       return false;
-    useState.reinitInsts.insert({user, *leafRange});
+    useState.recordReinitUse(user, *leafRange);
     return true;
   }
 
@@ -2733,9 +2747,7 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
   for (auto reinitPair : addressUseState.reinitInsts) {
     if (!isReinitToInitConvertibleInst(reinitPair.first))
       continue;
-    SmallBitVector bits(liveness.getNumSubElements());
-    reinitPair.second.setBits(bits);
-    if (!consumes.claimConsume(reinitPair.first, bits))
+    if (!consumes.claimConsume(reinitPair.first, reinitPair.second))
       convertMemoryReinitToInitForm(reinitPair.first, debugVar);
   }
 
@@ -2926,7 +2938,7 @@ void ExtendUnconsumedLiveness::run() {
       }
     }
     for (auto pair : addressUseState.reinitInsts) {
-      if (pair.second.contains(element)) {
+      if (pair.second.test(element)) {
         destroys[pair.first] = DestroyKind::Reinit;
       }
     }

test/SILOptimizer/moveonly_addresschecker_unmaximized.sil

@@ -16,7 +16,7 @@ struct M4 {
 sil @get_M4 : $@convention(thin) () -> @owned M4
 sil @end_2 : $@convention(thin) (@owned M, @owned M) -> ()
 sil @see_addr_2 : $@convention(thin) (@in_guaranteed M, @in_guaranteed M) -> ()
-
+sil @replace_2 : $@convention(thin) (@inout M, @inout M) -> ()
 
 /// Two non-contiguous fields (#M4.s2, #M4.s4) are borrowed by @see_addr_2.
 /// Two non-contiguous fields (#M4.s1, #M$.s3) are consumed by @end_2.
@@ -65,3 +65,45 @@ bb0:
   return %22 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @rdar111356251 : $@convention(thin) () -> () {
+// CHECK:         [[STACK:%[^,]+]] = alloc_stack $M4
+// CHECK:         [[GET_M4:%[^,]+]] = function_ref @get_M4 : $@convention(thin) () -> @owned M4
+// CHECK:         [[INSTANCE:%[^,]+]] = apply [[GET_M4]]() : $@convention(thin) () -> @owned M4
+// CHECK:         store [[INSTANCE]] to [init] [[STACK]] : $*M4
+// CHECK:         [[S2_ADDR:%[^,]+]] = struct_element_addr [[STACK]] : $*M4, #M4.s2
+// CHECK:         [[S4_ADDR:%[^,]+]] = struct_element_addr [[STACK]] : $*M4, #M4.s4
+// CHECK:         [[REPLACE_2:%[^,]+]] = function_ref @replace_2 : $@convention(thin) (@inout M, @inout M) -> ()
+// CHECK:         apply [[REPLACE_2]]([[S2_ADDR]], [[S4_ADDR]]) : $@convention(thin) (@inout M, @inout M) -> ()
+// CHECK:         [[S4_ADDR_2:%[^,]+]] = struct_element_addr [[STACK]] : $*M4, #M4.s4
+// CHECK:         destroy_addr [[S4_ADDR_2]] : $*M
+// CHECK:         [[S2_ADDR_2:%[^,]+]] = struct_element_addr [[STACK]] : $*M4, #M4.s2
+// CHECK:         destroy_addr [[S2_ADDR_2]] : $*M
+// CHECK:         [[S1_ADDR:%[^,]+]] = struct_element_addr [[STACK]] : $*M4, #M4.s1
+// CHECK:         [[S1:%[^,]+]] = load [take] [[S1_ADDR]] : $*M
+// CHECK:         [[S3_ADDR:%[^,]+]] = struct_element_addr [[STACK]] : $*M4, #M4.s3
+// CHECK:         [[S3:%[^,]+]] = load [take] [[S3_ADDR]] : $*M
+// CHECK:         [[END_2:%[^,]+]] = function_ref @end_2 : $@convention(thin) (@owned M, @owned M) -> ()
+// CHECK:         apply [[END_2]]([[S1]], [[S3]]) : $@convention(thin) (@owned M, @owned M) -> ()
+// CHECK-LABEL: } // end sil function 'rdar111356251'
+sil [ossa] @rdar111356251 : $@convention(thin) () -> () {
+bb0:
+  %stack_addr = alloc_stack $M4
+  %stack = mark_must_check [consumable_and_assignable] %stack_addr : $*M4
+  %get_M4 = function_ref @get_M4 : $@convention(thin) () -> @owned M4
+  %instance = apply %get_M4() : $@convention(thin) () -> @owned M4
+  store %instance to [init] %stack : $*M4
+  %s2_addr = struct_element_addr %stack : $*M4, #M4.s2
+  %s4_addr = struct_element_addr %stack : $*M4, #M4.s4
+  %replace_2 = function_ref @replace_2 : $@convention(thin) (@inout M, @inout M) -> ()
+  apply %replace_2(%s2_addr, %s4_addr) : $@convention(thin) (@inout M, @inout M) -> ()
+  %12 = struct_element_addr %stack : $*M4, #M4.s1
+  %13 = load [copy] %12 : $*M
+  %14 = struct_element_addr %stack : $*M4, #M4.s3
+  %15 = load [copy] %14 : $*M
+  %16 = function_ref @end_2 : $@convention(thin) (@owned M, @owned M) -> ()
+  %17 = apply %16(%13, %15) : $@convention(thin) (@owned M, @owned M) -> ()
+  destroy_addr %stack : $*M4
+  dealloc_stack %stack_addr : $*M4
+  %22 = tuple ()
+  return %22 : $()
+}