From 65333cdcae483eace027459b803a17f0f93a3be5 Mon Sep 17 00:00:00 2001
From: rk <59029880+rrr523@users.noreply.github.com>
Date: Thu, 21 Dec 2023 23:46:34 +0800
Subject: [PATCH] Feat/reed solomon (#446)

* chore: Update docs

* docs: Update

* docs: Update

* feat: Set Webworker

* feat: Clean InjectWorker

* feat: Clean InjectWorker
---
 .changeset/mighty-laws-suffer.md              |  5 ++
 .changeset/tidy-boats-laugh.md                |  5 ++
 packages/reed-solomon/README.md               | 40 ++++------
 packages/reed-solomon/examples/node.js        |  4 +-
 .../reed-solomon/examples/web-worker.html     | 29 +++----
 packages/reed-solomon/src/index.js            | 49 +++++++++---
 packages/reed-solomon/src/node.adapter.js     | 54 ++-----------
 packages/reed-solomon/src/web.adapter.js      | 79 +++++++++----------
 packages/reed-solomon/types/index.d.ts        | 10 +++
 packages/reed-solomon/types/web.adapter.d.ts  |  4 +-
 10 files changed, 138 insertions(+), 141 deletions(-)
 create mode 100644 .changeset/mighty-laws-suffer.md
 create mode 100644 .changeset/tidy-boats-laugh.md

diff --git a/.changeset/mighty-laws-suffer.md b/.changeset/mighty-laws-suffer.md
new file mode 100644
index 00000000..58d65e3b
--- /dev/null
+++ b/.changeset/mighty-laws-suffer.md
@@ -0,0 +1,5 @@
+---
+'@bnb-chain/reed-solomon': minor
+---
+
+feat: Clean InjectWorker
diff --git a/.changeset/tidy-boats-laugh.md b/.changeset/tidy-boats-laugh.md
new file mode 100644
index 00000000..2a5e6476
--- /dev/null
+++ b/.changeset/tidy-boats-laugh.md
@@ -0,0 +1,5 @@
+---
+'@bnb-chain/reed-solomon': minor
+---
+
+feat: Support setting worker counts
diff --git a/packages/reed-solomon/README.md b/packages/reed-solomon/README.md
index c62133d1..ef9ca2be 100644
--- a/packages/reed-solomon/README.md
+++ b/packages/reed-solomon/README.md
@@ -33,7 +33,6 @@ Use directly in the browser via script tag:
   get reed solomon
 </button>
 
-<script src="https://cdn.jsdelivr.net/npm/@bnb-chain/reed-solomon/dist/index.aio.js"></script>
 <script src="https://cdn.jsdelivr.net/npm/@bnb-chain/reed-solomon/dist/index.aio.js"></script>
 <script>
   const fileInput = document.getElementById('file');
@@ -55,42 +54,37 @@ Use directly in the browser via script tag:
 ### Browser(WebWorker)
 
 ```html
+<!-- prefetch js -->
+<link rel="prefetch" href="https://unpkg.com/@bnb-chain/reed-solomon/dist/web.adapter.aio.js" />
+<link rel="prefetch" href="https://unpkg.com/@bnb-chain/reed-solomon/dist/utils.aio.js" />
+<script src="https://unpkg.com/@bnb-chain/reed-solomon/dist/web.adapter.aio.js"></script>
 <script>
+  const rs = new WebAdapter.WebAdapterReedSolomon()
+  rs.initWorkers({
+    workerNum: 6,
+    injectWorker,
+  })
+
   document.getElementById('worker-btn').onclick = async function() {
     const selectFile = fileInput.files[0];
     const arrBuffer = await selectFile.arrayBuffer()
     if (!arrBuffer) alert('no file selected');
 
     const sourceData = new Uint8Array(arrBuffer)
-    const rs = new WebAdapter.WebAdapterReedSolomon()
-    const res = await rs.encodeInWorker(injectWorker, sourceData)
+    const res = await rs.encodeInWorker(sourceData)
   }
 
   // inject worker
   function injectWorker() {
-    // replace your CDN url
-    importScripts('http://localhost:9002/dist/web.adapter.aio.js');
-    importScripts('http://localhost:9002/dist/utils.aio.js');
+    // or download this file and put it to your CDN server
+    importScripts('https://unpkg.com/@bnb-chain/reed-solomon/dist/web.adapter.aio.js');
+    importScripts('https://unpkg.com/@bnb-chain/reed-solomon/dist/utils.aio.js');
 
     const rs = new WebAdapter.WebAdapterReedSolomon();
-
     onmessage = function (event) {
       const { index, chunk } = event.data;
-      const encodeShards = rs.encodeSegment(chunk);
-      let encodeDataHash = [];
-
-      for (let i = 0; i < encodeShards.length; i++) {
-        const priceHash = RSUtils.sha256(encodeShards[i]);
-        encodeDataHash.push(priceHash);
-      }
-
-      postMessage({
-        index,
-        segChecksum: RSUtils.sha256(chunk),
-        encodeDataHash,
-      });
-
-      self.close();
+      const encodeShard = rs.getEncodeShard(chunk, index)
+      postMessage(encodeShard);
     };
   }
 </script>
@@ -128,6 +122,6 @@ const res = await rs.encodeInWorker(__filename, Uint8Array.from(fileBuffer))
 
 [Code](./examples/node-worker.js)
 
-## Benchark
+## Benchmark
 
 [benchmark](./benchmark.md)
diff --git a/packages/reed-solomon/examples/node.js b/packages/reed-solomon/examples/node.js
index 372739bb..156e70f1 100644
--- a/packages/reed-solomon/examples/node.js
+++ b/packages/reed-solomon/examples/node.js
@@ -3,13 +3,13 @@ const fs = require('node:fs');
 const path = require('node:path');
 const { ReedSolomon } = require('../dist/index');
 
-const fileBuffer = fs.readFileSync('./README.md');
+const sourceData = fs.readFileSync('./README.md');
 
 (async () => {
   const rs = new ReedSolomon();
   console.log('file size', sourceData.length / 1024 / 1024, 'm');
   console.time('cost');
-  const res = await rs.encode(Uint8Array.from(fileBuffer));
+  const res = await rs.encode(Uint8Array.from(sourceData));
   console.log('res', res);
   console.timeEnd('cost');
 })();
diff --git a/packages/reed-solomon/examples/web-worker.html b/packages/reed-solomon/examples/web-worker.html
index eadede31..71989c9c 100644
--- a/packages/reed-solomon/examples/web-worker.html
+++ b/packages/reed-solomon/examples/web-worker.html
@@ -17,6 +17,12 @@
   <script type="module">
     const fileInput = document.getElementById('file');
 
+    const rs = new WebAdapter.WebAdapterReedSolomon()
+    rs.initWorkers({
+      workerNum: 6,
+      injectWorker,
+    })
+
     // use webworker
     document.getElementById('worker-btn').onclick = async function() {
       const selectFile = fileInput.files[0];
@@ -26,8 +32,7 @@
       const sourceData = new Uint8Array(arrBuffer)
       console.time('webworker cost')
       console.log('file size', sourceData.length / 1024 / 1024, 'm')
-      const rs = new WebAdapter.WebAdapterReedSolomon()
-      const res = await rs.encodeInWorker(injectWorker, sourceData)
+      const res = await rs.encodeInWorker(sourceData)
       console.log('res', res)
       console.timeEnd('webworker cost')
     }
@@ -35,26 +40,16 @@
     function injectWorker() {
       importScripts('http://localhost:9002/dist/web.adapter.aio.js');
       importScripts('http://localhost:9002/dist/utils.aio.js');
+      // or use public CDN
+      // importScripts('https://unpkg.com/@bnb-chain/reed-solomon/dist/web.adapter.aio.js');
+      // importScripts('https://unpkg.com/@bnb-chain/reed-solomon/dist/utils.aio.js');
 
       const rs = new WebAdapter.WebAdapterReedSolomon();
 
       onmessage = function (event) {
         const { index, chunk } = event.data;
-        const encodeShards = rs.encodeSegment(chunk);
-        let encodeDataHash = [];
-
-        for (let i = 0; i < encodeShards.length; i++) {
-          const priceHash = RSUtils.sha256(encodeShards[i]);
-          encodeDataHash.push(priceHash);
-        }
-
-        postMessage({
-          index,
-          segChecksum: RSUtils.sha256(chunk),
-          encodeDataHash,
-        });
-
-        self.close();
+        const encodeShard = rs.getEncodeShard(chunk, index)
+        postMessage(encodeShard);
       };
     }
 
diff --git a/packages/reed-solomon/src/index.js b/packages/reed-solomon/src/index.js
index d455603c..57fd93ee 100644
--- a/packages/reed-solomon/src/index.js
+++ b/packages/reed-solomon/src/index.js
@@ -161,19 +161,46 @@ export class ReedSolomon {
     for (let i = 0; i < encodeDataHashList.length; i++) {
       encodeDataHashList[i] = [];
     }
+
+    let encodeShards = chunkList.map((chunk, index) => {
+      return this.getEncodeShard(chunk, index);
+    });
+
+    return this.getChecksumsByEncodeShards(encodeShards);
+  }
+
+  getEncodeShard(chunk, index) {
+    const encodeShards = this.encodeSegment(chunk);
+    let encodeDataHash = [];
+    for (let i = 0; i < encodeShards.length; i++) {
+      const priceHash = sha256(encodeShards[i]);
+      encodeDataHash.push(priceHash);
+    }
+    return {
+      index,
+      segChecksum: sha256(chunk),
+      encodeDataHash,
+    };
+  }
+
+  /**
+   * @param {Array[{index, segChecksum, encodeDataHash}]} encodeShards
+   */
+  getChecksumsByEncodeShards(encodeShards) {
     let hashList = [];
     let segChecksumList = [];
+    let encodeDataHashList = new Array(this.totalShards);
+    for (let i = 0; i < encodeDataHashList.length; i++) {
+      encodeDataHashList[i] = [];
+    }
 
-    for (let i = 0; i < chunkList.length; i++) {
-      const data = chunkList[i];
-      // console.log('data i', i)
-      const encodeShards = this.encodeSegment(data);
-      // console.log('data done', i)
-      segChecksumList.push(sha256(data));
+    for (let i = 0; i < encodeShards.length; i++) {
+      segChecksumList.push(encodeShards[i].segChecksum);
+    }
 
-      for (let i = 0; i < encodeShards.length; i++) {
-        const priceHash = sha256(encodeShards[i]);
-        encodeDataHashList[i].push(priceHash);
+    for (let i = 0; i < encodeShards.length; i++) {
+      for (let j = 0; j < encodeDataHashList.length; j++) {
+        encodeDataHashList[j][i] = encodeShards[i].encodeDataHash[j];
       }
     }
 
@@ -185,4 +212,8 @@ export class ReedSolomon {
 
     return toBase64(hashList);
   }
+
+  sortByIndex(encodeShards) {
+    return encodeShards.sort((a, b) => a.index - b.index);
+  }
 }
diff --git a/packages/reed-solomon/src/node.adapter.js b/packages/reed-solomon/src/node.adapter.js
index 76a3d887..15349715 100644
--- a/packages/reed-solomon/src/node.adapter.js
+++ b/packages/reed-solomon/src/node.adapter.js
@@ -1,24 +1,15 @@
-import { sha256 } from 'ethereum-cryptography/sha256.js';
 import { isMainThread, parentPort, Worker, workerData } from 'node:worker_threads';
 import { ReedSolomon } from './index';
-import { getIntegrityUint8Array, toBase64, splitPrice } from './utils';
+import { splitPrice } from './utils';
 
 export class NodeAdapterReedSolomon extends ReedSolomon {
   async encodeInWorker(p, sourceData) {
     return new Promise((resolve, reject) => {
       if (isMainThread) {
+        // RES is `encodeShards` Array
+        let RES = [];
         const chunkList = splitPrice(sourceData, this.segmentSize);
-
-        let hashList = [];
-        let segChecksumList = [];
-        let encodeDataHashList = new Array(6);
-
-        for (let i = 0; i < encodeDataHashList.length; i++) {
-          encodeDataHashList[i] = [];
-        }
-
         const threads = new Set();
-        let RES = [];
 
         for (let i = 0; i < chunkList.length; i++) {
           const worker = new Worker(p, {
@@ -38,24 +29,8 @@ export class NodeAdapterReedSolomon extends ReedSolomon {
             threads.delete(w);
             // console.log(`Thread exiting, ${threads.size} running...`)
             if (threads.size === 0) {
-              for (let i = 0; i < RES.length; i++) {
-                segChecksumList.push(RES[i].segChecksum);
-              }
-
-              for (let i = 0; i < chunkList.length; i++) {
-                for (let j = 0; j < encodeDataHashList.length; j++) {
-                  encodeDataHashList[j][i] = RES[i].encodeDataHash[j];
-                }
-              }
-
-              hashList[0] = sha256(getIntegrityUint8Array(segChecksumList));
-
-              for (let i = 0; i < encodeDataHashList.length; i++) {
-                hashList[i + 1] = sha256(getIntegrityUint8Array(encodeDataHashList[i]));
-              }
-
-              const res = toBase64(hashList);
-              resolve(res);
+              const sortedRes = this.sortByIndex(RES);
+              resolve(this.getChecksumsByEncodeShards(sortedRes));
             }
           });
 
@@ -65,23 +40,10 @@ export class NodeAdapterReedSolomon extends ReedSolomon {
           });
         }
       } else {
-        const encodeShards = this.encodeSegment(workerData.chunk);
-        let encodeDataHash = [];
-
-        for (let i = 0; i < encodeShards.length; i++) {
-          const priceHash = sha256(encodeShards[i]);
-          encodeDataHash.push(priceHash);
-        }
-
-        // console.log('encodeShards', encodeShards.length)
-        // console.log('encodeDataHash', encodeDataHash.length)
-        // console.log('workerData.index', workerData.index)
+        const { chunk, index } = workerData;
 
-        parentPort.postMessage({
-          index: workerData.index,
-          segChecksum: sha256(workerData.chunk),
-          encodeDataHash: encodeDataHash,
-        });
+        const encodeShard = this.getEncodeShard(chunk, index);
+        parentPort.postMessage(encodeShard);
       }
     });
   }
diff --git a/packages/reed-solomon/src/web.adapter.js b/packages/reed-solomon/src/web.adapter.js
index 75cca210..e400befc 100644
--- a/packages/reed-solomon/src/web.adapter.js
+++ b/packages/reed-solomon/src/web.adapter.js
@@ -1,58 +1,51 @@
 import { ReedSolomon } from '.';
-import { sha256, getIntegrityUint8Array, toBase64, splitPrice } from './utils';
+import { splitPrice } from './utils';
 
 export class WebAdapterReedSolomon extends ReedSolomon {
-  async encodeInWorker(workerFn, sourceData) {
-    const chunkList = splitPrice(sourceData, this.segmentSize);
+  initWorkers({ injectWorker, workerNum = 10 }) {
+    this.workerNum = workerNum;
+    this.workers = [];
+    for (let i = 0; i < workerNum; i++) {
+      const worker = createWorker(injectWorker);
+      this.workers.push(worker);
+    }
+  }
 
-    const workers = [];
+  async encodeInWorker(sourceData) {
+    // RES is `encodeShards` Array
+    let RES = [];
+    const chunkList = splitPrice(sourceData, this.segmentSize);
+    const queue = [];
 
     for (let i = 0; i < chunkList.length; i++) {
-      // const worker = new Worker('worker.js');
-      const worker = createWorker(workerFn);
-      workers.push(worker);
-      worker.postMessage({
+      queue.push({
         index: i,
         chunk: chunkList[i],
       });
     }
 
-    const plist = workers.map(
-      (worker) =>
-        new Promise((resolve) => {
-          worker.onmessage = (e) => {
-            resolve(e.data);
-          };
-        }),
-    );
-
-    return Promise.all(plist).then((RES) => {
-      let hashList = [];
-      let segChecksumList = [];
-      let encodeDataHashList = new Array(this.totalShards);
-      for (let i = 0; i < encodeDataHashList.length; i++) {
-        encodeDataHashList[i] = [];
-      }
-
-      for (let i = 0; i < RES.length; i++) {
-        segChecksumList.push(RES[i].segChecksum);
+    return new Promise((resolve) => {
+      let completedWorkers = 0;
+
+      for (let i = 0; i < queue.length; i++) {
+        const worker = this.workers[i % this.workerNum];
+        worker.postMessage({
+          index: queue[i].index,
+          chunk: queue[i].chunk,
+        });
+
+        worker.onmessage = (e) => {
+          // console.log('worker data', e.data)
+          completedWorkers++;
+          RES.push(e.data);
+
+          if (completedWorkers === queue.length) {
+            // console.log('RES', RES)
+            const sortedRes = this.sortByIndex(RES);
+            resolve(this.getChecksumsByEncodeShards(sortedRes));
+          }
+        };
       }
-
-      for (let i = 0; i < chunkList.length; i++) {
-        for (let j = 0; j < encodeDataHashList.length; j++) {
-          encodeDataHashList[j][i] = RES[i].encodeDataHash[j];
-        }
-      }
-
-      hashList[0] = sha256(getIntegrityUint8Array(segChecksumList));
-
-      for (let i = 0; i < encodeDataHashList.length; i++) {
-        hashList[i + 1] = sha256(getIntegrityUint8Array(encodeDataHashList[i]));
-      }
-
-      const res = toBase64(hashList);
-
-      return res;
     });
   }
 }
diff --git a/packages/reed-solomon/types/index.d.ts b/packages/reed-solomon/types/index.d.ts
index abff184a..cd2c3a18 100644
--- a/packages/reed-solomon/types/index.d.ts
+++ b/packages/reed-solomon/types/index.d.ts
@@ -1,9 +1,19 @@
 declare module '@bnb-chain/reed-solomon' {
+  export type EncodeShard = {
+    index: number;
+    segChecksum: Uint8Array;
+    encodeDataHash: Uint8Array[];
+  };
+
   export class ReedSolomon {
     constructor(dataShards = 4, parityShards = 2, segmentSize = 16777216);
 
     encodeSegment(data: Uint8Array): Uint8Array[];
 
     encode(data: Uint8Arary): string[];
+
+    getEncodeShard(chunk: Uint8Array, index: number): EncodeShard;
+
+    getChecksumsByEncodeShards(encodeShards: EncodeShard[]): string[];
   }
 }
diff --git a/packages/reed-solomon/types/web.adapter.d.ts b/packages/reed-solomon/types/web.adapter.d.ts
index ecb48500..b079cd86 100644
--- a/packages/reed-solomon/types/web.adapter.d.ts
+++ b/packages/reed-solomon/types/web.adapter.d.ts
@@ -1,5 +1,7 @@
 declare module '@bnb-chain/reed-solomon/web.adapter' {
   export class WebAdapterReedSolomon {
-    encodeInWorker(workerFn: () => void, data: Uint8Array): Promise<string[]>;
+    initWorkers(workersOpts: { workerNum?: number; injectWorker: () => void }): void;
+
+    encodeInWorker(data: Uint8Array): Promise<string[]>;
   }
 }