Changed to keep generated classes in Context

Author: janario

tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/Context.java

@@ -10,6 +10,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import javax.annotation.processing.ProcessingEnvironment;
@@ -60,6 +61,9 @@ public final class Context {
 	private boolean addSuppressWarningsAnnotation;
 	private AccessType persistenceUnitDefaultAccessType;
 
+	// keep track of all classes for which model have been generated
+	private final Collection<String> generatedModelClasses = new HashSet<String>();
+
 	public Context(ProcessingEnvironment pe) {
 		this.pe = pe;
 
@@ -181,6 +185,14 @@ public TypeElement getTypeElementForFullyQualifiedName(String fqcn) {
 		return elementUtils.getTypeElement( fqcn );
 	}
 
+	void markGenerated(String name) {
+		generatedModelClasses.add( name );
+	}
+
+	boolean isAlreadyGenerated(String name) {
+		return generatedModelClasses.contains( name );
+	}
+
 	public void logMessage(Diagnostic.Kind type, String message) {
 		if ( !logDebug && type.equals( Diagnostic.Kind.OTHER ) ) {
 			return;

tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/JPAMetaModelEntityProcessor.java

@@ -137,13 +137,13 @@ public boolean process(final Set<? extends TypeElement> annotations, final Round
 	}
 
 	private void createMetaModelClasses() {
-		// keep track of all classes for which model have been generated
-		Collection<String> generatedModelClasses = new ArrayList<String>();
-
 		for ( MetaEntity entity : context.getMetaEntities() ) {
+			if ( context.isAlreadyGenerated( entity.getQualifiedName() ) ) {
+				continue;
+			}
 			context.logMessage( Diagnostic.Kind.OTHER, "Writing meta model for entity " + entity );
 			ClassWriter.writeFile( entity, context );
-			generatedModelClasses.add( entity.getQualifiedName() );
+			context.markGenerated( entity.getQualifiedName() );
 		}
 
 		// we cannot process the delayed entities in any order. There might be dependencies between them.
@@ -154,7 +154,7 @@ private void createMetaModelClasses() {
 			int toProcessCountBeforeLoop = toProcessEntities.size();
 			for ( MetaEntity entity : toProcessEntities ) {
 				// see METAGEN-36
-				if ( generatedModelClasses.contains( entity.getQualifiedName() ) ) {
+				if ( context.isAlreadyGenerated( entity.getQualifiedName() ) ) {
 					processedEntities.add( entity );
 					continue;
 				}
@@ -165,6 +165,7 @@ private void createMetaModelClasses() {
 						Diagnostic.Kind.OTHER, "Writing meta model for embeddable/mapped superclass" + entity
 				);
 				ClassWriter.writeFile( entity, context );
+				context.markGenerated( entity.getQualifiedName() );
 				processedEntities.add( entity );
 			}
 			toProcessEntities.removeAll( processedEntities );