refactor(experimental): support number and symbol discriminator value…

…s for getDataEnumCodec
solana-labs · Mar 27, 2024 · f29af61 · f29af61
1 parent bf029dd
commit f29af61
Show file tree

Hide file tree

Showing 6 changed files with 140 additions and 4 deletions.
diff --git a/.changeset/late-mangos-doubt.md b/.changeset/late-mangos-doubt.md
@@ -0,0 +1,28 @@
+---
+'@solana/codecs-data-structures': patch
+---
+
+DataEnum codecs can now use numbers or symbols as discriminator values
+
+```ts
+const codec = getDataEnumCodec([
+  [1, getStructCodec([[['one', u32]]])]
+  [2, getStructCodec([[['two', u32]]])]
+]);
+
+codec.encode({ __kind: 1, one: 42 });
+codec.encode({ __kind: 2, two: 42 });
+```
+
+This means you can also use enum values as discriminators, like so:
+
+```ts
+enum Event { Click, KeyPress }
+const codec = getDataEnumCodec([
+  [Event.Click, getStructCodec([[['x', u32], ['y', u32]]])],
+  [Event.KeyPress, getStructCodec([[['key', u32]]])]
+]);
+
+codec.encode({ __kind: Event.Click, x: 1, y: 2 });
+codec.encode({ __kind: Event.KeyPress, key: 3 });
+```
diff --git a/packages/codecs-data-structures/README.md b/packages/codecs-data-structures/README.md
@@ -339,6 +339,25 @@ messageCodec.encode({ message: 'Write', fields: ['Hi'] });
 messageCodec.encode({ message: 'Move', x: 5, y: 6 });
 ```
 
+Note that, the discriminator value of a variant may also be a `number`, `symbol` or a JavaScript `enum`. For instance, the following is also valid:
+
+```ts
+enum Message {
+    Quit,
+    Write,
+    Move,
+}
+const messageCodec = getDataEnumCodec([
+    [Message.Quit, getUnitCodec()],
+    [Message.Write, getStructCodec([...])],
+    [Message.Move, getStructCodec([...])],
+]);
+
+codec.encode({ __kind: Message.Quit });
+codec.encode({ __kind: Message.Write, fields: ['Hi'] });
+codec.encode({ __kind: Message.Move, x: 5, y: 6 });
+```
+
 Finally, note that separate `getDataEnumEncoder` and `getDataEnumDecoder` functions are available.
 
 ```ts

diff --git a/packages/codecs-data-structures/src/__tests__/data-enum-test.ts b/packages/codecs-data-structures/src/__tests__/data-enum-test.ts
@@ -150,6 +150,34 @@ describe('getDataEnumCodec', () => {
         expect(codec.read(b('012a000000'), 0)).toStrictEqual([{ size: 'large', value: 42 }, 5]);
     });
 
+    it('encodes data enums with number discriminator values', () => {
+        const codec = dataEnum([
+            [1, struct([['one', u8()]])],
+            [2, struct([['two', u32()]])],
+        ]);
+        expect(codec.encode({ __kind: 1, one: 42 })).toStrictEqual(b('002a'));
+        expect(codec.read(b('002a'), 0)).toStrictEqual([{ __kind: 1, one: 42 }, 2]);
+    });
+
+    it('encodes data enums with enum discriminator values', () => {
+        enum Event {
+            Click,
+            KeyPress,
+        }
+        const codec = dataEnum([
+            [
+                Event.Click,
+                struct([
+                    ['x', u8()],
+                    ['y', u8()],
+                ]),
+            ],
+            [Event.KeyPress, struct([['key', u32()]])],
+        ]);
+        expect(codec.encode({ __kind: Event.Click, x: 1, y: 2 })).toStrictEqual(b('000102'));
+        expect(codec.read(b('000102'), 0)).toStrictEqual([{ __kind: Event.Click, x: 1, y: 2 }, 3]);
+    });
+
     it('has the right sizes', () => {
         const webEvent = dataEnum(getWebEvent());
         expect(isVariableSize(webEvent)).toBe(true);

diff --git a/packages/codecs-data-structures/src/__typetests__/data-enum-typetest.ts b/packages/codecs-data-structures/src/__typetests__/data-enum-typetest.ts
@@ -25,6 +25,26 @@ import { getUnitCodec } from '../unit';
             { discriminator: 'myType' },
         ) satisfies Encoder<{ myType: 'A'; value: string } | { myType: 'B'; x: number; y: number }>;
     }
+
+    // It can use numbers as discriminator values.
+    {
+        getDataEnumEncoder([
+            [1, {} as Encoder<{ value: string }>],
+            [2, {} as Encoder<{ x: number; y: number }>],
+        ]) satisfies Encoder<{ __kind: 1; value: string } | { __kind: 2; x: number; y: number }>;
+    }
+
+    // It can use enums as discriminator values.
+    {
+        const enum Event {
+            Click,
+            KeyPress,
+        }
+        getDataEnumEncoder([
+            [Event.Click, {} as Encoder<{ x: number; y: number }>],
+            [Event.KeyPress, {} as Encoder<{ key: string }>],
+        ]) satisfies Encoder<{ __kind: Event.Click; x: number; y: number } | { __kind: Event.KeyPress; key: string }>;
+    }
 }
 
 // [DESCRIBE] getDataEnumDecoder.
@@ -47,6 +67,26 @@ import { getUnitCodec } from '../unit';
             { discriminator: 'myType' },
         ) satisfies Decoder<{ myType: 'A'; value: string } | { myType: 'B'; x: number; y: number }>;
     }
+
+    // It can use numbers as discriminator values.
+    {
+        getDataEnumDecoder([
+            [1, {} as Decoder<{ value: string }>],
+            [2, {} as Decoder<{ x: number; y: number }>],
+        ]) satisfies Decoder<{ __kind: 1; value: string } | { __kind: 2; x: number; y: number }>;
+    }
+
+    // It can use enums as discriminator values.
+    {
+        const enum Event {
+            Click,
+            KeyPress,
+        }
+        getDataEnumDecoder([
+            [Event.Click, {} as Decoder<{ x: number; y: number }>],
+            [Event.KeyPress, {} as Decoder<{ key: string }>],
+        ]) satisfies Decoder<{ __kind: Event.Click; x: number; y: number } | { __kind: Event.KeyPress; key: string }>;
+    }
 }
 
 // [DESCRIBE] getDataEnumCodec.
@@ -70,6 +110,26 @@ import { getUnitCodec } from '../unit';
         ) satisfies Codec<{ myType: 'A'; value: string } | { myType: 'B'; x: number; y: number }>;
     }
 
+    // It can use numbers as discriminator values.
+    {
+        getDataEnumCodec([
+            [1, {} as Codec<{ value: string }>],
+            [2, {} as Codec<{ x: number; y: number }>],
+        ]) satisfies Codec<{ __kind: 1; value: string } | { __kind: 2; x: number; y: number }>;
+    }
+
+    // It can use enums as discriminator values.
+    {
+        const enum Event {
+            Click,
+            KeyPress,
+        }
+        getDataEnumCodec([
+            [Event.Click, {} as Codec<{ x: number; y: number }>],
+            [Event.KeyPress, {} as Codec<{ key: string }>],
+        ]) satisfies Codec<{ __kind: Event.Click; x: number; y: number } | { __kind: Event.KeyPress; key: string }>;
+    }
+
     // It can infer complex data enum types from provided variants.
     {
         getDataEnumCodec(

diff --git a/packages/codecs-data-structures/src/data-enum.ts b/packages/codecs-data-structures/src/data-enum.ts
@@ -86,7 +86,8 @@ export type DataEnumCodecConfig<
     size?: TDiscriminatorSize;
 };
 
-type Variants<T> = readonly (readonly [string, T])[];
+type DiscriminatorValue = number | string | symbol;
+type Variants<T> = readonly (readonly [DiscriminatorValue, T])[];
 type ArrayIndices<T extends readonly unknown[]> = Exclude<Partial<T>['length'], T['length']> & number;
 
 type GetEncoderTypeFromVariants<
@@ -239,7 +240,7 @@ function getDataEnumMaxSize<const TVariants extends Variants<Decoder<any> | Enco
 
 function getVariantDiscriminator<const TVariants extends Variants<Decoder<any> | Encoder<any>>>(
     variants: TVariants,
-    discriminatorValue: string,
+    discriminatorValue: DiscriminatorValue,
 ) {
     const discriminator = variants.findIndex(([key]) => discriminatorValue === key);
     if (discriminator < 0) {

diff --git a/packages/errors/src/context.ts b/packages/errors/src/context.ts
@@ -281,8 +281,8 @@ export type SolanaErrorContext = DefaultUnspecifiedErrorContextToUndefined<
             expected: number;
         };
         [SOLANA_ERROR__CODECS__INVALID_DATA_ENUM_VARIANT]: {
-            value: string;
-            variants: string[];
+            value: number | string | symbol;
+            variants: (number | string | symbol)[];
         };
         [SOLANA_ERROR__CODECS__INVALID_NUMBER_OF_ITEMS]: {
             actual: bigint | number;