grpc-swift/Sources/GRPCCodeGen/Internal/StructuredSwift+Client.swift

/*
 * Copyright 2024, gRPC Authors All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

extension ClosureInvocationDescription {
  /// ```
  /// { response in
  ///   try response.message
  /// }
  /// ```
  static var defaultClientUnaryResponseHandler: Self {
    ClosureInvocationDescription(
      argumentNames: ["response"],
      body: [.expression(.try(.identifierPattern("response").dot("message")))]
    )
  }
}

extension FunctionSignatureDescription {
  /// ```
  /// func <Name><Result>(
  ///   request: GRPCCore.ClientRequest<Input>,
  ///   serializer: some GRPCCore.MessageSerializer<Input>,
  ///   deserializer: some GRPCCore.MessageDeserializer<Output>,
  ///   options: GRPCCore.CallOptions,
  ///   onResponse handleResponse: @Sendable @escaping (GRPCCore.ClientResponse<Input>) async throws -> Result
  /// ) async throws -> Result where Result: Sendable
  /// ```
  static func clientMethod(
    accessLevel: AccessModifier? = nil,
    name: String,
    input: String,
    output: String,
    streamingInput: Bool,
    streamingOutput: Bool,
    includeDefaults: Bool,
    includeSerializers: Bool
  ) -> Self {
    var signature = FunctionSignatureDescription(
      accessModifier: accessLevel,
      kind: .function(name: name, isStatic: false),
      generics: [.member("Result")],
      parameters: [],  // Populated below.
      keywords: [.async, .throws],
      returnType: .identifierPattern("Result"),
      whereClause: WhereClause(requirements: [.conformance("Result", "Sendable")])
    )

    signature.parameters.append(
      ParameterDescription(
        label: "request",
        type: .clientRequest(forType: input, streaming: streamingInput)
      )
    )

    if includeSerializers {
      signature.parameters.append(
        ParameterDescription(
          label: "serializer",
          // Type is optional, so be explicit about which 'some' to use
          type: ExistingTypeDescription.some(.serializer(forType: input))
        )
      )
      signature.parameters.append(
        ParameterDescription(
          label: "deserializer",
          // Type is optional, so be explicit about which 'some' to use
          type: ExistingTypeDescription.some(.deserializer(forType: output))
        )
      )
    }

    signature.parameters.append(
      ParameterDescription(
        label: "options",
        type: .callOptions,
        defaultValue: includeDefaults ? .memberAccess(.dot("defaults")) : nil
      )
    )

    signature.parameters.append(
      ParameterDescription(
        label: "onResponse",
        name: "handleResponse",
        type: .closure(
          ClosureSignatureDescription(
            parameters: [
              ParameterDescription(
                type: .clientResponse(forType: output, streaming: streamingOutput)
              )
            ],
            keywords: [.async, .throws],
            returnType: .identifierPattern("Result"),
            sendable: true,
            escaping: true
          )
        ),
        defaultValue: includeDefaults && !streamingOutput
          ? .closureInvocation(.defaultClientUnaryResponseHandler)
          : nil
      )
    )

    return signature
  }
}

extension FunctionDescription {
  /// ```
  /// func <Name><Result>(
  ///   request: GRPCCore.ClientRequest<Input>,
  ///   options: GRPCCore.CallOptions = .defaults,
  ///   onResponse handleResponse: @Sendable @escaping (GRPCCore.ClientResponse<Input>) async throws -> Result
  /// ) async throws -> Result where Result: Sendable {
  ///   try await self.<Name>(
  ///     request: request,
  ///     serializer: <Serializer>,
  ///     deserializer: <Deserializer>,
  ///     options: options
  ///     onResponse: handleResponse,
  ///   )
  /// }
  /// ```
  static func clientMethodWithDefaults(
    accessLevel: AccessModifier? = nil,
    name: String,
    input: String,
    output: String,
    streamingInput: Bool,
    streamingOutput: Bool,
    serializer: Expression,
    deserializer: Expression
  ) -> Self {
    FunctionDescription(
      signature: .clientMethod(
        accessLevel: accessLevel,
        name: name,
        input: input,
        output: output,
        streamingInput: streamingInput,
        streamingOutput: streamingOutput,
        includeDefaults: true,
        includeSerializers: false
      ),
      body: [
        .expression(
          .try(
            .await(
              .functionCall(
                calledExpression: .identifierPattern("self").dot(name),
                arguments: [
                  FunctionArgumentDescription(
                    label: "request",
                    expression: .identifierPattern("request")
                  ),
                  FunctionArgumentDescription(
                    label: "serializer",
                    expression: serializer
                  ),
                  FunctionArgumentDescription(
                    label: "deserializer",
                    expression: deserializer
                  ),
                  FunctionArgumentDescription(
                    label: "options",
                    expression: .identifierPattern("options")
                  ),
                  FunctionArgumentDescription(
                    label: "onResponse",
                    expression: .identifierPattern("handleResponse")
                  ),
                ]
              )
            )
          )
        )
      ]
    )
  }
}

extension ProtocolDescription {
  /// ```
  /// protocol <Name>: Sendable {
  ///   func foo<Result: Sendable>(
  ///     ...
  ///   ) async throws -> Result
  /// }
  /// ```
  static func clientProtocol(
    accessLevel: AccessModifier? = nil,
    name: String,
    methods: [MethodDescriptor]
  ) -> Self {
    ProtocolDescription(
      accessModifier: accessLevel,
      name: name,
      conformances: ["Sendable"],
      members: methods.map { method in
        .commentable(
          .preFormatted(method.documentation),
          .function(
            signature: .clientMethod(
              name: method.name.generatedLowerCase,
              input: method.inputType,
              output: method.outputType,
              streamingInput: method.isInputStreaming,
              streamingOutput: method.isOutputStreaming,
              includeDefaults: false,
              includeSerializers: true
            )
          )
        )
      }
    )
  }
}

extension ExtensionDescription {
  /// ```
  /// extension <Name> {
  ///   func foo<Result: Sendable>(
  ///     request: GRPCCore.ClientRequest<Input>,
  ///     options: GRPCCore.CallOptions = .defaults,
  ///     onResponse handleResponse: @Sendable @escaping (GRPCCore.ClientResponse<Input>) async throws -> Result
  ///   ) async throws -> Result where Result: Sendable {
  ///     // ...
  ///   }
  ///   // ...
  /// }
  /// ```
  static func clientMethodSignatureWithDefaults(
    accessLevel: AccessModifier? = nil,
    name: String,
    methods: [MethodDescriptor],
    serializer: (String) -> String,
    deserializer: (String) -> String
  ) -> Self {
    ExtensionDescription(
      onType: name,
      declarations: methods.map { method in
        .function(
          .clientMethodWithDefaults(
            accessLevel: accessLevel,
            name: method.name.generatedLowerCase,
            input: method.inputType,
            output: method.outputType,
            streamingInput: method.isInputStreaming,
            streamingOutput: method.isOutputStreaming,
            serializer: .identifierPattern(serializer(method.inputType)),
            deserializer: .identifierPattern(deserializer(method.outputType))
          )
        )
      }
    )
  }
}

extension FunctionSignatureDescription {
  /// ```
  /// func foo<Result>(
  ///   _ message: <Input>,
  ///   metadata: GRPCCore.Metadata = [:],
  ///   options: GRPCCore.CallOptions = .defaults,
  ///   onResponse handleResponse: @Sendable @escaping (GRPCCore.ClientResponse<Output>) async throws -> Result = { response in
  ///     try response.message
  ///   }
  /// ) async throws -> Result where Result: Sendable
  /// ```
  static func clientMethodExploded(
    accessLevel: AccessModifier? = nil,
    name: String,
    input: String,
    output: String,
    streamingInput: Bool,
    streamingOutput: Bool
  ) -> Self {
    var signature = FunctionSignatureDescription(
      accessModifier: accessLevel,
      kind: .function(name: name),
      generics: [.member("Result")],
      parameters: [],  // Populated below
      keywords: [.async, .throws],
      returnType: .identifierPattern("Result"),
      whereClause: WhereClause(requirements: [.conformance("Result", "Sendable")])
    )

    if !streamingInput {
      signature.parameters.append(
        ParameterDescription(label: "_", name: "message", type: .member(input))
      )
    }

    // metadata: GRPCCore.Metadata = [:]
    signature.parameters.append(
      ParameterDescription(
        label: "metadata",
        type: .metadata,
        defaultValue: .literal(.dictionary([]))
      )
    )

    // options: GRPCCore.CallOptions = .defaults
    signature.parameters.append(
      ParameterDescription(
        label: "options",
        type: .callOptions,
        defaultValue: .dot("defaults")
      )
    )

    if streamingInput {
      signature.parameters.append(
        ParameterDescription(
          label: "requestProducer",
          name: "producer",
          type: .closure(
            ClosureSignatureDescription(
              parameters: [ParameterDescription(type: .rpcWriter(forType: input))],
              keywords: [.async, .throws],
              returnType: .identifierPattern("Void"),
              sendable: true,
              escaping: true
            )
          )
        )
      )
    }

    signature.parameters.append(
      ParameterDescription(
        label: "onResponse",
        name: "handleResponse",
        type: .closure(
          ClosureSignatureDescription(
            parameters: [
              ParameterDescription(
                type: .clientResponse(forType: output, streaming: streamingOutput)
              )
            ],
            keywords: [.async, .throws],
            returnType: .identifierPattern("Result"),
            sendable: true,
            escaping: true
          )
        ),
        defaultValue: streamingOutput ? nil : .closureInvocation(.defaultClientUnaryResponseHandler)
      )
    )

    return signature
  }
}

extension [CodeBlock] {
  /// ```
  /// let request = GRPCCore.StreamingClientRequest<Input>(
  ///   metadata: metadata,
  ///   producer: producer
  /// )
  /// return try await self.foo(
  ///   request: request,
  ///   options: options,
  ///   onResponse: handleResponse
  /// )
  /// ```
  static func clientMethodExploded(
    name: String,
    input: String,
    streamingInput: Bool
  ) -> Self {
    func arguments(streaming: Bool) -> [FunctionArgumentDescription] {
      let metadata = FunctionArgumentDescription(
        label: "metadata",
        expression: .identifierPattern("metadata")
      )

      if streaming {
        return [
          metadata,
          FunctionArgumentDescription(
            label: "producer",
            expression: .identifierPattern("producer")
          ),
        ]
      } else {
        return [
          FunctionArgumentDescription(label: "message", expression: .identifierPattern("message")),
          metadata,
        ]
      }
    }

    return [
      CodeBlock(
        item: .declaration(
          .variable(
            kind: .let,
            left: .identifierPattern("request"),
            right: .functionCall(
              calledExpression: .identifierType(
                .clientRequest(forType: input, streaming: streamingInput)
              ),
              arguments: arguments(streaming: streamingInput)
            )
          )
        )
      ),
      CodeBlock(
        item: .expression(
          .return(
            .try(
              .await(
                .functionCall(
                  calledExpression: .identifierPattern("self").dot(name),
                  arguments: [
                    FunctionArgumentDescription(
                      label: "request",
                      expression: .identifierPattern("request")
                    ),
                    FunctionArgumentDescription(
                      label: "options",
                      expression: .identifierPattern("options")
                    ),
                    FunctionArgumentDescription(
                      label: "onResponse",
                      expression: .identifierPattern("handleResponse")
                    ),
                  ]
                )
              )
            )
          )
        )
      ),
    ]
  }
}

extension FunctionDescription {
  /// ```
  /// func foo<Result>(
  ///   _ message: <Input>,
  ///   metadata: GRPCCore.Metadata = [:],
  ///   options: GRPCCore.CallOptions = .defaults,
  ///   onResponse handleResponse: @Sendable @escaping (GRPCCore.ClientResponse<Output>) async throws -> Result = { response in
  ///     try response.message
  ///   }
  /// ) async throws -> Result where Result: Sendable {
  ///   // ...
  /// }
  /// ```
  static func clientMethodExploded(
    accessLevel: AccessModifier? = nil,
    name: String,
    input: String,
    output: String,
    streamingInput: Bool,
    streamingOutput: Bool
  ) -> Self {
    FunctionDescription(
      signature: .clientMethodExploded(
        accessLevel: accessLevel,
        name: name,
        input: input,
        output: output,
        streamingInput: streamingInput,
        streamingOutput: streamingOutput
      ),
      body: .clientMethodExploded(name: name, input: input, streamingInput: streamingInput)
    )
  }
}

extension ExtensionDescription {
  /// ```
  /// extension <Name> {
  ///   // (exploded client methods)
  /// }
  /// ```
  static func explodedClientMethods(
    accessLevel: AccessModifier? = nil,
    on extensionName: String,
    methods: [MethodDescriptor]
  ) -> ExtensionDescription {
    ExtensionDescription(
      onType: extensionName,
      declarations: methods.map { method in
        .commentable(
          .preFormatted(method.documentation),
          .function(
            .clientMethodExploded(
              accessLevel: accessLevel,
              name: method.name.generatedLowerCase,
              input: method.inputType,
              output: method.outputType,
              streamingInput: method.isInputStreaming,
              streamingOutput: method.isOutputStreaming
            )
          )
        )
      }
    )
  }
}

extension FunctionDescription {
  /// ```
  /// func <Name><Result>(
  ///   request: GRPCCore.ClientRequest<Input>,
  ///   serializer: some GRPCCore.MessageSerializer<Input>,
  ///   deserializer: some GRPCCore.MessageDeserializer<Output>,
  ///   options: GRPCCore.CallOptions = .default,
  ///   onResponse handleResponse: @Sendable @escaping (GRPCCore.ClientResponse<Output>) async throws -> Result
  /// ) async throws -> Result where Result: Sendable {
  ///   try await self.<Name>(...)
  /// }
  /// ```
  static func clientMethod(
    accessLevel: AccessModifier,
    name: String,
    input: String,
    output: String,
    serviceEnum: String,
    methodEnum: String,
    streamingInput: Bool,
    streamingOutput: Bool
  ) -> Self {
    let underlyingMethod: String
    switch (streamingInput, streamingOutput) {
    case (false, false):
      underlyingMethod = "unary"
    case (true, false):
      underlyingMethod = "clientStreaming"
    case (false, true):
      underlyingMethod = "serverStreaming"
    case (true, true):
      underlyingMethod = "bidirectionalStreaming"
    }

    return FunctionDescription(
      accessModifier: accessLevel,
      kind: .function(name: name),
      generics: [.member("Result")],
      parameters: [
        ParameterDescription(
          label: "request",
          type: .clientRequest(forType: input, streaming: streamingInput)
        ),
        ParameterDescription(
          label: "serializer",
          // Be explicit: 'type' is optional and '.some' resolves to Optional.some by default.
          type: ExistingTypeDescription.some(.serializer(forType: input))
        ),
        ParameterDescription(
          label: "deserializer",
          // Be explicit: 'type' is optional and '.some' resolves to Optional.some by default.
          type: ExistingTypeDescription.some(.deserializer(forType: output))
        ),
        ParameterDescription(
          label: "options",
          type: .callOptions,
          defaultValue: .dot("defaults")
        ),
        ParameterDescription(
          label: "onResponse",
          name: "handleResponse",
          type: .closure(
            ClosureSignatureDescription(
              parameters: [
                ParameterDescription(
                  type: .clientResponse(forType: output, streaming: streamingOutput)
                )
              ],
              keywords: [.async, .throws],
              returnType: .identifierPattern("Result"),
              sendable: true,
              escaping: true
            )
          ),
          defaultValue: streamingOutput
            ? nil
            : .closureInvocation(.defaultClientUnaryResponseHandler)
        ),
      ],
      keywords: [.async, .throws],
      returnType: .identifierPattern("Result"),
      whereClause: WhereClause(requirements: [.conformance("Result", "Sendable")]),
      body: [
        .try(
          .await(
            .functionCall(
              calledExpression: .identifierPattern("self").dot("client").dot(underlyingMethod),
              arguments: [
                FunctionArgumentDescription(
                  label: "request",
                  expression: .identifierPattern("request")
                ),
                FunctionArgumentDescription(
                  label: "descriptor",
                  expression: .identifierPattern(serviceEnum)
                    .dot("Method")
                    .dot(methodEnum)
                    .dot("descriptor")
                ),
                FunctionArgumentDescription(
                  label: "serializer",
                  expression: .identifierPattern("serializer")
                ),
                FunctionArgumentDescription(
                  label: "deserializer",
                  expression: .identifierPattern("deserializer")
                ),
                FunctionArgumentDescription(
                  label: "options",
                  expression: .identifierPattern("options")
                ),
                FunctionArgumentDescription(
                  label: "onResponse",
                  expression: .identifierPattern("handleResponse")
                ),
              ]
            )
          )
        )
      ]
    )
  }
}

extension StructDescription {
  /// ```
  /// struct <Name>: <ClientProtocol> {
  ///   private let client: GRPCCore.GRPCClient
  ///
  ///   init(wrapping client: GRPCCore.GRPCClient) {
  ///     self.client = client
  ///   }
  ///
  ///   // ...
  /// }
  /// ```
  static func client(
    accessLevel: AccessModifier,
    name: String,
    serviceEnum: String,
    clientProtocol: String,
    methods: [MethodDescriptor]
  ) -> Self {
    StructDescription(
      accessModifier: accessLevel,
      name: name,
      conformances: [clientProtocol],
      members: [
        .variable(accessModifier: .private, kind: .let, left: "client", type: .grpcClient),
        .function(
          accessModifier: accessLevel,
          kind: .initializer,
          parameters: [
            ParameterDescription(label: "wrapping", name: "client", type: .grpcClient)
          ],
          whereClause: nil,
          body: [
            .expression(
              .assignment(
                left: .identifierPattern("self").dot("client"),
                right: .identifierPattern("client")
              )
            )
          ]
        ),
      ]
        + methods.map { method in
          .commentable(
            .preFormatted(method.documentation),
            .function(
              .clientMethod(
                accessLevel: accessLevel,
                name: method.name.generatedLowerCase,
                input: method.inputType,
                output: method.outputType,
                serviceEnum: serviceEnum,
                methodEnum: method.name.generatedUpperCase,
                streamingInput: method.isInputStreaming,
                streamingOutput: method.isOutputStreaming
              )
            )
          )
        }
    )
  }
}