Teach witness thunks to hop to the actor when needed (#59456)

DougGregor · web-flow · commit 33818ff8e8b2 · 2022-06-15T18:37:59.000+09:00
* Teach witness thunks to hop to the actor when needed. When a synchronous, actor-isolated declaration witnesses an asynchronous, not-similarly-isolated requirement, emit an actor hop within the witness thunk to ensure that we properly enter the context of the actor. Fixes #58517 / rdar://92881539. (cherry picked from commit 25a7988) * Enter actor isolation for a distributed witness via the distributed thunk We cannot hop to a distributed actor directly, so call through the distributed thunk instead. (cherry picked from commit 840b7ea)
diff --git a/include/swift/AST/ProtocolConformance.h b/include/swift/AST/ProtocolConformance.h
@@ -601,6 +601,9 @@ class NormalProtocolConformance : public RootProtocolConformance,
   /// Set the witness for the given requirement.
   void setWitness(ValueDecl *requirement, Witness witness) const;
 
+  /// Override the witness for a given requirement.
+  void overrideWitness(ValueDecl *requirement, Witness newWitness);
+
   /// Retrieve the protocol conformances that satisfy the requirements of the
   /// protocol, which line up with the conformance constraints in the
   /// protocol's requirement signature.
diff --git a/include/swift/AST/Witness.h b/include/swift/AST/Witness.h
@@ -97,6 +97,7 @@ class Witness {
     /// The derivative generic signature, when the requirement is a derivative
     /// function.
     GenericSignature derivativeGenSig;
+    Optional<ActorIsolation> enterIsolation;
   };
 
   llvm::PointerUnion<ValueDecl *, StoredWitness *> storage;
@@ -125,10 +126,12 @@ class Witness {
   ///
   /// Deserialized witnesses do not have a synthetic environment.
   static Witness forDeserialized(ValueDecl *decl,
-                                 SubstitutionMap substitutions) {
+                                 SubstitutionMap substitutions,
+                                 Optional<ActorIsolation> enterIsolation) {
     // TODO: It's probably a good idea to have a separate 'deserialized' bit.
     return Witness(
-        decl, substitutions, nullptr, SubstitutionMap(), CanGenericSignature());
+        decl, substitutions, nullptr, SubstitutionMap(), CanGenericSignature(),
+        enterIsolation);
   }
 
   /// Create a witness that requires substitutions.
@@ -145,11 +148,15 @@ class Witness {
   ///
   /// \param derivativeGenSig The derivative generic signature, when the
   /// requirement is a derivative function.
+  ///
+  /// \param enterIsolation The actor isolation that the witness thunk will
+  /// need to hop to before calling the witness.
   Witness(ValueDecl *decl,
           SubstitutionMap substitutions,
           GenericEnvironment *syntheticEnv,
           SubstitutionMap reqToSyntheticEnvSubs,
-          GenericSignature derivativeGenSig);
+          GenericSignature derivativeGenSig,
+          Optional<ActorIsolation> enterIsolation);
 
   /// Retrieve the witness declaration reference, which includes the
   /// substitutions needed to use the witness from the synthetic environment
@@ -198,6 +205,17 @@ class Witness {
     return GenericSignature();
   }
 
+  Optional<ActorIsolation> getEnterIsolation() const {
+    if (auto *storedWitness = storage.dyn_cast<StoredWitness *>())
+      return storedWitness->enterIsolation;
+
+    return None;
+  }
+
+  /// Retrieve a copy of the witness with an actor isolation that the
+  /// witness thunk will need to hop to.
+  Witness withEnterIsolation(ActorIsolation enterIsolation) const;
+
   SWIFT_DEBUG_DUMP;
 
   void dump(llvm::raw_ostream &out) const;
diff --git a/lib/AST/ProtocolConformance.cpp b/lib/AST/ProtocolConformance.cpp
@@ -44,9 +44,10 @@ using namespace swift;
 Witness::Witness(ValueDecl *decl, SubstitutionMap substitutions,
                  GenericEnvironment *syntheticEnv,
                  SubstitutionMap reqToSynthesizedEnvSubs,
-                 GenericSignature derivativeGenSig) {
+                 GenericSignature derivativeGenSig,
+                 Optional<ActorIsolation> enterIsolation) {
   if (!syntheticEnv && substitutions.empty() &&
-      reqToSynthesizedEnvSubs.empty()) {
+      reqToSynthesizedEnvSubs.empty() && !enterIsolation) {
     storage = decl;
     return;
   }
@@ -56,11 +57,17 @@ Witness::Witness(ValueDecl *decl, SubstitutionMap substitutions,
   auto storedMem = ctx.Allocate(sizeof(StoredWitness), alignof(StoredWitness));
   auto stored = new (storedMem) StoredWitness{declRef, syntheticEnv,
                                               reqToSynthesizedEnvSubs,
-                                              derivativeGenSig};
+                                              derivativeGenSig, enterIsolation};
 
   storage = stored;
 }
 
+Witness Witness::withEnterIsolation(ActorIsolation enterIsolation) const {
+  return Witness(getDecl(), getSubstitutions(), getSyntheticEnvironment(),
+                 getRequirementToSyntheticSubs(),
+                 getDerivativeGenericSignature(), enterIsolation);
+}
+
 void Witness::dump() const { dump(llvm::errs()); }
 
 void Witness::dump(llvm::raw_ostream &out) const {
@@ -902,7 +909,7 @@ NormalProtocolConformance::getWitnessUncached(ValueDecl *requirement) const {
 
 Witness SelfProtocolConformance::getWitness(ValueDecl *requirement) const {
   return Witness(requirement, SubstitutionMap(), nullptr, SubstitutionMap(),
-                 GenericSignature());
+                 GenericSignature(), None);
 }
 
 ConcreteDeclRef
@@ -941,6 +948,12 @@ void NormalProtocolConformance::setWitness(ValueDecl *requirement,
   Mapping[requirement] = witness;
 }
 
+void NormalProtocolConformance::overrideWitness(ValueDecl *requirement,
+                                                Witness witness) {
+  assert(Mapping.count(requirement) == 1 && "Witness not known");
+  Mapping[requirement] = witness;
+}
+
 SpecializedProtocolConformance::SpecializedProtocolConformance(
     Type conformingType,
     ProtocolConformance *genericConformance,
diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -756,7 +756,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                            SILDeclRef witness,
                            SubstitutionMap witnessSubs,
                            IsFreeFunctionWitness_t isFree,
-                           bool isSelfConformance);
+                           bool isSelfConformance,
+                           Optional<ActorIsolation> enterIsolation);
 
   /// Generates subscript arguments for keypath. This function handles lowering
   /// of all index expressions including default arguments.
diff --git a/lib/SILGen/SILGenPoly.cpp b/lib/SILGen/SILGenPoly.cpp
@@ -4407,14 +4407,16 @@ emitOpenExistentialInSelfConformance(SILGenFunction &SGF, SILLocation loc,
                                    : AccessKind::Read);
 }
 
-void SILGenFunction::emitProtocolWitness(AbstractionPattern reqtOrigTy,
-                                         CanAnyFunctionType reqtSubstTy,
-                                         SILDeclRef requirement,
-                                         SubstitutionMap reqtSubs,
-                                         SILDeclRef witness,
-                                         SubstitutionMap witnessSubs,
-                                         IsFreeFunctionWitness_t isFree,
-                                         bool isSelfConformance) {
+void SILGenFunction::emitProtocolWitness(
+    AbstractionPattern reqtOrigTy,
+    CanAnyFunctionType reqtSubstTy,
+    SILDeclRef requirement,
+    SubstitutionMap reqtSubs,
+    SILDeclRef witness,
+    SubstitutionMap witnessSubs,
+    IsFreeFunctionWitness_t isFree,
+    bool isSelfConformance,
+    Optional<ActorIsolation> enterIsolation) {
   // FIXME: Disable checks that the protocol witness carries debug info.
   // Should we carry debug info for witnesses?
   F.setBare(IsBare);
@@ -4434,6 +4436,34 @@ void SILGenFunction::emitProtocolWitness(AbstractionPattern reqtOrigTy,
   SmallVector<ManagedValue, 8> origParams;
   collectThunkParams(loc, origParams);
 
+  // If we are supposed to enter the actor, do so now.
+  if (enterIsolation) {
+    if (enterIsolation->isDistributedActor()) {
+      // For a distributed actor, call through the distributed thunk.
+      witness = witness.asDistributed();
+    } else {
+      // For a non-distributed actor, hop to the actor.
+      Optional<ManagedValue> actorSelf;
+
+      // For an instance actor, get the actor 'self'.
+      if (*enterIsolation == ActorIsolation::ActorInstance) {
+        auto actorSelfVal = origParams.back();
+
+        if (actorSelfVal.getType().isAddress()) {
+          auto &actorSelfTL = getTypeLowering(actorSelfVal.getType());
+          if (!actorSelfTL.isAddressOnly()) {
+            actorSelfVal = emitManagedLoad(
+                *this, loc, actorSelfVal, actorSelfTL);
+          }
+        }
+
+        actorSelf = actorSelfVal;
+      }
+
+      emitHopToTargetActor(loc, enterIsolation, actorSelf);
+    }
+  }
+
   // Get the type of the witness.
   auto witnessInfo = getConstantInfo(getTypeExpansionContext(), witness);
   CanAnyFunctionType witnessSubstTy = witnessInfo.LoweredType;
diff --git a/lib/SILGen/SILGenType.cpp b/lib/SILGen/SILGenType.cpp
@@ -812,7 +812,8 @@ SILFunction *SILGenModule::emitProtocolWitness(
 
   SGF.emitProtocolWitness(AbstractionPattern(reqtOrigTy), reqtSubstTy,
                           requirement, reqtSubMap, witnessRef,
-                          witnessSubs, isFree, /*isSelfConformance*/ false);
+                          witnessSubs, isFree, /*isSelfConformance*/ false,
+                          witness.getEnterIsolation());
 
   emitLazyConformancesForFunction(f);
   return f;
@@ -884,7 +885,8 @@ static SILFunction *emitSelfConformanceWitness(SILGenModule &SGM,
 
   SGF.emitProtocolWitness(AbstractionPattern(reqtOrigTy), reqtSubstTy,
                           requirement, reqtSubs, requirement,
-                          witnessSubs, isFree, /*isSelfConformance*/ true);
+                          witnessSubs, isFree, /*isSelfConformance*/ true,
+                          None);
 
   SGM.emitLazyConformancesForFunction(f);
 
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
@@ -102,7 +102,7 @@ swift::Witness RequirementMatch::getWitness(ASTContext &ctx) const {
   auto syntheticEnv = ReqEnv->getSyntheticEnvironment();
   return swift::Witness(this->Witness, WitnessSubstitutions,
                         syntheticEnv, ReqEnv->getRequirementToSyntheticMap(),
-                        DerivativeGenSig);
+                        DerivativeGenSig, None);
 }
 
 AssociatedTypeDecl *
@@ -2932,7 +2932,7 @@ static bool hasExplicitGlobalActorAttr(ValueDecl *decl) {
   return !globalActorAttr->first->isImplicit();
 }
 
-bool ConformanceChecker::checkActorIsolation(
+Optional<ActorIsolation> ConformanceChecker::checkActorIsolation(
     ValueDecl *requirement, ValueDecl *witness) {
   /// Retrieve a concrete witness for Sendable checking.
   auto getConcreteWitness = [&] {
@@ -2972,12 +2972,12 @@ bool ConformanceChecker::checkActorIsolation(
     }
 
     // Otherwise, we're done.
-    return false;
+    return None;
 
   case ActorReferenceResult::ExitsActorToNonisolated:
     diagnoseNonSendableTypesInReference(
         getConcreteWitness(), DC, loc, SendableCheckReason::Conformance);
-    return false;
+    return None;
 
   case ActorReferenceResult::EntersActor:
     // Handled below.
@@ -3056,12 +3056,16 @@ bool ConformanceChecker::checkActorIsolation(
     // that is explicitly marked nonisolated.
     if (isa<ConstructorDecl>(witness) &&
         witness->getAttrs().hasAttribute<NonisolatedAttr>())
-      return false;
+      return None;
 
     diagnoseNonSendableTypesInReference(
         getConcreteWitness(), DC, loc, SendableCheckReason::Conformance);
 
-    return false;
+    if (refResult.isolation.isActorIsolated() && isAsyncDecl(requirement) &&
+        !isAsyncDecl(witness))
+      return refResult.isolation;
+
+    return None;
   }
 
   // Limit the behavior of the diagnostic based on context.
@@ -3080,7 +3084,6 @@ bool ConformanceChecker::checkActorIsolation(
 
   // Complain that this witness cannot conform to the requirement due to
   // actor isolation.
-  bool isError = (behavior == DiagnosticBehavior::Unspecified);
   witness->diagnose(diag::actor_isolated_witness,
                     isDistributed && !isDistributedDecl(witness),
                     refResult.isolation, witness->getDescriptiveKind(),
@@ -3168,7 +3171,7 @@ bool ConformanceChecker::checkActorIsolation(
     requirement->diagnose(diag::decl_declared_here, requirement->getName());
   }
 
-  return isError;
+  return None;
 }
 
 bool ConformanceChecker::checkObjCTypeErasedGenerics(
@@ -5071,8 +5074,13 @@ void ConformanceChecker::resolveValueWitnesses() {
 
       auto &C = witness->getASTContext();
 
-      if (checkActorIsolation(requirement, witness)) {
-        return;
+      // Check actor isolation. If we need to enter into the actor's
+      // isolation within the witness thunk, record that.
+      if (auto enteringIsolation = checkActorIsolation(requirement, witness)) {
+        Conformance->overrideWitness(
+            requirement,
+            Conformance->getWitnessUncached(requirement)
+              .withEnterIsolation(*enteringIsolation));
       }
 
       // Objective-C checking for @objc requirements.
diff --git a/lib/Sema/TypeCheckProtocol.h b/lib/Sema/TypeCheckProtocol.h
@@ -788,8 +788,11 @@ class ConformanceChecker : public WitnessChecker {
 
   /// Check that the witness and requirement have compatible actor contexts.
   ///
-  /// \returns true if an error occurred, false otherwise.
-  bool checkActorIsolation(ValueDecl *requirement, ValueDecl *witness);
+  /// \returns the isolation that needs to be enforced to invoke the witness
+  /// from the requirement, used when entering an actor-isolated synchronous
+  /// witness from an asynchronous requirement.
+  Optional<ActorIsolation>
+  checkActorIsolation(ValueDecl *requirement, ValueDecl *witness);
 
   /// Record a type witness.
   ///
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
@@ -6964,14 +6964,22 @@ void ModuleFile::finishNormalConformance(NormalProtocolConformance *conformance,
         fatal(witnessSubstitutions.takeError());
     }
 
+    // Determine whether we need to enter the actor isolation of the witness.
+    Optional<ActorIsolation> enterIsolation;
+    if (*rawIDIter++) {
+      enterIsolation = getActorIsolation(witness);
+    }
+
     // Handle opaque witnesses that couldn't be deserialized.
     if (isOpaque) {
       trySetOpaqueWitness();
       continue;
     }
 
     // Set the witness.
-    trySetWitness(Witness::forDeserialized(witness, witnessSubstitutions.get()));
+    trySetWitness(
+        Witness::forDeserialized(
+          witness, witnessSubstitutions.get(), enterIsolation));
   }
   assert(rawIDIter <= rawIDs.end() && "read too much");
   
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
@@ -56,7 +56,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 690; // cast ownership serialization
+const uint16_t SWIFTMODULE_VERSION_MINOR = 691; // witness enter isolation
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -1592,6 +1592,7 @@ void Serializer::writeLocalNormalProtocolConformance(
         subs = subs.mapReplacementTypesOutOfContext();
 
       data.push_back(addSubstitutionMapRef(subs));
+      data.push_back(witness.getEnterIsolation().hasValue() ? 1 : 0);
   });
 
   unsigned abbrCode
diff --git a/test/SILGen/distributed_thunk.swift b/test/SILGen/distributed_thunk.swift
@@ -17,3 +17,16 @@ extension DA {
 
   distributed func f() { }
 }
+
+protocol ServerProto {
+  func doSomething() async throws
+}
+
+extension DA: ServerProto {
+  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17distributed_thunk2DACAA11ServerProtoA2aDP11doSomethingyyYaKFTW : $@convention(witness_method: ServerProto) @async (@in_guaranteed DA) -> @error Error
+  // CHECK-NOT: hop_to_executor
+  // CHECK-NOT: return
+  // CHECK: function_ref @$s17distributed_thunk2DAC11doSomethingyyFTE
+  // CHECK: return
+  distributed func doSomething() { }
+}
diff --git a/test/SILGen/hop_to_executor_witness.swift b/test/SILGen/hop_to_executor_witness.swift
@@ -0,0 +1,22 @@
+// RUN: %target-swift-frontend -emit-silgen %s -module-name test -swift-version 5 | %FileCheck --enable-var-scope %s --implicit-check-not 'hop_to_executor {{%[0-9]+}}'
+// REQUIRES: concurrency
+
+@available(SwiftStdlib 5.1, *)
+protocol TestableAsync {
+  func test() async
+}
+
+@available(SwiftStdlib 5.1, *)
+actor TestActor : TestableAsync {
+  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s4test9TestActorCAA13TestableAsyncA2aDPAAyyYaFTW : $@convention(witness_method: TestableAsync) @async (@in_guaranteed TestActor) -> ()
+  // CHECK: [[ACTOR_INSTANCE:%.*]] = load_borrow %0 : $*TestActor
+  // CHECK-NEXT:  hop_to_executor [[ACTOR_INSTANCE]] : $TestActor
+  func test() { }
+}
+
+@available(SwiftStdlib 5.1, *)
+@MainActor class TestClass : TestableAsync {
+  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s4test9TestClassCAA13TestableAsyncA2aDPAAyyYaFTW : $@convention(witness_method: TestableAsync) @async (@in_guaranteed TestClass) -> ()
+  // CHECK: hop_to_executor {{%.*}} : $MainActor 
+  func test() { }
+}
