@cloudpss/ubjson/src/helper/decode-ae.ts

import type { DecodeOptions } from '../options.js';
import { constants } from './constants.js';
import { protoAction, constructorAction } from './decode.js';
import { unsupportedType } from './errors.js';
import { decode } from './string-decoder.js';
import { toUint8Array } from './utils.js';
const { fromCharCode } = String;

/**
 * 数据读取函数
 * `yield` 返回需要的字节数，`return` 返回解析结果
 */
export type DecodeFuncAe<T = unknown> = Generator<number, T, void>;

/** 数据包装 */
export interface DecodeCursorAe {
    /** 当前数据块 */
    readonly view: DataView;
    /** 当前数据块 */
    readonly data: Uint8Array<ArrayBuffer>;
    /** 当前数据大小 */
    readonly size: number;
    /** 当前读指针位置 */
    offset: number;
    /** 选项 */
    readonly options: DecodeOptions | undefined;
}

/** 创建数据包装 */
export function DecodeCursor(data: BufferSource, options?: DecodeOptions): DecodeCursorAe {
    const d = toUint8Array(data);
    const v = new DataView(d.buffer, d.byteOffset, d.byteLength);
    return {
        data: d,
        view: v,
        size: v.byteLength,
        offset: 0,
        options,
    };
}

/**
 * 读取第一个非 NOOP 的字节
 */
export function* readMarker(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    let marker: number;
    do {
        if (cursor.offset + 1 > cursor.size) {
            yield cursor.offset - cursor.size + 1;
        }
        marker = cursor.data[cursor.offset++]!;
    } while (marker === constants.NO_OP);
    return marker;
}

/** 读取一个大于 0 的整数 */
export function* readLength(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    let marker: number;
    do {
        if (cursor.offset + 1 > cursor.size) {
            yield cursor.offset - cursor.size + 1;
        }
        marker = cursor.data[cursor.offset++]!;
    } while (marker === constants.NO_OP);
    let length: number;
    switch (marker) {
        case constants.INT8:
            if (cursor.offset + 1 > cursor.size) {
                yield cursor.offset - cursor.size + 1;
            }
            length = cursor.view.getInt8(cursor.offset++);
            break;
        case constants.UINT8:
            length = yield* readUint8Data(cursor);
            break;
        case constants.INT16:
            length = yield* readInt16Data(cursor);
            break;
        case constants.INT32:
            length = yield* readInt32Data(cursor);
            break;
        case constants.INT64: {
            const l = yield* readInt64Data(cursor);
            if (l < 0 || l > Number.MAX_SAFE_INTEGER) {
                throw new Error('Invalid length');
            }
            length = Number(l);
            break;
        }
        default:
            throw new Error(`Unexpected marker '${fromCharCode(marker)}'(${marker}) for int length`);
    }
    if (length < 0) {
        throw new Error('Invalid length');
    }
    return length;
}

/** readInt8Data */
export function* readInt8Data(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    if (cursor.offset + 1 > cursor.size) {
        yield cursor.offset - cursor.size + 1;
    }
    return cursor.view.getInt8(cursor.offset++);
}
/** readUint8Data */
export function* readUint8Data(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    if (cursor.offset + 1 > cursor.size) {
        yield cursor.offset - cursor.size + 1;
    }
    return cursor.data[cursor.offset++]!;
}
/** readInt16Data */
export function* readInt16Data(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    if (cursor.offset + 2 > cursor.size) {
        yield cursor.offset - cursor.size + 2;
    }
    const value = cursor.view.getInt16(cursor.offset);
    cursor.offset += 2;
    return value;
}
/** readInt32Data */
export function* readInt32Data(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    if (cursor.offset + 4 > cursor.size) {
        yield cursor.offset - cursor.size + 4;
    }
    const value = cursor.view.getInt32(cursor.offset);
    cursor.offset += 4;
    return value;
}
/** readInt64Data */
export function* readInt64Data(cursor: DecodeCursorAe): DecodeFuncAe<bigint> {
    if (cursor.offset + 8 > cursor.size) {
        yield cursor.offset - cursor.size + 8;
    }
    const value = cursor.view.getBigInt64(cursor.offset);
    cursor.offset += 8;
    return value;
}
/** readFloat32Data */
export function* readFloat32Data(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    if (cursor.offset + 4 > cursor.size) {
        yield cursor.offset - cursor.size + 4;
    }
    const value = cursor.view.getFloat32(cursor.offset);
    cursor.offset += 4;
    return value;
}
/** readFloat64Data */
export function* readFloat64Data(cursor: DecodeCursorAe): DecodeFuncAe<number> {
    if (cursor.offset + 8 > cursor.size) {
        yield cursor.offset - cursor.size + 8;
    }
    const value = cursor.view.getFloat64(cursor.offset);
    cursor.offset += 8;
    return value;
}

/** 读取数据 */
export function* read(cursor: DecodeCursorAe): DecodeFuncAe<unknown> {
    let marker: number;
    do {
        if (cursor.offset + 1 > cursor.size) {
            yield cursor.offset - cursor.size + 1;
        }
        marker = cursor.data[cursor.offset++]!;
    } while (marker === constants.NO_OP);
    return yield* readData(cursor, marker);
}

/** 读取优化对象数据 */
function* readObjectOptimizedData(cursor: DecodeCursorAe, marker: OptimizedFormatMarkers): DecodeFuncAe<unknown> {
    const { count, type } = marker;
    const object: Record<string, unknown> = {};
    for (let i = 0; i < count; i++) {
        const key = yield* readKey(cursor);
        const value = yield* readData(cursor, type ?? (yield* readMarker(cursor)));
        if (key === '__proto__') {
            protoAction(cursor, object, value);
            continue;
        }
        if (key === 'constructor') {
            constructorAction(cursor, object, value);
            continue;
        }
        object[key] = value;
    }
    return object;
}

/** 根据标签读取后续数据 */
export function* readData(cursor: DecodeCursorAe, marker: number): DecodeFuncAe<unknown> {
    // 按照出现频率排序
    switch (marker) {
        case constants.STRING: {
            let marker: number;
            do {
                if (cursor.offset + 1 > cursor.size) {
                    yield cursor.offset - cursor.size + 1;
                }
                marker = cursor.data[cursor.offset++]!;
            } while (marker === constants.NO_OP);
            let length: number;
            switch (marker) {
                case constants.INT8:
                    if (cursor.offset + 1 > cursor.size) {
                        yield cursor.offset - cursor.size + 1;
                    }
                    length = cursor.view.getInt8(cursor.offset++);
                    break;
                case constants.UINT8:
                    length = yield* readUint8Data(cursor);
                    break;
                case constants.INT16:
                    length = yield* readInt16Data(cursor);
                    break;
                case constants.INT32:
                    length = yield* readInt32Data(cursor);
                    break;
                case constants.INT64: {
                    const l = yield* readInt64Data(cursor);
                    if (l < 0 || l > Number.MAX_SAFE_INTEGER) {
                        throw new Error('Invalid length');
                    }
                    length = Number(l);
                    break;
                }
                default:
                    throw new Error(`Unexpected marker '${fromCharCode(marker)}'(${marker}) for int length`);
            }
            if (length < 0) {
                throw new Error('Invalid length');
            }
            if (cursor.offset + length > cursor.size) {
                yield cursor.offset - cursor.size + length;
            }
            const begin = cursor.offset;
            const end = begin + length;
            cursor.offset = end;
            const { data } = cursor;
            return decode(data, begin, end);
        }
        case constants.OBJECT: {
            let type;
            let count;
            let marker: number;
            do {
                if (cursor.offset + 1 > cursor.size) {
                    yield cursor.offset - cursor.size + 1;
                }
                marker = cursor.data[cursor.offset++]!;
            } while (marker === constants.NO_OP);
            switch (marker) {
                case constants.TYPE_MARKER: {
                    do {
                        if (cursor.offset + 1 > cursor.size) {
                            yield cursor.offset - cursor.size + 1;
                        }
                        type = cursor.data[cursor.offset++]!;
                    } while (type === constants.NO_OP);
                    let marker2;
                    do {
                        if (cursor.offset + 1 > cursor.size) {
                            yield cursor.offset - cursor.size + 1;
                        }
                        marker2 = cursor.data[cursor.offset++]!;
                    } while (marker2 === constants.NO_OP);
                    if (marker2 !== constants.COUNT_MARKER) {
                        throw new Error('Expected count marker');
                    }
                }
                /* fall through */
                case constants.COUNT_MARKER: {
                    count = yield* readLength(cursor);
                    break;
                }
                default: {
                    // 不是 '$' 或 '#'
                    // 直到 '}'
                    const object: Record<string, unknown> = {};
                    for (;;) {
                        while (marker === constants.NO_OP) {
                            if (cursor.offset + 1 > cursor.size) {
                                yield cursor.offset - cursor.size + 1;
                            }
                            marker = cursor.data[cursor.offset++]!;
                        }
                        if (marker === constants.OBJECT_END) break;
                        let length: number;
                        switch (marker) {
                            case constants.INT8:
                                if (cursor.offset + 1 > cursor.size) {
                                    yield cursor.offset - cursor.size + 1;
                                }
                                length = cursor.view.getInt8(cursor.offset++);
                                break;
                            case constants.UINT8:
                                length = yield* readUint8Data(cursor);
                                break;
                            case constants.INT16:
                                length = yield* readInt16Data(cursor);
                                break;
                            case constants.INT32:
                                length = yield* readInt32Data(cursor);
                                break;
                            case constants.INT64: {
                                const l = yield* readInt64Data(cursor);
                                if (l < 0 || l > Number.MAX_SAFE_INTEGER) {
                                    throw new Error('Invalid length');
                                }
                                length = Number(l);
                                break;
                            }
                            default:
                                throw new Error(
                                    `Unexpected marker '${fromCharCode(marker)}'(${marker}) for int length`,
                                );
                        }
                        if (length < 0) {
                            throw new Error('Invalid length');
                        }
                        if (cursor.offset + length > cursor.size) {
                            yield cursor.offset - cursor.size + length;
                        }
                        marker = constants.NO_OP;
                        const begin = cursor.offset;
                        const end = begin + length;
                        cursor.offset = end;
                        const { data } = cursor;
                        const key = decode(data, begin, end);
                        let valueMarker: number;
                        do {
                            if (cursor.offset + 1 > cursor.size) {
                                yield cursor.offset - cursor.size + 1;
                            }
                            valueMarker = cursor.data[cursor.offset++]!;
                        } while (valueMarker === constants.NO_OP);
                        const value = yield* readData(cursor, valueMarker);
                        if (key === '__proto__') {
                            protoAction(cursor, object, value);
                            continue;
                        }
                        if (key === 'constructor') {
                            constructorAction(cursor, object, value);
                            continue;
                        }
                        object[key] = value;
                    }
                    return object;
                }
            }
            return yield* readObjectOptimizedData(cursor, { type, count });
        }
        case constants.ARRAY: {
            let type;
            let count;
            let marker: number;
            do {
                if (cursor.offset + 1 > cursor.size) {
                    yield cursor.offset - cursor.size + 1;
                }
                marker = cursor.data[cursor.offset++]!;
            } while (marker === constants.NO_OP);
            switch (marker) {
                case constants.TYPE_MARKER: {
                    do {
                        if (cursor.offset + 1 > cursor.size) {
                            yield cursor.offset - cursor.size + 1;
                        }
                        type = cursor.data[cursor.offset++]!;
                    } while (type === constants.NO_OP);
                    let marker2;
                    do {
                        if (cursor.offset + 1 > cursor.size) {
                            yield cursor.offset - cursor.size + 1;
                        }
                        marker2 = cursor.data[cursor.offset++]!;
                    } while (marker2 === constants.NO_OP);
                    if (marker2 !== constants.COUNT_MARKER) {
                        throw new Error('Expected count marker');
                    }
                }
                /* fall through */
                case constants.COUNT_MARKER: {
                    count = yield* readLength(cursor);
                    break;
                }
                default: {
                    // 不是 '$' 或 '#'
                    const array = [];
                    for (;;) {
                        while (marker === constants.NO_OP) {
                            if (cursor.offset + 1 > cursor.size) {
                                yield cursor.offset - cursor.size + 1;
                            }
                            marker = cursor.data[cursor.offset++]!;
                        }
                        // 直到 ']'
                        if (marker === constants.ARRAY_END) break;
                        array.push(yield* readData(cursor, marker));
                        marker = constants.NO_OP;
                    }
                    return array;
                }
            }
            switch (type) {
                case constants.UINT8: {
                    if (cursor.offset + count > cursor.size) {
                        yield cursor.offset - cursor.size + count;
                    }
                    const buf = new Uint8Array(
                        cursor.data.buffer,
                        cursor.data.byteOffset + cursor.offset,
                        count,
                    ).slice();
                    cursor.offset += count;
                    return buf;
                }
                case constants.INT8: {
                    if (cursor.offset + count > cursor.size) {
                        yield cursor.offset - cursor.size + count;
                    }
                    const buf = new Int8Array(
                        cursor.data.buffer,
                        cursor.data.byteOffset + cursor.offset,
                        count,
                    ).slice();
                    cursor.offset += count;
                    return buf;
                }
                case constants.INT16: {
                    if (cursor.offset + count * 2 > cursor.size) {
                        yield cursor.offset - cursor.size + count * 2;
                    }
                    const result = new Int16Array(count);
                    for (let i = 0; i < count; i++) {
                        result[i] = yield* readInt16Data(cursor);
                    }
                    return result;
                }
                case constants.INT32: {
                    if (cursor.offset + count * 4 > cursor.size) {
                        yield cursor.offset - cursor.size + count * 4;
                    }
                    const result = new Int32Array(count);
                    for (let i = 0; i < count; i++) {
                        result[i] = yield* readInt32Data(cursor);
                    }
                    return result;
                }
                case constants.FLOAT32: {
                    if (cursor.offset + count * 4 > cursor.size) {
                        yield cursor.offset - cursor.size + count * 4;
                    }
                    const result = new Float32Array(count);
                    for (let i = 0; i < count; i++) {
                        result[i] = yield* readFloat32Data(cursor);
                    }
                    return result;
                }
                case constants.FLOAT64: {
                    if (cursor.offset + count * 8 > cursor.size) {
                        yield cursor.offset - cursor.size + count * 8;
                    }
                    const result = new Float64Array(count);
                    for (let i = 0; i < count; i++) {
                        result[i] = yield* readFloat64Data(cursor);
                    }
                    return result;
                }
                case constants.INT64: {
                    if (cursor.offset + count * 8 > cursor.size) {
                        yield cursor.offset - cursor.size + count * 8;
                    }
                    const result = new BigInt64Array(count);
                    for (let i = 0; i < count; i++) {
                        result[i] = yield* readInt64Data(cursor);
                    }
                    return result;
                }
                case constants.NULL:
                    return Array.from({ length: count }).fill(null);
                case constants.TRUE:
                    return Array.from({ length: count }).fill(true);
                case constants.FALSE:
                    return Array.from({ length: count }).fill(false);
                case undefined:
                default:
                    break;
            }
            const array: unknown[] = [];
            array.length = count;
            for (let i = 0; i < count; i++) {
                array[i] = type === undefined ? yield* read(cursor) : yield* readData(cursor, type);
            }
            return array;
        }
        case constants.FLOAT64: {
            if (cursor.offset + 8 > cursor.size) {
                yield cursor.offset - cursor.size + 8;
            }
            const value = cursor.view.getFloat64(cursor.offset);
            cursor.offset += 8;
            return value;
        }
        case constants.UINT8: {
            if (cursor.offset + 1 > cursor.size) {
                yield cursor.offset - cursor.size + 1;
            }
            return cursor.data[cursor.offset++]!;
        }
        case constants.INT16: {
            if (cursor.offset + 2 > cursor.size) {
                yield cursor.offset - cursor.size + 2;
            }
            const value = cursor.view.getInt16(cursor.offset);
            cursor.offset += 2;
            return value;
        }
        case constants.FLOAT32: {
            if (cursor.offset + 4 > cursor.size) {
                yield cursor.offset - cursor.size + 4;
            }
            const value = cursor.view.getFloat32(cursor.offset);
            cursor.offset += 4;
            return value;
        }
        case constants.CHAR: {
            if (cursor.offset + 1 > cursor.size) {
                yield cursor.offset - cursor.size + 1;
            }
            return fromCharCode(cursor.data[cursor.offset++]!);
        }
        case constants.INT32: {
            if (cursor.offset + 4 > cursor.size) {
                yield cursor.offset - cursor.size + 4;
            }
            const value = cursor.view.getInt32(cursor.offset);
            cursor.offset += 4;
            return value;
        }
        case constants.INT8: {
            if (cursor.offset + 1 > cursor.size) {
                yield cursor.offset - cursor.size + 1;
            }
            return cursor.view.getInt8(cursor.offset++);
        }
        case constants.NULL:
            return null;
        case constants.TRUE:
            return true;
        case constants.FALSE:
            return false;
        case constants.INT64: {
            if (cursor.offset + 8 > cursor.size) {
                yield cursor.offset - cursor.size + 8;
            }
            const value = cursor.view.getBigInt64(cursor.offset);
            cursor.offset += 8;
            if (value < Number.MIN_SAFE_INTEGER || value > Number.MAX_SAFE_INTEGER) {
                return value;
            }
            return Number(value);
        }
        case constants.HIGH_PRECISION_NUMBER: {
            const length = yield* readLength(cursor);
            if (cursor.offset + length > cursor.size) {
                yield cursor.offset - cursor.size + length;
            }
            const begin = cursor.offset;
            const _buffer = new Uint8Array(cursor.data.buffer, begin, length).slice();
            cursor.offset = begin + length;
            // return _buffer
            unsupportedType('high precision number');
        }
    }
    throw new Error(`Unexpected marker '${fromCharCode(marker)}'(${marker})`);
}

/** readKey */
export function* readKey(cursor: DecodeCursorAe): DecodeFuncAe<string> {
    let marker: number;
    do {
        if (cursor.offset + 1 > cursor.size) {
            yield cursor.offset - cursor.size + 1;
        }
        marker = cursor.data[cursor.offset++]!;
    } while (marker === constants.NO_OP);
    let length: number;
    switch (marker) {
        case constants.INT8:
            if (cursor.offset + 1 > cursor.size) {
                yield cursor.offset - cursor.size + 1;
            }
            length = cursor.view.getInt8(cursor.offset++);
            break;
        case constants.UINT8:
            length = yield* readUint8Data(cursor);
            break;
        case constants.INT16:
            length = yield* readInt16Data(cursor);
            break;
        case constants.INT32:
            length = yield* readInt32Data(cursor);
            break;
        case constants.INT64: {
            const l = yield* readInt64Data(cursor);
            if (l < 0 || l > Number.MAX_SAFE_INTEGER) {
                throw new Error('Invalid length');
            }
            length = Number(l);
            break;
        }
        default:
            throw new Error(`Unexpected marker '${fromCharCode(marker)}'(${marker}) for int length`);
    }
    if (length < 0) {
        throw new Error('Invalid length');
    }
    if (cursor.offset + length > cursor.size) {
        yield cursor.offset - cursor.size + length;
    }
    const begin = cursor.offset;
    const end = begin + length;
    cursor.offset = end;
    const { data } = cursor;
    return decode(data, begin, end);
}

/** Optimized Format 数据 */
export type OptimizedFormatMarkers = { type?: number; count: number };