PackSpecialization: Fix result & type parameter handling

Refactor certain functions to make them simpler. and avoid calling
AST.Type.loweredType, which can fail. Instead, access the types of the
function's (SIL) arguments directly.

Correctly handle exploding packs that contain generic or opaque types by using
AST.Type.mapOutOfEnvironment().

@Substituted types cause the shouldExplode predicate to be unreliable for AST
types, so restrict it to just SIL.Type. Add test cases for functions that have
@Substituted types.

Re-enable PackSpecialization in FunctionPass pipeline.

Add a check to avoid emitting a destructure_tuple of the original function's
return tuple when it is void/().

Author: aidan-hall

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]
   }
 
@@ -600,10 +618,9 @@ private struct PackExplodedFunction {
 
     for (argument, parameterInfo) in zip(function.parameters, function.convention.parameters) {
       if argument.type.shouldExplode {
-
         let mappedParameterInfos = argument.type.packElements.map { element in
           ParameterInfo(
-            type: element.canonicalType,
+            type: element.hasArchetype ? element.rawType.mapOutOfEnvironment().canonical : element.canonicalType,
             convention: explodedPackElementArgumentConvention(
               pack: parameterInfo, type: element, in: function),
             options: parameterInfo.options,
@@ -631,36 +648,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 { element in
+        ResultInfo(
+          type: element.hasArchetype ? element.rawType.mapOutOfEnvironment().canonical : element.canonicalType,
+          convention: explodedPackElementResultConvention(in: function, type: element),
+          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 +749,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 +797,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,7 +1062,7 @@ 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.

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -527,7 +527,7 @@ void addFunctionPasses(SILPassPipelinePlan &P,
   // of embedded Swift.
   if (!P.getOptions().EmbeddedSwift) {
     P.addGenericSpecializer();
-    // P.addPackSpecialization();
+    P.addPackSpecialization();
     // Run devirtualizer after the specializer, because many
     // class_method/witness_method instructions may use concrete types now.
     P.addDevirtualizer();