Sort in the stage instead of using stage

src/Doctrine/Odm/Extension/OrderExtension.php

@@ -17,6 +17,7 @@
 use ApiPlatform\Doctrine\Odm\PropertyHelperTrait as MongoDbOdmPropertyHelperTrait;
 use ApiPlatform\Metadata\Operation;
 use Doctrine\ODM\MongoDB\Aggregation\Builder;
+use Doctrine\ODM\MongoDB\Aggregation\Stage\Search;
 use Doctrine\ODM\MongoDB\Aggregation\Stage\Sort;
 use Doctrine\Persistence\ManagerRegistry;
 
@@ -66,7 +67,8 @@ public function applyToCollection(Builder $aggregationBuilder, string $resourceC
                 if ($this->isPropertyNested($field, $resourceClass)) {
                     [$field] = $this->addLookupsForNestedProperty($field, $aggregationBuilder, $resourceClass, true);
                 }
-                $aggregationBuilder->sort(
+                $this->addSort(
+                    $aggregationBuilder,
                     $context['mongodb_odm_sort_fields'] = ($context['mongodb_odm_sort_fields'] ?? []) + [$field => $order]
                 );
             }
@@ -76,7 +78,8 @@ public function applyToCollection(Builder $aggregationBuilder, string $resourceC
 
         if (null !== $this->order) {
             foreach ($identifiers as $identifier) {
-                $aggregationBuilder->sort(
+                $this->addSort(
+                    $aggregationBuilder,
                     $context['mongodb_odm_sort_fields'] = ($context['mongodb_odm_sort_fields'] ?? []) + [$identifier => $this->order]
                 );
             }
@@ -88,26 +91,30 @@ protected function getManagerRegistry(): ManagerRegistry
         return $this->managerRegistry;
     }
 
-    private function hasSortStage(Builder $aggregationBuilder): bool
+    private function addSort(Builder $aggregationBuilder, array $sortFields): void
     {
-        $shouldStop = false;
-        $index = 0;
+        $firstStage = $aggregationBuilder->getPipeline(0);
+        if ($firstStage instanceof Search) {
+            // The $search stage supports "sort" for performance, it's always first if present
+            $firstStage->sort($sortFields);
+        } else {
+            // Append a $sort stage at the end of the pipeline
+            $aggregationBuilder->sort($sortFields);
+        }
+    }
 
-        do {
-            try {
+    private function hasSortStage(Builder $aggregationBuilder): bool
+    {
+        try {
+            for ($index = 0; true; $index++) {
                 if ($aggregationBuilder->getStage($index) instanceof Sort) {
                     // If at least one stage is sort, then it has sorting
                     return true;
                 }
-            } catch (\OutOfRangeException $outOfRangeException) {
-                // There is no more stages on the aggregation builder
-                $shouldStop = true;
             }
-
-            ++$index;
-        } while (!$shouldStop);
-
-        // No stage was sort, and we iterated through all stages
-        return false;
+        } catch (\OutOfRangeException) {
+            // There is no more stages on the aggregation builder
+            return false;
+        }
     }
 }