fix: subsequent null and undefined encoding failure (#24)

jacob-ebey · Apr 29, 2024 · 47adfe1 · 47adfe1
1 parent da85f75
commit 47adfe1
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 33 deletions.
diff --git a/src/flatten.ts b/src/flatten.ts
@@ -3,6 +3,7 @@ import {
   NAN,
   NEGATIVE_INFINITY,
   NEGATIVE_ZERO,
+  NULL,
   POSITIVE_INFINITY,
   UNDEFINED,
   TYPE_BIGINT,
@@ -24,9 +25,10 @@ export function flatten(this: ThisEncode, input: unknown): number {
   if (existing) return existing;
 
   if (input === undefined) return UNDEFINED;
+  if (input === null) return NULL;
   if (Number.isNaN(input)) return NAN;
-  if (input === Infinity) return POSITIVE_INFINITY;
-  if (input === -Infinity) return NEGATIVE_INFINITY;
+  if (input === Number.POSITIVE_INFINITY) return POSITIVE_INFINITY;
+  if (input === Number.NEGATIVE_INFINITY) return NEGATIVE_INFINITY;
   if (input === 0 && 1 / input < 0) return NEGATIVE_ZERO;
 
   const index = this.index++;
@@ -53,17 +55,18 @@ function stringify(this: ThisEncode, input: unknown, index: number) {
     case "bigint":
       str[index] = `["${TYPE_BIGINT}","${input}"]`;
       break;
-    case "symbol":
+    case "symbol": {
       const keyFor = Symbol.keyFor(input);
       if (!keyFor)
         throw new Error(
           "Cannot encode symbol unless created with Symbol.for()"
         );
       str[index] = `["${TYPE_SYMBOL}",${JSON.stringify(keyFor)}]`;
       break;
-    case "object":
+    }
+    case "object": {
       if (!input) {
-        str[index] = "null";
+        str[index] = `${NULL}`;
         break;
       }
 
@@ -77,8 +80,9 @@ function stringify(this: ThisEncode, input: unknown, index: number) {
             const [pluginIdentifier, ...rest] = pluginResult;
             str[index] = `[${JSON.stringify(pluginIdentifier)}`;
             if (rest.length > 0) {
-              str[index] +=
-                "," + rest.map((v) => flatten.call(this, v)).join(",");
+              str[index] += `,${rest
+                .map((v) => flatten.call(this, v))
+                .join(",")}`;
             }
             str[index] += "]";
             break;
@@ -93,7 +97,7 @@ function stringify(this: ThisEncode, input: unknown, index: number) {
             result +=
               (i ? "," : "") +
               (i in input ? flatten.call(this, input[i]) : HOLE);
-          str[index] = result + "]";
+          str[index] = `${result}]`;
         } else if (input instanceof Date) {
           str[index] = `["${TYPE_DATE}",${input.getTime()}]`;
         } else if (input instanceof URL) {
@@ -128,6 +132,7 @@ function stringify(this: ThisEncode, input: unknown, index: number) {
         }
       }
       break;
+    }
     default:
       throw new Error("Cannot encode function or unexpected type");
   }

diff --git a/src/turbo-stream.spec.ts b/src/turbo-stream.spec.ts
@@ -2,7 +2,7 @@ import { test } from "node:test";
 import { expect } from "expect";
 
 import { decode, encode } from "./turbo-stream.js";
-import { Deferred } from "./utils.js";
+import { Deferred, type DecodePlugin, type EncodePlugin } from "./utils.js";
 
 async function quickDecode(stream: ReadableStream<Uint8Array>) {
   const decoded = await decode(stream);
@@ -165,6 +165,22 @@ test("should encode and decode promise", async () => {
   await decoded.done;
 });
 
+test("should encode and decode subsequent null from promise in object value", async () => {
+  const input = { root: null, promise: Promise.resolve(null) };
+  const decoded = await decode(encode(input));
+  const value = decoded.value as typeof input;
+  expect(await value.promise).toEqual(await input.promise);
+  await decoded.done;
+});
+
+test("should encode and decode subsequent undefined from promise in object value", async () => {
+  const input = { root: undefined, promise: Promise.resolve(undefined) };
+  const decoded = await decode(encode(input));
+  const value = decoded.value as typeof input;
+  expect(await value.promise).toEqual(await input.promise);
+  await decoded.done;
+});
+
 test("should encode and decode rejected promise", async () => {
   const input = Promise.reject(new Error("foo"));
   const decoded = await decode(encode(input));
@@ -204,33 +220,44 @@ test("should encode and decode set with promises as values", async () => {
   await decoded.done;
 });
 
-test("should encode and decode custom type", async () => {
+test("should encode and decode custom type", async ({ mock }) => {
   class Custom {
+    child: Custom | undefined;
     constructor(public foo: string) {}
   }
   const input = new Custom("bar");
+  input.child = new Custom("baz");
+
+  const decoder = mock.fn<DecodePlugin>((type, foo, child) => {
+    if (type === "Custom") {
+      const value = new Custom(foo as string);
+      value.child = child as Custom | undefined;
+      return { value };
+    }
+  });
+
+  const encoder = mock.fn<EncodePlugin>((value) => {
+    if (value instanceof Custom) {
+      return ["Custom", value.foo, value.child];
+    }
+  });
+
   const decoded = await decode(
     encode(input, {
-      plugins: [
-        (value) => {
-          if (value instanceof Custom) {
-            return ["Custom", value.foo];
-          }
-        },
-      ],
+      plugins: [encoder],
     }),
     {
-      plugins: [
-        (type, foo) => {
-          if (type === "Custom") {
-            return { value: new Custom(foo as string) };
-          }
-        },
-      ],
+      plugins: [decoder],
     }
   );
-  expect(decoded.value).toBeInstanceOf(Custom);
-  expect(decoded.value).toEqual(input);
+  const value = decoded.value as Custom;
+  expect(value).toBeInstanceOf(Custom);
+  expect(value.foo).toEqual(input.foo);
+  expect(value.child).toBeInstanceOf(Custom);
+  expect(value.child?.foo).toEqual(input.child.foo);
+
+  expect(encoder.mock.callCount()).toBe(2);
+  expect(decoder.mock.callCount()).toBe(2);
 });
 
 test("should encode and decode custom type when nested alongside Promise", async () => {

diff --git a/src/unflatten.ts b/src/unflatten.ts
@@ -4,6 +4,7 @@ import {
   NAN,
   NEGATIVE_INFINITY,
   NEGATIVE_ZERO,
+  NULL,
   POSITIVE_INFINITY,
   UNDEFINED,
   TYPE_BIGINT,
@@ -46,6 +47,8 @@ function hydrate(this: ThisDecode, index: number) {
   switch (index) {
     case UNDEFINED:
       return;
+    case NULL:
+      return null;
     case NAN:
       return NaN;
     case POSITIVE_INFINITY:

diff --git a/src/utils.ts b/src/utils.ts
@@ -1,9 +1,10 @@
 export const HOLE = -1;
 export const NAN = -2;
-export const NEGATIVE_INFINITY = -4;
-export const NEGATIVE_ZERO = -5;
-export const POSITIVE_INFINITY = -3;
-export const UNDEFINED = -1;
+export const NEGATIVE_INFINITY = -3;
+export const NEGATIVE_ZERO = -4;
+export const NULL = -5;
+export const POSITIVE_INFINITY = -6;
+export const UNDEFINED = -7;
 
 export const TYPE_BIGINT = "B";
 export const TYPE_DATE = "D";
@@ -55,13 +56,13 @@ export class Deferred<T = unknown> {
 }
 
 export function createLineSplittingTransform() {
-  let decoder = new TextDecoder();
+  const decoder = new TextDecoder();
   let leftover = "";
 
   return new TransformStream({
     transform(chunk, controller) {
-      let str = decoder.decode(chunk, { stream: true });
-      let parts = (leftover + str).split("\n");
+      const str = decoder.decode(chunk, { stream: true });
+      const parts = (leftover + str).split("\n");
 
       // The last part might be a partial line, so keep it for the next chunk.
       leftover = parts.pop() || "";