Fix size computation of array-like nodes (#731)

Author: lorisleiva

.changeset/salty-beans-check.md

@@ -0,0 +1,5 @@
+---
+'@codama/visitors-core': patch
+---
+
+Fix size computation of array-like nodes (arrays, sets and maps).

packages/visitors-core/src/getByteSizeVisitor.ts

@@ -54,10 +54,10 @@ export function getByteSizeVisitor(
 
                 visitArrayType(node, { self }) {
                     if (!isNode(node.count, 'fixedCountNode')) return null;
-                    const fixedSize = node.count.value;
+                    const count = node.count.value;
+                    if (count === 0) return 0;
                     const itemSize = visit(node.item, self);
-                    const arraySize = itemSize !== null ? itemSize * fixedSize : null;
-                    return fixedSize === 0 ? 0 : arraySize;
+                    return itemSize !== null ? itemSize * count : null;
                 },
 
                 visitDefinedType(node, { self }) {
@@ -116,6 +116,15 @@ export function getByteSizeVisitor(
                     return visit(node.type, self);
                 },
 
+                visitMapType(node, { self }) {
+                    if (!isNode(node.count, 'fixedCountNode')) return null;
+                    const count = node.count.value;
+                    if (count === 0) return 0;
+                    const keySize = visit(node.key, self);
+                    const valueSize = visit(node.value, self);
+                    return keySize !== null && valueSize !== null ? (keySize + valueSize) * count : null;
+                },
+
                 visitNumberType(node) {
                     if (node.format === 'shortU16') return null;
                     return parseInt(node.format.slice(1), 10) / 8;
@@ -131,6 +140,14 @@ export function getByteSizeVisitor(
                 visitPublicKeyType() {
                     return 32;
                 },
+
+                visitSetType(node, { self }) {
+                    if (!isNode(node.count, 'fixedCountNode')) return null;
+                    const count = node.count.value;
+                    if (count === 0) return 0;
+                    const itemSize = visit(node.item, self);
+                    return itemSize !== null ? itemSize * count : null;
+                },
             }),
         v => recordNodeStackVisitor(v, stack),
     );

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

@@ -1,18 +1,24 @@
 import {
+    arrayTypeNode,
     definedTypeLinkNode,
     definedTypeNode,
     enumEmptyVariantTypeNode,
     enumStructVariantTypeNode,
     enumTupleVariantTypeNode,
     enumTypeNode,
+    fixedCountNode,
     fixedSizeTypeNode,
+    mapTypeNode,
     Node,
     NumberFormat,
     numberTypeNode,
+    prefixedCountNode,
     programLinkNode,
     programNode,
     publicKeyTypeNode,
+    remainderCountNode,
     rootNode,
+    setTypeNode,
     stringTypeNode,
     structFieldTypeNode,
     structTypeNode,
@@ -108,6 +114,27 @@ test('it gets the size of variable data enums', () => {
     );
 });
 
+test('it gets the size of fixed size array-like nodes with `FixCountNodes`', () => {
+    expectSize(arrayTypeNode(numberTypeNode('u16'), fixedCountNode(3)), 2 * 3);
+    expectSize(setTypeNode(numberTypeNode('u16'), fixedCountNode(3)), 2 * 3);
+    expectSize(mapTypeNode(numberTypeNode('u16'), numberTypeNode('u8'), fixedCountNode(3)), (2 + 1) * 3);
+});
+
+test('it returns null for array-like nodes with `PrefixedCountNodes`', () => {
+    expectSize(arrayTypeNode(numberTypeNode('u16'), prefixedCountNode(numberTypeNode('u32'))), null);
+    expectSize(setTypeNode(numberTypeNode('u16'), prefixedCountNode(numberTypeNode('u32'))), null);
+    expectSize(
+        mapTypeNode(numberTypeNode('u16'), numberTypeNode('u8'), prefixedCountNode(numberTypeNode('u32'))),
+        null,
+    );
+});
+
+test('it returns null for array-like nodes with `RemainderCountNodes`', () => {
+    expectSize(arrayTypeNode(numberTypeNode('u16'), remainderCountNode()), null);
+    expectSize(setTypeNode(numberTypeNode('u16'), remainderCountNode()), null);
+    expectSize(mapTypeNode(numberTypeNode('u16'), numberTypeNode('u8'), remainderCountNode()), 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.