flatbuffers/ts/flexbuffers/reference.ts

import {fromByteWidth} from './bit-width-util.js';
import {BitWidth} from './bit-width.js';
import {fromUTF8Array} from './flexbuffers-util.js';
import {
  indirect,
  keyForIndex,
  keyIndex,
  readFloat,
  readInt,
  readUInt,
} from './reference-util.js';
import {
  fixedTypedVectorElementSize,
  fixedTypedVectorElementType,
  isAVector,
  isFixedTypedVector,
  isIndirectNumber,
  isNumber,
  isTypedVector,
  packedType,
  typedVectorElementType,
} from './value-type-util.js';
import {ValueType} from './value-type.js';

export function toReference(buffer: ArrayBuffer): Reference {
  const len = buffer.byteLength;

  if (len < 3) {
    throw new Error('Buffer needs to be bigger than 3 bytes');
  }

  const dataView = new DataView(buffer);
  const rootByteWidth = dataView.getUint8(len - 1);
  const packedType = dataView.getUint8(len - 2);

  let parentWidth = fromByteWidth(rootByteWidth);
  let offset = len - rootByteWidth - 2;

  const rootValueType = packedType >> 2;
  if (
    rootValueType === ValueType.VECTOR ||
    rootValueType === ValueType.MAP ||
    rootValueType === ValueType.BLOB ||
    rootValueType === ValueType.STRING
  ) {
    // Ensure parent width is wide enough to address the buffer
    let w = 1;
    while ((1 << (w * 8)) <= len && w < 8) w <<= 1;
    parentWidth = fromByteWidth(w);
    offset = len - w - 2; // no extra hacks
  }

  return new Reference(dataView, offset, parentWidth, packedType, '/');
}

function valueForIndexWithKey(
  index: number,
  key: string,
  dataView: DataView,
  offset: number,
  parentWidth: number,
  byteWidth: number,
  length: number,
  path: string,
): Reference {
  const _indirect = indirect(dataView, offset, parentWidth);
  const elementOffset = _indirect + index * byteWidth;
  const packedType = dataView.getUint8(_indirect + length * byteWidth + index);
  return new Reference(
    dataView,
    elementOffset,
    fromByteWidth(byteWidth),
    packedType,
    `${path}/${key}`,
  );
}

export class Reference {
  private readonly byteWidth: number;
  private readonly valueType: ValueType;
  private _length = -1;
  constructor(
    private dataView: DataView,
    private offset: number,
    private parentWidth: number,
    private packedType: ValueType,
    private path: string,
  ) {
    this.byteWidth = 1 << (packedType & 3);
    this.valueType = packedType >> 2;
  }

  isNull(): boolean {
    return this.valueType === ValueType.NULL;
  }
  isNumber(): boolean {
    return isNumber(this.valueType) || isIndirectNumber(this.valueType);
  }
  isFloat(): boolean {
    return (
      ValueType.FLOAT === this.valueType ||
      ValueType.INDIRECT_FLOAT === this.valueType
    );
  }
  isInt(): boolean {
    return this.isNumber() && !this.isFloat();
  }
  isString(): boolean {
    return (
      ValueType.STRING === this.valueType || ValueType.KEY === this.valueType
    );
  }
  isBool(): boolean {
    return ValueType.BOOL === this.valueType;
  }
  isBlob(): boolean {
    return ValueType.BLOB === this.valueType;
  }
  isVector(): boolean {
    return isAVector(this.valueType);
  }
  isMap(): boolean {
    return ValueType.MAP === this.valueType;
  }

  boolValue(): boolean | null {
    if (this.isBool()) {
      return readInt(this.dataView, this.offset, this.parentWidth) > 0;
    }
    return null;
  }

  intValue(): number | bigint | null {
    if (this.valueType === ValueType.INT) {
      return readInt(this.dataView, this.offset, this.parentWidth);
    }
    if (this.valueType === ValueType.UINT) {
      return readUInt(this.dataView, this.offset, this.parentWidth);
    }
    if (this.valueType === ValueType.INDIRECT_INT) {
      return readInt(
        this.dataView,
        indirect(this.dataView, this.offset, this.parentWidth),
        fromByteWidth(this.byteWidth),
      );
    }
    if (this.valueType === ValueType.INDIRECT_UINT) {
      return readUInt(
        this.dataView,
        indirect(this.dataView, this.offset, this.parentWidth),
        fromByteWidth(this.byteWidth),
      );
    }
    return null;
  }

  floatValue(): number | null {
    if (this.valueType === ValueType.FLOAT) {
      return readFloat(this.dataView, this.offset, this.parentWidth);
    }
    if (this.valueType === ValueType.INDIRECT_FLOAT) {
      return readFloat(
        this.dataView,
        indirect(this.dataView, this.offset, this.parentWidth),
        fromByteWidth(this.byteWidth),
      );
    }
    return null;
  }

  numericValue(): number | bigint | null {
    return this.floatValue() || this.intValue();
  }

  stringValue(): string | null {
    if (
      this.valueType === ValueType.STRING ||
      this.valueType === ValueType.KEY
    ) {
      const begin = indirect(this.dataView, this.offset, this.parentWidth);
      return fromUTF8Array(
        new Uint8Array(this.dataView.buffer, begin, this.length()),
      );
    }
    return null;
  }

  blobValue(): Uint8Array | null {
    if (this.isBlob()) {
      const begin = indirect(this.dataView, this.offset, this.parentWidth);
      return new Uint8Array(this.dataView.buffer, begin, this.length());
    }
    return null;
  }

  get(key: number): Reference {
    const length = this.length();
    if (Number.isInteger(key) && isAVector(this.valueType)) {
      if (key >= length || key < 0) {
        throw new Error(`Key: [${key}] is not applicable on ${this.path} of ${this.valueType} length: ${length}`);
      }
      const _indirect = indirect(this.dataView, this.offset, this.parentWidth);
      const elementOffset = _indirect + key * this.byteWidth;

      let _packedType: ValueType;

      if (isTypedVector(this.valueType)) {
      // Root typed vector: derive type instead of reading from buffer
      const _valueType = typedVectorElementType(this.valueType);
      _packedType = packedType(_valueType, BitWidth.WIDTH8);
    } else if (isFixedTypedVector(this.valueType)) {
      const _valueType = fixedTypedVectorElementType(this.valueType);
      _packedType = packedType(_valueType, BitWidth.WIDTH8);
    } else {
      // Only read packed type from buffer if it exists
      const typeOffset = _indirect + length * this.byteWidth + key;
      if (typeOffset < this.dataView.byteLength) {
        _packedType = this.dataView.getUint8(typeOffset);
      } else {
        // fallback for edge cases (e.g., root vectors)
        _packedType = this.packedType;
      }
      }

      return new Reference(
        this.dataView,
        elementOffset,
        fromByteWidth(this.byteWidth),
        _packedType,
        `${this.path}[${key}]`,
      );
    }

    if (typeof key === 'string') {
      const index = keyIndex(
        key,
        this.dataView,
        this.offset,
        this.parentWidth,
        this.byteWidth,
        length,
      );
      if (index !== null) {
        return valueForIndexWithKey(
          index,
          key,
          this.dataView,
          this.offset,
          this.parentWidth,
          this.byteWidth,
          length,
          this.path,
        );
      }
    }

    throw new Error(`Key [${key}] is not applicable on ${this.path} of ${this.valueType}`);
  }

  length(): number {
    let size;
    if (this._length > -1) {
      return this._length;
    }
    if (isFixedTypedVector(this.valueType)) {
      this._length = fixedTypedVectorElementSize(this.valueType);
    } else if (
      this.valueType === ValueType.BLOB ||
      this.valueType === ValueType.MAP ||
      isAVector(this.valueType)
    ) {
      this._length = readUInt(
        this.dataView,
        indirect(this.dataView, this.offset, this.parentWidth) - this.byteWidth,
        fromByteWidth(this.byteWidth),
      ) as number;
    } else if (this.valueType === ValueType.NULL) {
      this._length = 0;
    } else if (this.valueType === ValueType.STRING) {
      const _indirect = indirect(this.dataView, this.offset, this.parentWidth);
      let sizeByteWidth = this.byteWidth;
      size = readUInt(
        this.dataView,
        _indirect - sizeByteWidth,
        fromByteWidth(this.byteWidth),
      );
      while (this.dataView.getInt8(_indirect + (size as number)) !== 0) {
        sizeByteWidth <<= 1;
        size = readUInt(
          this.dataView,
          _indirect - sizeByteWidth,
          fromByteWidth(this.byteWidth),
        );
      }
      this._length = size as number;
    } else if (this.valueType === ValueType.KEY) {
      const _indirect = indirect(this.dataView, this.offset, this.parentWidth);
      size = 1;
      while (this.dataView.getInt8(_indirect + size) !== 0) {
        size++;
      }
      this._length = size;
    } else {
      this._length = 1;
    }
    return Number(this._length);
  }

  toObject(): unknown {
    const length = this.length();
    if (this.isVector()) {
      const result = [];
      for (let i = 0; i < length; i++) {
        result.push(this.get(i).toObject());
      }
      return result;
    }
    if (this.isMap()) {
      const result: Record<string, unknown> = {};
      for (let i = 0; i < length; i++) {
        const key = keyForIndex(
          i,
          this.dataView,
          this.offset,
          this.parentWidth,
          this.byteWidth,
        );
        result[key] = valueForIndexWithKey(
          i,
          key,
          this.dataView,
          this.offset,
          this.parentWidth,
          this.byteWidth,
          length,
          this.path,
        ).toObject();
      }
      return result;
    }
    if (this.isNull()) {
      return null;
    }
    if (this.isBool()) {
      return this.boolValue();
    }
    if (this.isNumber()) {
      return this.numericValue();
    }
    return this.blobValue() || this.stringValue();
  }
}