From a2743475c0d7d6a1ba5f51f31405c93b9d940e6c Mon Sep 17 00:00:00 2001
From: Cayman <caymannava@gmail.com>
Date: Tue, 15 Oct 2024 16:52:17 -0400
Subject: [PATCH] chore: upgradeToNewType as utility function

---
 packages/ssz/src/index.ts                     |  1 +
 packages/ssz/src/util/upgrade.ts              | 25 +++++++++
 .../test/unit/byType/container/tree.test.ts   | 53 +++++++++++++++++++
 3 files changed, 79 insertions(+)
 create mode 100644 packages/ssz/src/util/upgrade.ts

diff --git a/packages/ssz/src/index.ts b/packages/ssz/src/index.ts
index ab3330c6..b3b69b26 100644
--- a/packages/ssz/src/index.ts
+++ b/packages/ssz/src/index.ts
@@ -39,3 +39,4 @@ export {BitArray, getUint8ByteToBitBooleanArray} from "./value/bitArray";
 export {fromHexString, toHexString, byteArrayEquals} from "./util/byteArray";
 export {Snapshot} from "./util/types";
 export {hash64, symbolCachedPermanentRoot} from "./util/merkleize";
+export {upgradeToNewType} from "./util/upgrade";
diff --git a/packages/ssz/src/util/upgrade.ts b/packages/ssz/src/util/upgrade.ts
new file mode 100644
index 00000000..992305df
--- /dev/null
+++ b/packages/ssz/src/util/upgrade.ts
@@ -0,0 +1,25 @@
+import {BranchNode, Node, zeroNode} from "@chainsafe/persistent-merkle-tree";
+import {ContainerType} from "../type/container";
+
+/** Upgrade the current View/ViewDU to a root node of new type */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function upgradeToNewType(rootNode: Node, oldType: ContainerType<any>, newType: ContainerType<any>): Node {
+  const newFieldsCount = newType.fieldsEntries.length;
+  const currentFieldsCount = oldType.fieldsEntries.length;
+  if (newFieldsCount < currentFieldsCount) {
+    throw Error(`Cannot convert to a type with fewer fields: ${newFieldsCount} < ${currentFieldsCount}`);
+  }
+
+  if (newType.depth === oldType.depth) {
+    // no need to grow the tree
+    return rootNode;
+  }
+
+  // grow the tree
+  let node = rootNode;
+  for (let depth = oldType.depth; depth < newType.depth; depth++) {
+    node = new BranchNode(node, zeroNode(depth));
+  }
+
+  return node;
+}
diff --git a/packages/ssz/test/unit/byType/container/tree.test.ts b/packages/ssz/test/unit/byType/container/tree.test.ts
index 6b545792..83fad27a 100644
--- a/packages/ssz/test/unit/byType/container/tree.test.ts
+++ b/packages/ssz/test/unit/byType/container/tree.test.ts
@@ -1,4 +1,5 @@
 import {expect} from "chai";
+import {Tree} from "@chainsafe/persistent-merkle-tree";
 import {
   BitArray,
   BitListType,
@@ -12,6 +13,7 @@ import {
   ListCompositeType,
   NoneType,
   toHexString,
+  Type,
   UintNumberType,
   UnionType,
   ValueOf,
@@ -20,6 +22,7 @@ import {
 } from "../../../../src";
 import {uint64NumInfType, uint64NumType} from "../../../utils/primitiveTypes";
 import {runViewTestMutation} from "../runViewTestMutation";
+import {upgradeToNewType} from "../../../../src/util/upgrade";
 
 // Test both ContainerType, ContainerNodeStructType only if
 // - All fields are immutable
@@ -635,3 +638,53 @@ describe("ContainerNodeStruct batchHashTreeRoot", function () {
     expect(viewDU.batchHashTreeRoot()).to.be.deep.equal(expectedRoot);
   });
 });
+
+describe("upgradeToNewType utility", () => {
+  const numFields = [2, 7, 15, 17, 31, 33, 63, 65, 127, 129];
+  for (const [i, numField] of numFields.entries()) {
+    it(`upgradeToNewType with ${numField} fields`, () => {
+      const fields: Record<string, Type<unknown>> = {};
+      for (let j = 0; j < numField; j++) {
+        fields[`f${j}`] = uint64NumInfType;
+      }
+      const oldType = new ContainerType(fields);
+      const view = oldType.defaultView();
+      const viewDU = oldType.defaultViewDU();
+      for (let j = 0; j < numField; j++) {
+        (view as Record<string, number>)[`f${j}`] = j;
+        (viewDU as Record<string, number>)[`f${j}`] = j;
+      }
+
+      for (let j = i + 1; j < numFields.length; j++) {
+        const newFields: Record<string, Type<unknown>> = {};
+        for (let k = 0; k < numFields[j]; k++) {
+          (newFields as Record<string, Type<unknown>>)[`f${k}`] = uint64NumInfType;
+        }
+
+        const newType = new ContainerType(newFields);
+        const newView = newType.getView(new Tree(upgradeToNewType(view.node, oldType, newType)));
+        // commit view DU to make sure the view is updated before accessing viewDU.node
+        viewDU.commit();
+        const newViewDU = newType.getViewDU(upgradeToNewType(viewDU.node, oldType, newType));
+        for (let k = i + 1; k < numFields[j]; k++) {
+          (newView as Record<string, number>)[`f${k}`] = k;
+          (newViewDU as Record<string, number>)[`f${k}`] = k;
+        }
+        newViewDU.commit();
+
+        const expectedValue = newType.defaultValue();
+        for (let k = 0; k < numFields[j]; k++) {
+          (expectedValue as Record<string, number>)[`f${k}`] = k;
+        }
+        const expectedViewDU = newType.toViewDU(expectedValue);
+
+        expect(newView.toValue()).to.be.deep.equal(expectedValue);
+        expect(newView.hashTreeRoot()).to.be.deep.equal(expectedViewDU.hashTreeRoot());
+        expect(newView.serialize()).to.be.deep.equal(expectedViewDU.serialize());
+        expect(newViewDU.toValue()).to.be.deep.equal(expectedValue);
+        expect(newViewDU.hashTreeRoot()).to.be.deep.equal(expectedViewDU.hashTreeRoot());
+        expect(newViewDU.serialize()).to.be.deep.equal(expectedViewDU.serialize());
+      }
+    });
+  }
+});