Merge pull request #248 from EddiG/alignment-fn

Allow the data alignment to support zero-copy decoding

Author: gfx

benchmark/msgpack-benchmark.js

@@ -3,13 +3,65 @@
 "use strict";
 require("ts-node/register");
 const Benchmark = require("benchmark");
-const fs = require("fs");
-const msgpack = require("../src");
+
+const msgpackEncode = require("..").encode;
+const msgpackDecode = require("..").decode;
+const ExtensionCodec = require("..").ExtensionCodec;
+
+const float32ArrayExtensionCodec = new ExtensionCodec();
+float32ArrayExtensionCodec.register({
+  type: 0x01,
+  encode: (object) => {
+    if (object instanceof Float32Array) {
+      return new Uint8Array(object.buffer, object.byteOffset, object.byteLength);
+    }
+    return null;
+  },
+  decode: (data) => {
+    const copy = new Uint8Array(data.byteLength);
+    copy.set(data);
+    return new Float32Array(copy.buffer);
+  },
+});
+
+const float32ArrayZeroCopyExtensionCodec = new ExtensionCodec();
+float32ArrayZeroCopyExtensionCodec.register({
+  type: 0x01,
+  encode: (object) => {
+    if (object instanceof Float32Array) {
+      return (pos) => {
+        const bpe = Float32Array.BYTES_PER_ELEMENT;
+        const padding = 1 + ((bpe - ((pos + 1) % bpe)) % bpe);
+        const data = new Uint8Array(object.buffer);
+        const result = new Uint8Array(padding + data.length);
+        result[0] = padding;
+        result.set(data, padding);
+        return result;
+      };
+    }
+    return null;
+  },
+  decode: (data) => {
+    const padding = data[0];
+    const bpe = Float32Array.BYTES_PER_ELEMENT;
+    const offset = data.byteOffset + padding;
+    const length = data.byteLength - padding;
+    return new Float32Array(data.buffer, offset, length / bpe);
+  },
+});
 
 const implementations = {
   "@msgpack/msgpack": {
-    encode: require("..").encode,
-    decode: require("..").decode,
+    encode: msgpackEncode,
+    decode: msgpackDecode,
+  },
+  "@msgpack/msgpack (Float32Array extension)": {
+    encode: (data) => msgpackEncode(data, { extensionCodec: float32ArrayExtensionCodec }),
+    decode: (data) => msgpackDecode(data, { extensionCodec: float32ArrayExtensionCodec }),
+  },
+  "@msgpack/msgpack (Float32Array with zero-copy extension)": {
+    encode: (data) => msgpackEncode(data, { extensionCodec: float32ArrayZeroCopyExtensionCodec }),
+    decode: (data) => msgpackDecode(data, { extensionCodec: float32ArrayZeroCopyExtensionCodec }),
   },
   "msgpack-lite": {
     encode: require("msgpack-lite").encode,
@@ -21,28 +73,52 @@ const implementations = {
   },
 };
 
-// exactly the same as:
-// https://raw.githubusercontent.com/endel/msgpack-benchmark/master/sample-large.json
-const sampleFiles = ["./sample-large.json"];
+const samples = [
+  {
+    // exactly the same as:
+    // https://raw.githubusercontent.com/endel/msgpack-benchmark/master/sample-large.json
+    name: "./sample-large.json",
+    data: require("./sample-large.json"),
+  },
+  {
+    name: "Large array of numbers",
+    data: [
+      {
+        position: new Array(1e3).fill(1.14),
+      },
+    ],
+  },
+  {
+    name: "Large Float32Array",
+    data: [
+      {
+        position: new Float32Array(1e3).fill(1.14),
+      },
+    ],
+  },
+];
 
 function validate(name, data, encoded) {
-  if (JSON.stringify(data) !== JSON.stringify(implementations[name].decode(encoded))) {
-    throw new Error("Bad implementation: " + name);
-  }
+  return JSON.stringify(data) === JSON.stringify(implementations[name].decode(encoded));
 }
 
-for (const sampleFile of sampleFiles) {
-  const data = require(sampleFile);
+for (const sample of samples) {
+  const { name: sampleName, data } = sample;
   const encodeSuite = new Benchmark.Suite();
   const decodeSuite = new Benchmark.Suite();
 
   console.log("");
-  console.log("**" + sampleFile + ":** (" + JSON.stringify(data).length + " bytes in JSON)");
+  console.log("**" + sampleName + ":** (" + JSON.stringify(data).length + " bytes in JSON)");
   console.log("");
 
   for (const name of Object.keys(implementations)) {
     implementations[name].toDecode = implementations[name].encode(data);
-    validate(name, data, implementations[name].toDecode);
+    if (!validate(name, data, implementations[name].toDecode)) {
+      console.log("```");
+      console.log("Not supported by " + name);
+      console.log("```");
+      continue;
+    }
     encodeSuite.add("(encode) " + name, () => {
       implementations[name].encode(data);
     });
@@ -60,7 +136,7 @@ for (const sampleFile of sampleFiles) {
 
   console.log("");
 
-  decodeSuite.on("cycle", function(event) {
+  decodeSuite.on("cycle", (event) => {
     console.log(String(event.target));
   });
 

src/Encoder.ts

@@ -396,6 +396,21 @@ export class Encoder<ContextType = undefined> {
   }
 
   private encodeExtension(ext: ExtData) {
+    if (typeof ext.data === "function") {
+      const data = ext.data(this.pos + 6);
+      const size = data.length;
+
+      if (size >= 0x100000000) {
+        throw new Error(`Too large extension object: ${size}`);
+      }
+
+      this.writeU8(0xc9);
+      this.writeU32(size);
+      this.writeI8(ext.type);
+      this.writeU8a(data);
+      return;
+    }
+
     const size = ext.data.length;
     if (size === 1) {
       // fixext 1

src/ExtData.ts

@@ -3,12 +3,9 @@
  */
 export class ExtData {
   readonly type: number;
-  readonly data: Uint8Array;
+  readonly data: Uint8Array | ((pos: number) => Uint8Array);
 
-  constructor(
-    type: number,
-    data: Uint8Array,
-  ) {
+  constructor(type: number, data: Uint8Array | ((pos: number) => Uint8Array)) {
     this.type = type;
     this.data = data;
   }

src/ExtensionCodec.ts

@@ -9,7 +9,10 @@ export type ExtensionDecoderType<ContextType> = (
   context: ContextType,
 ) => unknown;
 
-export type ExtensionEncoderType<ContextType> = (input: unknown, context: ContextType) => Uint8Array | null;
+export type ExtensionEncoderType<ContextType> = (
+  input: unknown,
+  context: ContextType,
+) => Uint8Array | ((dataPos: number) => Uint8Array) | null;
 
 // immutable interface to ExtensionCodec
 export type ExtensionCodecType<ContextType> = {

test/ExtensionCodec.test.ts

@@ -201,4 +201,43 @@ describe("ExtensionCodec", () => {
       ]);
     });
   });
+
+  context("custom extensions with alignment", () => {
+    const extensionCodec = new ExtensionCodec();
+
+    extensionCodec.register({
+      type: 0x01,
+      encode: (object: unknown) => {
+        if (object instanceof Float32Array) {
+          return (pos: number) => {
+            const bpe = Float32Array.BYTES_PER_ELEMENT;
+            const padding = 1 + ((bpe - ((pos + 1) % bpe)) % bpe);
+            const data = new Uint8Array(object.buffer);
+            const result = new Uint8Array(padding + data.length);
+            result[0] = padding;
+            result.set(data, padding);
+            return result;
+          };
+        }
+        return null;
+      },
+      decode: (data: Uint8Array) => {
+        const padding = data[0]!;
+        const bpe = Float32Array.BYTES_PER_ELEMENT;
+        const offset = data.byteOffset + padding;
+        const length = data.byteLength - padding;
+        return new Float32Array(data.buffer, offset, length / bpe);
+      },
+    });
+
+    it("encodes and decodes Float32Array type with zero-copy", () => {
+      const data = {
+        position: new Float32Array([1.1, 2.2, 3.3, 4.4, 5.5]),
+      };
+      const encoded = encode(data, { extensionCodec });
+      const decoded = decode(encoded, { extensionCodec });
+      assert.deepStrictEqual(decoded, data);
+      assert.strictEqual(decoded.position.buffer, encoded.buffer);
+    });
+  });
 });