@cloudpss/ubjson/src/helper/encode.ts

import { constants } from './constants.js';
import { unsupportedView } from './errors.js';

/** 数据包装 */
export interface EncodeCursor {
    /** 数据 */
    readonly view: DataView;
    /** 数据 */
    readonly data: Uint8Array<ArrayBuffer>;
    /** 当前写指针位置 */
    length: number;
    /** 确保 buffer 还有 capacity 的空闲空间 */
    ensureCapacity(capacity: number): void;
}

/** 创建数据包装 */
export function EncodeCursor(length: number): EncodeCursor {
    const data = new Uint8Array(length);
    const view = new DataView(data.buffer, data.byteOffset, data.byteLength);
    return {
        data,
        view,
        length: 0,
        ensureCapacity(capacity: number) {
            if (this.length + capacity > this.data.length) {
                throw new RangeError('Out of capacity');
            }
        },
    };
}

/** 写入标签 */
export function writeMarker(cursor: EncodeCursor, marker: constants): void {
    cursor.ensureCapacity(1);
    cursor.data[cursor.length++] = marker;
}

const { MAX_SAFE_INTEGER } = Number;
export const I8_MASK = constants.INT8 << 8;
export const U8_MASK = constants.UINT8 << 8;
/** 写入长度 */
export function writeLength(cursor: EncodeCursor, length: number): void {
    if (length < 0x80) {
        cursor.ensureCapacity(2);
        cursor.view.setUint16(cursor.length, I8_MASK | length);
        cursor.length += 2;
    } else if (length < 0x8000) {
        cursor.ensureCapacity(3);
        cursor.data[cursor.length++] = constants.INT16;
        cursor.view.setInt16(cursor.length, length);
        cursor.length += 2;
    } else if (length < 0x8000_0000) {
        cursor.ensureCapacity(5);
        cursor.data[cursor.length++] = constants.INT32;
        cursor.view.setInt32(cursor.length, length);
        cursor.length += 4;
    } else if (length < MAX_SAFE_INTEGER) {
        cursor.ensureCapacity(9);
        cursor.data[cursor.length++] = constants.INT64;
        cursor.view.setBigInt64(cursor.length, BigInt(length));
        cursor.length += 8;
    } else {
        throw new RangeError('Invalid length');
    }
}

const { isFinite } = Number;
const { fround } = Math;
/** 写入数字 */
export function writeNumber(cursor: EncodeCursor, value: number): void {
    // eslint-disable-next-line unicorn/prefer-math-trunc
    if (value >> 0 === value) {
        if (value >= 0 && value <= 0xff) {
            cursor.ensureCapacity(2);
            const { length } = cursor;
            cursor.view.setUint16(length, U8_MASK | value);
            cursor.length = length + 2;
        } else if (value < 0x80 && value >= -0x80) {
            cursor.ensureCapacity(2);
            const { length } = cursor;
            cursor.view.setUint16(length, I8_MASK | (value & 0xff));
            cursor.length = length + 2;
        } else if (value < 0x8000 && value >= -0x8000) {
            cursor.ensureCapacity(3);
            const { length } = cursor;
            cursor.data[length] = constants.INT16;
            cursor.view.setInt16(length + 1, value);
            cursor.length = length + 3;
        } else {
            // must be 32 bit
            cursor.ensureCapacity(5);
            const { length } = cursor;
            cursor.data[length] = constants.INT32;
            cursor.view.setInt32(length + 1, value);
            cursor.length = length + 5;
        }
    } else if (!isFinite(value) || fround(value) === value) {
        // 如果不会损失精度，使用 32 位浮点
        cursor.ensureCapacity(5);
        const { length } = cursor;
        cursor.data[length] = constants.FLOAT32;
        cursor.view.setFloat32(length + 1, value);
        cursor.length = length + 5;
    } else {
        cursor.ensureCapacity(9);
        const { length } = cursor;
        cursor.data[length] = constants.FLOAT64;
        cursor.view.setFloat64(length + 1, value);
        cursor.length = length + 9;
    }
    return;
}

/** TypedArray 类型 */
export type TypedArrayType =
    | constants.UINT8
    | constants.INT8
    | constants.INT16
    | constants.INT32
    | constants.INT64
    | constants.FLOAT32
    | constants.FLOAT64;
const T_ARR_HEADER = (type: TypedArrayType): number =>
    (constants.ARRAY << 24) | (constants.TYPE_MARKER << 16) | (type << 8) | constants.COUNT_MARKER;
export const U8_ARR_HEADER = T_ARR_HEADER(constants.UINT8);
export const I8_ARR_HEADER = T_ARR_HEADER(constants.INT8);
export const I16_ARR_HEADER = T_ARR_HEADER(constants.INT16);
export const I32_ARR_HEADER = T_ARR_HEADER(constants.INT32);
export const I64_ARR_HEADER = T_ARR_HEADER(constants.INT64);
export const F32_ARR_HEADER = T_ARR_HEADER(constants.FLOAT32);
export const F64_ARR_HEADER = T_ARR_HEADER(constants.FLOAT64);

/** 写入 TypedArray 前导，包括 marker 和长度 */
export function writeTypedArrayHeader(cursor: EncodeCursor, value: ArrayBufferView): TypedArrayType {
    // ARRAY(1) + TYPE_MARKER(1) + TYPE(1) + COUNT_MARKER(1) + COUNT(MIN2 MAX5) + DATA
    cursor.ensureCapacity(9);
    let type: TypedArrayType;
    const { length } = cursor;
    if (value instanceof Uint8Array) {
        cursor.view.setUint32(length, U8_ARR_HEADER);
        type = constants.UINT8;
    } else if (value instanceof Float64Array) {
        cursor.view.setUint32(length, F64_ARR_HEADER);
        type = constants.FLOAT64;
    } else if (value instanceof Int32Array) {
        cursor.view.setUint32(length, I32_ARR_HEADER);
        type = constants.INT32;
    } else if (value instanceof BigInt64Array) {
        cursor.view.setUint32(length, I64_ARR_HEADER);
        type = constants.INT64;
    } else if (value instanceof Float32Array) {
        cursor.view.setUint32(length, F32_ARR_HEADER);
        type = constants.FLOAT32;
    } else if (value instanceof Int8Array) {
        cursor.view.setUint32(length, I8_ARR_HEADER);
        type = constants.INT8;
    } else if (value instanceof Int16Array) {
        cursor.view.setUint32(length, I16_ARR_HEADER);
        type = constants.INT16;
    } else {
        unsupportedView(value);
    }
    cursor.length = length + 4;
    writeLength(cursor, value.length);
    return type;
}
/** 写入 TypedArray */
export function writeTypedArray(cursor: EncodeCursor, value: ArrayBufferView): void {
    cursor.ensureCapacity(9 + value.byteLength);
    const type = writeTypedArrayHeader(cursor, value);
    writeTypedArrayData(cursor, type, value);
}
/** 写入 TypedArray 数据 */
export function writeTypedArrayData(cursor: EncodeCursor, type: TypedArrayType, value: ArrayBufferView): void {
    const { byteLength } = value;
    cursor.ensureCapacity(byteLength);
    let pointer = cursor.length;
    cursor.length = pointer + byteLength;
    if (type === constants.UINT8 || type === constants.INT8) {
        // fast path for typed arrays with `BYTES_PER_ELEMENT` of 1
        cursor.data.set(value as Uint8Array | Int8Array, pointer);
        return;
    }

    const { view } = cursor;
    if (type === constants.FLOAT64) {
        const arrayLength = byteLength / 8;
        for (let i = 0; i < arrayLength; i++) {
            view.setFloat64(pointer, (value as Float64Array)[i]!);
            pointer += 8;
        }
    } else if (type === constants.INT32) {
        const arrayLength = byteLength / 4;
        for (let i = 0; i < arrayLength; i++) {
            view.setInt32(pointer, (value as Int32Array)[i]!);
            pointer += 4;
        }
    } else if (type === constants.INT64) {
        const arrayLength = byteLength / 8;
        for (let i = 0; i < arrayLength; i++) {
            view.setBigInt64(pointer, (value as BigInt64Array)[i]!);
            pointer += 8;
        }
    } else if (type === constants.FLOAT32) {
        const arrayLength = byteLength / 4;
        for (let i = 0; i < arrayLength; i++) {
            view.setFloat32(pointer, (value as Float32Array)[i]!);
            pointer += 4;
        }
    } else {
        (type) satisfies constants.INT16;
        const arrayLength = byteLength / 2;
        for (let i = 0; i < arrayLength; i++) {
            view.setInt16(pointer, (value as Int16Array)[i]!);
            pointer += 2;
        }
    }
}