From dbc6281fb0cf9aeb71be7c2d96d23208e1d4b42f Mon Sep 17 00:00:00 2001
From: Loris Leiva <loris.leiva@gmail.com>
Date: Tue, 9 Apr 2024 10:20:11 +0100
Subject: [PATCH] Add nested type guards

---
 src/nodes/typeNodes/NestedTypeNode.ts         | 12 +++++
 test/nodes/typeNodes/NestedTypeNode.test.ts   | 14 ++++++
 .../typeNodes/NestedTypeNode.typetest.ts      | 48 +++++++++++++++++++
 3 files changed, 74 insertions(+)

diff --git a/src/nodes/typeNodes/NestedTypeNode.ts b/src/nodes/typeNodes/NestedTypeNode.ts
index 01cc4ca6..5f13dd16 100644
--- a/src/nodes/typeNodes/NestedTypeNode.ts
+++ b/src/nodes/typeNodes/NestedTypeNode.ts
@@ -54,3 +54,15 @@ export function isNestedTypeNode<TKind extends TypeNode['kind']>(
   const resolved = resolveNestedTypeNode(node);
   return !!node && kinds.includes(resolved.kind as TKind);
 }
+
+export function assertIsNestedTypeNode<TKind extends TypeNode['kind']>(
+  node: Node | null | undefined,
+  kind: TKind | TKind[]
+): asserts node is NestedTypeNode<Extract<TypeNode, { kind: TKind }>> {
+  const kinds = Array.isArray(kind) ? kind : [kind];
+  if (!isNestedTypeNode(node, kinds)) {
+    throw new Error(
+      `Expected nested type of ${kinds.join(' | ')}, got ${node?.kind ?? 'null'}.`
+    );
+  }
+}
diff --git a/test/nodes/typeNodes/NestedTypeNode.test.ts b/test/nodes/typeNodes/NestedTypeNode.test.ts
index 0bb88546..6a46e0e1 100644
--- a/test/nodes/typeNodes/NestedTypeNode.test.ts
+++ b/test/nodes/typeNodes/NestedTypeNode.test.ts
@@ -1,5 +1,6 @@
 import test from 'ava';
 import {
+  assertIsNestedTypeNode,
   fixedSizeTypeNode,
   isNestedTypeNode,
   numberTypeNode,
@@ -52,3 +53,16 @@ test('it checks if a node is a nested type', (t) => {
   t.true(isNestedTypeNode(nestedNode, 'numberTypeNode'));
   t.false(isNestedTypeNode(nestedNode, 'stringTypeNode'));
 });
+
+test('it asserts that a node is a nested type', (t) => {
+  const flatNode = numberTypeNode('u64');
+  t.notThrows(() => assertIsNestedTypeNode(flatNode, 'numberTypeNode'));
+  t.throws(() => assertIsNestedTypeNode(flatNode, 'stringTypeNode'));
+
+  const nestedNode = sizePrefixTypeNode(
+    fixedSizeTypeNode(numberTypeNode('u64'), 32),
+    numberTypeNode('u8')
+  );
+  t.notThrows(() => assertIsNestedTypeNode(nestedNode, 'numberTypeNode'));
+  t.throws(() => assertIsNestedTypeNode(nestedNode, 'stringTypeNode'));
+});
diff --git a/test/nodes/typeNodes/NestedTypeNode.typetest.ts b/test/nodes/typeNodes/NestedTypeNode.typetest.ts
index 6bc76b10..9b2a9da3 100644
--- a/test/nodes/typeNodes/NestedTypeNode.typetest.ts
+++ b/test/nodes/typeNodes/NestedTypeNode.typetest.ts
@@ -6,6 +6,10 @@ import {
   NestedTypeNode,
   StringTypeNode,
   stringTypeNode,
+  transformNestedTypeNode,
+  Node,
+  isNestedTypeNode,
+  assertIsNestedTypeNode,
 } from '../../../src';
 
 {
@@ -31,3 +35,47 @@ import {
   // @ts-expect-error The nested type is not a string.
   resolveNestedTypeNode(numberNestedNode) satisfies StringTypeNode;
 }
+
+{
+  // [transformNestedTypeNode]: it transforms the nested type of a nested node.
+  const transformedNode = transformNestedTypeNode(
+    fixedSizeTypeNode(stringTypeNode(), 32),
+    () => numberTypeNode('u32')
+  );
+  transformedNode satisfies NestedTypeNode<NumberTypeNode>;
+  // @ts-expect-error The nested type is not a number.
+  transformedNode satisfies NestedTypeNode<StringTypeNode>;
+}
+
+{
+  // [isNestedTypeNode]: it narrows the type of a node to a nested type node.
+  const node = {} as Node;
+  if (isNestedTypeNode(node, 'numberTypeNode')) {
+    node satisfies NestedTypeNode<NumberTypeNode>;
+    // @ts-expect-error The nested type is not a string.
+    node satisfies NestedTypeNode<StringTypeNode>;
+  }
+  if (isNestedTypeNode(node, 'stringTypeNode')) {
+    node satisfies NestedTypeNode<StringTypeNode>;
+    // @ts-expect-error The nested type is not a number.
+    node satisfies NestedTypeNode<NumberTypeNode>;
+  }
+}
+
+{
+  // [assertIsNestedTypeNode]: it narrows the type of a node to a nested type node.
+  {
+    const node = {} as Node;
+    assertIsNestedTypeNode(node, 'numberTypeNode');
+    node satisfies NestedTypeNode<NumberTypeNode>;
+    // @ts-expect-error The nested type is not a string.
+    node satisfies NestedTypeNode<StringTypeNode>;
+  }
+  {
+    const node = {} as Node;
+    assertIsNestedTypeNode(node, 'stringTypeNode');
+    node satisfies NestedTypeNode<StringTypeNode>;
+    // @ts-expect-error The nested type is not a number.
+    node satisfies NestedTypeNode<NumberTypeNode>;
+  }
+}