Imhotep/packages/imhotep-solver/src/bindings.ts

/**
 * Variable binding environment for first-order logic evaluation.
 *
 * Bindings map variable names to subject ids in the geometry world.
 * Tuple bindings support multi-variable quantification by maintaining
 * a flat array of subject ids indexed by variable position.
 *
 * Invariant: the environment is immutable; every extension produces a
 * new BindingEnv sharing the parent map via prototype chain or copy.
 */

// ---------------------------------------------------------------------------
// Binding Entry
// ---------------------------------------------------------------------------

export interface BindingEntry {
  variableName: string;
  subjectId: number;
  domainId?: string;
}

// ---------------------------------------------------------------------------
// Binding Environment
// ---------------------------------------------------------------------------

export class BindingEnv {
  private readonly parent: BindingEnv | null;
  private readonly entries: Map<string, number>;
  private _size: number;

  constructor(parent: BindingEnv | null = null, newEntries: BindingEntry[] = []) {
    this.parent = parent;
    this.entries = new Map<string, number>();
    for (const entry of newEntries) {
      this.entries.set(entry.variableName, entry.subjectId);
    }
    this._size = parent ? parent._size + this.entries.size : this.entries.size;
  }

  /**
   * Look up the subject id bound to a variable name.
   * Returns undefined if the variable is not in scope.
   */
  lookup(variableName: string): number | undefined {
    let env: BindingEnv | null = this;
    while (env) {
      const value = env.entries.get(variableName);
      if (value !== undefined) {
        return value;
      }
      env = env.parent;
    }
    return undefined;
  }

  /**
   * Extend this environment with a single binding.
   */
  bind(variableName: string, subjectId: number): BindingEnv {
    return new BindingEnv(this, [{ variableName, subjectId }]);
  }

  /**
   * Extend this environment with multiple bindings (tuple binding).
   */
  bindTuple(newEntries: BindingEntry[]): BindingEnv {
    return new BindingEnv(this, newEntries);
  }

  /**
   * Return all bound variable names in this environment.
   */
  variables(): string[] {
    const vars = new Set<string>();
    let env: BindingEnv | null = this;
    while (env) {
      for (const name of env.entries.keys()) {
        vars.add(name);
      }
      env = env.parent;
    }
    return Array.from(vars);
  }

  /**
   * Return the number of bound variables.
   */
  size(): number {
    return this._size;
  }

  /**
   * Produce a plain object representation for diagnostics.
   */
  toObject(): Record<string, number> {
    const result: Record<string, number> = {};
    let env: BindingEnv | null = this;
    while (env) {
      for (const [name, id] of env.entries) {
        if (!(name in result)) {
          result[name] = id;
        }
      }
      env = env.parent;
    }
    return result;
  }
}

// ---------------------------------------------------------------------------
// Tuple Binding Logic
// ---------------------------------------------------------------------------

export interface TupleBindingSpec {
  variableNames: string[];
  domainSubjectIds: Uint32Array;
}

/**
 * Generate all tuple combinations from a set of domain bindings.
 *
 * This is data-oriented: it iterates flat domain arrays and yields
 * pre-allocated tuple arrays to avoid per-tuple allocation.
 *
 * @param specs - Array of binding specs, one per quantified variable.
 * @param onTuple - Callback invoked for each tuple combination.
 *                  Return false to short-circuit enumeration.
 */
export function enumerateTuples(
  specs: TupleBindingSpec[],
  onTuple: (tuple: number[]) => boolean,
): void {
  if (specs.length === 0) {
    onTuple([]);
    return;
  }

  const tuple: number[] = new Array(specs.length);
  const indices: number[] = new Array(specs.length).fill(0);

  // Pre-compute domain lengths for fast bounds checking.
  const lengths = specs.map((s) => s.domainSubjectIds.length);

  function step(depth: number): boolean {
    if (depth === specs.length) {
      return onTuple(tuple);
    }

    const spec = specs[depth];
    const len = lengths[depth];

    for (let i = 0; i < len; i++) {
      indices[depth] = i;
      tuple[depth] = spec.domainSubjectIds[i];
      if (!step(depth + 1)) {
        return false;
      }
    }

    return true;
  }

  step(0);
}

/**
 * Build a BindingEnv from a single tuple and its spec.
 */
export function tupleToBindings(
  specs: TupleBindingSpec[],
  tuple: number[],
): BindingEntry[] {
  const entries: BindingEntry[] = [];
  for (let i = 0; i < specs.length; i++) {
    const spec = specs[i];
    for (let j = 0; j < spec.variableNames.length; j++) {
      entries.push({
        variableName: spec.variableNames[j],
        subjectId: tuple[i],
      });
    }
  }
  return entries;
}