From ccc5bbacacc7e136cb98e5d6047f10793d4874b0 Mon Sep 17 00:00:00 2001
From: Ivan Babenko <babenko.i.a@gmail.com>
Date: Tue, 27 Sep 2022 19:40:13 +0200
Subject: [PATCH] Map inner hits into the Hit decorator

---
 src/Decorators/Hit.php                     | 10 ++++++++++
 src/Factories/LazyModelFactory.php         | 15 +++++++++++----
 tests/Integration/Queries/RawQueryTest.php | 14 ++++++++++++--
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/src/Decorators/Hit.php b/src/Decorators/Hit.php
index e0d7629..6819e4d 100644
--- a/src/Decorators/Hit.php
+++ b/src/Decorators/Hit.php
@@ -6,6 +6,7 @@
 use Elastic\ScoutDriverPlus\Factories\LazyModelFactory;
 use Illuminate\Contracts\Support\Arrayable;
 use Illuminate\Database\Eloquent\Model;
+use Illuminate\Support\Collection as BaseCollection;
 use Illuminate\Support\Traits\ForwardsCalls;
 
 /**
@@ -32,6 +33,15 @@ public function model(): ?Model
         );
     }
 
+    public function innerHits(): BaseCollection
+    {
+        return $this->baseHit->innerHits()->map(
+            fn (BaseCollection $baseInnerHits) => $baseInnerHits->map(
+                fn (BaseHit $baseInnerHit) => new self($baseInnerHit, $this->lazyModelFactory)
+            )
+        );
+    }
+
     /**
      * @return mixed
      */
diff --git a/src/Factories/LazyModelFactory.php b/src/Factories/LazyModelFactory.php
index 0e8cdc3..2ae9ae4 100644
--- a/src/Factories/LazyModelFactory.php
+++ b/src/Factories/LazyModelFactory.php
@@ -5,18 +5,25 @@
 use Elastic\Adapter\Search\Hit as BaseHit;
 use Elastic\Adapter\Search\SearchResult as BaseSearchResult;
 use Illuminate\Database\Eloquent\Model;
+use Illuminate\Support\Collection as BaseCollection;
 
 class LazyModelFactory
 {
-    private array $documentIds;
+    private array $documentIds = [];
     private ModelFactory $modelFactory;
     private array $models = [];
 
     public function __construct(BaseSearchResult $searchResult, ModelFactory $modelFactory)
     {
-        $this->documentIds = $searchResult->hits()->mapToGroups(
-            static fn (BaseHit $baseHit) => [$baseHit->indexName() => $baseHit->document()->id()]
-        )->toArray();
+        $searchResult->hits()->each(function (BaseHit $baseHit) {
+            $this->documentIds[$baseHit->indexName()][] = $baseHit->document()->id();
+
+            $baseHit->innerHits()->each(function (BaseCollection $baseInnerHits) {
+                $baseInnerHits->each(function (BaseHit $baseInnerHit) {
+                    $this->documentIds[$baseInnerHit->indexName()][] = $baseInnerHit->document()->id();
+                });
+            });
+        });
 
         $this->modelFactory = $modelFactory;
     }
diff --git a/tests/Integration/Queries/RawQueryTest.php b/tests/Integration/Queries/RawQueryTest.php
index fe67f7e..58274de 100644
--- a/tests/Integration/Queries/RawQueryTest.php
+++ b/tests/Integration/Queries/RawQueryTest.php
@@ -227,11 +227,21 @@ public function test_models_can_be_found_using_raw_field_collapsing(): void
 
         // find the cheapest books by author
         $found = Book::searchQuery(['match_all' => new stdClass()])
-            ->collapseRaw(['field' => 'author_id'])
-            ->sort('price', 'asc')
+            ->collapseRaw([
+                'field' => 'author_id',
+                'inner_hits' => [
+                    'name' => 'cheapest',
+                    'size' => 5,
+                    'sort' => [['price' => 'asc']],
+                ],
+            ])
+            ->sort('price')
             ->execute();
 
         $this->assertFoundModels(collect([$firstTarget, $secondTarget]), $found);
+
+        $this->assertCount(5, $found->hits()->first()->innerHits()->get('cheapest')->map->model());
+        $this->assertCount(1, $found->hits()->last()->innerHits()->get('cheapest')->map->model());
     }
 
     public function test_models_can_be_found_using_field_collapsing(): void