diff --git a/jadx-core/src/main/java/jadx/api/JadxDecompiler.java b/jadx-core/src/main/java/jadx/api/JadxDecompiler.java
index 315a2afbd..d1cccf062 100644
--- a/jadx-core/src/main/java/jadx/api/JadxDecompiler.java
+++ b/jadx-core/src/main/java/jadx/api/JadxDecompiler.java
@@ -197,8 +197,12 @@ public final class JadxDecompiler {
 				continue;
 			}
 			executor.execute(() -> {
-				cls.decompile();
-				SaveCode.save(outDir, args, cls.getClassNode());
+				try {
+					cls.decompile();
+					SaveCode.save(outDir, args, cls.getClassNode());
+				} catch (Exception e) {
+					LOG.error("Error saving class: {}", cls.getFullName(), e);
+				}
 			});
 		}
 	}
diff --git a/jadx-core/src/main/java/jadx/api/JavaClass.java b/jadx-core/src/main/java/jadx/api/JavaClass.java
index a24094b82..372c76841 100644
--- a/jadx-core/src/main/java/jadx/api/JavaClass.java
+++ b/jadx-core/src/main/java/jadx/api/JavaClass.java
@@ -64,6 +64,10 @@ public final class JavaClass implements JavaNode {
 		}
 	}
 
+	public synchronized void unload() {
+		cls.unload();
+	}
+
 	ClassNode getClassNode() {
 		return cls;
 	}
@@ -107,12 +111,7 @@ public final class JavaClass implements JavaNode {
 					rootDecompiler.getMethodsMap().put(m, javaMethod);
 				}
 			}
-			Collections.sort(mths, new Comparator<JavaMethod>() {
-				@Override
-				public int compare(JavaMethod o1, JavaMethod o2) {
-					return o1.getName().compareTo(o2.getName());
-				}
-			});
+			mths.sort(Comparator.comparing(JavaMethod::getName));
 			this.methods = Collections.unmodifiableList(mths);
 		}
 	}
diff --git a/jadx-core/src/main/java/jadx/core/ProcessClass.java b/jadx-core/src/main/java/jadx/core/ProcessClass.java
index 42214dd13..985510e9b 100644
--- a/jadx-core/src/main/java/jadx/core/ProcessClass.java
+++ b/jadx-core/src/main/java/jadx/core/ProcessClass.java
@@ -10,11 +10,9 @@ import jadx.core.dex.visitors.DepthTraversal;
 import jadx.core.dex.visitors.IDexTreeVisitor;
 import jadx.core.utils.ErrorsCounter;
 
-import static jadx.core.dex.nodes.ProcessState.GENERATED;
 import static jadx.core.dex.nodes.ProcessState.NOT_LOADED;
 import static jadx.core.dex.nodes.ProcessState.PROCESSED;
 import static jadx.core.dex.nodes.ProcessState.STARTED;
-import static jadx.core.dex.nodes.ProcessState.UNLOADED;
 
 public final class ProcessClass {
 
@@ -38,15 +36,9 @@ public final class ProcessClass {
 				if (cls.getState() == PROCESSED && codeGen != null) {
 					processDependencies(cls, passes);
 					codeGen.visit(cls);
-					cls.setState(GENERATED);
 				}
 			} catch (Exception e) {
 				ErrorsCounter.classError(cls, e.getClass().getSimpleName(), e);
-			} finally {
-				if (cls.getState() == GENERATED) {
-					cls.unload();
-					cls.setState(UNLOADED);
-				}
 			}
 		}
 	}
diff --git a/jadx-core/src/main/java/jadx/core/codegen/CodeGen.java b/jadx-core/src/main/java/jadx/core/codegen/CodeGen.java
index 8c33d8862..eae73cb09 100644
--- a/jadx-core/src/main/java/jadx/core/codegen/CodeGen.java
+++ b/jadx-core/src/main/java/jadx/core/codegen/CodeGen.java
@@ -1,12 +1,10 @@
 package jadx.core.codegen;
 
 import jadx.core.dex.nodes.ClassNode;
-import jadx.core.dex.visitors.AbstractVisitor;
 import jadx.core.utils.exceptions.CodegenException;
 
-public class CodeGen extends AbstractVisitor {
+public class CodeGen {
 
-	@Override
 	public boolean visit(ClassNode cls) throws CodegenException {
 		ClassGen clsGen = new ClassGen(cls, cls.root().getArgs());
 		CodeWriter clsCode = clsGen.makeClass();
diff --git a/jadx-core/src/main/java/jadx/core/dex/nodes/ClassNode.java b/jadx-core/src/main/java/jadx/core/dex/nodes/ClassNode.java
index e4adfe0c2..24ab78a42 100644
--- a/jadx-core/src/main/java/jadx/core/dex/nodes/ClassNode.java
+++ b/jadx-core/src/main/java/jadx/core/dex/nodes/ClassNode.java
@@ -39,6 +39,8 @@ import jadx.core.dex.nodes.parser.StaticValuesParser;
 import jadx.core.utils.exceptions.DecodeException;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
+import static jadx.core.dex.nodes.ProcessState.UNLOADED;
+
 public class ClassNode extends LineAttrNode implements ILoadable, IDexNode {
 	private static final Logger LOG = LoggerFactory.getLogger(ClassNode.class);
 
@@ -264,6 +266,7 @@ public class ClassNode extends LineAttrNode implements ILoadable, IDexNode {
 		for (ClassNode innerCls : getInnerClasses()) {
 			innerCls.unload();
 		}
+		setState(UNLOADED);
 	}
 
 	private void buildCache() {
diff --git a/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java b/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java
index 9a4318cd3..47b1a98e4 100644
--- a/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java
+++ b/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java
@@ -403,8 +403,8 @@ public class MethodNode extends LineAttrNode implements ILoadable, IDexNode {
 	}
 
 	public void finishBasicBlocks() {
-		((ArrayList<BlockNode>) blocks).trimToSize();
-		((ArrayList<BlockNode>) exitBlocks).trimToSize();
+		trimList(blocks);
+		trimList(exitBlocks);
 
 		blocks = Collections.unmodifiableList(blocks);
 		exitBlocks = Collections.unmodifiableList(exitBlocks);
@@ -414,6 +414,12 @@ public class MethodNode extends LineAttrNode implements ILoadable, IDexNode {
 		}
 	}
 
+	private void trimList(List<BlockNode> blocks) {
+		if (blocks instanceof ArrayList) {
+			((ArrayList<BlockNode>)blocks).trimToSize();
+		}
+	}
+
 	public List<BlockNode> getBasicBlocks() {
 		return blocks;
 	}
diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/SaveCode.java b/jadx-core/src/main/java/jadx/core/dex/visitors/SaveCode.java
index 5368a10b6..7b04b87ae 100644
--- a/jadx-core/src/main/java/jadx/core/dex/visitors/SaveCode.java
+++ b/jadx-core/src/main/java/jadx/core/dex/visitors/SaveCode.java
@@ -6,6 +6,7 @@ import jadx.api.JadxArgs;
 import jadx.core.codegen.CodeWriter;
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.nodes.ClassNode;
+import jadx.core.utils.exceptions.JadxRuntimeException;
 
 public class SaveCode {
 
@@ -16,6 +17,9 @@ public class SaveCode {
 			return;
 		}
 		CodeWriter clsCode = cls.getCode();
+		if (clsCode == null) {
+			throw new JadxRuntimeException("Code not generated for class " + cls.getFullName());
+		}
 		String fileName = cls.getClassInfo().getFullPath() + ".java";
 		if (args.isFallbackMode()) {
 			fileName += ".jadx";
diff --git a/jadx-gui/src/main/java/jadx/gui/JadxWrapper.java b/jadx-gui/src/main/java/jadx/gui/JadxWrapper.java
index a17221dae..5e79fc301 100644
--- a/jadx-gui/src/main/java/jadx/gui/JadxWrapper.java
+++ b/jadx-gui/src/main/java/jadx/gui/JadxWrapper.java
@@ -52,6 +52,8 @@ public class JadxWrapper {
 						Thread.sleep(500);
 					}
 					progressMonitor.close();
+					LOG.info("decompilation complete, freeing memory ...");
+					decompiler.getClasses().forEach(JavaClass::unload);
 					LOG.info("done");
 				} catch (InterruptedException e) {
 					LOG.error("Save interrupted", e);
diff --git a/jadx-gui/src/main/java/jadx/gui/treemodel/JClass.java b/jadx-gui/src/main/java/jadx/gui/treemodel/JClass.java
index c20060e8b..8998a3541 100644
--- a/jadx-gui/src/main/java/jadx/gui/treemodel/JClass.java
+++ b/jadx-gui/src/main/java/jadx/gui/treemodel/JClass.java
@@ -51,6 +51,7 @@ public class JClass extends JLoadableNode {
 	public synchronized void load() {
 		if (!loaded) {
 			cls.decompile();
+			cls.unload();
 			loaded = true;
 		}
 		update();