NEW Rollback elements on page rollback

- Rollback only if it's a Live => Stage rollback, because version based
  rollback is a little more difficult to determine
- Fixed an issue if the page class changes from one that has Elemental to one
  that doesn't

code/extensions/ElementPageExtension.php

@@ -116,25 +116,28 @@ public function onBeforeWrite() {
 			}
 		}
 
-		$elements = $this->owner->ElementArea();
+		if ($this->owner->hasMethod('ElementArea')) {
+			$elements = $this->owner->ElementArea();
 
-		if(!$elements->isInDB()) {
-			$elements->write();
-			$this->owner->ElementAreaID = $elements->ID;
-		}
-		else {
-			// Copy widgets content to Content to enable search
-			$searchableContent = array();
+			if(!$elements->isInDB()) {
+				$elements->write();
+				$this->owner->ElementAreaID = $elements->ID;
+			}
+			else {
+				// Copy widgets content to Content to enable search
+				$searchableContent = array();
 
-			foreach ($elements->Items() as $element) {
+				foreach ($elements->Items() as $element) {
 
-				if($element->config()->exclude_from_content) continue;
-				array_push($searchableContent, strip_tags($element->Content(), '<a>'));
+					if($element->config()->exclude_from_content) continue;
+					array_push($searchableContent, strip_tags($element->Content(), '<a>'));
 
-			}
+				}
 
-			$this->owner->Content = implode(' ', $searchableContent);
+				$this->owner->Content = implode(' ', $searchableContent);
+			}
 		}
+
 
 		// set theme_enabled back to what it was
 		Config::inst()->update('SSViewer', 'theme_enabled', $originalThemeEnabled);
@@ -208,4 +211,34 @@ public function onAfterPublish() {
 			}
 		}
 	}
+
+    /**
+     * Roll back all changes if the parent page has a rollback event
+     *
+     * Only do rollback if it's the 'cancel draft changes' rollback, not a specific version
+     * rollback. 
+     *
+     * @param string $version
+     * @return null
+     */
+    public function onBeforeRollback($version) {
+        if ($version !== 'Live') {
+            // we don't yet have a smart way of rolling back to a specific version
+            return;
+        }
+        if($id = $this->owner->ElementAreaID) {
+			$widgets = Versioned::get_by_stage('BaseElement', 'Live', "ParentID = '$id'");
+			$staged = array();
+
+			foreach($widgets as $widget) {
+				$staged[] = $widget->ID;
+
+                $widget->invokeWithExtensions('onBeforeRollback', $widget);
+
+                $widget->publish("Live", "Stage", false);
+
+                $widget->invokeWithExtensions('onAfterRollback', $widget);
+			}
+		}
+    }
 }