refactor: better handling of SourceFragment in CompilationUnit (#2806)

src/main/java/spoon/reflect/declaration/CtCompilationUnit.java

@@ -29,6 +29,7 @@
 import spoon.support.DerivedProperty;
 import spoon.support.Experimental;
 import spoon.support.UnsettableProperty;
+import spoon.support.util.ModelList;
 
 /**
  * Defines a compilation unit. In Java, a compilation unit can contain only one
@@ -196,7 +197,7 @@ enum UNIT_TYPE {
 	 */
 	@Experimental
 	@PropertyGetter(role = CtRole.DECLARED_IMPORT)
-	List<CtImport> getImports();
+	ModelList<CtImport> getImports();
 
 	@Override
 	CtCompilationUnit clone();

src/main/java/spoon/support/reflect/declaration/CtCompilationUnitImpl.java

@@ -40,6 +40,7 @@
 import spoon.reflect.declaration.ParentNotInitializedException;
 import spoon.reflect.path.CtRole;
 import spoon.reflect.reference.CtModuleReference;
+import spoon.reflect.reference.CtPackageReference;
 import spoon.reflect.reference.CtTypeReference;
 import spoon.reflect.visitor.CtVisitor;
 import spoon.reflect.visitor.DefaultJavaPrettyPrinter;
@@ -131,7 +132,7 @@ public UNIT_TYPE getUnitType() {
 			if (getDeclaredTypes().isEmpty()) {
 				if (getDeclaredModuleReference() != null) {
 					return UNIT_TYPE.MODULE_DECLARATION;
-				} else if (getPackageDeclaration() != null) {
+				} else if (packageDeclaration != null) {
 					return UNIT_TYPE.PACKAGE_DECLARATION;
 				} else {
 					return UNIT_TYPE.UNKNOWN;
@@ -245,6 +246,15 @@ public CtPackage getDeclaredPackage() {
 
 	@Override
 	public CtPackageDeclaration getPackageDeclaration() {
+		if (packageDeclaration == null) {
+			CtPackageReference packRef;
+			if (declaredTypeReferences.size() > 0) {
+				packRef = declaredTypeReferences.get(0).getPackage().clone();
+			} else {
+				packRef = getFactory().getModel().getRootPackage().getReference();
+			}
+			packageDeclaration = getFactory().Package().createPackageDeclaration(packRef);
+		}
 		return packageDeclaration;
 	}
 
@@ -267,6 +277,8 @@ public CtCompilationUnitImpl setPackageDeclaration(CtPackageDeclaration packageD
 	@Override
 	public CtCompilationUnitImpl setFile(File file) {
 		this.file = file;
+		//reset cached position (if any)
+		this.position = SourcePosition.NOPOSITION;
 		return this;
 	}
 
@@ -334,7 +346,7 @@ public String getOriginalSourceCode() {
 
 
 	@Override
-	public List<CtImport> getImports() {
+	public ModelList<CtImport> getImports() {
 		return this.imports;
 	}
 
@@ -354,12 +366,7 @@ public ElementSourceFragment getOriginalSourceFragment() {
 				throw new SpoonException("Root source fragment of compilation unit of package is not supported");
 			}
 			rootFragment = new ElementSourceFragment(this, null);
-			for (CtImport imprt : getImports()) {
-				rootFragment.addChild(new ElementSourceFragment(imprt, null /*TODO role for import of CU*/));
-			}
-			for (CtTypeReference<?> ctType : declaredTypeReferences) {
-				rootFragment.addTreeOfSourceFragmentsOfElement(ctType.getTypeDeclaration());
-			}
+			rootFragment.addTreeOfSourceFragmentsOfElement(this);
 		}
 		return rootFragment;
 	}
@@ -381,6 +388,9 @@ public SourcePosition getPosition() {
 			String sourceCode = getOriginalSourceCode();
 			if (sourceCode != null) {
 				position = getFactory().Core().createSourcePosition((CompilationUnit) this, 0, sourceCode.length() - 1, getLineSeparatorPositions());
+			} else {
+				//it is a virtual compilation unit (e.g. for Snippet)
+				position = getFactory().Core().createSourcePosition((CompilationUnit) this, 0, Integer.MAX_VALUE - 1, getLineSeparatorPositions());
 			}
 		}
 		return position;

src/main/java/spoon/support/sniper/SniperJavaPrettyPrinter.java

@@ -208,7 +208,7 @@ public void printSourceFragment(SourceFragment fragment, Boolean isModified) {
 	@Override
 	public SniperJavaPrettyPrinter scan(CtElement element) {
 		if (element != null) {
-			CtRole role = element.getRoleInParent();
+			CtRole role = getRoleInCompilationUnit(element);
 			onPrintEvent(new ElementPrinterEvent(role, element) {
 				@Override
 				public void print(Boolean muted) {
@@ -223,6 +223,14 @@ public void printSourceFragment(SourceFragment fragment, Boolean isModified) {
 		return this;
 	}
 
+	private CtRole getRoleInCompilationUnit(CtElement element) {
+		CtRole role = element.getRoleInParent();
+		if (role == CtRole.CONTAINED_TYPE) {
+			role = CtRole.DECLARED_TYPE;
+		}
+		return role;
+	}
+
 	/**
 	 * Called whenever {@link DefaultJavaPrettyPrinter} scans/prints an element or writes a token
 	 */

src/main/java/spoon/support/sniper/internal/ElementSourceFragment.java

@@ -22,13 +22,14 @@
 import spoon.reflect.cu.CompilationUnit;
 import spoon.reflect.cu.SourcePosition;
 import spoon.reflect.cu.SourcePositionHolder;
+import spoon.reflect.declaration.CtCompilationUnit;
 import spoon.reflect.declaration.CtElement;
 import spoon.reflect.declaration.CtModifiable;
 import spoon.reflect.meta.ContainerKind;
 import spoon.reflect.meta.RoleHandler;
 import spoon.reflect.meta.impl.RoleHandlerHelper;
 import spoon.reflect.path.CtRole;
-import spoon.reflect.visitor.CtScanner;
+import spoon.reflect.visitor.EarlyTerminatingScanner;
 import spoon.support.Experimental;
 import spoon.support.reflect.CtExtendedModifier;
 import spoon.support.reflect.cu.position.SourcePositionImpl;
@@ -142,15 +143,12 @@ public void addTreeOfSourceFragmentsOfElement(CtElement element) {
 		Deque<ElementSourceFragment> parents = new ArrayDeque<>();
 		parents.push(this);
 		//scan all children of `element` and build tree of SourceFragments
-		new CtScanner() {
-			CtRole scannedRole;
-			@Override
-			public void scan(CtRole role, CtElement element) {
-				scannedRole = role;
-				super.scan(role, element);
-			}
+		new EarlyTerminatingScanner<Void>() {
 			@Override
 			protected void enter(CtElement e) {
+				if (e instanceof CtCompilationUnit) {
+					return;
+				}
 				ElementSourceFragment newFragment = addChild(parents.peek(), scannedRole, e);
 				if (newFragment != null) {
 					parents.push(newFragment);
@@ -165,12 +163,17 @@ protected void enter(CtElement e) {
 			}
 			@Override
 			protected void exit(CtElement e) {
+				if (e instanceof CtCompilationUnit) {
+					return;
+				}
 				ElementSourceFragment topFragment = parents.peek();
 				if (topFragment != null && topFragment.getElement() == e) {
 					parents.pop();
 				}
 			}
-		}.scan(element.getRoleInParent(), element);
+		}
+		.setVisitCompilationUnitContent(true)
+		.scan(element.getRoleInParent(), element);
 	}
 	/**
 	 * @param parentFragment the parent {@link ElementSourceFragment}, which will receive {@link ElementSourceFragment} made for `otherElement`
@@ -235,9 +238,6 @@ private ElementSourceFragment addChild(ElementSourceFragment parentFragment, CtR
 
 	private RoleHandler getRoleHandler(CtRole roleInParent, SourcePositionHolder otherElement) {
 		SourcePositionHolder parent = element;
-		if (parent instanceof CompilationUnit) {
-			parent = null;
-		}
 		if (parent == null) {
 			if (otherElement instanceof CtElement) {
 				parent = ((CtElement) otherElement).getParent();
@@ -523,13 +523,15 @@ public List<SourceFragment> getGroupedChildrenFragments() {
 					i++;
 				}
 				//add suffix space
+				/*
 				if (i < flatChildren.size()) {
 					SourceFragment nextChild = flatChildren.get(i);
 					if (isSpaceFragment(nextChild)) {
 						childrenInSameCollection.add(nextChild);
 						i++;
 					}
 				}
+				*/
 				result.add(new CollectionSourceFragment(childrenInSameCollection));
 			} else {
 				throw new SpoonException("Unexpected SourceFragment of type " + child.getClass());

src/main/java/spoon/support/visitor/replace/ReplaceListListener.java

@@ -16,8 +16,8 @@
  */
 package spoon.support.visitor.replace;
 
-import java.util.List;
+import java.util.Collection;
 
-public interface ReplaceListListener<T extends List> {
+public interface ReplaceListListener<T extends Collection> {
 	void set(T replace);
 }

src/main/java/spoon/support/visitor/replace/ReplacementVisitor.java

@@ -15,8 +15,6 @@
  * knowledge of the CeCILL-C license and that you accept its terms.
  */
 package spoon.support.visitor.replace;
-
-
 /**
  * Used to replace an element by another one.
  *
@@ -1270,15 +1268,15 @@ public void set(spoon.reflect.declaration.CtPackageDeclaration replace) {
 	}
 
 	// auto-generated, see spoon.generating.ReplacementVisitorGenerator
-	class CtCompilationUnitImportsReplaceListener implements spoon.support.visitor.replace.ReplaceListListener<java.util.List> {
+	class CtCompilationUnitImportsReplaceListener implements spoon.support.visitor.replace.ReplaceListListener<java.util.Collection> {
 		private final spoon.reflect.declaration.CtCompilationUnit element;
 
 		CtCompilationUnitImportsReplaceListener(spoon.reflect.declaration.CtCompilationUnit element) {
 			this.element = element;
 		}
 
 		@java.lang.Override
-		public void set(java.util.List replace) {
+		public void set(java.util.Collection replace) {
 			this.element.setImports(replace);
 		}
 	}
@@ -1341,7 +1339,7 @@ public void set(spoon.reflect.reference.CtTypeReference replace) {
 
 	public static void replace(spoon.reflect.declaration.CtElement original, spoon.reflect.declaration.CtElement replace) {
 		try {
-			new spoon.support.visitor.replace.ReplacementVisitor(original, (replace == null ? spoon.support.visitor.replace.ReplacementVisitor.EMPTY : new spoon.reflect.declaration.CtElement[]{ replace })).scan(original.getParent());
+			new spoon.support.visitor.replace.ReplacementVisitor(original, (replace == null ? EMPTY : new spoon.reflect.declaration.CtElement[]{ replace })).scan(original.getParent());
 		} catch (spoon.support.visitor.replace.InvalidReplaceException e) {
 			throw e;
 		}
@@ -1363,7 +1361,7 @@ public static <E extends spoon.reflect.declaration.CtElement> void replace(spoon
 
 	private ReplacementVisitor(spoon.reflect.declaration.CtElement original, spoon.reflect.declaration.CtElement... replace) {
 		this.original = original;
-		this.replace = (replace == null) ? spoon.support.visitor.replace.ReplacementVisitor.EMPTY : replace;
+		this.replace = (replace == null) ? EMPTY : replace;
 	}
 
 	private <K, V extends spoon.reflect.declaration.CtElement> void replaceInMapIfExist(java.util.Map<K, V> mapProtected, spoon.support.visitor.replace.ReplaceMapListener listener) {
@@ -2220,4 +2218,3 @@ public void visitCtTypeMemberWildcardImportReference(spoon.reflect.reference.CtT
 		replaceElementIfExist(wildcardReference.getTypeReference(), new spoon.support.visitor.replace.ReplacementVisitor.CtTypeMemberWildcardImportReferenceTypeReferenceReplaceListener(wildcardReference));
 	}
 }
-

src/test/java/spoon/generating/replace/ReplaceScanner.java

@@ -41,8 +41,11 @@
 import spoon.reflect.reference.CtTypeReference;
 import spoon.reflect.visitor.CtScanner;
 import spoon.reflect.visitor.filter.TypeFilter;
+import spoon.support.util.ModelList;
+import spoon.support.util.ModelSet;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -104,10 +107,12 @@ public <T> void visitCtMethod(CtMethod<T> element) {
 		target.addMethod(clone);
 	}
 
+	private static Set<String> modelCollectionTypes = new HashSet<>(Arrays.asList(ModelList.class.getName(), ModelSet.class.getName()));
+
 	private <T> CtInvocation<?> createInvocation(Factory factory, CtMethod<T> candidate, List<CtExpression<?>> invArgs, CtInvocation getter, Class getterTypeClass) {
 		CtInvocation<?> invocation;
 		Type type;
-		if (getterTypeClass.equals(Collection.class) || getterTypeClass.equals(List.class)) {
+		if (getterTypeClass.equals(Collection.class) || List.class.isAssignableFrom(getterTypeClass)) {
 			invocation = factory.Code().createInvocation(null, this.list, invArgs);
 			type = Type.LIST;
 		} else if (getterTypeClass.equals(Map.class)) {
@@ -129,11 +134,15 @@ private <T> CtInvocation<?> createInvocation(Factory factory, CtMethod<T> candid
 			listener = listeners.get(listenerName);
 		} else {
 			final CtTypeReference getterType = getGetterType(factory, getter);
-			listener = createListenerClass(factory, listenerName, getterType, type);
+			CtTypeReference<?> setterParamType = getterType;
+			if (modelCollectionTypes.contains(setterParamType.getQualifiedName())) {
+				setterParamType = factory.Type().createReference(Collection.class);
+			}
+			listener = createListenerClass(factory, listenerName, setterParamType, type);
 			final CtMethod setter = getSetter(name, getter.getTarget().getType().getDeclaration());
 			final CtField field = updateField(listener, setter.getDeclaringType().getReference());
 			updateConstructor(listener, setter.getDeclaringType().getReference());
-			updateSetter(factory, (CtMethod<?>) listener.getMethodsByName("set").get(0), getterType, field, setter);
+			updateSetter(factory, (CtMethod<?>) listener.getMethodsByName("set").get(0), setterParamType, field, setter);
 			// Add auto-generated comment.
 			final CtComment comment = factory.Core().createComment();
 			comment.setCommentType(CtComment.CommentType.INLINE);

src/test/java/spoon/test/position/TestSourceFragment.java

@@ -184,7 +184,7 @@ public void testExactSourceFragments() throws Exception {
 				"/**\n" + 
 				"	 * c0\n" + 
 				"	 */", 
-				group("\n\t", "public", "\n\t", "@Deprecated", " ", "//c1 ends with tab and space\t ", "\n\t", "static", " "), "/*c2*/", " ",
+				group("\n\t", "public", "\n\t", "@Deprecated", " ", "//c1 ends with tab and space\t ", "\n\t", "static"), " ", "/*c2*/", " ",
 				"<", group("T", ",", " ", "U"), ">",
 				" ", "T", " ", "m3", "(", group("U param", ",", " ", "@Deprecated int p2"), ")", " ", "{\n" + 
 						"		return null;\n" +