brainy/src/transaction/operations/IndexOperations.ts

/**
 * Index Operations with Rollback Support
 *
 * Provides transactional operations for all indexes:
 * - JsHnswVectorIndex (unified vector index)
 * - MetadataIndexManager (roaring bitmap filtering)
 * - GraphAdjacencyIndex (LSM-tree graph storage)
 *
 * Each operation can be executed and rolled back atomically.
 */

import type { JsHnswVectorIndex } from '../../hnsw/hnswIndex.js'
import type { MetadataIndexManager } from '../../utils/metadataIndex.js'
import type { GraphIndexProvider } from '../../plugin.js'
import type { GraphVerb } from '../../coreTypes.js'
import type { Operation, RollbackAction } from '../types.js'

/**
 * Add to HNSW index with rollback support
 *
 * Rollback strategy:
 * - Remove item from index

 */
export class AddToHNSWOperation implements Operation {
  readonly name = 'AddToHNSW'

  constructor(
    private readonly index: JsHnswVectorIndex,
    private readonly id: string,
    private readonly vector: number[]
  ) {}

  async execute(): Promise<RollbackAction> {
    // Check if item already exists (for rollback decision)
    const existed = await this.itemExists(this.id)

    // Add to index
    await this.index.addItem({ id: this.id, vector: this.vector })

    // Return rollback action
    return async () => {
      if (!existed) {
        // Remove newly added item
        await this.index.removeItem(this.id)
      }
      // If item existed before, we don't rollback (update is OK)
      // This prevents index corruption from removing pre-existing items
    }
  }

  /**
   * Check if item exists in index.
   *
   * `getItem` is an optional, feature-detected provider capability — see the
   * VectorIndexProvider docs; it is intentionally absent from the required
   * contract (Brainy's JS HNSW index omits it). When the capability is
   * missing the answer must be `false`, not `true`: treating unknowable
   * pre-existence as "existed" made every rollback skip removeItem, leaving
   * phantom entries in the index after a failed transaction. The safe default
   * is to remove what this operation added — update flows pair this op with a
   * RemoveFromHNSWOperation whose own rollback restores the prior vector, so
   * reverse-order rollback reconstructs the original state either way.
   */
  private async itemExists(id: string): Promise<boolean> {
    const index = this.index as JsHnswVectorIndex & {
      getItem?: (id: string) => Promise<unknown>
    }
    if (typeof index.getItem !== 'function') return false
    try {
      const item = await index.getItem(id)
      return item !== undefined && item !== null
    } catch {
      return false
    }
  }
}

/**
 * Remove from HNSW index with rollback support
 *
 * Rollback strategy:
 * - Re-add item to index with original vector
 *
 * Note: Requires storing the vector for rollback
 */
export class RemoveFromHNSWOperation implements Operation {
  readonly name = 'RemoveFromHNSW'

  constructor(
    private readonly index: JsHnswVectorIndex,
    private readonly id: string,
    private readonly vector: number[] // Required for rollback
  ) {}

  async execute(): Promise<RollbackAction> {
    // Remove from index
    await this.index.removeItem(this.id)

    // Return rollback action
    return async () => {
      // Re-add item with original vector
      await this.index.addItem({ id: this.id, vector: this.vector })
    }
  }
}

/**
 * Add to metadata index with rollback support
 *
 * Rollback strategy:
 * - Remove item from index
 */
export class AddToMetadataIndexOperation implements Operation {
  readonly name = 'AddToMetadataIndex'

  constructor(
    private readonly index: MetadataIndexManager,
    private readonly id: string,
    private readonly entity: any // Entity or metadata structure
  ) {}

  async execute(): Promise<RollbackAction> {
    // Add to metadata index (skipFlush=true for transaction atomicity)
    await this.index.addToIndex(this.id, this.entity, true)

    // Return rollback action
    return async () => {
      // Remove from metadata index
      await this.index.removeFromIndex(this.id, this.entity)
    }
  }
}

/**
 * Remove from metadata index with rollback support
 *
 * Rollback strategy:
 * - Re-add item to index with original metadata
 */
export class RemoveFromMetadataIndexOperation implements Operation {
  readonly name = 'RemoveFromMetadataIndex'

  constructor(
    private readonly index: MetadataIndexManager,
    private readonly id: string,
    private readonly entity: any // Required for rollback
  ) {}

  async execute(): Promise<RollbackAction> {
    // Remove from metadata index
    await this.index.removeFromIndex(this.id, this.entity)

    // Return rollback action
    return async () => {
      // Re-add with original metadata (skipFlush=true)
      await this.index.addToIndex(this.id, this.entity, true)
    }
  }
}

/**
 * Add verb to graph index with rollback support
 *
 * Rollback strategy:
 * - Remove verb from graph index
 *
 * 8.0 u64 contract: the coordinator resolves both endpoint ints via the
 * shared idMapper (`getOrAssign`) and passes them alongside the verb; the
 * provider returns the interned verb int, which is surfaced through the
 * optional `onVerbInt` callback so the coordinator can feed its warm cache.
 *
 * Generation: `generationFn` is resolved at execute time (not construction) so
 * the edge is stamped at the transaction's in-flight commit generation — which
 * the generation store only assigns once the batch begins executing. The same
 * generation is reused for the rollback removal, so an add and its undo
 * reference one watermark in a provider's per-generation edge chain.
 */
/**
 * A verb's interned endpoint ints — eager (already resolved), or a thunk
 * evaluated when the operation EXECUTES. The lazy form exists for
 * `transact()` forward references: a relate whose endpoint is added in the
 * SAME batch cannot resolve ints at plan time (the entity does not exist yet
 * — a strict native id mapper rightly refuses to assign, and even a permissive
 * one would leak the assignment if the batch is rejected at precommit).
 * Deferring to execute time resolves after the batch's add operations have
 * applied, mirroring how `generationFn` is already evaluated lazily.
 */
export type VerbEndpointInts =
  | { sourceInt: bigint; targetInt: bigint }
  | (() => { sourceInt: bigint; targetInt: bigint })

/** Resolve a {@link VerbEndpointInts} at execution time. */
function resolveEndpointInts(
  endpoints: VerbEndpointInts
): { sourceInt: bigint; targetInt: bigint } {
  return typeof endpoints === 'function' ? endpoints() : endpoints
}

export class AddToGraphIndexOperation implements Operation {
  readonly name = 'AddToGraphIndex'

  /**
   * @param index - The graph-index provider (JS baseline or native).
   * @param verb - The verb to index (`sourceInt`/`targetInt` mirrored on it).
   * @param endpointInts - The endpoints' interned ints — eager, or a thunk
   *   evaluated at execute time (REQUIRED for transact forward references;
   *   see {@link VerbEndpointInts}).
   * @param generationFn - Resolves the commit generation to stamp this edge at,
   *   evaluated when the operation executes (see class note).
   * @param onVerbInt - Optional hook invoked with the interned verb int
   *   returned by the provider (feeds the coordinator's verb-int warm cache).
   */
  constructor(
    private readonly index: GraphIndexProvider,
    private readonly verb: GraphVerb,
    private readonly endpointInts: VerbEndpointInts,
    private readonly generationFn: () => bigint,
    private readonly onVerbInt?: (verbInt: bigint) => void
  ) {}

  async execute(): Promise<RollbackAction> {
    // Stamp this edge at the in-flight commit generation; reuse it for the
    // rollback so add + undo reference the same watermark. Endpoint ints
    // resolve HERE — after any same-batch adds have applied.
    const generation = this.generationFn()
    const { sourceInt, targetInt } = resolveEndpointInts(this.endpointInts)
    const verbInt = await this.index.addVerb(this.verb, sourceInt, targetInt, generation)
    this.onVerbInt?.(verbInt)

    // Return rollback action
    return async () => {
      // Remove verb from graph index
      await this.index.removeVerb(this.verb.id, generation)
    }
  }
}

/**
 * Remove verb from graph index with rollback support
 *
 * Rollback strategy:
 * - Re-add verb to graph index
 *
 * 8.0 u64 contract: rollback re-adds through `addVerb(verb, sourceInt,
 * targetInt, generation)`, so the coordinator resolves the endpoint ints up
 * front (while the entity → int mappings are guaranteed to still exist). The
 * removal generation is resolved at execute time and reused for the rollback
 * re-add, so the round trip references one watermark.
 */
export class RemoveFromGraphIndexOperation implements Operation {
  readonly name = 'RemoveFromGraphIndex'

  /**
   * @param index - The graph-index provider (JS baseline or native).
   * @param verb - The verb being removed (required for rollback re-add).
   * @param endpointInts - The endpoints' interned ints for the rollback
   *   re-add — eager, or a thunk evaluated at execute time (required when the
   *   verb or its endpoints were created in the SAME transact batch; see
   *   {@link VerbEndpointInts}).
   * @param generationFn - Resolves the commit generation for this removal,
   *   evaluated when the operation executes.
   */
  constructor(
    private readonly index: GraphIndexProvider,
    private readonly verb: GraphVerb, // Required for rollback
    private readonly endpointInts: VerbEndpointInts,
    private readonly generationFn: () => bigint
  ) {}

  async execute(): Promise<RollbackAction> {
    // Resolve the removal generation once; reuse it for the rollback re-add.
    // Endpoint ints resolve HERE (after any same-batch adds applied) and are
    // captured for the rollback, whose re-add must use the same mappings.
    const generation = this.generationFn()
    const { sourceInt, targetInt } = resolveEndpointInts(this.endpointInts)
    await this.index.removeVerb(this.verb.id, generation)

    // Return rollback action
    return async () => {
      // Re-add verb with original data
      await this.index.addVerb(this.verb, sourceInt, targetInt, generation)
    }
  }
}

/**
 * Batch operation: Add multiple items to HNSW index
 *
 * Useful for bulk imports with transaction support.
 * Rolls back all items if any fail.
 */
export class BatchAddToHNSWOperation implements Operation {
  readonly name = 'BatchAddToHNSW'

  private operations: AddToHNSWOperation[]

  constructor(
    index: JsHnswVectorIndex,
    items: Array<{ id: string; vector: number[] }>
  ) {
    this.operations = items.map(
      item => new AddToHNSWOperation(index, item.id, item.vector)
    )
  }

  async execute(): Promise<RollbackAction> {
    const rollbackActions: RollbackAction[] = []

    // Execute all operations
    for (const op of this.operations) {
      const rollback = await op.execute()
      if (rollback) {
        rollbackActions.push(rollback)
      }
    }

    // Return combined rollback action
    return async () => {
      // Execute all rollbacks in reverse order
      for (let i = rollbackActions.length - 1; i >= 0; i--) {
        await rollbackActions[i]()
      }
    }
  }
}

/**
 * Batch operation: Add multiple entities to metadata index
 *
 * Useful for bulk imports with transaction support.
 */
export class BatchAddToMetadataIndexOperation implements Operation {
  readonly name = 'BatchAddToMetadataIndex'

  private operations: AddToMetadataIndexOperation[]

  constructor(
    index: MetadataIndexManager,
    items: Array<{ id: string; entity: any }>
  ) {
    this.operations = items.map(
      item => new AddToMetadataIndexOperation(index, item.id, item.entity)
    )
  }

  async execute(): Promise<RollbackAction> {
    const rollbackActions: RollbackAction[] = []

    // Execute all operations
    for (const op of this.operations) {
      const rollback = await op.execute()
      if (rollback) {
        rollbackActions.push(rollback)
      }
    }

    // Return combined rollback action
    return async () => {
      // Execute all rollbacks in reverse order
      for (let i = rollbackActions.length - 1; i >= 0; i--) {
        await rollbackActions[i]()
      }
    }
  }
}