Sema: Record isolated parameters in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -66,8 +66,9 @@ CHANGE(RecordedContextualInfo)
 CHANGE(RecordedTarget)
 CHANGE(RecordedCaseLabelItemInfo)
 CHANGE(RecordedPotentialThrowSite)
+CHANGE(RecordedIsolatedParam)
 
-LAST_CHANGE(RecordedPotentialThrowSite)
+LAST_CHANGE(RecordedIsolatedParam)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE

include/swift/Sema/CSTrail.h

@@ -134,6 +134,7 @@ class SolverTrail {
       DeclContext *TheDeclContext;
       CaseLabelItem *TheItem;
       CatchNode TheCatchNode;
+      ParamDecl *TheParam;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -228,6 +229,9 @@ class SolverTrail {
     /// Create a change that recorded a potential throw site.
     static Change RecordedPotentialThrowSite(CatchNode catchNode);
 
+    /// Create a change that recorded an isolated parameter.
+    static Change RecordedIsolatedParam(ParamDecl *param);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1594,7 +1594,7 @@ class Solution {
   llvm::DenseMap<Expr *, ExprPattern *> exprPatterns;
 
   /// The set of parameters that have been inferred to be 'isolated'.
-  std::vector<ParamDecl *> isolatedParams;
+  llvm::DenseSet<ParamDecl *> isolatedParams;
 
   /// The set of closures that have been inferred to be "isolated by
   /// preconcurrency".
@@ -2292,7 +2292,7 @@ class ConstraintSystem {
   llvm::SmallDenseMap<Expr *, ExprPattern *, 2> exprPatterns;
 
   /// The set of parameters that have been inferred to be 'isolated'.
-  llvm::SmallSetVector<ParamDecl *, 2> isolatedParams;
+  llvm::SmallDenseSet<ParamDecl *, 2> isolatedParams;
 
   /// The set of closures that have been inferred to be "isolated by
   /// preconcurrency".
@@ -2855,9 +2855,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c isolatedParams.
-    unsigned numIsolatedParams;
-
     /// The length of \c PreconcurrencyClosures.
     unsigned numPreconcurrencyClosures;
 
@@ -4215,6 +4212,13 @@ class ConstraintSystem {
   bool resolveClosure(TypeVariableType *typeVar, Type contextualType,
                       ConstraintLocatorBuilder locator);
 
+  /// Used by the above to update isolatedParams and record a change in
+  /// the trail.
+  void recordIsolatedParam(ParamDecl *param);
+
+  /// Undo the above change.
+  void removeIsolatedParam(ParamDecl *param);
+
   /// Given the fact that contextual type is now available for the type
   /// variable representing a pack expansion type, let's resolve the expansion.
   ///

lib/Sema/CSApply.cpp

@@ -8539,11 +8539,9 @@ bool ExprRewriter::isDistributedThunk(ConcreteDeclRef ref, Expr *context) {
   // If this is a method reference on an potentially isolated
   // actor then it cannot be a remote thunk.
   bool isPotentiallyIsolated = isPotentiallyIsolatedActor(
-      actor,
-    [&](ParamDecl *P) {
-    return P->isIsolated() ||
-           llvm::is_contained(solution.isolatedParams, P);
-  });
+      actor, [&](ParamDecl *P) {
+        return P->isIsolated() || solution.isolatedParams.count(P);
+      });
 
   // Adjust the declaration context to the innermost context that is neither
   // a local function nor a closure, so that the actor reference is checked

lib/Sema/CSSimplify.cpp

@@ -11513,6 +11513,19 @@ static Type getOpenedResultBuilderTypeFor(ConstraintSystem &cs,
   return builderType;
 }
 
+void ConstraintSystem::recordIsolatedParam(ParamDecl *param) {
+  bool inserted = isolatedParams.insert(param).second;
+  ASSERT(inserted);
+
+  if (solverState)
+    recordChange(SolverTrail::Change::RecordedIsolatedParam(param));
+}
+
+void ConstraintSystem::removeIsolatedParam(ParamDecl *param) {
+  bool erased = isolatedParams.erase(param);
+  ASSERT(erased);
+}
+
 bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
                                       Type contextualType,
                                       ConstraintLocatorBuilder locator) {
@@ -11658,7 +11671,7 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
 
         // Note when a parameter is inferred to be isolated.
         if (contextualParam->isIsolated() && !flags.isIsolated() && paramDecl)
-          isolatedParams.insert(paramDecl);
+          recordIsolatedParam(paramDecl);
 
         // Carry-over the ownership specifier from the contextual parameter.
         auto paramOwnership =

lib/Sema/CSSolver.cpp

@@ -231,7 +231,7 @@ Solution ConstraintSystem::finalize() {
     solution.exprPatterns.insert(pattern);
 
   for (const auto &param : isolatedParams)
-    solution.isolatedParams.push_back(param);
+    solution.isolatedParams.insert(param);
 
   for (const auto &closure : preconcurrencyClosures)
     solution.preconcurrencyClosures.push_back(closure);
@@ -406,7 +406,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   for (auto param : solution.isolatedParams) {
-    isolatedParams.insert(param);
+    if (isolatedParams.count(param) == 0)
+      recordIsolatedParam(param);
   }
 
   for (auto &pair : solution.exprPatterns) {
@@ -697,7 +698,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numIsolatedParams = cs.isolatedParams.size();
   numPreconcurrencyClosures = cs.preconcurrencyClosures.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
   numArgumentLists = cs.ArgumentLists.size();
@@ -738,9 +738,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any isolated parameters.
-  truncate(cs.isolatedParams, numIsolatedParams);
-
   // Remove any preconcurrency closures.
   truncate(cs.preconcurrencyClosures, numPreconcurrencyClosures);
 

lib/Sema/CSTrail.cpp

@@ -319,6 +319,14 @@ SolverTrail::Change::RecordedPotentialThrowSite(CatchNode catchNode) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RecordedIsolatedParam(ParamDecl *param) {
+  Change result;
+  result.Kind = ChangeKind::RecordedIsolatedParam;
+  result.TheParam = param;
+  return result;
+}
+
 SyntacticElementTargetKey
 SolverTrail::Change::getSyntacticElementTargetKey() const {
   ASSERT(Kind == ChangeKind::RecordedTarget);
@@ -493,6 +501,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedExprPattern:
     cs.removeExprPatternFor(TheExpr);
     break;
+
+  case ChangeKind::RecordedIsolatedParam:
+    cs.removeIsolatedParam(TheParam);
+    break;
   }
 }
 
@@ -707,6 +719,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     // FIXME: Print something here
     out << "(RecordedPotentialThrowSite)\n";
     break;
+
+  case ChangeKind::RecordedIsolatedParam:
+    out << "(RecordedIsolatedParam ";
+    TheParam->dumpRef(out);
+    out << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -4495,7 +4495,7 @@ size_t Solution::getTotalMemory() const {
          size_in_bytes(targets) +
          size_in_bytes(caseLabelItems) +
          size_in_bytes(exprPatterns) +
-         (isolatedParams.size() * sizeof(void *)) +
+         size_in_bytes(isolatedParams) +
          (preconcurrencyClosures.size() * sizeof(void *)) +
          size_in_bytes(resultBuilderTransformed) +
          size_in_bytes(appliedPropertyWrappers) +