[region-isolation] Allow for unapplied isolated parameter ownership.

Given a function or a partial_apply with an isolated parameter, we do not know
immediately what the actual isolation is of the function or partial_apply since
we do not know which instance will be applied to the function or partial_apply.

In this commit, I introduce a new bit into SILIsolationInfo that tracks this
information upon construction and allows for it to merge with ownership that has
the appropriate type and a specific instance. Since the values that created the
two isolations, will be in the same region this should ensure that the value is
only ever in a flow sensitive manner in a region with only one actor instance
(since regions with isolations with differing actor instances are illegal).

Author: gottesmm

include/swift/SILOptimizer/Utils/SILIsolationInfo.h

@@ -114,6 +114,10 @@ class ActorInstance {
   }
 };
 
+/// The isolation info inferred for a specific SILValue. Use
+/// SILIsolationInfo::get() to compute these. It is intended to be a
+/// conservatively correct model that we expand over time with more pattern
+/// matching.
 class SILIsolationInfo {
 public:
   /// The lattice is:
@@ -128,6 +132,20 @@ class SILIsolationInfo {
     Actor,
   };
 
+  enum class Flag : uint8_t {
+    None,
+
+    /// If set, this means that the element that we derived this from was marked
+    /// with nonisolated(unsafe).
+    UnsafeNonIsolated = 0x1,
+
+    /// If set, this means that this actor isolation is from an isolated
+    /// parameter and should be allowed to merge into a self parameter.
+    UnappliedIsolatedAnyParameter = 0x2,
+  };
+
+  using Options = OptionSet<Flag>;
+
 private:
   /// The actor isolation if this value has one. The default unspecified case
   /// otherwise.
@@ -142,13 +160,13 @@ class SILIsolationInfo {
   ActorInstance actorInstance;
 
   unsigned kind : 8;
-  unsigned unsafeNonIsolated : 1;
+  unsigned options : 8;
 
   SILIsolationInfo(SILValue isolatedValue, SILValue actorInstance,
-                   ActorIsolation actorIsolation, bool isUnsafeNonIsolated)
+                   ActorIsolation actorIsolation, Options options = Options())
       : actorIsolation(actorIsolation), isolatedValue(isolatedValue),
         actorInstance(ActorInstance::getForValue(actorInstance)), kind(Actor),
-        unsafeNonIsolated(isUnsafeNonIsolated) {
+        options(options.toRaw()) {
     assert((!actorInstance ||
             (actorIsolation.getKind() == ActorIsolation::ActorInstance &&
              actorInstance->getType()
@@ -159,24 +177,22 @@ class SILIsolationInfo {
   }
 
   SILIsolationInfo(SILValue isolatedValue, ActorInstance actorInstance,
-                   ActorIsolation actorIsolation, bool isUnsafeNonIsolated)
+                   ActorIsolation actorIsolation, Options options = Options())
       : actorIsolation(actorIsolation), isolatedValue(isolatedValue),
-        actorInstance(actorInstance), kind(Actor),
-        unsafeNonIsolated(isUnsafeNonIsolated) {
+        actorInstance(actorInstance), kind(Actor), options(options.toRaw()) {
     assert(actorInstance);
     assert(actorIsolation.getKind() == ActorIsolation::ActorInstance);
   }
 
   SILIsolationInfo(Kind kind, SILValue isolatedValue)
-      : actorIsolation(), isolatedValue(isolatedValue), kind(kind),
-        unsafeNonIsolated(false) {}
+      : actorIsolation(), isolatedValue(isolatedValue), kind(kind), options(0) {
+  }
 
-  SILIsolationInfo(Kind kind, bool isUnsafeNonIsolated)
-      : actorIsolation(), kind(kind), unsafeNonIsolated(isUnsafeNonIsolated) {}
+  SILIsolationInfo(Kind kind, Options options = Options())
+      : actorIsolation(), kind(kind), options(options.toRaw()) {}
 
 public:
-  SILIsolationInfo()
-      : actorIsolation(), kind(Kind::Unknown), unsafeNonIsolated(false) {}
+  SILIsolationInfo() : actorIsolation(), kind(Kind::Unknown), options(0) {}
 
   operator bool() const { return kind != Kind::Unknown; }
 
@@ -188,12 +204,43 @@ class SILIsolationInfo {
   bool isActorIsolated() const { return kind == Kind::Actor; }
   bool isTaskIsolated() const { return kind == Kind::Task; }
 
-  bool isUnsafeNonIsolated() const { return unsafeNonIsolated; }
+  Options getOptions() const { return Options(options); }
+
+  void setOptions(Options newOptions) { options = newOptions.toRaw(); }
+
+  bool isUnsafeNonIsolated() const {
+    return getOptions().contains(Flag::UnsafeNonIsolated);
+  }
 
   SILIsolationInfo withUnsafeNonIsolated(bool newValue = true) const {
     assert(*this && "Cannot be unknown");
     auto self = *this;
-    self.unsafeNonIsolated = newValue;
+    if (newValue) {
+      self.options = (self.getOptions() | Flag::UnsafeNonIsolated).toRaw();
+    } else {
+      self.options =
+          self.getOptions().toRaw() & ~Options(Flag::UnsafeNonIsolated).toRaw();
+    }
+    return self;
+  }
+
+  /// Returns true if this actor isolation is derived from an unapplied
+  /// isolation parameter. When merging, we allow for this to be merged with a
+  /// more specific isolation kind.
+  bool isUnappliedIsolatedAnyParameter() const {
+    return getOptions().contains(Flag::UnappliedIsolatedAnyParameter);
+  }
+
+  SILIsolationInfo withUnappliedIsolatedParameter(bool newValue = true) const {
+    assert(*this && "Cannot be unknown");
+    auto self = *this;
+    if (newValue) {
+      self.options =
+          (self.getOptions() | Flag::UnappliedIsolatedAnyParameter).toRaw();
+    } else {
+      self.options = self.getOptions().toRaw() &
+                     ~Options(Flag::UnappliedIsolatedAnyParameter).toRaw();
+    }
     return self;
   }
 
@@ -243,7 +290,8 @@ class SILIsolationInfo {
   }
 
   static SILIsolationInfo getDisconnected(bool isUnsafeNonIsolated) {
-    return {Kind::Disconnected, isUnsafeNonIsolated};
+    return {Kind::Disconnected,
+            isUnsafeNonIsolated ? Flag::UnsafeNonIsolated : Flag::None};
   }
 
   /// Create an actor isolation for a value that we know is actor isolated to a
@@ -261,7 +309,7 @@ class SILIsolationInfo {
   getFlowSensitiveActorIsolated(SILValue isolatedValue,
                                 ActorIsolation actorIsolation) {
     return {isolatedValue, SILValue(), actorIsolation,
-            false /*nonisolated(unsafe)*/};
+            Flag::UnappliedIsolatedAnyParameter};
   }
 
   /// Only use this as a fallback if we cannot find better information.
@@ -270,8 +318,7 @@ class SILIsolationInfo {
     if (crossing.getCalleeIsolation().isActorIsolated()) {
       // SIL level, just let it through
       return SILIsolationInfo(SILValue(), SILValue(),
-                              crossing.getCalleeIsolation(),
-                              false /*nonisolated(unsafe)*/);
+                              crossing.getCalleeIsolation());
     }
 
     return {};
@@ -287,8 +334,7 @@ class SILIsolationInfo {
       return {};
     }
     return {isolatedValue, actorInstance,
-            ActorIsolation::forActorInstanceSelf(typeDecl),
-            false /*nonisolated(unsafe)*/};
+            ActorIsolation::forActorInstanceSelf(typeDecl)};
   }
 
   static SILIsolationInfo getActorInstanceIsolated(SILValue isolatedValue,
@@ -301,8 +347,7 @@ class SILIsolationInfo {
       return {};
     }
     return {isolatedValue, actorInstance,
-            ActorIsolation::forActorInstanceSelf(typeDecl),
-            false /*nonisolated(unsafe)*/};
+            ActorIsolation::forActorInstanceSelf(typeDecl)};
   }
 
   /// A special actor instance isolated for partial apply cases where we do not
@@ -318,14 +363,13 @@ class SILIsolationInfo {
     }
     return {isolatedValue, SILValue(),
             ActorIsolation::forActorInstanceSelf(typeDecl),
-            false /*nonisolated(unsafe)*/};
+            Flag::UnappliedIsolatedAnyParameter};
   }
 
   static SILIsolationInfo getGlobalActorIsolated(SILValue value,
                                                  Type globalActorType) {
     return {value, SILValue() /*no actor instance*/,
-            ActorIsolation::forGlobalActor(globalActorType),
-            false /*nonisolated(unsafe)*/};
+            ActorIsolation::forGlobalActor(globalActorType)};
   }
 
   static SILIsolationInfo getGlobalActorIsolated(SILValue value,
@@ -395,7 +439,7 @@ class SILDynamicMergedIsolationInfo {
 
 public:
   SILDynamicMergedIsolationInfo() : innerInfo() {}
-  explicit SILDynamicMergedIsolationInfo(SILIsolationInfo innerInfo)
+  SILDynamicMergedIsolationInfo(SILIsolationInfo innerInfo)
       : innerInfo(innerInfo) {}
 
   /// Returns nullptr only if both this isolation info and \p other are actor

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1635,8 +1635,9 @@ class PartitionOpTranslator {
               if (originalMergedInfo)
                   originalMergedInfo->printForDiagnostics(llvm::dbgs());
               else llvm::dbgs() << "nil";
-              llvm::dbgs() << "\nValue: "
+              llvm::dbgs() << "\nValue Rep: "
                            << value->getRepresentative().getValue();
+              llvm::dbgs() << "Original Src: " << src;
               llvm::dbgs() << "Value Info: ";
               value->getIsolationRegionInfo().printForDiagnostics(llvm::dbgs());
               llvm::dbgs() << "\n");

lib/SILOptimizer/Utils/SILIsolationInfo.cpp

@@ -875,15 +875,38 @@ SILIsolationInfo SILIsolationInfo::get(SILArgument *arg) {
 }
 
 void SILIsolationInfo::print(llvm::raw_ostream &os) const {
+  auto printOptions = [&] {
+    auto opts = getOptions();
+    if (!opts)
+      return;
+
+    os << ": ";
+
+    std::array<std::pair<Flag, StringLiteral>, 2> data = {
+        std::make_pair(Flag::UnsafeNonIsolated,
+                       StringLiteral("nonisolated(unsafe)")),
+        std::make_pair(Flag::UnappliedIsolatedAnyParameter,
+                       StringLiteral("unapplied_isolated_parameter")),
+    };
+
+    llvm::interleave(
+        data, os,
+        [&](const std::pair<Flag, StringLiteral> &value) {
+          opts -= value.first;
+          os << value.second;
+        },
+        ", ");
+
+    assert(!opts && "Unhandled flag?!");
+  };
+
   switch (Kind(*this)) {
   case Unknown:
     os << "unknown";
     return;
   case Disconnected:
     os << "disconnected";
-    if (unsafeNonIsolated) {
-      os << ": nonisolated(unsafe)";
-    }
+    printOptions();
     return;
   case Actor:
     if (ActorInstance instance = getActorInstance()) {
@@ -892,9 +915,7 @@ void SILIsolationInfo::print(llvm::raw_ostream &os) const {
         SILValue value = instance.getValue();
         if (auto name = VariableNameInferrer::inferName(value)) {
           os << "'" << *name << "'-isolated";
-          if (unsafeNonIsolated) {
-            os << ": nonisolated(unsafe)";
-          }
+          printOptions();
           os << "\n";
           os << "instance: " << *value;
 
@@ -904,9 +925,7 @@ void SILIsolationInfo::print(llvm::raw_ostream &os) const {
       }
       case ActorInstance::Kind::ActorAccessorInit:
         os << "'self'-isolated";
-        if (unsafeNonIsolated) {
-          os << ": nonisolated(unsafe)";
-        }
+        printOptions();
         os << '\n';
         os << "instance: actor accessor init\n";
         return;
@@ -916,23 +935,17 @@ void SILIsolationInfo::print(llvm::raw_ostream &os) const {
     if (getActorIsolation().getKind() == ActorIsolation::ActorInstance) {
       if (auto *vd = getActorIsolation().getActorInstance()) {
         os << "'" << vd->getBaseIdentifier() << "'-isolated";
-        if (unsafeNonIsolated) {
-          os << ": nonisolated(unsafe)";
-        }
+        printOptions();
         return;
       }
     }
 
     getActorIsolation().printForDiagnostics(os);
-    if (unsafeNonIsolated) {
-      os << ": nonisolated(unsafe)";
-    }
+    printOptions();
     return;
   case Task:
     os << "task-isolated";
-    if (unsafeNonIsolated) {
-      os << ": nonisolated(unsafe)";
-    }
+    printOptions();
     os << '\n';
     os << "instance: " << *getIsolatedValue();
     return;
@@ -959,8 +972,13 @@ bool SILIsolationInfo::hasSameIsolation(const SILIsolationInfo &other) const {
     ActorInstance actor1 = getActorInstance();
     ActorInstance actor2 = other.getActorInstance();
 
-    // If either are non-null, and the actor instance doesn't match, return
-    // false.
+    // If either have an actor instance, and the actor instance doesn't match,
+    // return false.
+    //
+    // This ensures that cases like comparing two global actor isolated things
+    // do not hit this path.
+    //
+    // It also catches cases where we have a missing actor instance.
     if ((actor1 || actor2) && actor1 != actor2)
       return false;
 
@@ -1054,7 +1072,7 @@ void SILIsolationInfo::printForOneLineLogging(llvm::raw_ostream &os) const {
     return;
   case Disconnected:
     os << "disconnected";
-    if (unsafeNonIsolated) {
+    if (getOptions().contains(Flag::UnsafeNonIsolated)) {
       os << ": nonisolated(unsafe)";
     }
     return;
@@ -1143,14 +1161,31 @@ std::optional<SILDynamicMergedIsolationInfo>
 SILDynamicMergedIsolationInfo::merge(SILIsolationInfo other) const {
   // If we are greater than the other kind, then we are further along the
   // lattice. We ignore the change.
-  if (unsigned(other.getKind()) < unsigned(innerInfo.getKind()))
+  if (unsigned(innerInfo.getKind() > unsigned(other.getKind())))
     return {*this};
 
-  // If we are both actor isolated and our isolations are not
-  // compatible... return None.
-  if (other.isActorIsolated() && innerInfo.isActorIsolated() &&
-      !innerInfo.hasSameIsolation(other))
+  // If we are both actor isolated...
+  if (innerInfo.isActorIsolated() && other.isActorIsolated()) {
+    // If both innerInfo and other have the same isolation, we are obviously
+    // done. Just return innerInfo since we could return either.
+    if (innerInfo.hasSameIsolation(other))
+      return {innerInfo};
+
+    // Ok, there is some difference in between innerInfo and other. Lets see if
+    // they are both actor instance isolated and if either are unapplied
+    // isolated any parameter. In such a case, take the one that is further
+    // along.
+    if (innerInfo.getActorIsolation().isActorInstanceIsolated() &&
+        other.getActorIsolation().isActorInstanceIsolated()) {
+      if (innerInfo.isUnappliedIsolatedAnyParameter())
+        return other;
+      if (other.isUnappliedIsolatedAnyParameter())
+        return innerInfo;
+    }
+
+    // Otherwise, they do not match... so return None to signal merge failure.
     return {};
+  }
 
   // If we are both disconnected and other has the unsafeNonIsolated bit set,
   // drop that bit and return that.
@@ -1159,11 +1194,11 @@ SILDynamicMergedIsolationInfo::merge(SILIsolationInfo other) const {
   // merging. These bits should not propagate through merging and should instead
   // always be associated with non-merged infos.
   if (other.isDisconnected() && other.isUnsafeNonIsolated()) {
-    return SILDynamicMergedIsolationInfo(other.withUnsafeNonIsolated(false));
+    return {other.withUnsafeNonIsolated(false)};
   }
 
   // Otherwise, just return other.
-  return SILDynamicMergedIsolationInfo(other);
+  return {other};
 }
 
 //===----------------------------------------------------------------------===//