[Distributed] Pass args from Invocation to HBuffer

Author: ktoso

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -9,15 +9,13 @@
 // See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 //===----------------------------------------------------------------------===//
-
 import Swift
 import _Concurrency
 
 @available(SwiftStdlib 5.6, *)
 public protocol DistributedActorSystem: Sendable {
   /// The identity used by actors that communicate via this transport
   associatedtype ActorID: Sendable & Hashable & Codable // TODO: make Codable conditional here
-
   /// The specific type of the argument builder to be used for remote calls.
   associatedtype Invocation: DistributedTargetInvocation
 
@@ -26,7 +24,6 @@ public protocol DistributedActorSystem: Sendable {
 
   // ==== ---------------------------------------------------------------------
   // - MARK: Resolving actors by identity
-
   /// Resolve a local or remote actor address to a real actor instance, or throw if unable to.
   /// The returned value is either a local actor or proxy to a remote actor.
   ///
@@ -48,12 +45,11 @@ public protocol DistributedActorSystem: Sendable {
   /// Detecting liveness of such remote actors shall be offered / by transport libraries
   /// by other means, such as "watching an actor for termination" or similar.
   func resolve<Act>(id: ActorID, as actorType: Act.Type) throws -> Act?
-      where Act: DistributedActor,
-            Act.ID == ActorID
+    where Act: DistributedActor,
+          Act.ID == ActorID
 
   // ==== ---------------------------------------------------------------------
   // - MARK: Actor Lifecycle
-
   /// Create an `ActorID` for the passed actor type.
   ///
   /// This function is invoked by an distributed actor during its initialization,
@@ -65,8 +61,8 @@ public protocol DistributedActorSystem: Sendable {
   /// `system.resolve(id: addr1, as: Greeter.self)` MUST return a reference
   /// to the same actor.
   func assignID<Act>(_ actorType: Act.Type) -> ActorID
-      where Act: DistributedActor,
-            Act.ID == ActorID
+    where Act: DistributedActor,
+          Act.ID == ActorID
 
   /// Invoked during a distributed actor's initialization, as soon as it becomes fully initialized.
   ///
@@ -84,8 +80,8 @@ public protocol DistributedActorSystem: Sendable {
   ///
   /// - Parameter actor: reference to the (local) actor that was just fully initialized.
   func actorReady<Act>(_ actor: Act)
-      where Act: DistributedActor,
-            Act.ID == ActorID
+    where Act: DistributedActor,
+          Act.ID == ActorID
 
   /// Called during when a distributed actor is deinitialized, or fails to initialize completely (e.g. by throwing
   /// out of an `init` that did not completely initialize all of the the actors stored properties yet).
@@ -102,10 +98,9 @@ public protocol DistributedActorSystem: Sendable {
 
   // ==== ---------------------------------------------------------------------
   // - MARK: Remote Method Invocations
-
   /// Invoked by the Swift runtime when a distributed remote call is about to be made.
   ///
-  /// The returned DistributedTargetInvocation will be populated with all
+  /// The returned `DistributedTargetInvocation` will be populated with all
   /// arguments, generic substitutions, and specific error and return types
   /// that are associated with this specific invocation.
   @inlinable
@@ -124,19 +119,17 @@ public protocol DistributedActorSystem: Sendable {
 //  func remoteCall<Act, Err, Res>(
 //      on actor: Act,
 //      target: RemoteCallTarget,
-//      arguments: Invocation,
+//      invocation: inout Invocation,
 //      throwing: Err.Type,
 //      returning: Res.Type
-//  ) async throws -> Res.Type
+//  ) async throws -> Res
 //      where Act: DistributedActor,
 //            Act.ID == ActorID,
 //            Res: SerializationRequirement
-
 }
 
 // ==== ----------------------------------------------------------------------------------------------------------------
 // MARK: Execute Distributed Methods
-
 @available(SwiftStdlib 5.6, *)
 extension DistributedActorSystem {
 
@@ -154,21 +147,20 @@ extension DistributedActorSystem {
   /// is that thanks to this approach it can avoid any existential boxing, and can serve the most
   /// latency sensitive-use-cases.
   public func executeDistributedTarget<Act, ResultHandler>(
-      on actor: Act,
-      mangledTargetName: String,
-      invocation: Self.Invocation,
-      handler: ResultHandler
+    on actor: Act,
+    mangledTargetName: String,
+    invocation: inout Invocation,
+    handler: ResultHandler
   ) async throws where Act: DistributedActor,
                        Act.ID == ActorID,
                        ResultHandler: DistributedTargetInvocationResultHandler {
     // NOTE: this implementation is not the most efficient, nor final, version of this func
     // we end up demangling the name multiple times, perform more heap allocations than
     // we truly need to etc. We'll eventually move this implementation to a specialized one
     // avoiding these issues.
-
     guard mangledTargetName.count > 0 && mangledTargetName.first == "$" else {
       throw ExecuteDistributedTargetError(
-          message: "Illegal mangledTargetName detected, must start with '$'")
+        message: "Illegal mangledTargetName detected, must start with '$'")
     }
 
     // Get the expected parameter count of the func
@@ -179,36 +171,36 @@ extension DistributedActorSystem {
 
     guard paramCount >= 0 else {
       throw ExecuteDistributedTargetError(
-          message: """
-                   Failed to decode distributed invocation target expected parameter count,
-                   error code: \(paramCount)
-                   mangled name: \(mangledTargetName)
-                   """)
+        message: """
+                 Failed to decode distributed invocation target expected parameter count,
+                 error code: \(paramCount)
+                 mangled name: \(mangledTargetName)
+                 """)
     }
 
     // Prepare buffer for the parameter types to be decoded into:
     let paramTypesBuffer = UnsafeMutableRawBufferPointer
-        .allocate(byteCount: MemoryLayout<Any.Type>.size * Int(paramCount),
-                  alignment: MemoryLayout<Any.Type>.alignment)
+      .allocate(byteCount: MemoryLayout<Any.Type>.size * Int(paramCount),
+        alignment: MemoryLayout<Any.Type>.alignment)
     defer {
       paramTypesBuffer.deallocate()
     }
 
     // Demangle and write all parameter types into the prepared buffer
     let decodedNum = nameUTF8.withUnsafeBufferPointer { nameUTF8 in
       __getParameterTypeInfo(
-          nameUTF8.baseAddress!, UInt(nameUTF8.endIndex),
-          paramTypesBuffer.baseAddress!._rawValue, Int(paramCount))
+        nameUTF8.baseAddress!, UInt(nameUTF8.endIndex),
+        paramTypesBuffer.baseAddress!._rawValue, Int(paramCount))
     }
 
     // Fail if the decoded parameter types count seems off and fishy
     guard decodedNum == paramCount else {
       throw ExecuteDistributedTargetError(
-          message: """
-                   Failed to decode the expected number of params of distributed invocation target, error code: \(decodedNum)
-                   (decoded: \(decodedNum), expected params: \(paramCount)
-                   mangled name: \(mangledTargetName)
-                   """)
+        message: """
+                 Failed to decode the expected number of params of distributed invocation target, error code: \(decodedNum)
+                 (decoded: \(decodedNum), expected params: \(paramCount)
+                 mangled name: \(mangledTargetName)
+                 """)
     }
 
     // Copy the types from the buffer into a Swift Array
@@ -224,21 +216,23 @@ extension DistributedActorSystem {
     func allocateReturnTypeBuffer<R>(_: R.Type) -> UnsafeRawPointer? {
       return UnsafeRawPointer(UnsafeMutablePointer<R>.allocate(capacity: 1))
     }
-    guard let returnType: Any.Type = _getReturnTypeInfo(mangledMethodName: mangledTargetName) else {
+
+    guard let returnTypeFromTypeInfo: Any.Type = _getReturnTypeInfo(mangledMethodName: mangledTargetName) else {
       throw ExecuteDistributedTargetError(
-          message: "Failed to decode distributed target return type")
+        message: "Failed to decode distributed target return type")
     }
 
-    guard let resultBuffer = _openExistential(returnType, do: allocateReturnTypeBuffer) else {
+    guard let resultBuffer = _openExistential(returnTypeFromTypeInfo, do: allocateReturnTypeBuffer) else {
       throw ExecuteDistributedTargetError(
-          message: "Failed to allocate buffer for distributed target return type")
+        message: "Failed to allocate buffer for distributed target return type")
     }
 
     func destroyReturnTypeBuffer<R>(_: R.Type) {
       resultBuffer.assumingMemoryBound(to: R.self).deallocate()
     }
+
     defer {
-      _openExistential(returnType, do: destroyReturnTypeBuffer)
+      _openExistential(returnTypeFromTypeInfo, do: destroyReturnTypeBuffer)
     }
 
     // Prepare the buffer to decode the argument values into
@@ -249,12 +243,37 @@ extension DistributedActorSystem {
     }
 
     do {
+      // Decode the invocation and pack arguments into the h-buffer
+      // TODO(distributed): decode the generics info
+      var argumentDecoder = invocation.makeArgumentDecoder()
+      var paramIdx = 0
+      for unsafeRawArgPointer in hargs {
+        guard paramIdx < paramCount else {
+          throw ExecuteDistributedTargetError(
+            message: "Unexpected attempt to decode more parameters than expected: \(paramIdx + 1)")
+        }
+        let paramType = paramTypes[paramIdx]
+        paramIdx += 1
+
+        // FIXME(distributed): func doDecode<Arg: SerializationRequirement>(_: Arg.Type) throws {
+        // FIXME:     but how would we call this...?
+        // FIXME:     > type 'Arg' constrained to non-protocol, non-class type 'Self.Invocation.SerializationRequirement'
+        func doDecodeArgument<Arg>(_: Arg.Type) throws {
+          let unsafeArgPointer = unsafeRawArgPointer
+            .bindMemory(to: Arg.self, capacity: 1)
+          try argumentDecoder.decodeNext(Arg.self, into: unsafeArgPointer)
+        }
+        try _openExistential(paramType, do: doDecodeArgument)
+      }
+
+      let returnType = try invocation.decodeReturnType() ?? returnTypeFromTypeInfo
+      // let errorType = try invocation.decodeErrorType() // TODO: decide how to use?
       // Execute the target!
       try await _executeDistributedTarget(
-          on: actor,
-          mangledTargetName, UInt(mangledTargetName.count),
-          argumentBuffer: hargs.buffer._rawValue,
-          resultBuffer: resultBuffer._rawValue
+        on: actor,
+        mangledTargetName, UInt(mangledTargetName.count),
+        argumentBuffer: hargs.buffer._rawValue,
+        resultBuffer: resultBuffer._rawValue
       )
 
       func onReturn<R>(_ resultTy: R.Type) async throws {
@@ -279,15 +298,14 @@ func _executeDistributedTarget(
 
 // ==== ----------------------------------------------------------------------------------------------------------------
 // MARK: Support types
-
 /// A distributed 'target' can be a `distributed func` or `distributed` computed property.
 @available(SwiftStdlib 5.6, *)
 public struct RemoteCallTarget {
   let mangledName: String
 
   // Only intended to be created by the _Distributed library.
-  internal init(mangledName: String) {
-    self.mangledName = mangledName
+  public init(_mangledName: String) {
+    self.mangledName = _mangledName
   }
 
   // <module>.Base.hello(hi:)
@@ -333,38 +351,26 @@ public protocol DistributedTargetInvocation {
   associatedtype SerializationRequirement
 
   // === Sending / recording  -------------------------------------------------
-
-  // RECIPIENT:
-  // when we get the gen args; we can check them against the where clause
-  // and throw, rather than blow up... we have the ability to check the where clause
-  // the same code as implements `as?` takes generic signature and checks type arguments against that
-  //
-  // accessor is not generic, gets the subs passed in, forms SubsMap and uses to look up other things
-  // pass to call emission which would do the right thing...
-  /// Ad-hoc requirement
-  ///
   /// The arguments must be encoded order-preserving, and once `decodeGenericSubstitutions`
   /// is called, the substitutions must be returned in the same order in which they were recorded.
-  mutating func recordGenericSubstitution<T>(mangledType: T.Type) throws
+  mutating func recordGenericSubstitution<T>(_ type: T.Type) throws
 
 //  /// Ad-hoc requirement
 //  ///
 //  /// Record an argument of `Argument` type in this arguments storage.
 //  mutating func recordArgument<Argument: SerializationRequirement>(argument: Argument) throws
+  mutating func recordErrorType<E: Error>(_ type: E.Type) throws
 
 //  /// Ad-hoc requirement
 //  ///
-//  mutating func recordReturnType<R: SerializationRequirement>(mangledType: R.Type) throws
-
-  mutating func recordErrorType<E: Error>(mangledType: E.Type) throws
-
+//  /// Record the return type of the distributed method.
+//  mutating func recordReturnType<R: SerializationRequirement>(_ type: R.Type) throws
   mutating func doneRecording() throws
 
   // === Receiving / decoding -------------------------------------------------
-
   mutating func decodeGenericSubstitutions() throws -> [Any.Type]
 
-  mutating func argumentDecoder() -> Self.ArgumentDecoder
+  func makeArgumentDecoder() -> Self.ArgumentDecoder
 
   mutating func decodeReturnType() throws -> Any.Type?
 
@@ -375,20 +381,6 @@ public protocol DistributedTargetInvocation {
 ///
 /// It will be called exactly `N` times where `N` is the known number of arguments
 /// to the target invocation.
-///
-/// ## Ad-hoc protocol requirement
-///
-/// Adopters of this protocol must defined the following method:
-///
-/// ```
-/// mutating func decodeNext<Argument: SerializationRequirement>(
-///     into pointer: UnsafeMutablePointer<Argument>
-/// ) throws
-/// ```
-///
-/// which will be invoked with the specific `Argument` type that the target invocation is expecting.
-///
-/// This method is allowed to invoke
 @available(SwiftStdlib 5.6, *)
 public protocol DistributedTargetInvocationArgumentDecoder {
   associatedtype SerializationRequirement
@@ -409,11 +401,18 @@ public protocol DistributedTargetInvocationArgumentDecoder {
 //  mutating func decodeNext<Argument: SerializationRequirement>(
 //      into pointer: UnsafeMutablePointer<Argument> // pointer to our hbuffer
 //  ) throws
+  // FIXME(distributed): remove this since it must have the ': SerializationRequirement'
+  mutating func decodeNext<Argument>(
+    _ argumentType: Argument.Type,
+    into pointer: UnsafeMutablePointer<Argument> // pointer to our hbuffer
+  ) throws
 }
 
 @available(SwiftStdlib 5.6, *)
 public protocol DistributedTargetInvocationResultHandler {
-  // FIXME: these must be ad-hoc protocol requirements, because Res: SerializationRequirement !!!
+  associatedtype SerializationRequirement
+
+  // FIXME(distributed): these must be ad-hoc protocol requirements, because Res: SerializationRequirement !!!
   func onReturn<Res>(value: Res) async throws
   func onThrow<Err: Error>(error: Err) async throws
 }
@@ -443,7 +442,7 @@ public struct DistributedActorCodingError: DistributedActorSystemError {
   }
 
   public static func missingActorSystemUserInfo<Act>(_ actorType: Act.Type) -> Self
-      where Act: DistributedActor {
+    where Act: DistributedActor {
     .init(message: "Missing DistributedActorSystem userInfo while decoding")
   }
 }