From fc7ecc8f6734870c368701f7830e8c439ab9a018 Mon Sep 17 00:00:00 2001
From: Konrad `ktoso` Malawski <konrad_malawski@apple.com>
Date: Fri, 16 Jun 2023 11:54:51 +0900
Subject: [PATCH] [Distributed] Harden typechecker against completely empty DAS
 types

---
 lib/AST/ASTContext.cpp                         | 18 ++++++++++++++++++
 ...tributed_imcomplete_system_dont_crash.swift |  8 ++++++++
 2 files changed, 26 insertions(+)

diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index 93c8f0f0959e4..2f1c6e4e624f3 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -1397,6 +1397,9 @@ FuncDecl *ASTContext::getMakeInvocationEncoderOnDistributedActorSystem(
 FuncDecl *
 ASTContext::getRecordGenericSubstitutionOnDistributedInvocationEncoder(
     NominalTypeDecl *nominal) const {
+  if (!nominal)
+    return nullptr;
+
   for (auto result : nominal->lookupDirect(Id_recordGenericSubstitution)) {
     auto *func = dyn_cast<FuncDecl>(result);
     if (func &&
@@ -1410,6 +1413,9 @@ ASTContext::getRecordGenericSubstitutionOnDistributedInvocationEncoder(
 
 AbstractFunctionDecl *ASTContext::getRecordArgumentOnDistributedInvocationEncoder(
     NominalTypeDecl *nominal) const {
+  if (!nominal)
+    return nullptr;
+
   return evaluateOrDefault(
       nominal->getASTContext().evaluator,
       GetDistributedTargetInvocationEncoderRecordArgumentFunctionRequest{nominal},
@@ -1418,6 +1424,9 @@ AbstractFunctionDecl *ASTContext::getRecordArgumentOnDistributedInvocationEncode
 
 AbstractFunctionDecl *ASTContext::getRecordReturnTypeOnDistributedInvocationEncoder(
     NominalTypeDecl *nominal) const {
+  if (!nominal)
+    return nullptr;
+
   return evaluateOrDefault(
       nominal->getASTContext().evaluator,
       GetDistributedTargetInvocationEncoderRecordReturnTypeFunctionRequest{nominal},
@@ -1426,6 +1435,9 @@ AbstractFunctionDecl *ASTContext::getRecordReturnTypeOnDistributedInvocationEnco
 
 AbstractFunctionDecl *ASTContext::getRecordErrorTypeOnDistributedInvocationEncoder(
     NominalTypeDecl *nominal) const {
+  if (!nominal)
+    return nullptr;
+
   return evaluateOrDefault(
       nominal->getASTContext().evaluator,
       GetDistributedTargetInvocationEncoderRecordErrorTypeFunctionRequest{nominal},
@@ -1434,6 +1446,9 @@ AbstractFunctionDecl *ASTContext::getRecordErrorTypeOnDistributedInvocationEncod
 
 AbstractFunctionDecl *ASTContext::getDecodeNextArgumentOnDistributedInvocationDecoder(
     NominalTypeDecl *nominal) const {
+  if (!nominal)
+    return nullptr;
+
   return evaluateOrDefault(
       nominal->getASTContext().evaluator,
       GetDistributedTargetInvocationDecoderDecodeNextArgumentFunctionRequest{nominal},
@@ -1442,6 +1457,9 @@ AbstractFunctionDecl *ASTContext::getDecodeNextArgumentOnDistributedInvocationDe
 
 AbstractFunctionDecl *ASTContext::getOnReturnOnDistributedTargetInvocationResultHandler(
     NominalTypeDecl *nominal) const {
+  if (!nominal)
+    return nullptr;
+
   return evaluateOrDefault(
       nominal->getASTContext().evaluator,
       GetDistributedTargetInvocationResultHandlerOnReturnFunctionRequest{nominal},
diff --git a/test/Distributed/distributed_imcomplete_system_dont_crash.swift b/test/Distributed/distributed_imcomplete_system_dont_crash.swift
index a2cc7d2ca3772..e4f5d3e7f9dce 100644
--- a/test/Distributed/distributed_imcomplete_system_dont_crash.swift
+++ b/test/Distributed/distributed_imcomplete_system_dont_crash.swift
@@ -36,4 +36,12 @@ public final class CompletelyHollowActorSystem: DistributedActorSystem {
     // expected-error@-1{{type 'CompletelyHollowActorSystem.ResultHandler' does not conform to protocol 'DistributedTargetInvocationResultHandler'}}
   }
 
+}
+
+public final class CompletelyHollowActorSystem_NotEvenTypes: DistributedActorSystem {
+  // expected-error@-1{{type 'CompletelyHollowActorSystem_NotEvenTypes' does not conform to protocol 'DistributedActorSystem'}}
+  // expected-error@-2{{class 'CompletelyHollowActorSystem_NotEvenTypes' is missing witness for protocol requirement 'remoteCallVoid'}}
+  // expected-error@-3{{class 'CompletelyHollowActorSystem_NotEvenTypes' is missing witness for protocol requirement 'remoteCall'}}
+  // expected-note@-4{{protocol 'DistributedActorSystem' requires function 'remoteCall' with signature:}}
+  // expected-note@-5{{protocol 'DistributedActorSystem' requires function 'remoteCallVoid' with signature:}}
 }
\ No newline at end of file