Use program names when unwrapping link nodes (#314)

lorisleiva · web-flow · commit e55bcd7ee746 · 2024-11-22T15:58:16.000Z
diff --git a/.changeset/giant-items-sniff.md b/.changeset/giant-items-sniff.md
@@ -0,0 +1,5 @@
+---
+'@codama/visitors': minor
+---
+
+Use program names when unwrapping link nodes
diff --git a/packages/visitors/src/getDefinedTypeHistogramVisitor.ts b/packages/visitors/src/getDefinedTypeHistogramVisitor.ts
@@ -1,8 +1,20 @@
 import { CamelCaseString } from '@codama/nodes';
-import { extendVisitor, interceptVisitor, mergeVisitor, pipe, visit, Visitor } from '@codama/visitors-core';
+import {
+    extendVisitor,
+    findProgramNodeFromPath,
+    interceptVisitor,
+    mergeVisitor,
+    NodeStack,
+    pipe,
+    recordNodeStackVisitor,
+    visit,
+    Visitor,
+} from '@codama/visitors-core';
+
+type DefinedTypeHistogramKey = CamelCaseString | `${CamelCaseString}.${CamelCaseString}`;
 
 export type DefinedTypeHistogram = {
-    [key: CamelCaseString]: {
+    [key: DefinedTypeHistogramKey]: {
         directlyAsInstructionArgs: number;
         inAccounts: number;
         inDefinedTypes: number;
@@ -33,6 +45,7 @@ function mergeHistograms(histograms: DefinedTypeHistogram[]): DefinedTypeHistogr
 }
 
 export function getDefinedTypeHistogramVisitor(): Visitor<DefinedTypeHistogram> {
+    const stack = new NodeStack();
     let mode: 'account' | 'definedType' | 'instruction' | null = null;
     let stackLevel = 0;
 
@@ -67,8 +80,10 @@ export function getDefinedTypeHistogramVisitor(): Visitor<DefinedTypeHistogram>
                 },
 
                 visitDefinedTypeLink(node) {
+                    const program = findProgramNodeFromPath(stack.getPath());
+                    const key = program ? `${program.name}.${node.name}` : node.name;
                     return {
-                        [node.name]: {
+                        [key]: {
                             directlyAsInstructionArgs: Number(mode === 'instruction' && stackLevel <= 1),
                             inAccounts: Number(mode === 'account'),
                             inDefinedTypes: Number(mode === 'definedType'),
@@ -88,5 +103,6 @@ export function getDefinedTypeHistogramVisitor(): Visitor<DefinedTypeHistogram>
                     return mergeHistograms([...dataHistograms, ...extraHistograms, ...subHistograms]);
                 },
             }),
+        v => recordNodeStackVisitor(v, stack),
     );
 }
diff --git 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,
+    findProgramNodeFromPath,
     getLastNodeFromPath,
     LinkableDictionary,
     NodeStack,
@@ -14,16 +15,25 @@ import {
 export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
     const linkables = new LinkableDictionary();
     const stack = new NodeStack();
-    const typesToInlineMainCased = typesToInline === '*' ? '*' : typesToInline.map(camelCase);
-    const shouldInline = (definedType: CamelCaseString): boolean =>
-        typesToInlineMainCased === '*' || typesToInlineMainCased.includes(definedType);
+    const typesToInlineCamelCased = (typesToInline === '*' ? [] : typesToInline).map(fullPath => {
+        if (!fullPath.includes('.')) return camelCase(fullPath);
+        const [programName, typeName] = fullPath.split('.');
+        return `${camelCase(programName)}.${camelCase(typeName)}`;
+    });
+    const shouldInline = (typeName: CamelCaseString, programName: CamelCaseString | undefined): boolean => {
+        if (typesToInline === '*') return true;
+        const fullPath = `${programName}.${typeName}`;
+        if (!!programName && typesToInlineCamelCased.includes(fullPath)) return true;
+        return typesToInlineCamelCased.includes(typeName);
+    };
 
     return pipe(
         nonNullableIdentityVisitor(),
         v =>
             extendVisitor(v, {
                 visitDefinedTypeLink(linkType, { self }) {
-                    if (!shouldInline(linkType.name)) {
+                    const programName = linkType.program?.name ?? findProgramNodeFromPath(stack.getPath())?.name;
+                    if (!shouldInline(linkType.name, programName)) {
                         return linkType;
                     }
                     const definedTypePath = linkables.getPathOrThrow(stack.getPath('definedTypeLinkNode'));
@@ -42,7 +52,7 @@ export function unwrapDefinedTypesVisitor(typesToInline: string[] | '*' = '*') {
                             .map(account => visit(account, self))
                             .filter(assertIsNodeFilter('accountNode')),
                         definedTypes: program.definedTypes
-                            .filter(definedType => !shouldInline(definedType.name))
+                            .filter(definedType => !shouldInline(definedType.name, program.name))
                             .map(type => visit(type, self))
                             .filter(assertIsNodeFilter('definedTypeNode')),
                         instructions: program.instructions
diff --git a/packages/visitors/src/unwrapInstructionArgsDefinedTypesVisitor.ts b/packages/visitors/src/unwrapInstructionArgsDefinedTypesVisitor.ts
@@ -1,24 +1,23 @@
-import { assertIsNode, CamelCaseString, getAllDefinedTypes, isNode } from '@codama/nodes';
-import { rootNodeVisitor, visit } from '@codama/visitors-core';
+import { assertIsNode, CamelCaseString, definedTypeLinkNode, isNode } from '@codama/nodes';
+import { getRecordLinkablesVisitor, LinkableDictionary, rootNodeVisitor, visit } from '@codama/visitors-core';
 
 import { getDefinedTypeHistogramVisitor } from './getDefinedTypeHistogramVisitor';
 import { unwrapDefinedTypesVisitor } from './unwrapDefinedTypesVisitor';
 
 export function unwrapInstructionArgsDefinedTypesVisitor() {
     return rootNodeVisitor(root => {
         const histogram = visit(root, getDefinedTypeHistogramVisitor());
-        const allDefinedTypes = getAllDefinedTypes(root);
+        const linkables = new LinkableDictionary();
+        visit(root, getRecordLinkablesVisitor(linkables));
 
-        const definedTypesToInline: string[] = Object.keys(histogram)
+        const definedTypesToInline = (Object.keys(histogram) as CamelCaseString[])
             // Get all defined types used exactly once as an instruction argument.
-            .filter(
-                name =>
-                    (histogram[name as CamelCaseString].total ?? 0) === 1 &&
-                    (histogram[name as CamelCaseString].directlyAsInstructionArgs ?? 0) === 1,
-            )
+            .filter(key => (histogram[key].total ?? 0) === 1 && (histogram[key].directlyAsInstructionArgs ?? 0) === 1)
             // Filter out enums which are better defined as external types.
-            .filter(name => {
-                const found = allDefinedTypes.find(type => type.name === name);
+            .filter(key => {
+                const names = key.split('.');
+                const link = names.length == 2 ? definedTypeLinkNode(names[1], names[0]) : definedTypeLinkNode(key);
+                const found = linkables.get([link]);
                 return found && !isNode(found.type, 'enumTypeNode');
             });
 
diff --git a/packages/visitors/test/getDefinedTypeHistogramVisitor.test.ts b/packages/visitors/test/getDefinedTypeHistogramVisitor.test.ts
@@ -5,7 +5,9 @@ import {
     enumTypeNode,
     instructionArgumentNode,
     instructionNode,
+    numberTypeNode,
     programNode,
+    rootNode,
     structFieldTypeNode,
     structTypeNode,
 } from '@codama/nodes';
@@ -65,14 +67,14 @@ test('it counts the amount of times defined types are used within the tree', ()
 
     // Then we expect the following histogram.
     expect(histogram).toEqual({
-        myEnum: {
+        'customProgram.myEnum': {
             directlyAsInstructionArgs: 0,
             inAccounts: 1,
             inDefinedTypes: 0,
             inInstructionArgs: 0,
             total: 1,
         },
-        myStruct: {
+        'customProgram.myStruct': {
             directlyAsInstructionArgs: 1,
             inAccounts: 1,
             inDefinedTypes: 0,
@@ -81,3 +83,47 @@ test('it counts the amount of times defined types are used within the tree', ()
         },
     });
 });
+
+test('it counts links from different programs separately', () => {
+    // Given a program node with a defined type used in another type.
+    const programA = programNode({
+        definedTypes: [
+            definedTypeNode({ name: 'myType', type: numberTypeNode('u8') }),
+            definedTypeNode({ name: 'myCopyType', type: definedTypeLinkNode('myType') }),
+        ],
+        name: 'programA',
+        publicKey: '1111',
+    });
+
+    // And another program with a defined type sharing the same name.
+    const programB = programNode({
+        definedTypes: [
+            definedTypeNode({ name: 'myType', type: numberTypeNode('u16') }),
+            definedTypeNode({ name: 'myCopyType', type: definedTypeLinkNode('myType') }),
+        ],
+        name: 'programB',
+        publicKey: '2222',
+    });
+
+    // When we unwrap the defined type from programA.
+    const node = rootNode(programA, [programB]);
+    const histogram = visit(node, getDefinedTypeHistogramVisitor());
+
+    // Then we expect programA to have been modified but not programB.
+    expect(histogram).toStrictEqual({
+        'programA.myType': {
+            directlyAsInstructionArgs: 0,
+            inAccounts: 0,
+            inDefinedTypes: 1,
+            inInstructionArgs: 0,
+            total: 1,
+        },
+        'programB.myType': {
+            directlyAsInstructionArgs: 0,
+            inAccounts: 0,
+            inDefinedTypes: 1,
+            inInstructionArgs: 0,
+            total: 1,
+        },
+    });
+});
diff --git a/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts b/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts
@@ -72,3 +72,42 @@ test('it follows linked nodes using the correct paths', () => {
         definedTypeNode({ name: 'typeA', type: numberTypeNode('u64') }),
     );
 });
+
+test('it does not unwrap types from the wrong programs', () => {
+    // Given a program node with a defined type used in another type.
+    const programA = programNode({
+        definedTypes: [
+            definedTypeNode({ name: 'myType', type: numberTypeNode('u8') }),
+            definedTypeNode({ name: 'myCopyType', type: definedTypeLinkNode('myType') }),
+        ],
+        name: 'programA',
+        publicKey: '1111',
+    });
+
+    // And another program with a defined type sharing the same name.
+    const programB = programNode({
+        definedTypes: [
+            definedTypeNode({ name: 'myType', type: numberTypeNode('u16') }),
+            definedTypeNode({ name: 'myCopyType', type: definedTypeLinkNode('myType') }),
+        ],
+        name: 'programB',
+        publicKey: '2222',
+    });
+
+    // When we unwrap the defined type from programA.
+    const node = rootNode(programA, [programB]);
+    const result = visit(node, unwrapDefinedTypesVisitor(['programA.myType']));
+
+    // Then we expect programA to have been modified but not programB.
+    assertIsNode(result, 'rootNode');
+    expect(result).toStrictEqual(
+        rootNode(
+            programNode({
+                definedTypes: [definedTypeNode({ name: 'myCopyType', type: numberTypeNode('u8') })],
+                name: 'programA',
+                publicKey: '1111',
+            }),
+            [programB],
+        ),
+    );
+});
diff --git a/packages/visitors/test/unwrapInstructionArgsDefinedTypesVisitor.test.ts b/packages/visitors/test/unwrapInstructionArgsDefinedTypesVisitor.test.ts
diff --git a/vitest.workspace.ts b/vitest.workspace.ts

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'@codama/visitors': minor
 +---
++
 +Use program names when unwrapping link nodes