[DI] Process type wrapper _storage before any other initialization

`_storage` is special because it emits assignments to compiler
synthesized stored property `$storage`, without that `self`
cannot be fully initialized.

Author: xedin

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -539,25 +539,8 @@ ConstructorDecl *DIMemoryObjectInfo::getActorInitSelf() const {
 }
 
 VarDecl *DIMemoryObjectInfo::getAsTypeWrapperLocalStorageVar() const {
-  if (!MemoryInst->isVar())
-    return nullptr;
-
-  if (!isa<TupleType>(getASTType()))
-    return nullptr;
-
-  auto *DC = MemoryInst->getFunction()->getDeclContext();
-  if (!DC)
-    return nullptr;
-
-  auto *ctor = dyn_cast_or_null<ConstructorDecl>(DC->getAsDecl());
-  if (!ctor || ctor->isImplicit())
-    return nullptr;
-
-  if (auto *var = getLoc().getAsASTNode<VarDecl>()) {
-    auto &ctx = var->getASTContext();
-    if (var->isImplicit() && var->getName() == ctx.Id_localStorageVar)
-      return var;
-  }
+  if (isTypeWrapperLocalStorageVar(getFunction(), MemoryInst))
+    return getLoc().getAsASTNode<VarDecl>();
 
   return nullptr;
 }
@@ -2015,3 +1998,32 @@ void swift::ownership::collectDIElementUsesFrom(
   collector.collectFrom(MemoryInfo.getUninitializedValue(),
                         /*collectDestroysOfContainer*/ true);
 }
+
+bool swift::ownership::canHaveTypeWrapperLocalStorageVar(SILFunction &F) {
+  auto *DC = F.getDeclContext();
+  if (!DC)
+    return false;
+
+  auto *ctor = dyn_cast_or_null<ConstructorDecl>(DC->getAsDecl());
+  if (!ctor || ctor->isImplicit() || !ctor->isDesignatedInit())
+    return false;
+
+  auto *parentType = ctor->getDeclContext()->getSelfNominalTypeDecl();
+  return parentType && parentType->hasTypeWrapper();
+}
+
+bool swift::ownership::isTypeWrapperLocalStorageVar(
+    SILFunction &F, MarkUninitializedInst *Inst) {
+  if (!Inst->isVar())
+    return false;
+
+  if (!canHaveTypeWrapperLocalStorageVar(F))
+    return false;
+
+  if (auto *var = Inst->getLoc().getAsASTNode<VarDecl>()) {
+    auto &ctx = var->getASTContext();
+    return var->isImplicit() && var->getName() == ctx.Id_localStorageVar;
+  }
+
+  return false;
+}

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h

@@ -357,6 +357,15 @@ struct DIElementUseInfo {
 void collectDIElementUsesFrom(const DIMemoryObjectInfo &MemoryInfo,
                               DIElementUseInfo &UseInfo);
 
+/// Check whether given function is a user-defined initializer of a
+/// type wrapped type.
+bool canHaveTypeWrapperLocalStorageVar(SILFunction &F);
+
+/// Check whether this instruction represents `_storage` variable
+/// injected by type-checker into user-defined designated initializer
+/// of a type wrapped type.
+bool isTypeWrapperLocalStorageVar(SILFunction &F, MarkUninitializedInst *Inst);
+
 } // end namespace ownership
 } // end namespace swift
 

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -1160,7 +1160,6 @@ void LifetimeChecker::doIt() {
   // Insert hop_to_executor instructions for actor initializers, if needed.
   injectActorHops();
 
-
   // If the memory object had any non-trivial stores that are init or assign
   // based on the control flow path reaching them, then insert dynamic control
   // logic and CFG diamonds to handle this.
@@ -3542,6 +3541,24 @@ static void processMemoryObject(MarkUninitializedInst *I,
   LifetimeChecker(MemInfo, UseInfo, blockStates).doIt();
 }
 
+static MarkUninitializedInst *findLocalTypeWrapperStorageVar(SILFunction &F) {
+  if (!canHaveTypeWrapperLocalStorageVar(F))
+    return nullptr;
+
+  auto BB = F.getEntryBlock();
+  if (!BB)
+    return nullptr;
+
+  for (auto &I : *BB) {
+    SILInstruction *Inst = &I;
+    auto *MUI = dyn_cast<MarkUninitializedInst>(Inst);
+    if (MUI && isTypeWrapperLocalStorageVar(F, MUI))
+      return MUI;
+  }
+
+  return nullptr;
+}
+
 /// Check that all memory objects that require initialization before use are
 /// properly set and transform the code as required for flow-sensitive
 /// properties.
@@ -3552,6 +3569,17 @@ static bool checkDefiniteInitialization(SILFunction &Fn) {
 
   BlockStates blockStates(&Fn);
 
+  // Special handling of _storage variable introduced by type wrapper.
+  // It has to be checked first because it injects initialization of
+  // `self.$storage`.
+  if (auto *storageVar = findLocalTypeWrapperStorageVar(Fn)) {
+    // Then process the memory object.
+    processMemoryObject(storageVar, blockStates);
+    storageVar->replaceAllUsesWith(storageVar->getOperand());
+    storageVar->eraseFromParent();
+    Changed = true;
+  }
+
   for (auto &BB : Fn) {
     for (auto I = BB.begin(), E = BB.end(); I != E;) {
       SILInstruction *Inst = &*I;