Support movesAsLike for the array variant of @_rawLayout

Azoy · Azoy · commit ca8c90ecc0bf · 2024-07-28T10:29:57.000-07:00
diff --git a/docs/ReferenceGuides/UnderscoredAttributes.md b/docs/ReferenceGuides/UnderscoredAttributes.md
@@ -1033,15 +1033,17 @@ forms are currently accepted:
 
 - `@_rawLayout(size: N, alignment: M)` specifies the type's size and alignment
   in bytes.
-- `@_rawLayout(like: T)` specifies the type's size and alignment should be
-  equal to the type `T`'s.
-- `@_rawLayout(likeArrayOf: T, count: N)` specifies the type's size should be
-  `MemoryLayout<T>.stride * N` and alignment should match `T`'s, like an
-  array of N contiguous elements of `T` in memory.
-- `@_rawLayout(like: T, movesAsLike)` specifies the type's size and alignment
-  should be equal to the type `T`'s. It also guarantees that moving a value of
-  this raw layout type will have the same move semantics as the type it's like.
-  This is important for things like ObjC weak references and non-trivial move
+- `@_rawLayout(like: T(, movesAsLike))` specifies the type's size and alignment
+  should be equal to the type `T`'s. An optional `movesAsLike` parameter can be
+  passed to guarantee that moving a value of this raw layout type will have the
+  same move semantics as the type it's like. This is important for things like
+  ObjC weak references and non-trivial move constructors in C++.
+- `@_rawLayout(likeArrayOf: T, count: N(, movesAsLike))` specifies the type's
+  size should be `MemoryLayout<T>.stride * N` and alignment should match `T`'s,
+  like an array of N contiguous elements of `T` in memory. An optional
+  `movesAsLike` parameter can be passed to guarantee that moving a value of this
+  raw layout type will have the same move semantics as the type it's like. This
+  is important for things like ObjC weak references and non-trivial move
   constructors in C++.
 
 A notable difference between `@_rawLayout(like: T)` and
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -2553,11 +2553,12 @@ class RawLayoutAttr final : public DeclAttribute {
         MovesAsLike(movesAsLike) {}
 
   /// Construct a `@_rawLayout(likeArrayOf: T, count: N)` attribute.
-  RawLayoutAttr(TypeRepr *LikeType, unsigned Count, SourceLoc AtLoc,
-                SourceRange Range)
+  RawLayoutAttr(TypeRepr *LikeType, unsigned Count, bool movesAsLike,
+                SourceLoc AtLoc, SourceRange Range)
       : DeclAttribute(DeclAttrKind::RawLayout, AtLoc, Range,
                       /*implicit*/ false),
-        LikeType(LikeType), SizeOrCount(Count), Alignment(0) {}
+        LikeType(LikeType), SizeOrCount(Count), Alignment(0),
+        MovesAsLike(movesAsLike) {}
 
   /// Construct a `@_rawLayout(size: N, alignment: M)` attribute.
   RawLayoutAttr(unsigned Size, unsigned Alignment, SourceLoc AtLoc,
diff --git a/lib/IRGen/GenRecord.h b/lib/IRGen/GenRecord.h
@@ -203,21 +203,9 @@ class RecordTypeInfoImpl : public Base,
     }
 
     if (auto rawLayout = T.getRawLayout()) {
-      // Because we have a rawlayout attribute, we know this has to be a struct.
-      auto structDecl = T.getStructOrBoundGenericStruct();
-
-      if (auto likeType = rawLayout->getResolvedScalarLikeType(structDecl)) {
-        if (rawLayout->shouldMoveAsLikeType()) {
-          auto astT = T.getASTType();
-          auto subs = astT->getContextSubstitutionMap();
-          auto loweredLikeType = IGF.IGM.getLoweredType(likeType->subst(subs));
-          auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
-
-          likeTypeInfo.assignWithTake(IGF, dest, src, loweredLikeType,
-                                      isOutlined);
-          return;
-        }
-      }
+      return takeRawLayout(IGF, dest, src, T, isOutlined,
+                           /* zeroizeIfSensitive */ false, rawLayout,
+                           /* isInit */ false);
     }
 
     if (isOutlined || T.hasParameterizedExistential()) {
@@ -280,20 +268,8 @@ class RecordTypeInfoImpl : public Base,
       // If the fields are not ABI-accessible, use the value witness table.
       return emitInitializeWithTakeCall(IGF, T, dest, src);
     } else if (auto rawLayout = T.getRawLayout()) {
-      // Because we have a rawlayout attribute, we know this has to be a struct.
-      auto structDecl = T.getStructOrBoundGenericStruct();
-
-      if (auto likeType = rawLayout->getResolvedScalarLikeType(structDecl)) {
-        if (rawLayout->shouldMoveAsLikeType()) {
-          auto astT = T.getASTType();
-          auto subs = astT->getContextSubstitutionMap();
-          auto loweredLikeType = IGF.IGM.getLoweredType(likeType->subst(subs));
-          auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
-
-          likeTypeInfo.initializeWithTake(IGF, dest, src, loweredLikeType,
-                                          isOutlined, zeroizeIfSensitive);
-        }
-      }
+      return takeRawLayout(IGF, dest, src, T, isOutlined, zeroizeIfSensitive,
+                           rawLayout, /* isInit */ true);
     } else if (isOutlined || T.hasParameterizedExistential()) {
       auto offsets = asImpl().getNonFixedOffsets(IGF, T);
       for (auto &field : getFields()) {
@@ -313,6 +289,90 @@ class RecordTypeInfoImpl : public Base,
       fillWithZerosIfSensitive(IGF, src, T);
   }
 
+  void takeRawLayout(IRGenFunction &IGF, Address dest, Address src, SILType T,
+                     bool isOutlined, bool zeroizeIfSensitive,
+                     RawLayoutAttr *rawLayout, bool isInit) 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);
+
+        if (isInit) {
+          likeTypeInfo.initializeWithTake(IGF, dest, src, loweredLikeType,
+                                        isOutlined, zeroizeIfSensitive);
+        } else {
+          likeTypeInfo.assignWithTake(IGF, dest, src, loweredLikeType,
+                                      isOutlined);
+        }
+      }
+
+      if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(structDecl)) {
+        auto likeType = likeArray->first;
+        unsigned count = likeArray->second;
+
+        auto astT = T.getASTType();
+        auto subs = astT->getContextSubstitutionMap();
+        auto loweredLikeType = IGF.IGM.getLoweredType(likeType.subst(subs));
+        auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
+
+        for (unsigned i = 0; i != count; i += 1) {
+          auto index = llvm::ConstantInt::get(IGF.IGM.SizeTy, i);
+
+          Address srcEltAddr;
+          Address destEltAddr;
+
+          // If we have a fixed type, we can use a typed GEP to index into the
+          // array raw layout. Otherwise, we need to advance by bytes given the
+          // stride from the VWT of the like type.
+          if (auto fixedLikeType = dyn_cast<FixedTypeInfo>(&likeTypeInfo)) {
+            srcEltAddr = Address(IGF.Builder.CreateInBoundsGEP(
+                                    fixedLikeType->getStorageType(),
+                                    src.getAddress(),
+                                    index),
+                                 fixedLikeType->getStorageType(),
+                                 src.getAlignment());
+            destEltAddr = Address(IGF.Builder.CreateInBoundsGEP(
+                                    fixedLikeType->getStorageType(),
+                                    dest.getAddress(),
+                                    index),
+                                 fixedLikeType->getStorageType(),
+                                 dest.getAlignment());
+          } else {
+            auto eltSize = likeTypeInfo.getStride(IGF, loweredLikeType);
+            auto offset = IGF.Builder.CreateMul(index, eltSize);
+
+            srcEltAddr = Address(IGF.Builder.CreateInBoundsGEP(
+                                    IGF.IGM.Int8Ty,
+                                    src.getAddress(),
+                                    offset),
+                                 IGF.IGM.Int8Ty,
+                                 src.getAlignment());
+            destEltAddr = Address(IGF.Builder.CreateInBoundsGEP(
+                                    IGF.IGM.Int8Ty,
+                                    dest.getAddress(),
+                                    offset),
+                                 IGF.IGM.Int8Ty,
+                                 dest.getAlignment());
+          }
+
+          if (isInit) {
+            likeTypeInfo.initializeWithTake(IGF, destEltAddr, srcEltAddr,
+                                            loweredLikeType, isOutlined,
+                                            zeroizeIfSensitive);
+          } else {
+            likeTypeInfo.assignWithTake(IGF, destEltAddr, srcEltAddr,
+                                        loweredLikeType, isOutlined);
+          }
+        }
+      }
+    }
+  }
+
   void destroy(IRGenFunction &IGF, Address addr, SILType T,
                bool isOutlined) const override {
     // If the fields are not ABI-accessible, use the value witness table.
diff --git a/lib/IRGen/StructLayout.cpp b/lib/IRGen/StructLayout.cpp
@@ -111,7 +111,19 @@ StructLayout::StructLayout(IRGenModule &IGM, std::optional<CanType> type,
       SpareBits.extendWithClearBits(MinimumSize.getValueInBits());
       IsFixedLayout = true;
       IsKnownAlwaysFixedSize = IsFixedSize;
-    } else if (auto likeType = rawLayout->getResolvedScalarLikeType(sd)) {
+    } else {
+      std::optional<Type> likeType = std::nullopt;
+      unsigned count = 0;
+
+      if (auto like = rawLayout->getResolvedScalarLikeType(sd)) {
+        likeType = like;
+      }
+
+      if (auto like = rawLayout->getResolvedArrayLikeTypeAndCount(sd)) {
+        likeType = like->first;
+        count = like->second;
+      }
+
       // 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.
@@ -121,15 +133,21 @@ StructLayout::StructLayout(IRGenModule &IGM, std::optional<CanType> type,
                                       
       // Take layout attributes from the like type.
       if (const FixedTypeInfo *likeFixedType = dyn_cast<FixedTypeInfo>(&likeTypeInfo)) {
-        MinimumSize = likeFixedType->getFixedSize();
+        // If we have no count, treat this as a scalar.
+        if (count == 0) {
+          MinimumSize = likeFixedType->getFixedSize();
+        } else {
+          MinimumSize = likeFixedType->getFixedStride() * count;
+        }
+
         SpareBits.extendWithClearBits(MinimumSize.getValueInBits());
         MinimumAlign = likeFixedType->getFixedAlignment();
         IsFixedLayout = true;
         IsKnownAlwaysFixedSize = IsFixedSize;
 
-        // @_rawLayout(like: T) has an optional `movesAsLike` which enforces that
-        // a value of this raw layout type should have the same move semantics
-        // as the like its like.
+        // @_rawLayout has an optional `movesAsLike` which enforces that a value
+        // of this raw layout type should have the same move semantics as the
+        // type its like.
         if (rawLayout->shouldMoveAsLikeType()) {
           IsKnownTriviallyDestroyable = likeFixedType->isTriviallyDestroyable(ResilienceExpansion::Maximal);
           IsKnownBitwiseTakable = likeFixedType->isBitwiseTakable(ResilienceExpansion::Maximal);
@@ -147,29 +165,6 @@ StructLayout::StructLayout(IRGenModule &IGM, std::optional<CanType> type,
           IsKnownBitwiseTakable = IsNotBitwiseTakable;
         }
       }
-    } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(sd)) {
-      auto elementType = likeArray->first;
-      unsigned count = likeArray->second;
-      
-      auto subs = (*type)->getContextSubstitutionMap();
-      auto loweredElementType = IGM.getLoweredType(elementType.subst(subs));
-      const TypeInfo &likeTypeInfo = IGM.getTypeInfo(loweredElementType);
-      
-      // Take layout attributes from the like type.
-      if (const FixedTypeInfo *likeFixedType = dyn_cast<FixedTypeInfo>(&likeTypeInfo)) {
-        MinimumSize = likeFixedType->getFixedStride() * count;
-        SpareBits.extendWithClearBits(MinimumSize.getValueInBits());
-        MinimumAlign = likeFixedType->getFixedAlignment();
-        IsFixedLayout = true;
-        IsKnownAlwaysFixedSize = IsFixedSize;
-      } else {
-        MinimumSize = Size(0);
-        MinimumAlign = Alignment(1);
-        IsFixedLayout = false;
-        IsKnownAlwaysFixedSize = IsNotFixedSize;
-      }
-    } else {
-      llvm_unreachable("unhandled raw layout variant?");
     }
     
     // Set the LLVM struct type for a fixed layout according to the stride and
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -4039,14 +4039,23 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
         return makeParserSuccess();
       }
       
+      bool movesAsLike = false;
+
+      // @_rawLayout(likeArrayOf: T, count: N, movesAsLike)
+      if (consumeIf(tok::comma) && Tok.isAny(tok::identifier) &&
+          !parseSpecificIdentifier("movesAsLike",
+            diag::attr_rawlayout_expected_label, "movesAsLike")) {
+        movesAsLike = true;
+      }
+
       SourceLoc rParenLoc;
       if (!consumeIf(tok::r_paren, rParenLoc)) {
         diagnose(Tok.getLoc(), diag::attr_expected_rparen,
                  AttrName, /*isModifier*/false);
         return makeParserSuccess();
       }
       
-      attr = new (Context) RawLayoutAttr(likeType.get(), count,
+      attr = new (Context) RawLayoutAttr(likeType.get(), count, movesAsLike,
                                          AtLoc, SourceRange(Loc, rParenLoc));
     } else {
       diagnose(Loc, diag::attr_rawlayout_expected_label,
diff --git a/test/IRGen/raw_layout.swift b/test/IRGen/raw_layout.swift
