[Concurrency] Custom executors with move-only Job (#63569)

Author: ktoso

include/swift/ABI/Task.h

@@ -136,6 +136,10 @@ class alignas(2 * alignof(void*)) Job :
     return Flags.getPriority();
   }
 
+  uint32_t getJobId() const {
+    return Id;
+  }
+
   /// Given that we've fully established the job context in the current
   /// thread, actually start running this job.  To establish the context
   /// correctly, call swift_job_run or runJobInExecutorContext.

include/swift/AST/Decl.h

@@ -3860,6 +3860,9 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// Find the 'RemoteCallArgument(label:name:value:)' initializer function.
   ConstructorDecl *getDistributedRemoteCallArgumentInitFunction() const;
 
+  /// Get the move-only `enqueue(Job)` protocol requirement function on the `Executor` protocol.
+  AbstractFunctionDecl *getExecutorOwnedEnqueueFunction() const;
+
   /// Collect the set of protocols to which this type should implicitly
   /// conform, such as AnyObject (for classes).
   void getImplicitProtocols(SmallVectorImpl<ProtocolDecl *> &protocols);

include/swift/AST/DiagnosticsSema.def

@@ -6408,6 +6408,11 @@ WARNING(hashvalue_implementation,Deprecation,
         "conform type %0 to 'Hashable' by implementing 'hash(into:)' instead",
         (Type))
 
+WARNING(executor_enqueue_unowned_implementation,Deprecation,
+        "'Executor.enqueue(UnownedJob)' is deprecated as a protocol requirement; "
+        "conform type %0 to 'Executor' by implementing 'func enqueue(Job)' instead",
+        (Type))
+
 //------------------------------------------------------------------------------
 // MARK: property wrapper diagnostics
 //------------------------------------------------------------------------------

include/swift/AST/KnownProtocols.def

@@ -76,6 +76,7 @@ PROTOCOL(SIMDScalar)
 PROTOCOL(BinaryInteger)
 PROTOCOL(FixedWidthInteger)
 PROTOCOL(RangeReplaceableCollection)
+PROTOCOL(Executor)
 PROTOCOL(SerialExecutor)
 PROTOCOL(GlobalActor)
 

include/swift/AST/KnownSDKTypes.def

@@ -39,6 +39,10 @@ KNOWN_SDK_TYPE_DECL(ObjectiveC, ObjCBool, StructDecl, 0)
 // standardized
 KNOWN_SDK_TYPE_DECL(Concurrency, UnsafeContinuation, NominalTypeDecl, 2)
 KNOWN_SDK_TYPE_DECL(Concurrency, MainActor, NominalTypeDecl, 0)
+KNOWN_SDK_TYPE_DECL(Concurrency, Job, StructDecl, 0)
+KNOWN_SDK_TYPE_DECL(Concurrency, UnownedJob, StructDecl, 0)
+KNOWN_SDK_TYPE_DECL(Concurrency, Executor, NominalTypeDecl, 0)
+KNOWN_SDK_TYPE_DECL(Concurrency, SerialExecutor, NominalTypeDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, UnownedSerialExecutor, NominalTypeDecl, 0)
 
 KNOWN_SDK_TYPE_DECL(Concurrency, TaskLocal, ClassDecl, 1)

include/swift/Runtime/Concurrency.h

@@ -715,6 +715,11 @@ bool swift_task_isOnExecutor(
     const Metadata *selfType,
     const SerialExecutorWitnessTable *wtable);
 
+/// Return the 64bit TaskID (if the job is an AsyncTask),
+/// or the 32bits of the job Id otherwise.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+uint64_t swift_task_getJobTaskId(Job *job);
+
 #if SWIFT_CONCURRENCY_ENABLE_DISPATCH
 
 /// Enqueue the given job on the main executor.

lib/AST/ASTContext.cpp

@@ -1120,6 +1120,7 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   case KnownProtocolKind::GlobalActor:
   case KnownProtocolKind::AsyncSequence:
   case KnownProtocolKind::AsyncIteratorProtocol:
+  case KnownProtocolKind::Executor:
   case KnownProtocolKind::SerialExecutor:
     M = getLoadedModule(Id_Concurrency);
     break;

lib/AST/Decl.cpp

@@ -5248,6 +5248,39 @@ VarDecl *NominalTypeDecl::getGlobalActorInstance() const {
                            nullptr);
 }
 
+AbstractFunctionDecl *
+NominalTypeDecl::getExecutorOwnedEnqueueFunction() const {
+  auto &C = getASTContext();
+
+  auto proto = dyn_cast<ProtocolDecl>(this);
+  if (!proto)
+    return nullptr;
+
+  llvm::SmallVector<ValueDecl *, 2> results;
+  lookupQualified(getSelfNominalTypeDecl(),
+                  DeclNameRef(C.Id_enqueue),
+                  NL_ProtocolMembers,
+                  results);
+
+  for (auto candidate: results) {
+    // we're specifically looking for the Executor protocol requirement
+    if (!isa<ProtocolDecl>(candidate->getDeclContext()))
+      continue;
+
+    if (auto *funcDecl = dyn_cast<AbstractFunctionDecl>(candidate)) {
+      if (funcDecl->getParameters()->size() != 1)
+        continue;
+
+      auto params = funcDecl->getParameters();
+      if (params->get(0)->getSpecifier() == ParamSpecifier::LegacyOwned) { // TODO: make this Consuming
+        return funcDecl;
+      }
+    }
+  }
+
+  return nullptr;
+}
+
 ClassDecl::ClassDecl(SourceLoc ClassLoc, Identifier Name, SourceLoc NameLoc,
                      ArrayRef<InheritedEntry> Inherited,
                      GenericParamList *GenericParams, DeclContext *Parent,

lib/IRGen/GenMeta.cpp

@@ -6231,6 +6231,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::CxxSequence:
   case KnownProtocolKind::UnsafeCxxInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:
+  case KnownProtocolKind::Executor:
   case KnownProtocolKind::SerialExecutor:
   case KnownProtocolKind::Sendable:
   case KnownProtocolKind::UnsafeSendable:

lib/Sema/TypeCheckConcurrency.cpp

@@ -1225,6 +1225,79 @@ void swift::diagnoseMissingExplicitSendable(NominalTypeDecl *nominal) {
   }
 }
 
+void swift::tryDiagnoseExecutorConformance(ASTContext &C,
+                                           const NominalTypeDecl *nominal,
+                                           ProtocolDecl *proto) {
+  assert(proto->isSpecificProtocol(KnownProtocolKind::Executor) ||
+         proto->isSpecificProtocol(KnownProtocolKind::SerialExecutor));
+
+  auto &diags = C.Diags;
+  auto module = nominal->getParentModule();
+  Type nominalTy = nominal->getDeclaredInterfaceType();
+
+  // enqueue(_: UnownedJob)
+  auto enqueueDeclName = DeclName(C, DeclBaseName(C.Id_enqueue), { Identifier() });
+
+  FuncDecl *unownedEnqueueRequirement = nullptr;
+  FuncDecl *moveOnlyEnqueueRequirement = nullptr;
+  for (auto req: proto->getProtocolRequirements()) {
+    auto *funcDecl = dyn_cast<FuncDecl>(req);
+    if (!funcDecl)
+      continue;
+
+    if (funcDecl->getName() != enqueueDeclName)
+      continue;
+
+
+    // look for the first parameter being a Job or UnownedJob
+    if (funcDecl->getParameters()->size() != 1)
+      continue;
+    if (auto param = funcDecl->getParameters()->front()) {
+      if (param->getType()->isEqual(C.getJobDecl()->getDeclaredInterfaceType())) {
+        assert(moveOnlyEnqueueRequirement == nullptr);
+        moveOnlyEnqueueRequirement = funcDecl;
+      } else if (param->getType()->isEqual(C.getUnownedJobDecl()->getDeclaredInterfaceType())) {
+        assert(unownedEnqueueRequirement == nullptr);
+        unownedEnqueueRequirement = funcDecl;
+      }
+    }
+
+    // if we found both, we're done here and break out of the loop
+    if (unownedEnqueueRequirement && moveOnlyEnqueueRequirement)
+      break; // we're done looking for the requirements
+  }
+
+
+  auto conformance = module->lookupConformance(nominalTy, proto);
+  auto concreteConformance = conformance.getConcrete();
+  auto unownedEnqueueWitness = concreteConformance->getWitnessDeclRef(unownedEnqueueRequirement);
+  auto moveOnlyEnqueueWitness = concreteConformance->getWitnessDeclRef(moveOnlyEnqueueRequirement);
+
+  if (auto enqueueUnownedDecl = unownedEnqueueWitness.getDecl()) {
+    // Old UnownedJob based impl is present, warn about it suggesting the new protocol requirement.
+    if (enqueueUnownedDecl->getLoc().isValid()) {
+      diags.diagnose(enqueueUnownedDecl->getLoc(), diag::executor_enqueue_unowned_implementation, nominalTy);
+    }
+  }
+
+  if (auto unownedEnqueueDecl = unownedEnqueueWitness.getDecl()) {
+    if (auto moveOnlyEnqueueDecl = moveOnlyEnqueueWitness.getDecl()) {
+      if (unownedEnqueueDecl && unownedEnqueueDecl->getLoc().isInvalid() &&
+          moveOnlyEnqueueDecl && moveOnlyEnqueueDecl->getLoc().isInvalid()) {
+        // Neither old nor new implementation have been found, but we provide default impls for them
+        // that are mutually recursive, so we must error and suggest implementing the right requirement.
+        auto ownedRequirement = C.getExecutorDecl()->getExecutorOwnedEnqueueFunction();
+        nominal->diagnose(diag::type_does_not_conform, nominalTy, proto->getDeclaredInterfaceType());
+        ownedRequirement->diagnose(diag::no_witnesses,
+                                   getProtocolRequirementKind(ownedRequirement),
+                                   ownedRequirement->getName(),
+                                   proto->getDeclaredInterfaceType(),
+                                   /*AddFixIt=*/true);
+      }
+    }
+  }
+}
+
 /// Determine whether this is the main actor type.
 static bool isMainActor(Type type) {
   if (auto nominal = type->getAnyNominal())