[Distributed] Target identifiers for protocol calls (#70928)

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -12,24 +12,26 @@
 
 #include "TypeCheckDistributed.h"
 
-#include "TypeChecker.h"
+#include "CodeSynthesis.h"
+#include "DerivedConformances.h"
 #include "TypeCheckType.h"
+#include "TypeChecker.h"
+#include "swift/AST/ASTMangler.h"
 #include "swift/AST/ASTPrinter.h"
 #include "swift/AST/Availability.h"
+#include "swift/AST/DistributedDecl.h"
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Initializer.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/TypeCheckRequests.h"
-#include "swift/AST/NameLookupRequests.h"
-#include "swift/AST/ASTMangler.h"
-#include "swift/AST/DistributedDecl.h"
 #include "swift/Basic/Defer.h"
 #include "swift/ClangImporter/ClangModule.h"
 #include "swift/Sema/ConstraintSystem.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
-#include "DerivedConformances.h"
 
 using namespace swift;
 
@@ -518,10 +520,33 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
   {
     // --- Mangle the thunk name
     Mangle::ASTMangler mangler;
-    auto mangled =
-        C.AllocateCopy(mangler.mangleDistributedThunk(cast<FuncDecl>(thunk)));
-    auto mangledTargetStringLiteral =
-        new (C) StringLiteralExpr(mangled, SourceRange(), implicit);
+
+    // FIXME: cleanup
+    StringLiteralExpr *mangledTargetStringLiteral = nullptr;
+    auto witnessedDistributedRequirements =
+        func->getDistributedMethodWitnessedProtocolRequirements();
+
+    if (witnessedDistributedRequirements.size() == 1) {
+      auto protocolFunc = witnessedDistributedRequirements.front();
+
+      // we expect to witness exactly one distributed requirement,
+      // otherwise we should have diagnosed errors about more than 1 already.
+      std::string mangledString =
+          mangler.mangleDistributedThunk(cast<FuncDecl>(protocolFunc));
+      // FIXME: make it THUNK so the mangling is right
+      // MUST BE LIKE: s4main28GreeterP_ConcreteSystem_StubC5greetSSyYaKFTE
+
+      StringRef mangled = C.AllocateCopy(mangledString);
+      mangledTargetStringLiteral =
+          new (C) StringLiteralExpr(mangled, SourceRange(), implicit);
+    } else {
+      // default mangling
+      auto mangled =
+          C.AllocateCopy(mangler.mangleDistributedThunk(cast<FuncDecl>(thunk)));
+      mangledTargetStringLiteral =
+          new (C) StringLiteralExpr(mangled, SourceRange(), implicit);
+    }
+    assert(mangledTargetStringLiteral && "must be initialized");
 
     // --- let target = RemoteCallTarget(<mangled name>)
     targetVar->setInterfaceType(remoteCallTargetTy);
@@ -640,31 +665,15 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
   return {body, /*isTypeChecked=*/false};
 }
 
-static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
+/// Create a new FuncDecl that has the same signature as the passed in func.
+/// This is used both to create stub witnesses as well as distributed thunks.
+///
+/// \param DC The declaration context of the newly created function
+static FuncDecl *
+createSameSignatureFunctionDecl(DeclContext *DC, FuncDecl *func,
+                                llvm::Optional<DeclName> nameOverride,
+                                bool forceAsync, bool forceThrows) {
   auto &C = func->getASTContext();
-  auto DC = func->getDeclContext();
-
-  // NOTE: So we don't need a thunk in the protocol, we should call the underlying
-  // thing instead, which MUST have a thunk, since it must be a distributed func as well...
-  if (isa<ProtocolDecl>(DC)) {
-    return nullptr;
-  }
-
-  assert(getConcreteReplacementForProtocolActorSystemType(func) &&
-         "Thunk synthesis must have concrete actor system type available");
-
-  DeclName thunkName;
-
-  // Since accessors don't have names, let's generate one based on
-  // the computed property.
-  if (auto *accessor = dyn_cast<AccessorDecl>(func)) {
-    auto *var = accessor->getStorage();
-    thunkName = DeclName(C, var->getBaseName(),
-                         /*argumentNames=*/ArrayRef<Identifier>());
-  } else {
-    // Let's use the name of a 'distributed func'
-    thunkName = func->getName();
-  }
 
   // --- Prepare generic parameters
   GenericParamList *genericParamList = nullptr;
@@ -689,9 +698,9 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
     }
 
     auto paramDecl = new (C)
-         ParamDecl(SourceLoc(),
-                   /*argumentNameLoc=*/SourceLoc(), funcParam->getArgumentName(),
-                   /*parameterNameLoc=*/SourceLoc(), paramName, DC);
+        ParamDecl(SourceLoc(),
+                  /*argumentNameLoc=*/SourceLoc(), funcParam->getArgumentName(),
+                  /*parameterNameLoc=*/SourceLoc(), paramName, DC);
 
     paramDecl->setImplicit(true);
     paramDecl->setSpecifier(funcParam->getSpecifier());
@@ -701,11 +710,56 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   }
   ParameterList *params = ParameterList::create(C, paramDecls); // = funcParams->clone(C);
 
-  FuncDecl *thunk = FuncDecl::createImplicit(
-      C, swift::StaticSpellingKind::None, thunkName, SourceLoc(),
-      /*async=*/true, /*throws=*/true, /*thrownType=*/Type(),
-      genericParamList, params, func->getResultInterfaceType(), DC);
+  DeclName funcName = nameOverride.value_or(func->getName());
 
+  FuncDecl *copy = FuncDecl::createImplicit(
+      C, swift::StaticSpellingKind::None, funcName, SourceLoc(),
+      /*async=*/forceAsync || func->hasAsync(),
+      /*throws=*/forceThrows || func->hasThrows(),
+      /*thrownType=*/Type(), genericParamList, params,
+      func->getResultInterfaceType(), DC);
+
+  copy->setSynthesized(true);
+
+  if (isa<ClassDecl>(DC))
+    copy->getAttrs().add(new (C) FinalAttr(/*isImplicit=*/true));
+
+  copy->setGenericSignature(baseSignature);
+  copy->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
+
+  return copy;
+}
+
+static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
+  auto &C = func->getASTContext();
+  auto DC = func->getDeclContext();
+
+  // NOTE: So we don't need a thunk in the protocol, we should call the
+  // underlying thing instead, which MUST have a thunk, since it must be a
+  // distributed func as well...
+  if (isa<ProtocolDecl>(DC)) {
+    return nullptr;
+  }
+
+  assert(getConcreteReplacementForProtocolActorSystemType(func) &&
+         "Thunk synthesis must have concrete actor system type available");
+
+  DeclName thunkName;
+
+  // Since accessors don't have names, let's generate one based on
+  // the computed property.
+  if (auto *accessor = dyn_cast<AccessorDecl>(func)) {
+    auto *var = accessor->getStorage();
+    thunkName = DeclName(C, var->getBaseName(),
+                         /*argumentNames=*/ArrayRef<Identifier>());
+  } else {
+    // Let's use the name of a 'distributed func'
+    thunkName = func->getName();
+  }
+
+  FuncDecl *thunk = createSameSignatureFunctionDecl(DC, func, thunkName,
+                                                    /*forceAsync=*/true,
+                                                    /*forceThrows=*/true);
   assert(thunk && "couldn't create a distributed thunk");
 
   thunk->setSynthesized(true);
@@ -713,13 +767,12 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   thunk->getAttrs().add(
       new (C) NonisolatedAttr(/*unsafe=*/false, /*implicit=*/true));
 
-  if (isa<ClassDecl>(DC))
-    thunk->getAttrs().add(new (C) FinalAttr(/*isImplicit=*/true));
-
-  thunk->setGenericSignature(baseSignature);
-  thunk->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
   thunk->setBodySynthesizer(deriveBodyDistributed_thunk, func);
 
+  /// Record which function this is a thunk for, we'll need this to link back
+  /// calls in case this is a distributed requirement witness.
+  thunk->getAttrs().add(new (C) DistributedThunkTargetAttr(func));
+
   return thunk;
 }
 
@@ -795,6 +848,44 @@ addDistributedActorCodableConformance(
   return conformance;
 }
 
+/******************************************************************************/
+/******************************************************************************/
+
+void swift::assertRequiredSynthesizedPropertyOrder(ASTContext &Context,
+                                                   NominalTypeDecl *nominal) {
+#ifndef NDEBUG
+  if (auto id = nominal->getDistributedActorIDProperty()) {
+    if (auto system = nominal->getDistributedActorSystemProperty()) {
+      if (auto classDecl = dyn_cast<ClassDecl>(nominal)) {
+        if (auto unownedExecutor = classDecl->getUnownedExecutorProperty()) {
+          int idIdx, actorSystemIdx, unownedExecutorIdx = 0;
+          int idx = 0;
+          for (auto member : nominal->getMembers()) {
+            if (auto binding = dyn_cast<PatternBindingDecl>(member)) {
+              if (binding->getSingleVar()->getName() == Context.Id_id) {
+                idIdx = idx;
+              } else if (binding->getSingleVar()->getName() ==
+                         Context.Id_actorSystem) {
+                actorSystemIdx = idx;
+              } else if (binding->getSingleVar()->getName() ==
+                         Context.Id_unownedExecutor) {
+                unownedExecutorIdx = idx;
+              }
+              idx += 1;
+            }
+          }
+          if (idIdx + actorSystemIdx + unownedExecutorIdx >= 0 + 1 + 2) {
+            // we have found all the necessary fields, let's assert their order
+            assert(idIdx < actorSystemIdx < unownedExecutorIdx &&
+                   "order of fields MUST be exact.");
+          }
+        }
+      }
+    }
+  }
+#endif
+}
+
 /******************************************************************************/
 /*********************** SYNTHESIS ENTRY POINTS *******************************/
 /******************************************************************************/