From a2edb5e8cab0a8638cb013eeb227c79ebee79429 Mon Sep 17 00:00:00 2001
From: sukhwinder-somar <sukhwinder@somar.co.nz>
Date: Thu, 12 Oct 2023 11:41:47 +1300
Subject: [PATCH] update PR with new supplied by reviewer

---
 .../ElementalContentControllerExtension.php   | 54 +++++++++----------
 1 file changed, 27 insertions(+), 27 deletions(-)

diff --git a/src/Extensions/ElementalContentControllerExtension.php b/src/Extensions/ElementalContentControllerExtension.php
index 90adae2a..cd642195 100644
--- a/src/Extensions/ElementalContentControllerExtension.php
+++ b/src/Extensions/ElementalContentControllerExtension.php
@@ -2,9 +2,10 @@
 
 namespace DNADesign\Elemental\Extensions;
 
+use SilverStripe\Core\Extension;
+use SilverStripe\ORM\DataObject;
 use DNADesign\Elemental\Models\ElementalArea;
 use DNADesign\Elemental\Extensions\ElementalAreasExtension;
-use SilverStripe\Core\Extension;
 
 class ElementalContentControllerExtension extends Extension
 {
@@ -24,52 +25,51 @@ public function handleElement()
             return false;
         }
 
-        /** @var SiteTree $elementOwner */
+        /** @var DataObject $elementOwner */
         $elementOwner = $this->owner->data();
 
-        $elementalAreaRelations = $this->owner->getElementalRelations();
+        if (!$elementOwner->hasExtension(ElementalAreasExtension::class)) {
+            user_error(get_class($elementOwner) . ' doesnt have the ElementalAreasExtension applied', E_USER_ERROR);
+            return false;
+        }
 
+        $elementalAreaRelations = $elementOwner->getElementalRelations();
         if (!$elementalAreaRelations) {
-            user_error(get_class($this->owner) . ' has no ElementalArea relationships', E_USER_ERROR);
+            user_error(get_class($elementOwner) . ' has no ElementalArea relationships', E_USER_ERROR);
             return false;
         }
 
-        foreach ($elementalAreaRelations as $elementalAreaRelation) {
-            $element = $this->findElement($elementOwner->{$elementalAreaRelation}()->Elements(), $id);
+        $element = $this->findElement($elementalAreaRelations, $elementOwner, $id);
 
-            if ($element) {
-                return $element->getController();
-            }
+        if ($element) {
+            return $element->getController();
         }
 
         user_error('Element $id not found for this page', E_USER_ERROR);
         return false;
     }
 
-    private function findElement($elements, $id)
+    private function findElement(iterable $elementalAreaRelations, DataObject $owner, $id): ?DataObject
     {
-        $element = $elements->filter('ID', $id)->First();
-
-        if ($element) {
-            return $element;
-        }
+        foreach ($elementalAreaRelations as $elementalAreaRelation) {
+            $elements = $owner->$elementalAreaRelation()->Elements();
+            $found = $elements->filter('ID', $id)->First();
 
-        foreach ($elements as $el) {
-            if (!$el->hasMethod('Elements')) {
-                continue;
+            if ($found) {
+                return $found;
             }
 
-            $subElementAreaRelations = $el->getElementalRelations();
-
-            if (!$subElementAreaRelations) {
-                continue;
-            }
+            /** @var BaseElement $element */
+            foreach ($elements as $element) {
+                if (!$element->hasExtension(ElementalAreasExtension::class)) {
+                    continue;
+                }
 
-            foreach ($subElementAreaRelations as $subElementalAreaRelation) {
-                $element = $this->findElement($el->{$subElementalAreaRelation}()->Elements(), $id);
+                /** @var BaseElement&ElementalAreasExtension $element */
+                $found = $this->findElement($element->getElementalRelations(), $element, $id);
 
-                if ($element) {
-                    return $element;
+                if ($found) {
+                    return $found;
                 }
             }
         }