PackSpecialization: Fix result & type parameter handling

aidan-hall · aidan-hall · commit f6d9ef7691ed · 2025-11-14T21:01:16.000Z
We can avoid issues with loweredType failing by not calling it. We should also only get the SILType for results that we actually intend to explode. Attempting to explode packs containing generic type parameters lead to compiler crashes. It is preferable to wait until generic types have been replaced with concrete types to explode packs, since we can only pass exploded pack elements directly if we know they are loadable. @Substituted types cause the shouldExplode predicate to be unreliable for AST types, so we have restricted it to just SIL types. We have added test cases for functions that have @Substituted types.
diff --git a/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/PackSpecialization.swift b/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/PackSpecialization.swift
@@ -353,36 +353,45 @@ private struct CallSiteSpecializer {
     var directResults = resultInstruction.results.makeIterator()
     var substitutedResultTupleElements = [any Value]()
     var mappedResultPacks = self.callee.resultMap.makeIterator()
+    var indirectResultIterator = self.apply.indirectResultOperands.makeIterator()
 
     for resultInfo in self.apply.functionConvention.results {
       // We only need to handle direct and pack results, since indirect results are handled above
       if !resultInfo.isSILIndirect {
         // Direct results of the original function are mapped to direct results of the specialized function.
         substitutedResultTupleElements.append(directResults.next()!)
 
-      } else if resultInfo.type.shouldExplode {
-        // Some elements of pack results may get mapped to direct results of the specialized function.
-        // We handle those here.
-        let mapped = mappedResultPacks.next()!
-
-        let originalPackArgument = self.apply.arguments[mapped.argumentIndex]
-        let packIndices = packArgumentIndices[mapped.argumentIndex]!
-
-        for (mappedDirectResult, (packIndex, elementType)) in zip(
-          mapped.expandedElements, zip(packIndices, originalPackArgument.type.packElements)
-        )
-            where !mappedDirectResult.isSILIndirect
-        {
-
-          let result = directResults.next()!
-          let outputResultAddress = builder.createPackElementGet(
-            packIndex: packIndex, pack: originalPackArgument,
-            elementType: elementType)
-
-          builder.createStore(
-            source: result, destination: outputResultAddress,
-            // The callee is responsible for initializing return pack elements
-            ownership: storeOwnership(for: result, normal: .initialize))
+      } else {
+        guard let indirectResult = indirectResultIterator.next()?.value,
+              indirectResult.type.shouldExplode
+        else {
+          continue
+        }
+
+        do {
+          // Some elements of pack results may get mapped to direct results of the specialized function.
+          // We handle those here.
+          let mapped = mappedResultPacks.next()!
+
+
+          let packIndices = packArgumentIndices[mapped.argumentIndex]!
+
+          for (mappedDirectResult, (packIndex, elementType)) in zip(
+                mapped.expandedElements, zip(packIndices, indirectResult.type.packElements)
+              )
+              where !mappedDirectResult.isSILIndirect
+          {
+
+            let result = directResults.next()!
+            let outputResultAddress = builder.createPackElementGet(
+              packIndex: packIndex, pack: indirectResult,
+              elementType: elementType)
+
+            builder.createStore(
+              source: result, destination: outputResultAddress,
+              // The callee is responsible for initializing return pack elements
+              ownership: storeOwnership(for: result, normal: .initialize))
+          }
         }
       }
     }
@@ -499,6 +508,14 @@ private struct PackExplodedFunction {
     /// Index of this pack in the function's result type tuple.
     /// For a non-tuple result, this is 0.
     let resultIndex: Int
+    /// ResultInfo for the results produced by exploding the original result.
+    ///
+    /// NOTE: The expandedElements members of MappedResult & MappedParameter
+    /// correspond to slices of the [ResultInfo] and [ParameterInfo] arrays
+    /// produced at the same time as the ResultMap & ParameterMap respectively.
+    /// Replacing these members with integer ranges or spans referring to those
+    /// full arrays could be an easy performance optimization if this pass
+    /// becomes a bottleneck.
     let expandedElements: [ResultInfo]
   }
 
@@ -510,6 +527,7 @@ private struct PackExplodedFunction {
   /// order.
   struct MappedParameter {
     let argumentIndex: Int
+    /// ParameterInfo for the parameters produced by exploding the original parameter.
     let expandedElements: [ParameterInfo]
   }
 
@@ -631,36 +649,42 @@ private struct PackExplodedFunction {
     var resultMap = ResultMap()
     var newResults = [ResultInfo]()
 
-    var indirectResultIdx = 0
-    for (resultIndex, resultInfo) in function.convention.results.enumerated() {
-      let silType = resultInfo.type.loweredType(in: function)
-      if silType.shouldExplode {
-        let mappedResultInfos = silType.packElements.map { elem in
-          // TODO: Determine correct values for options and hasLoweredAddress
-          ResultInfo(
-            type: elem.canonicalType,
-            convention: explodedPackElementResultConvention(in: function, type: elem),
-            options: resultInfo.options,
-            hasLoweredAddresses: resultInfo.hasLoweredAddresses)
-        }
+    var indirectResultIterator = function.arguments[0..<function.convention.indirectSILResultCount]
+      .lazy.enumerated().makeIterator()
 
-        resultMap.append(
-          MappedResult(
-            argumentIndex: indirectResultIdx, resultIndex: resultIndex,
-            expandedElements: mappedResultInfos))
-        newResults += mappedResultInfos
-      } else {
-        // Leave the original result unchanged
+    for (resultIndex, resultInfo) in function.convention.results.enumerated() {
+      assert(
+        !resultInfo.isSILIndirect || !indirectResultIterator.isEmpty,
+        "There must be exactly as many indirect results in the function convention and argument list."
+      )
+
+      guard resultInfo.isSILIndirect,
+        // There should always be a value here (expressed by the assert above).
+        let (indirectResultIdx, indirectResult) = indirectResultIterator.next(),
+        indirectResult.type.shouldExplode
+      else {
         newResults.append(resultInfo)
+        continue
       }
 
-      if resultInfo.isSILIndirect {
-        indirectResultIdx += 1
+      let mappedResultInfos = indirectResult.type.packElements.map { elem in
+        ResultInfo(
+          type: elem.canonicalType,
+          convention: explodedPackElementResultConvention(in: function, type: elem),
+          options: resultInfo.options,
+          hasLoweredAddresses: resultInfo.hasLoweredAddresses)
       }
+
+      resultMap.append(
+        MappedResult(
+          argumentIndex: indirectResultIdx, resultIndex: resultIndex,
+          expandedElements: mappedResultInfos))
+      newResults += mappedResultInfos
+
     }
 
     assert(
-      indirectResultIdx == function.argumentConventions.firstParameterIndex,
+      indirectResultIterator.isEmpty,
       "We should have walked through all the indirect results, and no further.")
 
     return (newResults, resultMap)
@@ -726,31 +750,36 @@ private struct PackExplodedFunction {
     let originalValue = originalReturn.returnedValue
 
     let originalReturnTupleElements: [Value]
-    if originalValue.type.isTuple {
+    if originalValue.type.isVoid {
+      originalReturnTupleElements = []
+    } else if originalValue.type.isTuple {
       originalReturnTupleElements = [Value](
         builder.createDestructureTuple(tuple: originalValue).results)
     } else {
       originalReturnTupleElements = [originalValue]
     }
 
-    var returnValues = [any Value]()
-
     // Thread together the original and exploded direct return values.
-    var resultMapIndex = 0
-    var originalReturnIndex = 0
-    for (i, originalResult) in self.original.convention.results.enumerated()
-        where originalResult.type.shouldExplode
-                || !originalResult.isSILIndirect
-    {
-      if !resultMap.indices.contains(resultMapIndex) || resultMap[resultMapIndex].resultIndex != i {
-        returnValues.append(originalReturnTupleElements[originalReturnIndex])
-        originalReturnIndex += 1
-
-      } else {
+    let theReturnValues: [any Value]
+    do {
+      var returnValues = [any Value]()
+      // The next original result to process.
+      var resultIndex = 0
+      var originalDirectResultIterator = originalReturnTupleElements.makeIterator()
+
+      for mappedResult in resultMap {
+
+        // Collect any direct results before the next mappedResult.
+        while resultIndex < mappedResult.resultIndex {
+          if !self.original.convention.results[resultIndex].isSILIndirect {
+            returnValues.append(originalDirectResultIterator.next()!)
+          }
+          resultIndex += 1
+        }
 
-        let mapped = resultMap[resultMapIndex]
+        assert(resultIndex == mappedResult.resultIndex, "The next pack result is not skipped.")
 
-        let argumentMapping = argumentMap[mapped.argumentIndex]!
+        let argumentMapping = argumentMap[mappedResult.argumentIndex]!
         for argument in argumentMapping.arguments {
 
           switch argument.resources {
@@ -769,18 +798,26 @@ private struct PackExplodedFunction {
           }
         }
 
-        resultMapIndex += 1
+        // We have finished processing mappedResult, so step forward.
+        resultIndex += 1
       }
+
+      // Collect any remaining original direct results.
+      while let directResult = originalDirectResultIterator.next() {
+        returnValues.append(directResult)
+      }
+
+      theReturnValues = returnValues
     }
 
-    // Return the single value directly, rather than constructing a single-element tuple for it.
-    if returnValues.count == 1 {
-      builder.createReturn(of: returnValues[0])
+    // Return a single return value directly, rather than constructing a single-element tuple for it.
+    if theReturnValues.count == 1 {
+      builder.createReturn(of: theReturnValues[0])
     } else {
-      let tupleElementTypes = returnValues.map { $0.type }
+      let tupleElementTypes = theReturnValues.map { $0.type }
       let tupleType = context.getTupleType(elements: tupleElementTypes).loweredType(
         in: specialized)
-      let tuple = builder.createTuple(type: tupleType, elements: returnValues)
+      let tuple = builder.createTuple(type: tupleType, elements: theReturnValues)
       builder.createReturn(of: tuple)
     }
 
@@ -1026,10 +1063,10 @@ private func loadOwnership(for value: any Value, normal: LoadInst.LoadOwnership)
   }
 }
 
-extension TypeProperties {
+extension Type {
   /// A pack argument can explode if it contains no pack expansion types
   fileprivate var shouldExplode: Bool {
     // For now, we only attempt to explode indirect packs, since these are the most common and inefficient.
-    return isSILPack && !containsSILPackExpansionType && isSILPackElementAddress
+    return isSILPack && !hasTypeParameter && !containsSILPackExpansionType && isSILPackElementAddress
   }
 }
diff --git a/test/SILOptimizer/pack_specialization.sil b/test/SILOptimizer/pack_specialization.sil
@@ -377,6 +377,27 @@ bb0(%0 : $*Pack{any P}, %1 : $*Pack{any P}):
   return %17
 }
 
+// Handling of generic type parameters in packs.
+// We should wait until these are fully specialized before exploding.
+sil [ossa] @copy_pack_generic : $@convention(thin) <T> (@pack_guaranteed Pack{T}) -> @pack_out Pack{T} {
+bb0(%0 : $*Pack{T}, %1 : $*Pack{T}):
+  %2 = tuple ()
+  return %2
+}
+
+// CHECK-LABEL: sil [ossa] @call_copy_pack_generic : $@convention(thin) <T> (@pack_guaranteed Pack{T}) -> @pack_out Pack{T} {
+// CHECK:       bb0(%0 : $*Pack{T}, %1 : $*Pack{T})
+// CHECK:         [[CPG:%[0-9]+]] = function_ref @copy_pack_generic
+// CHECK-NEXT:    [[RESULT:%[0-9]+]] = apply [[CPG]]
+// CHECK-NEXT:    return [[RESULT]]
+// CHECK-LABEL: } // end sil function 'call_copy_pack_generic'
+sil [ossa] @call_copy_pack_generic : $@convention(thin) <T> (@pack_guaranteed Pack{T}) -> @pack_out Pack{T} {
+bb0(%0 : $*Pack{T}, %1 : $*Pack{T}):
+  %16 = function_ref @copy_pack_generic : $@convention(thin) <T> (@pack_guaranteed Pack{T}) -> @pack_out Pack{T}
+  %17 = apply %16<T>(%0, %1) : $@convention(thin) <T> (@pack_guaranteed Pack{T}) -> @pack_out Pack{T}
+  return %17
+}
+
 // INTERLEAVING PACK AND NON-PACK ARGUMENTS TESTS:
 //
 // The new function arguments and return values corresponding to pack parameters
@@ -699,7 +720,8 @@ bb0(%0 : $*Pack{Int}):
 // BAIL OUT CONDITION TESTS:
 //
 // Only perform pack specialization on functions that have indirect pack
-// parameters that contain no pack expansions.
+// parameters and/or results that contain no pack expansions or generic type
+// parameters.
 //
 // 2025-10-15: We currently only explode packs with address elements
 // (SILPackType::isElementAddress), because these are the most common (since
@@ -789,3 +811,94 @@ bb0(%0 : $*Pack{Int}):
   %2 = apply %1(%0) : $@convention(thin) (@pack_guaranteed Pack{Int}) -> ()
   return %2
 }
+
+
+sil [ossa] @indirect_result_pack : $@convention(thin) () -> @pack_out Pack{Int} {
+bb0(%0 : $*Pack{Int}):
+  %1 = tuple ()
+  return %1
+}
+
+sil [ossa] @direct_result_pack : $@convention(thin) () -> @pack_out @direct Pack{Int} {
+bb0(%0 : $*@direct Pack{Int}):
+  %1 = tuple ()
+  return %1
+}
+
+sil [ossa] @result_pack_expansion : $@convention(thin) <each A> () -> @pack_out Pack{repeat each A} {
+bb0(%0 : $*Pack{repeat each A}):
+  %1 = tuple ()
+  return %1
+}
+
+sil [ossa] @mixed_result_packs : $@convention(thin) <each A> () -> (@pack_out Pack{Int}, @pack_out @direct Pack{Int}, @pack_out Pack{repeat each A}) {
+bb0(%0 : $*Pack{Int}, %1 : $*@direct Pack{Int}, %2 : $*Pack{repeat each A}):
+  %3 = tuple ()
+  return %3
+}
+
+// CHECK-LABEL: sil [ossa] @result_call_eligibility_tests : $@convention(thin) <each A> () -> (@pack_out Pack{Int}, @pack_out @direct Pack{Int}, @pack_out Pack{repeat each A}) {
+// CHECK: bb0(%0 : $*Pack{Int}, %1 : $*@direct Pack{Int}, %2 : $*Pack{repeat each A}):
+// CHECK: function_ref @$s20indirect_result_packTf8x_n : $@convention(thin) () -> Int
+// CHECK: function_ref @direct_result_pack : $@convention(thin) () -> @pack_out @direct Pack{Int}
+// CHECK: function_ref @result_pack_expansion : $@convention(thin) <each τ_0_0> () -> @pack_out Pack{repeat each τ_0_0}
+// CHECK: function_ref @$s18mixed_result_packsTf8xnn_n : $@convention(thin) <each τ_0_0> () -> (Int, @pack_out @direct Pack{Int}, @pack_out Pack{repeat each τ_0_0})
+// CHECK-LABEL: } // end sil function 'result_call_eligibility_tests'
+sil [ossa] @result_call_eligibility_tests : $@convention(thin) <each A> () -> (@pack_out Pack{Int}, @pack_out @direct Pack{Int}, @pack_out Pack{repeat each A}) {
+bb0(%1 : $*Pack{Int}, %2 : $*@direct Pack{Int}, %3 : $*Pack{repeat each A}):
+
+  %4 = function_ref @indirect_result_pack : $@convention(thin) () -> @pack_out Pack{Int}
+  %5 = apply %4(%1) : $@convention(thin) () -> @pack_out Pack{Int}
+
+  %8 = function_ref @direct_result_pack : $@convention(thin) () -> @pack_out @direct Pack{Int}
+  %9 = apply %8(%2) : $@convention(thin) () -> @pack_out @direct Pack{Int}
+
+  %10 = function_ref @result_pack_expansion : $@convention(thin) <each A> () -> @pack_out Pack{repeat each A}
+  %11 = apply %10<Pack{repeat each A}>(%3) : $@convention(thin) <each A> () -> @pack_out Pack{repeat each A}
+
+  %12 = function_ref @mixed_result_packs : $@convention(thin) <each A> () -> (@pack_out Pack{Int}, @pack_out @direct Pack{Int}, @pack_out Pack{repeat each A})
+  %13 = apply %12<Pack{repeat each A}>(%1, %2, %3) : $@convention(thin) <each A> () -> (@pack_out Pack{Int}, @pack_out @direct Pack{Int}, @pack_out Pack{repeat each A})
+
+  %99 = tuple ()
+  return %99
+}
+
+// @substituted TYPE TESTS:
+//
+// PackSubstitution must correctly identify the types of Pack arguments in the presence of @substituted types.
+
+// CHECK-LABEL: sil shared [ossa] @$s25substitute_parameter_typeTf8x_n : $@convention(thin) @substituted <τ_0_0> (Int32) -> () for <Int32> {
+// CHECK: bb0(%0 : $Int32):
+// CHECK-LABEL: } // end sil function '$s25substitute_parameter_typeTf8x_n'
+sil [ossa] @substitute_parameter_type : $@convention(thin) @substituted <T> (@pack_guaranteed Pack{T}) -> () for <Int32> {
+bb0(%0 : $*Pack{Int32}):
+  %99 = tuple ()
+  return %99
+}
+
+// CHECK-LABEL: sil [ossa] @call_substitute_parameter_type : $@convention(thin) (@pack_guaranteed Pack{Int32}) -> () {
+// CHECK-LABEL: } // end sil function 'call_substitute_parameter_type'
+sil [ossa] @call_substitute_parameter_type : $@convention(thin) (@pack_guaranteed Pack{Int32}) -> () {
+bb0(%0 : $*Pack{Int32}):
+  %1 = function_ref @substitute_parameter_type : $@convention(thin) @substituted <T> (@pack_guaranteed Pack{T}) -> () for <Int32>
+  %2 = apply %1(%0) : $@convention(thin) @substituted <T> (@pack_guaranteed Pack{T}) -> () for <Int32>
+  return %2
+}
+
+// CHECK-LABEL: sil shared [ossa] @$s22substitute_result_typeTf8x_n : $@convention(thin) @substituted <τ_0_0> () -> Int32 for <Int32> {
+// CHECK: bb0:
+// CHECK-LABEL: } // end sil function '$s22substitute_result_typeTf8x_n'
+sil [ossa] @substitute_result_type : $@convention(thin) @substituted <T> () -> @pack_out Pack{T} for <Int32> {
+bb0(%0 : $*Pack{Int32}):
+  %99 = tuple ()
+  return %99
+}
+
+// CHECK-LABEL: sil [ossa] @call_substitute_result_type : $@convention(thin) () -> @pack_out Pack{Int32} {
+// CHECK-LABEL: } // end sil function 'call_substitute_result_type'
+sil [ossa] @call_substitute_result_type : $@convention(thin) () -> @pack_out Pack{Int32} {
+bb0(%0 : $*Pack{Int32}):
+  %1 = function_ref @substitute_result_type : $@convention(thin) @substituted <T> () -> @pack_out Pack{T} for <Int32>
+  %2 = apply %1(%0) : $@convention(thin) @substituted <T> () -> @pack_out Pack{T} for <Int32>
+  return %2
+}
