Merge pull request swiftlang#76303 from Azoy/rawlayout-fixes

Azoy · web-flow · commit 3b741ee48893 · 2024-09-11T16:25:24.000-07:00
[IRGen] Collect metadata from like type and add type information when GEPing into raw layouts
diff --git a/include/swift/SIL/SILType.h b/include/swift/SIL/SILType.h
@@ -909,6 +909,12 @@ class SILType {
     return sd->getAttrs().getAttribute<RawLayoutAttr>();
   }
 
+  /// If this is a raw layout type, returns the substituted like type.
+  Type getRawLayoutSubstitutedLikeType() const;
+
+  /// If this is a raw layout type, returns the substituted count type.
+  Type getRawLayoutSubstitutedCountType() const;
+
   /// If this is a SILBoxType, return getSILBoxFieldType(). Otherwise, return
   /// SILType().
   ///
diff --git a/lib/IRGen/GenRecord.h b/lib/IRGen/GenRecord.h
@@ -297,31 +297,31 @@ class RecordTypeInfoImpl : public Base,
                        std::function<void
                           (const TypeInfo &, SILType, Address, Address)> body) const {
     if (rawLayout->shouldMoveAsLikeType()) {
-      // Because we have a rawlayout attribute, we know this has to be a struct.
-      auto structDecl = T.getStructOrBoundGenericStruct();
-
-      if (auto likeType = rawLayout->getResolvedScalarLikeType(structDecl)) {
-        auto astT = T.getASTType();
-        auto subs = astT->getContextSubstitutionMap();
-        auto loweredLikeType = IGF.IGM.getLoweredType(likeType->subst(subs));
-        auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
-
+      auto likeType = T.getRawLayoutSubstitutedLikeType();
+      auto loweredLikeType = IGF.IGM.getLoweredType(likeType);
+      auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
+
+      // Fixup src/dest address element types because currently they are in
+      // terms of the raw layout type's [n x i8] where we're at a point to use
+      // the like type's concrete storage type.
+      src = Address(src.getAddress(), likeTypeInfo.getStorageType(),
+                    src.getAlignment());
+      dest = Address(dest.getAddress(), likeTypeInfo.getStorageType(),
+                     dest.getAlignment());
+
+      // If we're a scalar, then we only need to run the body once.
+      if (rawLayout->getScalarLikeType()) {
         body(likeTypeInfo, loweredLikeType, dest, src);
       }
 
-      if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(structDecl)) {
-        auto likeType = likeArray->first;
-        auto countType = likeArray->second;
-
-        auto astT = T.getASTType();
-        auto subs = astT->getContextSubstitutionMap();
-        auto loweredLikeType = IGF.IGM.getLoweredType(likeType.subst(subs));
-        auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
-        countType = countType.subst(subs);
+      // Otherwise, emit a loop that calls body N times where N is the count
+      // of the array variant. This could be generic in which case we need to
+      // pull the value out of metadata or it could be a constant integer.
+      if (rawLayout->getArrayLikeTypeAndCount()) {
+        auto countType = T.getRawLayoutSubstitutedCountType()->getCanonicalType();
 
-        IGF.emitLoopOverElements(likeTypeInfo, loweredLikeType,
-                                 countType->getCanonicalType(), dest, src,
-            [&](Address dest, Address src) {
+        IGF.emitLoopOverElements(likeTypeInfo, loweredLikeType, countType,
+                                 dest, src, [&](Address dest, Address src) {
           body(likeTypeInfo, loweredLikeType, dest, src);
         });
       }
@@ -555,6 +555,23 @@ class RecordTypeInfoImpl : public Base,
       auto fType = field.getType(collector.IGF.IGM, T);
       field.getTypeInfo().collectMetadataForOutlining(collector, fType);
     }
+
+    // If we're a raw layout type, collect metadata from our like type and count
+    // as well.
+    if (auto likeType = T.getRawLayoutSubstitutedLikeType()) {
+      auto loweredLikeType = collector.IGF.IGM.getLoweredType(likeType);
+      collector.IGF.IGM.getTypeInfo(loweredLikeType)
+          .collectMetadataForOutlining(collector, loweredLikeType);
+
+      if (auto countType = T.getRawLayoutSubstitutedCountType()) {
+        if (countType->isValueParameter()) {
+          auto loweredCountType = collector.IGF.IGM.getLoweredType(countType);
+          collector.IGF.IGM.getTypeInfo(loweredCountType)
+            .collectMetadataForOutlining(collector, loweredCountType);
+        }
+      }
+    }
+
     collector.collectTypeMetadata(T);
   }
 };
diff --git a/lib/IRGen/GenStruct.cpp b/lib/IRGen/GenStruct.cpp
@@ -1007,30 +1007,23 @@ namespace {
         return IGM.typeLayoutCache.getOrCreateResilientEntry(T);
       }
 
-      auto decl = T.getASTType()->getStructOrBoundGenericStruct();
-      auto rawLayout = decl->getAttrs().getAttribute<RawLayoutAttr>();
-
       // If we have a raw layout struct who is fixed size, it means the
       // layout of the struct is fully concrete.
-      if (rawLayout) {
+      if (auto rawLayout = T.getRawLayout()) {
         // Defer to this fixed type info for type layout if the raw layout
         // specifies size and alignment.
         if (rawLayout->getSizeAndAlignment()) {
           return IGM.typeLayoutCache.getOrCreateTypeInfoBasedEntry(*this, T);
         }
 
-        // The given struct type T that we're building is fully concrete, but
-        // our like type is still in terms of the potential archetype of the
-        // type.
-        auto subs = T.getASTType()->getContextSubstitutionMap(decl);
-        auto likeType = rawLayout->getResolvedLikeType(decl).subst(subs);
-        auto loweredLikeType = IGM.getLoweredType(likeType->getCanonicalType());
+        auto likeType = T.getRawLayoutSubstitutedLikeType();
+        auto loweredLikeType = IGM.getLoweredType(likeType);
         auto likeTypeLayout = IGM.getTypeInfo(loweredLikeType)
             .buildTypeLayoutEntry(IGM, loweredLikeType, useStructLayouts);
 
         // If we're an array, use the ArrayLayoutEntry.
-        if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(decl)) {
-          auto countType = likeArray->second.subst(subs)->getCanonicalType();
+        if (rawLayout->getArrayLikeTypeAndCount()) {
+          auto countType = T.getRawLayoutSubstitutedCountType()->getCanonicalType();
           return IGM.typeLayoutCache.getOrCreateArrayEntry(likeTypeLayout,
                                                            loweredLikeType,
                                                            countType);
@@ -1133,29 +1126,22 @@ namespace {
         return IGM.typeLayoutCache.getOrCreateResilientEntry(T);
       }
 
-      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) {
+      if (auto rawLayout = T.getRawLayout()) {
         // Note: We don't have to handle the size and alignment case here for
         // raw layout because those are always fixed, so only dependent layouts
         // will be non-fixed.
 
-        // The given struct type T that we're building is fully concrete, but
-        // our like type is still in terms of the potential archetype of the
-        // type.
-        auto subs = T.getASTType()->getContextSubstitutionMap(decl);
-        auto likeType = rawLayout->getResolvedLikeType(decl).subst(subs);
+        auto likeType = T.getRawLayoutSubstitutedLikeType();
         auto loweredLikeType = IGM.getLoweredType(likeType->getCanonicalType());
         auto likeTypeLayout = IGM.getTypeInfo(loweredLikeType)
             .buildTypeLayoutEntry(IGM, loweredLikeType, useStructLayouts);
 
         // If we're an array, use the ArrayLayoutEntry.
-        if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(decl)) {
-          auto countType = likeArray->second.subst(subs)->getCanonicalType();
+        if (rawLayout->getArrayLikeTypeAndCount()) {
+          auto countType = T.getRawLayoutSubstitutedCountType()->getCanonicalType();
           return IGM.typeLayoutCache.getOrCreateArrayEntry(likeTypeLayout,
                                                            loweredLikeType,
                                                            countType);
diff --git a/lib/IRGen/StructLayout.cpp b/lib/IRGen/StructLayout.cpp
@@ -85,7 +85,6 @@ StructLayout::StructLayout(IRGenModule &IGM, std::optional<CanType> type,
     rawLayout = decl->getAttrs().getAttribute<RawLayoutAttr>();
   }
   if (rawLayout && type) {
-    auto sd = cast<StructDecl>(decl);
     IsKnownTriviallyDestroyable = triviallyDestroyable;
     IsKnownBitwiseTakable = bitwiseTakable;
     SpareBits.clear();
@@ -112,31 +111,19 @@ StructLayout::StructLayout(IRGenModule &IGM, std::optional<CanType> type,
       IsFixedLayout = true;
       IsKnownAlwaysFixedSize = IsFixedSize;
     } else {
-      std::optional<Type> likeType = std::nullopt;
-      std::optional<Type> countType = std::nullopt;
+      auto loweredType = IGM.getLoweredType(*type);
 
-      if (auto like = rawLayout->getResolvedScalarLikeType(sd)) {
-        likeType = like;
-      }
+      Type likeType = loweredType.getRawLayoutSubstitutedLikeType();
+      std::optional<Type> countType = std::nullopt;
 
-      if (auto like = rawLayout->getResolvedArrayLikeTypeAndCount(sd)) {
-        likeType = like->first;
-        countType = like->second;
+      if (rawLayout->getArrayLikeTypeAndCount()) {
+        countType = loweredType.getRawLayoutSubstitutedCountType();
       }
 
-      // If our likeType is dependent, then all calls to try and lay it out will
-      // be non-fixed, but in a concrete case we want a fixed layout, so try to
-      // substitute it out.
-      auto subs = (*type)->getContextSubstitutionMap();
-      auto loweredLikeType = IGM.getLoweredType(likeType->subst(subs));
+      auto loweredLikeType = IGM.getLoweredType(likeType);
       auto &likeTypeInfo = IGM.getTypeInfo(loweredLikeType);
       auto likeFixedType = dyn_cast<FixedTypeInfo>(&likeTypeInfo);
 
-      // Substitute our count type if we have one.
-      if (countType) {
-        countType = countType->subst(subs);
-      }
-
       // Take layout attributes from the like type.
       //
       // We can only fixup the type's layout when either this is a scalar and
diff --git a/lib/SIL/IR/SILType.cpp b/lib/SIL/IR/SILType.cpp
@@ -1256,6 +1256,38 @@ bool SILType::isSendable(SILFunction *fn) const {
   return getASTType()->isSendableType();
 }
 
+Type SILType::getRawLayoutSubstitutedLikeType() const {
+  auto rawLayout = getRawLayout();
+
+  if (!rawLayout)
+    return Type();
+
+  if (rawLayout->getSizeAndAlignment())
+    return Type();
+
+  auto structDecl = getStructOrBoundGenericStruct();
+  auto likeType = rawLayout->getResolvedLikeType(structDecl);
+  auto astT = getASTType();
+  auto subs = astT->getContextSubstitutionMap();
+  return likeType.subst(subs);
+}
+
+Type SILType::getRawLayoutSubstitutedCountType() const {
+  auto rawLayout = getRawLayout();
+
+  if (!rawLayout)
+    return Type();
+
+  if (rawLayout->getSizeAndAlignment() || rawLayout->getScalarLikeType())
+    return Type();
+
+  auto structDecl = getStructOrBoundGenericStruct();
+  auto countType = rawLayout->getResolvedCountType(structDecl);
+  auto astT = getASTType();
+  auto subs = astT->getContextSubstitutionMap();
+  return countType.subst(subs);
+}
+
 std::optional<DiagnosticBehavior>
 SILType::getConcurrencyDiagnosticBehavior(SILFunction *fn) const {
   auto declRef = fn->getDeclRef();
diff --git a/test/IRGen/stdlib/Mutex.swift b/test/IRGen/stdlib/Mutex.swift
@@ -0,0 +1,45 @@
+// RUN: %empty-directory(%t)
+// RUN: %{python} %utils/chex.py < %s > %t/Mutex.swift
+// RUN: %target-swift-frontend -enable-experimental-feature RawLayout -enable-experimental-feature ValueGenerics -emit-ir -disable-availability-checking -I %S/Inputs -cxx-interoperability-mode=upcoming-swift -module-name stdlib %t/Mutex.swift | %FileCheck %t/Mutex.swift --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
+
+// REQUIRES: synchronization
+
+import Synchronization
+
+struct GenericMutex<T>: ~Copyable {
+  let mutex: Mutex<T> // this previously crashed the compiler
+}
+
+// Force emission of GenericMutex...
+// CHECK: %T6stdlib12GenericMutexVyytG
+func forceGenericMutex() -> GenericMutex<Void> {
+  GenericMutex(mutex: Mutex(()))
+}
+
+final class Awaitable<Value, Failure>: Sendable where Value: Sendable, Failure: Error {
+  struct State {
+    var pendingConsumers: [CheckedContinuation<Value, Failure>] = []
+    var result: Result<Value, Failure>?
+  }
+
+  let state: Mutex<State> // This previously crashed the compiler...
+
+  init() {
+    self.state = Mutex(.init())
+  }
+}
+
+// CHECK: define {{.*}} ptr @"$s15Synchronization5MutexVy6stdlib9AwaitableC5StateVyxq__GGs8SendableRzs5ErrorR_r0_lWOb"(ptr [[SRC:%.*]], ptr [[DEST:%.*]], ptr {{%.*}}, ptr {{%.*}}, ptr {{%.*}}, ptr {{%.*}}, ptr {{%.*}}, ptr {{%.*}}, ptr [[MUTEX:%.*]])
+// CHECK:   [[DEST_HANDLE_PTR:%.*]] = getelementptr inbounds %T15Synchronization5MutexV, ptr [[DEST]], i32 0, i32 0
+// CHECK:   [[SRC_HANDLE_PTR:%.*]] = getelementptr inbounds %T15Synchronization5MutexV, ptr [[SRC]], i32 0, i32 0
+// CHECK:   call void @llvm.memcpy.p0.p0.i{{32|64}}(ptr {{.*}} [[DEST_HANDLE_PTR]], ptr {{.*}} [[SRC_HANDLE_PTR]], i{{32|64}} {{.*}}, i1 false)
+// CHECK:   [[DEST_MUTEX_VALUE_OFFSET_PTR:%.*]] = getelementptr inbounds i32, ptr [[MUTEX]], i{{32|64}} 7
+// CHECK:   [[DEST_MUTEX_VALUE_OFFSET:%.*]] = load i32, ptr [[DEST_MUTEX_VALUE_OFFSET_PTR]]
+// CHECK:   [[DEST_VALUE_PTR:%.*]] = getelementptr inbounds i8, ptr [[DEST]], i32 [[DEST_MUTEX_VALUE_OFFSET]]
+// CHECK:   [[SRC_MUTEX_VALUE_OFFSET_PTR:%.*]] = getelementptr inbounds i32, ptr [[MUTEX]], i{{32|64}} 7
+// CHECK:   [[SRC_MUTEX_VALUE_OFFSET:%.*]] = load i32, ptr [[SRC_MUTEX_VALUE_OFFSET_PTR]]
+// CHECK:   [[SRC_VALUE_PTR:%.*]] = getelementptr inbounds i8, ptr [[SRC]], i32 [[SRC_MUTEX_VALUE_OFFSET]]
+
+// These GEPs used to cause compiler crashes because they were incorrectly typed...
+// CHECK:   [[DEST_PENDING_CONSUMERS_PTR:%.*]] = getelementptr inbounds %T6stdlib9AwaitableC5StateV, ptr [[DEST_VALUE_PTR]], i32 0, i32 0
+// CHECK:   [[SRC_PENDING_CONSUMERS_PTR:%.*]] = getelementptr inbounds %T6stdlib9AwaitableC5StateV, ptr [[SRC_VALUE_PTR]], i32 0, i32 0
