From 8d6a038a48c919c7b1a09d7680aaf7cf58c0decd Mon Sep 17 00:00:00 2001 From: Loris Leiva Date: Fri, 1 Nov 2024 15:11:11 +0000 Subject: [PATCH] Fix LinkNode paths for getByteSizeVisitor --- .changeset/soft-beds-jam.md | 5 +++ .../src/getRenderMapVisitor.ts | 2 +- .../visitors-core/src/getByteSizeVisitor.ts | 14 ++++--- .../test/getByteSizeVisitor.test.ts | 42 ++++++++++++++++++- .../src/setFixedAccountSizesVisitor.ts | 14 ++----- 5 files changed, 58 insertions(+), 19 deletions(-) create mode 100644 .changeset/soft-beds-jam.md diff --git a/.changeset/soft-beds-jam.md b/.changeset/soft-beds-jam.md new file mode 100644 index 000000000..efd49d592 --- /dev/null +++ b/.changeset/soft-beds-jam.md @@ -0,0 +1,5 @@ +--- +'@codama/visitors-core': minor +--- + +Fix LinkNode paths for `getByteSizeVisitor` diff --git a/packages/renderers-js-umi/src/getRenderMapVisitor.ts b/packages/renderers-js-umi/src/getRenderMapVisitor.ts index ad5cddd59..a3862a9c8 100644 --- a/packages/renderers-js-umi/src/getRenderMapVisitor.ts +++ b/packages/renderers-js-umi/src/getRenderMapVisitor.ts @@ -63,7 +63,6 @@ export type GetRenderMapOptions = { export function getRenderMapVisitor(options: GetRenderMapOptions = {}): Visitor { const linkables = new LinkableDictionary(); const stack = new NodeStack(); - const byteSizeVisitor = getByteSizeVisitor(linkables, stack); let program: ProgramNode | null = null; const renderParentInstructions = options.renderParentInstructions ?? false; @@ -100,6 +99,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}): Visitor< }); const typeManifestVisitor = getTypeManifestVisitor(); const resolvedInstructionInputVisitor = getResolvedInstructionInputsVisitor(); + const byteSizeVisitor = getByteSizeVisitor(linkables, stack); function getInstructionAccountType(account: ResolvedInstructionAccount): string { if (account.isPda && account.isSigner === false) return 'Pda'; diff --git a/packages/visitors-core/src/getByteSizeVisitor.ts b/packages/visitors-core/src/getByteSizeVisitor.ts index e0ca181c5..743da04ca 100644 --- a/packages/visitors-core/src/getByteSizeVisitor.ts +++ b/packages/visitors-core/src/getByteSizeVisitor.ts @@ -3,6 +3,7 @@ import { isNode, isScalarEnum, REGISTERED_TYPE_NODE_KINDS, RegisteredTypeNode } import { extendVisitor } from './extendVisitor'; import { LinkableDictionary } from './LinkableDictionary'; import { mergeVisitor } from './mergeVisitor'; +import { getLastNodeFromPath } from './NodePath'; import { NodeStack } from './NodeStack'; import { pipe } from './pipe'; import { recordNodeStackVisitor } from './recordNodeStackVisitor'; @@ -69,11 +70,9 @@ export function getByteSizeVisitor( visitDefinedTypeLink(node, { self }) { // Fetch the linked type and return null if not found. // The validator visitor will throw a proper error later on. - // FIXME: Keep track of our own internal stack within this visitor (starting from a provided NodePath). - const linkedDefinedType = linkables.get([...stack.getPath(), node]); - if (!linkedDefinedType) { - return null; - } + const linkedDefinedPath = linkables.getPath(stack.getPath(node.kind)); + if (!linkedDefinedPath) return null; + const linkedDefinedType = getLastNodeFromPath(linkedDefinedPath); // This prevents infinite recursion by using assuming // cyclic types don't have a fixed size. @@ -81,7 +80,10 @@ export function getByteSizeVisitor( return null; } - return visit(linkedDefinedType, self); + stack.pushPath(linkedDefinedPath); + const result = visit(linkedDefinedType, self); + stack.popPath(); + return result; }, visitEnumEmptyVariantType() { diff --git a/packages/visitors-core/test/getByteSizeVisitor.test.ts b/packages/visitors-core/test/getByteSizeVisitor.test.ts index 84c87d3fe..b58abe503 100644 --- a/packages/visitors-core/test/getByteSizeVisitor.test.ts +++ b/packages/visitors-core/test/getByteSizeVisitor.test.ts @@ -1,4 +1,6 @@ import { + definedTypeLinkNode, + definedTypeNode, enumEmptyVariantTypeNode, enumStructVariantTypeNode, enumTupleVariantTypeNode, @@ -7,7 +9,10 @@ import { Node, NumberFormat, numberTypeNode, + programLinkNode, + programNode, publicKeyTypeNode, + rootNode, stringTypeNode, structFieldTypeNode, structTypeNode, @@ -15,7 +20,7 @@ import { } from '@codama/nodes'; import { expect, test } from 'vitest'; -import { getByteSizeVisitor, LinkableDictionary, NodeStack, visit, Visitor } from '../src'; +import { getByteSizeVisitor, getRecordLinkablesVisitor, LinkableDictionary, NodeStack, visit, Visitor } from '../src'; const expectSize = (node: Node, expectedSize: number | null) => { expect(visit(node, getByteSizeVisitor(new LinkableDictionary(), new NodeStack()) as Visitor)).toBe( @@ -104,3 +109,38 @@ test('it gets the size of variable data enums', () => { null, ); }); + +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]); + + // And given a recorded linkables dictionary. + const linkables = new LinkableDictionary(); + visit(root, getRecordLinkablesVisitor(linkables)); + + // When we visit the first defined type. + const visitor = getByteSizeVisitor(linkables, new NodeStack([root, programA])); + const result = visit(programA.definedTypes[0], visitor); + + // Then we expect the final linkable to be resolved. + expect(result).toBe(8); +}); diff --git a/packages/visitors/src/setFixedAccountSizesVisitor.ts b/packages/visitors/src/setFixedAccountSizesVisitor.ts index 95cf5d4f1..90db9e6f6 100644 --- a/packages/visitors/src/setFixedAccountSizesVisitor.ts +++ b/packages/visitors/src/setFixedAccountSizesVisitor.ts @@ -4,26 +4,22 @@ import { getLastNodeFromPath, isNodePath, LinkableDictionary, - NodeStack, pipe, recordLinkablesOnFirstVisitVisitor, - recordNodeStackVisitor, topDownTransformerVisitor, visit, } from '@codama/visitors-core'; export function setFixedAccountSizesVisitor() { const linkables = new LinkableDictionary(); - const stack = new NodeStack(); - const byteSizeVisitor = getByteSizeVisitor(linkables, stack); const visitor = topDownTransformerVisitor( [ { select: path => isNodePath(path, 'accountNode') && getLastNodeFromPath(path).size === undefined, - transform: node => { + transform: (node, stack) => { assertIsNode(node, 'accountNode'); - const size = visit(node.data, byteSizeVisitor); + const size = visit(node.data, getByteSizeVisitor(linkables, stack)); if (size === null) return node; return accountNode({ ...node, size }) as typeof node; }, @@ -32,9 +28,5 @@ export function setFixedAccountSizesVisitor() { ['rootNode', 'programNode', 'accountNode'], ); - return pipe( - visitor, - v => recordNodeStackVisitor(v, stack), - v => recordLinkablesOnFirstVisitVisitor(v, linkables), - ); + return pipe(visitor, v => recordLinkablesOnFirstVisitVisitor(v, linkables)); }