add SortedMap.lastOption & partition apis (#2652)

Co-authored-by: Tim <hello@timsmart.co>
Effect-TS · Apr 30, 2024 · c3c12c6 · c3c12c6
1 parent ba64ea6
commit c3c12c6
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 0 deletions.
diff --git a/.changeset/calm-penguins-kneel.md b/.changeset/calm-penguins-kneel.md
@@ -0,0 +1,5 @@
+---
+"effect": minor
+---
+
+add SortedMap.lastOption & partition apis
diff --git a/packages/effect/src/SortedMap.ts b/packages/effect/src/SortedMap.ts
@@ -250,3 +250,38 @@ export const entries = <K, V>(self: SortedMap<K, V>): IterableIterator<[K, V]> =
   iterator[Symbol.iterator] = () => entries(self)
   return iterator
 }
+
+/**
+ * @since 3.1.0
+ * @category elements
+ */
+export const lastOption = <K, V>(self: SortedMap<K, V>): Option.Option<[K, V]> => RBT.last(self.tree)
+
+/**
+ * @since 3.1.0
+ * @category filtering
+ */
+export const partition: {
+  <K, V>(
+    predicate: (a: Types.NoInfer<K>) => boolean
+  ): (self: SortedMap<K, V>) => [excluded: SortedMap<K, V>, satisfying: SortedMap<K, V>]
+  <K, V>(self: SortedMap<K, V>, predicate: (a: K) => boolean): [excluded: SortedMap<K, V>, satisfying: SortedMap<K, V>]
+} = Dual.dual(
+  2,
+  <K, V>(
+    self: SortedMap<K, V>,
+    predicate: (a: K) => boolean
+  ): [excluded: SortedMap<K, V>, satisfying: SortedMap<K, V>] => {
+    const ord = RBT.getOrder(self.tree)
+    let right = empty<K, V>(ord)
+    let left = empty<K, V>(ord)
+    for (const value of self) {
+      if (predicate(value[0])) {
+        right = set(right, value[0], value[1])
+      } else {
+        left = set(left, value[0], value[1])
+      }
+    }
+    return [left, right]
+  }
+)
diff --git a/packages/effect/test/SortedMap.test.ts b/packages/effect/test/SortedMap.test.ts
@@ -128,6 +128,14 @@ describe("SortedMap", () => {
     assert.deepEqual(SM.headOption(map2), O.none())
   })
 
+  it("lastOption", () => {
+    const map1 = makeSortedMap([0, 10], [1, 20], [2, 30])
+    const map2 = SM.empty<number, number>(N.Order)
+
+    assert.deepEqual(SM.lastOption(map1), O.some([key(2), value(30)] as const))
+    assert.deepEqual(SM.lastOption(map2), O.none())
+  })
+
   it("isEmpty", () => {
     const map1 = makeSortedMap([0, 10], [1, 20], [2, 30])
     const map2 = SM.empty<number, number>(N.Order)
@@ -172,6 +180,73 @@ describe("SortedMap", () => {
     )
   })
 
+  it("partition", () => {
+    const map1 = makeSortedMap([1, 10], [2, 20], [3, 30], [4, 40], [5, 50])
+
+    const [excl, satisfying] = pipe(
+      map1,
+      SM.partition((member) => member.id <= 3)
+    )
+
+    assert.deepEqual(
+      Array.from(satisfying),
+      [
+        [key(1), value(10)],
+        [key(2), value(20)],
+        [key(3), value(30)]
+      ]
+    )
+    assert.deepEqual(
+      Array.from(excl),
+      [
+        [key(4), value(40)],
+        [key(5), value(50)]
+      ]
+    )
+
+    const [excl2, satisfying2] = pipe(
+      map1,
+      SM.partition((member) => member.id <= 6)
+    )
+
+    assert.deepEqual(
+      Array.from(satisfying2),
+      [
+        [key(1), value(10)],
+        [key(2), value(20)],
+        [key(3), value(30)],
+        [key(4), value(40)],
+        [key(5), value(50)]
+      ]
+    )
+
+    assert.deepEqual(
+      Array.from(excl2),
+      []
+    )
+
+    const [excl3, satisfying3] = pipe(
+      map1,
+      SM.partition((member) => member.id === 0)
+    )
+
+    assert.deepEqual(
+      Array.from(excl3),
+      [
+        [key(1), value(10)],
+        [key(2), value(20)],
+        [key(3), value(30)],
+        [key(4), value(40)],
+        [key(5), value(50)]
+      ]
+    )
+
+    assert.deepEqual(
+      Array.from(satisfying3),
+      []
+    )
+  })
+
   it("reduce", () => {
     const map1 = makeSortedMap([0, 10], [1, 20], [2, 30])
     const result1 = pipe(map1, SM.reduce("", (acc, value) => acc + value.id))