Allow stored properties in extensions

include/swift/AST/DiagnosticsSema.def

@@ -1962,7 +1962,16 @@ ERROR(extension_metatype,none,
 ERROR(extension_placeholder,none,
       "cannot extend a type that contains placeholders", ())
 ERROR(extension_stored_property,none,
-      "extensions must not contain stored properties", ())
+      "stored property is not permitted in extension", ())
+ERROR(extension_stored_property_in_constrained_extension,none,
+      "stored property is not permitted in constrained extension", ())
+ERROR(extension_stored_property_on_protocol,none,
+      "stored property is not permitted in extension of protocol", ())
+ERROR(extension_stored_property_on_enum,none,
+      "stored property is not permitted in extension of enum", ())
+ERROR(extension_stored_property_in_different_file,none,
+      "only extensions defined in the same file as the extended type "
+      "are permitted to define stored properties", ())
 NOTE(extension_stored_property_fixit,none,
      "Remove '=' to make %0 a computed property", (Identifier))
 ERROR(extension_nongeneric_trailing_where,none,
@@ -6206,7 +6215,7 @@ ERROR(property_wrapper_not_single_var, none,
       "property wrapper can only apply to a single variable", ())
 ERROR(property_with_wrapper_in_bad_context,none,
       "%select{|non-static|non-static}1 %2 %0 declared inside "
-      "%select{a protocol|an extension|an enum}1 cannot have a wrapper",
+      "%select{a protocol|an enum}1 cannot have a wrapper",
       (Identifier, int, DescriptiveDeclKind))
 ERROR(property_with_wrapper_overrides,none,
       "property %0 with attached wrapper cannot override another property",

lib/AST/ASTVerifier.cpp

@@ -3697,7 +3697,7 @@ class Verifier : public ASTWalker {
       } else {
         llvm_unreachable("impossible parent node");
       }
-      
+
       if (!Ctx.SourceMgr.rangeContains(Enclosing, Current)) {
         Out << "child source range not contained within its parent: ";
         printEntity();

lib/IRGen/GenStruct.cpp

@@ -991,7 +991,12 @@ namespace {
     }
 
     SILType getType(VarDecl *field) {
-      assert(field->getDeclContext() == TheStruct->getAnyNominal());
+      auto DC = field->getDeclContext();
+      assert(DC == TheStruct->getAnyNominal()
+             // Stored properties in extensions within the same file are added
+             // as members to the defining type, so we shouldn't assert in that
+             // case
+             || DC->getContextKind() == DeclContextKind::ExtensionDecl);
       auto silType = SILType::getPrimitiveAddressType(TheStruct);
       return silType.getFieldType(
           field, IGM.getSILModule(),

lib/SIL/IR/SILDeclRef.cpp

@@ -563,8 +563,13 @@ SILLinkage SILDeclRef::getDefinitionLinkage() const {
 
   // Stored property initializers have linkage based on the access level of
   // their nominal.
-  if (isStoredPropertyInitializer())
-    decl = cast<NominalTypeDecl>(decl->getDeclContext());
+  if (isStoredPropertyInitializer()) {
+    if (auto extension = dyn_cast<ExtensionDecl>(decl->getDeclContext())) {
+      decl = extension->getExtendedNominal();
+    } else {
+      decl = cast<NominalTypeDecl>(decl->getDeclContext());
+    }
+  }
 
   // Compute the effective access level, taking e.g testable into consideration.
   auto effectiveAccess = decl->getEffectiveAccess();
@@ -790,7 +795,13 @@ IsSerialized_t SILDeclRef::isSerialized() const {
   // marked as @frozen.
   if (isStoredPropertyInitializer() || (isPropertyWrapperBackingInitializer() &&
                                         d->getDeclContext()->isTypeContext())) {
-    auto *nominal = cast<NominalTypeDecl>(d->getDeclContext());
+    NominalTypeDecl *nominal;
+    if (auto extension = dyn_cast<ExtensionDecl>(d->getDeclContext())) {
+      nominal = extension->getExtendedNominal();
+    } else {
+      nominal = cast<NominalTypeDecl>(d->getDeclContext());
+    }
+
     auto scope =
       nominal->getFormalAccessScope(/*useDC=*/nullptr,
                                     /*treatUsableFromInlineAsPublic=*/true);

lib/SIL/Verifier/SILVerifier.cpp

@@ -3264,7 +3264,12 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     require(EI->getField()->hasStorage(),
             "cannot get address of computed property with struct_element_addr");
 
-    require(EI->getField()->getDeclContext() == sd,
+    auto DC = EI->getField()->getDeclContext();
+    require(DC == sd
+                // Stored properties in extensions within the same file are
+                // added as members to the defining type, so we shouldn't assert
+                // in that case
+                || DC->getContextKind() == DeclContextKind::ExtensionDecl,
             "struct_element_addr field is not a member of the struct");
 
     if (EI->getModule().getStage() != SILStage::Lowered) {

lib/SILGen/SILGenConstructor.cpp

@@ -1061,11 +1061,18 @@ emitMemberInit(SILGenFunction &SGF, VarDecl *selfDecl, Pattern *pattern) {
       selfTy.getFieldType(field, SGF.SGM.M, SGF.getTypeExpansionContext());
     SILValue slot;
 
-    if (auto *structDecl = dyn_cast<StructDecl>(field->getDeclContext())) {
+    NominalTypeDecl *nominal;
+    if (auto extension = dyn_cast<ExtensionDecl>(field->getDeclContext())) {
+      nominal = extension->getExtendedNominal();
+    } else {
+      nominal = cast<NominalTypeDecl>(field->getDeclContext());
+    }
+
+    if (auto *structDecl = dyn_cast<StructDecl>(nominal)) {
       slot = SGF.B.createStructElementAddr(pattern, self.forward(SGF), field,
                                            fieldTy.getAddressType());
     } else {
-      assert(isa<ClassDecl>(field->getDeclContext()));
+      assert(isa<ClassDecl>(nominal));
       slot = SGF.B.createRefElementAddr(pattern, self.forward(SGF), field,
                                         fieldTy.getAddressType());
     }
@@ -1162,73 +1169,95 @@ void SILGenFunction::emitMemberInitializers(DeclContext *dc,
                                             NominalTypeDecl *nominal) {
   auto subs = getSubstitutionsForPropertyInitializer(dc, nominal);
 
+  // open question: is there a good "ordered set" type we could use
+  // here instead of a parallel array / unordered set?
+  llvm::SmallVector<PatternBindingDecl *, 4> patternDecls;
+  llvm::SmallDenseSet<PatternBindingDecl *, 4> seenPatterns;
+
   for (auto member : nominal->getMembers()) {
-    // Find instance pattern binding declarations that have initializers.
     if (auto pbd = dyn_cast<PatternBindingDecl>(member)) {
-      if (pbd->isStatic()) continue;
-
-      for (auto i : range(pbd->getNumPatternEntries())) {
-        auto init = pbd->getExecutableInit(i);
-        if (!init) continue;
+      patternDecls.push_back(pbd);
+      seenPatterns.insert(pbd);
+    }
+  }
 
-        auto *varPattern = pbd->getPattern(i);
+  // A type can also have stored properties that aren't direct members,
+  // e.g. if they were defined in an extension.
+  for (auto storedProperty : nominal->getStoredProperties()) {
+    if (auto pbd = storedProperty->getParentPatternBinding()) {
+      if (!seenPatterns.count(pbd)) {
+        patternDecls.push_back(pbd);
+        seenPatterns.insert(pbd);
+      }
+    }
+  }
 
-        // Cleanup after this initialization.
-        FullExpr scope(Cleanups, varPattern);
-
-        // Get the type of the initialization result, in terms
-        // of the constructor context's archetypes.
-        auto resultType = getInitializationTypeInContext(
-            pbd->getDeclContext(), dc, varPattern);
-        AbstractionPattern origType = resultType.first;
-        CanType substType = resultType.second;
-
-        // Figure out what we're initializing.
-        auto memberInit = emitMemberInit(*this, selfDecl, varPattern);
-
-        // This whole conversion thing is about eliminating the
-        // paired orig-to-subst subst-to-orig conversions that
-        // will happen if the storage is at a different abstraction
-        // level than the constructor. When emitApply() is used
-        // to call the stored property initializer, it naturally
-        // wants to convert the result back to the most substituted
-        // abstraction level. To undo this, we use a converting
-        // initialization and rely on the peephole that optimizes
-        // out the redundant conversion.
-        SILType loweredResultTy;
-        SILType loweredSubstTy;
-
-        // A converting initialization isn't necessary if the member is
-        // a property wrapper. Though the initial value can have a
-        // reabstractable type, the result of the initialization is
-        // always the property wrapper type, which is never reabstractable.
-        bool needsConvertingInit = false;
-        auto *singleVar = varPattern->getSingleVar();
-        if (!(singleVar && singleVar->getOriginalWrappedProperty())) {
-          loweredResultTy = getLoweredType(origType, substType);
-          loweredSubstTy = getLoweredType(substType);
-          needsConvertingInit = loweredResultTy != loweredSubstTy;
-        }
+  for (auto pbd : patternDecls) {
+    // Find instance pattern binding declarations that have initializers.
+    if (pbd->isStatic())
+      continue;
+
+    for (auto i : range(pbd->getNumPatternEntries())) {
+      auto init = pbd->getExecutableInit(i);
+      if (!init)
+        continue;
+
+      auto *varPattern = pbd->getPattern(i);
+
+      // Cleanup after this initialization.
+      FullExpr scope(Cleanups, varPattern);
+
+      // Get the type of the initialization result, in terms
+      // of the constructor context's archetypes.
+      auto resultType =
+          getInitializationTypeInContext(pbd->getDeclContext(), dc, varPattern);
+      AbstractionPattern origType = resultType.first;
+      CanType substType = resultType.second;
+
+      // Figure out what we're initializing.
+      auto memberInit = emitMemberInit(*this, selfDecl, varPattern);
+
+      // This whole conversion thing is about eliminating the
+      // paired orig-to-subst subst-to-orig conversions that
+      // will happen if the storage is at a different abstraction
+      // level than the constructor. When emitApply() is used
+      // to call the stored property initializer, it naturally
+      // wants to convert the result back to the most substituted
+      // abstraction level. To undo this, we use a converting
+      // initialization and rely on the peephole that optimizes
+      // out the redundant conversion.
+      SILType loweredResultTy;
+      SILType loweredSubstTy;
+
+      // A converting initialization isn't necessary if the member is
+      // a property wrapper. Though the initial value can have a
+      // reabstractable type, the result of the initialization is
+      // always the property wrapper type, which is never reabstractable.
+      bool needsConvertingInit = false;
+      auto *singleVar = varPattern->getSingleVar();
+      if (!(singleVar && singleVar->getOriginalWrappedProperty())) {
+        loweredResultTy = getLoweredType(origType, substType);
+        loweredSubstTy = getLoweredType(substType);
+        needsConvertingInit = loweredResultTy != loweredSubstTy;
+      }
 
-        if (needsConvertingInit) {
-          Conversion conversion = Conversion::getSubstToOrig(
-              origType, substType,
-              loweredResultTy);
+      if (needsConvertingInit) {
+        Conversion conversion =
+            Conversion::getSubstToOrig(origType, substType, loweredResultTy);
 
-          ConvertingInitialization convertingInit(conversion,
-                                                  SGFContext(memberInit.get()));
+        ConvertingInitialization convertingInit(conversion,
+                                                SGFContext(memberInit.get()));
 
-          emitAndStoreInitialValueInto(*this, varPattern, pbd, i, subs,
-                                       origType, substType, &convertingInit);
+        emitAndStoreInitialValueInto(*this, varPattern, pbd, i, subs, origType,
+                                     substType, &convertingInit);
 
-          auto finalValue = convertingInit.finishEmission(
-              *this, varPattern, ManagedValue::forInContext());
-          if (!finalValue.isInContext())
-            finalValue.forwardInto(*this, varPattern, memberInit.get());
-        } else {
-          emitAndStoreInitialValueInto(*this, varPattern, pbd, i, subs,
-                                       origType, substType, memberInit.get());
-        }
+        auto finalValue = convertingInit.finishEmission(
+            *this, varPattern, ManagedValue::forInContext());
+        if (!finalValue.isInContext())
+          finalValue.forwardInto(*this, varPattern, memberInit.get());
+      } else {
+        emitAndStoreInitialValueInto(*this, varPattern, pbd, i, subs, origType,
+                                     substType, memberInit.get());
       }
     }
   }

lib/SILGen/SILGenType.cpp

@@ -1076,6 +1076,7 @@ class SILGenType : public TypeMemberVisitor<SILGenType> {
 public:
   SILGenModule &SGM;
   NominalTypeDecl *theType;
+  llvm::SmallDenseSet<VarDecl *> visitedVarDecls;
 
   SILGenType(SILGenModule &SGM, NominalTypeDecl *theType)
     : SGM(SGM), theType(theType) {}
@@ -1087,6 +1088,14 @@ class SILGenType : public TypeMemberVisitor<SILGenType> {
     for (Decl *member : theType->getABIMembers())
       visit(member);
 
+    // Types can have stored properties that aren't included in
+    // `getABIMembers()`, so we have to handle those manually
+    for (VarDecl *VD : theType->getStoredProperties()) {
+      if (!visitedVarDecls.count(VD)) {
+        visit(VD);
+      }
+    }
+
     // Build a vtable if this is a class.
     if (auto theClass = dyn_cast<ClassDecl>(theType)) {
       SILGenVTable genVTable(SGM, theClass);
@@ -1176,6 +1185,8 @@ class SILGenType : public TypeMemberVisitor<SILGenType> {
   }
 
   void visitVarDecl(VarDecl *vd) {
+    visitedVarDecls.insert(vd);
+
     // Collect global variables for static properties.
     // FIXME: We can't statically emit a global variable for generic properties.
     if (vd->isStatic() && vd->hasStorage()) {
@@ -1305,7 +1316,13 @@ class SILGenExtension : public TypeMemberVisitor<SILGenExtension> {
     // Emit initializers for static variables.
     for (auto i : range(pd->getNumPatternEntries())) {
       if (pd->getExecutableInit(i)) {
-        assert(pd->isStatic() && "stored property in extension?!");
+        // Ignore any stored properties, since at codegen time
+        // these are handled as if they were defined in the type
+        // being extended rather than the extension itself
+        if (!pd->isStatic()) {
+          continue;
+        }
+
         SGM.emitGlobalInitialization(pd, i);
       }
     }
@@ -1315,8 +1332,14 @@ class SILGenExtension : public TypeMemberVisitor<SILGenExtension> {
     if (vd->hasStorage()) {
       bool hasDidSetOrWillSetDynamicReplacement =
           vd->hasDidSetOrWillSetDynamicReplacement();
-      assert((vd->isStatic() || hasDidSetOrWillSetDynamicReplacement) &&
-             "stored property in extension?!");
+
+      // Ignore any stored properties, since at codegen time
+      // these are handled as if they were defined in the type
+      // being extended rather than the extension itself
+      if (!(vd->isStatic() || hasDidSetOrWillSetDynamicReplacement)) {
+        return;
+      }
+
       if (!hasDidSetOrWillSetDynamicReplacement) {
         emitTypeMemberGlobalVariable(SGM, vd);
         visitAccessors(vd);

lib/Sema/CodeSynthesis.cpp

@@ -339,11 +339,28 @@ static ConstructorDecl *createImplicitConstructor(NominalTypeDecl *decl,
   if (ICK == ImplicitConstructorKind::Memberwise) {
     assert(isa<StructDecl>(decl) && "Only struct have memberwise constructor");
 
+    // open question: is there a good "ordered set" type we could use
+    // here instead of a parallel array / unordered set?
+    llvm::SmallVector<VarDecl *, 4> varDecls;
+    llvm::SmallDenseSet<VarDecl *, 4> seenDecls;
+
     for (auto member : decl->getMembers()) {
-      auto var = dyn_cast<VarDecl>(member);
-      if (!var)
-        continue;
+      if (auto var = dyn_cast<VarDecl>(member)) {
+        varDecls.push_back(var);
+        seenDecls.insert(var);
+      }
+    }
 
+    // A type can also have stored properties that aren't direct members,
+    // e.g. if they were defined in an extension.
+    for (auto storedDecl : decl->getStoredProperties()) {
+      if (!seenDecls.count(storedDecl)) {
+        varDecls.push_back(storedDecl);
+        seenDecls.insert(storedDecl);
+      }
+    }
+
+    for (auto var : varDecls) {
       if (!var->isMemberwiseInitialized(/*preferDeclaredProperties=*/true))
         continue;
 
@@ -1396,6 +1413,15 @@ HasMemberwiseInitRequest::evaluate(Evaluator &evaluator,
         return true;
     }
   }
+
+  // Also check the list of stored properties, which may include
+  // properties defined in extensions and not included as direct members
+  for (auto var : decl->getStoredProperties()) {
+    if (var->isMemberwiseInitialized(/*preferDeclaredProperties=*/true)) {
+      return true;
+    }
+  }
+
   return false;
 }