From 4a1aaad52c001f2e8f3e83203853d55a27d04dff Mon Sep 17 00:00:00 2001
From: Loris Leiva <loris.leiva@gmail.com>
Date: Mon, 1 Jan 2024 17:36:48 +0000
Subject: [PATCH] Update programLinkNode name when updating programNode name

---
 .../fragments/accountPdaHelpers.ts            |  2 +-
 src/visitors/getByteSizeVisitor.ts            |  2 +-
 src/visitors/updateDefinedTypesVisitor.ts     |  6 +--
 src/visitors/updateProgramsVisitor.ts         | 48 +++++++++++++++----
 4 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/src/renderers/js-experimental/fragments/accountPdaHelpers.ts b/src/renderers/js-experimental/fragments/accountPdaHelpers.ts
index d4e886d93..46d30c8a4 100644
--- a/src/renderers/js-experimental/fragments/accountPdaHelpers.ts
+++ b/src/renderers/js-experimental/fragments/accountPdaHelpers.ts
@@ -1,5 +1,5 @@
-import { LinkableDictionary } from 'src/shared';
 import { AccountNode, ProgramNode, isNodeFilter } from '../../../nodes';
+import { LinkableDictionary } from '../../../shared';
 import { NameApi } from '../nameTransformers';
 import { Fragment, fragment, fragmentFromTemplate } from './common';
 
diff --git a/src/visitors/getByteSizeVisitor.ts b/src/visitors/getByteSizeVisitor.ts
index 59d768303..2777a3657 100644
--- a/src/visitors/getByteSizeVisitor.ts
+++ b/src/visitors/getByteSizeVisitor.ts
@@ -1,10 +1,10 @@
-import { LinkableDictionary } from 'src/shared';
 import {
   REGISTERED_TYPE_NODE_KEYS,
   RegisteredTypeNodes,
   isNode,
   isScalarEnum,
 } from '../nodes';
+import { LinkableDictionary } from '../shared';
 import { mergeVisitor } from './mergeVisitor';
 import { Visitor, visit } from './visitor';
 
diff --git a/src/visitors/updateDefinedTypesVisitor.ts b/src/visitors/updateDefinedTypesVisitor.ts
index 67c8cbd4c..ac10ca77f 100644
--- a/src/visitors/updateDefinedTypesVisitor.ts
+++ b/src/visitors/updateDefinedTypesVisitor.ts
@@ -30,7 +30,7 @@ export function updateDefinedTypesVisitor(
             ? mainCase(updates.name)
             : undefined;
 
-        const transforms: BottomUpNodeTransformerWithSelector[] = [
+        const transformers: BottomUpNodeTransformerWithSelector[] = [
           {
             select: `${selectorStack.join('.')}.[definedTypeNode]${name}`,
             transform: (node) => {
@@ -56,7 +56,7 @@ export function updateDefinedTypesVisitor(
         ];
 
         if (newName) {
-          transforms.push({
+          transformers.push({
             select: `${selectorStack.join('.')}.[definedTypeLinkNode]${name}`,
             transform: (node) => {
               assertIsNode(node, 'definedTypeLinkNode');
@@ -66,7 +66,7 @@ export function updateDefinedTypesVisitor(
           });
         }
 
-        return transforms;
+        return transformers;
       }
     )
   );
diff --git a/src/visitors/updateProgramsVisitor.ts b/src/visitors/updateProgramsVisitor.ts
index 78631bd0f..edef34e88 100644
--- a/src/visitors/updateProgramsVisitor.ts
+++ b/src/visitors/updateProgramsVisitor.ts
@@ -1,4 +1,10 @@
-import { ProgramNodeInput, assertIsNode, programNode } from '../nodes';
+import {
+  ProgramNodeInput,
+  assertIsNode,
+  programLinkNode,
+  programNode,
+} from '../nodes';
+import { mainCase } from '../shared';
 import {
   BottomUpNodeTransformerWithSelector,
   bottomUpTransformerVisitor,
@@ -15,15 +21,37 @@ export type ProgramUpdates =
 
 export function updateProgramsVisitor(map: Record<string, ProgramUpdates>) {
   return bottomUpTransformerVisitor(
-    Object.entries(map).map(
-      ([name, updates]): BottomUpNodeTransformerWithSelector => ({
-        select: `[programNode]${name}`,
-        transform: (node) => {
-          assertIsNode(node, 'programNode');
-          if ('delete' in updates) return null;
-          return programNode({ ...node, ...updates });
-        },
-      })
+    Object.entries(map).flatMap(
+      ([name, updates]): BottomUpNodeTransformerWithSelector[] => {
+        const newName =
+          typeof updates === 'object' && 'name' in updates && updates.name
+            ? mainCase(updates.name)
+            : undefined;
+
+        const transformers: BottomUpNodeTransformerWithSelector[] = [
+          {
+            select: `[programNode]${name}`,
+            transform: (node) => {
+              assertIsNode(node, 'programNode');
+              if ('delete' in updates) return null;
+              return programNode({ ...node, ...updates });
+            },
+          },
+        ];
+
+        if (newName) {
+          transformers.push({
+            select: `[programLinkNode]${name}`,
+            transform: (node) => {
+              assertIsNode(node, 'programLinkNode');
+              if (node.importFrom) return node;
+              return programLinkNode(newName);
+            },
+          });
+        }
+
+        return transformers;
+      }
     )
   );
 }