From 544fbd418d62c78341888ff3568aadf7e936f492 Mon Sep 17 00:00:00 2001
From: Sam Minnee <sam@silverstripe.com>
Date: Fri, 21 Sep 2018 18:34:13 +1200
Subject: [PATCH 1/2] FIX: Use Hierarchy::prepopulateTreeDataCache() in CMS.

Requires https://github.com/silverstripe/silverstripe-framework/pull/8395
---
 src/Versioned.php | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/Versioned.php b/src/Versioned.php
index f5c8a140..601d1e13 100644
--- a/src/Versioned.php
+++ b/src/Versioned.php
@@ -2311,6 +2311,20 @@ public static function get_versionnumber_by_stage($class, $stage, $id, $cache =
         return $version ?: null;
     }
 
+    /**
+     * 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 = [])
+    {
+        $idList = $recordList ? $recordList->column('ID') : null;
+        self::prepopulate_versionnumber_cache($this->owner->baseClass(), Versioned::DRAFT, $idList);
+        self::prepopulate_versionnumber_cache($this->owner->baseClass(), Versioned::LIVE, $idList);
+    }
+
     /**
      * Pre-populate the cache for Versioned::get_versionnumber_by_stage() for
      * a list of record IDs, for more efficient database querying.  If $idList

From fb7219cdd894f201d4917028c517c1fa7a910113 Mon Sep 17 00:00:00 2001
From: Maxime Rainville <maxime@silverstripe.com>
Date: Tue, 25 Sep 2018 15:35:57 +1200
Subject: [PATCH 2/2] Implement PR feedback.

---
 src/Versioned.php                      | 7 ++++---
 tests/php/VersionedNumberCacheTest.php | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/Versioned.php b/src/Versioned.php
index 601d1e13..8c60d4d8 100644
--- a/src/Versioned.php
+++ b/src/Versioned.php
@@ -2314,13 +2314,14 @@ public static function get_versionnumber_by_stage($class, $stage, $id, $cache =
     /**
      * Hook into {@link Hierarchy::prepopulateTreeDataCache}.
      *
-     * @param DataList $recordList The list of records to prepopulate caches for. Null for all records.
+     * @param DataList|array $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 = [])
+    public function onPrepopulateTreeDataCache($recordList = null, array $options = [])
     {
-        $idList = $recordList ? $recordList->column('ID') : null;
+        $idList = is_array($recordList) ? $recordList :
+            ($recordList instanceof DataList ? $recordList->column('ID') : null);
         self::prepopulate_versionnumber_cache($this->owner->baseClass(), Versioned::DRAFT, $idList);
         self::prepopulate_versionnumber_cache($this->owner->baseClass(), Versioned::LIVE, $idList);
     }
diff --git a/tests/php/VersionedNumberCacheTest.php b/tests/php/VersionedNumberCacheTest.php
index bdc3a361..cd955533 100644
--- a/tests/php/VersionedNumberCacheTest.php
+++ b/tests/php/VersionedNumberCacheTest.php
@@ -100,7 +100,7 @@ public function testVersionNumberCache($stage, $ID, $cache, $expected)
      */
     public function testPrepopulatedVersionNumberCache($stage, $ID, $cache, $expected)
     {
-        Versioned::prepopulate_versionnumber_cache(TestObject::class, $stage);
+        TestObject::singleton()->onPrepopulateTreeDataCache();
         $actual = Versioned::get_versionnumber_by_stage(TestObject::class, $stage, self::${$ID}, $cache);
         $this->assertEquals(self::$expectedVersions[$expected], $actual);
     }