[AST/Sema] TypeWrappers: Augment GetTypeWrapper request to produce more info

The request is updated to return attribute, wrapper declaration,
the declaration its attached to and whether or not it has been
inferred (opposite to being declared directly).

Author: xedin

include/swift/AST/Decl.h

@@ -35,6 +35,7 @@
 #include "swift/AST/StorageImpl.h"
 #include "swift/AST/TypeAlignments.h"
 #include "swift/AST/TypeWalker.h"
+#include "swift/AST/TypeWrappers.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/ArrayRefView.h"
 #include "swift/Basic/Compiler.h"
@@ -3903,7 +3904,7 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   bool hasTypeWrapper() const { return bool(getTypeWrapper()); }
 
   /// Return a type wrapper (if any) associated with this type.
-  NominalTypeDecl *getTypeWrapper() const;
+  Optional<TypeWrapperInfo> getTypeWrapper() const;
 
   /// If this declaration has a type wrapper return a property that
   /// is used for all type wrapper related operations (mainly for

include/swift/AST/TypeCheckRequests.h

@@ -64,6 +64,7 @@ struct TypeWitnessAndDecl;
 class ValueDecl;
 enum class OpaqueReadOwnership: uint8_t;
 class StorageImplInfo;
+struct TypeWrapperInfo;
 
 /// Display a nominal type or extension thereof.
 void simple_display(
@@ -3522,15 +3523,17 @@ class GetSourceFileAsyncNode
 
 /// Return a type wrapper (if any) associated with the given declaration.
 class GetTypeWrapper
-    : public SimpleRequest<GetTypeWrapper, NominalTypeDecl *(NominalTypeDecl *),
+    : public SimpleRequest<GetTypeWrapper,
+                           Optional<TypeWrapperInfo>(NominalTypeDecl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  NominalTypeDecl *evaluate(Evaluator &evaluator, NominalTypeDecl *) const;
+  Optional<TypeWrapperInfo> evaluate(Evaluator &evaluator,
+                                     NominalTypeDecl *) const;
 
 public:
   bool isCached() const { return true; }

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -402,7 +402,7 @@ SWIFT_REQUEST(TypeChecker, GetSourceFileAsyncNode,
               AwaitExpr *(const SourceFile *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetTypeWrapper,
-              NominalTypeDecl *(NominalTypeDecl *),
+              Optional<TypeWrapperInfo>(NominalTypeDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetTypeWrapperType,
               Type(NominalTypeDecl *),

include/swift/AST/TypeWrappers.h

@@ -0,0 +1,40 @@
+//===--------- TypeWrappers.h - Type Wrapper ASTs ---------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines helper types for type wrappers.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_TYPE_WRAPPERS_H
+#define SWIFT_AST_TYPE_WRAPPERS_H
+
+namespace swift {
+
+struct TypeWrapperInfo {
+  CustomAttr *Attr;
+  NominalTypeDecl *Wrapper;
+  NominalTypeDecl *AttachedTo;
+  bool IsInferred;
+
+  TypeWrapperInfo(CustomAttr *attr, NominalTypeDecl *wrapperDecl,
+                  NominalTypeDecl *attachedTo, bool isInferred)
+      : Attr(attr), Wrapper(wrapperDecl), AttachedTo(attachedTo),
+        IsInferred(isInferred) {}
+
+  TypeWrapperInfo asInferred() const {
+    return {Attr, Wrapper, AttachedTo, true};
+  }
+};
+
+} // end namespace swift
+
+#endif // SWIFT_AST_TYPE_WRAPPERS_H

lib/AST/TypeWrapper.cpp

@@ -18,13 +18,15 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeResolutionStage.h"
+#include "swift/AST/TypeWrappers.h"
 
 using namespace swift;
 
-NominalTypeDecl *NominalTypeDecl::getTypeWrapper() const {
+Optional<TypeWrapperInfo> NominalTypeDecl::getTypeWrapper() const {
   auto *mutableSelf = const_cast<NominalTypeDecl *>(this);
   return evaluateOrDefault(getASTContext().evaluator,
-                           GetTypeWrapper{mutableSelf}, nullptr);
+                           GetTypeWrapper{mutableSelf},
+                           Optional<TypeWrapperInfo>());
 }
 
 VarDecl *ConstructorDecl::getLocalTypeWrapperStorageVar() const {

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -1257,8 +1257,9 @@ void LifetimeChecker::injectTypeWrapperStorageInitalization() {
       {
         bool isClass = isa<ClassDecl>(parentType);
 
-        auto typeWrapper = parentType->getTypeWrapper();
-        auto *typeWrapperInit = typeWrapper->getTypeWrapperInitializer();
+        auto typeWrapperInfo = parentType->getTypeWrapper();
+        auto *typeWrapperInit =
+            typeWrapperInfo->Wrapper->getTypeWrapperInitializer();
         SILValue typeWrapperInitRef = createInitRef(typeWrapperInit);
 
         auto *self = TheMemory.findUninitializedSelfValue();

lib/Sema/CodeSynthesis.cpp

@@ -364,7 +364,10 @@ static ConstructorDecl *createImplicitConstructor(NominalTypeDecl *decl,
   } else if (ICK == ImplicitConstructorKind::TypeWrapperStorage) {
     accessLevel = AccessLevel::Public;
 
-    auto *typeWrapper = decl->getTypeWrapper();
+    auto typeWrapperInfo = decl->getTypeWrapper();
+    assert(typeWrapperInfo);
+
+    auto *typeWrapper = typeWrapperInfo->Wrapper;
 
     auto *arg = new (ctx) ParamDecl(SourceLoc(), Loc, ctx.Id_storageWrapper,
                                     Loc, ctx.Id_storageWrapper, decl);

lib/Sema/TypeCheckTypeWrapper.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/Pattern.h"
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeCheckRequests.h"
+#include "swift/AST/TypeWrappers.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/SourceLoc.h"
 
@@ -77,12 +78,6 @@ VarDecl *VarDecl::getUnderlyingTypeWrapperStorage() const {
                            nullptr);
 }
 
-struct TypeWrapperAttrInfo {
-  CustomAttr *Attr;
-  NominalTypeDecl *Wrapper;
-  NominalTypeDecl *AttachedTo;
-};
-
 static void
 getDeclaredProtocolConformances(NominalTypeDecl *decl,
                                 SmallVectorImpl<ProtocolDecl *> &protocols) {
@@ -110,9 +105,8 @@ getDeclaredProtocolConformances(NominalTypeDecl *decl,
   }
 }
 
-static void
-getTypeWrappers(NominalTypeDecl *decl,
-                SmallVectorImpl<TypeWrapperAttrInfo> &typeWrappers) {
+static void getTypeWrappers(NominalTypeDecl *decl,
+                            SmallVectorImpl<TypeWrapperInfo> &typeWrappers) {
   auto &ctx = decl->getASTContext();
 
   // Attributes applied directly to the type.
@@ -126,7 +120,8 @@ getTypeWrappers(NominalTypeDecl *decl,
 
     auto *typeWrapper = nominal->getAttrs().getAttribute<TypeWrapperAttr>();
     if (typeWrapper && typeWrapper->isValid())
-      typeWrappers.push_back({mutableAttr, nominal, decl});
+      typeWrappers.push_back(
+          {mutableAttr, nominal, decl, /*isInferred=*/false});
   }
 
   // Do not allow transitive protocol inference between protocols.
@@ -139,33 +134,33 @@ getTypeWrappers(NominalTypeDecl *decl,
   getDeclaredProtocolConformances(decl, protocols);
 
   for (auto *protocol : protocols) {
-    SmallVector<TypeWrapperAttrInfo, 2> inferredAttrs;
+    SmallVector<TypeWrapperInfo, 2> inferredAttrs;
     getTypeWrappers(protocol, inferredAttrs);
 
     // De-duplicate inferred type wrappers. This also makes sure
     // that if both protocol and conforming type explicitly declare
     // the same type wrapper there is no clash between them.
     for (const auto &inferredAttr : inferredAttrs) {
-      if (llvm::find_if(typeWrappers, [&](const TypeWrapperAttrInfo &attr) {
+      if (llvm::find_if(typeWrappers, [&](const TypeWrapperInfo &attr) {
             return attr.Wrapper == inferredAttr.Wrapper;
           }) == typeWrappers.end())
-        typeWrappers.push_back(inferredAttr);
+        typeWrappers.push_back(inferredAttr.asInferred());
     }
   }
 }
 
-NominalTypeDecl *GetTypeWrapper::evaluate(Evaluator &evaluator,
-                                          NominalTypeDecl *decl) const {
+Optional<TypeWrapperInfo>
+GetTypeWrapper::evaluate(Evaluator &evaluator, NominalTypeDecl *decl) const {
   auto &ctx = decl->getASTContext();
 
   // Note that we don't actually care whether there are duplicates,
   // using the same type wrapper multiple times is still an error.
-  SmallVector<TypeWrapperAttrInfo, 2> typeWrappers;
+  SmallVector<TypeWrapperInfo, 2> typeWrappers;
 
   getTypeWrappers(decl, typeWrappers);
 
   if (typeWrappers.empty())
-    return nullptr;
+    return None;
 
   if (typeWrappers.size() != 1) {
     ctx.Diags.diagnose(decl, diag::cannot_use_multiple_type_wrappers,
@@ -183,25 +178,21 @@ NominalTypeDecl *GetTypeWrapper::evaluate(Evaluator &evaluator,
       }
     }
 
-    return nullptr;
+    return None;
   }
 
-  return typeWrappers.front().Wrapper;
+  return typeWrappers.front();
 }
 
 Type GetTypeWrapperType::evaluate(Evaluator &evaluator,
                                   NominalTypeDecl *decl) const {
-  SmallVector<TypeWrapperAttrInfo, 2> typeWrappers;
-
-  getTypeWrappers(decl, typeWrappers);
-
-  if (typeWrappers.size() != 1)
+  auto typeWrapperInfo = decl->getTypeWrapper();
+  if (!typeWrapperInfo)
     return Type();
 
-  auto *typeWrapperAttr = typeWrappers.front().Attr;
   auto type = evaluateOrDefault(
       evaluator,
-      CustomAttrTypeRequest{typeWrapperAttr, decl->getDeclContext(),
+      CustomAttrTypeRequest{typeWrapperInfo->Attr, decl->getDeclContext(),
                             CustomAttrTypeKind::TypeWrapper},
       Type());
 
@@ -325,7 +316,7 @@ GetTypeWrapperProperty::evaluate(Evaluator &evaluator,
                                  NominalTypeDecl *parent) const {
   auto &ctx = parent->getASTContext();
 
-  auto *typeWrapper = parent->getTypeWrapper();
+  auto typeWrapper = parent->getTypeWrapper();
   if (!typeWrapper)
     return nullptr;
 
@@ -339,10 +330,9 @@ GetTypeWrapperProperty::evaluate(Evaluator &evaluator,
 
   // $storage: Wrapper<<ParentType>, <ParentType>.$Storage>
   auto propertyTy = BoundGenericType::get(
-      typeWrapper, /*Parent=*/typeWrapperType->getParent(),
+      typeWrapper->Wrapper, /*Parent=*/typeWrapperType->getParent(),
       /*genericArgs=*/
-      {parent->getSelfInterfaceType(),
-       storage->getDeclaredInterfaceType()});
+      {parent->getSelfInterfaceType(), storage->getDeclaredInterfaceType()});
 
   return injectProperty(parent, ctx.Id_TypeWrapperProperty, propertyTy,
                         VarDecl::Introducer::Var,
@@ -426,7 +416,7 @@ static SubscriptExpr *subscriptTypeWrappedProperty(VarDecl *var,
         ctx, DeclBaseName::createSubscript(),
         {ctx.Id_wrappedSelf, ctx.Id_propertyKeyPath, ctx.Id_storageKeyPath});
 
-    auto *typeWrapper = parent->getTypeWrapper();
+    auto *typeWrapper = parent->getTypeWrapper()->Wrapper;
     auto candidates = typeWrapper->lookupDirect(subscriptName);
 
     if (!candidates.empty()) {