From a9b3e77874b1f188b8d16999841c9bd91fbc4015 Mon Sep 17 00:00:00 2001
From: I-Al-Istannen <I-Al-Istannen@users.noreply.github.com>
Date: Mon, 28 Aug 2023 18:17:14 +0200
Subject: [PATCH] fix: Remove compilation unit by reference equality if it has
 no file (#5400)

---
 .../factory/CompilationUnitFactory.java       | 10 +++++--
 .../declaration/CtCompilationUnitImpl.java    |  5 +++-
 .../test/refactoring/RefactoringTest.java     | 28 +++++++++++++++++++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/main/java/spoon/reflect/factory/CompilationUnitFactory.java b/src/main/java/spoon/reflect/factory/CompilationUnitFactory.java
index ae9c4c75ffa..90c70bc6f86 100644
--- a/src/main/java/spoon/reflect/factory/CompilationUnitFactory.java
+++ b/src/main/java/spoon/reflect/factory/CompilationUnitFactory.java
@@ -92,8 +92,14 @@ public CompilationUnit getOrCreate(CtPackage ctPackage) {
 	/**
 	 * remove a type from the list of types to be pretty-printed
 	 */
-	public void removeType(CtType type) {
-		cachedCompilationUnits.remove(type.getPosition().getCompilationUnit().getFile().getAbsolutePath());
+	public void removeType(CtType<?> type) {
+		File file = type.getPosition().getCompilationUnit().getFile();
+		if (file != null) {
+			cachedCompilationUnits.remove(type.getPosition().getCompilationUnit().getFile().getAbsolutePath());
+		} else {
+			// Virtual files do not have a file, so fall back to a slow scan, trying to find them by reference equality
+			cachedCompilationUnits.values().removeIf(it -> it == type.getPosition().getCompilationUnit());
+		}
 	}
 
 	/**
diff --git a/src/main/java/spoon/support/reflect/declaration/CtCompilationUnitImpl.java b/src/main/java/spoon/support/reflect/declaration/CtCompilationUnitImpl.java
index e2658149447..cfeb00f813b 100644
--- a/src/main/java/spoon/support/reflect/declaration/CtCompilationUnitImpl.java
+++ b/src/main/java/spoon/support/reflect/declaration/CtCompilationUnitImpl.java
@@ -426,6 +426,9 @@ public <E extends CtElement> E setParent(CtElement parent) {
 
 	@Override
 	public String toString() {
-		return this.file.getName();
+		if (this.file != null) {
+			return this.file.getName();
+		}
+		return "CompilationUnit<unknown file>";
 	}
 }
diff --git a/src/test/java/spoon/test/refactoring/RefactoringTest.java b/src/test/java/spoon/test/refactoring/RefactoringTest.java
index b70a35d2a74..615d839d072 100644
--- a/src/test/java/spoon/test/refactoring/RefactoringTest.java
+++ b/src/test/java/spoon/test/refactoring/RefactoringTest.java
@@ -17,6 +17,7 @@
 package spoon.test.refactoring;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -35,6 +36,7 @@
 import org.junit.jupiter.api.condition.JRE;
 import spoon.Launcher;
 import spoon.refactoring.Refactoring;
+import spoon.reflect.CtModel;
 import spoon.reflect.code.BinaryOperatorKind;
 import spoon.reflect.code.CtBinaryOperator;
 import spoon.reflect.code.CtInvocation;
@@ -45,6 +47,7 @@
 import spoon.reflect.visitor.filter.AbstractFilter;
 import spoon.reflect.visitor.filter.AbstractReferenceFilter;
 import spoon.reflect.visitor.filter.TypeFilter;
+import spoon.support.compiler.VirtualFile;
 import spoon.test.refactoring.processors.ThisTransformationProcessor;
 import spoon.test.refactoring.testclasses.AClass;
 
@@ -184,4 +187,29 @@ public void testRemoveDeprecatedMethods() {
 			// error is kinda okay
 		}
 	}
+
+	@Test
+	void renameClassInVirtualFile() {
+		// contract: Renaming classes defined in virtual files should work
+		Launcher spoon = new Launcher();
+		spoon.addInputResource(new VirtualFile(
+			"public class Test {\n" +
+			"  void foo() {\n" +
+			"  }\n" +
+			"}\n"));
+		CtModel model = spoon.buildModel();
+		CtClass<?> clazz = model.getElements(new TypeFilter<>(CtClass.class)).get(0);
+
+		Refactoring.changeTypeName(clazz, "Test2");
+
+		assertTrue(
+			clazz.prettyprint().contains("class Test2 "),
+			"Class was not renamed: '" + clazz.prettyprint() + "'"
+		);
+		assertFalse(
+			clazz.prettyprint().contains("class Test "),
+			"Class was not renamed: '" + clazz.prettyprint() + "'"
+		);
+	}
+
 }