Split enumDefault into scalar and data enums

Author: lorisleiva

packages/renderers-rust/README.md

@@ -57,7 +57,8 @@ The Rust renderer provides sensible default traits when generating the various R
 Using the `traitOptions` attribute, you may configure the default traits that will be applied to every Rust type. These default traits can be configured using 4 different attributes:
 
 -   `baseDefaults`: The default traits to implement for all types.
--   `enumDefaults`: The default traits to implement for all enum types, in addition to the `baseDefaults` traits.
+-   `dataEnumDefaults`: The default traits to implement for all data enum types, in addition to the `baseDefaults` traits. Data enums are enums with at least one non-unit variant — e.g. `pub enum Command { Write(String), Quit }`.
+-   `scalarEnumDefaults`: The default traits to implement for all scalar enum types, in addition to the `baseDefaults` traits. Scalar enums are enums with unit variants only — e.g. `pub enum Feedback { Good, Bad }`.
 -   `structDefaults`: The default traits to implement for all struct types, in addition to the `baseDefaults` traits.
 -   `aliasDefaults`: The default traits to implement for all type aliases, in addition to the `baseDefaults` traits.
 
@@ -67,7 +68,8 @@ Note that you must provide the fully qualified name of the traits you provide (e
 const traitOptions = {
     aliasDefaults: [],
     baseDefaults: ['borsh::BorshSerialize', 'borsh::BorshDeserialize', 'Clone', 'Debug', 'Eq', 'PartialEq'],
-    enumDefaults: ['Copy', 'PartialOrd', 'Hash', 'num_derive::FromPrimitive'],
+    dataEnumDefaults: [],
+    scalarEnumDefaults: ['Copy', 'PartialOrd', 'Hash', 'num_derive::FromPrimitive'],
     structDefaults: ['serde::Serialize', 'serde::Deserialize'],
 };
 ```

packages/renderers-rust/src/utils/traitOptions.ts

@@ -1,4 +1,4 @@
-import { AccountNode, assertIsNode, camelCase, DefinedTypeNode, isNode } from '@codama/nodes';
+import { AccountNode, assertIsNode, camelCase, DefinedTypeNode, isNode, isScalarEnum } from '@codama/nodes';
 
 import { ImportMap } from '../ImportMap';
 
@@ -7,8 +7,11 @@ export type TraitOptions = {
     aliasDefaults?: string[];
     /** The default traits to implement for all types. */
     baseDefaults?: string[];
-    /** The default traits to implement for enums only — on top of the base defaults. */
-    enumDefaults?: string[];
+    /**
+     * The default traits to implement for data enums only — on top of the base defaults.
+     * Data enums are enums with at least one non-unit variant.
+     */
+    dataEnumDefaults?: string[];
     /**
      * The mapping of feature flags to traits.
      * For each entry, the traits will be rendered within a
@@ -17,6 +20,11 @@ export type TraitOptions = {
     featureFlags?: Record<string, string[]>;
     /** The complete trait overrides of specific types. */
     overrides?: Record<string, string[]>;
+    /**
+     * The default traits to implement for scalar enums only — on top of the base defaults.
+     * Scalar enums are enums with no variants or only unit variants.
+     */
+    scalarEnumDefaults?: string[];
     /** The default traits to implement for structs only — on top of the base defaults. */
     structDefaults?: string[];
     /** Whether or not to use the fully qualified name for traits, instead of importing them. */
@@ -26,9 +34,10 @@ export type TraitOptions = {
 export const DEFAULT_TRAIT_OPTIONS: Required<TraitOptions> = {
     aliasDefaults: [],
     baseDefaults: ['borsh::BorshSerialize', 'borsh::BorshDeserialize', 'Clone', 'Debug', 'Eq', 'PartialEq'],
-    enumDefaults: ['Copy', 'PartialOrd', 'Hash', 'num_derive::FromPrimitive'],
+    dataEnumDefaults: [],
     featureFlags: { serde: ['serde::Serialize', 'serde::Deserialize'] },
     overrides: {},
+    scalarEnumDefaults: ['Copy', 'PartialOrd', 'Hash', 'num_derive::FromPrimitive'],
     structDefaults: ['serde::Serialize', 'serde::Deserialize'],
     useFullyQualifiedName: false,
 };
@@ -79,22 +88,29 @@ export function getTraitsFromNode(
     return { imports, render: traitLines.join('') };
 }
 
-function getNodeType(node: AccountNode | DefinedTypeNode): 'alias' | 'enum' | 'struct' {
+function getNodeType(node: AccountNode | DefinedTypeNode): 'alias' | 'dataEnum' | 'scalarEnum' | 'struct' {
     if (isNode(node, 'accountNode')) return 'struct';
     if (isNode(node.type, 'structTypeNode')) return 'struct';
-    if (isNode(node.type, 'enumTypeNode')) return 'enum';
+    if (isNode(node.type, 'enumTypeNode')) {
+        return isScalarEnum(node.type) ? 'scalarEnum' : 'dataEnum';
+    }
     return 'alias';
 }
 
 function getDefaultTraits(
-    nodeType: 'alias' | 'enum' | 'struct',
-    options: Pick<Required<TraitOptions>, 'aliasDefaults' | 'baseDefaults' | 'enumDefaults' | 'structDefaults'>,
+    nodeType: 'alias' | 'dataEnum' | 'scalarEnum' | 'struct',
+    options: Pick<
+        Required<TraitOptions>,
+        'aliasDefaults' | 'baseDefaults' | 'dataEnumDefaults' | 'scalarEnumDefaults' | 'structDefaults'
+    >,
 ): string[] {
     switch (nodeType) {
         case 'alias':
             return [...options.baseDefaults, ...options.aliasDefaults];
-        case 'enum':
-            return [...options.baseDefaults, ...options.enumDefaults];
+        case 'dataEnum':
+            return [...options.baseDefaults, ...options.dataEnumDefaults];
+        case 'scalarEnum':
+            return [...options.baseDefaults, ...options.scalarEnumDefaults];
         case 'struct':
         default:
             return [...options.baseDefaults, ...options.structDefaults];

packages/renderers-rust/test/utils/traitOptions.test.ts

@@ -2,6 +2,7 @@ import {
     accountNode,
     definedTypeNode,
     enumEmptyVariantTypeNode,
+    enumStructVariantTypeNode,
     enumTypeNode,
     numberTypeNode,
     structFieldTypeNode,
@@ -12,7 +13,30 @@ import { describe, expect, test } from 'vitest';
 import { getTraitsFromNode, TraitOptions } from '../../src/utils';
 
 describe('default values', () => {
-    test('it defaults to a set of traits for enums', () => {
+    test('it defaults to a set of traits for data enums', () => {
+        // Given a data enum defined type.
+        const node = definedTypeNode({
+            name: 'Command',
+            type: enumTypeNode([
+                enumStructVariantTypeNode(
+                    'Play',
+                    structTypeNode([structFieldTypeNode({ name: 'guess', type: numberTypeNode('u16') })]),
+                ),
+                enumEmptyVariantTypeNode('Quit'),
+            ]),
+        });
+
+        // When we get the traits from the node using the default options.
+        const { render, imports } = getTraitsFromNode(node);
+
+        // Then we expect the following traits to be rendered.
+        expect(render).toBe(`#[derive(BorshSerialize, BorshDeserialize, Clone, Debug, Eq, PartialEq)]\n`);
+
+        // And the following imports to be used.
+        expect([...imports.imports]).toStrictEqual(['borsh::BorshSerialize', 'borsh::BorshDeserialize']);
+    });
+
+    test('it defaults to a set of traits for scalar enums', () => {
         // Given a scalar enum defined type.
         const node = definedTypeNode({
             name: 'Feedback',
@@ -124,30 +148,55 @@ describe('default values', () => {
 const RESET_OPTIONS: Required<TraitOptions> = {
     aliasDefaults: [],
     baseDefaults: [],
-    enumDefaults: [],
+    dataEnumDefaults: [],
     featureFlags: {},
     overrides: {},
+    scalarEnumDefaults: [],
     structDefaults: [],
     useFullyQualifiedName: false,
 };
 
 describe('base traits', () => {
-    test('it uses both the base and enum traits', () => {
+    test('it uses both the base and data enum traits', () => {
+        // Given a data enum defined type.
+        const node = definedTypeNode({
+            name: 'Command',
+            type: enumTypeNode([
+                enumStructVariantTypeNode(
+                    'Play',
+                    structTypeNode([structFieldTypeNode({ name: 'guess', type: numberTypeNode('u16') })]),
+                ),
+                enumEmptyVariantTypeNode('Quit'),
+            ]),
+        });
+
+        // When we get the traits from the node using custom base and data enum defaults.
+        const { render } = getTraitsFromNode(node, {
+            ...RESET_OPTIONS,
+            baseDefaults: ['MyBaseTrait'],
+            dataEnumDefaults: ['MyDataEnumTrait'],
+        });
+
+        // Then we expect both the base and data enum traits to be rendered.
+        expect(render).toBe(`#[derive(MyBaseTrait, MyDataEnumTrait)]\n`);
+    });
+
+    test('it uses both the base and scalar enum traits', () => {
         // Given a scalar enum defined type.
         const node = definedTypeNode({
             name: 'Feedback',
             type: enumTypeNode([enumEmptyVariantTypeNode('Good'), enumEmptyVariantTypeNode('Bad')]),
         });
 
-        // When we get the traits from the node using custom base and enum defaults.
+        // When we get the traits from the node using custom base and scalar enum defaults.
         const { render } = getTraitsFromNode(node, {
             ...RESET_OPTIONS,
             baseDefaults: ['MyBaseTrait'],
-            enumDefaults: ['MyEnumTrait'],
+            scalarEnumDefaults: ['MyScalarEnumTrait'],
         });
 
-        // Then we expect both the base and enum traits to be rendered.
-        expect(render).toBe(`#[derive(MyBaseTrait, MyEnumTrait)]\n`);
+        // Then we expect both the base and scalar enum traits to be rendered.
+        expect(render).toBe(`#[derive(MyBaseTrait, MyScalarEnumTrait)]\n`);
     });
 
     test('it uses both the base and struct traits', () => {
@@ -197,23 +246,23 @@ describe('base traits', () => {
         });
 
         // When we get the traits from the node such that:
-        // - We provide custom base and enum defaults.
+        // - We provide custom base and scalar enum defaults.
         // - We provide custom feature flags for traits in both categories.
         const { render } = getTraitsFromNode(node, {
             ...RESET_OPTIONS,
             baseDefaults: ['MyBaseTrait', 'MyNonFeatureTrait'],
-            enumDefaults: ['MyEnumTrait'],
             featureFlags: {
                 base: ['MyBaseTrait'],
-                enum: ['MyEnumTrait'],
+                enum: ['MyScalarEnumTrait'],
             },
+            scalarEnumDefaults: ['MyScalarEnumTrait'],
         });
 
         // Then we expect both the base and enum traits to be rendered as separate feature flags.
         expect(render).toBe(
             `#[derive(MyNonFeatureTrait)]\n` +
                 `#[cfg(feature = "base", derive(MyBaseTrait))]\n` +
-                `#[cfg(feature = "enum", derive(MyEnumTrait))]\n`,
+                `#[cfg(feature = "enum", derive(MyScalarEnumTrait))]\n`,
         );
     });
 
@@ -229,16 +278,16 @@ describe('base traits', () => {
         const { render } = getTraitsFromNode(node, {
             ...RESET_OPTIONS,
             baseDefaults: ['MyBaseTrait'],
-            enumDefaults: ['MyEnumTrait'],
             featureFlags: {
                 base: ['MyBaseTrait'],
-                enum: ['MyEnumTrait'],
+                enum: ['MyScalarEnumTrait'],
             },
+            scalarEnumDefaults: ['MyScalarEnumTrait'],
         });
 
         // Then we expect the following traits to be rendered.
         expect(render).toBe(
-            `#[cfg(feature = "base", derive(MyBaseTrait))]\n#[cfg(feature = "enum", derive(MyEnumTrait))]\n`,
+            `#[cfg(feature = "base", derive(MyBaseTrait))]\n#[cfg(feature = "enum", derive(MyScalarEnumTrait))]\n`,
         );
     });
 });
@@ -257,8 +306,8 @@ describe('overridden traits', () => {
         const { render } = getTraitsFromNode(node, {
             ...RESET_OPTIONS,
             baseDefaults: ['MyBaseTrait'],
-            enumDefaults: ['MyEnumTrait'],
             overrides: { feedback: ['MyFeedbackTrait'] },
+            scalarEnumDefaults: ['MyScalarEnumTrait'],
         });
 
         // Then we expect only the feedback traits to be rendered.