[ConstraintSystem] Create opened pack element generic environments lazily

in the constraint system.

Author: hborla

include/swift/Sema/ConstraintLocator.h

@@ -690,21 +690,6 @@ class LocatorPathElt::PackElement final : public StoredIntegerElement<1> {
   }
 };
 
-class LocatorPathElt::OpenedPackElement final
-    : public StoredPointerElement<GenericEnvironment> {
-public:
-  OpenedPackElement(GenericEnvironment *env)
-      : StoredPointerElement(PathElementKind::OpenedPackElement, env) {}
-
-  GenericEnvironment *getGenericEnvironment() const {
-    return getStoredPointer();
-  }
-
-  static bool classof(const LocatorPathElt *elt) {
-    return elt->getKind() == PathElementKind::OpenedPackElement;
-  }
-};
-
 class LocatorPathElt::KeyPathComponent final : public StoredIntegerElement<1> {
 public:
   KeyPathComponent(unsigned index)

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -195,9 +195,6 @@ CUSTOM_LOCATOR_PATH_ELT(PackType)
 /// An element of a pack type - the T in <T, U, V, ...>
 CUSTOM_LOCATOR_PATH_ELT(PackElement)
 
-/// Stores the generic environment for an opened pack element.
-CUSTOM_LOCATOR_PATH_ELT(OpenedPackElement)
-
 /// The shape of a parameter pack.
 SIMPLE_LOCATOR_PATH_ELT(PackShape)
 

include/swift/Sema/ConstraintSystem.h

@@ -1494,6 +1494,10 @@ class Solution {
   llvm::DenseMap<ConstraintLocator *, OpenedArchetypeType *>
     OpenedExistentialTypes;
 
+  /// The pack expansion environment that can open pack elements for
+  /// a given locator.
+  llvm::DenseMap<ConstraintLocator *, UUID> PackExpansionEnvironments;
+
   /// The locators of \c Defaultable constraints whose defaults were used.
   llvm::SmallPtrSet<ConstraintLocator *, 2> DefaultedConstraints;
 
@@ -3021,6 +3025,8 @@ class ConstraintSystem {
   llvm::SmallMapVector<ConstraintLocator *, OpenedArchetypeType *, 4>
       OpenedExistentialTypes;
 
+  llvm::SmallMapVector<ConstraintLocator *, UUID, 4> PackExpansionEnvironments;
+
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
       resultBuilderTransformed;
@@ -3496,6 +3502,9 @@ class ConstraintSystem {
     /// The length of \c OpenedExistentialTypes.
     unsigned numOpenedExistentialTypes;
 
+    /// The length of \c PackExpansionEnvironments.
+    unsigned numPackExpansionEnvironments;
+
     /// The length of \c DefaultedConstraints.
     unsigned numDefaultedConstraints;
 
@@ -3974,6 +3983,12 @@ class ConstraintSystem {
   std::pair<Type, OpenedArchetypeType *> openExistentialType(
       Type type, ConstraintLocator *locator);
 
+  /// Add the given pack expansion as an opened pack element environment.
+  void addPackElementEnvironment(PackExpansionExpr *expr);
+
+  /// Get the opened element generic environment for the given locator.
+  GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator);
+
   /// Retrieve the constraint locator for the given anchor and
   /// path, uniqued and automatically infer the summary flags
   ConstraintLocator *
@@ -6124,12 +6139,12 @@ class HandlePlaceholderType {
 class OpenPackElementType {
   ConstraintSystem &cs;
   ConstraintLocator *locator;
-  GenericEnvironment *elementEnv;
+  PackExpansionExpr *elementEnv;
 
 public:
   explicit OpenPackElementType(ConstraintSystem &cs,
                                const ConstraintLocatorBuilder &locator,
-                               GenericEnvironment *elementEnv)
+                               PackExpansionExpr *elementEnv)
       : cs(cs), elementEnv(elementEnv) {
     this->locator = cs.getConstraintLocator(locator);
   }
@@ -6144,8 +6159,6 @@ class OpenPackElementType {
     auto *elementType = cs.createTypeVariable(locator,
                                               TVO_CanBindToHole |
                                               TVO_CanBindToNoEscape);
-    auto elementLoc = cs.getConstraintLocator(locator,
-        LocatorPathElt::OpenedPackElement(elementEnv));
 
     // If we're opening a pack element from an explicit type repr,
     // set the type repr types in the constraint system for generating
@@ -6156,7 +6169,7 @@ class OpenPackElementType {
     }
 
     cs.addConstraint(ConstraintKind::PackElementOf, elementType,
-                     packType, elementLoc);
+                     packType, cs.getConstraintLocator(elementEnv));
     return elementType;
   }
 };

lib/Sema/CSApply.cpp

@@ -3818,6 +3818,10 @@ namespace {
     }
 
     Expr *visitPackExpansionExpr(PackExpansionExpr *expr) {
+      auto *locator = cs.getConstraintLocator(expr);
+      auto *environment = cs.getPackElementEnvironment(locator);
+      expr->setGenericEnvironment(environment);
+
       return simplifyExprType(expr);
     }
 

lib/Sema/CSBindings.cpp

@@ -1414,10 +1414,6 @@ void PotentialBindings::infer(Constraint *constraint) {
     auto elementType = CS.simplifyType(constraint->getFirstType());
     auto packType = CS.simplifyType(constraint->getSecondType());
 
-    auto *loc = constraint->getLocator();
-    auto openedElement  =
-        loc->getLastElementAs<LocatorPathElt::OpenedPackElement>();
-
     if (elementType->isTypeVariableOrMember() && packType->isTypeVariableOrMember())
       break;
 
@@ -1428,7 +1424,7 @@ void PotentialBindings::infer(Constraint *constraint) {
       // Produce a potential binding to the opened element archetype corresponding
       // to the pack type.
       packType = packType->mapTypeOutOfContext();
-      auto *elementEnv = openedElement->getGenericEnvironment();
+      auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator());
       auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
       addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});
 

lib/Sema/CSGen.cpp

@@ -845,8 +845,8 @@ namespace {
     /// found during our walk.
     llvm::MapVector<UnresolvedMemberExpr *, Type> UnresolvedBaseTypes;
 
-    /// A stack of pack element generic environments.
-    llvm::SmallVector<GenericEnvironment *, 2> PackElementEnvironments;
+    /// A stack of pack expansions that can open pack elements.
+    llvm::SmallVector<PackExpansionExpr *, 2> PackElementEnvironments;
 
     /// Returns false and emits the specified diagnostic if the member reference
     /// base is a nil literal. Returns true otherwise.
@@ -1088,8 +1088,9 @@ namespace {
 
     ConstraintSystem &getConstraintSystem() const { return CS; }
 
-    void addPackElementEnvironment(GenericEnvironment *env) {
-      PackElementEnvironments.push_back(env);
+    void addPackElementEnvironment(PackExpansionExpr *expr) {
+      CS.addPackElementEnvironment(expr);
+      PackElementEnvironments.push_back(expr);
     }
 
     virtual Type visitErrorExpr(ErrorExpr *E) {
@@ -1424,7 +1425,7 @@ namespace {
       const auto placeholderHandler = HandlePlaceholderType(CS, locator);
 
       // Add a PackElementOf constraint for 'each T' type reprs.
-      GenericEnvironment *elementEnv = nullptr;
+      PackExpansionExpr *elementEnv = nullptr;
       if (!PackElementEnvironments.empty()) {
         options |= TypeResolutionFlags::AllowPackReferences;
         elementEnv = PackElementEnvironments.back();
@@ -1694,7 +1695,7 @@ namespace {
       auto options =
           TypeResolutionOptions(TypeResolverContext::InExpression);
       for (auto specializationArg : specializationArgs) {
-        GenericEnvironment *elementEnv = nullptr;
+        PackExpansionExpr *elementEnv = nullptr;
         if (!PackElementEnvironments.empty()) {
           options |= TypeResolutionFlags::AllowPackReferences;
           elementEnv = PackElementEnvironments.back();
@@ -1764,7 +1765,7 @@ namespace {
           auto options =
               TypeResolutionOptions(TypeResolverContext::InExpression);
           for (size_t i = 0, e = specializations.size(); i < e; ++i) {
-            GenericEnvironment *elementEnv = nullptr;
+            PackExpansionExpr *elementEnv = nullptr;
             if (!PackElementEnvironments.empty()) {
               options |= TypeResolutionFlags::AllowPackReferences;
               elementEnv = PackElementEnvironments.back();
@@ -3008,8 +3009,7 @@ namespace {
     }
 
     Type visitPackExpansionExpr(PackExpansionExpr *expr) {
-      auto *elementEnv = expr->getGenericEnvironment();
-      assert(PackElementEnvironments.back() == elementEnv);
+      assert(PackElementEnvironments.back() == expr);
       PackElementEnvironments.pop_back();
 
       auto *patternLoc =
@@ -3018,10 +3018,8 @@ namespace {
                                              TVO_CanBindToPack |
                                              TVO_CanBindToHole);
       auto elementResultType = CS.getType(expr->getPatternExpr());
-      auto *elementLoc = CS.getConstraintLocator(
-          expr, LocatorPathElt::OpenedPackElement(elementEnv));
       CS.addConstraint(ConstraintKind::PackElementOf, elementResultType,
-                       patternTy, elementLoc);
+                       patternTy, CS.getConstraintLocator(expr));
 
       auto *shapeLoc =
           CS.getConstraintLocator(expr, ConstraintLocator::PackShape);
@@ -4063,7 +4061,7 @@ namespace {
       }
 
       if (auto *expansion = dyn_cast<PackExpansionExpr>(expr)) {
-        CG.addPackElementEnvironment(expansion->getGenericEnvironment());
+        CG.addPackElementEnvironment(expansion);
       }
 
       return Action::Continue(expr);

lib/Sema/CSSolver.cpp

@@ -267,6 +267,11 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     OpenedExistentialTypes.insert(openedExistential);
   }
 
+  // Register the solutions's pack expansion environments.
+  for (const auto &expansion : solution.PackExpansionEnvironments) {
+    PackExpansionEnvironments.insert(expansion);
+  }
+
   // Register the defaulted type variables.
   DefaultedConstraints.insert(solution.DefaultedConstraints.begin(),
                               solution.DefaultedConstraints.end());
@@ -585,6 +590,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numArgumentMatchingChoices = cs.argumentMatchingChoices.size();
   numOpenedTypes = cs.OpenedTypes.size();
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
+  numPackExpansionEnvironments = cs.PackExpansionEnvironments.size();
   numDefaultedConstraints = cs.DefaultedConstraints.size();
   numAddedNodeTypes = cs.addedNodeTypes.size();
   numAddedKeyPathComponentTypes = cs.addedKeyPathComponentTypes.size();
@@ -662,6 +668,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any opened existential types.
   truncate(cs.OpenedExistentialTypes, numOpenedExistentialTypes);
 
+  // Remove any pack expansion environments.
+  truncate(cs.PackExpansionEnvironments, numPackExpansionEnvironments);
+
   // Remove any defaulted type variables.
   truncate(cs.DefaultedConstraints, numDefaultedConstraints);
 

lib/Sema/ConstraintLocator.cpp

@@ -97,7 +97,6 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   case ConstraintLocator::SyntacticElement:
   case ConstraintLocator::PackType:
   case ConstraintLocator::PackElement:
-  case ConstraintLocator::OpenedPackElement:
   case ConstraintLocator::PackShape:
   case ConstraintLocator::PackExpansionPattern:
   case ConstraintLocator::PatternBindingElement:
@@ -449,11 +448,6 @@ void LocatorPathElt::dump(raw_ostream &out) const {
     break;
   }
 
-  case ConstraintLocator::OpenedPackElement: {
-    out << "opened pack element";
-    break;
-  }
-
   case ConstraintLocator::PackShape: {
     out << "pack shape";
     break;

lib/Sema/ConstraintSystem.cpp

@@ -637,6 +637,26 @@ std::pair<Type, OpenedArchetypeType *> ConstraintSystem::openExistentialType(
   return {result, opened};
 }
 
+void ConstraintSystem::addPackElementEnvironment(PackExpansionExpr *expr) {
+  auto *locator = getConstraintLocator(expr);
+  PackExpansionEnvironments[locator] = UUID::fromTime();
+}
+
+GenericEnvironment *
+ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator) {
+  auto result = PackExpansionEnvironments.find(locator);
+  if (result == PackExpansionEnvironments.end())
+    return nullptr;
+
+  auto uuid = result->second;
+  auto &ctx = getASTContext();
+  auto elementSig = ctx.getOpenedElementSignature(
+      DC->getGenericSignatureOfContext().getCanonicalSignature());
+  auto *contextEnv = DC->getGenericEnvironmentOfContext();
+  auto contextSubs = contextEnv->getForwardingSubstitutionMap();
+  return GenericEnvironment::forOpenedElement(elementSig, uuid, contextSubs);
+}
+
 /// Extend the given depth map by adding depths for all of the subexpressions
 /// of the given expression.
 static void extendDepthMap(
@@ -5378,7 +5398,6 @@ void constraints::simplifyLocator(ASTNode &anchor,
       break;
 
     case ConstraintLocator::PackElement:
-    case ConstraintLocator::OpenedPackElement:
     case ConstraintLocator::PackShape:
       break;
 

lib/Sema/PreCheckExpr.cpp

@@ -399,21 +399,6 @@ static BinaryExpr *getCompositionExpr(Expr *expr) {
   return nullptr;
 }
 
-static void
-setPackElementEnvironment(DeclContext *dc, PackExpansionExpr *expansion) {
-  // FIXME: The generic environment should be created lazily when solving
-  // PackElementOf constraints.
-  auto &ctx = dc->getASTContext();
-  auto sig = ctx.getOpenedElementSignature(
-      dc->getGenericSignatureOfContext().getCanonicalSignature());
-  auto *contextEnv = dc->getGenericEnvironmentOfContext();
-  auto contextSubs = contextEnv->getForwardingSubstitutionMap();
-  auto *environment =
-      GenericEnvironment::forOpenedElement(sig, UUID::fromTime(),
-                                           contextSubs);
-  expansion->setGenericEnvironment(environment);
-}
-
 /// Bind an UnresolvedDeclRefExpr by performing name lookup and
 /// returning the resultant expression. Context is the DeclContext used
 /// for the lookup.
@@ -1160,12 +1145,6 @@ namespace {
         return Action::Continue(result);
       }
 
-      // Set the generic environment of a PackExpansionExpr.
-      if (auto *expansion = dyn_cast<PackExpansionExpr>(expr)) {
-        setPackElementEnvironment(DC, expansion);
-        return Action::Continue(expansion);
-      }
-
       // Type check the type parameters in an UnresolvedSpecializeExpr.
       if (auto *us = dyn_cast<UnresolvedSpecializeExpr>(expr)) {
         if (auto *typeExpr = simplifyUnresolvedSpecializeExpr(us))