Add discriminator nodes (#136)

* Add discriminator nodes

* Add byteDiscriminatorNodeFromBase58 helper

* Update @solana/codecs-strings to the latest version

* Fix typo

* Add DiscriminatorNodes to AccountNode

* Rename AnchorDiscriminator helper file

* Add DiscriminatorNodes to InstructionNode

* Update merge and identity visitors

* Add changeset

.changeset/gold-beans-listen.md

@@ -0,0 +1,5 @@
+---
+'@metaplex-foundation/kinobi': minor
+---
+
+Add discriminator nodes

package.json

@@ -59,6 +59,7 @@
   },
   "dependencies": {
     "@noble/hashes": "^1.1.5",
+    "@solana/codecs-strings": "2.0.0-experimental.a157265",
     "chalk": "^4.0.0",
     "json-stable-stringify": "^1.1.0",
     "nunjucks": "^3.2.3",

src/nodes/AccountNode.ts

@@ -1,11 +1,7 @@
 import type { IdlAccount } from '../idl';
-import {
-  AccountDiscriminator,
-  InvalidKinobiTreeError,
-  MainCaseString,
-  mainCase,
-} from '../shared';
+import { InvalidKinobiTreeError, MainCaseString, mainCase } from '../shared';
 import { assertIsNode } from './Node';
+import { DiscriminatorNode } from './discriminatorNodes';
 import { PdaLinkNode, pdaLinkNode } from './linkNodes';
 import { StructTypeNode, structTypeNode } from './typeNodes';
 import { createTypeNodeFromIdl } from './typeNodes/TypeNode';
@@ -16,9 +12,7 @@ export type AccountNode = {
   // Children.
   readonly data: StructTypeNode;
   readonly pda?: PdaLinkNode;
-
-  // Children to-be.
-  readonly discriminator?: AccountDiscriminator;
+  readonly discriminators?: DiscriminatorNode[];
 
   // Data.
   readonly name: MainCaseString;
@@ -37,13 +31,17 @@ export function accountNode(input: AccountNodeInput): AccountNode {
   }
   return {
     kind: 'accountNode',
+
+    // Children.
     data: input.data ?? structTypeNode([]),
     pda: input.pda,
+    discriminators: input.discriminators,
+
+    // Data.
     name: mainCase(input.name),
     idlName: input.idlName ?? input.name,
     docs: input.docs ?? [],
     size: input.size,
-    discriminator: input.discriminator,
   };
 }
 

src/nodes/InstructionNode.ts

@@ -14,6 +14,7 @@ import { InstructionRemainingAccountsNode } from './InstructionRemainingAccounts
 import { isNode } from './Node';
 import { ProgramNode } from './ProgramNode';
 import { RootNode } from './RootNode';
+import { DiscriminatorNode } from './discriminatorNodes';
 import { createTypeNodeFromIdl } from './typeNodes/TypeNode';
 import { numberValueNode } from './valueNodes';
 
@@ -24,9 +25,10 @@ export type InstructionNode = {
   readonly accounts: InstructionAccountNode[];
   readonly arguments: InstructionArgumentNode[];
   readonly extraArguments?: InstructionArgumentNode[];
-  readonly subInstructions?: InstructionNode[];
   readonly remainingAccounts?: InstructionRemainingAccountsNode[];
   readonly byteDeltas?: InstructionByteDeltaNode[];
+  readonly discriminators?: DiscriminatorNode[];
+  readonly subInstructions?: InstructionNode[];
 
   // Data.
   readonly name: MainCaseString;
@@ -49,15 +51,20 @@ export function instructionNode(input: InstructionNodeInput): InstructionNode {
   const name = mainCase(input.name);
   return {
     kind: 'instructionNode',
-    name,
+
+    // Children.
     accounts: input.accounts ?? [],
     arguments: input.arguments ?? [],
     extraArguments: input.extraArguments,
+    remainingAccounts: input.remainingAccounts,
+    byteDeltas: input.byteDeltas,
+    discriminators: input.discriminators,
     subInstructions: input.subInstructions,
+
+    // Data.
+    name,
     idlName: input.idlName ?? input.name,
     docs: input.docs ?? [],
-    remainingAccounts: input.remainingAccounts,
-    byteDeltas: input.byteDeltas,
     optionalAccountStrategy: input.optionalAccountStrategy ?? 'programId',
   };
 }

src/nodes/Node.ts

@@ -11,6 +11,7 @@ import type { PdaNode } from './PdaNode';
 import type { ProgramNode } from './ProgramNode';
 import type { RootNode } from './RootNode';
 import { REGISTERED_CONTEXTUAL_VALUE_NODES } from './contextualValueNodes/ContextualValueNode';
+import { REGISTERED_DISCRIMINATOR_NODES } from './discriminatorNodes/DiscriminatorNode';
 import { REGISTERED_LINK_NODES } from './linkNodes/LinkNode';
 import { REGISTERED_PDA_SEED_NODES } from './pdaSeedNodes/PdaSeedNode';
 import { REGISTERED_SIZE_NODES } from './sizeNodes/SizeNode';
@@ -34,6 +35,7 @@ const REGISTERED_NODES = {
 
   // Groups.
   ...REGISTERED_CONTEXTUAL_VALUE_NODES,
+  ...REGISTERED_DISCRIMINATOR_NODES,
   ...REGISTERED_LINK_NODES,
   ...REGISTERED_PDA_SEED_NODES,
   ...REGISTERED_SIZE_NODES,

src/nodes/discriminatorNodes/ByteDiscriminatorNode.ts

@@ -0,0 +1,26 @@
+import { getBase58Encoder } from '@solana/codecs-strings';
+
+export type ByteDiscriminatorNode = {
+  readonly kind: 'byteDiscriminatorNode';
+
+  // Data.
+  readonly bytes: number[];
+  readonly offset: number;
+};
+
+export function byteDiscriminatorNode(
+  bytes: number[],
+  offset: number = 0
+): ByteDiscriminatorNode {
+  return { kind: 'byteDiscriminatorNode', bytes, offset };
+}
+
+export function byteDiscriminatorNodeFromBase58(
+  base58Bytes: string,
+  offset: number = 0
+): ByteDiscriminatorNode {
+  return byteDiscriminatorNode(
+    [...getBase58Encoder().encode(base58Bytes)],
+    offset
+  );
+}

src/nodes/discriminatorNodes/DiscriminatorNode.ts

@@ -0,0 +1,25 @@
+import { getNodeKinds } from '../../shared/utils';
+import type { ByteDiscriminatorNode } from './ByteDiscriminatorNode';
+import type { FieldDiscriminatorNode } from './FieldDiscriminatorNode';
+import type { SizeDiscriminatorNode } from './SizeDiscriminatorNode';
+
+// Discriminator Node Registration.
+
+export const REGISTERED_DISCRIMINATOR_NODES = {
+  byteDiscriminatorNode: {} as ByteDiscriminatorNode,
+  fieldDiscriminatorNode: {} as FieldDiscriminatorNode,
+  sizeDiscriminatorNode: {} as SizeDiscriminatorNode,
+};
+
+export const REGISTERED_DISCRIMINATOR_NODE_KINDS = getNodeKinds(
+  REGISTERED_DISCRIMINATOR_NODES
+);
+export type RegisteredDiscriminatorNodeKind =
+  typeof REGISTERED_DISCRIMINATOR_NODE_KINDS[number];
+export type RegisteredDiscriminatorNode =
+  typeof REGISTERED_DISCRIMINATOR_NODES[RegisteredDiscriminatorNodeKind];
+
+// Discriminator Node Helpers.
+
+export type DiscriminatorNode = RegisteredDiscriminatorNode;
+export const DISCRIMINATOR_NODES = REGISTERED_DISCRIMINATOR_NODE_KINDS;

src/nodes/discriminatorNodes/FieldDiscriminatorNode.ts

@@ -0,0 +1,16 @@
+import { MainCaseString, mainCase } from '../../shared';
+
+export type FieldDiscriminatorNode = {
+  readonly kind: 'fieldDiscriminatorNode';
+
+  // Data.
+  readonly name: MainCaseString;
+  readonly offset: number;
+};
+
+export function fieldDiscriminatorNode(
+  name: string,
+  offset: number = 0
+): FieldDiscriminatorNode {
+  return { kind: 'fieldDiscriminatorNode', name: mainCase(name), offset };
+}

src/nodes/discriminatorNodes/SizeDiscriminatorNode.ts

@@ -0,0 +1,10 @@
+export type SizeDiscriminatorNode = {
+  readonly kind: 'sizeDiscriminatorNode';
+
+  // Data.
+  readonly size: number;
+};
+
+export function sizeDiscriminatorNode(size: number): SizeDiscriminatorNode {
+  return { kind: 'sizeDiscriminatorNode', size };
+}

src/nodes/discriminatorNodes/index.ts

@@ -0,0 +1,4 @@
+export * from './ByteDiscriminatorNode';
+export * from './DiscriminatorNode';
+export * from './FieldDiscriminatorNode';
+export * from './SizeDiscriminatorNode';

src/nodes/index.ts

@@ -12,6 +12,7 @@ export * from './ProgramNode';
 export * from './RootNode';
 
 export * from './contextualValueNodes';
+export * from './discriminatorNodes';
 export * from './linkNodes';
 export * from './pdaSeedNodes';
 export * from './sizeNodes';

src/renderers/js/getRenderMapVisitor.ts

@@ -4,6 +4,7 @@ import {
   Options as PrettierOptions,
 } from 'prettier';
 import {
+  FieldDiscriminatorNode,
   getAllAccounts,
   getAllDefinedTypes,
   getAllInstructionsWithSubs,
@@ -12,6 +13,7 @@ import {
   isNode,
   isNodeFilter,
   ProgramNode,
+  SizeDiscriminatorNode,
   structTypeNodeFromInstructionArgumentNodes,
   VALUE_NODES,
 } from '../../nodes';
@@ -315,12 +317,15 @@ export function getRenderMapVisitor(
             .addAlias('umi', 'publicKey', 'toPublicKey');
 
           // Discriminator.
-          const { discriminator } = node;
+          const discriminator =
+            (node.discriminators ?? []).find(
+              (d) => !isNode(d, 'byteDiscriminatorNode')
+            ) ?? null;
           let resolvedDiscriminator:
-            | { kind: 'size'; value: string }
-            | { kind: 'field'; name: string; value: string }
+            | SizeDiscriminatorNode
+            | (FieldDiscriminatorNode & { value: string })
             | null = null;
-          if (discriminator?.kind === 'field') {
+          if (isNode(discriminator, 'fieldDiscriminatorNode')) {
             const discriminatorField = node.data.fields.find(
               (f) => f.name === discriminator.name
             );
@@ -330,16 +335,12 @@ export function getRenderMapVisitor(
             if (discriminatorValue) {
               imports.mergeWith(discriminatorValue.imports);
               resolvedDiscriminator = {
-                kind: 'field',
-                name: discriminator.name,
+                ...discriminator,
                 value: discriminatorValue.render,
               };
             }
-          } else if (discriminator?.kind === 'size') {
-            resolvedDiscriminator =
-              node.size !== undefined
-                ? { kind: 'size', value: `${node.size}` }
-                : null;
+          } else if (isNode(discriminator, 'sizeDiscriminatorNode')) {
+            resolvedDiscriminator = discriminator;
           }
 
           // GPA Fields.

src/renderers/js/templates/accountsPage.njk

@@ -69,10 +69,10 @@ export function get{{ account.name | pascalCase }}GpaBuilder(context: Pick<Conte
   return gpaBuilder(context, programId)
     .registerFields<{{ gpaFields.type }}>({{ gpaFields.argument }})
     .deserializeUsing<{{ account.name | pascalCase }}>((account) => deserialize{{ account.name | pascalCase }}(account))
-    {%- if discriminator.kind === 'field' %}
+    {%- if discriminator.kind === 'fieldDiscriminatorNode' %}
       .whereField('{{ discriminator.name }}', {{ discriminator.value }})
-    {% elif discriminator.kind === 'size' %}
-      .whereSize({{ discriminator.value }})
+    {% elif discriminator.kind === 'sizeDiscriminatorNode' %}
+      .whereSize({{ discriminator.size }})
     {% endif -%}
   ;
 }

src/shared/AnchorDiscriminator.ts → src/shared/anchorHelpers.ts

@@ -1,19 +1,17 @@
 import { sha256 } from '@noble/hashes/sha256';
-import { pascalCase, snakeCase } from './utils';
 import { ArrayValueNode, arrayValueNode, numberValueNode } from '../nodes';
-
-export type AnchorDiscriminator = ArrayValueNode;
+import { pascalCase, snakeCase } from './utils';
 
 export const getAnchorInstructionDiscriminator = (
   idlName: string
-): AnchorDiscriminator => {
+): ArrayValueNode => {
   const hash = sha256(`global:${snakeCase(idlName)}`).slice(0, 8);
   return arrayValueNode([...hash].map((byte) => numberValueNode(byte)));
 };
 
 export const getAnchorAccountDiscriminator = (
   idlName: string
-): AnchorDiscriminator => {
+): ArrayValueNode => {
   const hash = sha256(`account:${pascalCase(idlName)}`).slice(0, 8);
   return arrayValueNode([...hash].map((byte) => numberValueNode(byte)));
 };

src/shared/index.ts

@@ -1,5 +1,4 @@
-export * from './AccountDiscriminator';
-export * from './AnchorDiscriminator';
+export * from './anchorHelpers';
 export * from './GpaField';
 export * from './ImportFrom';
 export * from './LinkableDictionary';