Make nicer method replacement code

3 files changed, 50 insertions, 16 deletions

diff --git a/javaparser-core-generators/src/main/java/com/github/javaparser/generator/Generator.java b/javaparser-core-generators/src/main/java/com/github/javaparser/generator/Generator.java
index 9db76f299..bb1c70f23 100644
--- a/javaparser-core-generators/src/main/java/com/github/javaparser/generator/Generator.java
+++ b/javaparser-core-generators/src/main/java/com/github/javaparser/generator/Generator.java
@@ -1,8 +1,15 @@
 package com.github.javaparser.generator;
 
-import com.github.javaparser.JavaParser;
+import com.github.javaparser.ast.Node;
+import com.github.javaparser.ast.body.CallableDeclaration;
+import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
+import com.github.javaparser.ast.nodeTypes.NodeWithAnnotations;
 import com.github.javaparser.utils.SourceRoot;
 
+import javax.annotation.Generated;
+
+import static com.github.javaparser.utils.CodeGenerationUtils.f;
+
 /**
  * A general pattern that the generators in this module will follow.
  */
@@ -14,4 +21,9 @@ public abstract class Generator {
     }
 
     public abstract void generate() throws Exception;
+
+    protected <T extends Node & NodeWithAnnotations> void markGenerated(T node) {
+        node.addSingleMemberAnnotation(Generated.class, f("\"%s\"", getClass().getName()));
+        node.tryAddImportToParentCompilationUnit(Generated.class);
+    }
 }
diff --git a/javaparser-core-generators/src/main/java/com/github/javaparser/generator/NodeGenerator.java b/javaparser-core-generators/src/main/java/com/github/javaparser/generator/NodeGenerator.java
index 1c3fb9d20..eb4f38cef 100644
--- a/javaparser-core-generators/src/main/java/com/github/javaparser/generator/NodeGenerator.java
+++ b/javaparser-core-generators/src/main/java/com/github/javaparser/generator/NodeGenerator.java
@@ -1,7 +1,7 @@
 package com.github.javaparser.generator;
 
-import com.github.javaparser.JavaParser;
 import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.body.CallableDeclaration;
 import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
 import com.github.javaparser.metamodel.BaseNodeMetaModel;
 import com.github.javaparser.metamodel.JavaParserMetaModel;
@@ -9,6 +9,7 @@ import com.github.javaparser.utils.Log;
 import com.github.javaparser.utils.SourceRoot;
 
 import java.io.IOException;
+import java.util.List;
 
 import static com.github.javaparser.utils.CodeGenerationUtils.f;
 
@@ -36,4 +37,35 @@ public abstract class NodeGenerator extends Generator {
     }
 
     protected abstract void generateNode(BaseNodeMetaModel nodeMetaModel, CompilationUnit nodeCu, ClassOrInterfaceDeclaration nodeCoid);
+
+    /**
+     * Utility method that looks for a method or constructor with an identical signature as "callable" and replaces it
+     * with callable. If not found, adds callable.
+     */
+    protected void addOrReplaceWhenSameSignature(ClassOrInterfaceDeclaration nodeCoid, CallableDeclaration callable) {
+        List<CallableDeclaration> existingCallables = nodeCoid.getCallablesWithSignature(callable.getSignature());
+        if (existingCallables.isEmpty()) {
+            nodeCoid.addMember(callable);
+            return;
+        }
+        if (existingCallables.size() > 1) {
+            throw new AssertionError(f("Wanted to regenerate a method with signature %s in %s, but found more than one.", callable.getSignature(), nodeCoid.getNameAsString()));
+        }
+        nodeCoid.getMembers().replace(existingCallables.get(0), callable);
+    }
+
+    /**
+     * Utility method that looks for a method or constructor with an identical signature as "callable" and replaces it
+     * with callable. If not found, fails.
+     */
+    protected void replaceWhenSameSignature(ClassOrInterfaceDeclaration nodeCoid, CallableDeclaration callable) {
+        List<CallableDeclaration> existingCallables = nodeCoid.getCallablesWithSignature(callable.getSignature());
+        if (existingCallables.isEmpty()) {
+            throw new AssertionError(f("Wanted to regenerate a method with signature %s in %s, but it wasn't there.", callable.getSignature(), nodeCoid.getNameAsString()));
+        }
+        if (existingCallables.size() > 1) {
+            throw new AssertionError(f("Wanted to regenerate a method with signature %s in %s, but found more than one.", callable.getSignature(), nodeCoid.getNameAsString()));
+        }
+        nodeCoid.getMembers().replace(existingCallables.get(0), callable);
+    }
 }
diff --git a/javaparser-core-generators/src/main/java/com/github/javaparser/generator/core/node/MainConstructorGenerator.java b/javaparser-core-generators/src/main/java/com/github/javaparser/generator/core/node/MainConstructorGenerator.java
index 24ef05494..fc3a39adc 100644
--- a/javaparser-core-generators/src/main/java/com/github/javaparser/generator/core/node/MainConstructorGenerator.java
+++ b/javaparser-core-generators/src/main/java/com/github/javaparser/generator/core/node/MainConstructorGenerator.java
@@ -6,7 +6,6 @@ import com.github.javaparser.ast.Node;
 import com.github.javaparser.ast.body.CallableDeclaration;
 import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
 import com.github.javaparser.ast.body.ConstructorDeclaration;
-import com.github.javaparser.ast.expr.SingleMemberAnnotationExpr;
 import com.github.javaparser.ast.stmt.BlockStmt;
 import com.github.javaparser.generator.NodeGenerator;
 import com.github.javaparser.metamodel.BaseNodeMetaModel;
@@ -14,9 +13,6 @@ import com.github.javaparser.metamodel.PropertyMetaModel;
 import com.github.javaparser.utils.SeparatedItemStringBuilder;
 import com.github.javaparser.utils.SourceRoot;
 
-import javax.annotation.Generated;
-import java.util.Optional;
-
 import static com.github.javaparser.JavaParser.parseExplicitConstructorInvocationStmt;
 import static com.github.javaparser.JavaParser.parseStatement;
 import static com.github.javaparser.utils.CodeGenerationUtils.f;
@@ -34,10 +30,9 @@ public class MainConstructorGenerator extends NodeGenerator {
         ConstructorDeclaration constructor = new ConstructorDeclaration()
                 .setPublic(true)
                 .setName(nodeCoid.getNameAsString())
-                .addParameter(Range.class, "range")
-                .addSingleMemberAnnotation(Generated.class, getClass().getName());
-
-
+                .addParameter(Range.class, "range");
+        markGenerated(constructor);
+        
         BlockStmt body = constructor.getBody();
 
         SeparatedItemStringBuilder superCall = new SeparatedItemStringBuilder("super(", ", ", ");");
@@ -53,12 +48,7 @@ public class MainConstructorGenerator extends NodeGenerator {
 
         body.getStatements().add(0, parseExplicitConstructorInvocationStmt(superCall.toString()));
 
-        nodeCoid.getMembers().stream()
-                .filter(m -> m instanceof CallableDeclaration)
-                .map(m -> ((CallableDeclaration) m))
-                .filter(m -> m.getSignature().equals(constructor.getSignature()))
-                .findFirst()
-                .ifPresent(m -> nodeCoid.getMembers().replace(m, constructor));
+        addOrReplaceWhenSameSignature(nodeCoid, constructor);
         nodeCu.addImport(Range.class);
     }
 }