diff --git a/packages/visitors/src/getDefinedTypeHistogramVisitor.ts b/packages/visitors/src/getDefinedTypeHistogramVisitor.ts index 067800c2..008bb764 100644 --- 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 { + const stack = new NodeStack(); let mode: 'account' | 'definedType' | 'instruction' | null = null; let stackLevel = 0; @@ -67,8 +80,10 @@ export function getDefinedTypeHistogramVisitor(): Visitor }, 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 return mergeHistograms([...dataHistograms, ...extraHistograms, ...subHistograms]); }, }), + v => recordNodeStackVisitor(v, stack), ); } diff --git a/packages/visitors/src/unwrapDefinedTypesVisitor.ts b/packages/visitors/src/unwrapDefinedTypesVisitor.ts index 56da06df..4e9b0cf1 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, + 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 index 0e9d2395..ec2faa6c 100644 --- a/packages/visitors/src/unwrapInstructionArgsDefinedTypesVisitor.ts +++ b/packages/visitors/src/unwrapInstructionArgsDefinedTypesVisitor.ts @@ -1,5 +1,5 @@ -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'; @@ -7,18 +7,17 @@ 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 index 85c6bed0..a429f544 100644 --- 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 index db57a991..4e27c00a 100644 --- a/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts +++ b/packages/visitors/test/unwrapDefinedTypesVisitor.test.ts @@ -96,7 +96,7 @@ test('it does not unwrap types from the wrong programs', () => { // When we unwrap the defined type from programA. const node = rootNode(programA, [programB]); - const result = visit(node, unwrapDefinedTypesVisitor(['myType'])); + const result = visit(node, unwrapDefinedTypesVisitor(['programA.myType'])); // Then we expect programA to have been modified but not programB. assertIsNode(result, 'rootNode');