[AST] Store original lazy VarDecl for backing storage

Store the original VarDecl in the same map we use for tracking the
original wrapper var for property wrappers. This will allow us to
use the same logic to determine the original var for both.

Author: hamishknight

include/swift/AST/Decl.h

@@ -6781,9 +6781,14 @@ class VarDecl : public AbstractStorageDecl {
   bool isLazyStorageProperty() const {
     return Bits.VarDecl.IsLazyStorageProperty;
   }
-  void setLazyStorageProperty(bool IsLazyStorage) {
-    Bits.VarDecl.IsLazyStorageProperty = IsLazyStorage;
-  }
+
+  /// Mark this VarDecl as backing storage for a the lazy variable `VD`.
+  void setLazyStorageFor(VarDecl *VD);
+
+  /// For a backing storage variable for either a property wrapper or `lazy`
+  /// variable, retrieves the original declared variable. Otherwise returns
+  /// \c nullptr.
+  VarDecl *getOriginalVarForBackingStorage() const;
 
   /// Retrieve the backing storage property for a lazy property.
   VarDecl *getLazyStorageProperty() const;

lib/AST/ASTContext.cpp

@@ -519,9 +519,9 @@ struct ASTContext::Implementation {
   /// Mapping from property declarations to the backing variable types.
   llvm::DenseMap<const VarDecl *, Type> PropertyWrapperBackingVarTypes;
 
-  /// A mapping from the backing storage of a property that has a wrapper
-  /// to the original property with the wrapper.
-  llvm::DenseMap<const VarDecl *, VarDecl *> OriginalWrappedProperties;
+  /// A mapping from the backing storage of a property that has a wrapper or
+  /// is `lazy` to the original property.
+  llvm::DenseMap<const VarDecl *, VarDecl *> OriginalVarsForBackingStorage;
 
   /// The builtin initializer witness for a literal. Used when building
   /// LiteralExprs in fully-checked AST.
@@ -917,7 +917,7 @@ void ASTContext::Implementation::dump(llvm::raw_ostream &os) const {
   SIZE_AND_BYTES(DefaultAssociatedConformanceWitnesses);
   SIZE_AND_BYTES(DefaultTypeRequestCaches);
   SIZE_AND_BYTES(PropertyWrapperBackingVarTypes);
-  SIZE_AND_BYTES(OriginalWrappedProperties);
+  SIZE_AND_BYTES(OriginalVarsForBackingStorage);
   SIZE_AND_BYTES(BuiltinInitWitness);
   SIZE_AND_BYTES(OriginalBodySourceRanges);
   SIZE_AND_BYTES(NextMacroDiscriminator);
@@ -7066,9 +7066,8 @@ VarDecl *VarDecl::getOriginalWrappedProperty(
   if (!Bits.VarDecl.IsPropertyWrapperBackingProperty)
     return nullptr;
 
-  ASTContext &ctx = getASTContext();
-  assert(ctx.getImpl().OriginalWrappedProperties.count(this) > 0);
-  auto original = ctx.getImpl().OriginalWrappedProperties[this];
+  auto *original = getOriginalVarForBackingStorage();
+  ASSERT(original);
   if (!kind)
     return original;
 
@@ -7086,8 +7085,24 @@ VarDecl *VarDecl::getOriginalWrappedProperty(
 void VarDecl::setOriginalWrappedProperty(VarDecl *originalProperty) {
   Bits.VarDecl.IsPropertyWrapperBackingProperty = true;
   ASTContext &ctx = getASTContext();
-  assert(ctx.getImpl().OriginalWrappedProperties.count(this) == 0);
-  ctx.getImpl().OriginalWrappedProperties[this] = originalProperty;
+  assert(ctx.getImpl().OriginalVarsForBackingStorage.count(this) == 0);
+  ctx.getImpl().OriginalVarsForBackingStorage[this] = originalProperty;
+}
+
+void VarDecl::setLazyStorageFor(VarDecl *VD) {
+  Bits.VarDecl.IsLazyStorageProperty = true;
+  ASTContext &ctx = getASTContext();
+  ASSERT(ctx.getImpl().OriginalVarsForBackingStorage.count(this) == 0);
+  ctx.getImpl().OriginalVarsForBackingStorage[this] = VD;
+}
+
+VarDecl *VarDecl::getOriginalVarForBackingStorage() const {
+  const auto &map = getASTContext().getImpl().OriginalVarsForBackingStorage;
+  auto iter = map.find(this);
+  if (iter == map.end())
+    return nullptr;
+
+  return iter->second;
 }
 
 #ifndef NDEBUG

lib/AST/Decl.cpp

@@ -8455,20 +8455,8 @@ bool VarDecl::isActorSelf() const {
   return nominal && nominal->isActor();
 }
 
-/// Whether the given variable is the backing storage property for
-/// a declared property that is either `lazy` or has an attached
-/// property wrapper.
-static bool isBackingStorageForDeclaredProperty(const VarDecl *var) {
-  if (var->isLazyStorageProperty())
-    return true;
-
-  if (var->getOriginalWrappedProperty())
-    return true;
-
-  return false;
-}
-
-/// Whether the given variable is a declared property that has separate backing storage.
+/// Whether the given variable is a declared property that has separate backing
+/// storage.
 static bool isDeclaredPropertyWithBackingStorage(const VarDecl *var) {
   if (var->getAttrs().hasAttribute<LazyAttr>())
     return true;
@@ -8488,8 +8476,7 @@ bool VarDecl::isMemberwiseInitialized(bool preferDeclaredProperties) const {
   // If this is a stored property, and not a backing property in a case where
   // we only want to see the declared properties, it can be memberwise
   // initialized.
-  if (hasStorage() && preferDeclaredProperties &&
-      isBackingStorageForDeclaredProperty(this))
+  if (preferDeclaredProperties && getOriginalVarForBackingStorage())
     return false;
 
   // If this is a computed property with `init` accessor, it's

lib/Sema/TypeCheckStorage.cpp

@@ -3031,7 +3031,7 @@ LazyStoragePropertyRequest::evaluate(Evaluator &evaluator,
                                         VD->getLoc(), StorageName,
                                         VD->getDeclContext());
   Storage->setInterfaceType(StorageInterfaceTy);
-  Storage->setLazyStorageProperty(true);
+  Storage->setLazyStorageFor(VD);
   Storage->setUserAccessible(false);
 
   // The storage is implicit and private.

lib/Serialization/Deserialization.cpp

@@ -4002,7 +4002,6 @@ class DeclDeserializer {
     bool isImplicit, isObjC, isStatic;
     uint8_t rawIntroducer;
     bool isGetterMutating, isSetterMutating;
-    bool isLazyStorageProperty;
     bool isTopLevelGlobal;
     DeclID lazyStorageID;
     unsigned numAccessors, numBackingProperties;
@@ -4018,7 +4017,6 @@ class DeclDeserializer {
     decls_block::VarLayout::readRecord(scratch, nameID, contextID,
                                        isImplicit, isObjC, isStatic, rawIntroducer,
                                        isGetterMutating, isSetterMutating,
-                                       isLazyStorageProperty,
                                        isTopLevelGlobal,
                                        lazyStorageID,
                                        opaqueReadOwnership,
@@ -4169,9 +4167,9 @@ class DeclDeserializer {
       VarDecl *storage = cast<VarDecl>(lazyStorageDecl.get());
       ctx.evaluator.cacheOutput(
           LazyStoragePropertyRequest{var}, std::move(storage));
+      storage->setLazyStorageFor(var);
     }
 
-    var->setLazyStorageProperty(isLazyStorageProperty);
     var->setTopLevelGlobal(isTopLevelGlobal);
 
     // If there are any backing properties, record them.

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 976; // Expand one type layout
+const uint16_t SWIFTMODULE_VERSION_MINOR = 977; // remove lazy storage bit
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -1730,7 +1730,6 @@ namespace decls_block {
     VarDeclIntroducerField,   // introducer
     BCFixed<1>,   // is getter mutating?
     BCFixed<1>,   // is setter mutating?
-    BCFixed<1>,   // is this the backing storage for a lazy property?
     BCFixed<1>,   // top level global?
     DeclIDField,  // if this is a lazy property, this is the backing storage
     OpaqueReadOwnershipField,   // opaque read ownership

lib/Serialization/Serialization.cpp

@@ -4768,7 +4768,6 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
                           rawIntroducer,
                           var->isGetterMutating(),
                           var->isSetterMutating(),
-                          var->isLazyStorageProperty(),
                           var->isTopLevelGlobal(),
                           S.addDeclRef(lazyStorage),
                           accessors.OpaqueReadOwnership,