PotentialStorage/src/document-storage/document-storage-in-memory.ts

import type { Doc } from '../utils/document-path.js'
import { DocumentPath } from '../utils/document-path.js'
import type { DocumentQuery, FindResult } from '../utils/document-query.js'
import { DocumentUpdate, type DocChange } from '../utils/document-update.js'
import { Json } from '../utils/json.js'
import { Keys } from '../utils/keys.js'
import type { AnySpec, Spec } from '../utils/schema.js'
import { BaseDocumentStorage, type IndexDeclarations } from './base-document-storage.js'

export interface DocumentStorageInMemoryOptions {
  /** Fields with an equality index, per collection pattern. */
  indexes?: IndexDeclarations
}

/** Field → JSON-encoded value → child keys. */
type CollectionIndex = Map<string, Map<string, Set<string>>>

/**
 * In-memory document storage. Declared index fields get an equality index per collection, used
 * when a filter has `{ field: value }` or `{ field: { $eq: value } }`; everything else is a scan.
 */
export class DocumentStorageInMemory<S extends Spec = AnySpec> extends BaseDocumentStorage<S> {
  private readonly collections = new Map<string, Map<string, Doc>>()
  private readonly indexes = new Map<string, CollectionIndex>()
  private readonly declarations: [pattern: string[], fields: string[]][]

  constructor(options: DocumentStorageInMemoryOptions = {}) {
    super()
    this.declarations = Object.entries(options.indexes ?? {}).map(([pattern, fields]) => [pattern === '' ? [] : pattern.split('/'), fields])
  }

  async readMany(paths: (readonly string[])[]): Promise<(Doc | undefined)[]> {
    return paths.map((path) => Json.clone(this.collection(path.slice(0, -1)).get(path.at(-1)!)))
  }

  async mutate(path: readonly string[], change: (current: Doc | undefined) => DocChange | null): Promise<Doc | undefined> {
    const parent = path.slice(0, -1)
    const key = path.at(-1)!
    const docs = this.collection(parent)
    const current = docs.get(key)
    const resolved = change(Json.clone(current))
    if (!resolved) return Json.clone(current)
    const next = DocumentUpdate.apply(current, resolved)
    if (next) docs.set(key, next)
    else docs.delete(key)
    this.reindex(parent, key, current, next)
    return Json.clone(next)
  }

  async *query(parent: readonly string[], query: DocumentQuery): AsyncGenerator<FindResult> {
    const docs = this.collection(parent)
    const keys = [...(this.lookup(parent, query) ?? docs.keys())].sort()
    yield* query.run(
      (async function* () {
        for (const key of keys) {
          const doc = docs.get(key)
          if (doc) yield { key, data: Json.clone(doc) }
        }
      })(),
    )
  }

  private collection(parent: readonly string[]): Map<string, Doc> {
    const id = Keys.join(parent)
    let docs = this.collections.get(id)
    if (!docs) this.collections.set(id, (docs = new Map()))
    return docs
  }

  /** Candidate keys from an equality index, if the filter allows one. */
  private lookup(parent: readonly string[], query: DocumentQuery): Iterable<string> | undefined {
    const index = this.indexes.get(Keys.join(parent))
    for (const field of this.fieldsFor(parent)) {
      let value = query.filter[field]
      if (Json.isObject(value) && Object.keys(value).length === 1 && '$eq' in value) value = value.$eq
      if (value === undefined || typeof value === 'object') continue
      return index?.get(field)?.get(JSON.stringify(value)) ?? []
    }
    return undefined
  }

  private reindex(parent: readonly string[], key: string, before: Doc | undefined, after: Doc | undefined): void {
    const fields = this.fieldsFor(parent)
    if (fields.length === 0) return
    const id = Keys.join(parent)
    let index = this.indexes.get(id)
    if (!index) this.indexes.set(id, (index = new Map()))
    for (const field of fields) {
      let byValue = index.get(field)
      if (!byValue) index.set(field, (byValue = new Map()))
      for (const value of DocumentStorageInMemory.indexValues(before, field)) byValue.get(value)?.delete(key)
      for (const value of DocumentStorageInMemory.indexValues(after, field)) {
        let keys = byValue.get(value)
        if (!keys) byValue.set(value, (keys = new Set()))
        keys.add(key)
      }
    }
  }

  private fieldsFor(parent: readonly string[]): string[] {
    return this.declarations
      .filter(([pattern]) => pattern.length === parent.length && pattern.every((s, i) => s === '*' || s === parent[i]))
      .flatMap(([, fields]) => fields)
  }

  /** Index keys for a field: its value, plus each element of an array (missing indexes as `null`). */
  private static indexValues(doc: Doc | undefined, field: string): string[] {
    if (!doc) return []
    const value = DocumentPath.get(doc, field) ?? null
    return Array.isArray(value) ? value.map((v) => JSON.stringify(v)) : [JSON.stringify(value)]
  }
}