Refactor TreeStateHelper part 2

- Store the root expanded tree elements when the DrillDownAdapter sets the tree model input

- When we save the tree state and we have drilled into the DrillDownAdapter we can save these root elements

com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelView.java

@@ -155,11 +155,6 @@ public void doCreatePartControl(Composite parent) {
         // Drill down
         fDrillDownAdapter = new DrillDownAdapter(fTreeViewer);
 
-        // Set drill down to home when disposing parent (not when the tree is disposed as the content provider is null at that point)
-        parent.addDisposeListener(e -> {
-            setDrillDownHome();
-        });
-
         // Listen to Double-click and press Return Action
         fTreeViewer.addDoubleClickListener(new IDoubleClickListener() {
             @Override
@@ -214,12 +209,6 @@ public void init(IViewSite site, IMemento memento) throws PartInitException {
 
     @Override
     public void saveState(IMemento memento) {
-        // saveState() is called periodically by Eclipse when auto-saving the workbench so only reset the drill-down when closing the Workbench.
-        // This period is set in org.eclipse.ui.internal.IPreferenceConstants#WORKBENCH_SAVE_INTERVAL and the default is 5 minutes.
-        if(PlatformUI.getWorkbench().isClosing()) {
-            setDrillDownHome();
-        }
-
         // Save expanded tree state
         TreeStateHelper.INSTANCE.saveStateToMemento(fTreeViewer, memento);
     }

com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelViewer.java

@@ -75,6 +75,11 @@ public class TreeModelViewer extends TreeViewer {
      */
     private SearchFilter searchFilter;
 
+    /**
+     * Last root expanded tree elements, saved when DrillDownAdapter is used.
+     */
+    private Object[] lastRootExpandedTreeElements;
+
     /**
      * Listener for theme font change
      */
@@ -197,6 +202,7 @@ public boolean canModify(Object element, String property) {
 
             viewpointFilterProvider = null;
             searchFilter = null;
+            lastRootExpandedTreeElements = null;
         });
     }
 
@@ -371,6 +377,15 @@ private String getAncestorFolderRenderText(IArchimateModelObject object) {
         return null;
     }
 
+    /**
+     * Get the last root expanded elements in case the DrillDownAdapter is active.
+     * Note that some of these elements might have been deleted when the tree was drilled into.
+     * Returns null if the DrillDownAdapter is at "home".
+     */
+    Object[] getLastRootVisibleExpandedElements() {
+        return lastRootExpandedTreeElements;
+    }
+
     // ========================= Model Providers =====================================
 
     /**
@@ -380,6 +395,14 @@ private class ModelTreeViewerContentProvider implements ITreeContentProvider {
 
         @Override
         public void inputChanged(Viewer v, Object oldInput, Object newInput) {
+            // The DrillDownAdapter sets the new input when calling DrillDownAdapter#goInto().
+            // We save the current expanded elements in the root state so we can persist these when we save the expanded tree state
+            if(oldInput == IEditorModelManager.INSTANCE) { // Drilldown has moved out of "Home"
+                lastRootExpandedTreeElements = getVisibleExpandedElements();
+            }
+            else if(newInput == IEditorModelManager.INSTANCE) { // Drilldown is "Home"
+                lastRootExpandedTreeElements = null;
+            }
         }
 
         @Override

com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeStateHelper.java

@@ -13,7 +13,6 @@
 import java.util.Map.Entry;
 
 import org.eclipse.emf.ecore.EObject;
-import org.eclipse.jface.viewers.TreeViewer;
 import org.eclipse.ui.IMemento;
 
 import com.archimatetool.editor.model.IEditorModelManager;
@@ -44,9 +43,6 @@ public class TreeStateHelper {
     // Expanded tree elements
     private List<Object> expandedElements;
 
-    // State has been restored from memento the first time the Tree is created
-    private boolean restoredFromMemento;
-
     private TreeStateHelper() {}
 
     /**
@@ -56,12 +52,10 @@ private TreeStateHelper() {}
      * This is called from {@link TreeModelView#init(org.eclipse.ui.IViewSite, IMemento)}, but we only want to do this once.
      */
     void setMemento(IMemento memento) {
-        if(restoredFromMemento || memento == null) {
+        if(memento == null) {
             return;
         }
 
-        restoredFromMemento = true;
-
         IMemento expandedMem = memento.getChild(MEMENTO_EXPANDED);
         if(expandedMem != null) {
             expandedElements = new ArrayList<>();
@@ -101,16 +95,7 @@ void setMemento(IMemento memento) {
      * Restore expanded elements on TreeView part creation.
      * This is called from {@link TreeModelView#doCreatePartControl(org.eclipse.swt.widgets.Composite)
      */
-    void restoreExpandedTreeElements(TreeViewer viewer) {
-        // Store expanded tree elements if the TreeViewer is closed so they can be restored when it's re-opened
-        // We could restore from the memento but this is more efficient and the drill-down is reset when closing the Tree.
-        viewer.getTree().addDisposeListener(e -> {
-            expandedElements = new ArrayList<>();
-            for(Object element : viewer.getVisibleExpandedElements()) {
-                expandedElements.add(element);
-            }
-        });
-
+    void restoreExpandedTreeElements(TreeModelViewer viewer) {
         if(expandedElements != null) {
             for(Object element : expandedElements) {
                 viewer.expandToLevel(element, 1);
@@ -123,11 +108,19 @@ void restoreExpandedTreeElements(TreeViewer viewer) {
      * Save expanded state of tree elements to a memento.
      * This is called from {@link TreeModelView#saveState(IMemento)}
      */
-    void saveStateToMemento(TreeViewer viewer, IMemento memento) {
+    void saveStateToMemento(TreeModelViewer viewer, IMemento memento) {
         Map<File, String> map = new HashMap<>();
 
-        for(Object object : viewer.getVisibleExpandedElements()) {
-            if(object instanceof IArchimateModelObject modelObject) {
+        // Get expanded elements from either:
+        // (1) last saved expanded elements when drilldown was applied or
+        // (2) current expanded elements if not in drilldown
+        Object[] elements = viewer.getLastRootVisibleExpandedElements();
+        if(elements == null) {
+            elements = viewer.getVisibleExpandedElements();
+        }
+
+        for(Object object : elements) {
+            if(object instanceof IArchimateModelObject modelObject && modelObject.getArchimateModel() != null) { // Check it wasn't deleted
                 // Only store if model has been saved to file
                 File file = modelObject.getArchimateModel().getFile();
                 if(file != null) {