From 3bfeb8c17b201036256a9f1906f6ef4bec861544 Mon Sep 17 00:00:00 2001
From: Loris Leiva <loris.leiva@gmail.com>
Date: Thu, 4 Jan 2024 23:55:00 +0000
Subject: [PATCH] Add a visitor that adds PDA nodes to programs (#139)

---
 .changeset/shaggy-paws-jam.md        |  5 ++
 src/visitors/addPdasVisitor.ts       | 37 +++++++++++++
 src/visitors/index.ts                |  1 +
 test/visitors/addPdasVisitor.test.ts | 81 ++++++++++++++++++++++++++++
 4 files changed, 124 insertions(+)
 create mode 100644 .changeset/shaggy-paws-jam.md
 create mode 100644 src/visitors/addPdasVisitor.ts
 create mode 100644 test/visitors/addPdasVisitor.test.ts

diff --git a/.changeset/shaggy-paws-jam.md b/.changeset/shaggy-paws-jam.md
new file mode 100644
index 000000000..c80288e04
--- /dev/null
+++ b/.changeset/shaggy-paws-jam.md
@@ -0,0 +1,5 @@
+---
+'@metaplex-foundation/kinobi': minor
+---
+
+Add a visitor that adds PDA nodes to programs
diff --git a/src/visitors/addPdasVisitor.ts b/src/visitors/addPdasVisitor.ts
new file mode 100644
index 000000000..6ab5787cd
--- /dev/null
+++ b/src/visitors/addPdasVisitor.ts
@@ -0,0 +1,37 @@
+import { KinobiError, mainCase } from '../shared';
+import { PdaSeedNode, assertIsNode, pdaNode, programNode } from '../nodes';
+import { bottomUpTransformerVisitor } from './bottomUpTransformerVisitor';
+
+export function addPdasVisitor(
+  pdas: Record<string, { name: string; seeds: PdaSeedNode[] }[]>
+) {
+  return bottomUpTransformerVisitor(
+    Object.entries(pdas).map(([uncasedProgramName, newPdas]) => {
+      const programName = mainCase(uncasedProgramName);
+      return {
+        select: `[programNode]${programName}`,
+        transform: (node) => {
+          assertIsNode(node, 'programNode');
+          const existingPdaNames = new Set(node.pdas.map((pda) => pda.name));
+          const newPdaNames = new Set(newPdas.map((pda) => pda.name));
+          const overlappingPdaNames = new Set(
+            [...existingPdaNames].filter((name) => newPdaNames.has(name))
+          );
+          if (overlappingPdaNames.size > 0) {
+            throw new KinobiError(
+              `Cannot add PDAs to program "${programName}" because the following PDA names ` +
+                `already exist: ${[...overlappingPdaNames].join(', ')}.`
+            );
+          }
+          return programNode({
+            ...node,
+            pdas: [
+              ...node.pdas,
+              ...newPdas.map((pda) => pdaNode(pda.name, pda.seeds)),
+            ],
+          });
+        },
+      };
+    })
+  );
+}
diff --git a/src/visitors/index.ts b/src/visitors/index.ts
index fd46d69e8..7cdb475a2 100644
--- a/src/visitors/index.ts
+++ b/src/visitors/index.ts
@@ -1,3 +1,4 @@
+export * from './addPdasVisitor';
 export * from './bottomUpTransformerVisitor';
 export * from './consoleLogVisitor';
 export * from './createSubInstructionsFromEnumArgsVisitor';
diff --git a/test/visitors/addPdasVisitor.test.ts b/test/visitors/addPdasVisitor.test.ts
new file mode 100644
index 000000000..ef8cd0ef7
--- /dev/null
+++ b/test/visitors/addPdasVisitor.test.ts
@@ -0,0 +1,81 @@
+import test from 'ava';
+import {
+  addPdasVisitor,
+  constantPdaSeedNodeFromString,
+  pdaNode,
+  programIdPdaSeedNode,
+  programNode,
+  publicKeyTypeNode,
+  variablePdaSeedNode,
+  visit,
+} from '../../src';
+
+test('it adds PDA nodes to a program', (t) => {
+  // Given a program with a single PDA.
+  const node = programNode({
+    name: 'myProgram',
+    publicKey: 'Epo9rxh99jpeeWabRZi4tpgUVxZQeVn9vbbDjUztJtu4',
+    pdas: [
+      pdaNode('associatedToken', [
+        variablePdaSeedNode('owner', publicKeyTypeNode()),
+        programIdPdaSeedNode(),
+        variablePdaSeedNode('mint', publicKeyTypeNode()),
+      ]),
+    ],
+  });
+
+  // When we add two more PDAs.
+  const newPdas = [
+    pdaNode('metadata', [
+      constantPdaSeedNodeFromString('metadata'),
+      programIdPdaSeedNode(),
+      variablePdaSeedNode('mint', publicKeyTypeNode()),
+    ]),
+    pdaNode('masterEdition', [
+      constantPdaSeedNodeFromString('metadata'),
+      programIdPdaSeedNode(),
+      variablePdaSeedNode('mint', publicKeyTypeNode()),
+      constantPdaSeedNodeFromString('edition'),
+    ]),
+  ];
+  const result = visit(node, addPdasVisitor({ myProgram: newPdas }));
+
+  // Then we expect the following program to be returned.
+  t.deepEqual(result, { ...node, pdas: [...node.pdas, ...newPdas] });
+});
+
+test('it fails to add a PDA if its name conflicts with an existing PDA on the program', (t) => {
+  // Given a program with a PDA named "myPda".
+  const node = programNode({
+    name: 'myProgram',
+    publicKey: 'Epo9rxh99jpeeWabRZi4tpgUVxZQeVn9vbbDjUztJtu4',
+    pdas: [
+      pdaNode('myPda', [
+        variablePdaSeedNode('owner', publicKeyTypeNode()),
+        programIdPdaSeedNode(),
+        variablePdaSeedNode('mint', publicKeyTypeNode()),
+      ]),
+    ],
+  });
+
+  // When we try to add another PDA with the same name.
+  const fn = () =>
+    visit(
+      node,
+      addPdasVisitor({
+        myProgram: [
+          pdaNode('myPda', [
+            constantPdaSeedNodeFromString('metadata'),
+            programIdPdaSeedNode(),
+            variablePdaSeedNode('mint', publicKeyTypeNode()),
+          ]),
+        ],
+      })
+    );
+
+  // Then we expect the following error to be thrown.
+  t.throws(fn, {
+    message:
+      'Cannot add PDAs to program "myProgram" because the following PDA names already exist: myPda.',
+  });
+});