From 613fa24af57d418a41963e022260a6a467b4ba84 Mon Sep 17 00:00:00 2001
From: Nicolas Reynis <ywg.circle@gmail.com>
Date: Wed, 11 Sep 2019 22:57:15 +0200
Subject: [PATCH] Implement flatMap (#25)

---
 README.md                  |  1 +
 src/Chain.php              |  2 +
 src/Link/FlatMap.php       | 32 ++++++++++++++++
 tests/ChainTest.php        |  1 +
 tests/Link/FlatMapTest.php | 75 ++++++++++++++++++++++++++++++++++++++
 5 files changed, 111 insertions(+)
 create mode 100644 src/Link/FlatMap.php
 create mode 100644 tests/Link/FlatMapTest.php

diff --git a/README.md b/README.md
index e6d6287..969bff6 100644
--- a/README.md
+++ b/README.md
@@ -127,6 +127,7 @@ All of these methods manipulate the array, but not all of them return an instanc
 - `->diff(array|Chain)`
 - `->filter(callable)`
 - `->find(callable)`
+- `->flatMap(callable)`
 - `->flip()`
 - `->intersect(array|Chain)`
 - `->intersectAssoc(array|Chain)`
diff --git a/src/Chain.php b/src/Chain.php
index 60b4217..2fea428 100644
--- a/src/Chain.php
+++ b/src/Chain.php
@@ -10,6 +10,7 @@
 use Cocur\Chain\Link\Filter;
 use Cocur\Chain\Link\First;
 use Cocur\Chain\Link\Find;
+use Cocur\Chain\Link\FlatMap;
 use Cocur\Chain\Link\Flip;
 use Cocur\Chain\Link\Intersect;
 use Cocur\Chain\Link\IntersectAssoc;
@@ -54,6 +55,7 @@ class Chain extends AbstractChain implements Countable
         Fill,
         Find,
         First,
+        FlatMap,
         Flip,
         Intersect,
         IntersectAssoc,
diff --git a/src/Link/FlatMap.php b/src/Link/FlatMap.php
new file mode 100644
index 0000000..52fae35
--- /dev/null
+++ b/src/Link/FlatMap.php
@@ -0,0 +1,32 @@
+<?php
+
+namespace Cocur\Chain\Link;
+
+use Cocur\Chain\Chain;
+
+/**
+ * FlatMap.
+ *
+ * @author    Nicolas Reynis
+ */
+trait FlatMap
+{
+    /**
+     * @return Chain
+     */
+    public function flatMap(callable $callback)
+    {
+        $flattened = [];
+        foreach ($this->array as $index => $element) {
+            $transformation = $callback($element, $index);
+            if (is_array($transformation)) {
+                array_push($flattened, ...$transformation);
+            } else {
+                $flattened[] = $transformation;
+            }
+        }
+        $this->array = $flattened;
+
+        return $this;
+    }
+}
diff --git a/tests/ChainTest.php b/tests/ChainTest.php
index df74089..de35b32 100644
--- a/tests/ChainTest.php
+++ b/tests/ChainTest.php
@@ -62,6 +62,7 @@ public function chainHasTraits()
         $this->assertTrue(method_exists($c, 'filter'));
         $this->assertTrue(method_exists($c, 'find'));
         $this->assertTrue(method_exists($c, 'first'));
+        $this->assertTrue(method_exists($c, 'flatMap'));
         $this->assertTrue(method_exists($c, 'flip'));
         $this->assertTrue(method_exists($c, 'intersect'));
         $this->assertTrue(method_exists($c, 'intersectAssoc'));
diff --git a/tests/Link/FlatMapTest.php b/tests/Link/FlatMapTest.php
new file mode 100644
index 0000000..43f6677
--- /dev/null
+++ b/tests/Link/FlatMapTest.php
@@ -0,0 +1,75 @@
+<?php
+
+namespace Cocur\Chain\Link;
+
+/**
+ * MapTest.
+ *
+ * @author    Nicolas Reynis
+ * @group     unit
+ */
+class FlatMapTest extends \PHPUnit\Framework\TestCase
+{
+    /**
+     * @test
+     * @covers Cocur\Chain\Link\FlatMap::flatMap()
+     */
+    public function flatMapAppliesMapToArray()
+    {
+        /** @var \Cocur\Chain\Link\FlatMap $mock */
+        $mock        = $this->getMockForTrait('Cocur\Chain\Link\FlatMap');
+        $mock->array = ['foobar', 'bar'];
+        $mock->flatMap(function ($v) { return [str_replace('bar', 'foo', $v)]; });
+
+        $this->assertEquals('foofoo', $mock->array[0]);
+        $this->assertEquals('foo', $mock->array[1]);
+    }
+
+    /**
+     * @test
+     * @covers Cocur\Chain\Link\FlatMap::flatMap()
+     */
+    public function flatMapCallbackReceivesAlsoArrayKeys()
+    {
+        /** @var \Cocur\Chain\Link\FlatMap $mock */
+        $mock        = $this->getMockForTrait('Cocur\Chain\Link\FlatMap');
+        $mock->array = ['foo' => 'fizz', 'bar' => 'buzz'];
+        $mock->flatMap(function ($v, $k) {
+            return $k == 'foo' ? [str_replace('fizz', 'bang', $v)] : [$v];
+        });
+
+        $this->assertEquals(['bang', 'buzz'], $mock->array);
+    }
+
+    /**
+     * @test
+     * @covers Cocur\Chain\Link\FlatMap::flatMap()
+     */
+    public function flatMapCallbackCanReturnScalarOrArray()
+    {
+        /** @var \Cocur\Chain\Link\FlatMap $mock */
+        $mock        = $this->getMockForTrait('Cocur\Chain\Link\FlatMap');
+        $mock->array = ['fizz', 'buzz'];
+        $mock->flatMap(function ($v, $k) {
+            return $k == 1 ? [$v] : $v;
+        });
+
+        $this->assertEquals(['fizz', 'buzz'], $mock->array);
+    }
+
+    /**
+     * @test
+     * @covers Cocur\Chain\Link\FlatMap::flatMap()
+     */
+    public function flatMapAcceptsEmptyArray()
+    {
+        /** @var \Cocur\Chain\Link\FlatMap $mock */
+        $mock        = $this->getMockForTrait('Cocur\Chain\Link\FlatMap');
+        $mock->array = [];
+        $mock->flatMap(function ($v) {
+            return $v;
+        });
+
+        $this->assertEquals([], $mock->array);
+    }
+}