From adfdd8c99ca685389cec9659f1ac9385deea2f1b Mon Sep 17 00:00:00 2001
From: Loris Leiva <loris.leiva@gmail.com>
Date: Fri, 1 Nov 2024 16:06:46 +0000
Subject: [PATCH] Fix LinkNode paths for unwrapDefinedTypesVisitor

---
 .changeset/healthy-zoos-share.md              |  5 ++
 .../visitors/src/unwrapDefinedTypesVisitor.ts | 11 ++-
 .../test/unwrapDefinedTypesVisitor.test.ts    | 74 +++++++++++++++++++
 3 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/healthy-zoos-share.md
 create mode 100644 packages/visitors/test/unwrapDefinedTypesVisitor.test.ts

diff --git a/.changeset/healthy-zoos-share.md b/.changeset/healthy-zoos-share.md
new file mode 100644
index 000000000..c32c58f2f
--- /dev/null
+++ b/.changeset/healthy-zoos-share.md
@@ -0,0 +1,5 @@
+---
+'@codama/visitors': minor
+---
+
+Fix LinkNode paths for `unwrapDefinedTypesVisitor`
diff --git a/packages/visitors/src/unwrapDefinedTypesVisitor.ts b/packages/visitors/src/unwrapDefinedTypesVisitor.ts
index 5e9d9f078..56da06df2 100644
--- a/packages/visitors/src/unwrapDefinedTypesVisitor.ts
+++ b/packages/visitors/src/unwrapDefinedTypesVisitor.ts
@@ -1,6 +1,7 @@
 import { assertIsNodeFilter, camelCase, CamelCaseString, programNode } from '@codama/nodes';
 import {
     extendVisitor,
+    getLastNodeFromPath,
     LinkableDictionary,
     NodeStack,
     nonNullableIdentityVisitor,
@@ -25,9 +26,13 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
                     if (!shouldInline(linkType.name)) {
                         return linkType;
                     }
-                    const definedType = linkables.getOrThrow(stack.getPath('definedTypeLinkNode'));
-                    // FIXME: Wrap in heap.pushStack() and heap.popStack().
-                    return visit(definedType.type, self);
+                    const definedTypePath = linkables.getPathOrThrow(stack.getPath('definedTypeLinkNode'));
+                    const definedType = getLastNodeFromPath(definedTypePath);
+
+                    stack.pushPath(definedTypePath);
+                    const result = visit(definedType.type, self);
+                    stack.popPath();
+                    return result;
                 },
 
                 visitProgram(program, { self }) {
diff --git a/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts b/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts
new file mode 100644
index 000000000..5810abe74
--- /dev/null
+++ b/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts
@@ -0,0 +1,74 @@
+import {
+    accountNode,
+    assertIsNode,
+    definedTypeLinkNode,
+    definedTypeNode,
+    numberTypeNode,
+    programLinkNode,
+    programNode,
+    rootNode,
+    structFieldTypeNode,
+    structTypeNode,
+} from '@codama/nodes';
+import { visit } from '@codama/visitors-core';
+import { expect, test } from 'vitest';
+
+import { unwrapDefinedTypesVisitor } from '../src';
+
+test('it unwraps defined types by following links', () => {
+    // Given a program node with an account that uses a defined type link.
+    const node = programNode({
+        accounts: [
+            accountNode({
+                data: structTypeNode([structFieldTypeNode({ name: 'value', type: definedTypeLinkNode('myType') })]),
+                name: 'myAccount',
+            }),
+        ],
+        definedTypes: [definedTypeNode({ name: 'myType', type: numberTypeNode('u64') })],
+        name: 'myProgram',
+        publicKey: '1111',
+    });
+
+    // When we unwrap the defined types.
+    const result = visit(node, unwrapDefinedTypesVisitor(['myType']));
+
+    // Then we expect the following tree.
+    assertIsNode(result, 'programNode');
+    expect(result.accounts[0].data).toStrictEqual(
+        structTypeNode([structFieldTypeNode({ name: 'value', type: numberTypeNode('u64') })]),
+    );
+});
+
+test('it follows linked nodes using the correct paths', () => {
+    // Given two link nodes designed so that the path would
+    // fail if we did not save and restored linked paths.
+    const programA = programNode({
+        definedTypes: [
+            definedTypeNode({
+                name: 'typeA',
+                type: definedTypeLinkNode('typeB1', programLinkNode('programB')),
+            }),
+        ],
+        name: 'programA',
+        publicKey: '1111',
+    });
+    const programB = programNode({
+        definedTypes: [
+            definedTypeNode({ name: 'typeB1', type: definedTypeLinkNode('typeB2') }),
+            definedTypeNode({ name: 'typeB2', type: numberTypeNode('u64') }),
+        ],
+        name: 'programB',
+        publicKey: '2222',
+    });
+    const root = rootNode(programA, [programB]);
+
+    // When we unwrap the defined types in programB.
+    const visitor = unwrapDefinedTypesVisitor(['typeB1', 'typeB2']);
+    const result = visit(root, visitor);
+
+    // Then we expect the final linkable to be resolved in programA.
+    assertIsNode(result, 'rootNode');
+    expect(result.program.definedTypes[0]).toStrictEqual(
+        definedTypeNode({ name: 'typeA', type: numberTypeNode('u64') }),
+    );
+});