Use name transformers for enum variant and data enum discriminator (#140

) * Use name transformers for enum variant and data enum discriminator * Add changeset
metaplex-foundation · Jan 5, 2024 · 69a9f3e · 69a9f3e
1 parent 3bfeb8c
commit 69a9f3e
Show file tree

Hide file tree

Showing 11 changed files with 216 additions and 193 deletions.
diff --git a/.changeset/sixty-gorillas-judge.md b/.changeset/sixty-gorillas-judge.md
@@ -0,0 +1,5 @@
+---
+'@metaplex-foundation/kinobi': patch
+---
+
+Use name transformers for enum variant and data enum discriminator
diff --git a/src/renderers/js-experimental/fragments/typeDataEnumHelpers.njk b/src/renderers/js-experimental/fragments/typeDataEnumHelpers.njk
@@ -1,23 +1,23 @@
 // Data Enum Helpers.
 {% for variant in typeNode.variants %}
   {% if variant.kind === 'enumStructVariantTypeNode' %}
-    export function {{ dataEnumFunction }}(kind: '{{ variant.name | pascalCase }}', data: GetDataEnumKindContent<{{ looseName }}, '{{ variant.name | pascalCase }}'>): GetDataEnumKind<{{ looseName }}, '{{ variant.name | pascalCase }}'>;
+    export function {{ dataEnumFunction }}(kind: '{{ getVariant(variant.name) }}', data: GetDataEnumKindContent<{{ looseName }}, '{{ getVariant(variant.name) }}'>): GetDataEnumKind<{{ looseName }}, '{{ getVariant(variant.name) }}'>;
   {% elif variant.kind === 'enumTupleVariantTypeNode' %}
-    export function {{ dataEnumFunction }}(kind: '{{ variant.name | pascalCase }}', data: GetDataEnumKindContent<{{ looseName }}, '{{ variant.name | pascalCase }}'>['fields']): GetDataEnumKind<{{ looseName }}, '{{ variant.name | pascalCase }}'>;
+    export function {{ dataEnumFunction }}(kind: '{{ getVariant(variant.name) }}', data: GetDataEnumKindContent<{{ looseName }}, '{{ getVariant(variant.name) }}'>['fields']): GetDataEnumKind<{{ looseName }}, '{{ getVariant(variant.name) }}'>;
   {% else %}
-    export function {{ dataEnumFunction }}(kind: '{{ variant.name | pascalCase }}'): GetDataEnumKind<{{ looseName }}, '{{ variant.name | pascalCase }}'>;
+    export function {{ dataEnumFunction }}(kind: '{{ getVariant(variant.name) }}'): GetDataEnumKind<{{ looseName }}, '{{ getVariant(variant.name) }}'>;
   {% endif %}
 {% endfor %}
-export function {{ dataEnumFunction }}<K extends {{ looseName }}['__kind']>(
+export function {{ dataEnumFunction }}<K extends {{ looseName }}['{{ dataEnumDiscriminator }}']>(
   kind: K,
   data?: any,
-): Extract<{{ looseName }}, { __kind: K }> {
-  return Array.isArray(data) ? { __kind: kind, fields: data } : { __kind: kind, ...(data ?? {}) };
+): Extract<{{ looseName }}, { {{ dataEnumDiscriminator }}: K }> {
+  return Array.isArray(data) ? { {{ dataEnumDiscriminator }}: kind, fields: data } : { {{ dataEnumDiscriminator }}: kind, ...(data ?? {}) };
 }
 
-export function {{ isDataEnumFunction }}<K extends {{ strictName }}['__kind']>(
+export function {{ isDataEnumFunction }}<K extends {{ strictName }}['{{ dataEnumDiscriminator }}']>(
   kind: K,
   value: {{ strictName }},
-): value is {{ strictName }} & { __kind: K } {
-  return value.__kind === kind;
+): value is {{ strictName }} & { {{ dataEnumDiscriminator }}: K } {
+  return value.{{ dataEnumDiscriminator }} === kind;
 };
diff --git a/src/renderers/js-experimental/fragments/typeDataEnumHelpers.ts b/src/renderers/js-experimental/fragments/typeDataEnumHelpers.ts
@@ -19,6 +19,8 @@ export function getTypeDataEnumHelpersFragment(
   return fragmentFromTemplate('typeDataEnumHelpers.njk', {
     strictName: nameApi.dataType(name),
     looseName: nameApi.dataArgsType(name),
+    dataEnumDiscriminator: nameApi.dataEnumDiscriminator(name),
+    getVariant: (variant: string) => nameApi.dataEnumVariant(variant),
     dataEnumFunction: nameApi.dataEnumFunction(name),
     isDataEnumFunction: nameApi.isDataEnumFunction(name),
     typeNode,

diff --git a/src/renderers/js-experimental/getTypeManifestVisitor.ts b/src/renderers/js-experimental/getTypeManifestVisitor.ts
@@ -7,7 +7,7 @@ import {
   structTypeNode,
   structTypeNodeFromInstructionArgumentNodes,
 } from '../../nodes';
-import { camelCase, pascalCase, pipe } from '../../shared';
+import { camelCase, pipe } from '../../shared';
 import { Visitor, extendVisitor, staticVisitor, visit } from '../../visitors';
 import { ImportMap } from './ImportMap';
 import { TypeManifest, mergeManifests } from './TypeManifest';
@@ -140,9 +140,6 @@ export function getTypeManifestVisitor(input: {
           const currentParentName = parentName;
           parentName = null;
 
-          const variantNames = enumType.variants.map(({ name }) =>
-            pascalCase(name)
-          );
           const encoderImports = new ImportMap();
           const decoderImports = new ImportMap();
           const encoderOptions: string[] = [];
@@ -173,6 +170,9 @@ export function getTypeManifestVisitor(input: {
                   'defined type that is a scalar enum through a visitor.'
               );
             }
+            const variantNames = enumType.variants.map(({ name }) =>
+              nameApi.scalarEnumVariant(name)
+            );
             return {
               isEnum: true,
               strictType: fragment(`{ ${variantNames.join(', ')} }`),
@@ -199,7 +199,7 @@ export function getTypeManifestVisitor(input: {
           }
 
           const variants = enumType.variants.map((variant): TypeManifest => {
-            const variantName = pascalCase(variant.name);
+            const variantName = nameApi.dataEnumVariant(variant.name);
             parentName = currentParentName
               ? {
                   strict: `GetDataEnumKindContent<${currentParentName.strict}, '${variantName}'>`,
@@ -242,8 +242,11 @@ export function getTypeManifestVisitor(input: {
         },
 
         visitEnumEmptyVariantType(enumEmptyVariantType) {
-          const name = pascalCase(enumEmptyVariantType.name);
-          const kindAttribute = `__kind: "${name}"`;
+          const discriminator = nameApi.dataEnumDiscriminator(
+            enumEmptyVariantType.name
+          );
+          const name = nameApi.dataEnumVariant(enumEmptyVariantType.name);
+          const kindAttribute = `${discriminator}: "${name}"`;
           return {
             isEnum: false,
             strictType: fragment(`{ ${kindAttribute} }`),
@@ -260,8 +263,11 @@ export function getTypeManifestVisitor(input: {
         },
 
         visitEnumStructVariantType(enumStructVariantType, { self }) {
-          const name = pascalCase(enumStructVariantType.name);
-          const kindAttribute = `__kind: "${name}"`;
+          const discriminator = nameApi.dataEnumDiscriminator(
+            enumStructVariantType.name
+          );
+          const name = nameApi.dataEnumVariant(enumStructVariantType.name);
+          const kindAttribute = `${discriminator}: "${name}"`;
           const structManifest = visit(enumStructVariantType.struct, self);
           structManifest.strictType.mapRender(
             (r) => `{ ${kindAttribute},${r.slice(1, -1)}}`
@@ -275,8 +281,11 @@ export function getTypeManifestVisitor(input: {
         },
 
         visitEnumTupleVariantType(enumTupleVariantType, { self }) {
-          const name = pascalCase(enumTupleVariantType.name);
-          const kindAttribute = `__kind: "${name}"`;
+          const discriminator = nameApi.dataEnumDiscriminator(
+            enumTupleVariantType.name
+          );
+          const name = nameApi.dataEnumVariant(enumTupleVariantType.name);
+          const kindAttribute = `${discriminator}: "${name}"`;
           const struct = structTypeNode([
             structFieldTypeNode({
               name: 'fields',

diff --git a/src/renderers/js-experimental/nameTransformers.ts b/src/renderers/js-experimental/nameTransformers.ts
@@ -1,10 +1,10 @@
 import {
-  capitalize,
-  titleCase,
-  pascalCase,
   camelCase,
+  capitalize,
   kebabCase,
+  pascalCase,
   snakeCase,
+  titleCase,
 } from '../../shared';
 
 export type NameTransformerHelpers = {
@@ -39,6 +39,9 @@ export type NameTransformerKey =
   | 'accountFetchFromSeedsFunction'
   | 'accountSafeFetchFromSeedsFunction'
   | 'accountGetSizeFunction'
+  | 'scalarEnumVariant'
+  | 'dataEnumDiscriminator'
+  | 'dataEnumVariant'
   | 'dataEnumFunction'
   | 'isDataEnumFunction'
   | 'instructionAsyncInputType'
@@ -109,6 +112,9 @@ export const DEFAULT_NAME_TRANSFORMERS: NameTransformers = {
   accountSafeFetchFromSeedsFunction: (name) =>
     `safeFetch${pascalCase(name)}FromSeeds`,
   accountGetSizeFunction: (name) => `get${pascalCase(name)}Size`,
+  scalarEnumVariant: (name) => `${pascalCase(name)}`,
+  dataEnumDiscriminator: () => '__kind',
+  dataEnumVariant: (name) => `${pascalCase(name)}`,
   dataEnumFunction: (name) => `${camelCase(name)}`,
   isDataEnumFunction: (name) => `is${pascalCase(name)}`,
   instructionAsyncInputType: (name) => `${pascalCase(name)}AsyncInput`,
@@ -132,11 +138,11 @@ export const DEFAULT_NAME_TRANSFORMERS: NameTransformers = {
     `${snakeCase(name).toUpperCase()}_PROGRAM_ADDRESS`,
   programCreateFunction: (name) => `create${pascalCase(name)}Program`,
   programAccountsEnum: (name) => `${pascalCase(name)}Account`,
-  programAccountsEnumVariant: (name) => `${snakeCase(name).toUpperCase()}`,
+  programAccountsEnumVariant: (name) => `${pascalCase(name)}`,
   programAccountsIdentifierFunction: (name) =>
     `identify${pascalCase(name)}Account`,
   programInstructionsEnum: (name) => `${pascalCase(name)}Instruction`,
-  programInstructionsEnumVariant: (name) => `${snakeCase(name).toUpperCase()}`,
+  programInstructionsEnumVariant: (name) => `${pascalCase(name)}`,
   programInstructionsIdentifierFunction: (name) =>
     `identify${pascalCase(name)}Instruction`,
   programErrorClass: (name) => `${pascalCase(name)}ProgramError`,

diff --git a/src/renderers/js-experimental/renderValueNodeVisitor.ts b/src/renderers/js-experimental/renderValueNodeVisitor.ts
@@ -1,5 +1,5 @@
 import { RegisteredValueNodeKind, isNode, isScalarEnum } from '../../nodes';
-import { LinkableDictionary, MainCaseString, pascalCase } from '../../shared';
+import { LinkableDictionary, MainCaseString } from '../../shared';
 import { Visitor, visit } from '../../visitors';
 import { Fragment, fragment, mergeFragments } from './fragments';
 import { NameApi } from './nameTransformers';
@@ -25,7 +25,6 @@ export function renderValueNodeVisitor(input: {
     visitEnumValue(node) {
       const enumName = nameApi.dataType(node.enum.name);
       const enumFunction = nameApi.dataEnumFunction(node.enum.name);
-      const variantName = pascalCase(node.variant);
       const importFrom = node.enum.importFrom ?? 'generatedTypes';
 
       const enumNode = linkables.get(node.enum);
@@ -35,12 +34,14 @@ export function renderValueNodeVisitor(input: {
           : !nonScalarEnums.includes(node.enum.name);
 
       if (!node.value && isScalar) {
+        const variantName = nameApi.scalarEnumVariant(node.variant);
         return fragment(`${enumName}.${variantName}`).addImports(
           importFrom,
           enumName
         );
       }
 
+      const variantName = nameApi.dataEnumVariant(node.variant);
       if (!node.value) {
         return fragment(`${enumFunction}('${variantName}')`).addImports(
           importFrom,

diff --git a/test/packages/js-experimental/src/generated/programs/mplCandyMachineCore.ts b/test/packages/js-experimental/src/generated/programs/mplCandyMachineCore.ts
@@ -35,31 +35,31 @@ export function createMplCandyMachineCoreProgram(): MplCandyMachineCoreProgram {
 }
 
 export enum MplCandyMachineCoreAccount {
-  CANDY_MACHINE,
+  CandyMachine,
 }
 
 export function identifyMplCandyMachineCoreAccount(
   account: { data: Uint8Array } | Uint8Array
 ): MplCandyMachineCoreAccount {
   const data = account instanceof Uint8Array ? account : account.data;
   if (memcmp(data, new Uint8Array([51, 173, 177, 113, 25, 241, 109, 189]), 0)) {
-    return MplCandyMachineCoreAccount.CANDY_MACHINE;
+    return MplCandyMachineCoreAccount.CandyMachine;
   }
   throw new Error(
     'The provided account could not be identified as a mplCandyMachineCore account.'
   );
 }
 
 export enum MplCandyMachineCoreInstruction {
-  DUMMY,
-  ADD_CONFIG_LINES,
-  INITIALIZE,
-  MINT_FROM_CANDY_MACHINE,
-  SET_AUTHORITY,
-  SET_COLLECTION,
-  SET_MINT_AUTHORITY,
-  UPDATE_CANDY_MACHINE,
-  WITHDRAW,
+  Dummy,
+  AddConfigLines,
+  Initialize,
+  MintFromCandyMachine,
+  SetAuthority,
+  SetCollection,
+  SetMintAuthority,
+  UpdateCandyMachine,
+  Withdraw,
 }
 
 export function identifyMplCandyMachineCoreInstruction(
@@ -68,33 +68,33 @@ export function identifyMplCandyMachineCoreInstruction(
   const data =
     instruction instanceof Uint8Array ? instruction : instruction.data;
   if (memcmp(data, new Uint8Array([167, 117, 211, 79, 251, 254, 47, 135]), 0)) {
-    return MplCandyMachineCoreInstruction.DUMMY;
+    return MplCandyMachineCoreInstruction.Dummy;
   }
   if (memcmp(data, new Uint8Array([223, 50, 224, 227, 151, 8, 115, 106]), 0)) {
-    return MplCandyMachineCoreInstruction.ADD_CONFIG_LINES;
+    return MplCandyMachineCoreInstruction.AddConfigLines;
   }
   if (memcmp(data, new Uint8Array([175, 175, 109, 31, 13, 152, 155, 237]), 0)) {
-    return MplCandyMachineCoreInstruction.INITIALIZE;
+    return MplCandyMachineCoreInstruction.Initialize;
   }
   if (memcmp(data, new Uint8Array([51, 57, 225, 47, 182, 146, 137, 166]), 0)) {
-    return MplCandyMachineCoreInstruction.MINT_FROM_CANDY_MACHINE;
+    return MplCandyMachineCoreInstruction.MintFromCandyMachine;
   }
   if (memcmp(data, new Uint8Array([133, 250, 37, 21, 110, 163, 26, 121]), 0)) {
-    return MplCandyMachineCoreInstruction.SET_AUTHORITY;
+    return MplCandyMachineCoreInstruction.SetAuthority;
   }
   if (memcmp(data, new Uint8Array([192, 254, 206, 76, 168, 182, 59, 223]), 0)) {
-    return MplCandyMachineCoreInstruction.SET_COLLECTION;
+    return MplCandyMachineCoreInstruction.SetCollection;
   }
   if (
     memcmp(data, new Uint8Array([67, 127, 155, 187, 100, 174, 103, 121]), 0)
   ) {
-    return MplCandyMachineCoreInstruction.SET_MINT_AUTHORITY;
+    return MplCandyMachineCoreInstruction.SetMintAuthority;
   }
   if (memcmp(data, new Uint8Array([219, 200, 88, 176, 158, 63, 253, 127]), 0)) {
-    return MplCandyMachineCoreInstruction.UPDATE_CANDY_MACHINE;
+    return MplCandyMachineCoreInstruction.UpdateCandyMachine;
   }
   if (memcmp(data, new Uint8Array([183, 18, 70, 156, 148, 109, 161, 34]), 0)) {
-    return MplCandyMachineCoreInstruction.WITHDRAW;
+    return MplCandyMachineCoreInstruction.Withdraw;
   }
   throw new Error(
     'The provided instruction could not be identified as a mplCandyMachineCore instruction.'

diff --git a/test/packages/js-experimental/src/generated/programs/mplTokenAuthRules.ts b/test/packages/js-experimental/src/generated/programs/mplTokenAuthRules.ts
@@ -37,25 +37,25 @@ export function createMplTokenAuthRulesProgram(): MplTokenAuthRulesProgram {
 }
 
 export enum MplTokenAuthRulesAccount {
-  FREQUENCY_ACCOUNT,
+  FrequencyAccount,
 }
 
 export function identifyMplTokenAuthRulesAccount(
   account: { data: Uint8Array } | Uint8Array
 ): MplTokenAuthRulesAccount {
   const data = account instanceof Uint8Array ? account : account.data;
   if (memcmp(data, getU64Encoder().encode(TaKey.Frequency), 0)) {
-    return MplTokenAuthRulesAccount.FREQUENCY_ACCOUNT;
+    return MplTokenAuthRulesAccount.FrequencyAccount;
   }
   throw new Error(
     'The provided account could not be identified as a mplTokenAuthRules account.'
   );
 }
 
 export enum MplTokenAuthRulesInstruction {
-  CREATE_RULE_SET,
-  VALIDATE,
-  CREATE_FREQUENCY_RULE,
+  CreateRuleSet,
+  Validate,
+  CreateFrequencyRule,
 }
 
 export function identifyMplTokenAuthRulesInstruction(
@@ -64,13 +64,13 @@ export function identifyMplTokenAuthRulesInstruction(
   const data =
     instruction instanceof Uint8Array ? instruction : instruction.data;
   if (memcmp(data, getU8Encoder().encode(0), 0)) {
-    return MplTokenAuthRulesInstruction.CREATE_RULE_SET;
+    return MplTokenAuthRulesInstruction.CreateRuleSet;
   }
   if (memcmp(data, getU8Encoder().encode(1), 0)) {
-    return MplTokenAuthRulesInstruction.VALIDATE;
+    return MplTokenAuthRulesInstruction.Validate;
   }
   if (memcmp(data, getU8Encoder().encode(2), 0)) {
-    return MplTokenAuthRulesInstruction.CREATE_FREQUENCY_RULE;
+    return MplTokenAuthRulesInstruction.CreateFrequencyRule;
   }
   throw new Error(
     'The provided instruction could not be identified as a mplTokenAuthRules instruction.'