[SIL][Optimizer] eliminate @pack_element types

* SwiftCompilerSources/Sources/AST/Type.swift
* SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyAllocStack.swift
* SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyUncheckedAddrCast.swift
* SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyWitnessMethod.swift
* SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift
* include/swift/AST/ASTBridging.h
* include/swift/AST/ASTBridgingImpl.h

Author: elsakeirouz

SwiftCompilerSources/Sources/AST/Type.swift

@@ -273,6 +273,10 @@ extension TypeProperties {
   public var isSILPackElementAddress: Bool {
     return rawType.bridged.isSILPackElementAddress()
   }
+
+  public func getPackElementType(_ index: Int) -> Type {
+    return Type(bridged: rawType.bridged.getPackElementType(index));
+  }
 }
 
 public struct TypeArray : RandomAccessCollection, CustomReflectable {

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyAllocStack.swift

@@ -19,6 +19,7 @@ extension AllocStackInst : Simplifiable, SILCombineSimplifiable {
       return
     }
     _ = optimizeExistential(context)
+    _ = optimizePackElement(context)
   }
 }
 
@@ -302,4 +303,37 @@ private extension AllocStackInst {
     }
     return concreteType
   }
+
+  /// Replaces an alloc_stack of an Pack Element by an alloc_stack of the concrete type.
+  ///
+  /// For example:
+  /// ```
+  ///   %0 = alloc_stack $@pack_element(...) ...
+  ///   use %1
+  /// ```
+  /// 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)    
+    for use in uses {
+        // FIXME: Are there any specific cases to handle?
+        use.set(to: newAlloc, context)
+    }
+    context.erase(instruction: self)
+    return true
+  }
+
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyUncheckedAddrCast.swift

@@ -45,6 +45,16 @@ extension UncheckedAddrCastInst : OnoneSimplifiable, SILCombineSimplifiable {
     //   %1 = vector_base_addr %0 : $*Builtin.FixedArray<N, Element>
     // ```
     _ = optimizeVectorBaseCast(context)
+
+    // ```
+    //   %1 = unchecked_addr_cast %0 to $*pack_element("...") ...
+    // ```
+    // ->
+    // ```
+    //   %1 = unchecked_addr_cast %0 to $T
+    // ```
+    _ = optimizePackElement(context)
+
   }
 }
 
@@ -78,4 +88,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

@@ -16,6 +16,7 @@ import SIL
 extension WitnessMethodInst : Simplifiable, SILCombineSimplifiable {
   func simplify(_ context: SimplifyContext) {
     _ = tryReplaceExistentialArchetype(of: self, context)
+    _ = tryReplacePackElementArchetype(of: self, context)
   }
 }
 
@@ -46,3 +47,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.getPackElementType(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

@@ -2994,6 +2994,7 @@ struct BridgedASTType {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance checkConformance(BridgedDeclObj proto) const;
   BRIDGED_INLINE bool containsSILPackExpansionType() const;
   BRIDGED_INLINE bool isSILPackElementAddress() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType getPackElementType(SwiftInt index) const;
 };
 
 class BridgedCanType {

include/swift/AST/ASTBridgingImpl.h

@@ -656,6 +656,11 @@ bool BridgedASTType::isSILPackElementAddress() const {
   return unbridged()->castTo<swift::SILPackType>()->isElementAddress();
 }
 
+BridgedASTType BridgedASTType::getPackElementType(SwiftInt index) const {
+    ASSERT(index >= 0);
+    return {unbridged()->castTo<swift::PackType>()->getElementType(index).getPointer()};
+}
+
 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);