@cloudpss/ubjson/src/stream-helper/encoder.ts

Opinionated UBJSON encoder/decoder for CloudPSS.

import { constants } from '../helper/constants.js';
import { encode, stringByteLength } from '../helper/string-encoder.js';
import { EncoderBase } from '../base/encoder.js';
import type { TypedArrayType } from '../helper/encode.js';
import type { EncodeOptions } from '../options.js';

const BLOCK_SIZE = 1024 * 64; // 64 KiB
const MAX_SIZE = 1024 * 1024 * 32; // 32 MiB

/** 保存一个内存池以减少重复分配 */
let POOL: Uint8Array | null = null;

/** 获取内存池 */
function alloc(size: number): Uint8Array {
    if (POOL == null || size !== BLOCK_SIZE) {
        return new Uint8Array(size);
    }
    const pool = POOL;
    POOL = null;
    return pool;
}

/** 归还内存池 */
function free(buf: Uint8Array): boolean {
    if (POOL == null && buf.byteLength === BLOCK_SIZE) {
        POOL = buf;
        return true;
    }
    return false;
}

/** 流式编码 UBJSON */
export class StreamEncoderHelper extends EncoderBase {
    constructor(
        options: EncodeOptions | null | undefined,
        protected readonly onChunk: (chunk: Uint8Array) => void,
    ) {
        super();
        this.sortObjectKeys = options?.sortObjectKeys ?? false;
        this.data = alloc(BLOCK_SIZE);
        this.view = new DataView(this.data.buffer);
    }
    /**
     * 销毁实例，释放内存池
     */
    destroy(): void {
        free(this.data);
        const self = this as unknown as { view: DataView | null; buffer: Uint8Array | null };
        self.view = null;
        self.buffer = null;
    }
    /**
     * 确保 buffer 还有 capacity 的空闲空间
     */
    protected ensureCapacity(capacity: number): void {
        if (capacity > MAX_SIZE) {
            // 超过最大尺寸限制
            throw new Error('Buffer has exceed max size');
        }
        // 无需扩容
        if (capacity >= 0 && this.data.byteLength >= this.length + capacity) return;

        const CURRENT_SIZE = this.data.byteLength;
        const NEXT_SIZE = Math.max(capacity, BLOCK_SIZE);
        const REUSE_BUF =
            CURRENT_SIZE >= NEXT_SIZE && // 满足容量需求
            CURRENT_SIZE - BLOCK_SIZE < NEXT_SIZE; // 不过于浪费
        // 提交目前的数据
        if (REUSE_BUF) {
            this.onChunk(this.data.slice(0, this.length));
        } else {
            if (free(this.data)) {
                // 归还内存池成功，buffer 可能重用，需要拷贝数据
                this.onChunk(this.data.slice(0, this.length));
            } else {
                this.onChunk(this.data.subarray(0, this.length));
            }
            // 重新分配缓冲区
            this.data = alloc(NEXT_SIZE);
            this.view = new DataView(this.data.buffer);
        }
        this.length = 0;
    }
    /** @inheritdoc */
    protected override writeLargeStringData(value: string): void {
        const strLen = value.length;
        const binLen = stringByteLength(value);
        this.ensureCapacity(5);
        this.data[this.length++] = constants.INT32;
        this.view.setInt32(this.length, binLen);
        this.length += 4;
        this.ensureCapacity(-1);

        // divide string to 64k chunks
        for (let i = 0; i < strLen; i += BLOCK_SIZE) {
            let end = i + BLOCK_SIZE;
            // avoid split surrogate pair
            const endAtSurrogate = end < strLen && (value.charCodeAt(end) & 0xfc00) === 0xdc00;
            if (endAtSurrogate) {
                end--;
            }
            const chunk = value.slice(i, end);
            this.onChunk(encode(chunk));
            if (endAtSurrogate) {
                i--;
            }
        }
    }
    /** @inheritdoc */
    protected override writeLargeTypedArrayData(type: TypedArrayType, value: ArrayBufferView): void {
        this.ensureCapacity(-1);
        const { byteLength } = value;
        if (type === constants.UINT8 || type === constants.INT8) {
            // fast path for typed arrays with `BYTES_PER_ELEMENT` of 1
            // divide buffer to 64k chunks
            const { buffer, byteOffset } = value;
            for (let i = 0; i < byteLength; i += BLOCK_SIZE) {
                this.onChunk(new Uint8Array(buffer.slice(byteOffset + i, byteOffset + i + BLOCK_SIZE)));
            }
            return;
        }
        if (type === constants.FLOAT64) {
            const arrayLength = byteLength / 8;
            for (let i = 0; i < arrayLength; i++) {
                this.ensureCapacity(8);
                this.view.setFloat64(this.length, (value as Float64Array)[i]!);
                this.length += 8;
            }
        } else if (type === constants.INT32) {
            const arrayLength = byteLength / 4;
            for (let i = 0; i < arrayLength; i++) {
                this.ensureCapacity(4);
                this.view.setInt32(this.length, (value as Int32Array)[i]!);
                this.length += 4;
            }
        } else if (type === constants.INT64) {
            const arrayLength = byteLength / 8;
            for (let i = 0; i < arrayLength; i++) {
                this.ensureCapacity(8);
                this.view.setBigInt64(this.length, (value as BigInt64Array)[i]!);
                this.length += 8;
            }
        } else if (type === constants.FLOAT32) {
            const arrayLength = byteLength / 4;
            for (let i = 0; i < arrayLength; i++) {
                this.ensureCapacity(4);
                this.view.setFloat32(this.length, (value as Float32Array)[i]!);
                this.length += 4;
            }
        } else {
            (type) satisfies constants.INT16;
            const arrayLength = byteLength / 2;
            for (let i = 0; i < arrayLength; i++) {
                this.ensureCapacity(2);
                this.view.setInt16(this.length, (value as Int16Array)[i]!);
                this.length += 2;
            }
        }
    }
    /** 获取写入结果 */
    encode(value: unknown): void {
        this.writeValue(value);
        this.ensureCapacity(-1);
    }
}