PotentialStorage/src/document-storage/base-document-storage.ts

import type { Doc } from '../utils/document-path.js'
import { DocumentQuery, type Filter, type FindOptions, type FindResult } from '../utils/document-query.js'
import { DocumentUpdate, type DocChange, type UpdateSpec } from '../utils/document-update.js'
import { Keys } from '../utils/keys.js'
import type { AnySpec, ChildOf, KeyOf, Spec, ValueAt, ValueOf } from '../utils/schema.js'
import { Ulid } from '../utils/ulid.js'

/**
 * Fields to index, per collection path pattern (`*` matches one segment), e.g.
 * `{ users: ['email', 'age'], 'users/*\/posts': ['createdAt'] }`. Each implementation decides
 * how (or whether) to use them; a native database maps them onto its own indexes.
 */
export type IndexDeclarations = Record<string, string[]>

export interface UpdateOptions {
  /** Treat a missing document as `{}`. */
  upsert?: boolean
}

/** Document type of a node's children. */
export type ChildDoc<S> = ValueOf<ChildOf<S, KeyOf<S>>>

/**
 * Document storage. Every node is a document (its own) and a collection of child documents.
 * Implementations provide three primitives over document paths; the user-facing API is
 * defined here once.
 */
export abstract class BaseDocumentStorage<S extends Spec = AnySpec> {
  /** Unescaped path of this node (`[]` for the root, which has no own document). */
  readonly path: readonly string[] = []
  protected readonly ids = new Ulid()

  abstract readMany(paths: (readonly string[])[]): Promise<(Doc | undefined)[]>
  /**
   * Maps the current state of one document to a change (`null` = no change), applies it and
   * returns the new state. Implementations decide how this is made safe against concurrent writers.
   */
  abstract mutate(path: readonly string[], change: (current: Doc | undefined) => DocChange | null): Promise<Doc | undefined>
  /** Children of `parent` matching the query, filtered, sorted and paged. */
  abstract query(parent: readonly string[], query: DocumentQuery): AsyncIterable<FindResult>

  child<K extends KeyOf<S>>(key: K): BaseDocumentStorage<ChildOf<S, K>> {
    const node = Object.create(this, { path: { value: [...this.path, Keys.segment(key)], enumerable: true } }) as BaseDocumentStorage<ChildOf<S, K>>
    // Frozen: child nodes share the root's state through the prototype, so an implementation that
    // reassigns a field (`this.x = ...`) while running on a child would silently shadow it.
    // Keep mutable state in objects (Maps, a `state` object) instead; this makes mistakes throw.
    Object.freeze(node)
    return node
  }

  as<S2 extends Spec>(): BaseDocumentStorage<S2> {
    return this as unknown as BaseDocumentStorage<S2>
  }

  async get<K extends KeyOf<S> | undefined = undefined>(key?: K): Promise<ValueAt<S, K> | undefined> {
    const [doc] = await this.readMany([this.pathOf(arguments.length > 0, key)])
    return doc as ValueAt<S, K> | undefined
  }

  async getMany<K extends KeyOf<S>>(keys: readonly K[]): Promise<(ValueOf<ChildOf<S, K>> | undefined)[]> {
    return (await this.readMany(keys.map((k) => this.pathOf(true, k)))) as (ValueOf<ChildOf<S, K>> | undefined)[]
  }

  /** `set(doc)` replaces this document; `set(key, doc)` a child document. */
  async set<K extends KeyOf<S>>(...args: [doc: ValueOf<S>] | [key: K, doc: ValueOf<ChildOf<S, K>>]): Promise<void> {
    const doc = args[args.length - 1]
    if (typeof doc !== 'object' || doc === null || Array.isArray(doc)) throw new TypeError('A document must be a plain object')
    await this.mutate(this.pathOf(args.length > 1, args[0]), () => ({ op: 'set', doc: doc as Doc }))
  }

  /** Adds a child document under a new time-sortable id and returns the id. */
  async insert(doc: ChildDoc<S>): Promise<string> {
    const id = this.ids.generate()
    await this.set(id as KeyOf<S>, doc as never)
    return id
  }

  /**
   * `update(spec, options?)` this document; `update(key, spec, options?)` a child.
   * Returns the new state, or `undefined` if the document doesn't exist (unless `upsert`).
   */
  async update<K extends KeyOf<S>>(
    ...args: [spec: UpdateSpec, options?: UpdateOptions] | [key: K, spec: UpdateSpec, options?: UpdateOptions]
  ): Promise<ValueAt<S, K> | undefined> {
    const hasKey = typeof args[0] !== 'object' || args[0] === null
    const spec = args[hasKey ? 1 : 0] as UpdateSpec
    const options = (args[hasKey ? 2 : 1] ?? {}) as UpdateOptions
    const next = await this.mutate(this.pathOf(hasKey, args[0]), (current) => (current || options.upsert ? DocumentUpdate.resolve(current ?? {}, spec) : null))
    return next as ValueAt<S, K> | undefined
  }

  /** Deletes this document (or a child). Returns whether one existed. */
  async delete(key?: KeyOf<S>): Promise<boolean> {
    let existed = false
    await this.mutate(this.pathOf(arguments.length > 0, key), (current) => ((existed = current !== undefined) ? { op: 'delete' } : null))
    return existed
  }

  /** Child documents matching `filter`. */
  find(filter: Filter = {}, options: FindOptions = {}): AsyncIterable<FindResult<ChildDoc<S>>> {
    return this.query(this.path, new DocumentQuery(filter, options)) as AsyncIterable<FindResult<ChildDoc<S>>>
  }

  async findOne(filter: Filter = {}, options: FindOptions = {}): Promise<FindResult<ChildDoc<S>> | undefined> {
    for await (const result of this.find(filter, { ...options, limit: 1 })) return result
    return undefined
  }

  /** Keys of child documents. */
  async *keys(): AsyncGenerator<string> {
    for await (const { key } of this.find()) yield key
  }

  /** Path of a child (`hasKey`) or this node; validates keys. */
  protected pathOf(hasKey: boolean, key: unknown): readonly string[] {
    if (hasKey) return [...this.path, Keys.segment(key)]
    if (this.path.length === 0) throw new TypeError('The root node has no own document')
    return this.path
  }
}