PotentialStorage/src/kv/memory.ts

import { BaseKVDriver, type BatchOp, type ScanOptions } from './driver.js'
import { ChunkedSortedMap, type SortedMap } from './sorted.js'

export interface InMemoryKVDriverOptions {
  /**
   * Deep-copy values on write and read so callers can't mutate stored state through a shared
   * reference (matching serialising drivers). Values must be JSON-compatible. Default `true`.
   * Disable for speed when values are treated as immutable.
   */
  clone?: boolean
}

/**
 * In-memory driver over a sorted map. Implements every optional capability:
 * `setIfAbsent`, atomic `batch` and `deleteRange`.
 */
export class InMemoryKVDriver extends BaseKVDriver {
  readonly #data: SortedMap<unknown> = new ChunkedSortedMap()
  readonly #clone: (v: unknown) => unknown

  constructor(opts: InMemoryKVDriverOptions = {}) {
    super()
    this.#clone = opts.clone === false ? identity : cloneValue
  }

  get size(): number {
    return this.#data.size
  }

  async getMany(keys: string[]): Promise<unknown[]> {
    const out = new Array<unknown>(keys.length)
    for (let i = 0; i < keys.length; i++) out[i] = this.#clone(this.#data.get(keys[i]!))
    return out
  }

  async setMany(entries: [string, unknown][]): Promise<void> {
    const clone = this.#clone
    this.#data.setMany(entries.map(([k, v]) => [k, clone(v)]))
  }

  async deleteMany(keys: string[]): Promise<void> {
    for (const key of keys) this.#data.delete(key)
  }

  async setIfAbsent(entries: [string, unknown][]): Promise<boolean[]> {
    return entries.map(([key, value]) => {
      if (this.#data.has(key)) return false
      this.#data.set(key, this.#clone(value))
      return true
    })
  }

  async batch(ops: BatchOp[]): Promise<void> {
    // Applied synchronously within one microtask, so no other operation can observe a partial batch.
    for (const op of ops) {
      if (op.type === 'set') this.#data.set(op.key, this.#clone(op.value))
      else this.#data.delete(op.key)
    }
  }

  override async deleteRange(gte: string, lt: string): Promise<void> {
    const doomed = [...this.#data.keys({ gte, lt })]
    for (const key of doomed) this.#data.delete(key)
  }

  async *scan(opts: ScanOptions): AsyncGenerator<[string, unknown]> {
    const { gte, lt, reverse = false, values = false } = opts
    let remaining = opts.limit ?? Infinity
    if (remaining <= 0) return
    // The key iterator re-seeks by itself if the map changes while we're suspended.
    for (const key of this.#data.keys({ gte, lt, reverse })) {
      yield [key, values ? this.#clone(this.#data.get(key)) : undefined]
      if (--remaining <= 0) return
    }
  }
}

function identity(v: unknown): unknown {
  return v
}

/** Deep copy of a JSON-compatible value (much faster than `structuredClone`). */
function cloneValue(v: unknown): unknown {
  if (v === null || typeof v !== 'object') return v
  if (Array.isArray(v)) {
    const out = new Array<unknown>(v.length)
    for (let i = 0; i < v.length; i++) out[i] = cloneValue(v[i])
    return out
  }
  const out: Record<string, unknown> = {}
  for (const k in v) out[k] = cloneValue((v as Record<string, unknown>)[k])
  return out
}