@cloudpss/ubjson/src/encoder.ts

Opinionated UBJSON encoder/decoder for CloudPSS.

import { EncoderBase } from './base/encoder.js';
import type { EncodeOptions } from './options.js';

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

/** 编码至 ubjson */
export class Encoder extends EncoderBase {
    private readonly flushedBuffers: Array<Uint8Array<ArrayBuffer>> = [];
    /** 通过内存池减少分配 */
    private readonly pool = new Uint8Array(BLOCK_SIZE);
    /** 缓存当前容量，避免对象访问耗时 */
    private capacity = 0;
    /**
     * 确保 buffer 还有 capacity 的空闲空间
     */
    protected ensureCapacity(capacity: number): void {
        if (capacity > MAX_SIZE) {
            // 超过最大尺寸限制
            throw new Error('Buffer has exceed max size');
        }
        if (this.capacity >= this.length + capacity) {
            // 无需扩容
            return;
        }
        this.flushBuffer(capacity);
    }

    /** 提交并扩容 */
    protected flushBuffer(capacity: number): void {
        // 提交目前的数据
        if (this.data === this.pool) {
            this.flushedBuffers.push(this.data.slice(0, this.length));
        } else {
            this.flushedBuffers.push(this.data.subarray(0, this.length));
        }

        // 重新分配缓冲区
        if (capacity < BLOCK_SIZE) capacity = BLOCK_SIZE;
        this.allocUnsafe(capacity);
    }

    /** 分配 buffer */
    private allocUnsafe(size: number): void {
        const data =
            size === BLOCK_SIZE
                ? // 从 pool 中获取
                  this.pool
                : // 新建 buffer
                  new Uint8Array(size);
        this.data = data;
        this.view = new DataView(data.buffer);
        this.length = 0;
        this.capacity = size;
    }

    /** 获取结果 */
    private getResult(): Uint8Array<ArrayBuffer> {
        if (this.flushedBuffers.length === 0) {
            // 缓冲区为空，复制当前 buffer
            return this.data.slice(0, this.length);
        }

        // 合并缓冲区
        const { length } = this;
        const activeBuffer = this.data.subarray(0, length);
        const flushedBuffers = this.flushedBuffers.splice(0);
        const size = flushedBuffers.reduce((sum, buffer) => sum + buffer.byteLength, length);
        const result = new Uint8Array(size);
        let offset = 0;
        for (const buffer of flushedBuffers) {
            result.set(buffer, offset);
            offset += buffer.byteLength;
        }
        result.set(activeBuffer, offset);
        return result;
    }

    /** 抛弃结果 */
    private cleanResult(): void {
        this.flushedBuffers.splice(0);
    }

    /** 获取写入结果 */
    encode(value: unknown): Uint8Array<ArrayBuffer> {
        try {
            this.allocUnsafe(BLOCK_SIZE);
            this.writeValue(value);
            return this.getResult();
        } catch (e) {
            this.cleanResult();
            throw e;
        }
    }

    /** 获取写入结果 */
    encodeMany(value: Iterable<unknown>): Uint8Array<ArrayBuffer> {
        try {
            this.allocUnsafe(BLOCK_SIZE);
            for (const v of value) {
                this.writeValue(v);
            }
            return this.getResult();
        } catch (e) {
            this.cleanResult();
            throw e;
        }
    }
}

let _ENCODER: Encoder | undefined;

/** 获取默认的编码器 */
export function getEncoder(options?: EncodeOptions): Encoder {
    _ENCODER ??= new Encoder();
    _ENCODER.sortObjectKeys = options?.sortObjectKeys ?? false;
    return _ENCODER;
}

/** 重置编码器, For testing only */
export function resetEncoder(): void {
    _ENCODER = undefined;
}