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
@@ -208,18 +203,12 @@ public void selectionChanged(SelectionChangedEvent event) {
     @Override
     public void init(IViewSite site, IMemento memento) throws PartInitException {
         super.init(site, memento);
-        // Restore expanded tree state from file when opening for first time
+        // Set memento with expanded tree state when creating the Tree
         TreeStateHelper.INSTANCE.setMemento(memento);
     }
 
     @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;
 
+    /**
+     * Root expanded tree elements, saved when DrillDownAdapter is used.
+     */
+    private Object[] rootExpandedElements;
+
     /**
      * Listener for theme font change
      */
@@ -197,6 +202,7 @@ public boolean canModify(Object element, String property) {
 
             viewpointFilterProvider = null;
             searchFilter = null;
+            rootExpandedElements = null;
         });
     }
 
@@ -371,6 +377,16 @@ private String getAncestorFolderRenderText(IArchimateModelObject object) {
         return null;
     }
 
+    /**
+     * Get the 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 either the root expanded elements that were visible before the DrillDownAdapter was drilled into,
+     * or the current visible expanded elements if the DrillDownAdapter is "Home".
+     */
+    Object[] getRootVisibleExpandedElements() {
+        return rootExpandedElements != null ? rootExpandedElements : getVisibleExpandedElements();
+    }
+
     // ========================= Model Providers =====================================
 
     /**
@@ -380,6 +396,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 && newInput != null) { // Drilldown has moved out of "Home"
+                rootExpandedElements = getVisibleExpandedElements();
+            }
+            else if(newInput == IEditorModelManager.INSTANCE) { // Drilldown is "Home"
+                rootExpandedElements = null;
+            }
         }
 
         @Override

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

@@ -6,14 +6,11 @@
 package com.archimatetool.editor.views.tree;
 
 import java.io.File;
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 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;
@@ -41,31 +38,31 @@ public class TreeStateHelper {
     private static final String MEMENTO_FILE = "file";
     private static final String MEMENTO_ELEMENTS = "elements";
 
-    // Expanded tree elements
-    private List<Object> expandedElements;
-
-    // State has been restored from memento the first time the Tree is created
-    private boolean restoredFromMemento;
+    private IMemento memento;
 
     private TreeStateHelper() {}
 
     /**
      * Set the Memento when the tree is first created.
-     * We store expanded elements now, as the tree has not yet been created.
-     * We will use the expanded elements later in restoreExpandedTreeElements.
-     * This is called from {@link TreeModelView#init(org.eclipse.ui.IViewSite, IMemento)}, but we only want to do this once.
+     * We will restore any expanded elements later in restoreExpandedTreeElements.
+     * This is called from {@link TreeModelView#init(org.eclipse.ui.IViewSite, IMemento)}.
      */
     void setMemento(IMemento memento) {
-        if(restoredFromMemento || memento == null) {
+        this.memento = memento;
+    }
+
+    /**
+     * Restore expanded elements on TreeView part creation.
+     * This is called from {@link TreeModelView#doCreatePartControl(org.eclipse.swt.widgets.Composite)
+     */
+    void restoreExpandedTreeElements(TreeModelViewer viewer) {
+        if(memento == null) {
             return;
         }
 
-        restoredFromMemento = true;
-
         IMemento expandedMem = memento.getChild(MEMENTO_EXPANDED);
+
         if(expandedMem != null) {
-            expandedElements = new ArrayList<>();
-
             try {
                 for(IMemento elementMem : expandedMem.getChildren(MEMENTO_MODEL)) {
                     String filePath = elementMem.getString(MEMENTO_FILE);
@@ -81,7 +78,7 @@ void setMemento(IMemento memento) {
                                 for(String id : elements.split(ELEMENT_SEP_CHAR)) {
                                     EObject object = objectMap.get(id);
                                     if(object != null) {
-                                        expandedElements.add(object);
+                                        viewer.expandToLevel(object, 1);
                                     }
                                 }
                                 break;
@@ -95,39 +92,19 @@ void setMemento(IMemento memento) {
                 Logger.logError("Error restoring tree state", ex);
             }
         }
-    }
 
-    /**
-     * 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);
-            }
-        });
-
-        if(expandedElements != null) {
-            for(Object element : expandedElements) {
-                viewer.expandToLevel(element, 1);
-            }
-            expandedElements = null;
-        }
+        memento = null;
     }
 
     /**
      * 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) {
+        for(Object object : viewer.getRootVisibleExpandedElements()) {
+            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) {