[pmo] Teach the use collector how to handle store [assign].

I also cleaned up the code there to make it explicit how each ownership
qualifier maps to PMOUseKind.

Author: gottesmm

lib/SILOptimizer/Mandatory/PMOMemoryUseCollector.cpp

@@ -239,21 +239,36 @@ bool ElementUseCollector::collectUses(SILValue Pointer) {
     }
 
     // Stores *to* the allocation are writes.
-    if (isa<StoreInst>(User) && UI->getOperandNumber() == 1) {
-      if (PointeeType.is<TupleType>()) {
-        UsesToScalarize.push_back(User);
+    if (auto *si = dyn_cast<StoreInst>(User)) {
+      if (UI->getOperandNumber() == StoreInst::Dest) {
+        if (PointeeType.is<TupleType>()) {
+          UsesToScalarize.push_back(User);
+          continue;
+        }
+
+        auto kind = ([&]() -> PMOUseKind {
+          switch (si->getOwnershipQualifier()) {
+          // Coming out of SILGen, we assume that raw stores are
+          // initializations, unless they have trivial type (which we classify
+          // as InitOrAssign).
+          case StoreOwnershipQualifier::Unqualified:
+            if (PointeeType.isTrivial(User->getModule()))
+              return PMOUseKind::InitOrAssign;
+            return PMOUseKind::Initialization;
+
+          case StoreOwnershipQualifier::Init:
+            return PMOUseKind::Initialization;
+
+          case StoreOwnershipQualifier::Assign:
+            return PMOUseKind::Assign;
+
+          case StoreOwnershipQualifier::Trivial:
+            return PMOUseKind::InitOrAssign;
+          }
+        })();
+        Uses.emplace_back(si, kind);
         continue;
       }
-
-      // Coming out of SILGen, we assume that raw stores are initializations,
-      // unless they have trivial type (which we classify as InitOrAssign).
-      auto Kind = ([&]() -> PMOUseKind {
-        if (PointeeType.isTrivial(User->getModule()))
-          return PMOUseKind::InitOrAssign;
-        return PMOUseKind::Initialization;
-      })();
-      Uses.emplace_back(User, Kind);
-      continue;
     }
 
     if (auto *CAI = dyn_cast<CopyAddrInst>(User)) {

test/SILOptimizer/predictable_memaccess_opts.sil

@@ -441,6 +441,27 @@ bb0(%0 : @owned $KlassAndTuple):
   return %4 : $Klass
 }
 
+// CHECK-LABEL: sil [ossa] @simple_assignstore : $@convention(thin) (@owned Klass, @owned Klass) -> @owned Klass {
+// CHECK: bb0([[ARG0:%.*]] : @owned $Klass, [[ARG1:%.*]] : @owned $Klass):
+// CHECK:   [[STACK:%.*]] = alloc_stack $Klass
+// CHECK:   store [[ARG0]] to [init] [[STACK]]
+// CHECK:   [[ARG1_COPY:%.*]] = copy_value [[ARG1]]
+// CHECK:   store [[ARG1]] to [assign] [[STACK]]
+// CHECK:   destroy_addr [[STACK]]
+// CHECK:   dealloc_stack [[STACK]]
+// CHECK:   return [[ARG1_COPY]]
+// CHECK: } // end sil function 'simple_assignstore'
+sil [ossa] @simple_assignstore : $@convention(thin) (@owned Klass, @owned Klass) -> @owned Klass {
+bb0(%0 : @owned $Klass, %1 : @owned $Klass):
+  %2 = alloc_stack $Klass
+  store %0 to [init] %2 : $*Klass
+  store %1 to [assign] %2 : $*Klass
+  %3 = load [copy] %2 : $*Klass
+  destroy_addr %2 : $*Klass
+  dealloc_stack %2 : $*Klass
+  return %3 : $Klass
+}
+
 ////////////////////
 // Negative Tests //
 ////////////////////