[Concurrency] Implement detecting isIsolatingCurrentContext user impls (#79946)

* [Concurrency] Detect non-default impls of isIsolatingCurrentContext

* [Concurrency] No need for trailing info about isIsolating... in conformance

* Apply changes from review

Author: ktoso

include/swift/ABI/Metadata.h

@@ -2930,6 +2930,14 @@ struct TargetProtocolConformanceDescriptor final
     return Demangle::makeSymbolicMangledNameStringRef(this->template getTrailingObjects<TargetGlobalActorReference<Runtime>>()->type);
   }
 
+  /// True if this is a conformance to 'SerialExecutor' which has a non-default
+  /// (i.e. not the stdlib's default implementation) witness. This means that
+  /// the developer has implemented this method explicitly and we should prefer
+  /// calling it.
+  bool hasNonDefaultSerialExecutorIsIsolatingCurrentContext() const {
+    return Flags.hasNonDefaultSerialExecutorIsIsolatingCurrentContext();
+  }
+
   /// Retrieve the protocol conformance of the global actor type to the
   /// GlobalActor protocol.
   const TargetProtocolConformanceDescriptor<Runtime> *

include/swift/ABI/MetadataValues.h

@@ -748,6 +748,14 @@ class ConformanceFlags {
     IsConformanceOfProtocolMask = 0x01u << 18,
     HasGlobalActorIsolation = 0x01u << 19,
 
+    // Used to detect if this is a conformance to SerialExecutor that has
+    // an user defined implementation of 'isIsolatingCurrentContext'. This
+    // requirement is special in the sense that if a non-default impl is present
+    // we will avoid calling the `checkIsolated` method which would lead to a
+    // crash. In other words, this API "soft replaces" 'checkIsolated' so we
+    // must at runtime the presence of a non-default implementation.
+    HasNonDefaultSerialExecutorIsIsolatingCurrentContext = 0x01u << 20,
+
     NumConditionalPackDescriptorsMask = 0xFFu << 24,
     NumConditionalPackDescriptorsShift = 24
   };
@@ -813,7 +821,15 @@ class ConformanceFlags {
                                  ? HasGlobalActorIsolation
                                  : 0));
   }
-  
+
+  ConformanceFlags withHasNonDefaultSerialExecutorIsIsolatingCurrentContext(
+                                           bool hasNonDefaultSerialExecutorIsIsolatingCurrentContext) const {
+    return ConformanceFlags((Value & ~HasNonDefaultSerialExecutorIsIsolatingCurrentContext)
+                            | (hasNonDefaultSerialExecutorIsIsolatingCurrentContext
+                                 ? HasNonDefaultSerialExecutorIsIsolatingCurrentContext
+                                 : 0));
+  }
+
   /// Retrieve the type reference kind kind.
   TypeReferenceKind getTypeReferenceKind() const {
     return TypeReferenceKind(
@@ -858,6 +874,10 @@ class ConformanceFlags {
     return Value & HasGlobalActorIsolation;
   }
 
+  bool hasNonDefaultSerialExecutorIsIsolatingCurrentContext() const {
+    return Value & HasNonDefaultSerialExecutorIsIsolatingCurrentContext;
+  }
+
   /// Retrieve the # of conditional requirements.
   unsigned getNumConditionalRequirements() const {
     return (Value & NumConditionalRequirementsMask)

include/swift/AST/KnownIdentifiers.def

@@ -118,6 +118,7 @@ IDENTIFIER(makeIterator)
 IDENTIFIER(makeAsyncIterator)
 IDENTIFIER(nestedContainer)
 IDENTIFIER(isEmpty)
+IDENTIFIER(isIsolatingCurrentContext)
 IDENTIFIER(isolation)
 IDENTIFIER(Iterator)
 IDENTIFIER(AsyncIterator)

include/swift/AST/Module.h

@@ -1158,6 +1158,9 @@ class ModuleDecl
   /// \returns true if this module is the "swift" standard library module.
   bool isStdlibModule() const;
 
+  /// \returns true if this module is the "_Concurrency" standard library module.
+  bool isConcurrencyModule() const;
+
   /// \returns true if this module has standard substitutions for mangling.
   bool hasStandardSubstitutions() const;
 

include/swift/Runtime/Bincompat.h

@@ -80,7 +80,7 @@ bool useLegacySwiftObjCHashing();
 /// - if not available, it will invoke the the *crashing* 'checkIsolated'
 ///
 /// This can be overridden by using `SWIFT_UNEXPECTED_EXECUTOR_LOG_LEVEL=1`
-/// or `SWIFT_IS_CURRENT_EXECUTOR_LEGACY_MODE_OVERRIDE=crash|nocrash|swift6|swift6.2`
+/// or `SWIFT_IS_CURRENT_EXECUTOR_LEGACY_MODE_OVERRIDE=crash|nocrash|swift6|isIsolatingCurrentContext`
 SWIFT_RUNTIME_STDLIB_SPI
 bool swift_bincompat_useLegacyNonCrashingExecutorChecks();
 

include/swift/Runtime/Concurrency.h

@@ -1023,6 +1023,7 @@ bool swift_task_isCurrentExecutor(SerialExecutorRef executor);
 /// this could be a pointer to a different enum instance if we need it to be.
 enum swift_task_is_current_executor_flag : uint64_t {
   /// We aren't passing any flags.
+  /// Effectively this is a backwards compatible mode.
   None = 0x0,
 
   /// This is not used today, but is just future ABI reservation.
@@ -1047,9 +1048,15 @@ enum swift_task_is_current_executor_flag : uint64_t {
   /// The routine should assert on failure.
   Assert = 0x8,
 
+  /// The routine MUST NOT assert on failure.
+  /// Even at the cost of not calling 'checkIsolated' if it is available.
+  MustNotAssert = 0x10,
+
   /// The routine should use 'isIsolatingCurrentContext' function on the
   /// 'expected' executor instead of `checkIsolated`.
-  HasIsIsolatingCurrentContext = 0x10,
+  ///
+  /// This is a variant of `MustNotAssert`
+  UseIsIsolatingCurrentContext = 0x20,
 };
 
 SWIFT_EXPORT_FROM(swift_Concurrency)

lib/AST/Module.cpp

@@ -2031,6 +2031,10 @@ bool ModuleDecl::isStdlibModule() const {
   return !getParent() && getName() == getASTContext().StdlibModuleName;
 }
 
+bool ModuleDecl::isConcurrencyModule() const {
+  return !getParent() && getName() == getASTContext().Id_Concurrency;
+}
+
 bool ModuleDecl::hasStandardSubstitutions() const {
   return !getParent() &&
       (getName() == getASTContext().StdlibModuleName ||
@@ -4285,11 +4289,6 @@ struct SwiftSettingsWalker : ASTWalker {
 
 } // namespace
 
-static bool isConcurrencyModule(DeclContext *dc) {
-  auto *m = dc->getParentModule();
-  return !m->getParent() && m->getName() == m->getASTContext().Id_Concurrency;
-}
-
 bool SwiftSettingsWalker::isSwiftSettingsMacroExpr(
     MacroExpansionExpr *macroExpr) {
   // First make sure we actually have a macro with the name SwiftSettings.
@@ -4406,7 +4405,8 @@ SwiftSettingsWalker::patternMatchDefaultIsolationMainActor(CallExpr *callExpr) {
   if (!nomDecl)
     return CanType();
   auto *nomDeclDC = nomDecl->getDeclContext();
-  if (!nomDeclDC->isModuleScopeContext() || !isConcurrencyModule(nomDeclDC))
+  auto *nomDeclModule = nomDecl->getParentModule();
+  if (!nomDeclDC->isModuleScopeContext() || !nomDeclModule->isConcurrencyModule())
     return CanType();
 
   return nomDecl->getDeclaredType()->getCanonicalType();

lib/IRGen/GenProto.cpp

@@ -2238,6 +2238,7 @@ namespace {
         Flags = Flags.withIsSynthesizedNonUnique(conf->isSynthesizedNonUnique());
         Flags = Flags.withIsConformanceOfProtocol(conf->isConformanceOfProtocol());
         Flags = Flags.withHasGlobalActorIsolation(isolation.isGlobalActor());
+        Flags = withSerialExecutorCheckingModeFlags(Flags, conf);
       } else {
         Flags = Flags.withIsRetroactive(false)
                      .withIsSynthesizedNonUnique(false);
@@ -2459,6 +2460,32 @@ namespace {
               rootGlobalActorConformance));
       B.addRelativeAddress(globalActorConformanceDescriptor);
     }
+
+    static ConformanceFlags
+    withSerialExecutorCheckingModeFlags(ConformanceFlags Flags, const NormalProtocolConformance *conf) {
+      ProtocolDecl *proto = conf->getProtocol();
+      auto &C = proto->getASTContext();
+
+      ConformanceFlags UpdatedFlags = Flags;
+      if (proto->isSpecificProtocol(swift::KnownProtocolKind::SerialExecutor)) {
+        conf->forEachValueWitness([&](const ValueDecl *req,
+                                      Witness witness) {
+          bool nameMatch = witness.getDecl()->getBaseIdentifier() == C.Id_isIsolatingCurrentContext;
+          if (nameMatch) {
+            if (DeclContext *NominalOrExtension = witness.getDecl()->getDeclContext()) {
+              // If the witness is NOT the default implementation in the _Concurrency library,
+              // we should record that this is an user provided implementation and we should call it.
+              bool hasNonDefaultIsIsolatingCurrentContext =
+                  !NominalOrExtension->getParentModule()->isConcurrencyModule();
+              UpdatedFlags = UpdatedFlags.withHasNonDefaultSerialExecutorIsIsolatingCurrentContext(
+                            hasNonDefaultIsIsolatingCurrentContext);
+            }
+          }
+        });
+      }
+
+      return UpdatedFlags;
+    }
   };
 }