refactor(experimental): add getLiteralUnionCodec to @solana/codecs-da…

…ta-structures (#2394)

This PR adds a new `getLiteralUnionCodec` that functions similarly to `getEnumCodec` but uses TypeScript unions to describe all possible values.

```ts
const codec = getLiteralUnionCodec(['left', 'right', 'up', 'down']);
// ^? FixedSizeCodec<"left" | "right" | "up" | "down">

const bytes = codec.encode('left'); // 0x00
const value = codec.decode(bytes); // 'left'
```

Fixes #2296

.changeset/tender-turtles-bake.md

@@ -0,0 +1,14 @@
+---
+'@solana/codecs-data-structures': patch
+'@solana/errors': patch
+---
+
+Added a new `getLiteralUnionCodec`
+
+```ts
+const codec = getLiteralUnionCodec(['left', 'right', 'up', 'down']);
+// ^? FixedSizeCodec<"left" | "right" | "up" | "down">
+
+const bytes = codec.encode('left'); // 0x00
+const value = codec.decode(bytes); // 'left'
+```

packages/codecs-data-structures/README.md

@@ -365,6 +365,40 @@ const bytes = getDiscriminatedUnionEncoder(variantEncoders).encode({ __kind: 'Qu
 const message = getDiscriminatedUnionDecoder(variantDecoders).decode(bytes);
 ```
 
+## Literal union codec
+
+The `getLiteralUnionCodec` function works similarly to the `getUnionCodec` function but does not require a JavaScript `enum` to exist.
+
+It accepts an array of literal values — such as `string`, `number`, `boolean`, etc. — and returns a codec that encodes and decodes such values using by using their index in the array. It uses TypeScript unions to represent all the possible values.
+
+```ts
+const codec = getLiteralUnionCodec(['left', 'right', 'up', 'down']);
+// ^? FixedSizeCodec<"left" | "right" | "up" | "down">
+
+const bytes = codec.encode('left'); // 0x00
+const value = codec.decode(bytes); // 'left'
+```
+
+As you can see, it uses a `u8` number by default to store the index of the value. However, you may provide a number codec as the `size` option of the `getLiteralUnionCodec` function to customise that behaviour.
+
+```ts
+const codec = getLiteralUnionCodec(['left', 'right', 'up', 'down'], {
+    size: getU32Codec(),
+});
+
+codec.encode('left'); // 0x00000000
+codec.encode('right'); // 0x01000000
+codec.encode('up'); // 0x02000000
+codec.encode('down'); // 0x03000000
+```
+
+Separate `getLiteralUnionEncoder` and `getLiteralUnionDecoder` functions are also available.
+
+```ts
+const bytes = getLiteralUnionEncoder(['left', 'right']).encode('left'); // 0x00
+const value = getLiteralUnionDecoder(['left', 'right']).decode(bytes); // 'left'
+```
+
 ## Boolean codec
 
 The `getBooleanCodec` function returns a `Codec<boolean>` that stores the boolean as `0` or `1` using a `u8` number by default.

packages/codecs-data-structures/src/__tests__/literal-union-test.ts

@@ -0,0 +1,133 @@
+import { assertIsFixedSize, assertIsVariableSize } from '@solana/codecs-core';
+import { getShortU16Codec, getU32Codec } from '@solana/codecs-numbers';
+import {
+    SOLANA_ERROR__CODECS__INVALID_LITERAL_UNION_VARIANT,
+    SOLANA_ERROR__CODECS__LITERAL_UNION_DISCRIMINATOR_OUT_OF_RANGE,
+    SolanaError,
+} from '@solana/errors';
+
+import { getLiteralUnionCodec } from '../literal-union';
+import { b } from './__setup__';
+
+describe('getLiteralUnionCodec', () => {
+    it('encodes string unions', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C']);
+        expect(codec.encode('A')).toStrictEqual(b('00'));
+        expect(codec.encode('B')).toStrictEqual(b('01'));
+        expect(codec.encode('C')).toStrictEqual(b('02'));
+    });
+
+    it('decodes string unions', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C']);
+        expect(codec.decode(b('00'))).toBe('A');
+        expect(codec.decode(b('01'))).toBe('B');
+        expect(codec.decode(b('02'))).toBe('C');
+    });
+
+    it('encodes number and bigint unions', () => {
+        const codec = getLiteralUnionCodec([1, 2n, 3]);
+        expect(codec.encode(1)).toStrictEqual(b('00'));
+        expect(codec.encode(2n)).toStrictEqual(b('01'));
+        expect(codec.encode(3)).toStrictEqual(b('02'));
+    });
+
+    it('decodes number and bigint unions', () => {
+        const codec = getLiteralUnionCodec([1, 2n, 3]);
+        expect(codec.decode(b('00'))).toBe(1);
+        expect(codec.decode(b('01'))).toBe(2n);
+        expect(codec.decode(b('02'))).toBe(3);
+    });
+
+    it('encodes boolean unions', () => {
+        const codec = getLiteralUnionCodec([true, false]);
+        expect(codec.encode(true)).toStrictEqual(b('00'));
+        expect(codec.encode(false)).toStrictEqual(b('01'));
+    });
+
+    it('decodes boolean unions', () => {
+        const codec = getLiteralUnionCodec([true, false]);
+        expect(codec.decode(b('00'))).toBe(true);
+        expect(codec.decode(b('01'))).toBe(false);
+    });
+
+    it('encodes null and undefined unions', () => {
+        const codec = getLiteralUnionCodec([null, undefined]);
+        expect(codec.encode(null)).toStrictEqual(b('00'));
+        expect(codec.encode(undefined)).toStrictEqual(b('01'));
+    });
+
+    it('decodes null and undefined unions', () => {
+        const codec = getLiteralUnionCodec([null, undefined]);
+        expect(codec.decode(b('00'))).toBeNull();
+        expect(codec.decode(b('01'))).toBeUndefined();
+    });
+
+    it('pushes the offset forward when writing', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C']);
+        expect(codec.write('A', new Uint8Array(10), 6)).toBe(7);
+    });
+
+    it('pushes the offset forward when reading', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C']);
+        expect(codec.read(b('ffff00'), 2)).toStrictEqual(['A', 3]);
+    });
+
+    it('encodes using a custom discriminator size', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C'], {
+            size: getU32Codec(),
+        });
+        expect(codec.encode('A')).toStrictEqual(b('00000000'));
+        expect(codec.encode('B')).toStrictEqual(b('01000000'));
+        expect(codec.encode('C')).toStrictEqual(b('02000000'));
+    });
+
+    it('decodes using a custom discriminator size', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C'], {
+            size: getU32Codec(),
+        });
+        expect(codec.decode(b('00000000'))).toBe('A');
+        expect(codec.decode(b('01000000'))).toBe('B');
+        expect(codec.decode(b('02000000'))).toBe('C');
+    });
+
+    it('throws when provided with an invalid variant', () => {
+        const codec = getLiteralUnionCodec(['one', 2, 3n, false, null]);
+        // @ts-expect-error Expected invalid variant.
+        expect(() => codec.encode('missing')).toThrow(
+            new SolanaError(SOLANA_ERROR__CODECS__INVALID_LITERAL_UNION_VARIANT, {
+                value: 'missing',
+                variants: ['one', 2, 3n, false, null],
+            }),
+        );
+    });
+
+    it('throws when provided with an invalid discriminator', () => {
+        const codec = getLiteralUnionCodec(['one', 2, 3n, false, null]);
+        expect(() => codec.decode(b('05'))).toThrow(
+            new SolanaError(SOLANA_ERROR__CODECS__LITERAL_UNION_DISCRIMINATOR_OUT_OF_RANGE, {
+                discriminator: 5,
+                maxRange: 4,
+                minRange: 0,
+            }),
+        );
+    });
+
+    it('returns the correct default fixed size', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C']);
+        assertIsFixedSize(codec);
+        expect(codec.fixedSize).toBe(1);
+    });
+
+    it('returns the correct custom fixed size', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C'], { size: getU32Codec() });
+        assertIsFixedSize(codec);
+        expect(codec.fixedSize).toBe(4);
+    });
+
+    it('returns the correct custom variable size', () => {
+        const codec = getLiteralUnionCodec(['A', 'B', 'C'], { size: getShortU16Codec() });
+        assertIsVariableSize(codec);
+        expect(codec.getSizeFromValue('A')).toBe(1);
+        expect(codec.maxSize).toBe(3);
+    });
+});

packages/codecs-data-structures/src/__typetests__/literal-union-typetest.ts

@@ -0,0 +1,42 @@
+import {
+    FixedSizeCodec,
+    FixedSizeDecoder,
+    FixedSizeEncoder,
+    VariableSizeCodec,
+    VariableSizeDecoder,
+    VariableSizeEncoder,
+} from '@solana/codecs-core';
+import { getU32Codec, getU32Decoder, getU32Encoder } from '@solana/codecs-numbers';
+
+import { getLiteralUnionCodec, getLiteralUnionDecoder, getLiteralUnionEncoder } from '../literal-union';
+
+{
+    // [getLiteralUnionEncoder]: It knows if the encoder is fixed size or variable size.
+    getLiteralUnionEncoder(['one', 2, 3n]) satisfies FixedSizeEncoder<'one' | 2 | 3n, 1>;
+    getLiteralUnionEncoder(['one', 2, 3n], { size: getU32Encoder() }) satisfies FixedSizeEncoder<'one' | 2 | 3n, 4>;
+    getLiteralUnionEncoder(['one', 2, 3n], { size: {} as VariableSizeEncoder<number> }) satisfies VariableSizeEncoder<
+        'one' | 2 | 3n
+    >;
+}
+
+{
+    // [getLiteralUnionDecoder]: It knows if the decoder is fixed size or variable size.
+    getLiteralUnionDecoder(['one', 2, 3n]) satisfies FixedSizeDecoder<'one' | 2 | 3n, 1>;
+    getLiteralUnionDecoder(['one', 2, 3n], { size: getU32Decoder() }) satisfies FixedSizeDecoder<'one' | 2 | 3n, 4>;
+    getLiteralUnionDecoder(['one', 2, 3n], { size: {} as VariableSizeDecoder<number> }) satisfies VariableSizeDecoder<
+        'one' | 2 | 3n
+    >;
+}
+
+{
+    // [getLiteralUnionCodec]: It knows if the codec is fixed size or variable size.
+    getLiteralUnionCodec(['one', 2, 3n]) satisfies FixedSizeCodec<'one' | 2 | 3n, 'one' | 2 | 3n, 1>;
+    getLiteralUnionCodec(['one', 2, 3n], { size: getU32Codec() }) satisfies FixedSizeCodec<
+        'one' | 2 | 3n,
+        'one' | 2 | 3n,
+        4
+    >;
+    getLiteralUnionCodec(['one', 2, 3n], { size: {} as VariableSizeCodec<number> }) satisfies VariableSizeCodec<
+        'one' | 2 | 3n
+    >;
+}

packages/codecs-data-structures/src/literal-union.ts

@@ -0,0 +1,134 @@
+import {
+    Codec,
+    combineCodec,
+    Decoder,
+    Encoder,
+    FixedSizeCodec,
+    FixedSizeDecoder,
+    FixedSizeEncoder,
+    mapDecoder,
+    mapEncoder,
+    VariableSizeCodec,
+    VariableSizeDecoder,
+    VariableSizeEncoder,
+} from '@solana/codecs-core';
+import {
+    FixedSizeNumberCodec,
+    FixedSizeNumberDecoder,
+    FixedSizeNumberEncoder,
+    getU8Decoder,
+    getU8Encoder,
+    NumberCodec,
+    NumberDecoder,
+    NumberEncoder,
+} from '@solana/codecs-numbers';
+import {
+    SOLANA_ERROR__CODECS__INVALID_LITERAL_UNION_VARIANT,
+    SOLANA_ERROR__CODECS__LITERAL_UNION_DISCRIMINATOR_OUT_OF_RANGE,
+    SolanaError,
+} from '@solana/errors';
+
+/** Defines the config for literal union codecs. */
+export type LiteralUnionCodecConfig<TDiscriminator = NumberCodec | NumberDecoder | NumberEncoder> = {
+    /**
+     * The number codec to use for the literal union discriminator.
+     * @defaultValue u8 discriminator.
+     */
+    size?: TDiscriminator;
+};
+
+type Variant = bigint | boolean | number | string | null | undefined;
+type GetTypeFromVariants<TVariants extends readonly Variant[]> = TVariants[number];
+
+/**
+ * Creates a literal union encoder.
+ *
+ * @param variants - The variant encoders of the literal union.
+ * @param config - A set of config for the encoder.
+ */
+export function getLiteralUnionEncoder<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+): FixedSizeEncoder<GetTypeFromVariants<TVariants>, 1>;
+export function getLiteralUnionEncoder<const TVariants extends readonly Variant[], TSize extends number>(
+    variants: TVariants,
+    config: LiteralUnionCodecConfig<NumberEncoder> & { size: FixedSizeNumberEncoder<TSize> },
+): FixedSizeEncoder<GetTypeFromVariants<TVariants>, TSize>;
+export function getLiteralUnionEncoder<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+    config?: LiteralUnionCodecConfig<NumberEncoder>,
+): VariableSizeEncoder<GetTypeFromVariants<TVariants>>;
+export function getLiteralUnionEncoder<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+    config: LiteralUnionCodecConfig<NumberEncoder> = {},
+): Encoder<GetTypeFromVariants<TVariants>> {
+    const discriminator = config.size ?? getU8Encoder();
+    return mapEncoder(discriminator, variant => {
+        const index = variants.indexOf(variant);
+        if (index < 0) {
+            throw new SolanaError(SOLANA_ERROR__CODECS__INVALID_LITERAL_UNION_VARIANT, {
+                value: variant,
+                variants,
+            });
+        }
+        return index;
+    });
+}
+
+/**
+ * Creates a literal union decoder.
+ *
+ * @param variants - The variant decoders of the literal union.
+ * @param config - A set of config for the decoder.
+ */
+export function getLiteralUnionDecoder<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+): FixedSizeDecoder<GetTypeFromVariants<TVariants>, 1>;
+export function getLiteralUnionDecoder<const TVariants extends readonly Variant[], TSize extends number>(
+    variants: TVariants,
+    config: LiteralUnionCodecConfig<NumberDecoder> & { size: FixedSizeNumberDecoder<TSize> },
+): FixedSizeDecoder<GetTypeFromVariants<TVariants>, TSize>;
+export function getLiteralUnionDecoder<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+    config?: LiteralUnionCodecConfig<NumberDecoder>,
+): VariableSizeDecoder<GetTypeFromVariants<TVariants>>;
+export function getLiteralUnionDecoder<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+    config: LiteralUnionCodecConfig<NumberDecoder> = {},
+): Decoder<GetTypeFromVariants<TVariants>> {
+    const discriminator = config.size ?? getU8Decoder();
+    return mapDecoder(discriminator, (index: bigint | number) => {
+        if (index < 0 || index >= variants.length) {
+            throw new SolanaError(SOLANA_ERROR__CODECS__LITERAL_UNION_DISCRIMINATOR_OUT_OF_RANGE, {
+                discriminator: index,
+                maxRange: variants.length - 1,
+                minRange: 0,
+            });
+        }
+        return variants[Number(index)];
+    });
+}
+
+/**
+ * Creates a literal union codec.
+ *
+ * @param variants - The variant codecs of the literal union.
+ * @param config - A set of config for the codec.
+ */
+
+export function getLiteralUnionCodec<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+): FixedSizeCodec<GetTypeFromVariants<TVariants>, GetTypeFromVariants<TVariants>, 1>;
+export function getLiteralUnionCodec<const TVariants extends readonly Variant[], TSize extends number>(
+    variants: TVariants,
+    config: LiteralUnionCodecConfig<NumberCodec> & { size: FixedSizeNumberCodec<TSize> },
+): FixedSizeCodec<GetTypeFromVariants<TVariants>, GetTypeFromVariants<TVariants>, TSize>;
+export function getLiteralUnionCodec<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+    config?: LiteralUnionCodecConfig<NumberCodec>,
+): VariableSizeCodec<GetTypeFromVariants<TVariants>>;
+export function getLiteralUnionCodec<const TVariants extends readonly Variant[]>(
+    variants: TVariants,
+    config: LiteralUnionCodecConfig<NumberCodec> = {},
+): Codec<GetTypeFromVariants<TVariants>> {
+    return combineCodec(getLiteralUnionEncoder(variants, config), getLiteralUnionDecoder(variants, config));
+}

packages/codecs/README.md

@@ -83,6 +83,7 @@ The `@solana/codecs` package is composed of several smaller packages, each with
     -   [Struct codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#struct-codec).
     -   [Enum codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#enum-codec).
     -   [Discriminated union codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#discriminated-union-codec).
+    -   [Literal union codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#literal-union-codec).
     -   [Boolean codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#boolean-codec).
     -   [Nullable codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#nullable-codec).
     -   [Bytes codec](https://github.com/solana-labs/solana-web3.js/tree/master/packages/codecs-data-structures#bytes-codec).

packages/errors/src/codes.ts

@@ -259,6 +259,8 @@ export const SOLANA_ERROR__CODECS__NUMBER_OUT_OF_RANGE = 8078011 as const;
 export const SOLANA_ERROR__CODECS__INVALID_STRING_FOR_BASE = 8078012 as const;
 export const SOLANA_ERROR__CODECS__EXPECTED_POSITIVE_BYTE_LENGTH = 8078013 as const;
 export const SOLANA_ERROR__CODECS__OFFSET_OUT_OF_RANGE = 8078014 as const;
+export const SOLANA_ERROR__CODECS__INVALID_LITERAL_UNION_VARIANT = 8078015 as const;
+export const SOLANA_ERROR__CODECS__LITERAL_UNION_DISCRIMINATOR_OUT_OF_RANGE = 8078016 as const;
 
 // RPC-related errors.
 // Reserve error codes in the range [8100000-8100999].
@@ -329,8 +331,10 @@ export type SolanaErrorCode =
     | typeof SOLANA_ERROR__CODECS__INVALID_BYTE_LENGTH
     | typeof SOLANA_ERROR__CODECS__INVALID_DISCRIMINATED_UNION_VARIANT
     | typeof SOLANA_ERROR__CODECS__INVALID_ENUM_VARIANT
+    | typeof SOLANA_ERROR__CODECS__INVALID_LITERAL_UNION_VARIANT
     | typeof SOLANA_ERROR__CODECS__INVALID_NUMBER_OF_ITEMS
     | typeof SOLANA_ERROR__CODECS__INVALID_STRING_FOR_BASE
+    | typeof SOLANA_ERROR__CODECS__LITERAL_UNION_DISCRIMINATOR_OUT_OF_RANGE
     | typeof SOLANA_ERROR__CODECS__NUMBER_OUT_OF_RANGE
     | typeof SOLANA_ERROR__CODECS__OFFSET_OUT_OF_RANGE
     | typeof SOLANA_ERROR__CRYPTO__RANDOM_VALUES_FUNCTION_UNIMPLEMENTED