1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18import { memcpy } from '../util/buffer';
19import { BigIntAvailable, BigInt64Array, BigUint64Array } from '../util/compat';
20import {
21    TypedArray, TypedArrayConstructor,
22    BigIntArray, BigIntArrayConstructor
23} from '../interfaces';
24
25/** @ignore */ type DataValue<T> = T extends TypedArray ? number : T extends BigIntArray ? WideValue<T> : T;
26/** @ignore */ type WideValue<T extends BigIntArray> = T extends BigIntArray ? bigint | Int32Array | Uint32Array : never;
27/** @ignore */ type ArrayCtor<T extends TypedArray | BigIntArray> =
28    T extends TypedArray  ? TypedArrayConstructor<T>  :
29    T extends BigIntArray ? BigIntArrayConstructor<T> :
30    any;
31
32/** @ignore */
33const roundLengthUpToNearest64Bytes = (len: number, BPE: number) => ((((len * BPE) + 63) & ~63) || 64) / BPE;
34/** @ignore */
35const sliceOrExtendArray = <T extends TypedArray | BigIntArray>(arr: T, len = 0) => (
36    arr.length >= len ? arr.subarray(0, len) : memcpy(new (arr.constructor as any)(len), arr, 0)
37) as T;
38
39/** @ignore */
40export interface BufferBuilder<T extends TypedArray | BigIntArray = any, TValue = DataValue<T>> {
41    readonly offset: number;
42}
43
44/** @ignore */
45export class BufferBuilder<T extends TypedArray | BigIntArray = any, TValue = DataValue<T>> {
46
47    constructor(buffer: T, stride = 1) {
48        this.buffer = buffer;
49        this.stride = stride;
50        this.BYTES_PER_ELEMENT = buffer.BYTES_PER_ELEMENT;
51        this.ArrayType = buffer.constructor as ArrayCtor<T>;
52        this._resize(this.length = buffer.length / stride | 0);
53    }
54
55    public buffer: T;
56    public length: number;
57    public readonly stride: number;
58    public readonly ArrayType: ArrayCtor<T>;
59    public readonly BYTES_PER_ELEMENT: number;
60
61    public get byteLength() { return this.length * this.stride * this.BYTES_PER_ELEMENT | 0; }
62    public get reservedLength() { return this.buffer.length / this.stride; }
63    public get reservedByteLength() { return this.buffer.byteLength; }
64
65    // @ts-ignore
66    public set(index: number, value: TValue) { return this; }
67    public append(value: TValue) { return this.set(this.length, value); }
68    public reserve(extra: number) {
69        if (extra > 0) {
70            this.length += extra;
71            const stride = this.stride;
72            const length = this.length * stride;
73            const reserved = this.buffer.length;
74            if (length >= reserved) {
75                this._resize(reserved === 0
76                    ? roundLengthUpToNearest64Bytes(length * 1, this.BYTES_PER_ELEMENT)
77                    : roundLengthUpToNearest64Bytes(length * 2, this.BYTES_PER_ELEMENT)
78                );
79            }
80        }
81        return this;
82    }
83    public flush(length = this.length) {
84        length = roundLengthUpToNearest64Bytes(length * this.stride, this.BYTES_PER_ELEMENT);
85        const array = sliceOrExtendArray<T>(this.buffer, length);
86        this.clear();
87        return array;
88    }
89    public clear() {
90        this.length = 0;
91        this._resize(0);
92        return this;
93    }
94    protected _resize(newLength: number) {
95        return this.buffer = <T> memcpy(new this.ArrayType(newLength), this.buffer);
96    }
97}
98
99(BufferBuilder.prototype as any).offset = 0;
100
101/** @ignore */
102export class DataBufferBuilder<T extends TypedArray> extends BufferBuilder<T, number> {
103    public last() { return this.get(this.length - 1); }
104    public get(index: number) { return this.buffer[index]; }
105    public set(index: number, value: number) {
106        this.reserve(index - this.length + 1);
107        this.buffer[index * this.stride] = value;
108        return this;
109    }
110}
111
112/** @ignore */
113export class BitmapBufferBuilder extends DataBufferBuilder<Uint8Array> {
114
115    constructor(data = new Uint8Array(0)) { super(data, 1 / 8); }
116
117    public numValid = 0;
118    public get numInvalid() { return this.length - this.numValid; }
119    public get(idx: number) { return this.buffer[idx >> 3] >> idx % 8 & 1; }
120    public set(idx: number, val: number) {
121        const { buffer } = this.reserve(idx - this.length + 1);
122        const byte = idx >> 3, bit = idx % 8, cur = buffer[byte] >> bit & 1;
123        // If `val` is truthy and the current bit is 0, flip it to 1 and increment `numValid`.
124        // If `val` is falsey and the current bit is 1, flip it to 0 and decrement `numValid`.
125        val ? cur === 0 && ((buffer[byte] |=  (1 << bit)), ++this.numValid)
126            : cur === 1 && ((buffer[byte] &= ~(1 << bit)), --this.numValid);
127        return this;
128    }
129    public clear() {
130        this.numValid = 0;
131        return super.clear();
132    }
133}
134
135/** @ignore */
136export class OffsetsBufferBuilder extends DataBufferBuilder<Int32Array> {
137    constructor(data = new Int32Array(1)) { super(data, 1); }
138    public append(value: number) {
139        return this.set(this.length - 1, value);
140    }
141    public set(index: number, value: number) {
142        const offset = this.length - 1;
143        const buffer = this.reserve(index - offset + 1).buffer;
144        if (offset < index++) {
145            buffer.fill(buffer[offset], offset, index);
146        }
147        buffer[index] = buffer[index - 1] + value;
148        return this;
149    }
150    public flush(length = this.length - 1) {
151        if (length > this.length) {
152            this.set(length - 1, 0);
153        }
154        return super.flush(length + 1);
155    }
156}
157
158/** @ignore */
159export class WideBufferBuilder<T extends TypedArray, R extends BigIntArray> extends BufferBuilder<T, DataValue<T>> {
160    public buffer64!: R;
161    protected _ArrayType64!: BigIntArrayConstructor<R>;
162    public get ArrayType64() {
163        return this._ArrayType64 || (this._ArrayType64 = <BigIntArrayConstructor<R>> (this.buffer instanceof Int32Array ? BigInt64Array : BigUint64Array));
164    }
165    public set(index: number, value: DataValue<T>) {
166        this.reserve(index - this.length + 1);
167        switch (typeof value) {
168            case 'bigint': this.buffer64[index] = value; break;
169            case 'number': this.buffer[index * this.stride] = value; break;
170            default: this.buffer.set(value as TypedArray, index * this.stride);
171        }
172        return this;
173    }
174    protected _resize(newLength: number) {
175        const data = super._resize(newLength);
176        const length = data.byteLength / (this.BYTES_PER_ELEMENT * this.stride);
177        if (BigIntAvailable) {
178            this.buffer64 = new this.ArrayType64(data.buffer, data.byteOffset, length);
179        }
180        return data;
181    }
182}
183