From 9cb863cb26fe113f3da31ed29cb8f6e622b0f1fc Mon Sep 17 00:00:00 2001
From: Guy Marriott <guy@scopey.co.nz>
Date: Fri, 21 Sep 2018 14:42:21 +1200
Subject: [PATCH] FIX Opt-in performance fix for many consecutive lookups using
 isLocalisedInStage

---
 src/Extension/FluentVersionedExtension.php | 80 ++++++++++++++++++++--
 1 file changed, 76 insertions(+), 4 deletions(-)

diff --git a/src/Extension/FluentVersionedExtension.php b/src/Extension/FluentVersionedExtension.php
index 6eced6f0..5838f101 100644
--- a/src/Extension/FluentVersionedExtension.php
+++ b/src/Extension/FluentVersionedExtension.php
@@ -3,8 +3,13 @@
 namespace TractorCow\Fluent\Extension;
 
 use InvalidArgumentException;
+use SilverStripe\CMS\Model\SiteTree;
+use SilverStripe\Core\Config\Config;
+use SilverStripe\Core\Injector\Injector;
+use SilverStripe\ORM\DataObject;
 use SilverStripe\ORM\DataQuery;
 use SilverStripe\ORM\DB;
+use SilverStripe\ORM\Queries\SQLExpression;
 use SilverStripe\ORM\Queries\SQLSelect;
 use SilverStripe\Versioned\Versioned;
 use TractorCow\Fluent\Model\Locale;
@@ -60,6 +65,21 @@ class FluentVersionedExtension extends FluentExtension
      */
     protected $localisedStageCache = [];
 
+    /**
+     * Array of objectIds keyed by table (ie. stage) and locale. Indicates which object IDs exist in which locales.
+     *
+     * static::$idsInLocaleCache[ $locale ][ $table(.self::SUFFIX_LIVE) ][ $objectId ] === true
+     *
+     * @var string[][][]
+     */
+    protected static $idsInLocaleCache = [];
+
+    /**
+     * @config
+     * @var array
+     */
+    private static $prepopulate_locale_object_types = [];
+
     protected function augmentDatabaseDontRequire($localisedTable)
     {
         DB::dont_require_table($localisedTable);
@@ -297,6 +317,12 @@ protected function isLocalisedInStage($stage, $locale = null)
             }
         }
 
+        $ownerClass = get_class($this->owner);
+        $prepopulateClasses = Config::inst()->get(self::class, 'prepopulate_locale_object_types');
+        if (is_array($prepopulateClasses) && in_array($ownerClass, $prepopulateClasses)) {
+            self::prepoulateIdsInLocale($locale, $ownerClass);
+        }
+
         // Get table
         $baseTable = $this->owner->baseTable();
         $table = $this->getLocalisedTable($baseTable);
@@ -304,21 +330,26 @@ protected function isLocalisedInStage($stage, $locale = null)
             $table .= self::SUFFIX_LIVE;
         }
 
+        if (isset(self::$idsInLocaleCache[$locale][$table])) {
+            return in_array($this->owner->ID, self::$idsInLocaleCache[$locale][$table]);
+        }
+
         // Check cache
         $key = $table . '/' . $locale . '/' . $this->owner->ID;
         if (isset($this->localisedStageCache[$key])) {
             return $this->localisedStageCache[$key];
         }
 
-        $query = new SQLSelect();
-        $query->selectField('"ID"');
+        $query = new SQLSelect(
+            '"ID"'
+        );
         $query->addFrom('"'. $table . '"');
         $query->addWhere([
             '"RecordID"' => $this->owner->ID,
             '"Locale"' => $locale,
         ]);
-        $query->firstRow();
-        $result = $query->execute()->value() !== null;
+
+        $result = $query->firstRow()->execute()->value() !== null;
 
         // Set cache
         $this->localisedStageCache[$key] = $result;
@@ -329,4 +360,45 @@ public function flushCache()
     {
         $this->localisedStageCache = [];
     }
+
+    /**
+     * Prepopulate the cache of
+     *
+     * @param string $locale
+     * @param string $dataObject
+     */
+    public static function prepoulateIdsInLocale($locale, $dataObjectClass, $populateLive = true, $populateDraft = true)
+    {
+        // Get the table for the given DataObject class
+        $self = new static();
+        $table = $self->getLocalisedTable(Injector::inst()->get($dataObjectClass)->baseTable());
+
+        // If we already have items then we've been here before...
+        if (isset(self::$idsInLocaleCache[$locale][$table])) {
+            return;
+        }
+
+        $tables = [];
+        if ($populateDraft) {
+            $tables[] = $table;
+        }
+        if ($populateLive) {
+            $tables[] = $table . self::SUFFIX_LIVE;
+        }
+
+        // Populate both the draft and live stages
+        foreach ($tables as $table) {
+            /** @var SQLSelect $select */
+            $select = SQLSelect::create(
+                ['"RecordID"'],
+                '"' . $table . '"',
+                ['Locale' => $locale]
+            );
+            $result = $select->execute();
+            $ids = $result->column('RecordID');
+
+            // We need to execute ourselves as the param is lost from the subSelect
+            self::$idsInLocaleCache[$locale][$table] = $ids;
+        }
+    }
 }