Add truncate to List.

Changelog.md

@@ -1,11 +1,12 @@
 ## Next
 
+* Add `truncate` to `List` (#409).
+* **Breaking:** Rename `sort` to `sortInPlace`, add `sort` (#405).
 * Add `isCleanReject` to `Error`, align reject code order with IC interface specification and improve comments (#401).
 * internal: updates `matchers` dev-dependency (#394).
 * Add `PriorityQueue` (#392).
 * Add support for Weak references (#388).
 * Clarify difference between `List` and `pure/List` in doc comments (#386).
-* **Breaking:** Rename `sort` to `sortInPlace`, add `sort` (#405).
 
 ## 1.0.0
 

src/List.mo

@@ -190,6 +190,49 @@ module {
     }
   };
 
+  /// Truncates the list to the specified size.
+  /// If the new size is larger than the current size, it will trap.
+  ///
+  /// Example:
+  /// ```motoko include=import
+  /// let list = List.fromArray<Nat>([1, 2, 3, 4, 5]);
+  /// List.truncate(list, 3); // list is now [1, 2, 3]
+  /// assert List.toArray(list) == [1, 2, 3];
+  /// ```
+  ///
+  /// Runtime: `O(size)`
+  ///
+  /// Space: `O(1)`
+  public func truncate<T>(list : List<T>, newSize : Nat) {
+    if (newSize > size(list)) Prim.trap "List.truncate: newSize is larger than current size";
+
+    let (blockIndex, elementIndex) = locate(newSize);
+    list.blockIndex := blockIndex;
+    list.elementIndex := elementIndex;
+    let newBlocksCount = new_index_block_length(Nat32.fromNat(if (elementIndex == 0) blockIndex - 1 else blockIndex));
+
+    let newBlocks = if (newBlocksCount < list.blocks.size()) {
+      let oldDataBlocks = list.blocks;
+      list.blocks := VarArray.tabulate<[var ?T]>(newBlocksCount, func(i) = oldDataBlocks[i]);
+      list.blocks
+    } else list.blocks;
+
+    var i = if (elementIndex == 0) blockIndex else blockIndex + 1;
+    while (i < newBlocksCount) {
+      newBlocks[i] := [var];
+      i += 1
+    };
+    if (elementIndex != 0) {
+      let block = newBlocks[blockIndex];
+      var i = elementIndex;
+      var to = block.size();
+      while (i < to) {
+        block[i] := null;
+        i += 1
+      }
+    }
+  };
+
   /// Resets the list to size 0, de-referencing all elements.
   ///
   /// Example:

test/List.test.mo

@@ -963,6 +963,38 @@ func testAddAll(n : Nat) : Bool {
   true
 };
 
+func testTruncate(n : Nat) : Bool {
+  for (i in Nat.range(0, n + 1)) {
+    let vec = List.fromArray<Nat>(Array.tabulate<Nat>(n, func j = j));
+    List.truncate(vec, i);
+    if (List.size(vec) != i) {
+      Debug.print("Truncate failed: expected size " # Nat.toText(i) # ", got " # Nat.toText(List.size(vec)));
+      return false
+    };
+    for (j in Nat.range(0, i)) {
+      if (List.at(vec, j) != j) {
+        Debug.print("Truncate failed at index " # Nat.toText(j) # ": expected " # Nat.toText(j) # ", got " # Nat.toText(List.at(vec, j)));
+        return false
+      }
+    };
+    let b = vec.blockIndex;
+    let e = vec.elementIndex;
+    let blocks = vec.blocks;
+    if (b < blocks.size()) {
+      let db = blocks[b];
+      var i = e;
+      while (i < db.size()) {
+        if (db[i] != null) {
+          Debug.print("Truncate failed: expected null at index " # Nat.toText(i) # ", got " # debug_show (db[i]));
+          return false
+        };
+        i += 1
+      }
+    }
+  };
+  true
+};
+
 func testRemoveLast(n : Nat) : Bool {
   let vec = List.fromArray<Nat>(Array.tabulate<Nat>(n, func(i) = i));
   var i = n;
@@ -1222,6 +1254,7 @@ func runAllTests() {
   runTest("testInit", testInit);
   runTest("testAdd", testAdd);
   runTest("testAddAll", testAddAll);
+  runTest("testTruncate", testTruncate);
   runTest("testRemoveLast", testRemoveLast);
   runTest("testAt", testAt);
   runTest("testGet", testGet);

validation/api/api.lock.json

@@ -565,6 +565,7 @@
       "public func toPure<T>(list : List<T>) : PureList.List<T>",
       "public func toText<T>(list : List<T>, f : T -> Text) : Text",
       "public func toVarArray<T>(list : List<T>) : [var T]",
+      "public func truncate<T>(list : List<T>, newSize : Nat)",
       "public func values<T>(list : List<T>) : Iter.Iter<T>"
     ]
   },