Store totalCount in paginator

Relay/Connection/Output/Connection.php

@@ -19,6 +19,9 @@ final class Connection
     /** @var PageInfo */
     public $pageInfo;
 
+    /** @var int */
+    public $totalCount;
+
     public function __construct(array $edges, PageInfo $pageInfo)
     {
         $this->edges = $edges;

Relay/Connection/Paginator.php

@@ -30,6 +30,11 @@ class Paginator
      */
     private $promise;
 
+    /**
+     * @var int
+     */
+    private $totalCount;
+
     /**
      * @param callable $fetcher
      * @param bool     $promise
@@ -49,7 +54,7 @@ public function __construct(callable $fetcher, $promise = self::MODE_REGULAR)
      */
     public function backward($args, $total, array $callableArgs = [])
     {
-        $total = is_callable($total) ? call_user_func_array($total, $callableArgs) : $total;
+        $total = $this->computeTotalCount($total, $callableArgs);
 
         $args = $this->protectArgs($args);
         $limit = $args['last'];
@@ -107,10 +112,20 @@ public function auto($args, $total, $callableArgs = [])
         $args = $this->protectArgs($args);
 
         if ($args['last']) {
-            return $this->backward($args, $total, $callableArgs);
+            $connection = $this->backward($args, $total, $callableArgs);
         } else {
-            return $this->forward($args);
+            $connection = $this->forward($args);
         }
+
+        if ($this->promise) {
+            $connection->then(function (Connection $connection) use ($total, $callableArgs) {
+                $connection->totalCount = $this->computeTotalCount($total, $callableArgs);
+            });
+        } else {
+            $connection->totalCount = $this->computeTotalCount($total, $callableArgs);
+        }
+
+        return $connection;
     }
 
     /**
@@ -137,4 +152,21 @@ private function protectArgs($args)
     {
         return $args instanceof Argument ? $args : new Argument($args);
     }
+
+    /**
+     * @param int   $total
+     * @param array $callableArgs
+     *
+     * @return int|mixed
+     */
+    private function computeTotalCount($total, array $callableArgs = [])
+    {
+        if ($this->totalCount !== null) {
+            return $this->totalCount;
+        }
+
+        $this->totalCount = is_callable($total) ? call_user_func_array($total, $callableArgs) : $total;
+
+        return $this->totalCount;
+    }
 }

Tests/Relay/Connection/PaginatorTest.php

@@ -279,6 +279,8 @@ public function testTotalCallableWithArguments()
             ['array' => $this->data]
         );
 
+        $this->assertSame(count($this->data), $result->totalCount);
+
         $this->assertCount(4, $result->edges);
         $this->assertSameEdgeNodeValue(['B', 'C', 'D', 'E'], $result);
         $this->assertTrue($result->pageInfo->hasPreviousPage);
@@ -292,7 +294,10 @@ public function testPromiseMode()
             ->getMock()
         ;
 
-        $promise->expects($this->once())->method('then');
+        $promise
+            ->expects($this->exactly(2))
+            ->method('then')
+            ->willReturnSelf();
 
         $paginator = new Paginator(function ($offset, $limit) use ($promise) {
             $this->assertSame(0, $offset);