[6.x] Only look among the entries found so far when a query narrows them further

src/Stache/Query/EntryQueryBuilder.php

@@ -16,6 +16,7 @@ class EntryQueryBuilder extends Builder implements QueryBuilder
 
     protected $collections = [];
     protected $nested = false;
+    private $lookingAmong;
 
     public function where($column, $operator = null, $value = null, $boolean = 'and')
     {
@@ -110,17 +111,48 @@ protected function getKeysFromCollections($collections)
 
     protected function getKeysFromCollectionsWithWheres($collections, $wheres)
     {
-        return collect($wheres)->reduce(function ($ids, $where) use ($collections) {
-            $keys = $where['type'] == 'Nested'
-                ? $this->getKeysFromCollectionsWithWheres($collections, $where['query']->wheres)
-                : $this->getKeysFromCollectionsWithWhere($collections, $where);
+        $within = $this->lookingAmong;
 
-            return $this->intersectKeysFromWhereClause($ids, $keys, $where);
-        });
+        try {
+            return collect($wheres)->reduce(function ($ids, $where) use ($collections, $within) {
+                // What an "and" finds is only kept where it's among the keys there are so far,
+                // so once those are few, they're all that needs looking through.
+                $this->lookingAmong = $this->keysToLookAmong($ids, $where, $within);
+
+                $keys = $where['type'] == 'Nested'
+                    ? $this->getKeysFromCollectionsWithWheres($collections, $where['query']->wheres)
+                    : $this->getKeysFromCollectionsWithWhere($collections, $where);
+
+                return $this->intersectKeysFromWhereClause($ids, $keys, $where);
+            });
+        } finally {
+            $this->lookingAmong = $within;
+        }
+    }
+
+    private function keysToLookAmong($ids, $where, $within)
+    {
+        if (! $ids || ($where['boolean'] === 'or' && $where['type'] !== 'NotIn')) {
+            return $within;
+        }
+
+        if ($ids->count() > 250) {
+            return $within;
+        }
+
+        return $ids
+            ->mapToGroups(function ($id) {
+                [$collection, $key] = explode('::', $id, 2);
+
+                return [$collection => $key];
+            })
+            ->map(fn ($keys) => $keys->flip());
     }
 
     protected function getKeysFromCollectionsWithWhere($collections, $where)
     {
+        $within = $this->lookingAmong;
+
         // A collection can be queried twice (e.g. by a status query), but its keys must only appear once.
         $collections = collect($collections)->unique()->values();
 
@@ -133,9 +165,17 @@ protected function getKeysFromCollectionsWithWhere($collections, $where)
         // These filters only look at values, so each store's index can be filtered
         // before its keys are prefixed, instead of prefixing every key first.
         if (in_array($where['type'], ['Basic', 'In'])) {
-            return $collections->flatMap(function ($collection) use ($where, $method) {
+            return $collections->flatMap(function ($collection) use ($where, $method, $within) {
+                if ($within && ! $within->has($collection)) {
+                    return [];
+                }
+
                 $items = $this->store->store($collection)->index($where['column'])->items();
 
+                if ($within) {
+                    $items = $items->intersectByKeys($within[$collection]);
+                }
+
                 return $this->{$method}($items, $where)->keys()->map(fn ($key) => "{$collection}::{$key}");
             })->values();
         }

tests/Data/Entries/EntryQueryBuilderTest.php

@@ -1700,6 +1700,65 @@ public function sorting_by_unsafe_method_does_not_invoke_it()
 
         $this->assertCount($count, Entry::all());
     }
+
+    #[Test]
+    public function a_where_after_a_few_entries_have_been_found_only_looks_at_those()
+    {
+        Collection::make('posts')->save();
+        collect(range(1, 20))->each(fn ($i) => EntryFactory::id('id-'.$i)->slug('post-'.$i)->collection('posts')->data(['title' => 'Post '.$i, 'colour' => $i % 2 ? 'red' : 'blue'])->create());
+
+        $query = new CountingEntryQueryBuilder(app(\Statamic\Stache\Stache::class)->store('entries'));
+
+        $entries = $query->where('collection', 'posts')->whereIn('id', ['id-3', 'id-4', 'id-5'])->where('colour', 'red')->get();
+
+        $this->assertEquals(['id-3', 'id-5'], $entries->map->id()->all());
+        $this->assertSame(3, $query->tested);
+    }
+
+    #[Test]
+    #[DataProvider('narrowedQueryProvider')]
+    public function wheres_after_a_few_entries_have_been_found_find_what_they_would_among_all_of_them($expected, $callback)
+    {
+        Collection::make('posts')->save();
+        Collection::make('things')->save();
+        collect(range(1, 12))->each(fn ($i) => EntryFactory::id('id-'.$i)->slug('post-'.$i)->collection($i > 9 ? 'things' : 'posts')->data(['title' => 'Post '.$i, 'colour' => $i % 2 ? 'red' : 'blue', 'size' => $i % 3 ? 'big' : 'small'])->create());
+
+        $query = Entry::query()->whereIn('id', ['id-6', 'id-3', 'id-4', 'id-5', 'id-11', 'id-12']);
+
+        $this->assertEquals($expected, $callback($query)->get()->map->id()->all());
+    }
+
+    public static function narrowedQueryProvider()
+    {
+        return [
+            'and' => [['id-3', 'id-5', 'id-11'], fn ($q) => $q->where('colour', 'red')],
+            'and, and' => [['id-5', 'id-11'], fn ($q) => $q->where('colour', 'red')->where('size', 'big')],
+            'and in' => [['id-3', 'id-6', 'id-12'], fn ($q) => $q->whereIn('size', ['small'])],
+            'and not in' => [['id-4', 'id-5', 'id-11'], fn ($q) => $q->whereNotIn('size', ['small'])],
+            'and, or' => [['id-3', 'id-5', 'id-11', 'id-1'], fn ($q) => $q->where('colour', 'red')->orWhere('title', 'Post 1')],
+            'and, or, and' => [['id-3', 'id-5', 'id-11', 'id-2'], fn ($q) => $q->where('colour', 'red')->orWhere('colour', 'blue')->whereIn('title', ['Post 3', 'Post 5', 'Post 11', 'Post 2'])],
+            'and, or not in' => [['id-5', 'id-11'], fn ($q) => $q->where('colour', 'red')->orWhereNotIn('size', ['small'])],
+            'and nested' => [['id-3', 'id-4', 'id-5', 'id-11'], fn ($q) => $q->where(fn ($q) => $q->where('colour', 'red')->orWhere('title', 'Post 4')->orWhere('title', 'Post 1'))],
+            'and nested with an and inside' => [['id-5', 'id-11'], fn ($q) => $q->where(fn ($q) => $q->where('colour', 'red')->where('size', 'big'))],
+            'or nested' => [['id-3', 'id-4', 'id-5', 'id-6', 'id-11', 'id-12', 'id-1', 'id-7'], fn ($q) => $q->orWhere(fn ($q) => $q->where('colour', 'red')->where('size', 'big'))],
+            'and collection' => [['id-11', 'id-12'], fn ($q) => $q->where('collection', 'things')],
+            'and status' => [['id-3', 'id-4', 'id-5', 'id-6', 'id-11', 'id-12'], fn ($q) => $q->whereStatus('published')],
+            'and like' => [['id-11', 'id-12'], fn ($q) => $q->where('title', 'like', 'Post 1%')],
+            'ordered' => [['id-3', 'id-5', 'id-11'], fn ($q) => $q->where('colour', 'red')->orderBy('title', 'asc')],
+        ];
+    }
+}
+
+class CountingEntryQueryBuilder extends \Statamic\Stache\Query\EntryQueryBuilder
+{
+    public $tested = 0;
+
+    protected function filterTestEquals($item, $value)
+    {
+        $this->tested++;
+
+        return parent::filterTestEquals($item, $value);
+    }
 }
 
 class CustomScope extends Scope