Add array chainWithIndex (#1256)

gcanti · Jun 29, 2020 · 21f23d0 · 21f23d0
1 parent 291c3bf
commit 21f23d0
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 2 deletions.
diff --git a/docs/modules/Array.ts.md b/docs/modules/Array.ts.md
@@ -135,6 +135,7 @@ Added in v2.0.0
   - [unsafeInsertAt](#unsafeinsertat)
   - [unsafeUpdateAt](#unsafeupdateat)
 - [utils](#utils)
+  - [chainWithIndex](#chainwithindex)
   - [chunksOf](#chunksof)
   - [comprehension](#comprehension)
   - [deleteAt](#deleteat)
@@ -1763,6 +1764,16 @@ Added in v2.0.0
 
 # utils
 
+## chainWithIndex
+
+**Signature**
+
+```ts
+export declare const chainWithIndex: <A, B>(f: (index: number, a: A) => B[]) => (ma: A[]) => B[]
+```
+
+Added in v2.7.0
+
 ## chunksOf
 
 Splits an array into length-`n` pieces. The last piece will be shorter if `n` does not evenly divide the length of

diff --git a/docs/modules/ReadonlyArray.ts.md b/docs/modules/ReadonlyArray.ts.md
@@ -127,6 +127,7 @@ Added in v2.5.0
   - [unsafeUpdateAt](#unsafeupdateat)
 - [utils](#utils)
   - [Spanned (interface)](#spanned-interface)
+  - [chainWithIndex](#chainwithindex)
   - [chunksOf](#chunksof)
   - [comprehension](#comprehension)
   - [deleteAt](#deleteat)
@@ -1556,6 +1557,16 @@ export interface Spanned<I, R> {
 
 Added in v2.5.0
 
+## chainWithIndex
+
+**Signature**
+
+```ts
+export declare const chainWithIndex: <A, B>(f: (i: number, a: A) => readonly B[]) => (ma: readonly A[]) => readonly B[]
+```
+
+Added in v2.7.0
+
 ## chunksOf
 
 Splits an array into length-`n` pieces. The last piece will be shorter if `n` does not evenly divide the length of

diff --git a/src/Array.ts b/src/Array.ts
@@ -1054,6 +1054,13 @@ export const apSecond: <B>(fb: Array<B>) => <A>(fa: Array<A>) => Array<B> = RA.a
  */
 export const chain: <A, B>(f: (a: A) => Array<B>) => (ma: Array<A>) => Array<B> = RA.chain as any
 
+/**
+ * @since 2.7.0
+ */
+export const chainWithIndex: <A, B>(
+  f: (index: number, a: A) => Array<B>
+) => (ma: Array<A>) => Array<B> = RA.chainWithIndex as any
+
 /**
  * Composes computations in sequence, using the return value of one computation to determine the next computation and
  * keeping only the result of the first.

diff --git a/src/ReadonlyArray.ts b/src/ReadonlyArray.ts
@@ -1431,13 +1431,16 @@ export const zero: Alternative1<URI>['zero'] = () => empty
 const map_: Monad1<URI>['map'] = (fa, f) => fa.map((a) => f(a))
 const mapWithIndex_: FunctorWithIndex1<URI, number>['mapWithIndex'] = (fa, f) => fa.map((a, i) => f(i, a))
 const ap_: Monad1<URI>['ap'] = (fab, fa) => flatten(map_(fab, (f) => map_(fa, f)))
-const chain_: Monad1<URI>['chain'] = (fa, f) => {
+const chainWithIndex_: <A, B>(fa: ReadonlyArray<A>, f: (i: number, a: A) => ReadonlyArray<B>) => ReadonlyArray<B> = (
+  fa,
+  f
+) => {
   let outLen = 0
   const l = fa.length
   const temp = new Array(l)
   for (let i = 0; i < l; i++) {
     const e = fa[i]
-    const arr = f(e)
+    const arr = f(i, e)
     outLen += arr.length
     temp[i] = arr
   }
@@ -1453,6 +1456,8 @@ const chain_: Monad1<URI>['chain'] = (fa, f) => {
   }
   return out
 }
+const chain_: <A, B>(fa: ReadonlyArray<A>, f: (a: A) => ReadonlyArray<B>) => ReadonlyArray<B> = (fa, f) =>
+  chainWithIndex_(fa, (_index, a) => f(a))
 const filterMap_: Filterable1<URI>['filterMap'] = (as, f) => filterMapWithIndex_(as, (_, a) => f(a))
 const filter_: Filter1<URI> = <A>(as: ReadonlyArray<A>, predicate: Predicate<A>) => as.filter(predicate)
 const partitionWithIndex_: PartitionWithIndex1<URI, number> = <A>(
@@ -1631,6 +1636,13 @@ export const apSecond = <B>(fb: ReadonlyArray<B>) => <A>(fa: ReadonlyArray<A>):
 export const chain: <A, B>(f: (a: A) => ReadonlyArray<B>) => (ma: ReadonlyArray<A>) => ReadonlyArray<B> = (f) => (ma) =>
   chain_(ma, f)
 
+/**
+ * @since 2.7.0
+ */
+export const chainWithIndex: <A, B>(
+  f: (i: number, a: A) => ReadonlyArray<B>
+) => (ma: ReadonlyArray<A>) => ReadonlyArray<B> = (f) => (ma) => chainWithIndex_(ma, f)
+
 /**
  * Composes computations in sequence, using the return value of one computation to determine the next computation and
  * keeping only the result of the first.

diff --git a/test/Array.ts b/test/Array.ts
@@ -131,6 +131,16 @@ describe('Array', () => {
       )
     })
 
+    it('chainWithIndex', () => {
+      assert.deepStrictEqual(
+        pipe(
+          [1, 2, 3],
+          _.chainWithIndex((i, n) => [n + i])
+        ),
+        [1, 3, 5]
+      )
+    })
+
     it('chainFirst', () => {
       assert.deepStrictEqual(
         pipe(

diff --git a/test/ReadonlyArray.ts b/test/ReadonlyArray.ts
@@ -127,6 +127,16 @@ describe('ReadonlyArray', () => {
       )
     })
 
+    it('chainWithIndex', () => {
+      assert.deepStrictEqual(
+        pipe(
+          [1, 2, 3],
+          _.chainWithIndex((i, n) => [n + i])
+        ),
+        [1, 3, 5]
+      )
+    })
+
     it('chainFirst', () => {
       assert.deepStrictEqual(
         pipe(