"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.OnnxruntimeWebAssemblySessionHandler = void 0;
const onnx_proto_1 = require("onnx-proto");
const onnxruntime_common_1 = require("onnxruntime-common");
const binding_1 = require("./binding");
let ortInit;
const tensorDataTypeStringToEnum = (type) => {
    switch (type) {
        case 'int8':
            return onnx_proto_1.onnx.TensorProto.DataType.INT8;
        case 'uint8':
            return onnx_proto_1.onnx.TensorProto.DataType.UINT8;
        case 'bool':
            return onnx_proto_1.onnx.TensorProto.DataType.BOOL;
        case 'int16':
            return onnx_proto_1.onnx.TensorProto.DataType.INT16;
        case 'uint16':
            return onnx_proto_1.onnx.TensorProto.DataType.UINT16;
        case 'int32':
            return onnx_proto_1.onnx.TensorProto.DataType.INT32;
        case 'uint32':
            return onnx_proto_1.onnx.TensorProto.DataType.UINT32;
        case 'float32':
            return onnx_proto_1.onnx.TensorProto.DataType.FLOAT;
        case 'float64':
            return onnx_proto_1.onnx.TensorProto.DataType.DOUBLE;
        case 'string':
            return onnx_proto_1.onnx.TensorProto.DataType.STRING;
        case 'int64':
            return onnx_proto_1.onnx.TensorProto.DataType.INT64;
        case 'uint64':
            return onnx_proto_1.onnx.TensorProto.DataType.UINT64;
        default:
            throw new Error(`unsupported data type: ${type}`);
    }
};
const tensorDataTypeEnumToString = (typeProto) => {
    switch (typeProto) {
        case onnx_proto_1.onnx.TensorProto.DataType.INT8:
            return 'int8';
        case onnx_proto_1.onnx.TensorProto.DataType.UINT8:
            return 'uint8';
        case onnx_proto_1.onnx.TensorProto.DataType.BOOL:
            return 'bool';
        case onnx_proto_1.onnx.TensorProto.DataType.INT16:
            return 'int16';
        case onnx_proto_1.onnx.TensorProto.DataType.UINT16:
            return 'uint16';
        case onnx_proto_1.onnx.TensorProto.DataType.INT32:
            return 'int32';
        case onnx_proto_1.onnx.TensorProto.DataType.UINT32:
            return 'uint32';
        case onnx_proto_1.onnx.TensorProto.DataType.FLOAT:
            return 'float32';
        case onnx_proto_1.onnx.TensorProto.DataType.DOUBLE:
            return 'float64';
        case onnx_proto_1.onnx.TensorProto.DataType.STRING:
            return 'string';
        case onnx_proto_1.onnx.TensorProto.DataType.INT64:
            return 'int32';
        case onnx_proto_1.onnx.TensorProto.DataType.UINT64:
            return 'uint32';
        default:
            throw new Error(`unsupported data type: ${onnx_proto_1.onnx.TensorProto.DataType[typeProto]}`);
    }
};
const numericTensorTypeToTypedArray = (type) => {
    switch (type) {
        case 'float32':
            return Float32Array;
        case 'uint8':
            return Uint8Array;
        case 'int8':
            return Int8Array;
        case 'uint16':
            return Uint16Array;
        case 'int16':
            return Int16Array;
        case 'int32':
            return Int32Array;
        case 'bool':
            return Uint8Array;
        case 'float64':
            return Float64Array;
        case 'uint32':
            return Uint32Array;
        case 'int64':
            return BigInt64Array;
        case 'uint64':
            return BigUint64Array;
        default:
            throw new Error(`unsupported type: ${type}`);
    }
};
class OnnxruntimeWebAssemblySessionHandler {
    loadModel(model) {
        const wasm = binding_1.getInstance();
        if (!ortInit) {
            wasm._OrtInit();
            ortInit = true;
        }
        const modelDataOffset = wasm._malloc(model.byteLength);
        try {
            wasm.HEAPU8.set(model, modelDataOffset);
            this.sessionHandle = wasm._OrtCreateSession(modelDataOffset, model.byteLength);
        }
        finally {
            wasm._free(modelDataOffset);
        }
        const inputCount = wasm._OrtGetInputCount(this.sessionHandle);
        const outputCount = wasm._OrtGetOutputCount(this.sessionHandle);
        this.inputNames = [];
        this.inputNamesUTF8Encoded = [];
        this.outputNames = [];
        this.outputNamesUTF8Encoded = [];
        for (let i = 0; i < inputCount; i++) {
            const name = wasm._OrtGetInputName(this.sessionHandle, i);
            this.inputNamesUTF8Encoded.push(name);
            this.inputNames.push(wasm.UTF8ToString(name));
        }
        for (let i = 0; i < outputCount; i++) {
            const name = wasm._OrtGetOutputName(this.sessionHandle, i);
            this.outputNamesUTF8Encoded.push(name);
            this.outputNames.push(wasm.UTF8ToString(name));
        }
    }
    async dispose() {
        const wasm = binding_1.getInstance();
        if (this.inputNamesUTF8Encoded) {
            this.inputNamesUTF8Encoded.forEach(str => wasm._OrtFree(str));
            this.inputNamesUTF8Encoded = [];
        }
        if (this.outputNamesUTF8Encoded) {
            this.outputNamesUTF8Encoded.forEach(str => wasm._OrtFree(str));
            this.outputNamesUTF8Encoded = [];
        }
        if (this.sessionHandle) {
            wasm._OrtReleaseSession(this.sessionHandle);
            this.sessionHandle = 0;
        }
    }
    async run(feeds, fetches, _options) {
        const wasm = binding_1.getInstance();
        const inputArray = [];
        const inputIndices = [];
        Object.entries(feeds).forEach(kvp => {
            const name = kvp[0];
            const tensor = kvp[1];
            const index = this.inputNames.indexOf(name);
            if (index === -1) {
                throw new Error(`invalid input '${name}'`);
            }
            if (tensor.type === 'string') {
                // TODO: support string tensor
                throw new TypeError('string tensor is not supported');
            }
            inputArray.push(tensor);
            inputIndices.push(index);
        });
        const outputIndices = [];
        Object.entries(fetches).forEach(kvp => {
            const name = kvp[0];
            // TODO: support pre-allocated output
            const index = this.outputNames.indexOf(name);
            if (index === -1) {
                throw new Error(`invalid output '${name}'`);
            }
            outputIndices.push(index);
        });
        const inputCount = inputIndices.length;
        const outputCount = outputIndices.length;
        const inputValues = [];
        const inputDataOffsets = [];
        // create input tensors
        for (let i = 0; i < inputCount; i++) {
            const data = inputArray[i].data;
            if (Array.isArray(data)) {
                // string tensor
                throw new TypeError('string tensor is not supported');
            }
            else {
                const dataOffset = wasm._malloc(data.byteLength);
                inputDataOffsets.push(dataOffset);
                wasm.HEAPU8.set(new Uint8Array(data.buffer, data.byteOffset, data.byteLength), dataOffset);
                const dims = inputArray[i].dims;
                const stack = wasm.stackSave();
                const dimsOffset = wasm.stackAlloc(4 * dims.length);
                try {
                    let dimIndex = dimsOffset / 4;
                    dims.forEach(d => wasm.HEAP32[dimIndex++] = d);
                    const tensor = wasm._OrtCreateTensor(tensorDataTypeStringToEnum(inputArray[i].type), dataOffset, data.byteLength, dimsOffset, dims.length);
                    inputValues.push(tensor);
                }
                finally {
                    wasm.stackRestore(stack);
                }
            }
        }
        const beforeRunStack = wasm.stackSave();
        const inputValuesOffset = wasm.stackAlloc(inputCount * 4);
        const inputNamesOffset = wasm.stackAlloc(inputCount * 4);
        const outputValuesOffset = wasm.stackAlloc(outputCount * 4);
        const outputNamesOffset = wasm.stackAlloc(outputCount * 4);
        try {
            let inputValuesIndex = inputValuesOffset / 4;
            let inputNamesIndex = inputNamesOffset / 4;
            let outputValuesIndex = outputValuesOffset / 4;
            let outputNamesIndex = outputNamesOffset / 4;
            for (let i = 0; i < inputCount; i++) {
                wasm.HEAPU32[inputValuesIndex++] = inputValues[i];
                wasm.HEAPU32[inputNamesIndex++] = this.inputNamesUTF8Encoded[inputIndices[i]];
            }
            for (let i = 0; i < outputCount; i++) {
                wasm.HEAPU32[outputValuesIndex++] = 0;
                wasm.HEAPU32[outputNamesIndex++] = this.outputNamesUTF8Encoded[outputIndices[i]];
            }
            // support RunOptions
            const errorCode = wasm._OrtRun(this.sessionHandle, inputNamesOffset, inputValuesOffset, inputCount, outputNamesOffset, outputCount, outputValuesOffset);
            const output = {};
            if (errorCode === 0) {
                for (let i = 0; i < outputCount; i++) {
                    const tensor = wasm.HEAPU32[outputValuesOffset / 4 + i];
                    const beforeGetTensorDataStack = wasm.stackSave();
                    // stack allocate 4 pointer value
                    const tensorDataOffset = wasm.stackAlloc(4 * 4);
                    try {
                        wasm._OrtGetTensorData(tensor, tensorDataOffset, tensorDataOffset + 4, tensorDataOffset + 8, tensorDataOffset + 12);
                        let tensorDataIndex = tensorDataOffset / 4;
                        const dataType = wasm.HEAPU32[tensorDataIndex++];
                        const dataOffset = wasm.HEAPU32[tensorDataIndex++];
                        const dimsOffset = wasm.HEAPU32[tensorDataIndex++];
                        const dimsLength = wasm.HEAPU32[tensorDataIndex++];
                        const dims = [];
                        for (let i = 0; i < dimsLength; i++) {
                            dims.push(wasm.HEAPU32[dimsOffset / 4 + i]);
                        }
                        wasm._OrtFree(dimsOffset);
                        const type = tensorDataTypeEnumToString(dataType);
                        if (type === 'string') {
                            // string tensor
                            throw new TypeError('string tensor is not supported');
                        }
                        else {
                            const typedArray = numericTensorTypeToTypedArray(type);
                            const size = dims.length === 0 ? 1 : dims.reduce((a, b) => a * b);
                            const t = new onnxruntime_common_1.Tensor(type, new typedArray(size), dims);
                            new Uint8Array(t.data.buffer, t.data.byteOffset, t.data.byteLength)
                                .set(wasm.HEAPU8.subarray(dataOffset, dataOffset + t.data.byteLength));
                            output[this.outputNames[outputIndices[i]]] = t;
                        }
                    }
                    finally {
                        wasm.stackRestore(beforeGetTensorDataStack);
                    }
                    wasm._OrtReleaseTensor(tensor);
                }
            }
            inputValues.forEach(t => wasm._OrtReleaseTensor(t));
            inputDataOffsets.forEach(i => wasm._free(i));
            if (errorCode === 0) {
                return output;
            }
            else {
                throw new Error(`failed to call OrtRun(). error code = ${errorCode}.`);
            }
        }
        finally {
            wasm.stackRestore(beforeRunStack);
        }
    }
    startProfiling() {
        // TODO: implement profiling
    }
    endProfiling() {
        // TODO: implement profiling
    }
}
exports.OnnxruntimeWebAssemblySessionHandler = OnnxruntimeWebAssemblySessionHandler;
//# sourceMappingURL=session-handler.js.map