feat: support for copy-on-read buffers (#14)

Cargo.toml

@@ -26,6 +26,7 @@ console_error_panic_hook = { version = "0.1.6", optional = true }
 # Unfortunately, `wee_alloc` requires nightly Rust when targeting wasm for now.
 wee_alloc = { version = "0.4.5", optional = true }
 web-sys = { version = "0.3.56", features = ["console"] }
+js-sys = "0.3.56"
 
 [dev-dependencies]
 wasm-bindgen-test = "0.3.13"

gen/wasm.d.ts

@@ -25,39 +25,27 @@ export function createWithPolynomial(mod_polynom: BigInt, bits: number, min_size
 */
 export function cut(rabin: Rabin, bytes: Uint8Array, end: boolean): Int32Array;
 /**
+* @param {Rabin} rabin
+* @param {any} buffer
+* @param {boolean} end
+* @returns {Int32Array}
 */
-export class Rabin {
-  free(): void;
-/**
-*/
-  max_size: number;
-/**
-*/
-  min_size: number;
-/**
-*/
-  polynom_shift: number;
+export function cut_buffer(rabin: Rabin, buffer: any, end: boolean): Int32Array;
 /**
 */
-  window_size: number;
+export class Rabin {
+  free(): void;
 }
 
 export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembly.Module;
 
 export interface InitOutput {
   readonly memory: WebAssembly.Memory;
+  readonly __wbg_rabin_free: (a: number) => void;
   readonly create: (a: number, b: number, c: number, d: number) => number;
   readonly createWithPolynomial: (a: number, b: number, c: number, d: number, e: number, f: number) => number;
   readonly cut: (a: number, b: number, c: number, d: number, e: number) => void;
-  readonly __wbg_rabin_free: (a: number) => void;
-  readonly __wbg_get_rabin_min_size: (a: number) => number;
-  readonly __wbg_set_rabin_min_size: (a: number, b: number) => void;
-  readonly __wbg_get_rabin_max_size: (a: number) => number;
-  readonly __wbg_set_rabin_max_size: (a: number, b: number) => void;
-  readonly __wbg_get_rabin_window_size: (a: number) => number;
-  readonly __wbg_set_rabin_window_size: (a: number, b: number) => void;
-  readonly __wbg_get_rabin_polynom_shift: (a: number) => number;
-  readonly __wbg_set_rabin_polynom_shift: (a: number, b: number) => void;
+  readonly cut_buffer: (a: number, b: number, c: number, d: number) => void;
   readonly __wbindgen_add_to_stack_pointer: (a: number) => number;
   readonly __wbindgen_malloc: (a: number) => number;
   readonly __wbindgen_free: (a: number, b: number) => void;

gen/wasm.js

@@ -4,6 +4,26 @@ export const activate = () => init(bytecode)
 
 let wasm;
 
+const heap = new Array(32).fill(undefined);
+
+heap.push(undefined, null, true, false);
+
+function getObject(idx) { return heap[idx]; }
+
+let heap_next = heap.length;
+
+function dropObject(idx) {
+    if (idx < 36) return;
+    heap[idx] = heap_next;
+    heap_next = idx;
+}
+
+function takeObject(idx) {
+    const ret = getObject(idx);
+    dropObject(idx);
+    return ret;
+}
+
 let cachedTextDecoder = new TextDecoder('utf-8', { ignoreBOM: true, fatal: true });
 
 cachedTextDecoder.decode();
@@ -19,6 +39,19 @@ function getUint8Memory0() {
 function getStringFromWasm0(ptr, len) {
     return cachedTextDecoder.decode(getUint8Memory0().subarray(ptr, ptr + len));
 }
+
+function getArrayU8FromWasm0(ptr, len) {
+    return getUint8Memory0().subarray(ptr / 1, ptr / 1 + len);
+}
+
+function addHeapObject(obj) {
+    if (heap_next === heap.length) heap.push(heap.length + 1);
+    const idx = heap_next;
+    heap_next = heap[idx];
+
+    heap[idx] = obj;
+    return idx;
+}
 /**
 * @param {number} bits
 * @param {number} min_size
@@ -100,6 +133,35 @@ export function cut(rabin, bytes, end) {
     }
 }
 
+let stack_pointer = 32;
+
+function addBorrowedObject(obj) {
+    if (stack_pointer == 1) throw new Error('out of js stack');
+    heap[--stack_pointer] = obj;
+    return stack_pointer;
+}
+/**
+* @param {Rabin} rabin
+* @param {any} buffer
+* @param {boolean} end
+* @returns {Int32Array}
+*/
+export function cut_buffer(rabin, buffer, end) {
+    try {
+        const retptr = wasm.__wbindgen_add_to_stack_pointer(-16);
+        _assertClass(rabin, Rabin);
+        wasm.cut_buffer(retptr, rabin.ptr, addBorrowedObject(buffer), end);
+        var r0 = getInt32Memory0()[retptr / 4 + 0];
+        var r1 = getInt32Memory0()[retptr / 4 + 1];
+        var v0 = getArrayI32FromWasm0(r0, r1).slice();
+        wasm.__wbindgen_free(r0, r1 * 4);
+        return v0;
+    } finally {
+        wasm.__wbindgen_add_to_stack_pointer(16);
+        heap[stack_pointer++] = undefined;
+    }
+}
+
 /**
 */
 export class Rabin {
@@ -122,54 +184,6 @@ export class Rabin {
         const ptr = this.__destroy_into_raw();
         wasm.__wbg_rabin_free(ptr);
     }
-    /**
-    */
-    get min_size() {
-        var ret = wasm.__wbg_get_rabin_min_size(this.ptr);
-        return ret >>> 0;
-    }
-    /**
-    * @param {number} arg0
-    */
-    set min_size(arg0) {
-        wasm.__wbg_set_rabin_min_size(this.ptr, arg0);
-    }
-    /**
-    */
-    get max_size() {
-        var ret = wasm.__wbg_get_rabin_max_size(this.ptr);
-        return ret >>> 0;
-    }
-    /**
-    * @param {number} arg0
-    */
-    set max_size(arg0) {
-        wasm.__wbg_set_rabin_max_size(this.ptr, arg0);
-    }
-    /**
-    */
-    get window_size() {
-        var ret = wasm.__wbg_get_rabin_window_size(this.ptr);
-        return ret >>> 0;
-    }
-    /**
-    * @param {number} arg0
-    */
-    set window_size(arg0) {
-        wasm.__wbg_set_rabin_window_size(this.ptr, arg0);
-    }
-    /**
-    */
-    get polynom_shift() {
-        var ret = wasm.__wbg_get_rabin_polynom_shift(this.ptr);
-        return ret >>> 0;
-    }
-    /**
-    * @param {number} arg0
-    */
-    set polynom_shift(arg0) {
-        wasm.__wbg_set_rabin_polynom_shift(this.ptr, arg0);
-    }
 }
 
 async function load(module, imports) {
@@ -209,6 +223,17 @@ async function init(input) {
     }
     const imports = {};
     imports.wbg = {};
+    imports.wbg.__wbg_copyTo_c2b6e812244947cd = function(arg0, arg1, arg2, arg3) {
+        var ret = getObject(arg0).copyTo(getArrayU8FromWasm0(arg1, arg2), arg3 >>> 0);
+        return addHeapObject(ret);
+    };
+    imports.wbg.__wbg_length_9cde16fa8cde1fcc = function(arg0) {
+        var ret = getObject(arg0).length;
+        return ret;
+    };
+    imports.wbg.__wbindgen_object_drop_ref = function(arg0) {
+        takeObject(arg0);
+    };
     imports.wbg.__wbindgen_throw = function(arg0, arg1) {
         throw new Error(getStringFromWasm0(arg0, arg1));
     };

lib.d.ts

@@ -22,3 +22,14 @@ export function createWithPolynom(
 ): Rabin | Promise<Rabin>
 
 export function cut(rabin: Rabin, bytes: Uint8Array, end: boolean): Int32Array
+
+export interface Buffer {
+  length: number
+  copyTo(buffer: Uint8Array, offset: number): any
+}
+
+export function cutBuffer(
+  rabin: Rabin,
+  buffer: Buffer,
+  end: boolean
+): Int32Array

lib.js

@@ -49,7 +49,6 @@ export const createWithPolynom = (
   })
 
 /**
- *
  * @param {Rabin} rabin
  * @param {Uint8Array} bytes
  * @param {boolean} [end=false]
@@ -61,4 +60,16 @@ export const cut = (rabin, bytes, end = false) =>
     ? none
     : Rabin.cut(rabin, bytes, end)
 
+/**
+ * @param {Rabin} rabin
+ * @param {{length:number, copyTo:(target:Uint8Array, offset:number) => unknown}} buffer
+ * @param {boolean} [end=false]
+ */
+export const cutBuffer = (rabin, buffer, end = false) =>
+  // If we have less then `maxSize` of bytes & it's not the end, there is no
+  // point to copy bytes into wasm as we'll get no chunks.
+  !end & (buffer.byteLength < rabin.maxSize)
+    ? none
+    : Rabin.cut_buffer(rabin, buffer, end)
+
 const none = new Uint32Array(0)

src/lib.rs

@@ -11,6 +11,17 @@ use wasm_bindgen::prelude::*;
 #[global_allocator]
 static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
 
+#[wasm_bindgen]
+extern "C" {
+    pub type Buffer;
+
+    #[wasm_bindgen(structural, method, js_name=copyTo)]
+    pub fn copy_to(this: &Buffer, target: &mut [u8], offset: usize) -> Buffer;
+
+    #[wasm_bindgen(method, getter)]
+    pub fn length(this: &Buffer) -> usize;
+}
+
 #[wasm_bindgen]
 pub fn create(bits: usize, min_size: usize, max_size: usize, window_size: usize) -> Rabin {
     Rabin::create(bits, min_size, max_size, window_size)
@@ -32,6 +43,13 @@ pub fn cut(rabin: &Rabin, bytes: &[u8], end: bool) -> Vec<i32> {
     rabin.split(bytes, end)
 }
 
+#[wasm_bindgen]
+pub fn cut_buffer(rabin: &Rabin, buffer: &Buffer, end: bool) -> Vec<i32> {
+    let mut bytes = vec![0; buffer.length()];
+    buffer.copy_to(&mut bytes, 0);
+    cut(&rabin, &bytes, end)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::*;

test/test.js

@@ -1,7 +1,7 @@
 // @ts-check
 
 import { encodeUTF8, read, sharbage } from "./util.js"
-import { create, cut, createWithPolynom } from "../lib.js"
+import { create, cut, cutBuffer, createWithPolynom } from "../lib.js"
 import { assert } from "chai"
 import * as FZSTD from "fzstd"
 
@@ -55,6 +55,31 @@ describe("rabin", () => {
     )
   })
 
+  it("shoud accept custom buffers", async () => {
+    const b1 = new Uint8Array(2 * 256)
+    const b2 = new Uint8Array(1 * 119)
+    const b3 = new Uint8Array(5 * 256)
+
+    b1.fill("a".charCodeAt(0))
+    b2.fill("b".charCodeAt(0))
+    b3.fill("c".charCodeAt(0))
+
+    const custom = {
+      length: b1.byteLength + b2.byteLength + b3.byteLength,
+      copyTo(buffer, offset) {
+        buffer.set(b1, offset)
+        buffer.set(b2, offset + b1.byteLength)
+        buffer.set(b3, offset + b1.byteLength + b2.byteLength)
+      },
+    }
+
+    const rabin = await create(6, 48, 192, 64)
+    assert.deepEqual(
+      [...cutBuffer(rabin, custom, true)],
+      [192, 192, 157, 64, 78, 192, 192, 192, 192, 192, 192, 76]
+    )
+  })
+
   it("chunks for rand_5MiB.zst", async () => {
     const buffer = FZSTD.decompress(await read("./rand_5MiB.zst"))