From 6ec7f789ef70fedcd84791ab669a5825ca3559f9 Mon Sep 17 00:00:00 2001
From: Skylot <skylot@gmail.com>
Date: Mon, 22 Nov 2021 13:56:15 +0000
Subject: [PATCH] fix: restore usage data after class reload (#1281)

---
 .../java/jadx/core/dex/nodes/ClassNode.java   | 22 +++++++++++++++++++
 .../src/main/java/jadx/core/utils/Utils.java  | 15 +++++++++++++
 2 files changed, 37 insertions(+)

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 aac5e64a..b68d5a76 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
@@ -108,6 +108,10 @@ public class ClassNode extends NotificationAttrNode implements ILoadable, ICodeN
 			ListConsumer<IFieldData, FieldNode> fieldsConsumer = new ListConsumer<>(fld -> FieldNode.build(this, fld));
 			ListConsumer<IMethodData, MethodNode> methodsConsumer = new ListConsumer<>(mth -> MethodNode.build(this, mth));
 			cls.visitFieldsAndMethods(fieldsConsumer, methodsConsumer);
+			if (this.fields != null && this.methods != null) {
+				// TODO: temporary solution for restore usage info in reloaded methods and fields
+				restoreUsageData(this.fields, this.methods, fieldsConsumer.getResult(), methodsConsumer.getResult());
+			}
 			this.fields = fieldsConsumer.getResult();
 			this.methods = methodsConsumer.getResult();
 
@@ -124,6 +128,24 @@ public class ClassNode extends NotificationAttrNode implements ILoadable, ICodeN
 		}
 	}
 
+	private void restoreUsageData(List<FieldNode> oldFields, List<MethodNode> oldMethods,
+			List<FieldNode> newFields, List<MethodNode> newMethods) {
+		Map<FieldInfo, FieldNode> oldFieldMap = Utils.groupBy(oldFields, FieldNode::getFieldInfo);
+		for (FieldNode newField : newFields) {
+			FieldNode oldField = oldFieldMap.get(newField.getFieldInfo());
+			if (oldField != null) {
+				newField.setUseIn(oldField.getUseIn());
+			}
+		}
+		Map<MethodInfo, MethodNode> oldMethodsMap = Utils.groupBy(oldMethods, MethodNode::getMethodInfo);
+		for (MethodNode newMethod : newMethods) {
+			MethodNode oldMethod = oldMethodsMap.get(newMethod.getMethodInfo());
+			if (oldMethod != null) {
+				newMethod.setUseIn(oldMethod.getUseIn());
+			}
+		}
+	}
+
 	private ArgType checkSuperType(IClassData cls) {
 		String superType = cls.getSuperType();
 		if (superType == null) {
diff --git a/jadx-core/src/main/java/jadx/core/utils/Utils.java b/jadx-core/src/main/java/jadx/core/utils/Utils.java
index 3aa5141f..38576a90 100644
--- a/jadx-core/src/main/java/jadx/core/utils/Utils.java
+++ b/jadx-core/src/main/java/jadx/core/utils/Utils.java
@@ -315,6 +315,21 @@ public class Utils {
 		return result;
 	}
 
+	/**
+	 * Build map from list of values with value to key mapping function
+	 * <br>
+	 * Similar to:
+	 * <br>
+	 * {@code list.stream().collect(Collectors.toMap(mapKey, Function.identity())); }
+	 */
+	public static <K, V> Map<K, V> groupBy(List<V> list, Function<V, K> mapKey) {
+		Map<K, V> map = new HashMap<>(list.size());
+		for (V v : list) {
+			map.put(mapKey.apply(v), v);
+		}
+		return map;
+	}
+
 	@Nullable
 	public static <T> T getOne(@Nullable List<T> list) {
 		if (list == null || list.size() != 1) {
-- 
GitLab