review feedback, cleanup getting decode func

Author: ktoso

include/swift/AST/Decl.h

@@ -3564,10 +3564,10 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   VarDecl *getDistributedActorIDProperty() const;
 
   /// Find the 'RemoteCallTarget.init(_:)' initializer function.
-  ConstructorDecl* getDistributedRemoteCallTargetInitFunction() const;
+  ConstructorDecl *getDistributedRemoteCallTargetInitFunction() const;
 
   /// Find the 'RemoteCallArgument(label:name:value:)' initializer function.
-  ConstructorDecl* getDistributedRemoteCallArgumentInitFunction() const;
+  ConstructorDecl *getDistributedRemoteCallArgumentInitFunction() const;
 
   /// Collect the set of protocols to which this type should implicitly
   /// conform, such as AnyObject (for classes).

include/swift/AST/DistributedDecl.h

@@ -53,6 +53,13 @@ Type getDistributedActorIDType(NominalTypeDecl *actor);
 Type getDistributedSerializationRequirementType(
     NominalTypeDecl *nominal, ProtocolDecl *protocol);
 
+/// Given a distributed thunk declaration, inside a 'distributed actor',
+/// finds the ad-hoc witness for 'decodeNextArgument' on the associated
+/// 'ActorSystem.InvocationDecoder' of the actor, or null.
+AbstractFunctionDecl *
+getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+    ValueDecl *thunk);
+
 /// Get the specific 'InvocationEncoder' type of a specific distributed actor
 /// system.
 Type getDistributedActorSystemInvocationEncoderType(NominalTypeDecl *system);

lib/AST/DistributedDecl.cpp

@@ -203,6 +203,32 @@ Type swift::getDistributedSerializationRequirementType(
   return conformance.getTypeWitnessByName(selfType, ctx.Id_SerializationRequirement);
 }
 
+AbstractFunctionDecl *
+swift::getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+    ValueDecl *thunk) {
+  assert(thunk);
+  auto &C = thunk->getASTContext();
+
+  auto *actor = thunk->getDeclContext()->getSelfNominalTypeDecl();
+  if (!actor)
+    return nullptr;
+  if (!actor->isDistributedActor())
+    return nullptr;
+
+  auto systemTy = getConcreteReplacementForProtocolActorSystemType(thunk);
+  if (!systemTy)
+    return nullptr;
+
+  auto decoderTy =
+      getDistributedActorSystemInvocationDecoderType(
+          systemTy->getAnyNominal());
+  if (!decoderTy)
+    return nullptr;
+
+  return C.getDecodeNextArgumentOnDistributedInvocationDecoder(
+      decoderTy->getAnyNominal());
+}
+
 Type ASTContext::getAssociatedTypeOfDistributedSystemOfActor(
     NominalTypeDecl *actor, Identifier member) {
   auto &ctx = actor->getASTContext();

lib/IRGen/GenDistributed.cpp

@@ -813,6 +813,7 @@ ArgumentDecoderInfo DistributedAccessor::findArgumentDecoder(
     decoder = instance.claimNext();
   }
 
+  // TODO(distributed): this can be removed most likely now
   if (isa<StructDecl>(decoderDecl) || isa<EnumDecl>(decoderDecl) ||
       decoderDecl->isFinal()) {
     auto *decodeSIL = IGM.getSILModule().lookUpFunction(SILDeclRef(decodeFn));

lib/SIL/IR/SILFunctionBuilder.cpp

@@ -220,25 +220,14 @@ void SILFunctionBuilder::addFunctionAttributes(
 
       F->setDynamicallyReplacedFunction(replacedFunc);
     }
-  }
-
-  if (constant.isDistributedThunk()) {
-    auto *actor = decl->getDeclContext()->getSelfNominalTypeDecl();
-    if (actor && actor->isDistributedActor()) {
-      auto &C = decl->getASTContext();
-      auto systemTy = getConcreteReplacementForProtocolActorSystemType(decl);
-      assert(systemTy);
-
-      auto decoderTy =
-          getDistributedActorSystemInvocationDecoderType(
-              systemTy->getAnyNominal());
-      assert(decoderTy);
-
-      auto decodeFunc = C.getDecodeNextArgumentOnDistributedInvocationDecoder(
-          decoderTy->getAnyNominal());
-      auto decodeRef = SILDeclRef(decodeFunc);
-      auto *adHocWitness = getOrCreateDeclaration(decodeFunc, decodeRef);
-      F->setReferencedAdHocRequirementWitnessFunction(adHocWitness);
+  } else if (constant.isDistributedThunk()) {
+    if (auto decodeFuncDecl =
+            getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+                decl)) {
+
+      auto decodeRef = SILDeclRef(decodeFuncDecl);
+      auto *adHocFunc = getOrCreateDeclaration(decodeFuncDecl, decodeRef);
+      F->setReferencedAdHocRequirementWitnessFunction(adHocFunc);
     }
   }
 }