[DI] TypeWrapper: Adjust immutability check to handle _storage as $Storage

`_storage` is used primary to support member-by-member initialization
of `$Storage`, so we need to refer to `$Storage` to determine whether
a partial field of `_storage` is `let` or not.

Author: xedin

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -468,22 +468,35 @@ DIMemoryObjectInfo::getPathStringToElement(unsigned Element,
 
 /// If the specified value is a 'let' property in an initializer, return true.
 bool DIMemoryObjectInfo::isElementLetProperty(unsigned Element) const {
-  // If we aren't representing 'self' in a non-delegating initializer, then we
-  // can't have 'let' properties.
-  if (!isNonDelegatingInit())
-    return IsLet;
-
-  auto &Module = MemoryInst->getModule();
+  NullablePtr<NominalTypeDecl> NTD;
+
+  // If this is an element of a `_storage` tuple, we need to
+  // check the `$Storage` to determine whether underlying storage
+  // backing element is immutable.
+  if (auto *storageVar = getAsTypeWrapperLocalStorageVar()) {
+    auto *wrappedType = cast<NominalTypeDecl>(
+        storageVar->getDeclContext()->getInnermostTypeContext());
+    assert(wrappedType && "_storage reference without type wrapper");
+    NTD = wrappedType->getTypeWrapperStorageDecl();
+  } else {
+    // If we aren't representing 'self' in a non-delegating initializer, then we
+    // can't have 'let' properties.
+    if (!isNonDelegatingInit())
+      return IsLet;
+
+    NTD = MemorySILType.getNominalOrBoundGenericNominal();
+  }
 
-  auto *NTD = MemorySILType.getNominalOrBoundGenericNominal();
   if (!NTD) {
     // Otherwise, we miscounted elements?
     assert(Element == 0 && "Element count problem");
     return false;
   }
 
+  auto &Module = MemoryInst->getModule();
+
   auto expansionContext = TypeExpansionContext(*MemoryInst->getFunction());
-  for (auto *VD : NTD->getStoredProperties()) {
+  for (auto *VD : NTD.get()->getStoredProperties()) {
     auto FieldType = MemorySILType.getFieldType(VD, Module, expansionContext);
     unsigned NumFieldElements =
         getElementCountRec(expansionContext, Module, FieldType, false);

test/SILOptimizer/type_wrapper_definite_init_diagnostics.swift

@@ -0,0 +1,91 @@
+// RUN: %target-swift-frontend -enable-experimental-feature TypeWrappers -enable-copy-propagation=requested-passes-only -emit-sil -primary-file %s -o /dev/null -verify
+
+// REQUIRES: asserts
+
+@typeWrapper
+struct Wrapper<S> {
+  var underlying: S
+
+  init(memberwise: S) { self.underlying = memberwise }
+
+  subscript<V>(storageKeyPath path: KeyPath<S, V>) -> V {
+    get { underlying[keyPath: path] }
+  }
+
+  subscript<V>(storageKeyPath path: WritableKeyPath<S, V>) -> V {
+    get { underlying[keyPath: path] }
+    set { underlying[keyPath: path] = newValue }
+  }
+}
+
+do {
+  @Wrapper
+  struct ImmutableTest1 {
+    let x: Int
+
+    init(x: Int) {
+      self.x = x
+      self.x = 0 // expected-error {{immutable value 'self.x' may only be initialized once}}
+    }
+  }
+
+  @Wrapper
+  class ImmutableTest2 {
+    let a: Int
+    let b: String
+
+    init(a: Int, b: String) {
+      self.a = 0
+      self.a = a // expected-error {{immutable value 'self.a' may only be initialized once}}
+
+      self.b = b
+      self.b = "" // expected-error {{immutable value 'self.b' may only be initialized once}}
+    }
+  }
+
+  // FIXME: Diagnostics should mention `self.b` or `self.a` and not `self.$_storage`
+
+  @Wrapper
+  struct ImmutableTest3 {
+    let a: Int
+    let b: String
+
+    init() {
+      self.b = "" // Ok
+      self.a = 42 // Ok
+    }
+
+    init(a: Int = 42) {
+      self.a = a
+      print(self.a) // expected-error {{'self' used before all stored properties are initialized}} expected-note {{'self.$_storage' not initialized}}
+      self.b = ""
+    }
+
+    /* FIXME: This test is crashing in findFullInitializationPoints
+    init(b: String) {
+      self.b = b
+      print(self.a)
+    }
+    */
+
+    init(otherB: String?) {
+      self.a = 0
+
+      if let b = otherB {
+        self.b = b
+      }
+    } // expected-error {{return from initializer without initializing all stored properties}} expected-note {{'self.$_storage' not initialized}}
+
+    init(optB: String? = nil) {
+      if let optB {
+        self.b = optB
+        print(self.b) // expected-error {{'self' used before all stored properties are initialized}} expected-note {{'self.$_storage' not initialized}}
+      } else {
+        self.b = ""
+        print(self.b) // expected-error {{'self' used before all stored properties are initialized}} expected-note {{'self.$_storage' not initialized}}
+      }
+
+      self.a = 0
+    }
+  }
+}