Added support for custom encoding/decoding context for keeping track of objects

grantila · grantila · commit 38ad5373c950 · 2020-01-14T01:22:13.000+01:00
diff --git a/README.md b/README.md
@@ -50,6 +50,7 @@ deepStrictEqual(decode(encoded), object);
   - [`decodeArrayStream(stream: AsyncIterable<ArrayLike<number>> | ReadableStream<ArrayLike<number>>, options?: DecodeAsyncOptions): AsyncIterable<unknown>`](#decodearraystreamstream-asynciterablearraylikenumber--readablestreamarraylikenumber-options-decodeasyncoptions-asynciterableunknown)
   - [`decodeStream(stream: AsyncIterable<ArrayLike<number>> | ReadableStream<ArrayLike<number>>, options?: DecodeAsyncOptions): AsyncIterable<unknown>`](#decodestreamstream-asynciterablearraylikenumber--readablestreamarraylikenumber-options-decodeasyncoptions-asynciterableunknown)
   - [Extension Types](#extension-types)
+    - [Codec context](#codec-context)
     - [Handling BigInt with ExtensionCodec](#handling-bigint-with-extensioncodec)
     - [The temporal module as timestamp extensions](#the-temporal-module-as-timestamp-extensions)
 - [MessagePack Specification](#messagepack-specification)
@@ -115,6 +116,7 @@ maxDepth | number | `100`
 initialBufferSize | number | `2048`
 sortKeys | boolean | false
 forceFloat32 | boolean | false
+context | user-defined | -
 
 ### `decode(buffer: ArrayLike<number> | ArrayBuffer, options?: DecodeOptions): unknown`
 
@@ -144,6 +146,7 @@ maxBinLength | number | `4_294_967_295` (UINT32_MAX)
 maxArrayLength | number | `4_294_967_295` (UINT32_MAX)
 maxMapLength | number | `4_294_967_295` (UINT32_MAX)
 maxExtLength | number | `4_294_967_295` (UINT32_MAX)
+context | user-defined | -
 
 You can use `max${Type}Length` to limit the length of each type decoded.
 
@@ -261,6 +264,50 @@ const decoded = decode(encoded, { extensionCodec });
 
 Not that extension types for custom objects must be `[0, 127]`, while `[-1, -128]` is reserved for MessagePack itself.
 
+#### Codec context
+
+When using an extension codec, it may be necessary to keep encoding/decoding state, to keep track of which objects got encoded/re-created. To do this, pass a `context` to the `EncodeOptions` and `DecodeOptions` (and if using typescript, type the `ExtensionCodec` too). Don't forget to pass the `{extensionCodec, context}` along recursive encoding/decoding:
+
+```typescript
+import { encode, decode, ExtensionCodec } from "@msgpack/msgpack";
+
+class MyContext {
+  track(object: any) { /*...*/ }
+}
+
+class MyType { /* ... */ }
+
+const extensionCodec = new ExtensionCodec<MyContext>();
+
+// MyType
+const MYTYPE_EXT_TYPE = 0 // Any in 0-127
+extensionCodec.register({
+  type: MYTYPE_EXT_TYPE,
+  encode: (object, context) => {
+    if (object instanceof MyType) {
+      context.track(object); // <-- like this
+      return encode(object.toJSON(), { extensionCodec, context });
+    } else {
+      return null;
+    }
+  },
+  decode: (data, extType, context) => {
+    const decoded = decode(data, { extensionCodec, context });
+    const my = new MyType(decoded);
+    context.track(my); // <-- and like this
+    return my;
+  },
+});
+
+// and later
+import { encode, decode } from "@msgpack/msgpack";
+
+const context = new MyContext();
+
+const encoded = = encode({myType: new MyType<any>()}, { extensionCodec, context });
+const decoded = decode(encoded, { extensionCodec, context });
+```
+
 #### Handling BigInt with ExtensionCodec
 
 This library does not handle BigInt by default, but you can handle it with `ExtensionCodec` like this:
diff --git a/src/Decoder.ts b/src/Decoder.ts
@@ -1,5 +1,5 @@
 import { prettyByte } from "./utils/prettyByte";
-import { ExtensionCodec } from "./ExtensionCodec";
+import { ExtensionCodec, ExtensionCodecType } from "./ExtensionCodec";
 import { getInt64, getUint64 } from "./utils/int";
 import { utf8DecodeJs, TEXT_ENCODING_AVAILABLE, TEXT_DECODER_THRESHOLD, utf8DecodeTD } from "./utils/utf8";
 import { createDataView, ensureUint8Array } from "./utils/typedArrays";
@@ -60,7 +60,7 @@ const DEFAULT_MAX_LENGTH = 0xffff_ffff; // uint32_max
 
 const sharedCachedKeyDecoder = new CachedKeyDecoder();
 
-export class Decoder {
+export class Decoder<ContextType> {
   totalPos = 0;
   pos = 0;
 
@@ -70,7 +70,8 @@ export class Decoder {
   readonly stack: Array<StackState> = [];
 
   constructor(
-    readonly extensionCodec = ExtensionCodec.defaultCodec,
+    readonly extensionCodec: ExtensionCodecType<ContextType> = ExtensionCodec.defaultCodec as any,
+    readonly context: ContextType,
     readonly maxStrLength = DEFAULT_MAX_LENGTH,
     readonly maxBinLength = DEFAULT_MAX_LENGTH,
     readonly maxArrayLength = DEFAULT_MAX_LENGTH,
@@ -519,7 +520,7 @@ export class Decoder {
 
     const extType = this.view.getInt8(this.pos + headOffset);
     const data = this.decodeBinary(size, headOffset + 1 /* extType */);
-    return this.extensionCodec.decode(data, extType);
+    return this.extensionCodec.decode(data, extType, this.context);
   }
 
   lookU8() {
diff --git a/src/Encoder.ts b/src/Encoder.ts
@@ -1,19 +1,20 @@
 import { utf8EncodeJs, utf8Count, TEXT_ENCODING_AVAILABLE, TEXT_ENCODER_THRESHOLD, utf8EncodeTE } from "./utils/utf8";
-import { ExtensionCodec } from "./ExtensionCodec";
+import { ExtensionCodec, ExtensionCodecType } from "./ExtensionCodec";
 import { setInt64, setUint64 } from "./utils/int";
 import { ensureUint8Array } from "./utils/typedArrays";
 import { ExtData } from "./ExtData";
 
 export const DEFAULT_MAX_DEPTH = 100;
 export const DEFAULT_INITIAL_BUFFER_SIZE = 2048;
 
-export class Encoder {
+export class Encoder<ContextType> {
   private pos = 0;
   private view = new DataView(new ArrayBuffer(this.initialBufferSize));
   private bytes = new Uint8Array(this.view.buffer);
 
   constructor(
-    readonly extensionCodec = ExtensionCodec.defaultCodec,
+    readonly extensionCodec: ExtensionCodecType<ContextType> = ExtensionCodec.defaultCodec as any,
+    readonly context: ContextType,
     readonly maxDepth = DEFAULT_MAX_DEPTH,
     readonly initialBufferSize = DEFAULT_INITIAL_BUFFER_SIZE,
     readonly sortKeys = false,
@@ -173,7 +174,7 @@ export class Encoder {
 
   encodeObject(object: unknown, depth: number) {
     // try to encode objects with custom codec first of non-primitives
-    const ext = this.extensionCodec.tryToEncode(object);
+    const ext = this.extensionCodec.tryToEncode(object, this.context);
     if (ext != null) {
       this.encodeExtension(ext);
     } else if (Array.isArray(object)) {
diff --git a/src/ExtensionCodec.ts b/src/ExtensionCodec.ts
@@ -3,26 +3,37 @@
 import { ExtData } from "./ExtData";
 import { timestampExtension } from "./timestamp";
 
-export type ExtensionDecoderType = (data: Uint8Array, extensionType: number) => unknown;
+export type ExtensionDecoderType<ContextType> = (
+  data: Uint8Array,
+  extensionType: number,
+  context: ContextType,
+) => unknown;
 
-export type ExtensionEncoderType = (input: unknown) => Uint8Array | null;
+export type ExtensionEncoderType<ContextType> = (input: unknown, context: ContextType) => Uint8Array | null;
 
 // immutable interfce to ExtensionCodec
-export type ExtensionCodecType = {
-  tryToEncode(object: unknown): ExtData | null;
-  decode(data: Uint8Array, extType: number): unknown;
+export type ExtensionCodecType<ContextType> = {
+  dummy: ContextType;
+  tryToEncode(object: unknown, context: ContextType): ExtData | null;
+  decode(data: Uint8Array, extType: number, context: ContextType): unknown;
 };
 
-export class ExtensionCodec implements ExtensionCodecType {
-  public static readonly defaultCodec: ExtensionCodecType = new ExtensionCodec();
+const typeDummy: any = undefined;
+
+export class ExtensionCodec<ContextType = undefined> implements ExtensionCodecType<ContextType> {
+  public static readonly defaultCodec: ExtensionCodecType<undefined> = new ExtensionCodec();
+
+  // ensures ExtensionCodecType<X> matches ExtensionCodec<X>
+  // this will make type errors a lot more clear
+  dummy: ContextType = typeDummy;
 
   // built-in extensions
-  private readonly builtInEncoders: Array<ExtensionEncoderType | undefined | null> = [];
-  private readonly builtInDecoders: Array<ExtensionDecoderType | undefined | null> = [];
+  private readonly builtInEncoders: Array<ExtensionEncoderType<ContextType> | undefined | null> = [];
+  private readonly builtInDecoders: Array<ExtensionDecoderType<ContextType> | undefined | null> = [];
 
   // custom extensions
-  private readonly encoders: Array<ExtensionEncoderType | undefined | null> = [];
-  private readonly decoders: Array<ExtensionDecoderType | undefined | null> = [];
+  private readonly encoders: Array<ExtensionEncoderType<ContextType> | undefined | null> = [];
+  private readonly decoders: Array<ExtensionDecoderType<ContextType> | undefined | null> = [];
 
   public constructor() {
     this.register(timestampExtension);
@@ -34,8 +45,8 @@ export class ExtensionCodec implements ExtensionCodecType {
     decode,
   }: {
     type: number;
-    encode: ExtensionEncoderType;
-    decode: ExtensionDecoderType;
+    encode: ExtensionEncoderType<ContextType>;
+    decode: ExtensionDecoderType<ContextType>;
   }): void {
     if (type >= 0) {
       // custom extensions
@@ -49,12 +60,12 @@ export class ExtensionCodec implements ExtensionCodecType {
     }
   }
 
-  public tryToEncode(object: unknown): ExtData | null {
+  public tryToEncode(object: unknown, context: ContextType): ExtData | null {
     // built-in extensions
     for (let i = 0; i < this.builtInEncoders.length; i++) {
       const encoder = this.builtInEncoders[i];
       if (encoder != null) {
-        const data = encoder(object);
+        const data = encoder(object, context);
         if (data != null) {
           const type = -1 - i;
           return new ExtData(type, data);
@@ -66,7 +77,7 @@ export class ExtensionCodec implements ExtensionCodecType {
     for (let i = 0; i < this.encoders.length; i++) {
       const encoder = this.encoders[i];
       if (encoder != null) {
-        const data = encoder(object);
+        const data = encoder(object, context);
         if (data != null) {
           const type = i;
           return new ExtData(type, data);
@@ -81,10 +92,10 @@ export class ExtensionCodec implements ExtensionCodecType {
     return null;
   }
 
-  public decode(data: Uint8Array, type: number): unknown {
+  public decode(data: Uint8Array, type: number, context: ContextType): unknown {
     const decoder = type < 0 ? this.builtInDecoders[-1 - type] : this.decoders[type];
     if (decoder) {
-      return decoder(data, type);
+      return decoder(data, type, context);
     } else {
       // decode() does not fail, returns ExtData instead.
       return new ExtData(type, data);
diff --git a/src/context.ts b/src/context.ts
@@ -0,0 +1,11 @@
+export type SplitTypes<T, U> = U extends T ? U : Exclude<T, U>;
+export type SplitUndefined<T> = SplitTypes<T, undefined>;
+
+export type ContextOf<ContextType> = ContextType extends undefined
+  ? {}
+  : {
+      /**
+       * Custom user-defined data, read/writable
+       */
+      context: ContextType;
+    };
diff --git a/src/decode.ts b/src/decode.ts
@@ -1,9 +1,10 @@
 import { ExtensionCodecType } from "./ExtensionCodec";
 import { Decoder } from "./Decoder";
+import { ContextOf, SplitUndefined } from "./context";
 
-export type DecodeOptions = Partial<
-  Readonly<{
-    extensionCodec: ExtensionCodecType;
+export type DecodeOptions<ContextType = undefined> = Readonly<
+  Partial<{
+    extensionCodec: ExtensionCodecType<ContextType>;
 
     /**
      * Maximum string length.
@@ -31,7 +32,8 @@ export type DecodeOptions = Partial<
      */
     maxExtLength: number;
   }>
->;
+> &
+  ContextOf<ContextType>;
 
 export const defaultDecodeOptions: DecodeOptions = {};
 
@@ -40,12 +42,13 @@ export const defaultDecodeOptions: DecodeOptions = {};
  *
  * This is a synchronous decoding function. See other variants for asynchronous decoding: `decodeAsync()`, `decodeStream()`, `decodeArrayStream()`.
  */
-export function decode(
+export function decode<ContextType>(
   buffer: ArrayLike<number> | ArrayBuffer,
-  options: DecodeOptions = defaultDecodeOptions,
+  options: DecodeOptions<SplitUndefined<ContextType>> = defaultDecodeOptions as any,
 ): unknown {
-  const decoder = new Decoder(
+  const decoder = new Decoder<ContextType>(
     options.extensionCodec,
+    (options as typeof options & { context: any }).context,
     options.maxStrLength,
     options.maxBinLength,
     options.maxArrayLength,
diff --git a/src/decodeAsync.ts b/src/decodeAsync.ts
@@ -1,15 +1,17 @@
 import { Decoder } from "./Decoder";
 import { defaultDecodeOptions, DecodeOptions } from "./decode";
 import { ensureAsyncIterabe, ReadableStreamLike } from "./utils/stream";
+import { SplitUndefined } from "./context";
 
-export async function decodeAsync(
+export async function decodeAsync<ContextType>(
   streamLike: ReadableStreamLike<ArrayLike<number>>,
-  options: DecodeOptions = defaultDecodeOptions,
+  options: DecodeOptions<SplitUndefined<ContextType>> = defaultDecodeOptions as any,
 ): Promise<unknown> {
   const stream = ensureAsyncIterabe(streamLike);
 
-  const decoder = new Decoder(
+  const decoder = new Decoder<ContextType>(
     options.extensionCodec,
+    (options as typeof options & { context: any }).context,
     options.maxStrLength,
     options.maxBinLength,
     options.maxArrayLength,
@@ -19,14 +21,15 @@ export async function decodeAsync(
   return decoder.decodeSingleAsync(stream);
 }
 
-export function decodeArrayStream(
+export function decodeArrayStream<ContextType>(
   streamLike: ReadableStreamLike<ArrayLike<number>>,
-  options: DecodeOptions = defaultDecodeOptions,
+  options: DecodeOptions<SplitUndefined<ContextType>> = defaultDecodeOptions as any,
 ) {
   const stream = ensureAsyncIterabe(streamLike);
 
-  const decoder = new Decoder(
+  const decoder = new Decoder<ContextType>(
     options.extensionCodec,
+    (options as typeof options & { context: any }).context,
     options.maxStrLength,
     options.maxBinLength,
     options.maxArrayLength,
@@ -37,14 +40,15 @@ export function decodeArrayStream(
   return decoder.decodeArrayStream(stream);
 }
 
-export function decodeStream(
+export function decodeStream<ContextType>(
   streamLike: ReadableStreamLike<ArrayLike<number>>,
-  options: DecodeOptions = defaultDecodeOptions,
+  options: DecodeOptions<SplitUndefined<ContextType>> = defaultDecodeOptions as any,
 ) {
   const stream = ensureAsyncIterabe(streamLike);
 
-  const decoder = new Decoder(
+  const decoder = new Decoder<ContextType>(
     options.extensionCodec,
+    (options as typeof options & { context: any }).context,
     options.maxStrLength,
     options.maxBinLength,
     options.maxArrayLength,
diff --git a/src/encode.ts b/src/encode.ts
@@ -1,9 +1,10 @@
 import { ExtensionCodecType } from "./ExtensionCodec";
 import { Encoder } from "./Encoder";
+import { ContextOf, SplitUndefined } from "./context";
 
-export type EncodeOptions = Partial<
+export type EncodeOptions<ContextType = undefined> = Partial<
   Readonly<{
-    extensionCodec: ExtensionCodecType;
+    extensionCodec: ExtensionCodecType<ContextType>;
     maxDepth: number;
     initialBufferSize: number;
     sortKeys: boolean;
@@ -15,19 +16,24 @@ export type EncodeOptions = Partial<
      */
     forceFloat32: boolean;
   }>
->;
+> &
+  ContextOf<ContextType>;
 
-const defaultEncodeOptions = {};
+const defaultEncodeOptions: EncodeOptions = {};
 
 /**
  * It encodes `value` in the MessagePack format and
  * returns a byte buffer.
  *
  * The returned buffer is a slice of a larger `ArrayBuffer`, so you have to use its `#byteOffset` and `#byteLength` in order to convert it to another typed arrays including NodeJS `Buffer`.
  */
-export function encode(value: unknown, options: EncodeOptions = defaultEncodeOptions): Uint8Array {
-  const encoder = new Encoder(
+export function encode<ContextType>(
+  value: unknown,
+  options: EncodeOptions<SplitUndefined<ContextType>> = defaultEncodeOptions as any,
+): Uint8Array {
+  const encoder = new Encoder<ContextType>(
     options.extensionCodec,
+    (options as typeof options & { context: any }).context,
     options.maxDepth,
     options.initialBufferSize,
     options.sortKeys,
diff --git a/test/ExtensionCodec.test.ts b/test/ExtensionCodec.test.ts
