MemoryLeakMarker.swift

import struct Foundation.UUID
import os
import libkern

public enum MemoryLeaks {

  public static func listActiveMarkers() -> Array<Marker> {
    let markersRegistryCopy: Dictionary<AnyHashable, WeakMarker> = markersRegistryLock
      .withLock { Self.markersRegistry }
    return markersRegistryCopy.values.compactMap(\.marker)
  }
}

extension MemoryLeaks {

  private static let logger: os.Logger = .init(
    subsystem: "Memory",
    category: "LeakCheck"
  )

  fileprivate static let markersRegistryLock: SpinLock = .init()
  fileprivate static var markersRegistry: Dictionary<AnyHashable, WeakMarker> = .init()

  fileprivate static func register(
    marker: Marker
  ) {
    let id: AnyHashable = marker.id
    let weakMarker: WeakMarker = .init(marker: marker)
    markersRegistryLock
      .withLock {
        Self.markersRegistry[id] = weakMarker
      }
    Self.logger.debug("NewMarker:\(id, privacy: .public)")
  }

  fileprivate static func unregister(
    markerWithID markerID: AnyHashable
  ) {
    markersRegistryLock
      .withLock {
        Self.markersRegistry[markerID] = .none
      }
    Self.logger.debug("RemovedMarker:\(markerID, privacy: .public)")
  }
}

extension MemoryLeaks {

  fileprivate final class WeakMarker {

    fileprivate weak var marker: Marker?

    fileprivate init(
      marker: Marker
    ) {
      self.marker = marker
    }
  }

  fileprivate final class SpinLock {

    private let ptr: UnsafeMutablePointer<atomic_flag>

    fileprivate init() {
      self.ptr = .allocate(capacity: 1)
      self.ptr.pointee = atomic_flag()
    }

    deinit {
      self.ptr.deinitialize(count: 1)
      self.ptr.deallocate()
    }

    fileprivate func lock() {
      while !atomic_flag_test_and_set(self.ptr) {
        // keep spinning
      }
    }

    fileprivate func unlock() {
      atomic_flag_clear(self.ptr)
    }

    fileprivate func withLock<R>(
      _ execute: () throws -> R
    ) rethrows -> R {
      self.lock()
      defer { self.unlock() }
      return try execute()
    }
  }
}

extension MemoryLeaks {

  public final class Marker {

    fileprivate let id: String
    private let customID: String
    private let uuid: String
    private let timestamp: Int
    private let file: StaticString
    private let line: UInt

    fileprivate init(
      customID: String,
      file: StaticString,
      line: UInt
    ) {
      let uuidString: String = UUID().uuidString
      let timestamp: Int = time(nil)
      self.id = "\(file):\(line)[\(uuidString)]<\(timestamp)>-\(customID)"
      self.customID = customID
      self.uuid = uuidString
      self.timestamp = timestamp
      self.file = file
      self.line = line
      MemoryLeaks
        .register(marker: self)
    }

    deinit {
      MemoryLeaks
        .unregister(markerWithID: self.id)
    }

    fileprivate func wrappedReturn<R>(
      _ value: @autoclosure () throws -> R
    ) rethrows -> R {
      try value()
    }
  }
}

extension MemoryLeaks.Marker: CustomStringConvertible {

  public var description: String {
    "MemoryLeaksMarker:\(self.id)"
  }
}

extension MemoryLeaks.Marker: CustomDebugStringConvertible {

  public var debugDescription: String {
    "MemoryLeaksMarker:\(self.id)"
  }
}

extension MemoryLeaks.Marker: CustomReflectable {

  public var customMirror: Mirror {
    .init(
      self,
      children: ["id": self.id],
      displayStyle: .tuple,
      ancestorRepresentation: .suppressed
    )
  }
}

extension MemoryLeaks {

  public static func mark<R>(
    _ function: @escaping () -> R,
    customID: String = "",
    file: StaticString = #fileID,
    line: UInt = #line
  ) -> () -> R {
#if DEBUG
    let marker: MemoryLeaks.Marker = .init(
      customID: customID,
      file: file,
      line: line
    )
    return {
      marker.wrappedReturn(function())
    }
#else
    return function
#endif
  }

  public static func markThrowing<R>(
    _ function: @escaping () throws -> R,
    customID: String = "",
    file: StaticString = #fileID,
    line: UInt = #line
  ) -> () throws -> R {
#if DEBUG
    let marker: MemoryLeaks.Marker = .init(
      customID: customID,
      file: file,
      line: line
    )
    return {
      try marker.wrappedReturn(try function())
    }
#else
    return function
#endif
  }

  public static func mark<R, A1>(
    _ function: @escaping (A1) -> R,
    customID: String = "",
    file: StaticString = #fileID,
    line: UInt = #line
  ) -> (A1) -> R {
#if DEBUG
    let marker: MemoryLeaks.Marker = .init(
      customID: customID,
      file: file,
      line: line
    )
    return { a1 in
      marker.wrappedReturn(function(a1))
    }
#else
    return function
#endif
  }

  public static func mark<R, A1, A2>(
    _ function: @escaping (A1, A2) -> R,
    customID: String = "",
    file: StaticString = #fileID,
    line: UInt = #line
  ) -> (A1, A2) -> R {
#if DEBUG
    let marker: MemoryLeaks.Marker = .init(
      customID: customID,
      file: file,
      line: line
    )
    return { a1, a2 in
      marker.wrappedReturn(function(a1, a2))
    }
#else
    return function
#endif
  }
}