diff --git a/Changelog.md b/Changelog.md index ce31f3b0..57149764 100644 --- a/Changelog.md +++ b/Changelog.md @@ -1,6 +1,7 @@ ## Next * Improve CI for external contributions (#413). +* Add `truncate` to `List` (#409). * Add `isSorted` to `List` (#408). * **Breaking:** Rename `sort` to `sortInPlace`, add `sort` (#405). * Add `isCleanReject` to `Error`, align reject code order with IC interface specification and improve comments (#401). diff --git a/src/List.mo b/src/List.mo index 62755632..22c0ef96 100644 --- a/src/List.mo +++ b/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([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(list : List, 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: diff --git a/test/List.test.mo b/test/List.test.mo index ba74c5c7..dc1d81fc 100644 --- a/test/List.test.mo +++ b/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(Array.tabulate(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(Array.tabulate(n, func(i) = i)); var i = n; @@ -1238,6 +1270,7 @@ func runAllTests() { runTest("testInit", testInit); runTest("testAdd", testAdd); runTest("testAddAll", testAddAll); + runTest("testTruncate", testTruncate); runTest("testRemoveLast", testRemoveLast); runTest("testAt", testAt); runTest("testGet", testGet); diff --git a/validation/api/api.lock.json b/validation/api/api.lock.json index b7e25ca8..9c15c150 100644 --- a/validation/api/api.lock.json +++ b/validation/api/api.lock.json @@ -566,6 +566,7 @@ "public func toPure(list : List) : PureList.List", "public func toText(list : List, f : T -> Text) : Text", "public func toVarArray(list : List) : [var T]", + "public func truncate(list : List, newSize : Nat)", "public func values(list : List) : Iter.Iter" ] },