[SILOptimizer] TypeWrappers: Implement self.$storage initialization injection

Inject a call to $Storage.init(...) which is then used to initialize
type wrapper property `$storage` via `self.$storage = <TypeWrapper>(memberwise: <storage>)`
call at each point where `_storage` becomes fully initialized.

Author: xedin

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -459,6 +459,13 @@ namespace {
     /// corresponds to `self`.
     void injectActorHops();
 
+    /// Injects `self.$storage = .init(memberwise: $Storage(...))`
+    /// assignment instructions into the function after each point
+    /// where `_storage` becomes fully initialized via `assign_by_wrapper`.
+    /// This is only necessary only for user-defined initializers of a
+    /// type wrapped types.
+    void injectTypeWrapperStorageInitalization();
+
     void emitSelfConsumedDiagnostic(SILInstruction *Inst);
 
     LiveOutBlockState &getBlockInfo(SILBasicBlock *BB) {
@@ -1070,6 +1077,284 @@ void LifetimeChecker::injectActorHops() {
     injectExecutorHopAfter(point);
 }
 
+void LifetimeChecker::injectTypeWrapperStorageInitalization() {
+  auto *storageVar = TheMemory.getAsTypeWrapperLocalStorageVar();
+  if (!storageVar)
+    return;
+
+  SmallVector<SILInstruction *> points;
+  findFullInitializationPoints(points);
+
+  // `_storage` has not been initialized at all.
+  if (points.empty())
+    return;
+
+  auto *ctor = cast<ConstructorDecl>(storageVar->getDeclContext()->getAsDecl());
+  auto *parentType = ctor->getDeclContext()->getSelfNominalTypeDecl();
+  auto *storageDecl = parentType->getTypeWrapperStorageDecl();
+
+  for (auto *point : points) {
+    SILBuilderWithScope::insertAfter(point, [&](SILBuilder &b) {
+      SILLocation loc = SILLocation::invalid().asAutoGenerated();
+
+      SmallVector<SILValue, 4> allocations;
+
+      auto allocStack = [&](SILType type) -> SILValue {
+        auto value = b.createAllocStack(loc, type);
+        allocations.push_back(value);
+        return value;
+      };
+
+      auto createInitRef = [&](ConstructorDecl *ctor) -> SILValue {
+        auto *init = F.getModule().lookUpFunction(SILDeclRef(ctor));
+        assert(init);
+        return b.createFunctionRef(loc, init);
+      };
+
+      // The instance type of $Storage declaration.
+      auto storageType = F.getLoweredType(
+          ctor->mapTypeIntoContext(storageDecl->getDeclaredInterfaceType()));
+
+      // Argument value to use in call to <TypeWrapper>.init(memberwise:)
+      SILValue storageObj = allocStack(storageType);
+
+      // let storageObj = $Storage(<destructured _storage tuple>)
+      {
+        SILValue localStorageVar = TheMemory.getUninitializedValue();
+
+        SmallVector<SILValue> storageInitArgs;
+
+        SILValue storageInitRef;
+        if (auto *memberwiseInit = storageDecl->getMemberwiseInitializer()) {
+          storageInitRef = createInitRef(memberwiseInit);
+        } else {
+          assert(storageDecl->hasDefaultInitializer());
+          storageInitRef = createInitRef(storageDecl->getDefaultInitializer());
+        }
+
+        CanSILFunctionType storageInitTy =
+            storageInitRef->getType().castTo<SILFunctionType>();
+        SILFunctionConventions convention(storageInitTy, F.getModule());
+
+        // If `.init` produces result indirectly, let's load it
+        // into the prepared buffer.
+        if (convention.hasIndirectSILResults()) {
+          assert(convention.getNumIndirectSILResults() == 1);
+          storageInitArgs.push_back(storageObj);
+        }
+
+        // Prepare arguments for $Storage.init(...) call. `_storage`
+        // tuple needs to be flattened and its elements have to be
+        // either copied or loaded based on a particular argument convention.
+        {
+          auto localStorageRef =
+              b.createBeginAccess(loc, localStorageVar, SILAccessKind::Read,
+                                  SILAccessEnforcement::Unsafe,
+                                  /*noNestedConflict=*/false,
+                                  /*fromBuiltin=*/false);
+
+          // There could be only one indirect result her - $Storage,
+          // which is also verified above, actual arguments start
+          // at 1 in such case.
+          unsigned argIdx = convention.getNumIndirectSILResults();
+
+          std::function<void(SILValue, SmallVectorImpl<SILValue> &)>
+            prepareArguments =
+            [&](SILValue val, SmallVectorImpl<SILValue> &results) {
+              // Destructure the tuple into individual elements recursively.
+              if (auto tupleType = val->getType().getAs<TupleType>()) {
+                for (unsigned i = 0, n = tupleType->getNumElements();
+                     i != n; ++i) {
+                  SILValue elt = b.createTupleElementAddr(loc, val, i);
+                  prepareArguments(elt, results);
+                }
+                return;
+              }
+
+              switch (convention.getSILArgumentConvention(argIdx)) {
+              case SILArgumentConvention::Indirect_In:
+              case SILArgumentConvention::Indirect_In_Constant: {
+                // The only way to load opaque type is to allocate a temporary
+                // variable on the stack for it and initialize from the element
+                // address.
+                SILValue arg = allocStack(val->getType());
+                b.createCopyAddr(loc, val, arg, IsNotTake, IsInitialization);
+                results.push_back(arg);
+                break;
+              }
+
+              case SILArgumentConvention::Indirect_In_Guaranteed:
+                // The argument is +0, so we can use the address of the
+                // element as an argument.
+                results.push_back(val);
+                break;
+
+              case SILArgumentConvention::Direct_Owned: {
+                // Copy the value out at +1.
+                results.push_back(b.createTrivialLoadOr(
+                    loc, val, LoadOwnershipQualifier::Copy));
+                break;
+              }
+
+              case SILArgumentConvention::Direct_Unowned:
+              case SILArgumentConvention::Direct_Guaranteed: {
+                results.push_back(b.createTrivialLoadOr(
+                    loc, val, LoadOwnershipQualifier::Take));
+                break;
+              }
+
+              case SILArgumentConvention::Indirect_Out:
+                llvm_unreachable("indirect result is handled separately");
+
+              case SILArgumentConvention::Indirect_Inout:
+              case SILArgumentConvention::Indirect_InoutAliasable:
+                llvm_unreachable("inout arguments are not supported by $Storage.");
+              }
+
+              ++argIdx;
+            };
+
+          prepareArguments(localStorageRef, storageInitArgs);
+
+          b.createEndAccess(loc, localStorageRef, /*abort=*/false);
+        }
+
+        // append $Storage.Type as the last argument to $Storage.init()
+        {
+          auto storageMetatypeType = F.getLoweredType(
+              ctor->mapTypeIntoContext(storageDecl->getInterfaceType()));
+          storageInitArgs.push_back(b.createMetatype(loc, storageMetatypeType));
+        }
+
+        SubstitutionMap storageInitSubs;
+        // Generic signature of $Storage is appropriate because
+        // the call is to either memberwise or default initializer
+        // which cannot introduce any additional generic parameters.
+        if (auto genericSig = storageDecl->getGenericSignature()) {
+          storageInitSubs = SubstitutionMap::get(
+              genericSig,
+              [&](SubstitutableType *type) {
+                return ctor->mapTypeIntoContext(type);
+              },
+              LookUpConformanceInModule(
+                  ctor->getDeclContext()->getParentModule()));
+        }
+
+        // <storage> = $Storage.init(<flattened `_storage` elements>)
+        auto storageInitResult = b.createApply(
+            loc, storageInitRef, storageInitSubs, storageInitArgs);
+
+        // If result was indirect it would already be loaded into
+        // \c storageObj by the call itself.
+        if (!convention.hasIndirectSILResults())
+          b.createTrivialStoreOr(loc, storageInitResult, storageObj,
+                                 StoreOwnershipQualifier::Init);
+      }
+
+      // self.$storage = <TypeWrapper>(memberwise: storageObj))
+      {
+        bool isClass = isa<ClassDecl>(parentType);
+
+        auto typeWrapper = parentType->getTypeWrapper();
+        auto *typeWrapperInit = typeWrapper->getTypeWrapperInitializer();
+        SILValue typeWrapperInitRef = createInitRef(typeWrapperInit);
+
+        auto *self = TheMemory.findUninitializedSelfValue();
+
+        SILValue selfRef;
+        if (isClass) {
+          selfRef = b.emitBeginBorrowOperation(loc, self);
+        } else {
+          selfRef = b.createBeginAccess(loc, self, SILAccessKind::Modify,
+                                        SILAccessEnforcement::Static,
+                                        /*noNestedConflict=*/false,
+                                        /*fromBuiltin=*/false);
+        }
+
+        CanSILFunctionType wrapperInitTy =
+            typeWrapperInitRef->getType().castTo<SILFunctionType>();
+        SILFunctionConventions convention(wrapperInitTy, F.getModule());
+
+        // Argument is always passed indirectly because it's the
+        // generic parameter `S`.
+
+        // Reference to self.$storage
+        SILValue storagePropRef;
+        if (isClass) {
+          storagePropRef = b.createRefElementAddr(
+              loc, selfRef, parentType->getTypeWrapperProperty());
+        } else {
+          assert(isa<StructDecl>(parentType));
+          storagePropRef = b.createStructElementAddr(
+              loc, selfRef, parentType->getTypeWrapperProperty());
+        }
+
+        auto typeWrapperType =
+            b.createMetatype(loc, F.getLoweredType(MetatypeType::get(
+                                      storagePropRef->getType().getASTType())));
+
+        SmallVector<SILValue, 2> wrapperInitArgs;
+
+        Optional<SILValue> localWrapperObj;
+
+        if (convention.hasIndirectSILResults()) {
+          assert(convention.getNumIndirectSILResults() == 1);
+          localWrapperObj = allocStack(storagePropRef->getType());
+          wrapperInitArgs.push_back(*localWrapperObj);
+        }
+
+        wrapperInitArgs.push_back(storageObj);
+        wrapperInitArgs.push_back(typeWrapperType);
+
+        // <wrapper-var> = <TypeWrapper>.init(memberwise: tmpStorage)
+        auto wrapperInitResult = b.createApply(
+            loc, typeWrapperInitRef,
+            SubstitutionMap::get(typeWrapperInit->getGenericSignature(),
+                                 /*substitutions=*/{storageType.getASTType()},
+                                 /*conformances=*/{}),
+            wrapperInitArgs);
+
+        // self.$storage is a property access so it has to has to be wrapped
+        // in begin/end access instructions.
+        {
+          auto storagePropAccess =
+              b.createBeginAccess(loc, storagePropRef, SILAccessKind::Modify,
+                                  SILAccessEnforcement::Dynamic,
+                                  /*noNestedConflict=*/false,
+                                  /*fromBuiltin=*/false);
+
+          // self.$storage = <wrapper-var>
+          {
+            // If the result of `init` is indirect, let's just
+            // initialize the property with it.
+            if (localWrapperObj) {
+              b.createCopyAddr(loc, *localWrapperObj, storagePropAccess, IsTake,
+                               IsInitialization);
+            } else {
+              b.createAssign(loc, wrapperInitResult, storagePropAccess,
+                             AssignOwnershipQualifier::Init);
+            }
+          }
+
+          b.createEndAccess(loc, storagePropAccess, /*abort=*/false);
+        }
+
+        if (isClass) {
+          b.emitEndBorrowOperation(loc, selfRef);
+        } else {
+          b.createEndAccess(loc, selfRef, /*aborted=*/false);
+        }
+
+        // Dealloc all stack allocated data.
+        for (auto alloca = allocations.rbegin(); alloca != allocations.rend();
+             ++alloca) {
+          b.createDeallocStack(loc, *alloca);
+        }
+      }
+    });
+  }
+}
+
 void LifetimeChecker::doIt() {
   // With any escapes tallied up, we can work through all the uses, checking
   // for definitive initialization, promoting loads, rewriting assigns, and
@@ -1160,6 +1445,10 @@ void LifetimeChecker::doIt() {
   // Insert hop_to_executor instructions for actor initializers, if needed.
   injectActorHops();
 
+  // Insert `self.$storage` initialization for user-defined initializers
+  // declared in a type wrapped type.
+  injectTypeWrapperStorageInitalization();
+
   // 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.