Merge pull request swiftlang#81678 from meg-gupta/fixsubst

meg-gupta · web-flow · commit 1998f416a591 · 2025-05-27T12:47:29.000-07:00
Fix use-after-free on substituting function type involving conditional ~Escapable with Escapable type
diff --git a/include/swift/AST/ExtInfo.h b/include/swift/AST/ExtInfo.h
@@ -740,12 +740,19 @@ class ASTExtInfoBuilder {
                              globalActor, thrownError, lifetimeDependencies);
   }
 
+  /// \p lifetimeDependencies should be arena allocated and not a temporary
+  /// Function types are allocated on the are arena and their ExtInfo should be
+  /// valid throughout their lifetime.
   [[nodiscard]] ASTExtInfoBuilder withLifetimeDependencies(
       llvm::ArrayRef<LifetimeDependenceInfo> lifetimeDependencies) const {
     return ASTExtInfoBuilder(bits, clangTypeInfo, globalActor, thrownError,
                              lifetimeDependencies);
   }
 
+  [[nodiscard]] ASTExtInfoBuilder withLifetimeDependencies(
+      SmallVectorImpl<LifetimeDependenceInfo> lifetimeDependencies) const =
+      delete;
+
   [[nodiscard]]
   ASTExtInfoBuilder withIsolation(FunctionTypeIsolation isolation) const {
     return ASTExtInfoBuilder(
@@ -920,11 +927,18 @@ class ASTExtInfo {
       .build();
   }
 
+  /// \p lifetimeDependencies should be arena allocated and not a temporary
+  /// Function types are allocated on the are arena and their ExtInfo should be
+  /// valid throughout their lifetime.
   [[nodiscard]] ASTExtInfo withLifetimeDependencies(
       ArrayRef<LifetimeDependenceInfo> lifetimeDependencies) const {
     return builder.withLifetimeDependencies(lifetimeDependencies).build();
   }
 
+  [[nodiscard]] ASTExtInfo withLifetimeDependencies(
+      SmallVectorImpl<LifetimeDependenceInfo> lifetimeDependencies) const =
+      delete;
+
   void Profile(llvm::FoldingSetNodeID &ID) const { builder.Profile(ID); }
 
   bool isEqualTo(ASTExtInfo other, bool useClangTypes) const {
@@ -1227,11 +1241,19 @@ class SILExtInfoBuilder {
     return SILExtInfoBuilder(bits, ClangTypeInfo(type).getCanonical(),
                              lifetimeDependencies);
   }
+
+  /// \p lifetimeDependencies should be arena allocated and not a temporary
+  /// Function types are allocated on the are arena and their ExtInfo should be
+  /// valid throughout their lifetime.
   [[nodiscard]] SILExtInfoBuilder withLifetimeDependencies(
       ArrayRef<LifetimeDependenceInfo> lifetimeDependenceInfo) const {
     return SILExtInfoBuilder(bits, clangTypeInfo, lifetimeDependenceInfo);
   }
 
+  [[nodiscard]] ASTExtInfoBuilder withLifetimeDependencies(
+      SmallVectorImpl<LifetimeDependenceInfo> lifetimeDependencies) const =
+      delete;
+
   void Profile(llvm::FoldingSetNodeID &ID) const {
     ID.AddInteger(bits);
     ID.AddPointer(clangTypeInfo.getType());
@@ -1368,11 +1390,17 @@ class SILExtInfo {
     return builder.withUnimplementable(isUnimplementable).build();
   }
 
+  /// \p lifetimeDependencies should be arena allocated and not a temporary
+  /// Function types are allocated on the are arena and their ExtInfo should be
+  /// valid throughout their lifetime.
   SILExtInfo withLifetimeDependencies(
       ArrayRef<LifetimeDependenceInfo> lifetimeDependencies) const {
     return builder.withLifetimeDependencies(lifetimeDependencies);
   }
 
+  SILExtInfo withLifetimeDependencies(SmallVectorImpl<LifetimeDependenceInfo>
+                                          lifetimeDependencies) const = delete;
+
   void Profile(llvm::FoldingSetNodeID &ID) const { builder.Profile(ID); }
 
   bool isEqualTo(SILExtInfo other, bool useClangTypes) const {
diff --git a/include/swift/AST/TypeTransform.h b/include/swift/AST/TypeTransform.h
@@ -358,7 +358,8 @@ case TypeKind::Id:
             }
           });
         if (didRemoveLifetimeDependencies) {
-          extInfo = extInfo.withLifetimeDependencies(substDependenceInfos);
+          extInfo = extInfo.withLifetimeDependencies(
+              ctx.AllocateCopy(substDependenceInfos));
         }
       }
 
@@ -939,7 +940,8 @@ case TypeKind::Id:
           });
 
         if (didRemoveLifetimeDependencies) {
-          extInfo = extInfo->withLifetimeDependencies(substDependenceInfos);
+          extInfo = extInfo->withLifetimeDependencies(
+              ctx.AllocateCopy(substDependenceInfos));
         }
       }
 
diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h
@@ -5210,12 +5210,12 @@ class SILFunctionType final
   /// Given an existing ExtInfo, and a set of interface parameters and results
   /// destined for a new SILFunctionType, return a new ExtInfo with only the
   /// lifetime dependencies relevant after substitution.
-  static ExtInfo
-  getSubstLifetimeDependencies(GenericSignature genericSig,
-                               ExtInfo origExtInfo,
-                               ArrayRef<SILParameterInfo> params,
-                               ArrayRef<SILYieldInfo> yields,
-                               ArrayRef<SILResultInfo> results);
+  static ExtInfo getSubstLifetimeDependencies(GenericSignature genericSig,
+                                              ExtInfo origExtInfo,
+                                              ASTContext &context,
+                                              ArrayRef<SILParameterInfo> params,
+                                              ArrayRef<SILYieldInfo> yields,
+                                              ArrayRef<SILResultInfo> results);
 
   /// Return a structurally-identical function type with a slightly tweaked
   /// ExtInfo.
diff --git a/lib/SIL/IR/SILFunctionType.cpp b/lib/SIL/IR/SILFunctionType.cpp
@@ -59,12 +59,10 @@ SILType SILFunctionType::substInterfaceType(SILModule &M,
   return interfaceType;
 }
 
-SILFunctionType::ExtInfo
-SILFunctionType::getSubstLifetimeDependencies(GenericSignature genericSig,
-                                              ExtInfo origExtInfo,
-                                              ArrayRef<SILParameterInfo> params,
-                                              ArrayRef<SILYieldInfo> yields,
-                                              ArrayRef<SILResultInfo> results) {
+SILFunctionType::ExtInfo SILFunctionType::getSubstLifetimeDependencies(
+    GenericSignature genericSig, ExtInfo origExtInfo, ASTContext &context,
+    ArrayRef<SILParameterInfo> params, ArrayRef<SILYieldInfo> yields,
+    ArrayRef<SILResultInfo> results) {
   if (origExtInfo.getLifetimeDependencies().empty()) {
     return origExtInfo;
   }
@@ -87,7 +85,8 @@ SILFunctionType::getSubstLifetimeDependencies(GenericSignature genericSig,
       }
     });
   if (didRemoveLifetimeDependencies) {
-    return origExtInfo.withLifetimeDependencies(substLifetimeDependencies);
+    return origExtInfo.withLifetimeDependencies(
+        context.AllocateCopy(substLifetimeDependencies));
   }
   return origExtInfo;
 }
@@ -131,10 +130,10 @@ CanSILFunctionType SILFunctionType::getUnsubstitutedType(SILModule &M) const {
 
   auto signature = isPolymorphic() ? getInvocationGenericSignature()
                                    : CanGenericSignature();
-  
-  auto extInfo = getSubstLifetimeDependencies(signature, getExtInfo(),
-                                              params, yields, results);
-  
+
+  auto extInfo = getSubstLifetimeDependencies(
+      signature, getExtInfo(), getASTContext(), params, yields, results);
+
   return SILFunctionType::get(signature,
                               extInfo,
                               getCoroutineKind(),
@@ -2855,7 +2854,7 @@ static CanSILFunctionType getSILFunctionType(
           .withSendable(isSendable)
           .withAsync(isAsync)
           .withUnimplementable(unimplementable)
-          .withLifetimeDependencies(loweredLifetimes)
+          .withLifetimeDependencies(TC.Context.AllocateCopy(loweredLifetimes))
           .build();
 
   return SILFunctionType::get(genericSig, silExtInfo, coroutineKind,
diff --git a/lib/SIL/IR/SILTypeSubstitution.cpp b/lib/SIL/IR/SILTypeSubstitution.cpp
@@ -240,11 +240,10 @@ class SILTypeSubstituter :
     auto genericSig = IFS.shouldSubstituteOpaqueArchetypes()
                         ? origType->getInvocationGenericSignature()
                         : nullptr;
-                        
-    extInfo = SILFunctionType::getSubstLifetimeDependencies(genericSig, extInfo,
-                                                            substParams,
-                                                            substYields,
-                                                            substResults);
+
+    extInfo = SILFunctionType::getSubstLifetimeDependencies(
+        genericSig, extInfo, TC.Context, substParams, substYields,
+        substResults);
 
     return SILFunctionType::get(genericSig, extInfo,
                                 origType->getCoroutineKind(),
diff --git a/test/Serialization/Inputs/ne_types.swift b/test/Serialization/Inputs/ne_types.swift
@@ -0,0 +1,98 @@
+@_addressableForDependencies
+public struct Something {
+  public struct View<T: BitwiseCopyable> : ~Copyable, ~Escapable {
+    var ptr: UnsafeBufferPointer<T>
+    
+    @lifetime(borrow ptr)
+    public init(ptr: borrowing UnsafeBufferPointer<T>) {
+      self.ptr = copy ptr
+    }
+    
+    public var span: Span<T> {
+      @lifetime(borrow self)
+      borrowing get {
+        Span(_unsafeElements: ptr)
+      }
+    }
+  }
+  
+  public struct MutableView<T: BitwiseCopyable> : ~Copyable, ~Escapable {
+    var ptr: UnsafeMutableBufferPointer<T>
+    
+    @lifetime(borrow ptr)
+    public init(ptr: borrowing UnsafeMutableBufferPointer<T>) {
+      self.ptr = copy ptr
+    }
+    
+    public var mutableSpan: MutableSpan<T> {
+      @lifetime(&self)
+      mutating get {
+        MutableSpan(_unsafeElements: ptr)
+      }
+    }
+  }
+  
+  var ptr: UnsafeMutableRawBufferPointer
+  public init(ptr: UnsafeMutableRawBufferPointer) {
+    self.ptr = ptr
+  }
+  
+  @lifetime(borrow self)
+  public func view<T>(of type: T.Type = T.self) -> View<T> {
+    let tp = ptr.assumingMemoryBound(to: T.self)
+    return __overrideLifetime(View(ptr: .init(tp)), borrowing: self)
+  }
+  
+  @lifetime(&self)
+  public mutating func mutableView<T>(of type: T.Type = T.self) -> MutableView<T> {
+    let tp = ptr.assumingMemoryBound(to: T.self)
+    return __overrideLifetime(MutableView(ptr: tp), mutating: &self)
+  }
+}
+
+@unsafe
+@_unsafeNonescapableResult
+@_alwaysEmitIntoClient
+@_transparent
+@lifetime(borrow source)
+public func __overrideLifetime<
+  T: ~Copyable & ~Escapable, U: ~Copyable & ~Escapable
+>(
+  _ dependent: consuming T, borrowing source: borrowing U
+) -> T {
+  dependent
+}
+
+/// Unsafely discard any lifetime dependency on the `dependent` argument. Return
+/// a value identical to `dependent` that inherits all lifetime dependencies from
+/// the `source` argument.
+@unsafe
+@_unsafeNonescapableResult
+@_alwaysEmitIntoClient
+@_transparent
+@lifetime(copy source)
+public func __overrideLifetime<
+  T: ~Copyable & ~Escapable, U: ~Copyable & ~Escapable
+>(
+  _ dependent: consuming T, copying source: borrowing U
+) -> T {
+  dependent
+}
+
+/// Unsafely discard any lifetime dependency on the `dependent` argument.
+/// Return a value identical to `dependent` with a lifetime dependency
+/// on the caller's exclusive borrow scope of the `source` argument.
+@unsafe
+@_unsafeNonescapableResult
+@_alwaysEmitIntoClient
+@_transparent
+@lifetime(&source)
+public func __overrideLifetime<
+  T: ~Copyable & ~Escapable, U: ~Copyable & ~Escapable
+>(
+  _ dependent: consuming T,
+  mutating source: inout U
+) -> T {
+  dependent
+}
+
diff --git a/test/Serialization/non_escapable_subst_test.swift b/test/Serialization/non_escapable_subst_test.swift
@@ -0,0 +1,26 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -o %t  %S/Inputs/ne_types.swift \
+// RUN: -enable-experimental-feature LifetimeDependence \
+// RUN: -enable-experimental-feature AddressableTypes \
+// RUN: -enable-library-evolution \
+// RUN: -emit-module-path %t/ne_types.swiftmodule 
+
+// RUN: %target-swift-frontend -emit-silgen -I %t %s \
+// RUN: -enable-experimental-feature LifetimeDependence
+
+// REQUIRES: swift_feature_LifetimeDependence
+// REQUIRES: swift_feature_AddressableTypes
+
+import ne_types
+
+// Ensure no memory crashes during SILGen
+
+struct NonEscapableFrameworkThingTests {
+  func example() async throws {
+    let ptr = UnsafeMutableRawBufferPointer.allocate(byteCount: 40, alignment: 1)
+    defer { ptr.deallocate() }
+    var something = Something(ptr: ptr)
+    var mutableView = something.mutableView(of: Float32.self)
+    var mutableSpan = mutableView.mutableSpan
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -358,7 +358,8 @@ case TypeKind::Id:`
`358`	`358`	`}`
`359`	`359`	`});`
`360`	`360`	`if (didRemoveLifetimeDependencies) {`
`361`		`- extInfo = extInfo.withLifetimeDependencies(substDependenceInfos);`
	`361`	`+ extInfo = extInfo.withLifetimeDependencies(`
	`362`	`+ ctx.AllocateCopy(substDependenceInfos));`
`362`	`363`	`}`
`363`	`364`	`}`
`364`	`365`
`@@ -939,7 +940,8 @@ case TypeKind::Id:`
`939`	`940`	`});`
`940`	`941`
`941`	`942`	`if (didRemoveLifetimeDependencies) {`
`942`		`- extInfo = extInfo->withLifetimeDependencies(substDependenceInfos);`
	`943`	`+ extInfo = extInfo->withLifetimeDependencies(`
	`944`	`+ ctx.AllocateCopy(substDependenceInfos));`
`943`	`945`	`}`
`944`	`946`	`}`
`945`	`947`