[AST] Add a way to retrieve $Storage decl from a type wrapped type

Author: xedin

include/swift/AST/Decl.h

@@ -3775,6 +3775,11 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// applicable property access routing).
   VarDecl *getTypeWrapperProperty() const;
 
+  /// If this declaration has a type wrapper, return `$Storage`
+  /// declaration that contains all the stored properties managed
+  /// by the wrapper.
+  NominalTypeDecl *getTypeWrapperStorageDecl() const;
+
   /// Get an initializer that could be used to instantiate a
   /// type wrapped type.
   ConstructorDecl *getTypeWrapperInitializer() const;

lib/Sema/TypeCheckTypeWrapper.cpp

@@ -158,6 +158,12 @@ ConstructorDecl *NominalTypeDecl::getTypeWrapperInitializer() const {
                            nullptr);
 }
 
+NominalTypeDecl *NominalTypeDecl::getTypeWrapperStorageDecl() const {
+  auto *mutableSelf = const_cast<NominalTypeDecl *>(this);
+  return evaluateOrDefault(getASTContext().evaluator,
+                           GetTypeWrapperStorage{mutableSelf}, nullptr);
+}
+
 NominalTypeDecl *
 GetTypeWrapperStorage::evaluate(Evaluator &evaluator,
                                 NominalTypeDecl *parent) const {
@@ -190,8 +196,8 @@ GetTypeWrapperProperty::evaluate(Evaluator &evaluator,
   if (!typeWrapper)
     return nullptr;
 
-  auto *storage =
-      evaluateOrDefault(ctx.evaluator, GetTypeWrapperStorage{parent}, nullptr);
+  auto *storage = parent->getTypeWrapperStorageDecl();
+  assert(storage);
 
   auto *typeWrapperType =
       evaluateOrDefault(ctx.evaluator, GetTypeWrapperType{parent}, Type())
@@ -217,10 +223,7 @@ VarDecl *GetTypeWrapperStorageForProperty::evaluate(Evaluator &evaluator,
   if (!property->isAccessedViaTypeWrapper())
     return nullptr;
 
-  auto &ctx = wrappedType->getASTContext();
-
-  auto *storage = evaluateOrDefault(
-      ctx.evaluator, GetTypeWrapperStorage{wrappedType}, nullptr);
+  auto *storage = wrappedType->getTypeWrapperStorageDecl();
   assert(storage);
 
   return injectProperty(storage, property->getName(),
@@ -394,8 +397,7 @@ SynthesizeTypeWrapperInitializerBody::evaluate(Evaluator &evaluator,
   auto *parent = ctor->getDeclContext()->getSelfNominalTypeDecl();
 
   // self.$_storage = .init(memberwise: $Storage(...))
-  auto *storageType =
-      evaluateOrDefault(ctx.evaluator, GetTypeWrapperStorage{parent}, nullptr);
+  auto *storageType = parent->getTypeWrapperStorageDecl();
   assert(storageType);
 
   auto *typeWrapperVar = parent->getTypeWrapperProperty();
@@ -493,8 +495,7 @@ VarDecl *SynthesizeLocalVariableForTypeWrapperStorage::evaluate(
   if (!(DC && DC->hasTypeWrapper()))
     return nullptr;
 
-  auto *storageDecl =
-      evaluateOrDefault(evaluator, GetTypeWrapperStorage{DC}, nullptr);
+  auto *storageDecl = DC->getTypeWrapperStorageDecl();
   assert(storageDecl);
 
   SmallVector<TupleTypeElt, 4> members;