From fc3242132ff2fbd10fe97b3100b6941de90d6bc5 Mon Sep 17 00:00:00 2001
From: Loris Leiva <loris.leiva@gmail.com>
Date: Sun, 31 Dec 2023 12:58:59 +0000
Subject: [PATCH] Add PdaNode (#126)

---
 .changeset/stale-yaks-arrive.md |  5 +++++
 src/nodes/Node.ts               |  2 ++
 src/nodes/PdaNode.ts            | 15 +++++++++++++++
 src/nodes/ProgramNode.ts        |  3 +++
 src/nodes/index.ts              |  1 +
 src/visitors/identityVisitor.ts | 15 +++++++++++++++
 src/visitors/mergeVisitor.ts    |  7 +++++++
 7 files changed, 48 insertions(+)
 create mode 100644 .changeset/stale-yaks-arrive.md
 create mode 100644 src/nodes/PdaNode.ts

diff --git a/.changeset/stale-yaks-arrive.md b/.changeset/stale-yaks-arrive.md
new file mode 100644
index 000000000..a09475eb4
--- /dev/null
+++ b/.changeset/stale-yaks-arrive.md
@@ -0,0 +1,5 @@
+---
+'@metaplex-foundation/kinobi': patch
+---
+
+Add PdaNode
diff --git a/src/nodes/Node.ts b/src/nodes/Node.ts
index b1c89a6d3..fef44d805 100644
--- a/src/nodes/Node.ts
+++ b/src/nodes/Node.ts
@@ -6,6 +6,7 @@ import type { InstructionAccountNode } from './InstructionAccountNode';
 import type { InstructionDataArgsNode } from './InstructionDataArgsNode';
 import type { InstructionExtraArgsNode } from './InstructionExtraArgsNode';
 import type { InstructionNode } from './InstructionNode';
+import type { PdaNode } from './PdaNode';
 import type { ProgramNode } from './ProgramNode';
 import type { RootNode } from './RootNode';
 import { REGISTERED_PDA_SEED_NODES } from './pdaSeedNodes';
@@ -18,6 +19,7 @@ import { REGISTERED_VALUE_NODES } from './valueNodes';
 const REGISTERED_NODES = {
   rootNode: {} as RootNode,
   programNode: {} as ProgramNode,
+  pdaNode: {} as PdaNode,
   accountNode: {} as AccountNode,
   accountDataNode: {} as AccountDataNode,
   instructionNode: {} as InstructionNode,
diff --git a/src/nodes/PdaNode.ts b/src/nodes/PdaNode.ts
new file mode 100644
index 000000000..f4c7890bb
--- /dev/null
+++ b/src/nodes/PdaNode.ts
@@ -0,0 +1,15 @@
+import { InvalidKinobiTreeError, MainCaseString, mainCase } from '../shared';
+import { PdaSeedNode } from './pdaSeedNodes';
+
+export type PdaNode = {
+  readonly kind: 'pdaNode';
+  readonly name: MainCaseString;
+  readonly seeds: PdaSeedNode[];
+};
+
+export function pdaNode(name: string, seeds: PdaSeedNode[]): PdaNode {
+  if (!name) {
+    throw new InvalidKinobiTreeError('PdaNode must have a name.');
+  }
+  return { kind: 'pdaNode', name: mainCase(name), seeds };
+}
diff --git a/src/nodes/ProgramNode.ts b/src/nodes/ProgramNode.ts
index a1a63ec30..14b5e9adf 100644
--- a/src/nodes/ProgramNode.ts
+++ b/src/nodes/ProgramNode.ts
@@ -4,9 +4,11 @@ import { AccountNode, accountNodeFromIdl } from './AccountNode';
 import { DefinedTypeNode, definedTypeNodeFromIdl } from './DefinedTypeNode';
 import { ErrorNode, errorNodeFromIdl } from './ErrorNode';
 import { InstructionNode, instructionNodeFromIdl } from './InstructionNode';
+import { PdaNode } from './PdaNode';
 
 export type ProgramNode = {
   readonly kind: 'programNode';
+  readonly pdas: PdaNode[];
   readonly accounts: AccountNode[];
   readonly instructions: InstructionNode[];
   readonly definedTypes: DefinedTypeNode[];
@@ -38,6 +40,7 @@ export type ProgramNodeInput = Omit<
 export function programNode(input: ProgramNodeInput): ProgramNode {
   return {
     kind: 'programNode',
+    pdas: input.pdas ?? [],
     accounts: input.accounts,
     instructions: input.instructions,
     definedTypes: input.definedTypes,
diff --git a/src/nodes/index.ts b/src/nodes/index.ts
index da53c713a..ff7787b8e 100644
--- a/src/nodes/index.ts
+++ b/src/nodes/index.ts
@@ -7,6 +7,7 @@ export * from './InstructionDataArgsNode';
 export * from './InstructionExtraArgsNode';
 export * from './InstructionNode';
 export * from './Node';
+export * from './PdaNode';
 export * from './ProgramNode';
 export * from './RootNode';
 
diff --git a/src/visitors/identityVisitor.ts b/src/visitors/identityVisitor.ts
index 004d69bf9..d89b69c8d 100644
--- a/src/visitors/identityVisitor.ts
+++ b/src/visitors/identityVisitor.ts
@@ -28,6 +28,7 @@ import {
   mapTypeNode,
   mapValueNode,
   optionTypeNode,
+  pdaNode,
   prefixedSizeNode,
   programNode,
   removeNullAndAssertIsNodeFilter,
@@ -76,6 +77,9 @@ export function identityVisitor<
     visitor.visitProgram = function visitProgram(node) {
       return programNode({
         ...node,
+        pdas: node.pdas
+          .map((account) => visit(this)(account))
+          .filter(removeNullAndAssertIsNodeFilter('pdaNode')),
         accounts: node.accounts
           .map((account) => visit(this)(account))
           .filter(removeNullAndAssertIsNodeFilter('accountNode')),
@@ -92,6 +96,17 @@ export function identityVisitor<
     };
   }
 
+  if (castedNodeKeys.includes('pdaNode')) {
+    visitor.visitPda = function visitPda(node) {
+      return pdaNode(
+        node.name,
+        node.seeds
+          .map((type) => visit(this)(type))
+          .filter(removeNullAndAssertIsNodeFilter(PDA_SEED_NODES))
+      );
+    };
+  }
+
   if (castedNodeKeys.includes('accountNode')) {
     visitor.visitAccount = function visitAccount(node) {
       const data = visit(this)(node.data);
diff --git a/src/visitors/mergeVisitor.ts b/src/visitors/mergeVisitor.ts
index f1856bdae..ac0b62056 100644
--- a/src/visitors/mergeVisitor.ts
+++ b/src/visitors/mergeVisitor.ts
@@ -26,6 +26,7 @@ export function mergeVisitor<
   if (castedNodeKeys.includes('programNode')) {
     visitor.visitProgram = function visitProgram(node) {
       return merge(node, [
+        ...node.pdas.flatMap(visit(this)),
         ...node.accounts.flatMap(visit(this)),
         ...node.instructions.flatMap(visit(this)),
         ...node.definedTypes.flatMap(visit(this)),
@@ -34,6 +35,12 @@ export function mergeVisitor<
     };
   }
 
+  if (castedNodeKeys.includes('pdaNode')) {
+    visitor.visitPda = function visitPda(node) {
+      return merge(node, node.seeds.flatMap(visit(this)));
+    };
+  }
+
   if (castedNodeKeys.includes('accountNode')) {
     visitor.visitAccount = function visitAccount(node) {
       return merge(node, [