Fix enum consts not being renamed by AstMapper

Col-E · Jan 19, 2025 · e267dab · e267dab
1 parent ba1ae17
commit e267dab
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 3 deletions.
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/source/AstMapper.java b/recaf-core/src/main/java/software/coley/recaf/services/source/AstMapper.java
@@ -4,6 +4,7 @@
 import jakarta.annotation.Nullable;
 import software.coley.recaf.services.mapping.IntermediateMappings;
 import software.coley.recaf.services.mapping.Mappings;
+import software.coley.recaf.util.AccessFlag;
 import software.coley.recaf.util.StringUtil;
 import software.coley.sourcesolver.model.ClassModel;
 import software.coley.sourcesolver.model.CompilationUnitModel;
@@ -226,6 +227,22 @@ private Range extractRelevantRange(@Nonnull Model model, @Nonnull String source)
 			// The name doesn't have an associated model, but is after the declared type.
 			String name = variableModel.getName();
 			int begin = variableModel.getType().getRange().end();
+
+			// Enum constants don't have an AST model for their type, so the range is "unknown".
+			// If we're confident this is an enum constant, then we'll make the beginning range the start of the field name instead.
+			if (begin == -1) {
+				ClassModel declaringClass = variableModel.getParentOfType(ClassModel.class);
+				if (declaringClass != null) {
+					if (declaringClass.resolve(resolver) instanceof ClassResolution declaringResolution
+							&& AccessFlag.isEnum(declaringResolution.getClassEntry().getAccess())) {
+						Resolution type = variableModel.getType().resolve(resolver);
+						if (type.matches(declaringResolution)) {
+							begin = variableModel.getRange().begin();
+						}
+					}
+				}
+			}
+
 			int end = source.indexOf(name, begin) + name.length();
 			if (end < begin)
 				return Range.UNKNOWN;

diff --git a/recaf-core/src/test/java/software/coley/recaf/services/source/AstServiceTest.java b/recaf-core/src/test/java/software/coley/recaf/services/source/AstServiceTest.java
@@ -17,8 +17,21 @@
 import software.coley.recaf.services.mapping.Mappings;
 import software.coley.recaf.test.TestBase;
 import software.coley.recaf.test.TestClassUtils;
-import software.coley.recaf.test.dummy.*;
-import software.coley.recaf.util.JavaVersion;
+import software.coley.recaf.test.dummy.AccessibleFields;
+import software.coley.recaf.test.dummy.AnnotationImpl;
+import software.coley.recaf.test.dummy.ClassWithConstructor;
+import software.coley.recaf.test.dummy.ClassWithExceptions;
+import software.coley.recaf.test.dummy.ClassWithMethodReference;
+import software.coley.recaf.test.dummy.ClassWithMultipleMethods;
+import software.coley.recaf.test.dummy.ClassWithStaticInit;
+import software.coley.recaf.test.dummy.DummyEnum;
+import software.coley.recaf.test.dummy.HelloWorld;
+import software.coley.recaf.test.dummy.OverlapClassAB;
+import software.coley.recaf.test.dummy.OverlapInterfaceA;
+import software.coley.recaf.test.dummy.OverlapInterfaceB;
+import software.coley.recaf.test.dummy.StringList;
+import software.coley.recaf.test.dummy.StringListUser;
+import software.coley.recaf.test.dummy.StringSupplier;
 import software.coley.recaf.util.StringUtil;
 import software.coley.recaf.util.Types;
 import software.coley.recaf.workspace.model.Workspace;
@@ -526,6 +539,25 @@ enum DummyEnum {
 			});
 		}
 
+		@Test
+		void renameField_EnumConst() {
+			String source = """
+					package software.coley.recaf.test.dummy;
+										
+					enum DummyEnum {
+						ONE, TWO, THREE
+					}
+					""";
+			handleUnit(source, unit -> {
+				IntermediateMappings mappings = new IntermediateMappings();
+				String owner = DummyEnum.class.getName().replace('.', '/');
+				mappings.addField(owner, "L" + owner + ";", "ONE", "FIRST");
+
+				String modified = applyMappings(unit, mappings);
+				assertTrue(modified.contains("FIRST,"));
+			});
+		}
+
 		@Test
 		void renameClass_ReplacesCast() {
 			String source = """
@@ -572,7 +604,7 @@ static void main(String[] args) {
 		void renameClass_ReplacePackageImport() {
 			String source = """
 					package software.coley.recaf.test.dummy;
-					
+										
 					import software.coley.recaf.util.*;
 										
 					class HelloWorld {