Allow the data alignment to support zero-copy decoding

Author: EddiG

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

@@ -4,6 +4,6 @@
 export class ExtData {
   constructor(
     readonly type: number,
-    readonly data: Uint8Array,
+    readonly data: Uint8Array | ((pos: number) => Uint8Array),
   ) {}
 }

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

@@ -196,4 +196,42 @@ 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 align = 4;
+            const padding = 1 + ((align - ((pos + 1) % align)) % align);
+            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 offset = data.byteOffset + padding;
+        const length = data.byteLength - padding;
+        return new Float32Array(data.buffer, offset, length / Float32Array.BYTES_PER_ELEMENT);
+      },
+    });
+
+    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);
+    });
+  });
 });