Merge pull request #74298 from nate-chandler/rdar129593468

[NoncopyableWrapperElim] Process undef values.

Author: nate-chandler

include/swift/SIL/SILFunction.h

@@ -28,6 +28,8 @@
 #include "swift/SIL/SILDeclRef.h"
 #include "swift/SIL/SILLinkage.h"
 #include "swift/SIL/SILPrintContext.h"
+#include "swift/SIL/SILUndef.h"
+#include "llvm/ADT/MapVector.h"
 
 namespace swift {
 
@@ -337,12 +339,8 @@ class SILFunction
 
   PerformanceConstraints perfConstraints = PerformanceConstraints::None;
 
-  /// This is the set of undef values we've created, for uniquing purposes.
-  ///
-  /// We use a SmallDenseMap since in most functions, we will have only one type
-  /// of undef if we have any at all. In that case, by staying small we avoid
-  /// needing a heap allocation.
-  llvm::SmallDenseMap<SILType, SILUndef *, 1> undefValues;
+  /// The undefs of each type in the function.
+  llvm::SmallMapVector<SILType, SILUndef *, 1> undefValues;
 
   /// This is the number of uses of this SILFunction inside the SIL.
   /// It does not include references from debug scopes.
@@ -1637,6 +1635,10 @@ class SILFunction
         decl->getLifetimeAnnotation());
   }
 
+  ArrayRef<std::pair<SILType, SILUndef *>> getUndefValues() {
+    return {undefValues.begin(), undefValues.end()};
+  }
+
   /// verify - Run the SIL verifier to make sure that the SILFunction follows
   /// invariants.
   void verify(CalleeCache *calleeCache = nullptr,

include/swift/SIL/SILInstruction.h

@@ -9162,7 +9162,7 @@ class MoveOnlyWrapperToCopyableBoxInst
                                    ValueOwnershipKind forwardingOwnershipKind)
       : UnaryInstructionBase(
             DebugLoc, operand,
-            operand->getType().removingMoveOnlyWrapperToBoxedType(
+            operand->getType().removingMoveOnlyWrapperFromBoxedType(
                 operand->getFunction()),
             forwardingOwnershipKind) {
     assert(

include/swift/SIL/SILType.h

@@ -871,7 +871,15 @@ class SILType {
   ///
   /// DISCUSSION: This is separate from removingMoveOnlyWrapper since this API
   /// requires a SILFunction * and is specialized.
-  SILType removingMoveOnlyWrapperToBoxedType(const SILFunction *fn);
+  SILType removingMoveOnlyWrapperFromBoxedType(const SILFunction *fn);
+
+  /// Whether there's a direct wrapper or a wrapper inside a box.
+  bool hasAnyMoveOnlyWrapping(const SILFunction *fn) {
+    return isMoveOnlyWrapped() || isBoxedMoveOnlyWrappedType(fn);
+  }
+
+  /// Removes a direct wrapper from a type or a wrapper from a type in a box.
+  SILType removingAnyMoveOnlyWrapping(const SILFunction *fn);
 
   /// Returns a SILType with any archetypes mapped out of context.
   SILType mapTypeOutOfContext() const;

include/swift/SIL/SILValue.h

@@ -599,14 +599,9 @@ class ValueBase : public SILNode, public SILAllocated<ValueBase> {
   /// NOTE: Please do not use this directly! It is only meant to be used by the
   /// optimizer pass: SILMoveOnlyWrappedTypeEliminator.
   bool unsafelyEliminateMoveOnlyWrapper(const SILFunction *fn) {
-    if (!Type.isMoveOnlyWrapped() && !Type.isBoxedMoveOnlyWrappedType(fn))
+    if (!Type.hasAnyMoveOnlyWrapping(fn))
       return false;
-    if (Type.isMoveOnlyWrapped()) {
-      Type = Type.removingMoveOnlyWrapper();
-    } else {
-      assert(Type.isBoxedMoveOnlyWrappedType(fn));
-      Type = Type.removingMoveOnlyWrapperToBoxedType(fn);
-    }
+    Type = Type.removingAnyMoveOnlyWrapping(fn);
     return true;
   }
 

lib/SIL/IR/SILType.cpp

@@ -1271,7 +1271,7 @@ SILType SILType::addingMoveOnlyWrapperToBoxedType(const SILFunction *fn) {
   return SILType::getPrimitiveObjectType(newBoxType);
 }
 
-SILType SILType::removingMoveOnlyWrapperToBoxedType(const SILFunction *fn) {
+SILType SILType::removingMoveOnlyWrapperFromBoxedType(const SILFunction *fn) {
   auto boxTy = castTo<SILBoxType>();
   auto *oldLayout = boxTy->getLayout();
   auto oldField = oldLayout->getFields()[0];
@@ -1289,6 +1289,17 @@ SILType SILType::removingMoveOnlyWrapperToBoxedType(const SILFunction *fn) {
   return SILType::getPrimitiveObjectType(newBoxType);
 }
 
+SILType SILType::removingAnyMoveOnlyWrapping(const SILFunction *fn) {
+  if (!isMoveOnlyWrapped() && !isBoxedMoveOnlyWrappedType(fn))
+    return *this;
+
+  if (isMoveOnlyWrapped())
+    return removingMoveOnlyWrapper();
+
+  assert(isBoxedMoveOnlyWrappedType(fn));
+  return removingMoveOnlyWrapperFromBoxedType(fn);
+}
+
 bool SILType::isSendable(SILFunction *fn) const {
   return getASTType()->isSendableType();
 }

lib/SIL/Verifier/SILVerifier.cpp

@@ -6548,9 +6548,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
             "Operand value should be an object");
     require(cvt->getOperand()->getType().isBoxedMoveOnlyWrappedType(cvt->getFunction()),
             "Operand should be move only wrapped");
-    require(cvt->getType() ==
-            cvt->getOperand()->getType().removingMoveOnlyWrapperToBoxedType(cvt->getFunction()),
-            "Result and operand must have the same type, today.");
+    require(
+        cvt->getType() ==
+            cvt->getOperand()->getType().removingMoveOnlyWrapperFromBoxedType(
+                cvt->getFunction()),
+        "Result and operand must have the same type, today.");
   }
 
   void checkCopyableToMoveOnlyWrapperValueInst(

lib/SILOptimizer/Mandatory/MoveOnlyWrappedTypeEliminator.cpp

@@ -51,12 +51,6 @@ namespace {
 
 struct SILMoveOnlyWrappedTypeEliminatorVisitor
     : SILInstructionVisitor<SILMoveOnlyWrappedTypeEliminatorVisitor, bool> {
-  const llvm::SmallSetVector<SILArgument *, 8> &touchedArgs;
-
-  SILMoveOnlyWrappedTypeEliminatorVisitor(
-      const llvm::SmallSetVector<SILArgument *, 8> &touchedArgs)
-      : touchedArgs(touchedArgs) {}
-
   bool visitSILInstruction(SILInstruction *inst) {
     llvm::errs() << "Unhandled SIL Instruction: " << *inst;
     llvm_unreachable("error");
@@ -294,56 +288,76 @@ static bool isMoveOnlyWrappedTrivial(SILValue value) {
 }
 
 bool SILMoveOnlyWrappedTypeEliminator::process() {
-  bool madeChange = true;
+  bool madeChange = false;
 
-  llvm::SmallSetVector<SILArgument *, 8> touchedArgs;
   llvm::SmallSetVector<SILInstruction *, 8> touchedInsts;
 
+  // For each value whose type is move-only wrapped:
+  // - rewrite the value's type
+  // - record its users for later visitation
+  auto visitValue = [&touchedInsts, fn = fn,
+                     trivialOnly = trivialOnly](SILValue value) -> bool {
+    if (!value->getType().hasAnyMoveOnlyWrapping(fn))
+      return false;
+
+    // If we are looking at trivial only, skip non-trivial function args.
+    if (trivialOnly && !isMoveOnlyWrappedTrivial(value))
+      return false;
+
+    for (auto *use : value->getNonTypeDependentUses())
+      touchedInsts.insert(use->getUser());
+
+    if (isa<SILUndef>(value))
+      value->replaceAllUsesWith(
+          SILUndef::get(fn, value->getType().removingAnyMoveOnlyWrapping(fn)));
+    else
+      value->unsafelyEliminateMoveOnlyWrapper(fn);
+
+    return true;
+  };
+
   for (auto &bb : *fn) {
     for (auto *arg : bb.getArguments()) {
-      if (!arg->getType().isMoveOnlyWrapped() &&
-          !arg->getType().isBoxedMoveOnlyWrappedType(fn))
+      bool relevant = visitValue(arg);
+      if (!relevant)
         continue;
 
-      // If we are looking at trivial only, skip non-trivial function args.
-      if (trivialOnly &&
-          !arg->getType().removingMoveOnlyWrapper().isTrivial(*fn))
-        continue;
-
-      arg->unsafelyEliminateMoveOnlyWrapper(fn);
-
       // If our new type is trivial, convert the arguments ownership to
       // None. Otherwise, preserve the ownership kind of the argument.
       if (arg->getType().isTrivial(*fn))
         arg->setOwnershipKind(OwnershipKind::None);
-      touchedArgs.insert(arg);
-      for (auto *use : arg->getNonTypeDependentUses())
-        touchedInsts.insert(use->getUser());
+
+      madeChange = true;
     }
 
     for (auto &ii : bb) {
-      for (SILValue v : ii.getResults()) {
-        if (!v->getType().isMoveOnlyWrapped() &&
-            !v->getType().isBoxedMoveOnlyWrappedType(fn))
-          continue;
-
-        if (trivialOnly &&
-            !isMoveOnlyWrappedTrivial(v))
+      bool touched = false;
+      for (SILValue value : ii.getResults()) {
+        bool relevant = visitValue(value);
+        if (!relevant)
           continue;
 
-        v->unsafelyEliminateMoveOnlyWrapper(fn);
-        touchedInsts.insert(&ii);
-
-        // Add all users as well. This ensures we visit things like
-        // destroy_value and end_borrow.
-        for (auto *use : v->getNonTypeDependentUses())
-          touchedInsts.insert(use->getUser());
-        madeChange = true;
+        touched = true;
       }
+      if (!touched)
+        continue;
+      touchedInsts.insert(&ii);
+
+      madeChange = true;
     }
   }
+  // SILFunction::undefValues may grow during the loop.
+  SmallVector<std::pair<SILType, SILUndef *>, 4> originalUndefs(
+      fn->getUndefValues());
+  for (auto pair : originalUndefs) {
+    bool relevant = visitValue(pair.second);
+    if (!relevant)
+      continue;
+
+    madeChange = true;
+  }
 
-  SILMoveOnlyWrappedTypeEliminatorVisitor visitor(touchedArgs);
+  SILMoveOnlyWrappedTypeEliminatorVisitor visitor;
   while (!touchedInsts.empty()) {
     visitor.visit(touchedInsts.pop_back_val());
   }

test/SILOptimizer/moveonly_type_eliminator.sil

@@ -562,3 +562,19 @@ bb3(%result : @owned $@moveOnly FakeOptional<Klass>):
   %result2 = moveonlywrapper_to_copyable [owned] %result : $@moveOnly FakeOptional<Klass>
   return %result2 : $FakeOptional<Klass>
 }
+
+// CHECK-LABEL: sil [ossa] @debug_value_undef : {{.*}} {
+// CHECK:         debug_value [moveable_value_debuginfo] undef : $*Klass, var, name "s"
+// CHECK-LABEL: } // end sil function 'debug_value_undef'
+sil  [ossa] @debug_value_undef : $@convention(thin) (@owned Klass) -> () {
+bb0(%x : @owned $Klass):
+  %addr = alloc_stack $@moveOnly Klass
+  %unwrapped_addr = moveonlywrapper_to_copyable_addr %addr : $*@moveOnly Klass
+  store %x to [init] %unwrapped_addr : $*Klass
+  debug_value %addr : $*@moveOnly Klass, var, name "s", argno 1, expr op_deref
+  destroy_addr %addr : $*@moveOnly Klass
+  debug_value undef : $*@moveOnly Klass, var, name "s", argno 1, expr op_deref
+  dealloc_stack %addr : $*@moveOnly Klass
+  %retval = tuple ()
+  return %retval : $()
+}