onenote-tools/src/onenote/onepkg/quantum/model.ts

/*
 * SPDX-License-Identifier: MIT
 *
 * TypeScript port of compcol `src/quantum/model.rs` at commit
 * 04a6db2aa7bd487a89c559631d79d1b384139f50.
 * Copyright (c) 2026 Karpeles Lab Inc., MIT License.
 * Modified to account for decoder work and report structured CAB errors.
 */

import { fail } from '../errors.js';
import { QuantumBitReader } from './bit-reader.js';
import { QuantumBudget } from './budget.js';

export class QuantumModel {
  shiftsLeft = 4;
  readonly symbols = new Uint16Array(65);
  readonly cumulativeFrequencies = new Uint16Array(65);

  constructor(
    start: number,
    readonly entries: number,
    private readonly budget: QuantumBudget
  ) {
    if (!Number.isInteger(entries) || entries < 1 || entries > 64) {
      fail('quantum-invalid-model', 'Quantum model size is out of range.', {
        structure: 'Quantum probability model',
      });
    }
    for (let index = 0; index <= entries; index += 1) {
      this.symbols[index] = start + index;
      this.cumulativeFrequencies[index] = entries - index;
    }
  }

  update(): void {
    this.shiftsLeft -= 1;
    if (this.shiftsLeft !== 0) {
      for (let index = this.entries - 1; index >= 0; index -= 1) {
        let frequency = this.cumulativeFrequencies[index]! >>> 1;
        if (frequency <= this.cumulativeFrequencies[index + 1]!) {
          frequency = this.cumulativeFrequencies[index + 1]! + 1;
        }
        this.cumulativeFrequencies[index] = frequency;
        this.budget.consume();
      }
      return;
    }

    this.shiftsLeft = 50;
    for (let index = 0; index < this.entries; index += 1) {
      this.cumulativeFrequencies[index] =
        (this.cumulativeFrequencies[index]! -
          this.cumulativeFrequencies[index + 1]! +
          1) >>>
        1;
      this.budget.consume();
    }
    // Preserve compcol/libmspack's selection-sort tie behavior exactly.
    for (let left = 0; left < this.entries - 1; left += 1) {
      for (let right = left + 1; right < this.entries; right += 1) {
        if (
          this.cumulativeFrequencies[left]! < this.cumulativeFrequencies[right]!
        ) {
          const frequency = this.cumulativeFrequencies[left]!;
          this.cumulativeFrequencies[left] = this.cumulativeFrequencies[right]!;
          this.cumulativeFrequencies[right] = frequency;
          const symbol = this.symbols[left]!;
          this.symbols[left] = this.symbols[right]!;
          this.symbols[right] = symbol;
        }
        this.budget.consume();
      }
    }
    for (let index = this.entries - 1; index >= 0; index -= 1) {
      this.cumulativeFrequencies[index] =
        this.cumulativeFrequencies[index]! +
        this.cumulativeFrequencies[index + 1]!;
      this.budget.consume();
    }
  }
}

export class QuantumArithmeticDecoder {
  private high = 0;
  private low = 0;
  private current = 0;

  constructor(private readonly budget: QuantumBudget) {}

  initializeFrame(reader: QuantumBitReader, input: Uint8Array): void {
    this.high = 0xffff;
    this.low = 0;
    this.current = reader.readBits(16, input);
  }

  getSymbol(
    model: QuantumModel,
    reader: QuantumBitReader,
    input: Uint8Array
  ): number {
    const high = this.high >>> 0;
    const low = this.low >>> 0;
    const current = this.current >>> 0;
    const range = (((high - low) >>> 0) & 0xffff) + 1;
    const total = model.cumulativeFrequencies[0]!;
    if (total === 0) {
      fail('quantum-invalid-model', 'Quantum probability model is exhausted.', {
        structure: 'Quantum probability model',
      });
    }
    const scaledCurrent = (current - low + 1) >>> 0;
    const numerator = (Math.imul(scaledCurrent, total) - 1) >>> 0;
    const symbolFrequency = Math.floor(numerator / range) & 0xffff;

    let index = 1;
    while (
      index < model.entries &&
      model.cumulativeFrequencies[index]! > symbolFrequency
    ) {
      index += 1;
      this.budget.consume();
    }
    const symbol = model.symbols[index - 1]!;

    const range2 = (high - low + 1) >>> 0;
    const frequencyLow = model.cumulativeFrequencies[index - 1]!;
    const frequencyHigh = model.cumulativeFrequencies[index]!;
    let nextHigh =
      (low + Math.floor((Math.imul(frequencyLow, range2) >>> 0) / total) - 1) &
      0xffff;
    let nextLow =
      (low + Math.floor((Math.imul(frequencyHigh, range2) >>> 0) / total)) &
      0xffff;

    for (let update = index - 1; update >= 0; update -= 1) {
      model.cumulativeFrequencies[update] =
        (model.cumulativeFrequencies[update]! + 8) & 0xffff;
      this.budget.consume();
    }
    if (model.cumulativeFrequencies[0]! > 3_800) model.update();

    let nextCurrent = current & 0xffff;
    while (true) {
      if ((nextLow & 0x8000) !== (nextHigh & 0x8000)) {
        if ((nextLow & 0x4000) !== 0 && (nextHigh & 0x4000) === 0) {
          nextCurrent ^= 0x4000;
          nextLow &= 0x3fff;
          nextHigh |= 0x4000;
        } else {
          break;
        }
      }
      nextLow = (nextLow << 1) & 0xffff;
      nextHigh = ((nextHigh << 1) | 1) & 0xffff;
      reader.ensureBits(1, input);
      const bit = reader.peekBits(1);
      reader.removeBits(1);
      nextCurrent = ((nextCurrent << 1) | bit) & 0xffff;
      this.budget.consume();
    }

    this.high = nextHigh;
    this.low = nextLow;
    this.current = nextCurrent;
    this.budget.consume();
    return symbol;
  }
}