fix(gui): new implementation for tree state save/load (#2399)

jadx-core/src/main/java/jadx/core/utils/ListUtils.java

@@ -4,6 +4,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Enumeration;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -213,4 +214,14 @@ public class ListUtils {
 		return false;
 	}
 
+	public static <T> List<T> enumerationToList(Enumeration<T> enumeration) {
+		if (enumeration == null || enumeration == Collections.emptyEnumeration()) {
+			return Collections.emptyList();
+		}
+		List<T> list = new ArrayList<>();
+		while (enumeration.hasMoreElements()) {
+			list.add(enumeration.nextElement());
+		}
+		return list;
+	}
 }

jadx-gui/src/main/java/jadx/gui/jobs/TaskWithExtraOnFinish.java

@@ -0,0 +1,90 @@
+package jadx.gui.jobs;
+
+import java.util.function.Consumer;
+
+import org.jetbrains.annotations.Nullable;
+
+import jadx.api.utils.tasks.ITaskExecutor;
+
+/**
+ * Add additional `onFinish` action to the existing task
+ */
+public class TaskWithExtraOnFinish implements IBackgroundTask {
+	private final IBackgroundTask task;
+	private final Consumer<TaskStatus> extraOnFinish;
+
+	public TaskWithExtraOnFinish(IBackgroundTask task, Consumer<TaskStatus> extraOnFinish) {
+		this.task = task;
+		this.extraOnFinish = extraOnFinish;
+	}
+
+	public TaskWithExtraOnFinish(IBackgroundTask task, Runnable extraOnFinish) {
+		this(task, s -> extraOnFinish.run());
+	}
+
+	@Override
+	public void onFinish(ITaskInfo taskInfo) {
+		task.onFinish(taskInfo);
+		extraOnFinish.accept(taskInfo.getStatus());
+	}
+
+	@Override
+	public String getTitle() {
+		return task.getTitle();
+	}
+
+	@Override
+	public ITaskExecutor scheduleTasks() {
+		return task.scheduleTasks();
+	}
+
+	@Override
+	public void onDone(ITaskInfo taskInfo) {
+		task.onDone(taskInfo);
+	}
+
+	@Override
+	public @Nullable Consumer<ITaskProgress> getProgressListener() {
+		return task.getProgressListener();
+	}
+
+	@Override
+	public @Nullable ITaskProgress getTaskProgress() {
+		return task.getTaskProgress();
+	}
+
+	@Override
+	public boolean canBeCanceled() {
+		return task.canBeCanceled();
+	}
+
+	@Override
+	public boolean isCanceled() {
+		return task.isCanceled();
+	}
+
+	@Override
+	public void cancel() {
+		task.cancel();
+	}
+
+	@Override
+	public int timeLimit() {
+		return task.timeLimit();
+	}
+
+	@Override
+	public boolean checkMemoryUsage() {
+		return task.checkMemoryUsage();
+	}
+
+	@Override
+	public int getCancelTimeoutMS() {
+		return task.getCancelTimeoutMS();
+	}
+
+	@Override
+	public int getShutdownTimeoutMS() {
+		return task.getShutdownTimeoutMS();
+	}
+}

jadx-gui/src/main/java/jadx/gui/settings/JadxProject.java

@@ -5,15 +5,12 @@ import java.io.Writer;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 import java.util.StringJoiner;
-import java.util.TreeSet;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -51,7 +48,6 @@ import static jadx.core.utils.GsonUtils.interfaceReplace;
 public class JadxProject {
 	private static final Logger LOG = LoggerFactory.getLogger(JadxProject.class);
 
-	private static final int CURRENT_PROJECT_VERSION = 1;
 	public static final String PROJECT_EXTENSION = "jadx";
 
 	private static final int SEARCH_HISTORY_LIMIT = 30;
@@ -136,25 +132,13 @@ public class JadxProject {
 		changed();
 	}
 
-	public List<String[]> getTreeExpansions() {
-		return data.getTreeExpansions();
+	public void setTreeExpansions(List<String> list) {
+		data.setTreeExpansionsV2(list);
+		changed();
 	}
 
-	public void addTreeExpansion(String[] expansion) {
-		data.getTreeExpansions().removeIf(arr -> Arrays.equals(arr, expansion));
-		data.getTreeExpansions().add(expansion);
-	}
-
-	public void removeTreeExpansion(String[] expansion) {
-		data.getTreeExpansions().removeIf(strings -> isParentOfExpansion(expansion, strings));
-	}
-
-	private void reduceTreeExpansions() {
-		// remove same entries before a project file save
-		// this is mostly needed for old projects ('add' guard don't work for existed entries)
-		Set<String[]> set = new TreeSet<>((a, b) -> Arrays.equals(a, b) ? 0 : Integer.compare(Arrays.hashCode(a), Arrays.hashCode(b)));
-		set.addAll(data.getTreeExpansions());
-		data.setTreeExpansions(new ArrayList<>(set));
+	public List<String> getTreeExpansions() {
+		return data.getTreeExpansionsV2();
 	}
 
 	private boolean isParentOfExpansion(String[] parent, String[] child) {
@@ -287,10 +271,6 @@ public class JadxProject {
 		mainWindow.updateProject(this);
 	}
 
-	private void onSave() {
-		reduceTreeExpansions();
-	}
-
 	public String getName() {
 		return name;
 	}
@@ -314,7 +294,6 @@ public class JadxProject {
 	}
 
 	public void save() {
-		onSave();
 		Path savePath = getProjectPath();
 		if (savePath != null) {
 			Path basePath = savePath.toAbsolutePath().getParent();
@@ -333,7 +312,6 @@ public class JadxProject {
 			project.data = loadProjectData(path);
 			project.saved = true;
 			project.setProjectPath(path);
-			project.upgrade();
 			return project;
 		} catch (Exception e) {
 			LOG.error("Error loading project", e);
@@ -359,20 +337,4 @@ public class JadxProject {
 				.registerTypeAdapter(IJavaCodeRef.class, interfaceReplace(JadxCodeRef.class))
 				.create();
 	}
-
-	private void upgrade() {
-		int fromVersion = data.getProjectVersion();
-		if (fromVersion == CURRENT_PROJECT_VERSION) {
-			return;
-		}
-		LOG.debug("upgrade project settings from version: {} to {}", fromVersion, CURRENT_PROJECT_VERSION);
-		if (fromVersion == 0) {
-			fromVersion++;
-		}
-		if (fromVersion != CURRENT_PROJECT_VERSION) {
-			throw new JadxRuntimeException("Project update failed");
-		}
-		data.setProjectVersion(CURRENT_PROJECT_VERSION);
-		save();
-	}
 }

jadx-gui/src/main/java/jadx/gui/settings/data/ProjectData.java

@@ -13,10 +13,9 @@ import org.jetbrains.annotations.Nullable;
 import jadx.api.data.impl.JadxCodeData;
 
 public class ProjectData {
-
-	private int projectVersion = 1;
+	private int projectVersion = 2;
 	private List<Path> files = new ArrayList<>();
-	private List<String[]> treeExpansions = new ArrayList<>();
+	private List<String> treeExpansionsV2 = new ArrayList<>();
 	private JadxCodeData codeData = new JadxCodeData();
 	private List<TabViewState> openTabs = Collections.emptyList();
 	private @Nullable Path mappingsPath;
@@ -33,12 +32,12 @@ public class ProjectData {
 		this.files = Objects.requireNonNull(files);
 	}
 
-	public List<String[]> getTreeExpansions() {
-		return treeExpansions;
+	public List<String> getTreeExpansionsV2() {
+		return treeExpansionsV2;
 	}
 
-	public void setTreeExpansions(List<String[]> treeExpansions) {
-		this.treeExpansions = treeExpansions;
+	public void setTreeExpansionsV2(List<String> treeExpansionsV2) {
+		this.treeExpansionsV2 = treeExpansionsV2;
 	}
 
 	public JadxCodeData getCodeData() {

jadx-gui/src/main/java/jadx/gui/tree/TreeExpansionService.java

@@ -0,0 +1,212 @@
+package jadx.gui.tree;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import javax.swing.JTree;
+import javax.swing.tree.DefaultTreeModel;
+import javax.swing.tree.TreeNode;
+import javax.swing.tree.TreePath;
+
+import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import jadx.api.metadata.ICodeNodeRef;
+import jadx.core.dex.nodes.RootNode;
+import jadx.core.utils.Utils;
+import jadx.core.utils.exceptions.JadxRuntimeException;
+import jadx.gui.treemodel.JClass;
+import jadx.gui.treemodel.JNode;
+import jadx.gui.treemodel.JPackage;
+import jadx.gui.treemodel.JRoot;
+import jadx.gui.ui.MainWindow;
+import jadx.gui.utils.JNodeCache;
+import jadx.gui.utils.NLS;
+import jadx.gui.utils.UiUtils;
+
+public class TreeExpansionService {
+	private static final Logger LOG = LoggerFactory.getLogger(TreeExpansionService.class);
+	private static final boolean DEBUG = UiUtils.JADX_GUI_DEBUG;
+
+	private static final Comparator<TreePath> PATH_LENGTH_REVERSE = Comparator.comparingInt(p -> -p.getPathCount());
+
+	private final MainWindow mainWindow;
+	private final JTree tree;
+	private final JNodeCache nodeCache;
+
+	public TreeExpansionService(MainWindow mainWindow, JTree tree) {
+		this.mainWindow = mainWindow;
+		this.tree = tree;
+		this.nodeCache = mainWindow.getCacheObject().getNodeCache();
+	}
+
+	public List<String> save() {
+		if (tree.getRowCount() == 0 || mainWindow.getWrapper().getCurrentDecompiler().isEmpty()) {
+			return Collections.emptyList();
+		}
+		List<TreePath> expandedPaths = collectExpandedPaths(tree);
+		List<String> list = new ArrayList<>();
+		for (TreePath expandedPath : expandedPaths) {
+			list.add(savePath(expandedPath));
+		}
+		if (DEBUG) {
+			LOG.debug("Saving tree expansions:\n {}", Utils.listToString(list, "\n "));
+		}
+		return list;
+	}
+
+	public void load(List<String> treeExpansions) {
+		List<TreePath> expandedPaths = new ArrayList<>();
+		mainWindow.getBackgroundExecutor().execute(NLS.str("progress.load"),
+				() -> {
+					loadPaths(treeExpansions, expandedPaths);
+					// send expand event to load sub-nodes and wait for completion
+					UiUtils.uiRunAndWait(() -> expandedPaths.forEach(path -> {
+						try {
+							tree.fireTreeWillExpand(path);
+						} catch (Exception e) {
+							throw new JadxRuntimeException("Tree expand error", e);
+						}
+					}));
+				},
+				s -> {
+					// expand paths after a loading task is finished
+					expandedPaths.forEach(tree::expandPath);
+				});
+	}
+
+	private void loadPaths(List<String> treeExpansions, List<TreePath> expandedPaths) {
+		if (DEBUG) {
+			LOG.debug("Restoring tree expansions:\n {}", Utils.listToString(treeExpansions, "\n "));
+		}
+		for (String treeExpansion : treeExpansions) {
+			try {
+				TreePath treePath = loadPath(treeExpansion);
+				if (treePath != null) {
+					expandedPaths.add(treePath);
+				}
+			} catch (Exception e) {
+				LOG.warn("Failed to load tree expansion entry: {}", treeExpansion, e);
+			}
+		}
+		if (DEBUG) {
+			LOG.debug("Restored expanded tree paths:\n {}", Utils.listToString(expandedPaths, "\n "));
+		}
+	}
+
+	private String savePath(TreePath path) {
+		JNode node = (JNode) path.getLastPathComponent();
+		if (node instanceof JPackage) {
+			return "p:" + ((JPackage) node).getPkg().getRawFullName();
+		}
+		if (node instanceof JClass) {
+			return "c:" + ((JClass) node).getCls().getRawName();
+		}
+		return Arrays.stream(path.getPath())
+				.map(p -> ((JNode) p).getID())
+				.skip(1) // skip root
+				.collect(Collectors.joining("//", "t:", ""));
+	}
+
+	private @Nullable TreePath loadPath(String pathStr) {
+		String pathData = pathStr.substring(2);
+		switch (pathStr.charAt(0)) {
+			case 'c':
+				return getTreePathForRef(getRoot().resolveRawClass(pathData));
+			case 'p':
+				return getTreePathForRef(getRoot().resolvePackage(pathData));
+			case 't':
+				return resolveTreePath(pathData.split("//"));
+
+			default:
+				throw new JadxRuntimeException("Unknown tree expansion path type: " + pathStr);
+		}
+	}
+
+	private @Nullable TreePath resolveTreePath(String[] pathArr) {
+		JNode current = (JNode) tree.getModel().getRoot();
+		for (String nodeStr : pathArr) {
+			JNode node = current.searchNode(n -> n.getID().equals(nodeStr));
+			if (node == null) {
+				if (DEBUG) {
+					List<String> children = current.childrenList().stream()
+							.map(n -> ((JNode) n).getID())
+							.collect(Collectors.toList());
+					LOG.warn("Failed to restore path: {}, node '{}' not found in '{}' children: {}",
+							Arrays.toString(pathArr), nodeStr, current, children);
+				}
+				return null;
+			}
+			current = node;
+		}
+		return new TreePath(current.getPath());
+	}
+
+	private @Nullable TreePath getTreePathForRef(@Nullable ICodeNodeRef ref) {
+		if (ref == null) {
+			return null;
+		}
+		JNode node = nodeCache.makeFrom(ref);
+		if (node.getParent() == null) {
+			if (DEBUG) {
+				LOG.warn("Resolving node not from tree: {}", node);
+			}
+			JNode treeNode = ((JRoot) tree.getModel().getRoot()).searchNode(node);
+			if (treeNode == null) {
+				if (DEBUG) {
+					LOG.error("Node not found in tree: {}", node);
+				}
+				return null;
+			}
+			node = treeNode;
+		}
+		TreeNode[] pathNodes = ((DefaultTreeModel) tree.getModel()).getPathToRoot(node);
+		if (pathNodes == null) {
+			return null;
+		}
+		return new TreePath(pathNodes);
+	}
+
+	private static List<TreePath> collectExpandedPaths(JTree tree) {
+		TreePath root = tree.getPathForRow(0);
+		Enumeration<TreePath> expandedDescendants = tree.getExpandedDescendants(root);
+		if (expandedDescendants == null) {
+			return Collections.emptyList();
+		}
+		List<TreePath> expandedPaths = new ArrayList<>();
+		while (expandedDescendants.hasMoreElements()) {
+			TreePath path = expandedDescendants.nextElement();
+			if (path.getPathCount() > 1) {
+				expandedPaths.add(path);
+			}
+		}
+		// filter out sub-paths
+		expandedPaths.sort(PATH_LENGTH_REVERSE); // put the longest paths before sub-paths
+		List<TreePath> result = new ArrayList<>();
+		for (TreePath path : expandedPaths) {
+			if (!isSubPath(result, path)) {
+				result.add(path);
+			}
+		}
+		return result;
+	}
+
+	private static boolean isSubPath(List<TreePath> paths, TreePath path) {
+		for (TreePath addedPath : paths) {
+			if (path.isDescendant(addedPath)) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+	private RootNode getRoot() {
+		return mainWindow.getWrapper().getDecompiler().getRoot();
+	}
+}

jadx-gui/src/main/java/jadx/gui/treemodel/JInputFiles.java

@@ -38,6 +38,11 @@ public class JInputFiles extends JNode {
 		return INPUT_FILES_ICON;
 	}
 
+	@Override
+	public String getID() {
+		return "JInputFiles";
+	}
+
 	@Override
 	public String makeString() {
 		return NLS.str("tree.input_files");

jadx-gui/src/main/java/jadx/gui/treemodel/JInputScripts.java

@@ -39,6 +39,11 @@ public class JInputScripts extends JNode {
 		return INPUT_SCRIPTS_ICON;
 	}
 
+	@Override
+	public String getID() {
+		return "JInputScripts";
+	}
+
 	@Override
 	public String makeString() {
 		return NLS.str("tree.input_scripts");

jadx-gui/src/main/java/jadx/gui/treemodel/JInputs.java

@@ -45,6 +45,11 @@ public class JInputs extends JNode {
 		return INPUTS_ICON;
 	}
 
+	@Override
+	public String getID() {
+		return "JInputs";
+	}
+
 	@Override
 	public String makeString() {
 		return NLS.str("tree.inputs_title");

jadx-gui/src/main/java/jadx/gui/treemodel/JLoadableNode.java

@@ -1,5 +1,9 @@
 package jadx.gui.treemodel;
 
+import java.util.function.Predicate;
+
+import org.jetbrains.annotations.Nullable;
+
 import jadx.gui.jobs.IBackgroundTask;
 
 public abstract class JLoadableNode extends JNode {
@@ -7,5 +11,17 @@ public abstract class JLoadableNode extends JNode {
 
 	public abstract void loadNode();
 
-	public abstract IBackgroundTask getLoadTask();
+	public abstract @Nullable IBackgroundTask getLoadTask();
+
+	@Override
+	public @Nullable JNode searchNode(Predicate<JNode> filter) {
+		loadNode();
+		return super.searchNode(filter);
+	}
+
+	@Override
+	public @Nullable JNode removeNode(Predicate<JNode> filter) {
+		loadNode();
+		return super.removeNode(filter);
+	}
 }

jadx-gui/src/main/java/jadx/gui/treemodel/JNode.java

@@ -2,11 +2,13 @@ package jadx.gui.treemodel;
 
 import java.util.Comparator;
 import java.util.Enumeration;
+import java.util.List;
 import java.util.function.Predicate;
 
 import javax.swing.Icon;
 import javax.swing.JPopupMenu;
 import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.TreeNode;
 
 import org.fife.ui.rsyntaxtextarea.SyntaxConstants;
 import org.jetbrains.annotations.NotNull;
@@ -15,6 +17,7 @@ import org.jetbrains.annotations.Nullable;
 import jadx.api.ICodeInfo;
 import jadx.api.JavaNode;
 import jadx.api.metadata.ICodeNodeRef;
+import jadx.core.utils.ListUtils;
 import jadx.gui.ui.MainWindow;
 import jadx.gui.ui.panel.ContentPanel;
 import jadx.gui.ui.tab.TabbedPane;
@@ -48,7 +51,6 @@ public abstract class JNode extends DefaultMutableTreeNode implements Comparable
 		return SyntaxConstants.SYNTAX_STYLE_NONE;
 	}
 
-	@NotNull
 	public ICodeInfo getCodeInfo() {
 		return ICodeInfo.EMPTY;
 	}
@@ -75,6 +77,15 @@ public abstract class JNode extends DefaultMutableTreeNode implements Comparable
 		return null;
 	}
 
+	/**
+	 * JNode identifier.
+	 * Should be locale independent.
+	 * TODO: implement list or enum of custom tree nodes to allow extension from plugins
+	 */
+	public String getID() {
+		return makeString();
+	}
+
 	public abstract String makeString();
 
 	public String makeStringHtml() {
@@ -139,6 +150,10 @@ public abstract class JNode extends DefaultMutableTreeNode implements Comparable
 		return null;
 	}
 
+	public List<TreeNode> childrenList() {
+		return ListUtils.enumerationToList(this.children());
+	}
+
 	private static final Comparator<JNode> COMPARATOR = Comparator
 			.comparing(JNode::makeLongString)
 			.thenComparingInt(JNode::getPos);

jadx-gui/src/main/java/jadx/gui/treemodel/JResource.java

@@ -20,7 +20,6 @@ import jadx.api.impl.SimpleCodeInfo;
 import jadx.core.utils.ListUtils;
 import jadx.core.utils.Utils;
 import jadx.core.xmlgen.ResContainer;
-import jadx.gui.jobs.IBackgroundTask;
 import jadx.gui.jobs.SimpleTask;
 import jadx.gui.ui.MainWindow;
 import jadx.gui.ui.codearea.BinaryContentPanel;
@@ -63,7 +62,7 @@ public class JResource extends JLoadableNode {
 
 	private transient volatile boolean loaded;
 	private transient List<JResource> subNodes = Collections.emptyList();
-	private transient ICodeInfo content;
+	private transient ICodeInfo content = ICodeInfo.EMPTY;
 
 	public JResource(ResourceFile resFile, String name, JResType type) {
 		this(resFile, name, name, type);
@@ -77,7 +76,7 @@ public class JResource extends JLoadableNode {
 		this.loaded = false;
 	}
 
-	public final void update() {
+	public synchronized void update() {
 		removeAllChildren();
 		if (Utils.isEmpty(subNodes)) {
 			if (type == JResType.DIR || type == JResType.ROOT
@@ -91,6 +90,10 @@ public class JResource extends JLoadableNode {
 				res.update();
 				add(res);
 			}
+			if (type != JResType.FILE) {
+				// no content, nothing to load
+				loaded = true;
+			}
 		}
 	}
 
@@ -101,7 +104,10 @@ public class JResource extends JLoadableNode {
 	}
 
 	@Override
-	public synchronized IBackgroundTask getLoadTask() {
+	public synchronized SimpleTask getLoadTask() {
+		if (loaded) {
+			return null;
+		}
 		return new SimpleTask(NLS.str("progress.load"), this::getCodeInfo, this::update);
 	}
 
@@ -316,6 +322,14 @@ public class JResource extends JLoadableNode {
 		return null;
 	}
 
+	@Override
+	public String getID() {
+		if (type == JResType.ROOT) {
+			return "JResources";
+		}
+		return makeString();
+	}
+
 	@Override
 	public String makeString() {
 		return shortName;

jadx-gui/src/main/java/jadx/gui/treemodel/JRoot.java

@@ -161,6 +161,11 @@ public class JRoot extends JNode {
 		return null;
 	}
 
+	@Override
+	public String getID() {
+		return "JRoot";
+	}
+
 	@Override
 	public String makeString() {
 		JadxProject project = wrapper.getProject();

jadx-gui/src/main/java/jadx/gui/treemodel/JSources.java

@@ -44,6 +44,11 @@ public class JSources extends JNode {
 		return null;
 	}
 
+	@Override
+	public String getID() {
+		return "JSources";
+	}
+
 	@Override
 	public String makeString() {
 		return NLS.str("tree.sources_title");

jadx-gui/src/main/java/jadx/gui/ui/MainWindow.java

@@ -27,7 +27,6 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.List;
@@ -102,7 +101,9 @@ import jadx.gui.events.types.JadxGuiEventsImpl;
 import jadx.gui.jobs.BackgroundExecutor;
 import jadx.gui.jobs.DecompileTask;
 import jadx.gui.jobs.ExportTask;
+import jadx.gui.jobs.IBackgroundTask;
 import jadx.gui.jobs.TaskStatus;
+import jadx.gui.jobs.TaskWithExtraOnFinish;
 import jadx.gui.logs.LogCollector;
 import jadx.gui.logs.LogOptions;
 import jadx.gui.logs.LogPanel;
@@ -113,8 +114,8 @@ import jadx.gui.settings.JadxProject;
 import jadx.gui.settings.JadxSettings;
 import jadx.gui.settings.ui.JadxSettingsWindow;
 import jadx.gui.settings.ui.plugins.PluginSettings;
+import jadx.gui.tree.TreeExpansionService;
 import jadx.gui.treemodel.ApkSignature;
-import jadx.gui.treemodel.JClass;
 import jadx.gui.treemodel.JLoadableNode;
 import jadx.gui.treemodel.JNode;
 import jadx.gui.treemodel.JPackage;
@@ -199,6 +200,7 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 	private final transient CacheManager cacheManager;
 	private final transient BackgroundExecutor backgroundExecutor;
 	private final transient JadxGuiEventsImpl events = new JadxGuiEventsImpl();
+	private final transient TreeExpansionService treeExpansionService;
 
 	private final TabsController tabsController;
 	private final NavigationController navController;
@@ -271,6 +273,7 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 		initUI();
 		this.editorSyncManager = new EditorSyncManager(this, tabbedPane);
 		this.backgroundExecutor = new BackgroundExecutor(settings, progressPane);
+		this.treeExpansionService = new TreeExpansionService(this, tree);
 		initMenuAndToolbar();
 		UiUtils.setWindowIcons(this);
 		this.shortcutsController.registerMouseEventListener(this);
@@ -603,6 +606,7 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 
 	private void saveAll() {
 		saveOpenTabs();
+		project.setTreeExpansions(treeExpansionService.save());
 		BreakpointManager.saveAndExit();
 	}
 
@@ -645,7 +649,7 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 		initTree();
 		updateLiveReload(project.isEnableLiveReload());
 		BreakpointManager.init(project.getFilePaths().get(0).toAbsolutePath().getParent());
-
+		treeExpansionService.load(project.getTreeExpansions());
 		List<EditorViewState> openTabs = project.getOpenTabs(this);
 		backgroundExecutor.execute(NLS.str("progress.load"),
 				() -> preLoadOpenTabs(openTabs),
@@ -839,41 +843,14 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 	public void reloadTree() {
 		treeReloading = true;
 		treeUpdateListener.forEach(listener -> listener.accept(treeRoot));
-
 		treeModel.reload();
-		List<String[]> treeExpansions = project.getTreeExpansions();
-		if (!treeExpansions.isEmpty()) {
-			expand(treeRoot, treeExpansions);
-		} else {
-			tree.expandRow(1);
-		}
-
 		treeReloading = false;
 	}
 
 	public void rebuildPackagesTree() {
-		cacheObject.setPackageHelper(null);
 		treeRoot.update();
 	}
 
-	private void expand(TreeNode node, List<String[]> treeExpansions) {
-		TreeNode[] pathNodes = treeModel.getPathToRoot(node);
-		if (pathNodes == null) {
-			return;
-		}
-		TreePath path = new TreePath(pathNodes);
-		String[] pathExpansion = getPathExpansion(path);
-		for (String[] expansion : treeExpansions) {
-			if (Arrays.equals(expansion, pathExpansion)) {
-				tree.expandPath(path);
-				break;
-			}
-		}
-		for (int i = node.getChildCount() - 1; i >= 0; i--) {
-			expand(node.getChildAt(i), treeExpansions);
-		}
-	}
-
 	private void toggleFlattenPackage() {
 		setFlattenPackage(!isFlattenPackage);
 	}
@@ -1216,6 +1193,8 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 					ExceptionDialog.throwTestException();
 				}
 			});
+		}
+		if (UiUtils.JADX_GUI_DEBUG) {
 			JCheckBoxMenuItem uiWatchDog = new JCheckBoxMenuItem(new ActionHandler("UI WatchDog", UIWatchDog::toggle));
 			uiWatchDog.setState(UIWatchDog.onStart());
 			help.add(uiWatchDog);
@@ -1374,19 +1353,14 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 				Object node = path.getLastPathComponent();
 				if (node instanceof JLoadableNode) {
 					JLoadableNode treeNode = (JLoadableNode) node;
-					backgroundExecutor.execute(treeNode.getLoadTask());
-					// schedule update for expanded nodes in a tree
-					backgroundExecutor.execute(NLS.str("progress.load"),
-							UiUtils.EMPTY_RUNNABLE,
+					IBackgroundTask loadTask = treeNode.getLoadTask();
+					if (loadTask != null) {
+						backgroundExecutor.execute(new TaskWithExtraOnFinish(loadTask,
 								status -> {
 									if (!treeReloading) {
 										treeModel.nodeStructureChanged(treeNode);
-									project.addTreeExpansion(getPathExpansion(event.getPath()));
 									}
-							});
-				} else {
-					if (!treeReloading) {
-						project.addTreeExpansion(getPathExpansion(event.getPath()));
+								}));
 					}
 				}
 			}
@@ -1394,7 +1368,6 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 			@Override
 			public void treeWillCollapse(TreeExpansionEvent event) {
 				if (!treeReloading) {
-					project.removeTreeExpansion(getPathExpansion(event.getPath()));
 					update();
 				}
 			}
@@ -1444,35 +1417,6 @@ public class MainWindow extends JFrame implements ExportProjectDialog.ExportProj
 		setTitle(DEFAULT_TITLE);
 	}
 
-	private static String[] getPathExpansion(TreePath path) {
-		List<String> pathList = new ArrayList<>();
-		while (path != null) {
-			Object node = path.getLastPathComponent();
-			String name;
-			if (node instanceof JClass) {
-				name = ((JClass) node).getCls().getClassNode().getClassInfo().getFullName();
-			} else {
-				name = node.toString();
-			}
-			pathList.add(name);
-			path = path.getParentPath();
-		}
-		return pathList.toArray(new String[0]);
-	}
-
-	public static void getExpandedPaths(JTree tree, TreePath path, List<TreePath> list) {
-		if (tree.isExpanded(path)) {
-			list.add(path);
-
-			TreeNode node = (TreeNode) path.getLastPathComponent();
-			for (int i = node.getChildCount() - 1; i >= 0; i--) {
-				TreeNode n = node.getChildAt(i);
-				TreePath child = path.pathByAddingChild(n);
-				getExpandedPaths(tree, child, list);
-			}
-		}
-	}
-
 	public void setLocationAndPosition() {
 		if (settings.loadWindowPos(this)) {
 			return;

jadx-gui/src/main/java/jadx/gui/ui/tab/TabsController.java

@@ -5,7 +5,6 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import org.jetbrains.annotations.Nullable;
@@ -17,12 +16,13 @@ import jadx.api.metadata.ICodeAnnotation;
 import jadx.api.metadata.ICodeNodeRef;
 import jadx.api.metadata.annotations.NodeDeclareRef;
 import jadx.gui.jobs.SimpleTask;
-import jadx.gui.jobs.TaskStatus;
+import jadx.gui.jobs.TaskWithExtraOnFinish;
 import jadx.gui.treemodel.JClass;
 import jadx.gui.treemodel.JNode;
 import jadx.gui.ui.MainWindow;
 import jadx.gui.ui.codearea.EditorViewState;
 import jadx.gui.utils.JumpPosition;
+import jadx.gui.utils.UiUtils;
 
 public class TabsController {
 	private static final Logger LOG = LoggerFactory.getLogger(TabsController.class);
@@ -113,16 +113,12 @@ public class TabsController {
 
 	private void loadCodeWithUIAction(JClass cls, Runnable action) {
 		SimpleTask loadTask = cls.getLoadTask();
-		mainWindow.getBackgroundExecutor().execute(
-				new SimpleTask(loadTask.getTitle(),
-						loadTask.getJobs(),
-						status -> {
-							Consumer<TaskStatus> onFinish = loadTask.getOnFinish();
-							if (onFinish != null) {
-								onFinish.accept(status);
+		if (loadTask == null) {
+			// already loaded
+			UiUtils.uiRun(action);
+			return;
 		}
-							action.run();
-						}));
+		mainWindow.getBackgroundExecutor().execute(new TaskWithExtraOnFinish(loadTask, action));
 	}
 
 	/**

jadx-gui/src/main/java/jadx/gui/update/JadxUpdate.kt

@@ -6,7 +6,6 @@ import jadx.core.Jadx
 import jadx.core.plugins.versions.VersionComparator
 import jadx.core.utils.GsonUtils.buildGson
 import jadx.gui.settings.JadxUpdateChannel
-import org.jetbrains.kotlin.konan.file.use
 import java.io.InputStreamReader
 import java.net.HttpURLConnection
 import java.net.URI

jadx-gui/src/main/java/jadx/gui/utils/UiUtils.java

@@ -40,6 +40,7 @@ import org.slf4j.LoggerFactory;
 
 import com.formdev.flatlaf.extras.FlatSVGIcon;
 
+import jadx.commons.app.JadxCommonEnv;
 import jadx.core.dex.info.AccessInfo;
 import jadx.core.dex.instructions.args.ArgType;
 import jadx.core.utils.StringUtils;
@@ -51,6 +52,8 @@ import jadx.gui.ui.codearea.AbstractCodeArea;
 public class UiUtils {
 	private static final Logger LOG = LoggerFactory.getLogger(UiUtils.class);
 
+	public static final boolean JADX_GUI_DEBUG = JadxCommonEnv.getBool("JADX_GUI_DEBUG", false);
+
 	/**
 	 * The minimum about of memory in bytes we are trying to keep free, otherwise the application may
 	 * run out of heap

jadx-gui/src/main/java/jadx/gui/utils/dbg/UIWatchDog.java

@@ -20,17 +20,12 @@ import jadx.gui.utils.UiUtils;
 public class UIWatchDog {
 	private static final Logger LOG = LoggerFactory.getLogger(UIWatchDog.class);
 
-	private static final boolean RUN_ON_START = true;
-
 	private static final int UI_MAX_DELAY_MS = 200;
 	private static final int CHECK_INTERVAL_MS = 50;
 
 	public static boolean onStart() {
-		if (RUN_ON_START) {
-			UiUtils.uiRun(UIWatchDog::toggle);
-			return true;
-		}
-		return false;
+		UiUtils.uiRunAndWait(UIWatchDog::toggle);
+		return INSTANCE.isEnabled();
 	}
 
 	public static synchronized void toggle() {
@@ -69,6 +64,10 @@ public class UIWatchDog {
 		}
 	}
 
+	private boolean isEnabled() {
+		return enabled.get();
+	}
+
 	@SuppressWarnings("BusyWait")
 	private void start(Thread uiThread) {
 		LOG.debug("UI watchdog started");