@idooel/runtime-context/src/core/DataPool.ts

Runtime data pool with namespaces, stackable contexts, subscriptions and optional persistence. Vue adapter included.

import { Context } from './Context'
import {
  DEFAULT_NAMESPACE_PREFIX,
} from './types'
import type {
  DataEvent,
  DataKey,
  Namespace,
  SetOptions,
  Subscriber,
  SubscriptionFilter,
  Unsubscribe,
  SnapshotOptions,
} from './types'

type NsMap = Map<DataKey, unknown>

interface Subscription {
  filter?: SubscriptionFilter
  fn: (event: DataEvent<any>) => void
}

function matchesFilter(filter: SubscriptionFilter | undefined, ns: Namespace, key: DataKey): boolean {
  if (!filter) return true
  if (filter.ns && filter.ns !== ns) return false
  if (filter.key && filter.key !== key) return false
  return true
}

export class DataPool<TSchema = unknown> {
  private readonly globalStore: Map<Namespace, NsMap> = new Map()
  private readonly contextStack: Context[] = []
  private readonly subscribers: Set<Subscription> = new Set()

  constructor() {}

  /**
   * Add idooel prefix to namespace if not already present
   */
  private prefixNamespace(ns: Namespace): Namespace {
    if (!ns && ns !== '') return `${DEFAULT_NAMESPACE_PREFIX}.undefined`
    return ns.startsWith(`${DEFAULT_NAMESPACE_PREFIX}.`) ? ns : `${DEFAULT_NAMESPACE_PREFIX}.${ns}`
  }

  private currentContext(): Context | null {
    return this.contextStack.length ? this.contextStack[this.contextStack.length - 1] : null
  }

  private ensureNsMap(target: Map<Namespace, NsMap>, ns: Namespace): NsMap {
    let nsMap = target.get(ns)
    if (!nsMap) {
      nsMap = new Map()
      target.set(ns, nsMap)
    }
    return nsMap
  }

  private notify<T>(event: DataEvent<T>): void {
    for (const sub of this.subscribers) {
      if (matchesFilter(sub.filter, event.ns, event.key)) {
        try {
          sub.fn(event)
        } catch {
          // ignore subscriber error
        }
      }
    }
  }

  get<T = unknown>(ns: Namespace, key: DataKey, fallback?: T): T | undefined {
    const prefixedNs = this.prefixNamespace(ns)
    const actualKey = !key && key !== '' ? 'undefined' : key
    // read from contexts (top to bottom)
    for (let i = this.contextStack.length - 1; i >= 0; i--) {
      const ctx = this.contextStack[i]
      const nsMap = ctx.store.get(prefixedNs)
      if (nsMap && nsMap.has(actualKey)) return nsMap.get(actualKey) as T
    }
    // read from global
    const globalNs = this.globalStore.get(prefixedNs)
    if (globalNs && globalNs.has(actualKey)) return globalNs.get(actualKey) as T
    return fallback
  }

  set<T = unknown>(ns: Namespace, key: DataKey, value: T, options?: SetOptions): void {
    const prefixedNs = this.prefixNamespace(ns)
    const actualKey = !key && key !== '' ? 'undefined' : key
    const targetStore = (options?.target === 'global' || !this.currentContext())
      ? this.globalStore
      : this.currentContext()!.store
    const nsMap = this.ensureNsMap(targetStore, prefixedNs)
    nsMap.set(actualKey, value)
    this.notify<T>({
      ns: prefixedNs,
      key: actualKey,
      value,
      source: targetStore === this.globalStore ? 'global' : 'context',
      contextId: this.currentContext()?.id ?? null,
    })
  }

  delete(ns: Namespace, key: DataKey): void {
    const prefixedNs = this.prefixNamespace(ns)
    const actualKey = !key && key !== '' ? 'undefined' : key
    let deleted = false
    // delete from contexts (top to bottom, first hit)
    for (let i = this.contextStack.length - 1; i >= 0; i--) {
      const ctx = this.contextStack[i]
      const nsMap = ctx.store.get(prefixedNs)
      if (nsMap && nsMap.has(actualKey)) {
        nsMap.delete(actualKey)
        deleted = true
        break
      }
    }
    // else delete from global
    if (!deleted) {
      const globalNs = this.globalStore.get(prefixedNs)
      if (globalNs && globalNs.has(actualKey)) {
        globalNs.delete(actualKey)
        deleted = true
      }
    }
    if (deleted) {
      this.notify({
        ns: prefixedNs,
        key: actualKey,
        value: undefined,
        source: 'delete',
        contextId: this.currentContext()?.id ?? null,
      })
    }
  }

  clear(ns?: Namespace): void {
    const prefixedNs = ns ? this.prefixNamespace(ns) : undefined
    if (prefixedNs) {
      this.globalStore.delete(prefixedNs)
      for (const ctx of this.contextStack) {
        ctx.store.delete(prefixedNs)
      }
    } else {
      this.globalStore.clear()
      for (const ctx of this.contextStack) {
        ctx.store.clear()
      }
    }
    this.notify({
      ns: prefixedNs ?? '*',
      key: '*',
      value: undefined,
      source: 'clear',
      contextId: this.currentContext()?.id ?? null,
    })
  }

  subscribe<T = unknown>(filter: SubscriptionFilter | undefined, listener: Subscriber<T>): Unsubscribe {
    // 为过滤器添加前缀，用于内部匹配
    const internalFilter = filter ? {
      ns: filter.ns ? this.prefixNamespace(filter.ns) : undefined,
      key: filter.key
    } : undefined
    
    const sub: Subscription = { 
      filter: internalFilter,
      fn: (event) => {
        // 在事件中恢复原始命名空间
        const originalEvent = {
          ...event,
          ns: event.ns.replace(/^idooel\./, '')
        }
        listener(originalEvent)
      }
    }
    this.subscribers.add(sub)
    return () => {
      this.subscribers.delete(sub)
    }
  }

  createContext(init?: Record<string, Record<string, unknown>>): Context {
    const prefixedInit = init ? Object.fromEntries(
      Object.entries(init).map(([ns, data]) => [this.prefixNamespace(ns), data])
    ) : undefined
    return new Context(prefixedInit)
  }

  enter(ctx: Context): { token: symbol } {
    this.contextStack.push(ctx)
    return { token: Symbol(ctx.id) }
  }

  exit(token: { token: symbol }): void {
    // pop until matched or empty
    while (this.contextStack.length) {
      const ctx = this.contextStack[this.contextStack.length - 1]
      const expected = Symbol.for(ctx.id)
      // cannot reconstruct previous symbol; do a simple pop
      this.contextStack.pop()
      if (token) break
    }
  }

  runInContext<T>(initOrCtx: Context | Record<string, Record<string, unknown>>, fn: () => T): T {
    const ctx = initOrCtx instanceof Context ? initOrCtx : this.createContext(initOrCtx)
    const { token } = this.enter(ctx)
    try {
      return fn()
    } finally {
      this.exit({ token })
    }
  }

  withContext<TArgs extends any[], TReturn>(initOrCtx: Context | Record<string, Record<string, unknown>>) {
    return (fn: (...args: TArgs) => TReturn) => {
      return (...args: TArgs) => this.runInContext(initOrCtx, () => fn(...args))
    }
  }

  wrapHandler<TArgs extends any[], TReturn>(handler: (...args: TArgs) => TReturn) {
    const captured = this.currentContext()
    if (!captured) return handler
    return (...args: TArgs) => this.runInContext(captured, () => handler(...args))
  }

  snapshot(options?: SnapshotOptions): Record<string, any> {
    const merged = options?.merged !== false
    if (merged) {
      const out: Record<string, any> = {}
      // start from global then overlay contexts bottom->top then top->bottom? We'll overlay from global then each context in order they are pushed.
      for (const [ns, nsMap] of this.globalStore) {
        out[ns] = Object.fromEntries(nsMap.entries())
      }
      for (const ctx of this.contextStack) {
        for (const [ns, nsMap] of ctx.store) {
          out[ns] = { ...(out[ns] ?? {}), ...Object.fromEntries(nsMap.entries()) }
        }
      }
      return out
    } else {
      return {
        global: Object.fromEntries(
          Array.from(this.globalStore.entries()).map(([ns, nsMap]) => [ns, Object.fromEntries(nsMap.entries())])
        ),
        contexts: this.contextStack.map(ctx => ({
          id: ctx.id,
          data: Object.fromEntries(
            Array.from(ctx.store.entries()).map(([ns, nsMap]) => [ns, Object.fromEntries(nsMap.entries())])
          ),
        })),
      }
    }
  }

}

export function createDataPool<TSchema = unknown>() {
  return new DataPool<TSchema>()
}