ARK-OS/arkos/arkrt/swift-argument-parser/ArgumentParser/Completions/CompletionsGenerator.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift Argument Parser open source project
//
// Copyright (c) 2020 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
//
//===----------------------------------------------------------------------===//

#if compiler(>=6.0)
// import ArgumentParserToolInfo
#else
// import ArgumentParserToolInfo
#endif

/// A shell for which the parser can generate a completion script.
public struct CompletionShell: RawRepresentable, Hashable, CaseIterable,
  Sendable
{
  public var rawValue: String

  /// Creates a new instance from the given string.
  public init?(rawValue: String) {
    switch rawValue {
    case "zsh", "bash", "fish":
      self.rawValue = rawValue
    default:
      return nil
    }
  }

  /// An instance representing `zsh`.
  public static var zsh: CompletionShell {
    // swift-format-ignore: NeverForceUnwrap
    // Statically known valid raw value.
    CompletionShell(rawValue: "zsh")!
  }

  /// An instance representing `bash`.
  public static var bash: CompletionShell {
    // swift-format-ignore: NeverForceUnwrap
    // Statically known valid raw value.
    CompletionShell(rawValue: "bash")!
  }

  /// An instance representing `fish`.
  public static var fish: CompletionShell {
    // swift-format-ignore: NeverForceUnwrap
    // Statically known valid raw value.
    CompletionShell(rawValue: "fish")!
  }

  /// Returns an instance representing the current shell, if recognized.
  public static func autodetected() -> CompletionShell? {
    Platform.shellName.flatMap(CompletionShell.init(rawValue:))
  }

  /// An array of all supported shells for completion scripts.
  public static var allCases: [CompletionShell] {
    [.zsh, .bash, .fish]
  }

  static let _requesting = ArgParserMutex<CompletionShell?>(nil)

  /// The shell for which completions will be or are being requested.
  ///
  /// `CompletionShell.requesting` is non-`nil` only while generating a shell
  /// completion script or while a Swift custom completion function is executing
  /// to offer completions for a word from a command line (that is, while
  /// `customCompletion` from `@Option(completion: .custom(customCompletion))`
  /// executes).
  public static var requesting: CompletionShell? {
    Self._requesting.withLock { $0 }
  }

  static let _requestingVersion = ArgParserMutex<String?>(nil)

  /// The shell version for which completions will be or are being requested.
  ///
  /// `CompletionShell.requestingVersion` is non-`nil` only while generating a
  /// shell completion script or while a Swift custom completion function is
  /// running (that is, while `customCompletion` from
  /// `@Option(completion: .custom(customCompletion))` executes).
  public static var requestingVersion: String? {
    Self._requestingVersion.withLock { $0 }
  }

  func format(completions: [String]) -> String {
    var completions = completions
    if self == .zsh {
      // This pseudo-completion is removed by the zsh completion script.
      // It allows trailing empty string completions to work in zsh.
      // zsh completion scripts generated by older SAP versions ignore spaces.
      completions.append(" ")
    }
    return completions.joined(separator: "\n")
  }
}

struct CompletionsGenerator {
  var shell: CompletionShell
  var command: ParsableCommand.Type

  init(command: ParsableCommand.Type, shell: CompletionShell?) throws {
    guard let _shell = shell ?? .autodetected() else {
      throw ParserError.unsupportedShell()
    }

    self.shell = _shell
    self.command = command
  }

  init(command: ParsableCommand.Type, shellName: String?) throws {
    if let shellName = shellName {
      guard let shell = CompletionShell(rawValue: shellName) else {
        throw ParserError.unsupportedShell(shellName)
      }
      try self.init(command: command, shell: shell)
    } else {
      try self.init(command: command, shell: nil)
    }
  }

  /// Generates a shell completion script for this generator's shell and command.
  func generateCompletionScript() -> String {
    CompletionShell._requesting.withLock { $0 = shell }
    switch shell {
    case .zsh:
      return ToolInfoV0(commandStack: [command]).zshCompletionScript
    case .bash:
      return ToolInfoV0(commandStack: [command]).bashCompletionScript
    case .fish:
      return ToolInfoV0(commandStack: [command]).fishCompletionScript
    default:
      fatalError("Invalid CompletionShell: \(shell)")
    }
  }
}

extension String {
  func shellEscapeForSingleQuotedString(iterationCount: UInt64 = 1) -> Self {
    iterationCount == 0
      ? self
      : self
        .replacing("'", with: "'\\''")
        .shellEscapeForSingleQuotedString(iterationCount: iterationCount - 1)
  }

  func shellEscapeForVariableName() -> Self {
    self.replacing("-", with: "_")
  }

  func replacing(_ old: Self, with new: Self) -> Self {
    guard !old.isEmpty else { return self }

    var result = ""
    var startIndex = self.startIndex

    // Look for occurrences of the old string.
    while let matchRange = self.firstMatch(of: old, at: startIndex) {
      // Add the substring before the match.
      result.append(contentsOf: self[startIndex..<matchRange.start])

      // Add the replacement string.
      result.append(contentsOf: new)

      // Move past the matched portion.
      startIndex = matchRange.end
    }

    // No more matches found, add the rest of the string.
    result.append(contentsOf: self[startIndex..<self.endIndex])

    return result
  }

  func firstMatch(
    of match: Self,
    at startIndex: Self.Index
  ) -> (start: Self.Index, end: Self.Index)? {
    guard !match.isEmpty else { return nil }
    guard match.count <= self.count else { return nil }

    var startIndex = startIndex
    while startIndex < self.endIndex {
      // Check if theres a match.
      if let endIndex = self.matches(match, at: startIndex) {
        // Return the match.
        return (startIndex, endIndex)
      }

      // Move to the next of index.
      self.formIndex(after: &startIndex)
    }

    return nil
  }

  func matches(
    _ match: Self,
    at startIndex: Self.Index
  ) -> Self.Index? {
    var selfIndex = startIndex
    var matchIndex = match.startIndex

    while true {
      // Only continue checking if there is more match to check
      guard matchIndex < match.endIndex else { return selfIndex }

      // Exit early if there is no more "self" to check.
      guard selfIndex < self.endIndex else { return nil }

      // Check match and self are the the same.
      guard self[selfIndex] == match[matchIndex] else { return nil }

      // Move to the next pair of indices.
      self.formIndex(after: &selfIndex)
      match.formIndex(after: &matchIndex)
    }
  }
}

extension CommandInfoV0 {
  var commandContext: [String] {
    (superCommands ?? []) + [commandName]
  }

  var initialCommand: String {
    superCommands?.first ?? commandName
  }

  var positionalArguments: [ArgumentInfoV0] {
    (arguments ?? []).filter { $0.kind == .positional }
  }

  var completionFunctionName: String {
    "_" + commandContext.joined(separator: "_")
  }

  var completionFunctionPrefix: String {
    "__\(initialCommand)"
  }
}

extension ArgumentInfoV0 {
  /// Returns a string with the arguments for the callback to generate custom
  /// completions for this argument.
  func commonCustomCompletionCall(command: CommandInfoV0) -> String {
    let subcommandNames =
      command.commandContext.dropFirst().map { "\($0) " }.joined()

    let argumentName: String
    switch kind {
    case .positional:
      if let index = command.positionalArguments.firstIndex(of: self) {
        argumentName = "positional@\(index)"
      } else {
        argumentName = "---"
      }
    default:
      argumentName = preferredName?.commonCompletionSynopsisString() ?? "---"
    }
    return "---completion \(subcommandNames)-- \(argumentName)"
  }
}

extension ArgumentInfoV0.NameInfoV0 {
  func commonCompletionSynopsisString() -> String {
    switch kind {
    case .long:
      return "--\(name)"
    case .short, .longWithSingleDash:
      return "-\(name)"
    }
  }
}