[Distributed] Don't crash in thunk generation when missing SR conformance

Author: ktoso

include/swift/AST/DistributedDecl.h

@@ -123,9 +123,11 @@ getDistributedSerializationRequirements(
 /// Given any set of generic requirements, locate those which are about the
 /// `SerializationRequirement`. Those need to be applied in the parameter and
 /// return type checking of distributed targets.
-llvm::SmallPtrSet<ProtocolDecl *, 2>
+void
 extractDistributedSerializationRequirements(
-    ASTContext &C, ArrayRef<Requirement> allRequirements);
+    ASTContext &C,
+    ArrayRef<Requirement> allRequirements,
+    llvm::SmallPtrSet<ProtocolDecl *, 2> &into);
 
 }
 

lib/AST/DistributedDecl.cpp

@@ -1268,10 +1268,11 @@ AbstractFunctionDecl::isDistributedTargetInvocationResultHandlerOnReturn() const
     return true;
 }
 
-llvm::SmallPtrSet<ProtocolDecl *, 2>
+void
 swift::extractDistributedSerializationRequirements(
-    ASTContext &C, ArrayRef<Requirement> allRequirements) {
-  llvm::SmallPtrSet<ProtocolDecl *, 2> serializationReqs;
+    ASTContext &C,
+    ArrayRef<Requirement> allRequirements,
+    llvm::SmallPtrSet<ProtocolDecl *, 2> &into) {
   auto DA = C.getDistributedActorDecl();
   auto daSerializationReqAssocType =
       DA->getAssociatedType(C.Id_SerializationRequirement);
@@ -1293,7 +1294,7 @@ swift::extractDistributedSerializationRequirements(
         auto requirementProto = req.getSecondType();
         if (auto proto = dyn_cast_or_null<ProtocolDecl>(
                 requirementProto->getAnyNominal())) {
-          serializationReqs.insert(proto);
+          into.insert(proto);
         } else {
           auto serialReqType = requirementProto->castTo<ExistentialType>()
                                    ->getConstraintType()
@@ -1302,14 +1303,12 @@ swift::extractDistributedSerializationRequirements(
               flattenDistributedSerializationTypeToRequiredProtocols(
                   serialReqType);
           for (auto p : flattenedRequirements) {
-            serializationReqs.insert(p);
+            into.insert(p);
           }
         }
       }
     }
   }
-
-  return serializationReqs;
 }
 
 /******************************************************************************/

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -841,9 +841,15 @@ FuncDecl *GetDistributedThunkRequest::evaluate(Evaluator &evaluator,
     if (!distributedTarget->isDistributed())
       return nullptr;
   }
-
   assert(distributedTarget);
 
+  // This evaluation type-check by now was already computed and cached;
+  // We need to check in order to avoid emitting a THUNK for a distributed func
+  // which had errors; as the thunk then may also cause un-addressable issues and confusion.
+  if (swift::checkDistributedFunction(distributedTarget)) {
+    return nullptr;
+  }
+
   auto &C = distributedTarget->getASTContext();
 
   if (!getConcreteReplacementForProtocolActorSystemType(distributedTarget)) {

lib/Sema/TypeCheckDeclOverride.cpp

@@ -2067,7 +2067,7 @@ static bool checkSingleOverride(ValueDecl *override, ValueDecl *base) {
       return (prop &&
               prop->isFinal() &&
               isa<ClassDecl>(prop->getDeclContext()) &&
-              cast<ClassDecl>(prop->getDeclContext())->isActor() &&
+              cast<ClassDecl>(prop->getDeclContext())->isAnyActor() &&
               !prop->isStatic() &&
               prop->getName() == ctx.Id_unownedExecutor &&
               prop->getInterfaceType()->getAnyNominal() == ctx.getUnownedSerialExecutorDecl());

lib/Sema/TypeCheckDistributed.cpp

@@ -498,8 +498,22 @@ bool CheckDistributedFunctionRequest::evaluate(
   // SerializationRequirement
   llvm::SmallPtrSet<ProtocolDecl *, 2> serializationRequirements;
   if (auto extension = dyn_cast<ExtensionDecl>(DC)) {
-    serializationRequirements = extractDistributedSerializationRequirements(
-        C, extension->getGenericRequirements());
+    auto actorOrProtocol = extension->getExtendedNominal();
+    if (auto actor = dyn_cast<ClassDecl>(actorOrProtocol)) {
+      assert(actor->isAnyActor());
+      serializationRequirements = getDistributedSerializationRequirementProtocols(
+          getDistributedActorSystemType(actor)->getAnyNominal(),
+          C.getProtocol(KnownProtocolKind::DistributedActorSystem));
+    } else if (auto protocol = dyn_cast<ProtocolDecl>(actorOrProtocol)) {
+      extractDistributedSerializationRequirements(
+          C, protocol->getGenericRequirements(),
+          /*into=*/serializationRequirements);
+      extractDistributedSerializationRequirements(
+          C, extension->getGenericRequirements(),
+          /*into=*/serializationRequirements);
+    } else {
+      // ignore
+    }
   } else if (auto actor = dyn_cast<ClassDecl>(DC)) {
     serializationRequirements = getDistributedSerializationRequirementProtocols(
         getDistributedActorSystemType(actor)->getAnyNominal(),
@@ -546,6 +560,7 @@ bool CheckDistributedFunctionRequest::evaluate(
         if (auto paramNominalTy = paramTy->getAnyNominal()) {
           addCodableFixIt(paramNominalTy, diag);
         } // else, no nominal type to suggest the fixit for, e.g. a closure
+
         return true;
       }
     }
@@ -795,11 +810,11 @@ void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal
   (void)nominal->getDistributedActorIDProperty();
 }
 
-void TypeChecker::checkDistributedFunc(FuncDecl *func) {
+bool TypeChecker::checkDistributedFunc(FuncDecl *func) {
   if (!func->isDistributed())
-    return;
+    return false;
 
-  swift::checkDistributedFunction(func);
+  return swift::checkDistributedFunction(func);
 }
 
 llvm::SmallPtrSet<ProtocolDecl *, 2>

lib/Sema/TypeCheckStmt.cpp

@@ -2778,6 +2778,14 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
   // So, build out the body now.
   ASTScope::expandFunctionBody(AFD);
 
+  if (AFD->isDistributedThunk()) {
+    if (auto func = dyn_cast<FuncDecl>(AFD)) {
+      if (TypeChecker::checkDistributedFunc(func)) {
+        return errorBody();
+      }
+    }
+  }
+
   // Type check the function body if needed.
   bool hadError = false;
   if (!alreadyTypeChecked) {

lib/Sema/TypeChecker.h

@@ -1132,7 +1132,9 @@ diagnosePotentialUnavailability(SourceRange ReferenceRange,
 void checkDistributedActor(SourceFile *SF, NominalTypeDecl *decl);
 
 /// Type check a single 'distributed func' declaration.
-void checkDistributedFunc(FuncDecl *func);
+///
+/// Returns `true` if there was an error.
+bool checkDistributedFunc(FuncDecl *func);
 
 bool checkAvailability(SourceRange ReferenceRange,
                        AvailabilityContext Availability,

test/Distributed/distributed_actor_func_param_not_conforming_req_full.swift

@@ -0,0 +1,41 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-build-swift -module-name main -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s %S/Inputs/FakeDistributedActorSystems.swift 2> %t/output.txt || echo 'failed expectedly'
+// RUN: %FileCheck %s < %t/output.txt
+
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+import Distributed
+
+// Notes:
+// This test specifically is not just a -typecheck -verify test but attempts to generate the whole module.
+// This is because we may be emitting errors but otherwise still attempt to emit a thunk for an "error-ed"
+// distributed function, which would then crash in later phases of compilation when we try to get types
+// of the `func` the THUNK is based on.
+
+typealias DefaultDistributedActorSystem = LocalTestingDistributedActorSystem
+
+distributed actor Service {
+}
+
+extension Service {
+  distributed func boombox(_ id: Box) async throws {}
+  // CHECK: parameter '' of type 'Box' in distributed instance method does not conform to serialization requirement 'Codable'
+
+  distributed func boxIt() async throws -> Box { fatalError() }
+  // CHECK: result type 'Box' of distributed instance method 'boxIt' does not conform to serialization requirement 'Codable'
+}
+
+public enum Box: Hashable { case boom }
+
+@main struct Main {
+  static func main() async {
+    try? await Service(actorSystem: .init()).boombox(Box.boom)
+  }
+}
+