Skip to content

Commit 0bb08e6

Browse files
lorisleivalorisleiva
authored
Add getMaxByteSizeVisitor helper (#840)
1 parent 8638fa7 commit 0bb08e6

File tree

6 files changed

+771
-7
lines changed

6 files changed

+771
-7
lines changed

‎.changeset/purple-sites-behave.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@codama/visitors-core': patch
3+
---
4+
5+
Add `getMaxByteSizeVisitor` helper

‎packages/visitors-core/src/getByteSizeVisitor.ts‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -89,16 +89,12 @@ export function getByteSizeVisitor(
8989

9090
visitDefinedTypeLink(node, { self }) {
9191
// Fetch the linked type and return null if not found.
92-
// The validator visitor will throw a proper error later on.
9392
const linkedDefinedPath = linkables.getPath(stack.getPath(node.kind));
9493
if (!linkedDefinedPath) return null;
9594
const linkedDefinedType = getLastNodeFromPath(linkedDefinedPath);
9695

97-
// This prevents infinite recursion by using assuming
98-
// cyclic types don't have a fixed size.
99-
if (definedTypeStack.includes(linkedDefinedType.name)) {
100-
return null;
101-
}
96+
// This prevents infinite recursion by assuming cyclic types don't have a fixed size.
97+
if (definedTypeStack.includes(linkedDefinedType.name)) return null;
10298

10399
stack.pushPath(linkedDefinedPath);
104100
const result = visit(linkedDefinedType, self);
@@ -111,7 +107,8 @@ export function getByteSizeVisitor(
111107
},
112108

113109
visitEnumType(node, { self }) {
114-
const prefix = visit(node.size, self) ?? 1;
110+
const prefix = visit(node.size, self);
111+
if (prefix === null) return null;
115112
if (isScalarEnum(node)) return prefix;
116113
const variantSizes = node.variants.map(v => visit(v, self));
117114
const allVariantHaveTheSameFixedSize = variantSizes.every((one, _, all) => one === all[0]);

‎packages/visitors-core/src/getMaxByteSizeVisitor.ts‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
import { CountNode, isNode, isScalarEnum, REGISTERED_TYPE_NODE_KINDS } from '@codama/nodes';
2+
3+
import { extendVisitor } from './extendVisitor';
4+
import { ByteSizeVisitorKeys } from './getByteSizeVisitor';
5+
import { LinkableDictionary } from './LinkableDictionary';
6+
import { mergeVisitor } from './mergeVisitor';
7+
import { getLastNodeFromPath } from './NodePath';
8+
import { NodeStack } from './NodeStack';
9+
import { pipe } from './pipe';
10+
import { recordNodeStackVisitor } from './recordNodeStackVisitor';
11+
import { visit, Visitor } from './visitor';
12+
13+
export function getMaxByteSizeVisitor(
14+
linkables: LinkableDictionary,
15+
options: { stack?: NodeStack } = {},
16+
): Visitor<number | null, ByteSizeVisitorKeys> {
17+
const stack = options.stack ?? new NodeStack();
18+
19+
const visitedDefinedTypes = new Map<string, number | null>();
20+
const definedTypeStack: string[] = [];
21+
22+
const sumSizes = (values: (number | null)[]): number | null =>
23+
values.reduce((all, one) => (all === null || one === null ? null : all + one), 0 as number | null);
24+
25+
const baseVisitor = mergeVisitor(
26+
() => null as number | null,
27+
(_, values) => sumSizes(values),
28+
{
29+
keys: [
30+
...REGISTERED_TYPE_NODE_KINDS,
31+
'accountNode',
32+
'constantValueNode',
33+
'definedTypeLinkNode',
34+
'definedTypeNode',
35+
'instructionArgumentNode',
36+
'instructionNode',
37+
],
38+
},
39+
);
40+
41+
return pipe(
42+
baseVisitor,
43+
v =>
44+
extendVisitor(v, {
45+
visitAccount(node, { self }) {
46+
return visit(node.data, self);
47+
},
48+
49+
visitArrayType(node, { self }) {
50+
return getArrayLikeSize(node.count, visit(node.item, self), self);
51+
},
52+
53+
visitConstantValue(node, { self }) {
54+
const typeSize = visit(node.type, self);
55+
if (typeSize !== null) return typeSize;
56+
if (isNode(node.value, 'bytesValueNode') && node.value.encoding === 'base16') {
57+
return Math.ceil(node.value.data.length / 2);
58+
}
59+
if (
60+
isNode(node.type, 'stringTypeNode') &&
61+
node.type.encoding === 'base16' &&
62+
isNode(node.value, 'stringValueNode')
63+
) {
64+
return Math.ceil(node.value.string.length / 2);
65+
}
66+
// Technically, we could still identify other fixed-size constants
67+
// but we'd need to import @solana/codecs to compute them.
68+
return null;
69+
},
70+
71+
visitDefinedType(node, { self }) {
72+
if (visitedDefinedTypes.has(node.name)) {
73+
return visitedDefinedTypes.get(node.name)!;
74+
}
75+
definedTypeStack.push(node.name);
76+
const child = visit(node.type, self);
77+
definedTypeStack.pop();
78+
visitedDefinedTypes.set(node.name, child);
79+
return child;
80+
},
81+
82+
visitDefinedTypeLink(node, { self }) {
83+
// Fetch the linked type and return null if not found.
84+
const linkedDefinedPath = linkables.getPath(stack.getPath(node.kind));
85+
if (!linkedDefinedPath) return null;
86+
const linkedDefinedType = getLastNodeFromPath(linkedDefinedPath);
87+
88+
// This prevents infinite recursion by assuming cyclic types don't have a fixed size.
89+
if (definedTypeStack.includes(linkedDefinedType.name)) return null;
90+
91+
stack.pushPath(linkedDefinedPath);
92+
const result = visit(linkedDefinedType, self);
93+
stack.popPath();
94+
return result;
95+
},
96+
97+
visitEnumEmptyVariantType() {
98+
return 0;
99+
},
100+
101+
visitEnumType(node, { self }) {
102+
const prefix = visit(node.size, self);
103+
if (prefix === null) return null;
104+
if (isScalarEnum(node)) return prefix;
105+
const variantSizes = node.variants.map(v => visit(v, self));
106+
if (variantSizes.includes(null)) return null;
107+
const maxVariantSize = Math.max(...(variantSizes as number[]));
108+
return prefix + maxVariantSize;
109+
},
110+
111+
visitFixedSizeType(node) {
112+
return node.size;
113+
},
114+
115+
visitInstruction(node, { self }) {
116+
return sumSizes(node.arguments.map(arg => visit(arg, self)));
117+
},
118+
119+
visitInstructionArgument(node, { self }) {
120+
return visit(node.type, self);
121+
},
122+
123+
visitMapType(node, { self }) {
124+
const innerSize = sumSizes([visit(node.key, self), visit(node.value, self)]);
125+
return getArrayLikeSize(node.count, innerSize, self);
126+
},
127+
128+
visitNumberType(node) {
129+
if (node.format === 'shortU16') return 3;
130+
return parseInt(node.format.slice(1), 10) / 8;
131+
},
132+
133+
visitOptionType(node, { self }) {
134+
return sumSizes([visit(node.prefix, self), visit(node.item, self)]);
135+
},
136+
137+
visitPostOffsetType(node, { self }) {
138+
const typeSize = visit(node.type, self);
139+
return node.strategy === 'padded' ? sumSizes([typeSize, node.offset]) : typeSize;
140+
},
141+
142+
visitPreOffsetType(node, { self }) {
143+
const typeSize = visit(node.type, self);
144+
return node.strategy === 'padded' ? sumSizes([typeSize, node.offset]) : typeSize;
145+
},
146+
147+
visitPublicKeyType() {
148+
return 32;
149+
},
150+
151+
visitRemainderOptionType(node, { self }) {
152+
const itemSize = visit(node.item, self);
153+
return itemSize === 0 ? 0 : null;
154+
},
155+
156+
visitSetType(node, { self }) {
157+
return getArrayLikeSize(node.count, visit(node.item, self), self);
158+
},
159+
160+
visitZeroableOptionType(node, { self }) {
161+
const itemSize = visit(node.item, self);
162+
if (!node.zeroValue) return itemSize;
163+
const zeroSize = visit(node.zeroValue, self);
164+
if (itemSize === null || zeroSize === null) return null;
165+
return Math.max(itemSize, zeroSize);
166+
},
167+
}),
168+
v => recordNodeStackVisitor(v, stack),
169+
);
170+
}
171+
172+
function getArrayLikeSize(
173+
count: CountNode,
174+
innerSize: number | null,
175+
self: Visitor<number | null, ByteSizeVisitorKeys>,
176+
): number | null {
177+
if (innerSize === 0 && isNode(count, 'prefixedCountNode')) return visit(count.prefix, self);
178+
if (innerSize === 0) return 0;
179+
if (!isNode(count, 'fixedCountNode')) return null;
180+
if (count.value === 0) return 0;
181+
return innerSize !== null ? innerSize * count.value : null;
182+
}

‎packages/visitors-core/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ export * from './deleteNodesVisitor';
44
export * from './extendVisitor';
55
export * from './getByteSizeVisitor';
66
export * from './getDebugStringVisitor';
7+
export * from './getMaxByteSizeVisitor';
78
export * from './getResolvedInstructionInputsVisitor';
89
export * from './getUniqueHashStringVisitor';
910
export * from './identityVisitor';

‎packages/visitors-core/test/getByteSizeVisitor.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -401,6 +401,9 @@ describe('numberTypeNode', () => {
401401
])('it returns the size of %s numbers', (format, expectedSize) => {
402402
expectSize(numberTypeNode(format as NumberFormat), expectedSize);
403403
});
404+
test('it returns null if the format is shortU16', () => {
405+
expectSize(numberTypeNode('shortU16'), null);
406+
});
404407
});
405408

406409
describe('optionTypeNode', () => {

0 commit comments

Comments
 (0)