Add option for raw layout to move as its like type

Author: Azoy

docs/ReferenceGuides/UnderscoredAttributes.md

@@ -955,12 +955,14 @@ the memory of the annotated type:
   threads, writes don't overlap with reads or writes coming from the same
   thread, and that the pointer is not used after the value is moved or
   consumed.
-- When the value is moved, a bitwise copy of its memory is performed to the new
-  address of the value in its new owner. As currently implemented, raw storage
-  types are not suitable for storing values which are not bitwise-movable, such
-  as nontrivial C++ types, Objective-C weak references, and data structures
-  such as `pthread_mutex_t` which are implemented in C as always requiring a
-  fixed address.
+- By default, when the value is moved a bitwise copy of its memory is performed
+  to the new address of the value in its new owner. This makes it unsuitable to
+  store not bitwise-movable types such as nontrivial C++ types, Objective-C weak
+  references, and data structures such as `pthread_mutex_t` which are
+  implemented in C as always requiring a fixed address. However, you can provide
+  `movesAsLike` to the `like:` version of this attribute to enforce that moving
+  the value will defer its move semantics to the type it's like. This makes it
+  suitable for storing such values that are not bitwise-movable.
 
 Using the `@_rawLayout` attribute will suppress the annotated type from
 being implicitly `Sendable`. If the type is safe to access across threads, it
@@ -987,6 +989,11 @@ forms are currently accepted:
 - `@_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
+  constructors in C++.
 
 A notable difference between `@_rawLayout(like: T)` and
 `@_rawLayout(likeArrayOf: T, count: 1)` is that the latter will pad out the

include/swift/AST/Attr.h

@@ -2536,17 +2536,21 @@ class RawLayoutAttr final : public DeclAttribute {
   unsigned SizeOrCount;
   /// If `LikeType` is null, the alignment in bytes to use for the raw storage.
   unsigned Alignment;
+  /// If a value of this raw layout type should move like its `LikeType`.
+  bool MovesAsLike = false;
   /// The resolved like type.
   mutable Type CachedResolvedLikeType = Type();
 
   friend class ResolveRawLayoutLikeTypeRequest;
 
 public:
   /// Construct a `@_rawLayout(like: T)` attribute.
-  RawLayoutAttr(TypeRepr *LikeType, SourceLoc AtLoc, SourceRange Range)
+  RawLayoutAttr(TypeRepr *LikeType, bool movesAsLike, SourceLoc AtLoc,
+                SourceRange Range)
       : DeclAttribute(DeclAttrKind::RawLayout, AtLoc, Range,
                       /*implicit*/ false),
-        LikeType(LikeType), SizeOrCount(0), Alignment(~0u) {}
+        LikeType(LikeType), SizeOrCount(0), Alignment(~0u),
+        MovesAsLike(movesAsLike) {}
 
   /// Construct a `@_rawLayout(likeArrayOf: T, count: N)` attribute.
   RawLayoutAttr(TypeRepr *LikeType, unsigned Count, SourceLoc AtLoc,
@@ -2618,6 +2622,11 @@ class RawLayoutAttr final : public DeclAttribute {
     return std::make_pair(getResolvedLikeType(sd), SizeOrCount);
   }
 
+  /// Whether a value of this raw layout should move like its `LikeType`.
+  bool shouldMoveAsLikeType() const {
+    return MovesAsLike;
+  }
+
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DeclAttrKind::RawLayout;
   }

include/swift/SIL/SILType.h

@@ -891,6 +891,17 @@ class SILType {
   /// Returns true if this SILType is a differentiable type.
   bool isDifferentiable(SILModule &M) const;
 
+  /// Returns the @_rawLayout attribute on this type if it has one.
+  RawLayoutAttr *getRawLayout() const {
+    auto sd = getStructOrBoundGenericStruct();
+
+    if (!sd) {
+      return nullptr;
+    }
+
+    return sd->getAttrs().getAttribute<RawLayoutAttr>();
+  }
+
   /// If this is a SILBoxType, return getSILBoxFieldType(). Otherwise, return
   /// SILType().
   ///

lib/IRGen/GenRecord.h

@@ -202,6 +202,25 @@ class RecordTypeInfoImpl : public Base,
       return emitAssignWithTakeCall(IGF, T, dest, src);
     }
 
+    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(IGF.IGM.getSwiftModule(),
+                                                      structDecl);
+          auto loweredLikeType = IGF.IGM.getLoweredType(likeType->subst(subs));
+          auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
+
+          likeTypeInfo.assignWithTake(IGF, dest, src, loweredLikeType,
+                                      isOutlined);
+          return;
+        }
+      }
+    }
+
     if (isOutlined || T.hasParameterizedExistential()) {
       auto offsets = asImpl().getNonFixedOffsets(IGF, T);
       for (auto &field : getFields()) {
@@ -261,7 +280,22 @@ class RecordTypeInfoImpl : public Base,
     } else if (!AreFieldsABIAccessible) {
       // 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(IGF.IGM.getSwiftModule(),
+                                                      structDecl);
+          auto loweredLikeType = IGF.IGM.getLoweredType(likeType->subst(subs));
+          auto &likeTypeInfo = IGF.IGM.getTypeInfo(loweredLikeType);
+
+          likeTypeInfo.initializeWithTake(IGF, dest, src, loweredLikeType,
+                                          isOutlined, zeroizeIfSensitive);
+        }
+      }
     } else if (isOutlined || T.hasParameterizedExistential()) {
       auto offsets = asImpl().getNonFixedOffsets(IGF, T);
       for (auto &field : getFields()) {

lib/IRGen/StructLayout.cpp

@@ -125,11 +125,26 @@ StructLayout::StructLayout(IRGenModule &IGM, std::optional<CanType> type,
         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.
+        if (rawLayout->shouldMoveAsLikeType()) {
+          IsKnownTriviallyDestroyable = likeFixedType->isTriviallyDestroyable(ResilienceExpansion::Maximal);
+          IsKnownBitwiseTakable = likeFixedType->isBitwiseTakable(ResilienceExpansion::Maximal);
+        }
       } else {
         MinimumSize = Size(0);
         MinimumAlign = Alignment(1);
         IsFixedLayout = false;
         IsKnownAlwaysFixedSize = IsNotFixedSize;
+
+        // We don't know our like type, so assume we're not known to be bitwise
+        // takable.
+        if (rawLayout->shouldMoveAsLikeType()) {
+          IsKnownTriviallyDestroyable = IsNotTriviallyDestroyable;
+          IsKnownBitwiseTakable = IsNotBitwiseTakable;
+        }
       }
     } else if (auto likeArray = rawLayout->getResolvedArrayLikeTypeAndCount(sd)) {
       auto elementType = likeArray->first;

lib/Parse/ParseDecl.cpp

@@ -3989,14 +3989,24 @@ ParserStatus Parser::parseNewDeclAttribute(DeclAttributes &Attributes,
       if (likeType.isNull()) {
         return makeParserSuccess();
       }
+
+      bool movesAsLike = false;
+
+      // @_rawLayout(like: T, 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(),
+      attr = new (Context) RawLayoutAttr(likeType.get(), movesAsLike,
                                          AtLoc, SourceRange(Loc, rParenLoc));
     } else if (firstLabel.is("likeArrayOf")) {
       // @_rawLayout(likeArrayOf: T, count: N)

lib/Serialization/Deserialization.cpp

@@ -6292,8 +6292,9 @@ llvm::Error DeclDeserializer::deserializeDeclCommon() {
         TypeID typeID;
         uint32_t rawSize;
         uint8_t rawAlign;
+        bool movesAsLike;
         serialization::decls_block::RawLayoutDeclAttrLayout::
-          readRecord(scratch, isImplicit, typeID, rawSize, rawAlign);
+          readRecord(scratch, isImplicit, typeID, rawSize, rawAlign, movesAsLike);
 
         if (typeID) {
           auto type = MF.getTypeChecked(typeID);
@@ -6303,6 +6304,7 @@ llvm::Error DeclDeserializer::deserializeDeclCommon() {
           auto typeRepr = new (ctx) FixedTypeRepr(type.get(), SourceLoc());
           if (rawAlign == 0) {
             Attr = new (ctx) RawLayoutAttr(typeRepr,
+                                           movesAsLike,
                                            SourceLoc(),
                                            SourceRange());
             break;

lib/Serialization/ModuleFormat.h

@@ -2165,7 +2165,8 @@ namespace decls_block {
     BCFixed<1>, // implicit
     TypeIDField, // like type
     BCVBR<32>, // size
-    BCVBR<8> // alignment
+    BCVBR<8>, // alignment
+    BCFixed<1> // movesAsLike
   >;
 
   using SwiftNativeObjCRuntimeBaseDeclAttrLayout = BCRecordLayout<

lib/Serialization/Serialization.cpp

@@ -3345,7 +3345,7 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
 
       RawLayoutDeclAttrLayout::emitRecord(
         S.Out, S.ScratchRecord, abbrCode, attr->isImplicit(),
-        typeID, rawSize, rawAlign);
+        typeID, rawSize, rawAlign, attr->shouldMoveAsLikeType());
     }
     }
   }

test/IRGen/Inputs/module.modulemap

@@ -39,3 +39,7 @@ module PointerAuth {
 module CFBridgedType {
   header "CFBridgedType.h"
 }
+
+module RawLayoutCXX {
+  header "raw_layout_cxx.h"
+}