swift-mirror/utils/swift-dev-utils/Sources/Utils/Misc.swift

//===--- Misc.swift -------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2024 - 2025 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
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

import Foundation

public extension Dictionary {
  @inline(__always)
  mutating func withValue<R>(
    for key: Key, default defaultValue: Value, body: (inout Value) throws -> R
  ) rethrows -> R {
    try body(&self[key, default: defaultValue])
  }

  mutating func insertValue(
    _ newValue: @autoclosure () -> Value, for key: Key
  ) -> Bool {
    if self[key] == nil {
      self[key] = newValue()
      return true
    }
    return false
  }
}
public extension Sequence {
  func sorted<T: Comparable>(by keyPath: KeyPath<Element, T>) -> [Element] {
    sorted(by: { $0[keyPath: keyPath] < $1[keyPath: keyPath] })
  }
}

public extension Collection where Element: Equatable {
  func commonPrefix(with other: some Collection<Element>) -> SubSequence {
    var (i, j) = (self.startIndex, other.startIndex)
    while i < self.endIndex, j < other.endIndex {
      guard self[i] == other[j] else { break }
      self.formIndex(after: &i)
      other.formIndex(after: &j)
    }
    return self[..<i]
  }
}

public extension String {
  init(utf8 buffer: UnsafeRawBufferPointer) {
    guard !buffer.isEmpty else {
      self = ""
      return
    }
    self = String(unsafeUninitializedCapacity: buffer.count,
                  initializingUTF8With: { dest in
      _ = dest.initialize(from: buffer)
      return buffer.count
    })
  }

  init(utf8 buffer: UnsafeBufferPointer<UInt8>) {
    self.init(utf8: UnsafeRawBufferPointer(buffer))
  }

  init(utf8 slice: Slice<UnsafeRawBufferPointer>) {
    self = String(utf8: .init(rebasing: slice))
  }

  init(utf8 buffer: ByteScanner.Bytes) {
    self = buffer.withUnsafeBytes(String.init(utf8:))
  }

  func scanningUTF8<R>(_ scan: (inout ByteScanner) throws -> R) rethrows -> R {
    var tmp = self
    return try tmp.withUTF8 { utf8 in
      var scanner = ByteScanner(utf8)
      return try scan(&scanner)
    }
  }

  func tryDropPrefix(_ prefix: String) -> String? {
    guard hasPrefix(prefix) else { return nil }
    return String(dropFirst(prefix.count))
  }

  func escaped(addQuotesIfNeeded: Bool) -> String {
    scanningUTF8 { scanner in
      var needsQuotes = false
      let result = scanner.consumeWhole { consumer in
        switch consumer.peek {
        case "\\", "\"":
          consumer.append("\\")
        case " ", "$": // $ is potentially a variable reference
          needsQuotes = true
        default:
          break
        }
      }
      let escaped = result.isUnchanged ? self : String(utf8: result)
      return addQuotesIfNeeded && needsQuotes ? "\"\(escaped)\"" : escaped
    }
  }

  var escaped: String {
    escaped(addQuotesIfNeeded: true)
  }

  init(_ str: StaticString) {
    self = str.withUTF8Buffer { utf8 in
      String(utf8: utf8)
    }
  }

  var isASCII: Bool {
    // Thanks, @testable interface!
    _classify()._isASCII
  }

  /// A more efficient version of replacingOccurrences(of:with:)/replacing(_:with:),
  /// since the former involves bridging, and the latter currently has no fast
  /// paths for strings.
  func replacing(_ other: String, with replacement: String) -> String {
    guard !other.isEmpty else {
      return self
    }
    guard isASCII else {
      // Not ASCII, fall back to slower method.
      return replacingOccurrences(of: other, with: replacement)
    }
    let otherUTF8 = other.utf8
    return scanningUTF8 { scanner in
      let bytes = scanner.consumeWhole { consumer in
        guard otherUTF8.count <= consumer.remaining.count else {
          // If there's no way we can eat the string, eat the remaining.
          consumer.eatRemaining()
          return
        }
        while consumer.trySkip(otherUTF8) {
          consumer.append(utf8: replacement)
        }
      }
      return bytes.isUnchanged ? self : String(utf8: bytes)
    }
  }
}

/// Pattern match by `is` property. E.g. `case \.isNewline: ...`
public func ~= <T>(keyPath: KeyPath<T, Bool>, subject: T) -> Bool {
  return subject[keyPath: keyPath]
}

public func ~= <T>(keyPath: KeyPath<T, Bool>, subject: T?) -> Bool {
  return subject?[keyPath: keyPath] == true
}

extension Sequence where Element: Sendable {
  public func parallelUnorderedMap<T: Sendable>(
    @_inheritActorContext _ transform: @escaping @Sendable (Element) async -> T
  ) async -> [T] {
    let worklist = TaskWorklist<T>()
    for elt in self {
      worklist.addTask {
        await transform(elt)
      }
    }
    return await worklist.results.all
  }
}

extension Task {
  /// Awaits the value of the result.
  ///
  /// If the current task is cancelled, this will cancel the subtask as well.
  public var valuePropagatingCancellation: Success {
    get async throws {
      try await withTaskCancellationHandler {
        return try await self.value
      } onCancel: {
        self.cancel()
      }
    }
  }
}

extension Task where Failure == Never {
  /// Awaits the value of the result.
  ///
  /// If the current task is cancelled, this will cancel the subtask as well.
  public var valuePropagatingCancellation: Success {
    get async {
      await withTaskCancellationHandler {
        return await self.value
      } onCancel: {
        self.cancel()
      }
    }
  }
}

extension AsyncSequence {
  public var all: [Element] {
    get async throws {
      var res: [Element] = []
      for try await elt in self {
        res.append(elt)
      }
      return res
    }
  }
}

extension AsyncSequence where Failure == Never {
  public var all: [Element] {
    get async {
      var res: [Element] = []
      for await elt in self {
        res.append(elt)
      }
      return res
    }
  }
}

extension RandomAccessCollection where Element: Equatable {
  public func findRepeatedSlice(minRepeats: Int = 1) -> SubSequence? {
    guard !isEmpty else { return nil }

    var numRepeats = 0
    var matchEnd = index(after: startIndex)
    var cursor = matchEnd
    while true {
      let match = startIndex ..< matchEnd
      let matchLen = distance(from: match.lowerBound, to: match.upperBound)
      guard let scanEnd = index(
        cursor, offsetBy: matchLen, limitedBy: endIndex
      ) else {
        return nil
      }
      let scanRange = cursor ..< scanEnd
      let scan = self[scanRange]
      guard self[match].elementsEqual(scan) else {
        numRepeats = 0
        formIndex(after: &matchEnd)
        cursor = matchEnd
        continue
      }
      numRepeats += 1
      if numRepeats >= minRepeats {
        return scan
      }
      cursor = scanEnd
    }
  }
}