Merge pull request #73541 from slavapestov/pack-expansion-closures-6.0

[6.0] Support closures that capture opened pack element types

Author: slavapestov

include/swift/AST/ASTMangler.h

@@ -657,7 +657,8 @@ class ASTMangler : public Mangler {
   void appendClosureEntity(const AbstractClosureExpr *closure);
 
   void appendClosureComponents(Type Ty, unsigned discriminator, bool isImplicit,
-                               const DeclContext *parentContext);
+                               const DeclContext *parentContext,
+                               ArrayRef<GenericEnvironment *> capturedEnvs);
 
   void appendDefaultArgumentEntity(const DeclContext *ctx, unsigned index);
 

include/swift/AST/CaptureInfo.h

@@ -43,6 +43,7 @@ class ValueDecl;
 class FuncDecl;
 class OpaqueValueExpr;
 class VarDecl;
+class GenericEnvironment;
 
 /// CapturedValue includes both the declaration being captured, along with flags
 /// that indicate how it is captured.
@@ -140,19 +141,27 @@ class DynamicSelfType;
 /// Stores information about captured variables.
 class CaptureInfo {
   class CaptureInfoStorage final
-      : public llvm::TrailingObjects<CaptureInfoStorage, CapturedValue> {
+      : public llvm::TrailingObjects<CaptureInfoStorage,
+                                     CapturedValue,
+                                     GenericEnvironment *> {
 
     DynamicSelfType *DynamicSelf;
     OpaqueValueExpr *OpaqueValue;
-    unsigned Count;
+    unsigned NumCapturedValues;
+    unsigned NumGenericEnvironments;
+
   public:
-    explicit CaptureInfoStorage(unsigned count, DynamicSelfType *dynamicSelf,
-                                OpaqueValueExpr *opaqueValue)
-      : DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue), Count(count) { }
+    explicit CaptureInfoStorage(DynamicSelfType *dynamicSelf,
+                                OpaqueValueExpr *opaqueValue,
+                                unsigned numCapturedValues,
+                                unsigned numGenericEnvironments)
+      : DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue),
+        NumCapturedValues(numCapturedValues),
+        NumGenericEnvironments(numGenericEnvironments) { }
 
-    ArrayRef<CapturedValue> getCaptures() const {
-      return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), Count);
-    }
+    ArrayRef<CapturedValue> getCaptures() const;
+
+    ArrayRef<GenericEnvironment *> getGenericEnvironments() const;
 
     DynamicSelfType *getDynamicSelfType() const {
       return DynamicSelf;
@@ -161,6 +170,10 @@ class CaptureInfo {
     OpaqueValueExpr *getOpaqueValue() const {
       return OpaqueValue;
     }
+
+    unsigned numTrailingObjects(OverloadToken<CapturedValue>) const {
+      return NumCapturedValues;
+    }
   };
 
   enum class Flags : unsigned {
@@ -173,9 +186,11 @@ class CaptureInfo {
 public:
   /// The default-constructed CaptureInfo is "not yet computed".
   CaptureInfo() = default;
-  CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
+  CaptureInfo(ASTContext &ctx,
+              ArrayRef<CapturedValue> captures,
               DynamicSelfType *dynamicSelf, OpaqueValueExpr *opaqueValue,
-              bool genericParamCaptures);
+              bool genericParamCaptures,
+              ArrayRef<GenericEnvironment *> genericEnv=ArrayRef<GenericEnvironment*>());
 
   /// A CaptureInfo representing no captures at all.
   static CaptureInfo empty();
@@ -189,6 +204,7 @@ class CaptureInfo {
            !hasDynamicSelfCapture() && !hasOpaqueValueCapture();
   }
 
+  /// Returns all captured values and opaque expressions.
   ArrayRef<CapturedValue> getCaptures() const {
     // FIXME: Ideally, everywhere that synthesizes a function should include
     // its capture info.
@@ -206,7 +222,14 @@ class CaptureInfo {
   /// \returns true if getLocalCaptures() will return a non-empty list.
   bool hasLocalCaptures() const;
 
-  /// \returns true if the function captures any generic type parameters.
+  /// Returns all captured pack element environments.
+  ArrayRef<GenericEnvironment *> getGenericEnvironments() const {
+    assert(hasBeenComputed());
+    return StorageAndFlags.getPointer()->getGenericEnvironments();
+  }
+
+  /// \returns true if the function captures the primary generic environment
+  /// from its innermost declaration context.
   bool hasGenericParamCaptures() const {
     // FIXME: Ideally, everywhere that synthesizes a function should include
     // its capture info.

include/swift/AST/Decl.h

@@ -6053,10 +6053,6 @@ enum class PropertyWrapperSynthesizedPropertyKind {
 class VarDecl : public AbstractStorageDecl {
   friend class NamingPatternRequest;
   NamedPattern *NamingPattern = nullptr;
-  /// When the variable is declared in context of a for-in loop over the elements of 
-  /// a parameter pack, this is the opened element environment of the pack expansion
-  /// to use as the variable's context generic environment.
-  GenericEnvironment *OpenedElementEnvironment = nullptr;
 
 public:
   enum class Introducer : uint8_t {
@@ -6195,13 +6191,6 @@ class VarDecl : public AbstractStorageDecl {
   NamedPattern *getNamingPattern() const;
   void setNamingPattern(NamedPattern *Pat);
 
-  GenericEnvironment *getOpenedElementEnvironment() const {
-    return OpenedElementEnvironment;
-  }
-  void setOpenedElementEnvironment(GenericEnvironment *Env) {
-    OpenedElementEnvironment = Env;
-  }
-
   /// If this is a VarDecl that does not belong to a CaseLabelItem's pattern,
   /// return this. Otherwise, this VarDecl must belong to a CaseStmt's
   /// CaseLabelItem. In that case, return the first case label item of the first

include/swift/AST/GenericSignature.h

@@ -197,6 +197,11 @@ class GenericSignature {
   /// array of the generic parameters for the innermost generic type.
   ArrayRef<GenericTypeParamType *> getInnermostGenericParams() const;
 
+  /// Returns the depth that a generic parameter at the next level of
+  /// nesting would have. This is zero for the empty signature,
+  /// and one plus the depth of the final generic parameter otherwise.
+  unsigned getNextDepth() const;
+
   /// Retrieve the requirements.
   ArrayRef<Requirement> getRequirements() const;
 
@@ -307,6 +312,9 @@ class alignas(1 << TypeAlignInBits) GenericSignatureImpl final
     return Mem;
   }
 
+  /// Returns the depth of the last generic parameter.
+  unsigned getMaxDepth() const;
+
   /// Transform the requirements into a form where implicit Copyable and
   /// Escapable conformances are omitted, and their absence is explicitly
   /// noted.

include/swift/AST/LocalArchetypeRequirementCollector.h

@@ -0,0 +1,79 @@
+//===--- LocalArchetypeRequirementCollector.h -------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file has utility code for extending a generic signature with opened
+// existentials and shape classes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_LOCAL_ARCHETYPE_REQUIREMENT_COLLECTOR_H
+#define SWIFT_AST_LOCAL_ARCHETYPE_REQUIREMENT_COLLECTOR_H
+
+#include "swift/AST/ASTContext.h"
+#include "swift/AST/GenericSignature.h"
+#include "swift/AST/Requirement.h"
+#include "swift/AST/Types.h"
+
+namespace swift {
+
+struct LocalArchetypeRequirementCollector {
+  const ASTContext &Context;
+  GenericSignature OuterSig;
+  unsigned Depth;
+
+  /// The lists of new parameters and requirements to add to the signature.
+  SmallVector<GenericTypeParamType *, 2> Params;
+  SmallVector<Requirement, 2> Requirements;
+
+  LocalArchetypeRequirementCollector(const ASTContext &ctx, GenericSignature sig);
+
+  void addOpenedExistential(Type constraint);
+  void addOpenedElement(CanGenericTypeParamType shapeClass);
+
+  GenericTypeParamType *addParameter();
+};
+
+struct MapLocalArchetypesOutOfContext {
+  GenericSignature baseGenericSig;
+  ArrayRef<GenericEnvironment *> capturedEnvs;
+
+  MapLocalArchetypesOutOfContext(GenericSignature baseGenericSig,
+                                 ArrayRef<GenericEnvironment *> capturedEnvs)
+      : baseGenericSig(baseGenericSig), capturedEnvs(capturedEnvs) {}
+
+  Type operator()(SubstitutableType *type) const;
+};
+
+struct MapIntoLocalArchetypeContext {
+  GenericEnvironment *baseGenericEnv;
+  ArrayRef<GenericEnvironment *> capturedEnvs;
+
+  MapIntoLocalArchetypeContext(GenericEnvironment *baseGenericEnv,
+                               ArrayRef<GenericEnvironment *> capturedEnvs)
+      : baseGenericEnv(baseGenericEnv), capturedEnvs(capturedEnvs) {}
+
+  Type operator()(SubstitutableType *type) const;
+};
+
+GenericSignature buildGenericSignatureWithCapturedEnvironments(
+    ASTContext &ctx,
+    GenericSignature sig,
+    ArrayRef<GenericEnvironment *> capturedEnvs);
+
+SubstitutionMap buildSubstitutionMapWithCapturedEnvironments(
+    SubstitutionMap baseSubMap,
+    GenericSignature genericSigWithCaptures,
+    ArrayRef<GenericEnvironment *> capturedEnvs);
+
+}
+
+#endif // SWIFT_AST_LOCAL_ARCHETYPE_REQUIREMENT_COLLECTOR_H

include/swift/AST/Types.h

@@ -128,10 +128,9 @@ class RecursiveTypeProperties {
     /// This type expression contains a TypeVariableType.
     HasTypeVariable      = 0x01,
 
-    /// This type expression contains a context-dependent archetype, either a
-    /// \c PrimaryArchetypeType, \c OpenedArchetypeType,
-    /// \c ElementArchetypeType, or \c PackArchetype.
-    HasArchetype         = 0x02,
+    /// This type expression contains a PrimaryArchetypeType
+    /// or PackArchetypeType.
+    HasPrimaryArchetype  = 0x02,
 
     /// This type expression contains a GenericTypeParamType.
     HasTypeParameter     = 0x04,
@@ -171,7 +170,7 @@ class RecursiveTypeProperties {
     /// This type contains a parameterized existential type \c any P<T>.
     HasParameterizedExistential = 0x2000,
 
-    /// This type contains an ElementArchetype.
+    /// This type contains an ElementArchetypeType.
     HasElementArchetype = 0x4000,
 
     /// Whether the type is allocated in the constraint solver arena. This can
@@ -183,7 +182,7 @@ class RecursiveTypeProperties {
     /// Contains a PackType.
     HasPack = 0x10000,
 
-    /// Contains a PackArchetypeType.
+    /// Contains a PackArchetypeType. Also implies HasPrimaryArchetype.
     HasPackArchetype = 0x20000,
 
     Last_Property = HasPackArchetype
@@ -205,9 +204,10 @@ class RecursiveTypeProperties {
   /// variable?
   bool hasTypeVariable() const { return Bits & HasTypeVariable; }
 
-  /// Does a type with these properties structurally contain a
-  /// context-dependent archetype (that is, a Primary- or OpenedArchetype)?
-  bool hasArchetype() const { return Bits & HasArchetype; }
+  /// Does a type with these properties structurally contain a primary
+  /// or pack archetype? These are the archetypes instantiated from a
+  /// primary generic environment.
+  bool hasPrimaryArchetype() const { return Bits & HasPrimaryArchetype; }
 
   /// Does a type with these properties structurally contain an
   /// archetype from an opaque type declaration?
@@ -696,9 +696,23 @@ class alignas(1 << TypeAlignInBits) TypeBase
     return getRecursiveProperties().hasPlaceholder();
   }
 
-  /// Determine whether the type involves a context-dependent archetype.
+  /// Determine whether the type involves a PrimaryArchetypeType *or* a
+  /// PackArchetypeType. These are the archetypes instantiated from a
+  /// primary generic environment.
+  bool hasPrimaryArchetype() const {
+    return getRecursiveProperties().hasPrimaryArchetype();
+  }
+
+  /// Whether the type contains a PackArchetypeType.
+  bool hasPackArchetype() const {
+    return getRecursiveProperties().hasPackArchetype();
+  }
+
+  /// Determine whether the type involves a primary, pack or local archetype.
+  ///
+  /// FIXME: Replace all remaining callers with a more precise check.
   bool hasArchetype() const {
-    return getRecursiveProperties().hasArchetype();
+    return hasPrimaryArchetype() || hasLocalArchetype();
   }
 
   /// Determine whether the type involves an opened existential archetype.
@@ -727,11 +741,6 @@ class alignas(1 << TypeAlignInBits) TypeBase
     return getRecursiveProperties().hasPack();
   }
 
-  /// Whether the type contains a PackArchetypeType.
-  bool hasPackArchetype() const {
-    return getRecursiveProperties().hasPackArchetype();
-  }
-
   /// Whether the type has any flavor of pack.
   bool hasAnyPack() const {
     return hasParameterPack() || hasPack() || hasPackArchetype();

include/swift/SIL/SILFunction.h

@@ -237,6 +237,9 @@ class SILFunction
   /// The context archetypes of the function.
   GenericEnvironment *GenericEnv = nullptr;
 
+  /// Captured local generic environments.
+  ArrayRef<GenericEnvironment *> CapturedEnvs;
+
   /// The information about specialization.
   /// Only set if this function is a specialization of another function.
   const GenericSpecializationInformation *SpecializationInfo = nullptr;
@@ -1276,8 +1279,22 @@ class SILFunction
   GenericEnvironment *getGenericEnvironment() const {
     return GenericEnv;
   }
-  void setGenericEnvironment(GenericEnvironment *env) {
+
+  /// Return any captured local generic environments, currently used for pack
+  /// element environments only. After SILGen, these are rewritten into
+  /// primary archetypes.
+  ArrayRef<GenericEnvironment *> getCapturedEnvironments() const {
+    return CapturedEnvs;
+  }
+
+  void setGenericEnvironment(GenericEnvironment *env);
+
+  void setGenericEnvironment(GenericEnvironment *env,
+                             ArrayRef<GenericEnvironment *> capturedEnvs,
+                             SubstitutionMap forwardingSubs) {
     GenericEnv = env;
+    CapturedEnvs = capturedEnvs;
+    ForwardingSubMap = forwardingSubs;
   }
 
   /// Retrieve the generic signature from the generic environment of this
@@ -1306,9 +1323,30 @@ class SILFunction
   /// responsibility of the caller.
   void eraseAllBlocks();
 
-  /// Return the identity substitutions necessary to forward this call if it is
-  /// generic.
-  SubstitutionMap getForwardingSubstitutionMap();
+  /// A substitution map that sends the generic parameters of the invocation
+  /// generic signature to some combination of primar and local archetypes.
+  ///
+  /// CAUTION: If this is a SILFunction that captures pack element environments,
+  /// then at SILGen time, this is not actually the forwarding substitution map
+  /// of the SILFunction's generic environment. This is because:
+  ///
+  /// 1) The SILFunction's generic signature includes extra generic parameters,
+  ///    to model captured pack elements;
+  /// 2) The SILFunction's generic environment is the AST generic environment,
+  ///    so it's based on the original generic signature;
+  /// 3) SILGen uses this AST generic environment together with local archetypes
+  ///    for lowering SIL instructions.
+  ///
+  /// Therefore, the SILFunction's forwarding substitution map takes the extended
+  /// generic signature (1). It maps the original generic parameters to the
+  /// archetypes of (2), and the extended generic parameters to the local archetypes
+  /// of (3).
+  ///
+  /// After SILGen, all archetypes are re-instantiated inside the SIL function,
+  /// and the forwarding substitution map and generic environment then align.
+  SubstitutionMap getForwardingSubstitutionMap() const {
+    return ForwardingSubMap;
+  }
 
   /// Returns true if this SILFunction must be a defer statement.
   ///

include/swift/SIL/SILModule.h

@@ -477,6 +477,11 @@ class SILModule {
   /// This should only be the case during parsing or deserialization.
   bool hasUnresolvedLocalArchetypeDefinitions();
 
+  /// If we added any instructions that reference unresolved local archetypes
+  /// and then deleted those instructions without resolving those archetypes,
+  /// we must reclaim those unresolved local archetypes.
+  void reclaimUnresolvedLocalArchetypeDefinitions();
+
   /// Get a unique index for a struct or class field in layout order.
   ///
   /// Precondition: \p decl must be a non-resilient struct or class.