From 9752eb99533adad412fe86899bcf01ab876db2eb Mon Sep 17 00:00:00 2001
From: Tuyen Nguyen <vutuyen2636@gmail.com>
Date: Mon, 18 Sep 2023 16:07:57 +0700
Subject: [PATCH] feat: implement sliceTo() for ListBasicTreeViewDU

---
 packages/ssz/src/viewDU/listBasic.ts          | 39 ++++++++++++++++++-
 .../test/unit/byType/listBasic/tree.test.ts   | 23 +++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/packages/ssz/src/viewDU/listBasic.ts b/packages/ssz/src/viewDU/listBasic.ts
index 85ed4e32..7d2c509b 100644
--- a/packages/ssz/src/viewDU/listBasic.ts
+++ b/packages/ssz/src/viewDU/listBasic.ts
@@ -1,4 +1,11 @@
-import {LeafNode, Node, zeroNode} from "@chainsafe/persistent-merkle-tree";
+import {
+  LeafNode,
+  Node,
+  getNodeAtDepth,
+  setNodesAtDepth,
+  treeZeroAfterIndex,
+  zeroNode,
+} from "@chainsafe/persistent-merkle-tree";
 import {ValueOf} from "../type/abstract";
 import {BasicType} from "../type/basic";
 import {ListBasicType} from "../view/listBasic";
@@ -30,4 +37,34 @@ export class ListBasicTreeViewDU<ElementType extends BasicType<unknown>> extends
 
     this.set(index, value);
   }
+
+  sliceTo(index: number): this {
+    // Commit before getting rootNode to ensure all pending data is in the rootNode
+    this.commit();
+
+    // All nodes beyond length are already zero
+    // Array of length 2: [X,X,0,0], for index >= 1 no action needed
+    if (index >= this._length - 1) {
+      return this;
+    }
+
+    const rootNode = this._rootNode;
+    const chunkIndex = Math.floor(index / this.type.itemsPerChunk);
+    const nodePrev = (this.nodes[chunkIndex] ?? getNodeAtDepth(rootNode, this.type.depth, chunkIndex)) as LeafNode;
+
+    const nodeChanged = nodePrev.clone();
+    // set remaining items in the same chunk to 0
+    for (let i = index + 1; i < (chunkIndex + 1) * this.type.itemsPerChunk; i++) {
+      this.type.elementType.tree_setToPackedNode(nodeChanged, i, 0);
+    }
+    const chunksNode = this.type.tree_getChunksNode(this._rootNode);
+    let newChunksNode = setNodesAtDepth(chunksNode, this.type.chunkDepth, [chunkIndex], [nodeChanged]);
+    // also do the treeZeroAfterIndex operation on the chunks tree
+    newChunksNode = treeZeroAfterIndex(newChunksNode, this.type.chunkDepth, chunkIndex);
+
+    // Must set new length and commit to tree to restore the same tree at that index
+    const newLength = index + 1;
+    const newRootNode = this.type.tree_setChunksNode(rootNode, newChunksNode, newLength);
+    return this.type.getViewDU(newRootNode) as this;
+  }
 }
diff --git a/packages/ssz/test/unit/byType/listBasic/tree.test.ts b/packages/ssz/test/unit/byType/listBasic/tree.test.ts
index 54ed292c..e6b15672 100644
--- a/packages/ssz/test/unit/byType/listBasic/tree.test.ts
+++ b/packages/ssz/test/unit/byType/listBasic/tree.test.ts
@@ -210,3 +210,26 @@ describe("ListBasicType drop caches", () => {
     expect(bytesAfter).to.equal(bytesBefore, "view retained changes");
   });
 });
+
+describe("ListBasicType.sliceTo", () => {
+  // same to BeaconState InactivityScores
+  it("Slice List at multiple length", () => {
+    const listType = new ListBasicType(new UintNumberType(8), 100);
+    const listView = listType.defaultViewDU();
+    const listRoots: string[] = [];
+    const listSerialized: string[] = [];
+
+    for (let i = 0; i < 16; i++) {
+      listView.push(i);
+      listSerialized[i] = toHexString(listView.serialize());
+      listRoots[i] = toHexString(listView.hashTreeRoot());
+    }
+
+    for (let i = 0; i < 16; i++) {
+      const listSlice = listView.sliceTo(i);
+      expect(listSlice.length).to.equal(i + 1, `Wrong length at .sliceTo(${i})`);
+      expect(toHexString(listSlice.serialize())).equals(listSerialized[i], `Wrong serialize at .sliceTo(${i})`);
+      expect(toHexString(listSlice.hashTreeRoot())).equals(listRoots[i], `Wrong root at .sliceTo(${i})`);
+    }
+  });
+});