FIX Opt-in performance fix for many consecutive lookups using isLocalisedInStage

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,24 @@ 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 = [];
+
+    /**
+     * Used to enable or disable the prepopulation of the locale content cache
+     * Defaults to true.
+     *
+     * @config
+     * @var boolean
+     */
+    private static $prepopulate_localecontent_cache = true;
+
     protected function augmentDatabaseDontRequire($localisedTable)
     {
         DB::dont_require_table($localisedTable);
@@ -304,21 +327,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 +357,70 @@ public function flushCache()
     {
         $this->localisedStageCache = [];
     }
+
+    /**
+     * Hook into {@link Hierarchy::prepopulateTreeDataCache}.
+     *
+     * @param DataList $recordList The list of records to prepopulate caches for. Null for all records.
+     * @param array $options A map of hints about what should be cached. "numChildrenMethod" and
+     *                       "childrenMethod" are allowed keys.
+     */
+    public function onPrepopulateTreeDataCache(DataList $recordList = null, array $options = [])
+    {
+        if (!Config::inst()->get(static::class, 'prepopulate_localecontent_cache')) {
+            return;
+        }
+
+        // This functionality hasn't been implemented; it's not actually used right now, so assume you aren't going to
+        // need it.
+        if ($recordList) {
+            user_error(
+                "FluentVersionedExtension::onPrepopulateTreeDataCache not currently optimised for recordList use",
+                E_USER_WARNING
+            );
+        }
+
+        self::prepoulateIdsInLocale(FluentState::singleton()->getLocale(), $this->owner->baseClass());
+    }
+
+    /**
+     * Prepopulate the cache of IDs in a locale, to optimise batch calls to isLocalisedInStage.
+     *
+     * @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;
+        }
+    }
 }