Merge pull request swiftlang#67611 from Azoy/raw-types-are-cool!

[WIP] Implement dependent layouts for raw types

Author: Azoy

include/swift/AST/ASTContext.h

@@ -941,6 +941,10 @@ class ASTContext final {
   /// descriptors.
   AvailabilityContext getSignedDescriptorAvailability();
 
+  /// Get the runtime availability of the swift_initRawStructMetadata entrypoint
+  /// that fixes up the value witness table of @_rawLayout dependent types.
+  AvailabilityContext getInitRawStructMetadataAvailability();
+
   /// Get the runtime availability of features introduced in the Swift 5.2
   /// compiler for the target platform.
   AvailabilityContext getSwift52Availability();

include/swift/Runtime/Metadata.h

@@ -1038,6 +1038,14 @@ SWIFT_RUNTIME_STDLIB_SPI
 void _swift_registerConcurrencyStandardTypeDescriptors(
     const ConcurrencyStandardTypeDescriptors *descriptors);
 
+/// Initialize the value witness table for a struct using the provided like type
+/// as the basis for the layout.
+SWIFT_RUNTIME_EXPORT
+void swift_initRawStructMetadata(StructMetadata *self,
+                                 StructLayoutFlags flags,
+                                 const TypeLayout *likeType,
+                                 int32_t count);
+
 #pragma clang diagnostic pop
 
 } // end namespace swift

include/swift/Runtime/RuntimeFunctions.def

@@ -2486,6 +2486,18 @@ FUNCTION(GenericInstantiateLayoutString,
          ATTRS(NoUnwind),
          EFFECT(MetaData))
 
+// void swift_initRawStructMetadata(Metadata *structType,
+//                                  StructLayoutFlags flags,
+//                                  const TypeLayout *likeType,
+//                                  int32_t count);
+FUNCTION(InitRawStructMetadata,
+         swift_initRawStructMetadata,
+         C_CC, AlwaysAvailable,
+         RETURNS(VoidTy),
+         ARGS(TypeMetadataPtrTy, SizeTy, Int8PtrPtrTy->getPointerTo(0), Int32Ty),
+         ATTRS(NoUnwind),
+         EFFECT(MetaData))
+
 #undef RETURNS
 #undef ARGS
 #undef ATTRS

lib/AST/Availability.cpp

@@ -573,6 +573,11 @@ ASTContext::getSignedDescriptorAvailability() {
   return getSwift59Availability();
 }
 
+AvailabilityContext
+ASTContext::getInitRawStructMetadataAvailability() {
+  return getSwiftFutureAvailability();
+}
+
 AvailabilityContext ASTContext::getSwift52Availability() {
   auto target = LangOpts.Target;
 

lib/IRGen/GenCall.cpp

@@ -134,8 +134,8 @@ AsyncContextLayout::AsyncContextLayout(
     IRGenModule &IGM, LayoutStrategy strategy, ArrayRef<SILType> fieldTypes,
     ArrayRef<const TypeInfo *> fieldTypeInfos, CanSILFunctionType originalType,
     CanSILFunctionType substitutedType, SubstitutionMap substitutionMap)
-    : StructLayout(IGM, /*decl=*/nullptr, LayoutKind::NonHeapObject, strategy,
-                   fieldTypeInfos, /*typeToFill*/ nullptr),
+    : StructLayout(IGM, /*type=*/ llvm::None, LayoutKind::NonHeapObject,
+                   strategy, fieldTypeInfos, /*typeToFill*/ nullptr),
       originalType(originalType), substitutedType(substitutedType),
       substitutionMap(substitutionMap)  {
   assert(fieldTypeInfos.size() == fieldTypes.size() &&

lib/IRGen/GenDiffFunc.cpp

@@ -211,7 +211,7 @@ class DifferentiableFuncTypeBuilder
   }
 
   StructLayout performLayout(ArrayRef<const TypeInfo *> fieldTypes) {
-    return StructLayout(IGM, /*decl=*/nullptr, LayoutKind::NonHeapObject,
+    return StructLayout(IGM, /*type=*/ llvm::None, LayoutKind::NonHeapObject,
                         LayoutStrategy::Universal, fieldTypes);
   }
 };
@@ -383,7 +383,7 @@ class LinearFuncTypeBuilder
   }
 
   StructLayout performLayout(ArrayRef<const TypeInfo *> fieldTypes) {
-    return StructLayout(IGM, /*decl=*/nullptr, LayoutKind::NonHeapObject,
+    return StructLayout(IGM, /*type=*/ llvm::None, LayoutKind::NonHeapObject,
                         LayoutStrategy::Universal, fieldTypes);
   }
 };

lib/IRGen/GenHeap.cpp

@@ -274,7 +274,7 @@ HeapLayout::HeapLayout(IRGenModule &IGM, LayoutStrategy strategy,
                        ArrayRef<const TypeInfo *> fieldTypeInfos,
                        llvm::StructType *typeToFill,
                        NecessaryBindings &&bindings, unsigned bindingsIndex)
-    : StructLayout(IGM, /*decl=*/nullptr, LayoutKind::HeapObject, strategy,
+    : StructLayout(IGM, /*type=*/ llvm::None, LayoutKind::HeapObject, strategy,
                    fieldTypeInfos, typeToFill),
       ElementTypes(fieldTypes.begin(), fieldTypes.end()),
       Bindings(std::move(bindings)), BindingsIndex(bindingsIndex) {

lib/IRGen/GenMeta.cpp

@@ -2983,6 +2983,116 @@ static void emitInitializeFieldOffsetVectorWithLayoutString(
                                 IGM.getPointerSize() * numFields);
 }
 
+static void emitInitializeRawLayoutOld(IRGenFunction &IGF, SILType likeType,
+                                       int32_t count, SILType T,
+                                       llvm::Value *metadata,
+                                       MetadataDependencyCollector *collector) {
+  auto &IGM = IGF.IGM;
+
+  // This is the list of field type layouts that we're going to pass to the init
+  // function. This will only ever hold 1 field which is the temporary one we're
+  // going to build up from our like type's layout.
+  auto fieldLayouts = IGF.createAlloca(
+                          llvm::ArrayType::get(IGM.TypeLayoutTy->getPointerTo(), 1),
+                          IGM.getPointerAlignment(), "fieldLayouts");
+  IGF.Builder.CreateLifetimeStart(fieldLayouts, IGM.getPointerSize());
+
+  // We're going to pretend that this is our field offset vector for the init to
+  // write to. We don't actually have fields, so we don't want to write a field
+  // offset in our metadata.
+  auto fieldOffsets = IGF.createAlloca(IGM.Int32Ty, Alignment(4), "fieldOffsets");
+  IGF.Builder.CreateLifetimeStart(fieldOffsets, Size(4));
+
+  // We need to make a temporary type layout with most of the same information
+  // from the type we're like.
+  auto ourTypeLayout = IGF.createAlloca(IGM.TypeLayoutTy, 
+                                        IGM.getPointerAlignment(),
+                                        "ourTypeLayout");
+  IGF.Builder.CreateLifetimeStart(ourTypeLayout, IGM.getPointerSize());
+
+  // Put our temporary type layout in the list of layouts we're using to
+  // initialize.
+  IGF.Builder.CreateStore(ourTypeLayout.getAddress(), fieldLayouts);
+
+  // Get the like type's type layout.
+  auto likeTypeLayout = emitTypeLayoutRef(IGF, likeType, collector);
+
+  // Grab the size, stride, and alignmentMask out of the layout.
+  auto loadedTyLayout = IGF.Builder.CreateLoad(
+      Address(likeTypeLayout, IGM.TypeLayoutTy, IGM.getPointerAlignment()),
+      "typeLayout");
+  auto size = IGF.Builder.CreateExtractValue(loadedTyLayout, 0, "size");
+  auto stride = IGF.Builder.CreateExtractValue(loadedTyLayout, 1, "stride");
+  auto flags = IGF.Builder.CreateExtractValue(loadedTyLayout, 2, "flags");
+  auto xi = IGF.Builder.CreateExtractValue(loadedTyLayout, 3, "xi");
+
+  // This will zero out the other bits.
+  auto alignMask = IGF.Builder.CreateAnd(flags,
+                                         ValueWitnessFlags::AlignmentMask,
+                                         "alignMask");
+
+  // Set the isNonPOD bit. This is important because older runtimes will attempt
+  // to replace various vwt functions with more optimized ones. In this case, we
+  // want to preserve the fact that noncopyable types have unreachable copy vwt
+  // functions.
+  auto vwtFlags = IGF.Builder.CreateOr(alignMask,
+                                       ValueWitnessFlags::IsNonPOD,
+                                       "vwtFlags");
+
+  // Count is only ever -1 if we're not an array like layout.
+  if (count != -1) {
+    stride = IGF.Builder.CreateMul(stride, IGM.getSize(Size(count)));
+    size = stride;
+  }
+
+  llvm::Value *resultAgg = llvm::UndefValue::get(IGM.TypeLayoutTy);
+  resultAgg = IGF.Builder.CreateInsertValue(resultAgg, size, 0);
+  resultAgg = IGF.Builder.CreateInsertValue(resultAgg, stride, 1);
+  resultAgg = IGF.Builder.CreateInsertValue(resultAgg, vwtFlags, 2);
+  resultAgg = IGF.Builder.CreateInsertValue(resultAgg, xi, 3);
+
+  IGF.Builder.CreateStore(resultAgg, ourTypeLayout);
+
+  StructLayoutFlags fnFlags = StructLayoutFlags::Swift5Algorithm;
+
+  // Call swift_initStructMetadata().
+  IGF.Builder.CreateCall(IGM.getInitStructMetadataFunctionPointer(),
+                         {metadata, IGM.getSize(Size(uintptr_t(fnFlags))),
+                          IGM.getSize(Size(1)), fieldLayouts.getAddress(),
+                          fieldOffsets.getAddress()});
+
+  IGF.Builder.CreateLifetimeEnd(ourTypeLayout, IGM.getPointerSize());
+  IGF.Builder.CreateLifetimeEnd(fieldOffsets, Size(4));
+  IGF.Builder.CreateLifetimeEnd(fieldLayouts, IGM.getPointerSize());
+}
+
+static void emitInitializeRawLayout(IRGenFunction &IGF, SILType likeType,
+                                    int32_t count, SILType T,
+                                    llvm::Value *metadata,
+                                    MetadataDependencyCollector *collector) {
+  // If our deployment target doesn't contain the new swift_initRawStructMetadata,
+  // emit a call to the swift_initStructMetadata tricking it into thinking
+  // we have a single field.
+  auto deploymentAvailability =
+      AvailabilityContext::forDeploymentTarget(IGF.IGM.Context);
+  auto initRawAvail = IGF.IGM.Context.getInitRawStructMetadataAvailability();
+
+  if (!IGF.IGM.Context.LangOpts.DisableAvailabilityChecking &&
+      !deploymentAvailability.isContainedIn(initRawAvail)) {
+    emitInitializeRawLayoutOld(IGF, likeType, count, T, metadata, collector);
+    return;
+  }
+
+  auto &IGM = IGF.IGM;
+  auto likeTypeLayout = emitTypeLayoutRef(IGF, likeType, collector);
+  StructLayoutFlags flags = StructLayoutFlags::Swift5Algorithm;
+
+  // Call swift_initRawStructMetadata().
+  IGF.Builder.CreateCall(IGM.getInitRawStructMetadataFunctionPointer(),
+                         {metadata, IGM.getSize(Size(uintptr_t(flags))),
+                          likeTypeLayout, IGM.getInt32(count)});
+}
+
 static void emitInitializeValueMetadata(IRGenFunction &IGF,
                                         NominalTypeDecl *nominalDecl,
                                         llvm::Value *metadata,
@@ -2996,10 +3106,33 @@ static void emitInitializeValueMetadata(IRGenFunction &IGF,
         IGM.getOptions().EnableLayoutStringValueWitnesses &&
         IGM.getOptions().EnableLayoutStringValueWitnessesInstantiation;
 
-  if (isa<StructDecl>(nominalDecl)) {
+  if (auto sd = dyn_cast<StructDecl>(nominalDecl)) {
     auto &fixedTI = IGM.getTypeInfo(loweredTy);
     if (isa<FixedTypeInfo>(fixedTI)) return;
 
+    // Use a different runtime function to initialize the value witness table
+    // if the struct has a raw layout. The existing swift_initStructMetadata
+    // is the wrong thing for these types.
+    if (auto rawLayout = nominalDecl->getAttrs().getAttribute<RawLayoutAttr>()) {
+      SILType loweredLikeType;
+      int32_t count = -1;
+
+      if (auto likeType = rawLayout->getResolvedScalarLikeType(sd)) {
+        loweredLikeType = IGM.getLoweredType(AbstractionPattern::getOpaque(),
+                                             *likeType);
+      } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(sd)) {
+        auto likeType = likeArray->first;
+        loweredLikeType = IGM.getLoweredType(AbstractionPattern::getOpaque(),
+                                             likeType);
+
+        count = likeArray->second;
+      }
+
+      emitInitializeRawLayout(IGF, loweredLikeType, count, loweredTy, metadata,
+                              collector);
+      return;
+    }
+
     if (useLayoutStrings) {
       emitInitializeFieldOffsetVectorWithLayoutString(IGF, loweredTy, metadata,
                                                       isVWTMutable, collector);

lib/IRGen/GenRecord.h

@@ -863,7 +863,6 @@ class RecordTypeBuilder {
     fields.reserve(astFields.size());
     fieldTypesForLayout.reserve(astFields.size());
 
-    bool loadable = true;
     auto fieldsABIAccessible = FieldsAreABIAccessible;
 
     unsigned explosionSize = 0;
@@ -880,7 +879,6 @@ class RecordTypeBuilder {
 
       auto loadableFieldTI = dyn_cast<LoadableTypeInfo>(&fieldTI);
       if (!loadableFieldTI) {
-        loadable = false;
         continue;
       }
 
@@ -902,7 +900,7 @@ class RecordTypeBuilder {
     }
 
     // Create the type info.
-    if (loadable) {
+    if (layout.isLoadable()) {
       assert(layout.isFixedLayout());
       assert(fieldsABIAccessible);
       return asImpl()->createLoadable(fields, std::move(layout), explosionSize);

lib/IRGen/GenStruct.cpp

@@ -1088,6 +1088,35 @@ namespace {
         fields.push_back(
             field.getTypeInfo().buildTypeLayoutEntry(IGM, fieldTy, useStructLayouts));
       }
+
+      auto decl = T.getASTType()->getStructOrBoundGenericStruct();
+      auto rawLayout = decl->getAttrs().getAttribute<RawLayoutAttr>();
+
+      // If we have a raw layout struct who is non-fixed size, it means the
+      // layout of the struct is dependent on the archetype of the thing it's
+      // like.
+      if (rawLayout) {
+        SILType loweredLikeType;
+
+        if (auto likeType = rawLayout->getResolvedScalarLikeType(decl)) {
+          loweredLikeType = IGM.getLoweredType(*likeType);
+        } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(decl)) {
+          loweredLikeType = IGM.getLoweredType(likeArray->first);
+        }
+
+        // The given struct type T that we're building may be in a generic
+        // environment that is different than that which was built our
+        // resolved rawLayout like type. Map our like type into the given
+        // environment.
+        auto subs = T.getASTType()->getContextSubstitutionMap(
+          IGM.getSwiftModule(), decl);
+
+        loweredLikeType = loweredLikeType.subst(IGM.getSILModule(), subs);
+
+        return IGM.getTypeInfo(loweredLikeType).buildTypeLayoutEntry(IGM,
+          loweredLikeType, useStructLayouts);
+      }
+
       assert(!fields.empty() &&
              "Empty structs should not be NonFixedStructTypeInfo");
 
@@ -1244,8 +1273,7 @@ namespace {
     }
 
     StructLayout performLayout(ArrayRef<const TypeInfo *> fieldTypes) {
-      return StructLayout(IGM, TheStruct->getAnyNominal(),
-                          LayoutKind::NonHeapObject,
+      return StructLayout(IGM, TheStruct, LayoutKind::NonHeapObject,
                           LayoutStrategy::Optimal, fieldTypes, StructTy);
     }
   };

lib/IRGen/GenTuple.cpp

@@ -512,7 +512,7 @@ namespace {
     }
 
     StructLayout performLayout(ArrayRef<const TypeInfo *> fieldTypes) {
-      return StructLayout(IGM, /*decl=*/nullptr, LayoutKind::NonHeapObject,
+      return StructLayout(IGM, /*type=*/ llvm::None, LayoutKind::NonHeapObject,
                           LayoutStrategy::Universal, fieldTypes);
     }
   };

lib/IRGen/GenType.cpp

@@ -2437,6 +2437,17 @@ namespace {
       if (IGM.isResilient(decl, ResilienceExpansion::Maximal))
         return true;
 
+      auto rawLayout = decl->getAttrs().getAttribute<RawLayoutAttr>();
+
+      // If our struct has a raw layout, it may be dependent on the like type.
+      if (rawLayout) {
+        if (auto likeType = rawLayout->getResolvedScalarLikeType(decl)) {
+          return visit((*likeType)->getCanonicalType());
+        } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(decl)) {
+          return visit(likeArray->first->getCanonicalType());
+        }
+      }
+
       for (auto field : decl->getStoredProperties()) {
         if (visit(field->getInterfaceType()->getCanonicalType()))
           return true;