Merge pull request swiftlang#67588 from jckarter/raw-layout-multi-file

Requestify raw layout "like" type resolution.

Author: jckarter

include/swift/AST/Attr.h

@@ -2496,8 +2496,10 @@ class RawLayoutAttr final : public DeclAttribute {
   /// If `LikeType` is null, the alignment in bytes to use for the raw storage.
   unsigned Alignment;
   /// The resolved like type.
-  Type ResolvedLikeType = Type();
-  
+  mutable Type CachedResolvedLikeType = Type();
+
+  friend class ResolveRawLayoutLikeTypeRequest;
+
 public:
   /// Construct a `@_rawLayout(like: T)` attribute.
   RawLayoutAttr(TypeRepr *LikeType,
@@ -2553,32 +2555,29 @@ class RawLayoutAttr final : public DeclAttribute {
     return std::make_pair(SizeOrCount, Alignment);
   }
 
-  /// Set the resolved type.
-  void setResolvedLikeType(Type ty) {
-    assert(LikeType && "doesn't have a like type");
-    ResolvedLikeType = ty;
-  }
-  
+  Type getResolvedLikeType(StructDecl *sd) const;
+
   /// Return the type whose single-element layout the attribute type should get
   /// its layout from. Returns None if the attribute specifies an array or manual
   /// layout.
-  llvm::Optional<Type> getResolvedScalarLikeType() const {
+  llvm::Optional<Type> getResolvedScalarLikeType(StructDecl *sd) const {
     if (!LikeType)
       return llvm::None;
     if (Alignment != ~0u)
       return llvm::None;
-    return ResolvedLikeType;
+    return getResolvedLikeType(sd);
   }
 
   /// Return the type whose array layout the attribute type should get its
   /// layout from, along with the size of that array. Returns None if the
   /// attribute specifies scalar or manual layout.
-  llvm::Optional<std::pair<Type, unsigned>> getResolvedArrayLikeTypeAndCount() const {
+  llvm::Optional<std::pair<Type, unsigned>>
+  getResolvedArrayLikeTypeAndCount(StructDecl *sd) const {
     if (!LikeType)
       return llvm::None;
     if (Alignment == ~0u)
       return llvm::None;
-    return std::make_pair(ResolvedLikeType, SizeOrCount);
+    return std::make_pair(getResolvedLikeType(sd), SizeOrCount);
   }
 
   static bool classof(const DeclAttribute *DA) {

include/swift/AST/TypeCheckRequests.h

@@ -3376,6 +3376,28 @@ class ResolveTypeEraserTypeRequest
   void cacheResult(Type value) const;
 };
 
+class ResolveRawLayoutLikeTypeRequest
+    : public SimpleRequest<ResolveRawLayoutLikeTypeRequest,
+                           Type (StructDecl*, RawLayoutAttr *),
+                           RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  Type evaluate(Evaluator &evaluator,
+                StructDecl *sd,
+                RawLayoutAttr *attr) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  llvm::Optional<Type> getCachedResult() const;
+  void cacheResult(Type value) const;
+};
+
 /// Determines whether this is a "simple" didSet i.e one that either does not
 /// use the implicit oldValue parameter in the body or does not take an explicit
 /// parameter (ex: 'didSet(oldValue)').

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -375,6 +375,9 @@ SWIFT_REQUEST(TypeChecker, ConformanceHasEffectRequest,
 SWIFT_REQUEST(TypeChecker, ResolveTypeRequest,
               Type (const TypeResolution *, TypeRepr *, GenericParamList *),
               Uncached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ResolveRawLayoutLikeTypeRequest,
+              Type(StructDecl *, RawLayoutAttr *),
+              SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SPIGroupsRequest,
               llvm::ArrayRef<Identifier>(Decl *),
               Cached, NoLocationInfo)

lib/AST/Attr.cpp

@@ -1885,6 +1885,14 @@ Type TypeEraserAttr::getResolvedType(const ProtocolDecl *PD) const {
                            ErrorType::get(ctx));
 }
 
+Type RawLayoutAttr::getResolvedLikeType(StructDecl *sd) const {
+  auto &ctx = sd->getASTContext();
+  return evaluateOrDefault(ctx.evaluator,
+                           ResolveRawLayoutLikeTypeRequest{sd,
+                               const_cast<RawLayoutAttr *>(this)},
+                           ErrorType::get(ctx));
+}
+
 AvailableAttr *
 AvailableAttr::createPlatformAgnostic(ASTContext &C,
                                    StringRef Message,

lib/AST/TypeCheckRequests.cpp

@@ -1387,6 +1387,24 @@ void ResolveTypeEraserTypeRequest::cacheResult(Type value) const {
   }
 }
 
+//----------------------------------------------------------------------------//
+// ResolveRawLayoutLikeTypeRequest computation.
+//----------------------------------------------------------------------------//
+
+llvm::Optional<Type> ResolveRawLayoutLikeTypeRequest::getCachedResult() const {
+  auto Ty = std::get<1>(getStorage())->CachedResolvedLikeType;
+  if (!Ty) {
+    return llvm::None;
+  }
+  return Ty;
+}
+
+void ResolveRawLayoutLikeTypeRequest::cacheResult(Type value) const {
+  assert(value && "Resolved type erasure type to null type!");
+  auto *attr = std::get<1>(getStorage());
+  attr->CachedResolvedLikeType = value;
+}
+
 //----------------------------------------------------------------------------//
 // TypeCheckSourceFileRequest computation.
 //----------------------------------------------------------------------------//

lib/IRGen/StructLayout.cpp

@@ -74,6 +74,7 @@ StructLayout::StructLayout(IRGenModule &IGM,
     rawLayout = decl->getAttrs().getAttribute<RawLayoutAttr>();
   }
   if (rawLayout) {
+    auto sd = cast<StructDecl>(decl);
     IsKnownTriviallyDestroyable = deinit;
     IsKnownBitwiseTakable = IsBitwiseTakable;
     SpareBits.clear();
@@ -98,7 +99,7 @@ StructLayout::StructLayout(IRGenModule &IGM,
       SpareBits.extendWithClearBits(MinimumSize.getValueInBits());
       IsFixedLayout = true;
       IsKnownAlwaysFixedSize = IsFixedSize;
-    } else if (auto likeType = rawLayout->getResolvedScalarLikeType()) {
+    } else if (auto likeType = rawLayout->getResolvedScalarLikeType(sd)) {
       const TypeInfo &likeTypeInfo
         = IGM.getTypeInfoForUnlowered(AbstractionPattern::getOpaque(),
                                       *likeType);
@@ -116,7 +117,7 @@ StructLayout::StructLayout(IRGenModule &IGM,
         IsFixedLayout = false;
         IsKnownAlwaysFixedSize = IsNotFixedSize;
       }
-    } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount()) {
+    } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(sd)) {
       auto elementType = likeArray->first;
       unsigned count = likeArray->second;
 

lib/Sema/TypeCheckAttr.cpp

@@ -3366,6 +3366,22 @@ ResolveTypeEraserTypeRequest::evaluate(Evaluator &evaluator,
   }
 }
 
+Type
+ResolveRawLayoutLikeTypeRequest::evaluate(Evaluator &evaluator,
+                                          StructDecl *sd,
+                                          RawLayoutAttr *attr) const {
+  assert(attr->LikeType);
+  // Resolve the like type in the struct's context.
+  return TypeResolution::resolveContextualType(
+        attr->LikeType, sd, llvm::None,
+        // Unbound generics and placeholders
+        // are not allowed within this
+        // attribute.
+        /*unboundTyOpener*/ nullptr,
+        /*placeholderHandler*/ nullptr,
+        /*packElementOpener*/ nullptr);
+}
+
 bool
 TypeEraserHasViableInitRequest::evaluate(Evaluator &evaluator,
                                          TypeEraserAttr *attr,
@@ -7195,30 +7211,10 @@ void AttributeChecker::visitRawLayoutAttr(RawLayoutAttr *attr) {
       diagnoseAndRemoveAttr(attr, diag::alignment_not_power_of_two);
       return;
     }
-  } else if (auto likeType = attr->getScalarLikeType()) {
-    // Resolve the like type in the struct's context.
-    auto resolvedType = TypeResolution::resolveContextualType(
-        likeType, sd, llvm::None,
-        // Unbound generics and placeholders
-        // are not allowed within this
-        // attribute.
-        /*unboundTyOpener*/ nullptr,
-        /*placeholderHandler*/ nullptr,
-        /*packElementOpener*/ nullptr);
-        
-    attr->setResolvedLikeType(resolvedType);
-  } else if (auto arrayType = attr->getArrayLikeTypeAndCount()) {
-    // Resolve the like type in the struct's context.
-    auto resolvedType = TypeResolution::resolveContextualType(
-        arrayType->first, sd, llvm::None,
-        // Unbound generics and placeholders
-        // are not allowed within this
-        // attribute.
-        /*unboundTyOpener*/ nullptr,
-        /*placeholderHandler*/ nullptr,
-        /*packElementOpener*/ nullptr);
-        
-    attr->setResolvedLikeType(resolvedType);
+  } else if (attr->getScalarLikeType()) {
+    (void)attr->getResolvedLikeType(sd);
+  } else if (attr->getArrayLikeTypeAndCount()) {
+    (void)attr->getResolvedLikeType(sd);
   } else {
     llvm_unreachable("new unhandled rawLayout attribute form?");
   }

lib/Serialization/Serialization.cpp

@@ -3199,15 +3199,19 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       uint8_t rawAlign;
       TypeID typeID;
 
+      auto SD = const_cast<StructDecl*>(cast<StructDecl>(D));
+      
       if (auto sizeAndAlign = attr->getSizeAndAlignment()) {
         typeID = 0;
         rawSize = sizeAndAlign->first;
         rawAlign = sizeAndAlign->second;
-      } else if (auto likeType = attr->getResolvedScalarLikeType()) {
+      } else if (auto likeType
+                   = attr->getResolvedScalarLikeType(SD)) {
         typeID = S.addTypeRef(*likeType);
         rawSize = 0;
         rawAlign = 0;
-      } else if (auto likeArrayTypeAndCount = attr->getResolvedArrayLikeTypeAndCount()) {
+      } else if (auto likeArrayTypeAndCount
+                   = attr->getResolvedArrayLikeTypeAndCount(SD)) {
         typeID = S.addTypeRef(likeArrayTypeAndCount->first);
         rawSize = likeArrayTypeAndCount->second;
         rawAlign = static_cast<uint8_t>(~0u);

test/IRGen/Inputs/raw_layout_multifile_b.swift

@@ -0,0 +1,5 @@
+extension Foo where T == Int32 {
+
+}
+
+public func foo(_: borrowing Foo<Int32>) {}

test/IRGen/raw_layout_multifile.swift

@@ -0,0 +1,7 @@
+// RUN: %target-swift-frontend -enable-experimental-feature RawLayout -emit-ir -primary-file %s %S/Inputs/raw_layout_multifile_b.swift
+// RUN: %target-swift-frontend -enable-experimental-feature RawLayout -emit-ir %s -primary-file %S/Inputs/raw_layout_multifile_b.swift
+
+import Swift
+
+@_rawLayout(like: Int32)
+public struct Foo<T>: ~Copyable {}