From bf029dd90230405b3d59f70aedd57fc0117b926e Mon Sep 17 00:00:00 2001 From: Loris Leiva Date: Wed, 27 Mar 2024 13:34:21 +0000 Subject: [PATCH] refactor(experimental): support custom discriminator property for getDataEnumCodec (#2380) This PR adds support for custom discriminator properties for `getDataEnumCodecs`. ```ts // Before. const codec = getDataEnumCodec([ ['click', getStructCodec([[['x', u32], ['y', u32]]])], ['keyPress', getStructCodec([[['key', u32]]])] ]); codec.encode({ __kind: 'click', x: 1, y: 2 }); codec.encode({ __kind: 'keyPress', key: 3 }); // After. const codec = getDataEnumCodec([ ['click', getStructCodec([[['x', u32], ['y', u32]]])], ['keyPress', getStructCodec([[['key', u32]]])] ], { discriminator: 'event' }); codec.encode({ event: 'click', x: 1, y: 2 }); codec.encode({ event: 'keyPress', key: 3 }); ``` Note that, to make this work, this PR updates a few exported types such as `GetDataEnumKind` or `GetDataEnumKindContent` which is used by Kinobi. I'll make sure to have a PR ready on Kinobi's side when this is ready to be merged. --- .changeset/violet-brooms-report.md | 15 ++ packages/codecs-data-structures/README.md | 22 ++- .../src/__tests__/data-enum-test.ts | 30 ++-- .../src/__typetests__/data-enum-typetest.ts | 135 ++++++++++++------ .../codecs-data-structures/src/data-enum.ts | 100 ++++++++----- 5 files changed, 211 insertions(+), 91 deletions(-) create mode 100644 .changeset/violet-brooms-report.md diff --git a/.changeset/violet-brooms-report.md b/.changeset/violet-brooms-report.md new file mode 100644 index 000000000000..e9da8d046036 --- /dev/null +++ b/.changeset/violet-brooms-report.md @@ -0,0 +1,15 @@ +--- +'@solana/codecs-data-structures': patch +--- + +DataEnum codecs now support custom discriminator properties + +```ts +const codec = getDataEnumCodec([ + ['click', getStructCodec([[['x', u32], ['y', u32]]])], + ['keyPress', getStructCodec([[['key', u32]]])] +], { discriminator: 'event' }); + +codec.encode({ event: 'click', x: 1, y: 2 }); +codec.encode({ event: 'keyPress', key: 3 }); +``` diff --git a/packages/codecs-data-structures/README.md b/packages/codecs-data-structures/README.md index 420049307763..7e8aa73803e6 100644 --- a/packages/codecs-data-structures/README.md +++ b/packages/codecs-data-structures/README.md @@ -253,7 +253,7 @@ In Rust, enums are powerful data types whose variants can be one of the followin Whilst we do not have such powerful enums in JavaScript, we can emulate them in TypeScript using a union of objects such that each object is differentiated by a specific field. **We call this a data enum**. -We use a special field named `__kind` to distinguish between the different variants of a data enum. Additionally, since all variants are objects, we use a `fields` property to wrap the array of tuple variants. Here is an example. +We use a special field named `__kind` to distinguish between the different variants of a data enum. Additionally, since all variants are objects, we can use a `fields` property to wrap the array of tuple variants. Here is an example. ```ts type Message = @@ -264,7 +264,7 @@ type Message = The `getDataEnumCodec` function helps us encode and decode these data enums. -It requires the name and codec of each variant as a first argument. Similarly to the struct codec, these are defined as an array of variant tuples where the first item is the name of the variant and the second item is its codec. Since empty variants do not have data to encode, they simply use the unit codec — documented below — which does nothing. +It requires the discriminator and codec of each variant as a first argument. Similarly to the struct codec, these are defined as an array of variant tuples where the first item is the discriminator of the variant and the second item is its codec. Since empty variants do not have data to encode, they simply use the unit codec — documented below — which does nothing. Here is how we can create a data enum codec for our previous example. @@ -274,12 +274,12 @@ const messageCodec = getDataEnumCodec([ ['Quit', getUnitCodec()], // Tuple variant. - ['Write', getStructCodec<{ fields: [string] }>([['fields', getTupleCodec([getStringCodec()])]])], + ['Write', getStructCodec([['fields', getTupleCodec([getStringCodec()])]])], // Struct variant. [ 'Move', - getStructCodec<{ x: number; y: number }>([ + getStructCodec([ ['x', getI32Codec()], ['y', getI32Codec()], ]), @@ -327,7 +327,19 @@ u32MessageCodec.encode({ __kind: 'Move', x: 5, y: 6 }); // └------┘ 4-byte discriminator (Index 2). ``` -Separate `getDataEnumEncoder` and `getDataEnumDecoder` functions are also available. +You may also customize the discriminator property — which defaults to `__kind` — by providing the desired property name as the `discriminator` option like so: + +```ts +const messageCodec = getDataEnumCodec([...], { + discriminator: 'message', +}); + +messageCodec.encode({ message: 'Quit' }); +messageCodec.encode({ message: 'Write', fields: ['Hi'] }); +messageCodec.encode({ message: 'Move', x: 5, y: 6 }); +``` + +Finally, note that separate `getDataEnumEncoder` and `getDataEnumDecoder` functions are available. ```ts const bytes = getDataEnumEncoder(variantEncoders).encode({ __kind: 'Quit' }); diff --git a/packages/codecs-data-structures/src/__tests__/data-enum-test.ts b/packages/codecs-data-structures/src/__tests__/data-enum-test.ts index 9a783b02035b..b4070af4a452 100644 --- a/packages/codecs-data-structures/src/__tests__/data-enum-test.ts +++ b/packages/codecs-data-structures/src/__tests__/data-enum-test.ts @@ -124,16 +124,30 @@ describe('getDataEnumCodec', () => { }); it('encodes data enums with different From and To types', () => { - const x = dataEnum(getU64Enum()); - expect(x.encode({ __kind: 'B', value: 2 })).toStrictEqual(b('010200000000000000')); - expect(x.encode({ __kind: 'B', value: 2n })).toStrictEqual(b('010200000000000000')); - expect(x.read(b('010200000000000000'), 0)).toStrictEqual([{ __kind: 'B', value: 2n }, 9]); + const codec = dataEnum(getU64Enum()); + expect(codec.encode({ __kind: 'B', value: 2 })).toStrictEqual(b('010200000000000000')); + expect(codec.encode({ __kind: 'B', value: 2n })).toStrictEqual(b('010200000000000000')); + expect(codec.read(b('010200000000000000'), 0)).toStrictEqual([{ __kind: 'B', value: 2n }, 9]); }); - it('encodes data enums with custom prefix', () => { - const x = dataEnum(getSameSizeVariants(), { size: u32() }); - expect(x.encode({ __kind: 'A', value: 42 })).toStrictEqual(b('000000002a00')); - expect(x.read(b('000000002a00'), 0)).toStrictEqual([{ __kind: 'A', value: 42 }, 6]); + it('encodes data enums with a custom prefix', () => { + const codec = dataEnum(getSameSizeVariants(), { size: u32() }); + expect(codec.encode({ __kind: 'A', value: 42 })).toStrictEqual(b('000000002a00')); + expect(codec.read(b('000000002a00'), 0)).toStrictEqual([{ __kind: 'A', value: 42 }, 6]); + }); + + it('encodes data enums with a custom discriminator property', () => { + const codec = dataEnum( + [ + ['small', struct([['value', u8()]])], + ['large', struct([['value', u32()]])], + ], + { discriminator: 'size' }, + ); + expect(codec.encode({ size: 'small', value: 42 })).toStrictEqual(b('002a')); + expect(codec.read(b('002a'), 0)).toStrictEqual([{ size: 'small', value: 42 }, 2]); + expect(codec.encode({ size: 'large', value: 42 })).toStrictEqual(b('012a000000')); + expect(codec.read(b('012a000000'), 0)).toStrictEqual([{ size: 'large', value: 42 }, 5]); }); it('has the right sizes', () => { diff --git a/packages/codecs-data-structures/src/__typetests__/data-enum-typetest.ts b/packages/codecs-data-structures/src/__typetests__/data-enum-typetest.ts index 79f5861af336..0b843cc626a2 100644 --- a/packages/codecs-data-structures/src/__typetests__/data-enum-typetest.ts +++ b/packages/codecs-data-structures/src/__typetests__/data-enum-typetest.ts @@ -5,58 +5,103 @@ import { getDataEnumCodec, getDataEnumDecoder, getDataEnumEncoder } from '../dat import { getStructCodec } from '../struct'; import { getUnitCodec } from '../unit'; +// [DESCRIBE] getDataEnumEncoder. { - // [getDataEnumEncoder]: It constructs data enums from a list of encoder variants. - getDataEnumEncoder([ - ['A', {} as Encoder<{ value: string }>], - ['B', {} as Encoder<{ x: number; y: number }>], - ]) satisfies Encoder<{ __kind: 'A'; value: string } | { __kind: 'B'; x: number; y: number }>; -} + // It constructs data enums from a list of encoder variants. + { + getDataEnumEncoder([ + ['A', {} as Encoder<{ value: string }>], + ['B', {} as Encoder<{ x: number; y: number }>], + ]) satisfies Encoder<{ __kind: 'A'; value: string } | { __kind: 'B'; x: number; y: number }>; + } -{ - // [getDataEnumDecoder]: It constructs data enums from a list of decoder variants. - getDataEnumDecoder([ - ['A', {} as Decoder<{ value: string }>], - ['B', {} as Decoder<{ x: number; y: number }>], - ]) satisfies Decoder<{ __kind: 'A'; value: string } | { __kind: 'B'; x: number; y: number }>; + // It can use a custom discriminator property. + { + getDataEnumEncoder( + [ + ['A', {} as Encoder<{ value: string }>], + ['B', {} as Encoder<{ x: number; y: number }>], + ], + { discriminator: 'myType' }, + ) satisfies Encoder<{ myType: 'A'; value: string } | { myType: 'B'; x: number; y: number }>; + } } +// [DESCRIBE] getDataEnumDecoder. { - // [getDataEnumCodec]: It constructs data enums from a list of codec variants. - getDataEnumCodec([ - ['A', {} as Codec<{ value: string }>], - ['B', {} as Codec<{ x: number; y: number }>], - ]) satisfies Codec<{ __kind: 'A'; value: string } | { __kind: 'B'; x: number; y: number }>; -} + // It constructs data enums from a list of decoder variants. + { + getDataEnumDecoder([ + ['A', {} as Decoder<{ value: string }>], + ['B', {} as Decoder<{ x: number; y: number }>], + ]) satisfies Decoder<{ __kind: 'A'; value: string } | { __kind: 'B'; x: number; y: number }>; + } -{ - // [getDataEnumCodec]: It can infer complex data enum types from provided variants. - getDataEnumCodec([ - ['PageLoad', {} as Codec], - [ - 'Click', - getStructCodec([ - ['x', {} as Codec], - ['y', {} as Codec], - ]), - ], - ['KeyPress', getStructCodec([['fields', {} as Codec<[string]>]])], - ['PageUnload', {} as Codec], - ]) satisfies Codec< - | { __kind: 'Click'; x: number; y: number } - | { __kind: 'KeyPress'; fields: [string] } - | { __kind: 'PageLoad' } - | { __kind: 'PageUnload' } - >; + // It can use a custom discriminator property. + { + getDataEnumDecoder( + [ + ['A', {} as Decoder<{ value: string }>], + ['B', {} as Decoder<{ x: number; y: number }>], + ], + { discriminator: 'myType' }, + ) satisfies Decoder<{ myType: 'A'; value: string } | { myType: 'B'; x: number; y: number }>; + } } +// [DESCRIBE] getDataEnumCodec. { - // [getDataEnumCodec]: It can infer codec data enum with different from and to types. - getDataEnumCodec([ - ['A', getUnitCodec()], - ['B', getStructCodec([['value', getU64Codec()]])], - ]) satisfies Codec< - { __kind: 'A' } | { __kind: 'B'; value: bigint | number }, - { __kind: 'A' } | { __kind: 'B'; value: bigint } - >; + // It constructs data enums from a list of codec variants. + { + getDataEnumCodec([ + ['A', {} as Codec<{ value: string }>], + ['B', {} as Codec<{ x: number; y: number }>], + ]) satisfies Codec<{ __kind: 'A'; value: string } | { __kind: 'B'; x: number; y: number }>; + } + + // It can use a custom discriminator property. + { + getDataEnumCodec( + [ + ['A', {} as Codec<{ value: string }>], + ['B', {} as Codec<{ x: number; y: number }>], + ], + { discriminator: 'myType' }, + ) satisfies Codec<{ myType: 'A'; value: string } | { myType: 'B'; x: number; y: number }>; + } + + // It can infer complex data enum types from provided variants. + { + getDataEnumCodec( + [ + ['PageLoad', {} as Codec], + [ + 'Click', + getStructCodec([ + ['x', {} as Codec], + ['y', {} as Codec], + ]), + ], + ['KeyPress', getStructCodec([['fields', {} as Codec<[string]>]])], + ['PageUnload', {} as Codec], + ], + { discriminator: 'event' }, + ) satisfies Codec< + | { event: 'Click'; x: number; y: number } + | { event: 'KeyPress'; fields: [string] } + | { event: 'PageLoad' } + | { event: 'PageUnload' } + >; + } + + // It can infer codec data enum with different from and to types. + { + getDataEnumCodec([ + ['A', getUnitCodec()], + ['B', getStructCodec([['value', getU64Codec()]])], + ]) satisfies Codec< + { __kind: 'A' } | { __kind: 'B'; value: bigint | number }, + { __kind: 'A' } | { __kind: 'B'; value: bigint } + >; + } } diff --git a/packages/codecs-data-structures/src/data-enum.ts b/packages/codecs-data-structures/src/data-enum.ts index 6f7715019d72..41eb0f7d737a 100644 --- a/packages/codecs-data-structures/src/data-enum.ts +++ b/packages/codecs-data-structures/src/data-enum.ts @@ -29,7 +29,9 @@ import { DrainOuterGeneric, getMaxSize, maxCodecSizes, sumCodecSizes } from './u * | { __kind: 'click', x: number, y: number }; * ``` */ -export type DataEnum = { __kind: string }; +export type DataEnum = { + [P in TDiscriminatorProperty]: TDiscriminatorValue; +}; /** * Extracts a variant from a data enum. @@ -39,11 +41,15 @@ export type DataEnum = { __kind: string }; * type WebPageEvent = * | { __kind: 'pageview', url: string } * | { __kind: 'click', x: number, y: number }; - * type ClickEvent = GetDataEnumKind; + * type ClickEvent = GetDataEnumKind; * // -> { __kind: 'click', x: number, y: number } * ``` */ -export type GetDataEnumKind = Extract; +export type GetDataEnumKind< + TDataEnum extends DataEnum, + TDiscriminatorProperty extends string, + TDiscriminatorValue extends TDataEnum[TDiscriminatorProperty], +> = Extract>; /** * Extracts a variant from a data enum without its discriminator. @@ -53,41 +59,56 @@ export type GetDataEnumKind = Extract * type WebPageEvent = * | { __kind: 'pageview', url: string } * | { __kind: 'click', x: number, y: number }; - * type ClickEvent = GetDataEnumKindContent; + * type ClickEvent = GetDataEnumKindContent; * // -> { x: number, y: number } * ``` */ -export type GetDataEnumKindContent = Omit< - Extract, - '__kind' ->; +export type GetDataEnumKindContent< + TDataEnum extends DataEnum, + TDiscriminatorProperty extends string, + TDiscriminatorValue extends TDataEnum[TDiscriminatorProperty], +> = Omit, TDiscriminatorProperty>; /** Defines the config for data enum codecs. */ -export type DataEnumCodecConfig = { +export type DataEnumCodecConfig< + TDiscriminatorProperty extends string = '__kind', + TDiscriminatorSize = NumberCodec | NumberDecoder | NumberEncoder, +> = { + /** + * The property name of the discriminator. + * @defaultValue `__kind`. + */ + discriminator?: TDiscriminatorProperty; /** * The codec to use for the enum discriminator prefixing the variant. * @defaultValue u8 prefix. */ - size?: TDiscriminator; + size?: TDiscriminatorSize; }; type Variants = readonly (readonly [string, T])[]; type ArrayIndices = Exclude['length'], T['length']> & number; -type GetEncoderTypeFromVariants>> = DrainOuterGeneric<{ +type GetEncoderTypeFromVariants< + TVariants extends Variants>, + TDiscriminatorProperty extends string, +> = DrainOuterGeneric<{ [I in ArrayIndices]: (TVariants[I][1] extends Encoder ? TFrom extends object ? TFrom : object - : never) & { __kind: TVariants[I][0] }; + : never) & { [P in TDiscriminatorProperty]: TVariants[I][0] }; }>[ArrayIndices]; -type GetDecoderTypeFromVariants>> = DrainOuterGeneric<{ +type GetDecoderTypeFromVariants< + TVariants extends Variants>, + TDiscriminatorProperty extends string, +> = DrainOuterGeneric<{ [I in ArrayIndices]: (TVariants[I][1] extends Decoder ? TTo extends object ? TTo : object - : never) & { __kind: TVariants[I][0] }; + : never) & { [P in TDiscriminatorProperty]: TVariants[I][0] }; }>[ArrayIndices]; /** @@ -96,11 +117,15 @@ type GetDecoderTypeFromVariants>> = Drai * @param variants - The variant encoders of the data enum. * @param config - A set of config for the encoder. */ -export function getDataEnumEncoder>>( +export function getDataEnumEncoder< + const TVariants extends Variants>, + const TDiscriminatorProperty extends string = '__kind', +>( variants: TVariants, - config: DataEnumCodecConfig = {}, -): Encoder> { - type TFrom = GetEncoderTypeFromVariants; + config: DataEnumCodecConfig = {}, +): Encoder> { + type TFrom = GetEncoderTypeFromVariants; + const discriminatorProperty = (config.discriminator ?? '__kind') as TDiscriminatorProperty; const prefix = config.size ?? getU8Encoder(); const fixedSize = getDataEnumFixedSize(variants, prefix); return createEncoder({ @@ -108,7 +133,7 @@ export function getDataEnumEncoder ? { fixedSize } : { getSizeFromValue: (variant: TFrom) => { - const discriminator = getVariantDiscriminator(variants, variant); + const discriminator = getVariantDiscriminator(variants, variant[discriminatorProperty]); const variantEncoder = variants[discriminator][1]; return ( getEncodedSize(discriminator, prefix) + @@ -118,7 +143,7 @@ export function getDataEnumEncoder maxSize: getDataEnumMaxSize(variants, prefix), }), write: (variant: TFrom, bytes, offset) => { - const discriminator = getVariantDiscriminator(variants, variant); + const discriminator = getVariantDiscriminator(variants, variant[discriminatorProperty]); offset = prefix.write(discriminator, bytes, offset); const variantEncoder = variants[discriminator][1]; return variantEncoder.write(variant as TFrom & void, bytes, offset); @@ -132,11 +157,15 @@ export function getDataEnumEncoder * @param variants - The variant decoders of the data enum. * @param config - A set of config for the decoder. */ -export function getDataEnumDecoder>>( +export function getDataEnumDecoder< + const TVariants extends Variants>, + const TDiscriminatorProperty extends string = '__kind', +>( variants: TVariants, - config: DataEnumCodecConfig = {}, -): Decoder> { - type TTo = GetDecoderTypeFromVariants; + config: DataEnumCodecConfig = {}, +): Decoder> { + type TTo = GetDecoderTypeFromVariants; + const discriminatorProperty = config.discriminator ?? '__kind'; const prefix = config.size ?? getU8Decoder(); const fixedSize = getDataEnumFixedSize(variants, prefix); return createDecoder({ @@ -155,7 +184,7 @@ export function getDataEnumDecoder } const [variant, vOffset] = variantField[1].read(bytes, offset); offset = vOffset; - return [{ __kind: variantField[0], ...(variant ?? {}) } as TTo, offset]; + return [{ [discriminatorProperty]: variantField[0], ...(variant ?? {}) } as TTo, offset]; }, }); } @@ -166,17 +195,22 @@ export function getDataEnumDecoder * @param variants - The variant codecs of the data enum. * @param config - A set of config for the codec. */ -export function getDataEnumCodec>>( +export function getDataEnumCodec< + const TVariants extends Variants>, + const TDiscriminatorProperty extends string = '__kind', +>( variants: TVariants, - config: DataEnumCodecConfig = {}, + config: DataEnumCodecConfig = {}, ): Codec< - GetEncoderTypeFromVariants, - GetDecoderTypeFromVariants & GetEncoderTypeFromVariants + GetEncoderTypeFromVariants, + GetDecoderTypeFromVariants & + GetEncoderTypeFromVariants > { return combineCodec( getDataEnumEncoder(variants, config), getDataEnumDecoder(variants, config) as Decoder< - GetDecoderTypeFromVariants & GetEncoderTypeFromVariants + GetDecoderTypeFromVariants & + GetEncoderTypeFromVariants >, ); } @@ -205,12 +239,12 @@ function getDataEnumMaxSize | Enco function getVariantDiscriminator | Encoder>>( variants: TVariants, - variant: { __kind: string }, + discriminatorValue: string, ) { - const discriminator = variants.findIndex(([key]) => variant.__kind === key); + const discriminator = variants.findIndex(([key]) => discriminatorValue === key); if (discriminator < 0) { throw new SolanaError(SOLANA_ERROR__CODECS__INVALID_DATA_ENUM_VARIANT, { - value: variant.__kind, + value: discriminatorValue, variants: variants.map(([key]) => key), }); }