Merge pull request #499 from streamich/resp-improvements

Adds RESP v2 support

src/json-pack/resp/RespEncoder.ts

@@ -1,14 +1,18 @@
 import {Writer} from '../../util/buffers/Writer';
 import {RESP} from './constants';
 import {utf8Size} from '../../util/strings/utf8';
-import {RespAttributes, RespPush} from './extensions';
+import {RespAttributes, RespPush, RespVerbatimString} from './extensions';
+import {JsonPackExtension} from '../JsonPackExtension';
 import type {IWriter, IWriterGrowable} from '../../util/buffers';
 import type {BinaryJsonEncoder, StreamingBinaryJsonEncoder, TlvBinaryJsonEncoder} from '../types';
 import type {Slice} from '../../util/buffers/Slice';
 
 const REG_RN = /[\r\n]/;
 const isSafeInteger = Number.isSafeInteger;
 
+/**
+ * Implements RESP3 encoding.
+ */
 export class RespEncoder<W extends IWriter & IWriterGrowable = IWriter & IWriterGrowable>
   implements BinaryJsonEncoder, StreamingBinaryJsonEncoder, TlvBinaryJsonEncoder
 {
@@ -38,8 +42,11 @@ export class RespEncoder<W extends IWriter & IWriterGrowable = IWriter & IWriter
         if (value instanceof Uint8Array) return this.writeBin(value);
         if (value instanceof Error) return this.writeErr(value.message);
         if (value instanceof Set) return this.writeSet(value);
-        if (value instanceof RespPush) return this.writePush(value.val);
-        if (value instanceof RespAttributes) return this.writeAttr(value.val);
+        if (value instanceof JsonPackExtension) {
+          if (value instanceof RespPush) return this.writePush(value.val);
+          if (value instanceof RespVerbatimString) return this.writeVerbatimStr('txt', value.val);
+          if (value instanceof RespAttributes) return this.writeAttr(value.val);
+        }
         return this.writeObj(value as Record<string, unknown>);
       }
       case 'undefined':
@@ -52,31 +59,23 @@ export class RespEncoder<W extends IWriter & IWriterGrowable = IWriter & IWriter
   }
 
   protected writeLength(length: number): void {
-    let digits = 1;
-    if (length < 10000) {
-      if (length < 100) {
-        if (length < 10) digits = 1;
-        else digits = 2;
-      } else {
-        if (length < 1000) digits = 3;
-        else digits = 4;
-      }
-    } else if (length < 100000000) {
-      if (length < 1000000) {
-        if (length < 100000) digits = 5;
-        else digits = 6;
-      } else {
-        if (length < 10000000) digits = 7;
-        else digits = 8;
-      }
-    } else {
-      let pow = 10;
-      while (length >= pow) {
-        digits++;
-        pow *= 10;
+    const writer = this.writer;
+    if (length < 100) {
+      if (length < 10) {
+        writer.u8(length + 48);
+        return;
       }
+      const octet1 = length % 10;
+      const octet2 = (length - octet1) / 10;
+      writer.u16(((octet2 + 48) << 8) + octet1 + 48);
+      return;
+    }
+    let digits = 1;
+    let pow = 10;
+    while (length >= pow) {
+      digits++;
+      pow *= 10;
     }
-    const writer = this.writer;
     writer.ensureCapacity(digits);
     const uint8 = writer.uint8;
     const x = writer.x;
@@ -248,6 +247,13 @@ export class RespEncoder<W extends IWriter & IWriterGrowable = IWriter & IWriter
     writer.u16(RESP.RN); // \r\n
   }
 
+  public writeSimpleStrAscii(str: string): void {
+    const writer = this.writer;
+    writer.u8(RESP.STR_SIMPLE); // +
+    writer.ascii(str);
+    writer.u16(RESP.RN); // \r\n
+  }
+
   public writeBulkStr(str: string): void {
     const writer = this.writer;
     const size = utf8Size(str);

src/json-pack/resp/RespEncoderLegacy.ts

@@ -0,0 +1,96 @@
+import {RESP} from './constants';
+import {RespAttributes, RespPush, RespVerbatimString} from './extensions';
+import {JsonPackExtension} from '../JsonPackExtension';
+import {RespEncoder} from './RespEncoder';
+import type {IWriter, IWriterGrowable} from '../../util/buffers';
+
+const REG_RN = /[\r\n]/;
+const isSafeInteger = Number.isSafeInteger;
+
+/**
+ * Implements RESP v2 encoding.
+ */
+export class RespEncoderLegacy<W extends IWriter & IWriterGrowable = IWriter & IWriterGrowable> extends RespEncoder<W> {
+  public writeAny(value: unknown): void {
+    switch (typeof value) {
+      case 'number':
+        return this.writeNumber(value as number);
+      case 'string':
+        return this.writeStr(value);
+      case 'boolean':
+        return this.writeSimpleStr(value ? 'TRUE' : 'FALSE');
+      case 'object': {
+        if (!value) return this.writeNull();
+        if (value instanceof Array) return this.writeArr(value);
+        if (value instanceof Uint8Array) return this.writeBin(value);
+        if (value instanceof Error) return this.writeErr(value.message);
+        if (value instanceof Set) return this.writeSet(value);
+        if (value instanceof JsonPackExtension) {
+          if (value instanceof RespPush) return this.writeArr(value.val);
+          if (value instanceof RespVerbatimString) return this.writeStr(value.val);
+          if (value instanceof RespAttributes) return this.writeObj(value.val);
+        }
+        return this.writeObj(value as Record<string, unknown>);
+      }
+      case 'undefined':
+        return this.writeUndef();
+      case 'bigint':
+        return this.writeSimpleStrAscii(value + '');
+      default:
+        return this.writeUnknown(value);
+    }
+  }
+
+  public writeNumber(num: number): void {
+    if (isSafeInteger(num)) this.writeInteger(num);
+    else this.writeSimpleStrAscii(num + '');
+  }
+
+  public writeStr(str: string): void {
+    const length = str.length;
+    if (length < 64 && !REG_RN.test(str)) this.writeSimpleStr(str);
+    else this.writeBulkStr(str);
+  }
+
+  public writeNull(): void {
+    this.writeNullArr();
+  }
+
+  public writeErr(str: string): void {
+    if (str.length < 64 && !REG_RN.test(str)) this.writeSimpleErr(str);
+    else this.writeBulkStr(str);
+  }
+
+  public writeSet(set: Set<unknown>): void {
+    this.writeArr([...set]);
+  }
+
+  public writeArr(arr: unknown[]): void {
+    const writer = this.writer;
+    const length = arr.length;
+    writer.u8(RESP.ARR); // *
+    this.writeLength(length);
+    writer.u16(RESP.RN); // \r\n
+    for (let i = 0; i < length; i++) {
+      const val = arr[i];
+      if (val === null) this.writeNullStr();
+      else this.writeAny(val);
+    }
+  }
+
+  public writeObj(obj: Record<string, unknown>): void {
+    const writer = this.writer;
+    const keys = Object.keys(obj);
+    const length = keys.length;
+    writer.u8(RESP.ARR); // %
+    this.writeLength(length << 1);
+    writer.u16(RESP.RN); // \r\n
+    for (let i = 0; i < length; i++) {
+      const key = keys[i];
+      this.writeStr(key);
+      const val = obj[key];
+      if (val === null) this.writeNullStr();
+      else this.writeAny(val);
+    }
+  }
+}

src/json-pack/resp/__tests__/RespEncoder.spec.ts

@@ -1,5 +1,6 @@
 import {bufferToUint8Array} from '../../../util/buffers/bufferToUint8Array';
 import {RespEncoder} from '../RespEncoder';
+import {RespVerbatimString} from '../extensions';
 const Parser = require('redis-parser');
 
 const parse = (uint8: Uint8Array): unknown => {
@@ -76,6 +77,12 @@ describe('strings', () => {
       const encoded = encoder.writer.flush();
       expect(toStr(encoded)).toBe('=8\r\ntxt:asdf\r\n');
     });
+
+    test('can encode verbatim string using RespVerbatimString', () => {
+      const encoder = new RespEncoder();
+      const encoded = encoder.encode(new RespVerbatimString('asdf'));
+      expect(toStr(encoded)).toBe('=8\r\ntxt:asdf\r\n');
+    });
   });
 });
 

src/json-pack/resp/__tests__/RespEncoderLegacy.spec.ts

@@ -0,0 +1,52 @@
+import {RespEncoderLegacy} from '../RespEncoderLegacy';
+
+const encode = (value: unknown): string => {
+  const encoder = new RespEncoderLegacy();
+  const encoded = encoder.encode(value);
+  return Buffer.from(encoded).toString();
+};
+
+test('can encode simple strings', () => {
+  expect(encode('')).toBe('+\r\n');
+  expect(encode('asdf')).toBe('+asdf\r\n');
+});
+
+test('can encode simple errors', () => {
+  expect(encode(new Error('asdf'))).toBe('-asdf\r\n');
+});
+
+test('can encode integers', () => {
+  expect(encode(0)).toBe(':0\r\n');
+  expect(encode(123)).toBe(':123\r\n');
+  expect(encode(-422469777)).toBe(':-422469777\r\n');
+});
+
+test('can encode bulk strings', () => {
+  expect(encode('ab\nc')).toBe('$4\r\nab\nc\r\n');
+  expect(encode(new Uint8Array([65]))).toBe('$1\r\nA\r\n');
+});
+
+test('can encode arrays', () => {
+  expect(encode(['a', 1])).toBe('*2\r\n+a\r\n:1\r\n');
+});
+
+test('encodes null as nullable array', () => {
+  expect(encode(null)).toBe('*-1\r\n');
+});
+
+test('encodes null in nested structure as nullable string', () => {
+  expect(encode(['a', 'b', null])).toBe('*3\r\n+a\r\n+b\r\n$-1\r\n');
+});
+
+test('encodes booleans as strings', () => {
+  expect(encode(true)).toBe('+TRUE\r\n');
+  expect(encode(false)).toBe('+FALSE\r\n');
+});
+
+test('encodes floats as strings', () => {
+  expect(encode(1.23)).toBe('+1.23\r\n');
+});
+
+test('encodes objects as 2-tuple arrays', () => {
+  expect(encode({foo: 'bar'})).toBe('*2\r\n+foo\r\n+bar\r\n');
+});

src/json-pack/resp/extensions.ts

@@ -11,3 +11,9 @@ export class RespAttributes extends JsonPackExtension<Record<string, unknown>> {
     super(2, val);
   }
 }
+
+export class RespVerbatimString extends JsonPackExtension<string> {
+  constructor(public readonly val: string) {
+    super(3, val);
+  }
+}

src/json-pack/resp/index.ts

@@ -1,4 +1,6 @@
 export * from './constants';
 export * from './extensions';
 export * from './RespEncoder';
+export * from './RespEncoderLegacy';
 export * from './RespDecoder';
+export * from './RespStreamingDecoder';