tornado-core/src/merkleTreeWorker.ts

/* eslint-disable @typescript-eslint/no-explicit-any */
import workerThreads from 'worker_threads';
import { MerkleTree, Element, TreeEdge, PartialMerkleTree } from '@tornado/fixed-merkle-tree';
import { mimc } from './mimc';
import { isNode } from './utils';

interface WorkData {
  merkleTreeHeight: number;
  edge?: TreeEdge;
  elements: Element[];
  zeroElement: string;
}

async function nodePostWork() {
  const { hash: hashFunction } = await mimc.getHash();
  const { merkleTreeHeight, edge, elements, zeroElement } = workerThreads.workerData as WorkData;

  if (edge) {
    const merkleTree = new PartialMerkleTree(merkleTreeHeight, edge, elements, {
      zeroElement,
      hashFunction,
    });

    (workerThreads.parentPort as workerThreads.MessagePort).postMessage(merkleTree.toString());
    return;
  }

  const merkleTree = new MerkleTree(merkleTreeHeight, elements, {
    zeroElement,
    hashFunction,
  });

  (workerThreads.parentPort as workerThreads.MessagePort).postMessage(merkleTree.toString());
}

if (isNode && workerThreads) {
  nodePostWork();
} else if (!isNode && typeof addEventListener === 'function' && typeof postMessage === 'function') {
  addEventListener('message', async (e: any) => {
    let data;

    if (e.data) {
      data = e.data;
    } else {
      data = e;
    }

    const { hash: hashFunction } = await mimc.getHash();
    const { merkleTreeHeight, edge, elements, zeroElement } = data as WorkData;

    if (edge) {
      const merkleTree = new PartialMerkleTree(merkleTreeHeight, edge, elements, {
        zeroElement,
        hashFunction,
      });

      postMessage(merkleTree.toString());
      return;
    }

    const merkleTree = new MerkleTree(merkleTreeHeight, elements, {
      zeroElement,
      hashFunction,
    });

    postMessage(merkleTree.toString());
  });
} else {
  throw new Error('This browser / environment does not support workers!');
}