[Distributed] Reimplement distributed call thunks completely in AST (#41616)

* [Distributed] dist actor always has default executor (currently)

* [Distributed] extra test for missing makeEncoder

* [DistributedDecl] Add DistributedActorSystem to known SDK types

* [DistributedActor] ok progress on getting the system via witness

* [Distributed] allow hop-to `let any: any X` where X is DistActor

* [Distributed] AST: Add an accessor to determine whether type is distributed actor

- Classes have specialized method on their declarations
- Archetypes and existentials check their conformances for
  presence of `DistributedActor` protocol.

* [Distributed] AST: Account for distributed members declared in class extensions

`getConcreteReplacementForProtocolActorSystemType` should use `getSelfClassDecl`
otherwise it wouldn't be able to find actor if the member is declared in an extension.

* [Distributed] fix ad-hoc requirement checks for 'mutating'

[PreChecker] LookupDC might be null, so account for that

* [Distributed] Completed AST synthesis for dist thunk

* [Distributed][ASTDumper] print pretty distributed in right color in AST dumps

* wip on making the local/remote calls

* using the _local to mark the localCall as known local

* [Distributed] fix passing Never when not throwing

* fix lifetime of mangled string

* [Distributed] Implement recordGenericSubstitution

* [Distributed] Dont add .

* [Distributed] dont emit thunk when func broken

* [Distributed] fix tests; cleanups

* [Distributed] cleanup, move is... funcs to DistributedDecl

* [Distributed] Remove SILGen for distributed thunks, it is in Sema now!

* [Distributed]  no need to check stored props in protocols

* remote not used flag

* fix mangling test

* [Distributed] Synthesis: Don't re-use AST nodes for `decodeArgument` references

* [Distributed] Synthesis: Make sure that each thunk parameter has an internal name

* [Distributed/Synthesis] NFC: Add a comment regarding empty internal parameter names

* [Distributed] NFC: Adjust distributed thunk manglings in the accessor section test-cases

* cleanup

* [Distributed] NFC: Adjust distributed thunk manglings in the accessor thunk test-cases

* review follow ups

* xfail some linux tests for now so we can land the AST thunk

* Update distributed_actor_remote_functions.swift

Co-authored-by: Pavel Yaskevich <[email protected]>

Author: ktoso

docs/ReferenceGuides/UnderscoredAttributes.md

@@ -830,10 +830,20 @@ the compiler.
 
 This `async` function uses the pre-SE-0338 semantics of unsafely inheriting the caller's executor.  This is an underscored feature because the right way of inheriting an executor is to pass in the required executor and switch to it.  Unfortunately, there are functions in the standard library which need to inherit their caller's executor but cannot change their ABI because they were not defined as `@_alwaysEmitIntoClient` in the initial release.
 
+
 ## `@_spi_available(platform, version)`
 
 Like `@available`, this attribute indicates a decl is available only as an SPI.
 This implies several behavioral changes comparing to regular `@available`:
 1. Type checker diagnoses when a client accidently exposes such a symbol in library APIs.
 2. When emitting public interfaces, `@_spi_available` is printed as `@available(platform, unavailable)`.
 3. ClangImporter imports ObjC macros `SPI_AVAILABLE` and `__SPI_AVAILABLE` to this attribute.
+
+
+## `_local`
+
+A distributed actor can be marked as "known to be local" which allows avoiding 
+the distributed actor isolation checks. This is used for things like `whenLocal`
+where the actor passed to the closure is known-to-be-local, and similarly a 
+`self` of obtained from an _isolated_ function inside a distributed actor is 
+also guaranteed to be local by construction.

include/swift/AST/ASTContext.h

@@ -658,12 +658,11 @@ class ASTContext final {
 
   /// Retrieve the declaration of DistributedActorSystem.make().
   ///
-  /// \param actorOrSystem distributed actor or actor system to get the
-  /// remoteCall function for. Since the method we're looking for is an ad-hoc
-  /// requirement, a specific type MUST be passed here as it is not possible
-  /// to obtain the decl from just the `DistributedActorSystem` protocol type.
+  /// \param thunk the function from which we'll be invoking things on the obtained
+  /// actor system; This way we'll always get the right type, taking care of any
+  /// where clauses etc.
   FuncDecl *getMakeInvocationEncoderOnDistributedActorSystem(
-      NominalTypeDecl *actorOrSystem) const;
+      AbstractFunctionDecl *thunk) const;
 
   // Retrieve the declaration of
   // DistributedInvocationEncoder.recordGenericSubstitution(_:).

include/swift/AST/Attr.def

@@ -57,6 +57,7 @@ TYPE_ATTR(async)
 TYPE_ATTR(Sendable)
 TYPE_ATTR(unchecked)
 TYPE_ATTR(_typeSequence)
+TYPE_ATTR(_local)
 
 // SIL-specific attributes
 TYPE_ATTR(block_storage)
@@ -720,6 +721,13 @@ DECL_ATTR(_backDeploy, BackDeploy,
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIBreakingToRemove,
   129)
 
+CONTEXTUAL_SIMPLE_DECL_ATTR(_local, KnownToBeLocal,
+  DeclModifier | OnFunc | OnParam | OnVar |
+  DistributedOnly |
+  ABIBreakingToAdd | ABIBreakingToRemove |
+  APIBreakingToAdd | APIBreakingToRemove,
+  130)
+
 // If you're adding a new underscored attribute here, please document it in
 // docs/ReferenceGuides/UnderscoredAttributes.md.
 

include/swift/AST/Decl.h

@@ -4490,10 +4490,6 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// semantics but has no corresponding witness table.
   bool isMarkerProtocol() const;
 
-  /// Is a protocol that can only be conformed by distributed actors.
-  /// Such protocols are allowed to contain distributed functions.
-  bool inheritsFromDistributedActor() const;
-
 private:
   void computeKnownProtocolKind() const;
 
@@ -5254,6 +5250,10 @@ class VarDecl : public AbstractStorageDecl {
   /// Does this have a 'distributed' modifier?
   bool isDistributed() const;
 
+  /// Is this var known to be a "local" distributed actor,
+  /// if so the implicit throwing ans some isolation checks can be skipped.
+  bool isKnownToBeLocal() const;
+
   /// Is this a stored property that will _not_ trigger any user-defined code
   /// upon any kind of access?
   bool isOrdinaryStoredProperty() const;
@@ -6308,6 +6308,13 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Returns 'true' if the function is distributed.
   bool isDistributed() const;
 
+  /// For a 'distributed' target (func or computed property),
+  /// get the 'thunk' responsible for performing the 'remoteCall'.
+  ///
+  /// \return the synthesized thunk, or null if the base of the call has
+  ///         diagnosed errors during type checking.
+  FuncDecl *getDistributedThunk() const;
+  
   /// Returns 'true' if the function has the @c @_backDeploy attribute.
   bool isBackDeployed() const;
 
@@ -6455,6 +6462,16 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// 'DistributedActorSystem' protocol.
   bool isDistributedActorSystemRemoteCall(bool isVoidReturn) const;
 
+  /// Determines whether this function is a 'makeInvocationEncoder' function,
+  /// which is used as ad-hoc protocol requirement by the
+  /// 'DistributedActorSystem' protocol.
+  bool isDistributedActorSystemMakeInvocationEncoder() const;
+
+  /// Determines if this function is a 'recordGenericSubstitution' function,
+  /// which is used as ad-hoc protocol requirement by the
+  /// 'DistributedTargetInvocationEncoder' protocol.
+  bool isDistributedTargetInvocationEncoderRecordGenericSubstitution() const;
+
   /// Determines if this function is a 'recordArgument' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationEncoder' protocol.

include/swift/AST/DiagnosticsSema.def

@@ -4595,6 +4595,9 @@ NOTE(note_distributed_actor_isolated_method,none,
 ERROR(distributed_actor_isolated_method,none,
       "only 'distributed' instance methods can be called on a potentially remote distributed actor",
       ())
+ERROR(distributed_local_cannot_be_used,none,
+      "'local' cannot be used in user-defined code currently",
+      ())
 ERROR(distributed_actor_func_param_not_codable,none,
       "parameter '%0' of type %1 in %2 does not conform to serialization requirement '%3'",
       (StringRef, Type, DescriptiveDeclKind, StringRef))

include/swift/AST/DistributedDecl.h

@@ -31,6 +31,16 @@ class DeclContext;
 class FuncDecl;
 class NominalTypeDecl;
 
+/// Determine the concrete type of 'ActorSystem' as seen from the member.
+/// E.g. when in a protocol, and trying to determine what the actor system was
+/// constrained to.
+///
+/// \param member the member from which context the lookup should be performed,
+///        e.g. a function or computed property.
+/// \return the concrete type of the ActorSystem to be used by this member,
+///         or null if no concrete actor system was found.
+Type getConcreteReplacementForProtocolActorSystemType(ValueDecl *member);
+
 /// Determine the `ActorSystem` type for the given actor.
 Type getDistributedActorSystemType(NominalTypeDecl *actor);
 
@@ -43,14 +53,12 @@ Type getDistributedActorIDType(NominalTypeDecl *actor);
 Type getDistributedSerializationRequirementType(
     NominalTypeDecl *nominal, ProtocolDecl *protocol);
 
-///// Determine the serialization requirement for the given actor, actor system
-///// or other type that has the SerializationRequirement associated type.
-//Type getDistributedSerializationRequirementType(
-//    NominalTypeDecl *nominal, ProtocolDecl *protocol);
-
-Type getDistributedActorSystemActorIDRequirementType(
-    NominalTypeDecl *system);
+/// Get the specific 'InvocationEncoder' type of a specific distributed actor
+/// system.
+Type getDistributedActorSystemInvocationEncoderType(NominalTypeDecl *system);
 
+/// Get the 'ActorID' type of a specific distributed actor system.
+Type getDistributedActorSystemActorIDRequirementType(NominalTypeDecl *system);
 
 /// Get the specific protocols that the `SerializationRequirement` specifies,
 /// and all parameters / return types of distributed targets must conform to.

include/swift/AST/Expr.h

@@ -313,7 +313,7 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     NumCaptures : 32
   );
 
-  SWIFT_INLINE_BITFIELD(ApplyExpr, Expr, 1+1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD(ApplyExpr, Expr, 1+1+1+1+1+1+1,
     ThrowsIsSet : 1,
     Throws : 1,
     ImplicitlyAsync : 1,
@@ -1908,6 +1908,11 @@ class TryExpr : public AnyTryExpr {
   static bool classof(const Expr *e) {
     return e->getKind() == ExprKind::Try;
   }
+
+  static TryExpr *createImplicit(ASTContext &ctx, SourceLoc tryLoc, Expr *sub,
+                                 Type type = Type()) {
+    return new (ctx) TryExpr(tryLoc, sub, type, /*implicit=*/true);
+  }
 };
 
 /// ForceTryExpr - A 'try!' surrounding an expression, marking that
@@ -2066,7 +2071,11 @@ class AwaitExpr final : public IdentityExpr {
             bool implicit = false)
     : IdentityExpr(ExprKind::Await, sub, type, implicit), AwaitLoc(awaitLoc) {
   }
-  
+
+  static AwaitExpr *createImplicit(ASTContext &ctx, SourceLoc awaitLoc, Expr *sub, Type type = Type()) {
+    return new (ctx) AwaitExpr(awaitLoc, sub, type, /*implicit=*/true);
+  }
+
   SourceLoc getLoc() const { return AwaitLoc; }
 
   SourceLoc getAwaitLoc() const { return AwaitLoc; }
@@ -4509,11 +4518,11 @@ class ApplyExpr : public Expr {
   }
 
   /// Is this application _implicitly_ required to be a throwing call?
-  /// This can happen if the function is actually a proxy function invocation,
-  /// which may throw, regardless of the target function throwing, e.g.
-  /// a distributed instance method call on a 'remote' actor, may throw due to network
-  /// issues reported by the transport, regardless if the actual target function
-  /// can throw.
+  /// This can happen if the function is actually a distributed thunk
+  /// invocation, which may throw, regardless of the target function throwing,
+  /// e.g. a distributed instance method call on a 'remote' actor, may throw due
+  /// to network issues reported by the transport, regardless if the actual
+  /// target function can throw.
   bool implicitlyThrows() const {
     return Bits.ApplyExpr.ImplicitlyThrows;
   }

include/swift/AST/KnownIdentifiers.def

@@ -293,6 +293,7 @@ IDENTIFIER(target)
 IDENTIFIER(throwing)
 IDENTIFIER(using)
 IDENTIFIER(InvocationDecoder)
+IDENTIFIER(InvocationEncoder)
 IDENTIFIER(whenLocal)
 IDENTIFIER(decodeNextArgument)
 IDENTIFIER(SerializationRequirement)

include/swift/AST/KnownSDKTypes.def

@@ -45,6 +45,9 @@ KNOWN_SDK_TYPE_DECL(Concurrency, TaskLocal, ClassDecl, 1)
 
 // Distributed actors
 KNOWN_SDK_TYPE_DECL(Distributed, DistributedActor, ProtocolDecl, 0)
+KNOWN_SDK_TYPE_DECL(Distributed, DistributedActorSystem, ProtocolDecl, 0)
+KNOWN_SDK_TYPE_DECL(Distributed, DistributedTargetInvocationEncoder, ProtocolDecl, 0)
+KNOWN_SDK_TYPE_DECL(Distributed, DistributedTargetInvocationDecoder, ProtocolDecl, 0)
 KNOWN_SDK_TYPE_DECL(Distributed, RemoteCallTarget, StructDecl, 0)
 
 // String processing

include/swift/AST/TypeCheckRequests.h

@@ -1184,6 +1184,27 @@ class GetDistributedRemoteCallTargetInitFunctionRequest :
     bool isCached() const { return true; }
 };
 
+/// Obtain the 'distributed thunk' for the passed-in function.
+///
+/// The thunk is responsible for invoking 'remoteCall' when invoked on a remote
+/// 'distributed actor'.
+class GetDistributedThunkRequest :
+    public SimpleRequest<GetDistributedThunkRequest,
+                         FuncDecl *(AbstractFunctionDecl *),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  FuncDecl *evaluate(Evaluator &evaluator, AbstractFunctionDecl *distributedFunc) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+};
+
 /// Obtain the 'id' property of a 'distributed actor'.
 class GetDistributedActorIDPropertyRequest :
     public SimpleRequest<GetDistributedActorIDPropertyRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -125,6 +125,9 @@ SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationDecoderDecodeNextArgume
 SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationResultHandlerOnReturnFunctionRequest,
               AbstractFunctionDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, GetDistributedThunkRequest,
+              FuncDecl *(AbstractFunctionDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetDistributedActorIDPropertyRequest,
               VarDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)

include/swift/AST/Types.h

@@ -773,6 +773,10 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Determines whether this type is an actor type.
   bool isActorType();
 
+  /// Determines whether this type conforms or inherits (if it's a protocol
+  /// type) from `DistributedActor`.
+  bool isDistributedActor();
+
   /// Determines the element type of a known
   /// [Autoreleasing]Unsafe[Mutable][Raw]Pointer variant, or returns null if the
   /// type is not a pointer.