fix: move class unload to synchronized block (#977)

Signed-off-by: NSkylot <skylot@gmail.com>

jadx-core/src/main/java/jadx/core/ProcessClass.java

 package jadx.core;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import jadx.api.ICodeInfo;
 import jadx.core.codegen.CodeGen;
+import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.nodes.ClassNode;
+import jadx.core.dex.nodes.LoadStage;
 import jadx.core.dex.visitors.DepthTraversal;
 import jadx.core.dex.visitors.IDexTreeVisitor;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
-import static jadx.core.dex.nodes.ProcessState.GENERATED;
+import static jadx.core.dex.nodes.ProcessState.GENERATED_AND_UNLOADED;
 import static jadx.core.dex.nodes.ProcessState.LOADED;
 import static jadx.core.dex.nodes.ProcessState.NOT_LOADED;
 import static jadx.core.dex.nodes.ProcessState.PROCESS_COMPLETE;
@@ -20,18 +23,27 @@ public final class ProcessClass {
 	private ProcessClass() {
 	}
 
-	public static void process(ClassNode cls) {
-		ClassNode topParentClass = cls.getTopParentClass();
-		if (topParentClass != cls) {
-			process(topParentClass);
-			return;
-		}
-		if (cls.getState().isProcessed()) {
+	@Nullable
+	private static ICodeInfo process(ClassNode cls, boolean codegen) {
+		if (!codegen && cls.getState() == PROCESS_COMPLETE) {
 			// nothing to do
-			return;
+			return null;
 		}
 		synchronized (cls.getClassInfo()) {
 			try {
+				if (codegen) {
+					if (cls.getState() == GENERATED_AND_UNLOADED) {
+						// allow to run code generation again
+						cls.setState(NOT_LOADED);
+					}
+					cls.setLoadStage(LoadStage.CODEGEN_STAGE);
+					if (cls.contains(AFlag.RELOAD_AT_CODEGEN_STAGE)) {
+						cls.remove(AFlag.RELOAD_AT_CODEGEN_STAGE);
+						cls.unload();
+					}
+				} else {
+					cls.setLoadStage(LoadStage.PROCESS_STAGE);
+				}
 				if (cls.getState() == NOT_LOADED) {
 					cls.load();
 				}
@@ -42,9 +54,18 @@ public final class ProcessClass {
 					}
 					cls.setState(PROCESS_COMPLETE);
 				}
+				if (codegen) {
+					ICodeInfo code = CodeGen.generate(cls);
+					if (!cls.contains(AFlag.DONT_UNLOAD_CLASS)) {
+						cls.unload();
+						cls.setState(GENERATED_AND_UNLOADED);
+					}
+					return code;
+				}
 			} catch (Throwable e) {
 				cls.addError("Class process error: " + e.getClass().getSimpleName(), e);
 			}
+			return null;
 		}
 	}
 
@@ -54,21 +75,14 @@ public final class ProcessClass {
 		if (topParentClass != cls) {
 			return generateCode(topParentClass);
 		}
-		if (cls.getState() == GENERATED) {
-			// allow to run code generation again
-			cls.setState(NOT_LOADED);
-		}
 		try {
 			for (ClassNode depCls : cls.getDependencies()) {
-				depCls.startProcessStage();
-				process(depCls);
+				process(depCls, false);
+			}
+			ICodeInfo code = process(cls, true);
+			if (code == null) {
+				throw new JadxRuntimeException("Codegen failed");
 			}
-			cls.startCodegenStage();
-			process(cls);
-
-			ICodeInfo code = CodeGen.generate(cls);
-			cls.setState(GENERATED);
-			cls.unload();
 			return code;
 		} catch (Throwable e) {
 			throw new JadxRuntimeException("Failed to generate code for class: " + cls.getFullName(), e);

jadx-core/src/main/java/jadx/core/dex/attributes/AFlag.java

@@ -75,4 +75,6 @@ public enum AFlag {
 	INCONSISTENT_CODE, // warning about incorrect decompilation
 
 	REQUEST_IF_REGION_OPTIMIZE, // run if region visitor again
+
+	DONT_UNLOAD_CLASS, // don't unload class after code generation (only for tests and debug!)
 }

jadx-core/src/main/java/jadx/core/dex/nodes/ClassNode.java

@@ -40,6 +40,7 @@ import jadx.core.utils.exceptions.JadxRuntimeException;
 
 import static jadx.core.dex.nodes.ProcessState.LOADED;
 import static jadx.core.dex.nodes.ProcessState.NOT_LOADED;
+import static jadx.core.dex.nodes.ProcessState.PROCESS_COMPLETE;
 
 public class ClassNode extends NotificationAttrNode implements ILoadable, ICodeNode, Comparable<ClassNode> {
 	private static final Logger LOG = LoggerFactory.getLogger(ClassNode.class);
@@ -209,10 +210,10 @@ public class ClassNode extends NotificationAttrNode implements ILoadable, ICodeN
 
 	public void ensureProcessed() {
 		ClassNode topClass = getTopParentClass();
-		ProcessState topState = topClass.getState();
-		if (!topState.isProcessed()) {
+		ProcessState state = topClass.getState();
+		if (state != PROCESS_COMPLETE) {
 			throw new JadxRuntimeException("Expected class to be processed at this point,"
-					+ " class: " + topClass + ", state: " + topState);
+					+ " class: " + topClass + ", state: " + state);
 		}
 	}
 
@@ -583,16 +584,8 @@ public class ClassNode extends NotificationAttrNode implements ILoadable, ICodeN
 		return loadStage;
 	}
 
-	public void startProcessStage() {
-		this.loadStage = LoadStage.PROCESS_STAGE;
-	}
-
-	public void startCodegenStage() {
-		this.loadStage = LoadStage.CODEGEN_STAGE;
-		if (contains(AFlag.RELOAD_AT_CODEGEN_STAGE)) {
-			unload();
-			remove(AFlag.RELOAD_AT_CODEGEN_STAGE);
-		}
+	public void setLoadStage(LoadStage loadStage) {
+		this.loadStage = loadStage;
 	}
 
 	public void reloadAtCodegenStage() {

jadx-core/src/main/java/jadx/core/dex/nodes/ProcessState.java

@@ -5,9 +5,5 @@ public enum ProcessState {
 	LOADED,
 	PROCESS_STARTED,
 	PROCESS_COMPLETE,
-	GENERATED;
-
-	public boolean isProcessed() {
-		return this == PROCESS_COMPLETE || this == GENERATED;
-	}
+	GENERATED_AND_UNLOADED;
 }

jadx-core/src/main/java/jadx/core/dex/visitors/PrepareForCodeGen.java

@@ -58,11 +58,10 @@ public class PrepareForCodeGen extends AbstractVisitor {
 
 	@Override
 	public void visit(MethodNode mth) throws JadxException {
-		List<BlockNode> blocks = mth.getBasicBlocks();
-		if (blocks == null) {
+		if (mth.isNoCode()) {
 			return;
 		}
-		for (BlockNode block : blocks) {
+		for (BlockNode block : mth.getBasicBlocks()) {
 			if (block.contains(AFlag.DONT_GENERATE)) {
 				continue;
 			}

jadx-core/src/test/java/jadx/tests/api/IntegrationTest.java

@@ -28,8 +28,6 @@ import jadx.api.ICodeInfo;
 import jadx.api.JadxArgs;
 import jadx.api.JadxDecompiler;
 import jadx.api.JadxInternalAccess;
-import jadx.core.ProcessClass;
-import jadx.core.codegen.CodeGen;
 import jadx.core.codegen.CodeWriter;
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
@@ -190,11 +188,10 @@ public abstract class IntegrationTest extends TestUtils {
 	}
 
 	protected void decompileAndCheck(JadxDecompiler d, List<ClassNode> clsList) {
-		if (unloadCls) {
-			clsList.forEach(ClassNode::decompile);
-		} else {
-			clsList.forEach(cls -> decompileWithoutUnload(d, cls));
+		if (!unloadCls) {
+			clsList.forEach(cls -> cls.add(AFlag.DONT_UNLOAD_CLASS));
 		}
+		clsList.forEach(ClassNode::decompile);
 
 		for (ClassNode cls : clsList) {
 			System.out.println("-----------------------------------------------------------");
@@ -251,22 +248,6 @@ public abstract class IntegrationTest extends TestUtils {
 		root.processResources(resStorage);
 	}
 
-	protected void decompileWithoutUnload(JadxDecompiler jadx, ClassNode cls) {
-		ProcessClass.process(cls);
-		generateClsCode(cls);
-		// don't unload class
-	}
-
-	protected void generateClsCode(ClassNode cls) {
-		try {
-			ICodeInfo code = CodeGen.generate(cls);
-			cls.root().getCodeCache().add(cls.getTopParentClass().getRawName(), code);
-		} catch (Exception e) {
-			e.printStackTrace();
-			fail(e.getMessage());
-		}
-	}
-
 	protected void checkCode(ClassNode cls) {
 		assertFalse(hasErrors(cls), "Inconsistent cls: " + cls);
 		for (MethodNode mthNode : cls.getMethods()) {