Merge pull request #81043 from meg-gupta/fixcow

Insert end_cow_mutation_addr for lifetime dependent values dependent on mutable addresses

Author: meg-gupta

SwiftCompilerSources/Sources/AST/Type.swift

@@ -182,6 +182,7 @@ extension TypeProperties {
   public var hasLocalArchetype: Bool { rawType.bridged.hasLocalArchetype() }
   public var isEscapable: Bool { rawType.bridged.isEscapable() }
   public var isNoEscape: Bool { rawType.bridged.isNoEscape() }
+  public var isBuiltinType: Bool { rawType.bridged.isBuiltinType() }
   public var archetypeRequiresClass: Bool { rawType.bridged.archetypeRequiresClass() }
 
   public var representationOfMetatype: AST.`Type`.MetatypeRepresentation {

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -124,11 +124,102 @@ let lifetimeDependenceScopeFixupPass = FunctionPass(
     }
     let args = scopeExtension.findArgumentDependencies()
 
+    // If the scope cannot be extended to the caller, this must be the outermost dependency level.
+    // Insert end_cow_mutation_addr if needed.
+    if args.isEmpty {
+      createEndCOWMutationIfNeeded(lifetimeDep: newLifetimeDep, context)
+    }
+
     // Redirect the dependence base to the function arguments. This may create additional mark_dependence instructions.
     markDep.redirectFunctionReturn(to: args, context)
   }
 }
 
+private extension Type {
+  func mayHaveMutableSpan(in function: Function, _ context: FunctionPassContext) -> Bool {
+    if hasArchetype {
+      return true
+    }
+    if isBuiltinType {
+      return false
+    }
+    // Only result types that are nominal can have a MutableSpan derived from an inout array access.
+    if nominal == nil {
+      return false
+    }
+    if nominal == context.swiftMutableSpan {
+      return true
+    }
+    if isStruct {
+      guard let fields = getNominalFields(in: function) else {
+        return false
+      }
+      return fields.contains { $0.mayHaveMutableSpan(in: function, context) }
+    }
+    if isTuple {
+      return tupleElements.contains { $0.mayHaveMutableSpan(in: function, context) }
+    }
+    if isEnum {
+      guard let cases = getEnumCases(in: function) else {
+        return true
+      }
+      return cases.contains { $0.payload?.mayHaveMutableSpan(in: function, context) ?? false }
+    }
+    // Classes cannot be ~Escapable, therefore cannot hold a MutableSpan.
+    if isClass {
+      return false
+    }
+    return false
+  }
+}
+
+/// Insert end_cow_mutation_addr for lifetime dependent values that maybe of type MutableSpan and depend on a mutable address.
+private func createEndCOWMutationIfNeeded(lifetimeDep: LifetimeDependence, _ context: FunctionPassContext) {
+  var scoped : ScopedInstruction
+
+  // Handle cases which generate mutable addresses: begin_access [modify] and yield &
+  switch lifetimeDep.scope {
+    case let .access(beginAccess):
+      if beginAccess.accessKind != .modify {
+        return
+      }
+      scoped = beginAccess
+    case let .yield(value):
+      let beginApply = value.definingInstruction as! BeginApplyInst
+      if value == beginApply.token {
+        return
+      }
+      if beginApply.convention(of: value as! MultipleValueInstructionResult) != .indirectInout {
+        return
+      }
+      scoped = beginApply
+    // None of the below cases can generate a mutable address.
+    case let .owned:
+      fallthrough
+    case let .borrowed:
+      fallthrough
+    case let .local:
+      fallthrough
+    case let .initialized:
+      fallthrough
+    case let .caller:
+      fallthrough
+    case let .global:
+      fallthrough
+    case let .unknown:
+      return
+  }
+
+  guard lifetimeDep.dependentValue.type.mayHaveMutableSpan(in: lifetimeDep.dependentValue.parentFunction, context) else {
+    return
+  }
+
+  for endInstruction in scoped.endInstructions {
+    let builder = Builder(before: endInstruction, context)
+    builder.createEndCOWMutationAddr(address: lifetimeDep.parentValue)
+  }
+}
+
 private extension MarkDependenceInstruction {
   /// Rewrite the mark_dependence base operand to ignore inner borrow scopes (begin_borrow, load_borrow).
   ///
@@ -194,7 +285,7 @@ private extension MarkDependenceAddrInst {
   }
 }
 
-/// A scope extension is a set of nested scopes and their owners. The owner is a value that represents ownerhip of
+/// A scope extension is a set of nested scopes and their owners. The owner is a value that represents ownership of
 /// the outermost scopes, which cannot be extended; it limits how far the nested scopes can be extended.
 private struct ScopeExtension {
   let context: FunctionPassContext

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -326,6 +326,10 @@ struct FunctionPassContext : MutatingContext {
     _bridged.getSwiftArrayDecl().getAs(NominalTypeDecl.self)
   }
 
+  var swiftMutableSpan: NominalTypeDecl {
+    _bridged.getSwiftMutableSpanDecl().getAs(NominalTypeDecl.self)
+  }
+
   func loadFunction(name: StaticString, loadCalleesRecursively: Bool) -> Function? {
     return name.withUTF8Buffer { (nameBuffer: UnsafeBufferPointer<UInt8>) in
       let nameStr = BridgedStringRef(data: nameBuffer.baseAddress, count: nameBuffer.count)

SwiftCompilerSources/Sources/Optimizer/TestPasses/MemBehaviorDumper.swift

@@ -63,6 +63,7 @@ private extension Instruction {
          is BeginAccessInst,
          is EndAccessInst,
          is EndCOWMutationInst,
+         is EndCOWMutationAddrInst,
          is CopyValueInst,
          is DestroyValueInst,
          is StrongReleaseInst,

SwiftCompilerSources/Sources/Optimizer/Utilities/AddressUtils.swift

@@ -131,7 +131,7 @@ extension AddressUseVisitor {
          is DestroyAddrInst, is DeallocStackInst, 
          is DeinitExistentialAddrInst,
          is IsUniqueInst, is MarkFunctionEscapeInst,
-         is PackElementSetInst:
+         is PackElementSetInst, is EndCOWMutationAddrInst:
       return leafAddressUse(of: operand)
 
     case is LoadInst, is LoadUnownedInst,  is LoadWeakInst, is ValueMetatypeInst, is ExistentialMetatypeInst,

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -575,6 +575,11 @@ public struct Builder {
     return notifyNew(endMutation.getAs(EndCOWMutationInst.self))
   }
 
+  public func createEndCOWMutationAddr(address: Value) -> EndCOWMutationAddrInst {
+    let endMutation = bridged.createEndCOWMutationAddr(address.bridged)
+    return notifyNew(endMutation.getAs(EndCOWMutationAddrInst.self))
+  }
+
   public func createMarkDependence(value: Value, base: Value, kind: MarkDependenceKind) -> MarkDependenceInst {
     let markDependence = bridged.createMarkDependence(value.bridged, base.bridged,
                                                       BridgedInstruction.MarkDependenceKind(rawValue: kind.rawValue)!)

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1242,6 +1242,10 @@ final public class EndCOWMutationInst : SingleValueInstruction, UnaryInstruction
   public var doKeepUnique: Bool { bridged.EndCOWMutationInst_doKeepUnique() }
 }
 
+final public class EndCOWMutationAddrInst : Instruction, UnaryInstruction {
+  public var address: Value { operand.value }
+}
+
 final public
 class ClassifyBridgeObjectInst : SingleValueInstruction, UnaryInstruction {}
 

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -215,6 +215,7 @@ public func registerSILClasses() {
   register(MoveValueInst.self)
   register(DropDeinitInst.self)
   register(EndCOWMutationInst.self)
+  register(EndCOWMutationAddrInst.self)
   register(ClassifyBridgeObjectInst.self)
   register(PartialApplyInst.self)
   register(ApplyInst.self)

docs/SIL/Instructions.md

@@ -2146,6 +2146,22 @@ not replace this reference with a not uniquely reference object.
 
 For details see [Copy-on-Write Representation](SIL.md#Copy-on-Write-Representation).
 
+### end_cow_mutation_addr
+
+```
+sil-instruction ::= 'end_cow_mutation_addr' sil-operand
+
+end_cow_mutation_addr %0 : $*T
+// %0 must be of an address $*T type
+```
+
+This instruction marks the end of mutation of an address. The address could be
+an opaque archetype, a struct, tuple or enum type and the end_cow_mutation_addr
+will apply to all members contained within it.
+It is currently only generated in cases where we maybe deriving a MutableSpan from
+`%0` since it is not possible to schedule an `end_cow_mutation` in the standard
+library automatically for Array.mutableSpan etc.
+
 ### destroy_not_escaped_closure
 
 ```

include/swift/AST/ASTBridging.h

@@ -3106,6 +3106,7 @@ struct BridgedASTType {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType getBuiltinVectorElementType() const;
   BRIDGED_INLINE bool isBuiltinFixedWidthInteger(SwiftInt width) const;
   BRIDGED_INLINE bool isOptional() const;
+  BRIDGED_INLINE bool isBuiltinType() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj getNominalOrBoundGenericNominal() const;
   BRIDGED_INLINE TraitResult canBeClass() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj getAnyNominal() const;

include/swift/AST/ASTBridgingImpl.h

@@ -556,6 +556,10 @@ bool BridgedASTType::isUnownedStorageType() const {
   return unbridged()->is<swift::UnownedStorageType>();
 }
 
+bool BridgedASTType::isBuiltinType() const {
+  return unbridged()->isBuiltinType();
+}
+
 OptionalBridgedDeclObj BridgedASTType::getNominalOrBoundGenericNominal() const {
   return {unbridged()->getNominalOrBoundGenericNominal()};
 }

include/swift/AST/KnownStdlibTypes.def

@@ -103,4 +103,6 @@ KNOWN_STDLIB_TYPE_DECL(InlineArray, NominalTypeDecl, 2)
 
 KNOWN_STDLIB_TYPE_DECL(SwiftSetting, NominalTypeDecl, 0)
 
+KNOWN_STDLIB_TYPE_DECL(MutableSpan, NominalTypeDecl, 1)
+
 #undef KNOWN_STDLIB_TYPE_DECL

include/swift/AST/Types.h

@@ -1054,6 +1054,9 @@ class alignas(1 << TypeAlignInBits) TypeBase
   bool is##NAME();
   #include "swift/AST/KnownStdlibTypes.def"
 
+  /// Check if this type is from the Builtin module.
+  bool isBuiltinType();
+
   /// Check if this type is equal to Builtin.IntN.
   bool isBuiltinIntegerType(unsigned bitWidth);
 
@@ -8199,6 +8202,10 @@ inline GenericTypeDecl *TypeBase::getAnyGeneric() {
   return getCanonicalType().getAnyGeneric();
 }
 
+inline bool TypeBase::isBuiltinType() {
+  return isa<BuiltinType>(getCanonicalType());
+}
+
 inline bool TypeBase::isBuiltinIntegerType(unsigned n) {
   if (auto intTy = dyn_cast<BuiltinIntegerType>(getCanonicalType()))
     return intTy->getWidth().isFixedWidth()

include/swift/SIL/AddressWalker.h

@@ -229,7 +229,7 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) {
         isa<PackElementSetInst>(user) || isa<PackElementGetInst>(user) ||
         isa<DeinitExistentialAddrInst>(user) || isa<LoadBorrowInst>(user) ||
         isa<TupleAddrConstructorInst>(user) || isa<DeallocPackInst>(user) ||
-        isa<MergeIsolationRegionInst>(user)) {
+        isa<MergeIsolationRegionInst>(user) || isa<EndCOWMutationAddrInst>(user)) {
       callVisitUse(op);
       continue;
     }

include/swift/SIL/SILBridging.h

@@ -1283,6 +1283,7 @@ struct BridgedBuilder{
                                           BridgedASTType::MetatypeRepresentation representation) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createEndCOWMutation(BridgedValue instance,
                                                                              bool keepUnique) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createEndCOWMutationAddr(BridgedValue instance) const;                                                                             
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createMarkDependence(
     BridgedValue value, BridgedValue base, BridgedInstruction::MarkDependenceKind dependenceKind) const;
 

include/swift/SIL/SILBridgingImpl.h

@@ -2516,6 +2516,12 @@ BridgedInstruction BridgedBuilder::createEndCOWMutation(BridgedValue instance, b
                                            keepUnique)};
 }
 
+BridgedInstruction
+BridgedBuilder::createEndCOWMutationAddr(BridgedValue instance) const {
+  return {unbridged().createEndCOWMutationAddr(regularLoc(),
+                                               instance.getSILValue())};
+}
+
 BridgedInstruction BridgedBuilder::createMarkDependence(BridgedValue value, BridgedValue base, BridgedInstruction::MarkDependenceKind kind) const {
   return {unbridged().createMarkDependence(regularLoc(), value.getSILValue(), base.getSILValue(), swift::MarkDependenceKind(kind))};
 }

include/swift/SIL/SILBuilder.h

@@ -2383,6 +2383,11 @@ class SILBuilder {
     return insert(new (getModule()) EndCOWMutationInst(getSILDebugLocation(Loc),
                                                   operand, keepUnique));
   }
+  EndCOWMutationAddrInst *createEndCOWMutationAddr(SILLocation Loc,
+                                                   SILValue operand) {
+    return insert(new (getModule()) EndCOWMutationAddrInst(
+        getSILDebugLocation(Loc), operand));
+  }
   DestroyNotEscapedClosureInst *createDestroyNotEscapedClosure(SILLocation Loc,
                                                  SILValue operand,
                                                  unsigned VerificationType) {

include/swift/SIL/SILCloner.h

@@ -3170,6 +3170,14 @@ void SILCloner<ImplClass>::visitEndCOWMutationInst(EndCOWMutationInst *Inst) {
                         getOpValue(Inst->getOperand()), Inst->doKeepUnique()));
 }
 template <typename ImplClass>
+void SILCloner<ImplClass>::visitEndCOWMutationAddrInst(
+    EndCOWMutationAddrInst *Inst) {
+  getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+  recordClonedInstruction(
+      Inst, getBuilder().createEndCOWMutationAddr(
+                getOpLocation(Inst->getLoc()), getOpValue(Inst->getOperand())));
+}
+template <typename ImplClass>
 void SILCloner<ImplClass>::visitDestroyNotEscapedClosureInst(
     DestroyNotEscapedClosureInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));

include/swift/SIL/SILInstruction.h

@@ -9580,6 +9580,14 @@ class EndCOWMutationInst
     sharedUInt8().EndCOWMutationInst.keepUnique = keepUnique;
   }
 };
+class EndCOWMutationAddrInst
+    : public UnaryInstructionBase<SILInstructionKind::EndCOWMutationAddrInst,
+                                  NonValueInstruction> {
+  friend SILBuilder;
+
+  EndCOWMutationAddrInst(SILDebugLocation debugLoc, SILValue address)
+      : UnaryInstructionBase(debugLoc, address) {}
+};
 
 /// Given an escaping closure return true iff it has a non-nil context and the
 /// context has a strong reference count greater than 1.

include/swift/SIL/SILNodes.def

@@ -528,7 +528,6 @@ ABSTRACT_VALUE_AND_INST(SingleValueInstruction, ValueBase, SILInstruction)
 
   SINGLE_VALUE_INST(EndCOWMutationInst, end_cow_mutation,
                             SingleValueInstruction, None, DoesNotRelease)
-
   SINGLE_VALUE_INST(DestroyNotEscapedClosureInst, destroy_not_escaped_closure,
                     SingleValueInstruction, MayHaveSideEffects, MayRelease)
 
@@ -908,7 +907,9 @@ NON_VALUE_INST(IncrementProfilerCounterInst, increment_profiler_counter,
 NON_VALUE_INST(MarkDependenceAddrInst, mark_dependence_addr,
                SILInstruction, None, DoesNotRelease)
 
-NODE_RANGE(NonValueInstruction, UnreachableInst, MarkDependenceAddrInst)
+NON_VALUE_INST(EndCOWMutationAddrInst, end_cow_mutation_addr,
+               SILInstruction, MayHaveSideEffects, DoesNotRelease)
+NODE_RANGE(NonValueInstruction, UnreachableInst, EndCOWMutationAddrInst)
 
 ABSTRACT_INST(MultipleValueInstruction, SILInstruction)
 MULTIPLE_VALUE_INST(BeginApplyInst, begin_apply,

include/swift/SILOptimizer/OptimizerBridging.h

@@ -216,6 +216,7 @@ struct BridgedPassContext {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDomTree getDomTree() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedPostDomTree getPostDomTree() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclObj getSwiftArrayDecl() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclObj getSwiftMutableSpanDecl() const;
 
   // AST
 

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -219,6 +219,11 @@ BridgedDeclObj BridgedPassContext::getSwiftArrayDecl() const {
   return {mod->getASTContext().getArrayDecl()};
 }
 
+BridgedDeclObj BridgedPassContext::getSwiftMutableSpanDecl() const {
+  swift::SILModule *mod = invocation->getPassManager()->getModule();
+  return {mod->getASTContext().getMutableSpanDecl()};
+}
+
 // AST
 
 SWIFT_IMPORT_UNSAFE BRIDGED_INLINE

lib/IRGen/IRGenSIL.cpp

@@ -1438,6 +1438,7 @@ class IRGenSILFunction :
   void visitIsUniqueInst(IsUniqueInst *i);
   void visitBeginCOWMutationInst(BeginCOWMutationInst *i);
   void visitEndCOWMutationInst(EndCOWMutationInst *i);
+  void visitEndCOWMutationAddrInst(EndCOWMutationAddrInst *i);
   void visitDestroyNotEscapedClosureInst(DestroyNotEscapedClosureInst *i);
   void visitDeallocStackInst(DeallocStackInst *i);
   void visitDeallocStackRefInst(DeallocStackRefInst *i);
@@ -6366,6 +6367,10 @@ void IRGenSILFunction::visitEndCOWMutationInst(EndCOWMutationInst *i) {
   setLoweredExplosion(i, v);
 }
 
+void IRGenSILFunction::visitEndCOWMutationAddrInst(EndCOWMutationAddrInst *i) {
+  // end_cow_mutation_addr is purely for SIL.
+}
+
 void IRGenSILFunction::visitDestroyNotEscapedClosureInst(
     swift::DestroyNotEscapedClosureInst *i) {
   // The closure operand is allowed to be an optional closure.

lib/SIL/IR/OperandOwnership.cpp

@@ -301,6 +301,7 @@ OPERAND_OWNERSHIP(DestroyingConsume, DestroyNotEscapedClosure)
 OPERAND_OWNERSHIP(DestroyingConsume, EndLifetime)
 OPERAND_OWNERSHIP(DestroyingConsume, BeginCOWMutation)
 OPERAND_OWNERSHIP(DestroyingConsume, EndCOWMutation)
+OPERAND_OWNERSHIP(TrivialUse, EndCOWMutationAddr)
 OPERAND_OWNERSHIP(DestroyingConsume, EndInitLetRef)
 // The move_value instruction creates a distinct lifetime.
 OPERAND_OWNERSHIP(DestroyingConsume, MoveValue)

lib/SIL/IR/SILPrinter.cpp

@@ -2763,6 +2763,9 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
       *this << "[keep_unique] ";
     *this << getIDAndType(ECMI->getOperand());
   }
+  void visitEndCOWMutationAddrInst(EndCOWMutationAddrInst *ECMI) {
+    *this << getIDAndType(ECMI->getOperand());
+  }
   void visitEndInitLetRefInst(EndInitLetRefInst *I) {
     *this << getIDAndType(I->getOperand());
   }