[SIL][Optimizer] eliminate @pack_element types

SwiftCompilerSources/Sources/AST/Type.swift

@@ -273,6 +273,10 @@ extension TypeProperties {
   public var isSILPackElementAddress: Bool {
     return rawType.bridged.isSILPackElementAddress()
   }
+
+  public var packElementTypes: TypeArray {
+    return TypeArray(bridged: rawType.bridged.getPackElementTypes())
+  }
 }
 
 public struct TypeArray : RandomAccessCollection, CustomReflectable {

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyAllocStack.swift

@@ -18,7 +18,10 @@ extension AllocStackInst : Simplifiable, SILCombineSimplifiable {
     if optimizeEnum(context) {
       return
     }
-    _ = optimizeExistential(context)
+    if optimizeExistential(context) {
+      return
+    }
+    _ = optimizePackElement(context)
   }
 }
 
@@ -302,4 +305,38 @@ private extension AllocStackInst {
     }
     return concreteType
   }
+
+  /// Replaces an alloc_stack of an Pack Element by an alloc_stack of the concrete type.
+  ///
+  /// For example:
+  /// ```
+  ///  %1 = dynamic_pack_index %0 of $Pack{T, U}
+  ///  %2 = open_pack_element %1 of <...> at <Pack{T, U}> ...
+  ///  %3 = alloc_stack $@pack_element(...) ...
+  ///  use %3
+  /// ```
+  /// is transformed to
+  /// ```
+  ///   %0 = alloc_stack $T
+  ///   use %0
+  /// ```
+  ///
+  private func optimizePackElement(_ context: SimplifyContext) -> Bool {
+    guard let concreteType = concreteTypeOfDependentPackElementArchetype else {
+      return false
+    }
+
+    let builder = Builder(before: self, context)
+    let newAlloc = builder.createAllocStack(concreteType.loweredType(in: parentFunction),
+                                            hasDynamicLifetime: hasDynamicLifetime,
+                                            isLexical: isLexical,
+                                            isFromVarDecl: isFromVarDecl,
+                                            usesMoveableValueDebugInfo: usesMoveableValueDebugInfo)    
+    let builderAfter = Builder(after: self, context)
+    let addrCast = builderAfter.createUncheckedAddrCast(from: newAlloc, to: self.type.addressType)
+    uses.replaceAll(with: addrCast, context)
+    self.replace(with: newAlloc, context)
+    return true
+  }
+
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyUncheckedAddrCast.swift

@@ -44,7 +44,21 @@ extension UncheckedAddrCastInst : OnoneSimplifiable, SILCombineSimplifiable {
     // ```
     //   %1 = vector_base_addr %0 : $*Builtin.FixedArray<N, Element>
     // ```
-    _ = optimizeVectorBaseCast(context)
+    if optimizeVectorBaseCast(context) {
+      return
+    }
+
+    // ```
+    //   %2 = dynamic_pack_index %1 of $Pack{T, U}
+    //   %3 = open_pack_element %2 of <...> at <Pack{T, U}> ...
+    //   %4 = pack_element_get %3 of %11 as $*Int
+    //   %5 = unchecked_addr_cast %4 to $*@pack_element(..)
+    // ```
+    // ->
+    // ```
+    //   %1 = unchecked_addr_cast %0 to $T
+    // ```
+    _ = optimizePackElement(context)
   }
 }
 
@@ -78,4 +92,16 @@ private extension UncheckedAddrCastInst {
     }
     return false
   }
+
+  private func optimizePackElement(_ context: SimplifyContext) -> Bool {
+    guard let concreteType = concreteTypeOfDependentPackElementArchetype else {
+      return false
+    }
+
+    let builder = Builder(before: self, context)
+    let newCast = builder.createUncheckedAddrCast(from: fromAddress,
+      to: concreteType.loweredType(in: parentFunction).addressType)
+    self.replace(with: newCast, context)
+    return true
+  }
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyWitnessMethod.swift

@@ -15,7 +15,10 @@ import SIL
 
 extension WitnessMethodInst : Simplifiable, SILCombineSimplifiable {
   func simplify(_ context: SimplifyContext) {
-    _ = tryReplaceExistentialArchetype(of: self, context)
+    if tryReplaceExistentialArchetype(of: self, context) {
+      return
+    }
+    _ = tryReplacePackElementArchetype(of: self, context)
   }
 }
 
@@ -46,3 +49,31 @@ private func tryReplaceExistentialArchetype(of witnessMethod: WitnessMethodInst,
   witnessMethod.replace(with: newWmi, context)
   return true
 }
+
+/// If the witness_method operates on a pack element archetype (`@pack_element("...")`) and the concrete
+/// type is known, replace the pack element archetype with the concrete type.
+/// For example:
+/// ```
+///   %3 = witness_method $@pack_element("...") Self, #P.foo, %2
+/// ```
+/// ->
+/// ```
+///   %3 = witness_method $ConcreteType, #P.foo, %2
+/// ```
+private func tryReplacePackElementArchetype(of witnessMethod: WitnessMethodInst, _ context: SimplifyContext) -> Bool {
+  guard let concreteType = witnessMethod.concreteTypeOfDependentPackElementArchetype else {
+    return false
+  }
+  let conf = concreteType.checkConformance(to: witnessMethod.lookupProtocol)
+  guard conf.isValid else {
+    return false
+  }
+
+  let builder = Builder(before: witnessMethod, context)
+  let newWmi = builder.createWitnessMethod(lookupType: concreteType,
+                                           conformance: conf,
+                                           member: witnessMethod.member,
+                                           methodType: witnessMethod.type)
+  witnessMethod.replace(with: newWmi, context)
+  return true
+}

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -547,6 +547,24 @@ extension Instruction {
     return nil
   }
 
+  var concreteTypeOfDependentPackElementArchetype: CanonicalType? {
+    // For simplicity only support a single type dependent operand, which is true in most of the cases anyway.
+    if let openArchetypeOp = typeDependentOperands.singleElement,
+       // Match the sequence
+       //   %1 = integer_literal $Builtin.Word, (integer value)
+       //   %2 = dynamic_pack_index %1 of $Pack{...}
+       //   %3 = open_pack_element %2 of ...
+       //   this_instruction_which_uses $*@pack_element(...)  // type-defs: %3
+       let ope = openArchetypeOp.value as? OpenPackElementInst,
+       let dpi = ope.operands.first?.value as? DynamicPackIndexInst,
+       let ili = dpi.operands.first?.value as? IntegerLiteralInst,
+       let index = ili.value
+    {
+      return dpi.indexedPackType.packElementTypes[index].canonical
+    }
+    return nil
+  }
+
   /// Returns true if a destroy of `type` must not be moved across this instruction.
   func isBarrierForDestroy(of type: Type, _ context: some Context) -> Bool {
     let instEffects = memoryEffects

include/swift/AST/ASTBridging.h

@@ -2911,6 +2911,15 @@ enum ENUM_EXTENSIBILITY_ATTR(open) BridgedMacroDefinitionKind : size_t {
   BridgedBuiltinIsolationMacro,
 };
 
+struct BridgedASTTypeArray {
+  BridgedArrayRef typeArray;
+
+  SwiftInt getCount() const { return SwiftInt(typeArray.Length); }
+
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
+  BridgedASTType getAt(SwiftInt index) const;
+};
+
 struct BridgedASTType {
   enum class TraitResult {
     IsNot,
@@ -2990,6 +2999,7 @@ struct BridgedASTType {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance checkConformance(BridgedDeclObj proto) const;
   BRIDGED_INLINE bool containsSILPackExpansionType() const;
   BRIDGED_INLINE bool isSILPackElementAddress() const;
+  BRIDGED_INLINE BridgedASTTypeArray getPackElementTypes() const;
 };
 
 class BridgedCanType {
@@ -3002,15 +3012,6 @@ class BridgedCanType {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType getRawType() const;
 };
 
-struct BridgedASTTypeArray {
-  BridgedArrayRef typeArray;
-
-  SwiftInt getCount() const { return SwiftInt(typeArray.Length); }
-
-  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
-  BridgedASTType getAt(SwiftInt index) const;
-};
-
 struct BridgedConformance {
   void * _Nullable opaqueValue;
 

include/swift/AST/ASTBridgingImpl.h

@@ -656,6 +656,10 @@ bool BridgedASTType::isSILPackElementAddress() const {
   return unbridged()->castTo<swift::SILPackType>()->isElementAddress();
 }
 
+BRIDGED_INLINE BridgedASTTypeArray BridgedASTType::getPackElementTypes() const {
+    return {unbridged()->castTo<swift::PackType>()->getElementTypes()};
+}
+
 static_assert((int)BridgedASTType::TraitResult::IsNot == (int)swift::TypeTraitResult::IsNot);
 static_assert((int)BridgedASTType::TraitResult::CanBe == (int)swift::TypeTraitResult::CanBe);
 static_assert((int)BridgedASTType::TraitResult::Is == (int)swift::TypeTraitResult::Is);