Import current class's package when patching.

Signed-off-by: Ross Allan <rallanpcl@gmail.com>
MinimallyCorrect · Jul 3, 2013 · 6e662a4 · 6e662a4
1 parent 722ec7b
commit 6e662a4
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 18 deletions.
diff --git a/src/common/javassist/RemappingPool.java b/src/common/javassist/RemappingPool.java
@@ -2,12 +2,14 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 
 import nallar.collections.PartiallySynchronizedMap;
 import nallar.tickthreading.Log;
 import nallar.tickthreading.patcher.remapping.ByteSource;
 import nallar.tickthreading.patcher.remapping.StringExtractor;
 import nallar.tickthreading.patcher.remapping.Transformer;
+import nallar.tickthreading.util.ReflectUtil;
 
 public class RemappingPool extends ClassPool {
 	private final PartiallySynchronizedMap<String, CtClass> srgClasses = new PartiallySynchronizedMap<String, CtClass>();
@@ -146,4 +148,24 @@ protected CtClass createCtClass(String className, boolean useCache) {
 	public void markChanged(final String className) {
 		srgClasses.remove(className);
 	}
+
+	public boolean addCurrentPackage(final String packageName) {
+		if (packageName == null) {
+			return false;
+		}
+		ArrayList<String> importedPackages = getImportedPackages_();
+		if (importedPackages.contains(packageName)) {
+			return false;
+		}
+		importedPackages.add(0, packageName);
+		return true;
+	}
+
+	public void removeCurrentPackage() {
+		getImportedPackages_().remove(0);
+	}
+
+	private ArrayList<String> getImportedPackages_() {
+		return ReflectUtil.get(this, "importedPackages");
+	}
 }
diff --git a/src/common/nallar/tickthreading/patcher/PatchManager.java b/src/common/nallar/tickthreading/patcher/PatchManager.java
@@ -24,11 +24,13 @@
 import java.util.TreeMap;
 
 import javassist.CannotCompileException;
+import javassist.ClassPool;
 import javassist.CtBehavior;
 import javassist.CtClass;
 import javassist.CtConstructor;
 import javassist.CtMethod;
 import javassist.NotFoundException;
+import javassist.RemappingPool;
 import nallar.tickthreading.Log;
 import nallar.tickthreading.mappings.ClassDescription;
 import nallar.tickthreading.mappings.FieldDescription;
@@ -231,26 +233,38 @@ public void runPatches(Mappings mappings) {
 					Log.info(className + " will not be patched, as it was not found.");
 					continue;
 				}
-				List<Element> patchElements = DomUtil.elementList(classElement.getChildNodes());
-				boolean patched = false;
-				for (Element patchElement : patchElements) {
-					PatchMethodDescriptor patch = patches.get(patchElement.getTagName());
-					if (patch == null) {
-						Log.severe("Patch " + patchElement.getTagName() + " was not found.");
-						continue;
-					}
-					try {
-						Object result = patch.run(patchElement, ctClass);
-						patched = true;
-						if (result instanceof CtClass) {
-							ctClass = (CtClass) result;
+				ClassPool classPool = classRegistry.classes;
+				boolean shouldRemove = false;
+				if (classPool instanceof RemappingPool) {
+					RemappingPool remappingPool = (RemappingPool) classPool;
+					shouldRemove = remappingPool.addCurrentPackage(ctClass.getPackageName());
+				}
+				try {
+					List<Element> patchElements = DomUtil.elementList(classElement.getChildNodes());
+					boolean patched = false;
+					for (Element patchElement : patchElements) {
+						PatchMethodDescriptor patch = patches.get(patchElement.getTagName());
+						if (patch == null) {
+							Log.severe("Patch " + patchElement.getTagName() + " was not found.");
+							continue;
+						}
+						try {
+							Object result = patch.run(patchElement, ctClass);
+							patched = true;
+							if (result instanceof CtClass) {
+								ctClass = (CtClass) result;
+							}
+						} catch (Exception e) {
+							Log.severe("Failed to patch " + ctClass.getName() + " with " + patch.name, e);
 						}
-					} catch (Exception e) {
-						Log.severe("Failed to patch " + ctClass.getName() + " with " + patch.name, e);
 					}
-				}
-				if (patched) {
-					patchingClasses.put(className, ctClass);
+					if (patched) {
+						patchingClasses.put(className, ctClass);
+					}
+				} finally {
+					if (shouldRemove) {
+						((RemappingPool) classPool).removeCurrentPackage();
+					}
 				}
 				if (!isSrg_) {
 					classRegistry.classes.markChanged(className);