Generator refactor

src/main/java/com/github/junkfactory/innerbuilder/JavaInnerBuilderHandler.java

@@ -6,6 +6,7 @@
 import com.github.junkfactory.innerbuilder.generators.PsiParams;
 import com.github.junkfactory.innerbuilder.generators.Utils;
 import com.github.junkfactory.innerbuilder.ui.JavaInnerBuilderOption;
+import com.github.junkfactory.innerbuilder.ui.JavaInnerBuilderOptionSelector;
 import com.intellij.ide.util.PropertiesComponent;
 import com.intellij.lang.LanguageCodeInsightActionHandler;
 import com.intellij.openapi.application.ApplicationManager;
@@ -16,14 +17,13 @@
 import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiJavaFile;
+import com.intellij.psi.codeStyle.JavaCodeStyleManager;
 import com.intellij.util.AstLoadingFilter;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.EnumSet;
 import java.util.Set;
 
-import static com.github.junkfactory.innerbuilder.ui.JavaInnerBuilderOptionSelector.selectFieldsAndOptions;
-
 class JavaInnerBuilderHandler implements LanguageCodeInsightActionHandler {
 
     private static final GeneratorFactory generatorFactory = GeneratorFactory.create();
@@ -48,7 +48,7 @@ public boolean startInWriteAction() {
     }
 
     private static boolean isApplicable(final PsiFile file, final Editor editor) {
-        return FieldCollector.builder()
+        return file instanceof PsiJavaFile && FieldCollector.builder()
                 .file(file)
                 .editor(editor)
                 .build()
@@ -77,7 +77,11 @@ public void invoke(@NotNull final Project project, @NotNull final Editor editor,
                 return;
             }
 
-            var selectedFields = selectFieldsAndOptions(existingFields, project);
+            var optionsDialog = JavaInnerBuilderOptionSelector.builder()
+                    .project(project)
+                    .members(existingFields)
+                    .build();
+            var selectedFields = optionsDialog.selectFieldsAndOptions();
             if (selectedFields.isEmpty()) {
                 return;
             }
@@ -86,6 +90,7 @@ public void invoke(@NotNull final Project project, @NotNull final Editor editor,
                     .file(file)
                     .selectedFields(selectedFields)
                     .factory(JavaPsiFacade.getElementFactory(project))
+                    .codeStyleManager(JavaCodeStyleManager.getInstance(project))
                     .build();
             var generatorParams = GeneratorParams.builder()
                     .project(project)

src/main/java/com/github/junkfactory/innerbuilder/generators/AbstractGenerator.java

@@ -2,12 +2,17 @@
 
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiJavaFile;
 import com.intellij.psi.PsiMethod;
+import com.intellij.psi.PsiType;
+import com.intellij.psi.util.PsiUtil;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-abstract class AbstractGenerator implements Runnable {
+import java.util.Objects;
+
+abstract class AbstractGenerator {
 
     @NonNls
     static final String BUILDER_CLASS_NAME = "Builder";
@@ -19,6 +24,10 @@ abstract class AbstractGenerator implements Runnable {
     static final String EMPTY = "";
     @NonNls
     static final String SPACE = " ";
+    @NonNls
+    static final String THIS_DOT = "this.";
+    @NonNls
+    static final String RETURN_THIS = "return this;";
 
     protected final GeneratorFactory generatorFactory;
     protected final GeneratorParams generatorParams;
@@ -49,6 +58,12 @@ protected PsiElement addMethod(@NotNull final PsiClass target, @Nullable final P
         return existingMethod;
     }
 
+    protected boolean addImport(PsiType psiType) {
+        var psiClass = Objects.requireNonNull(PsiUtil.resolveClassInType(psiType),
+                "Unable to resolve " + psiType.toString());
+        return generatorParams.psi().codeStyleManager().addImport((PsiJavaFile) generatorParams.psi().file(), psiClass);
+    }
+
     private PsiMethod findConstructor(PsiClass target, PsiMethod newMethod) {
         for (var constructor : target.getConstructors()) {
             if (Utils.areParameterListsEqual(constructor.getParameterList(), newMethod.getParameterList())) {

src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderClassGenerator.java

@@ -4,7 +4,7 @@
 import com.intellij.psi.PsiModifier;
 import com.intellij.psi.util.PsiUtil;
 
-class BuilderClassGenerator extends AbstractGenerator {
+class BuilderClassGenerator extends AbstractGenerator implements Generator {
 
     private final BuilderClassParams builderClassParams;
 
@@ -16,18 +16,18 @@ class BuilderClassGenerator extends AbstractGenerator {
     }
 
     @Override
-    public void run() {
+    public GenerationResult generate() {
         //builder constructor
         var builderClass = builderClassParams.builderClass();
         var builderConstructor = generateBuilderConstructor();
         addMethod(builderClass, null, builderConstructor, false);
 
         var fieldsGenerator = generatorFactory.createBuilderFieldsGenerator(generatorParams, builderClassParams);
-        fieldsGenerator.run();
+        fieldsGenerator.generate();
 
         var methodsGenerator = generatorFactory.createBuilderMethodsGenerator(generatorParams,
                 builderClassParams, fieldsGenerator);
-        methodsGenerator.run();
+        return methodsGenerator.generate();
     }
 
     private PsiMethod generateBuilderConstructor() {

src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderFieldsGenerator.java

@@ -27,13 +27,14 @@ public List<PsiField> getFields() {
     }
 
     @Override
-    public void run() {
+    public GenerationResult generate() {
         PsiField lastAddedField = null;
         for (var fieldMember : generatorParams.psi().selectedFields()) {
             lastAddedField = createOrUpdateField(builderClassParams.builderClass(), fieldMember, lastAddedField);
             fields.add(lastAddedField);
         }
         cleanupFields(builderClassParams.builderClass());
+        return GenerationResult.NO_RESULT;
     }
 
     private void cleanupFields(PsiClass builderClass) {
@@ -44,8 +45,8 @@ private void cleanupFields(PsiClass builderClass) {
         }
     }
 
-    private PsiField createOrUpdateField(final PsiClass builderClass, final PsiFieldMember member,
-                                         @Nullable final PsiElement last) {
+    private PsiField createOrUpdateField(PsiClass builderClass, PsiFieldMember member,
+                                         @Nullable PsiElement last) {
         var psiFactory = generatorParams.psi().factory();
         var field = member.getElement();
         var fieldName = field.getName();

src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderMethodsGenerator.java

@@ -2,6 +2,7 @@
 
 import com.github.junkfactory.innerbuilder.ui.JavaInnerBuilderOption;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiField;
 import com.intellij.psi.PsiMethod;
@@ -15,6 +16,9 @@ class BuilderMethodsGenerator extends AbstractGenerator implements MethodsGenera
 
     private final BuilderClassParams builderClassParams;
     private final FieldsGenerator fieldsGenerator;
+    private final GenerationResult generationResult;
+
+    private boolean isPublic;
 
     BuilderMethodsGenerator(GeneratorFactory generatorFactory,
                             GeneratorParams generatorParams,
@@ -23,11 +27,15 @@ class BuilderMethodsGenerator extends AbstractGenerator implements MethodsGenera
         super(generatorFactory, generatorParams);
         this.builderClassParams = builderClassParams;
         this.fieldsGenerator = fieldsGenerator;
+        this.generationResult = new GenerationResult();
     }
 
     @Override
-    public void run() {
+    public GenerationResult generate() {
         var builderClass = builderClassParams.builderClass();
+        var targetClass = builderClassParams.targetClass();
+        var targetModifierList = Objects.requireNonNull(targetClass.getModifierList());
+        isPublic = targetModifierList.hasModifierProperty(PsiModifier.PUBLIC);
         PsiElement lastAddedElement = null;
         for (var field : fieldsGenerator.getFields()) {
             var setterMethod = generateFieldMethod(field);
@@ -41,8 +49,9 @@ public void run() {
             addMethod(builderClass, lastAddedElement, validateMethod, false);
         }
 
-        var buildMethod = generateBuildMethod();
+        var buildMethod = generateBuildMethod(targetClass);
         addMethod(builderClass, null, buildMethod, builderClassParams.targetClass().isRecord());
+        return generationResult;
     }
 
     private PsiMethod generateValidateMethod() {
@@ -71,58 +80,108 @@ private PsiMethod generatePutToMap(PsiField field, PsiMethod fieldPutMethod) {
         //resolve the generic type of the map via the parameter type of the put method
         var param1 = Objects.requireNonNull(fieldPutMethod.getParameterList().getParameter(0));
         var param1Type = Utils.resolveGenericParameterType(field.getType(), param1);
+        var importAdded = addImport(param1Type);
+
         var param2 = Objects.requireNonNull(fieldPutMethod.getParameterList().getParameter(1));
         var param2Type = Utils.resolveGenericParameterType(field.getType(), param2);
+        importAdded = addImport(param2Type) || importAdded;
+        if (importAdded) {
+            generationResult.set(GenerationResult.Code.IMPORTS_ADDED);
+        }
 
-        //now build the put method
         var methodName = "putTo" + StringUtil.capitalize(field.getName());
-        var methodText = """
-                public %s %s(%s key, %s value) {
-                    this.%s.put(key, value);
-                    return this;
-                }""".formatted(BUILDER_CLASS_NAME, methodName, param1Type.getPresentableText(),
-                param2Type.getPresentableText(), field.getName());
+        var methodText = new StringBuilder();
+        if (isPublic) {
+            methodText.append(PsiModifier.PUBLIC).append(' ');
+        }
+        methodText.append(BUILDER_CLASS_NAME)
+                .append(' ')
+                .append(methodName)
+                .append('(')
+                .append(param1Type.getPresentableText())
+                .append(' ')
+                .append(param1.getName().toLowerCase())
+                .append(", ")
+                .append(param2Type.getPresentableText())
+                .append(' ')
+                .append(param2.getName().toLowerCase())
+                .append(") {")
+                .append(THIS_DOT)
+                .append(field.getName())
+                .append(".put(")
+                .append(param1.getName().toLowerCase())
+                .append(", ")
+                .append(param2.getName().toLowerCase())
+                .append(");")
+                .append(RETURN_THIS)
+                .append('}');
         var psiElementFactory = generatorParams.psi().factory();
-        return psiElementFactory.createMethodFromText(methodText, field);
+        return psiElementFactory.createMethodFromText(methodText.toString(), field);
     }
 
     private PsiMethod generateAddToCollection(PsiField field, PsiMethod fieldAddMethod) {
         //resolve the generic type of the collection via the parameter type of the add method
         var param = Objects.requireNonNull(fieldAddMethod.getParameterList().getParameter(0));
         var paramType = Utils.resolveGenericParameterType(field.getType(), param);
+        if (addImport(paramType)) {
+            generationResult.set(GenerationResult.Code.IMPORTS_ADDED);
+        }
 
-        //now build the add method
         var methodName = "addTo" + StringUtil.capitalize(field.getName());
-        var methodText = """
-                public %s %s(%s %s) {
-                    this.%s.add(%s);
-                    return this;
-                }""".formatted(BUILDER_CLASS_NAME, methodName, paramType.getPresentableText(),
-                param.getName().toLowerCase(), field.getName(), param.getName());
+        var methodText = new StringBuilder();
+        if (isPublic) {
+            methodText.append(PsiModifier.PUBLIC).append(' ');
+        }
+        methodText.append(BUILDER_CLASS_NAME)
+                .append(' ')
+                .append(methodName)
+                .append('(')
+                .append(paramType.getPresentableText())
+                .append(' ')
+                .append(param.getName().toLowerCase())
+                .append(") {")
+                .append(THIS_DOT)
+                .append(field.getName())
+                .append(".add(")
+                .append(param.getName())
+                .append(");")
+                .append(RETURN_THIS)
+                .append('}');
+
         var psiElementFactory = generatorParams.psi().factory();
-        return psiElementFactory.createMethodFromText(methodText, field);
+        return psiElementFactory.createMethodFromText(methodText.toString(), field);
     }
 
     private PsiMethod generateBuilderSetter(PsiField field) {
         var fieldType = field.getType();
         var fieldName = Utils.hasOneLetterPrefix(field.getName()) ?
                 Character.toLowerCase(field.getName().charAt(1)) + field.getName().substring(2) : field.getName();
 
-        var methodText = """
-                public %s %s(%s %s) {
-                    this.%s = %s;
-                    return this;
-                }""".formatted(BUILDER_CLASS_NAME, fieldName, fieldType.getPresentableText(),
-                fieldName, field.getName(), fieldName);
+        var methodText = new StringBuilder();
+        if (isPublic) {
+            methodText.append(PsiModifier.PUBLIC).append(' ');
+        }
+        methodText.append(BUILDER_CLASS_NAME)
+                .append(' ')
+                .append(fieldName)
+                .append('(')
+                .append(fieldType.getPresentableText())
+                .append(' ')
+                .append(fieldName)
+                .append(") {")
+                .append(THIS_DOT)
+                .append(field.getName())
+                .append(" = ")
+                .append(fieldName)
+                .append(";")
+                .append(RETURN_THIS)
+                .append('}');
+
         var psiElementFactory = generatorParams.psi().factory();
-        return psiElementFactory.createMethodFromText(methodText, field);
+        return psiElementFactory.createMethodFromText(methodText.toString(), field);
     }
 
-    private PsiMethod generateBuildMethod() {
-        var targetClass = builderClassParams.targetClass();
-        var targetModifierList = Objects.requireNonNull(targetClass.getModifierList());
-        boolean isPublic = targetModifierList.hasModifierProperty(PsiModifier.PUBLIC);
-
+    private PsiMethod generateBuildMethod(PsiClass targetClass) {
         var buildMethod = new StringBuilder()
                 .append(isPublic ? PsiModifier.PUBLIC : EMPTY)
                 .append(isPublic ? SPACE : EMPTY)
@@ -149,4 +208,5 @@ private PsiMethod generateBuildMethod() {
 
         return generatorParams.psi().factory().createMethodFromText(buildMethod.toString(), targetClass);
     }
+
 }

src/main/java/com/github/junkfactory/innerbuilder/generators/FieldsGenerator.java

@@ -4,6 +4,6 @@
 
 import java.util.List;
 
-public interface FieldsGenerator extends Runnable {
+public interface FieldsGenerator extends Generator {
     List<PsiField> getFields();
 }

src/main/java/com/github/junkfactory/innerbuilder/generators/GenerationResult.java

@@ -0,0 +1,37 @@
+package com.github.junkfactory.innerbuilder.generators;
+
+import java.util.BitSet;
+
+public class GenerationResult {
+
+    static final GenerationResult NO_RESULT = new GenerationResult();
+
+    public enum Code {
+        IMPORTS_ADDED,
+        ANNOTATIONS_ADDED
+    }
+
+    private final BitSet result;
+
+    GenerationResult() {
+        this.result = new BitSet();
+    }
+
+    void set(Code code) {
+        result.set(code.ordinal());
+    }
+
+    public boolean isEmpty() {
+        return !result.isEmpty();
+    }
+
+    public boolean did(Code code) {
+        return result.get(code.ordinal());
+    }
+
+    public void when(Code code, Runnable runnable) {
+        if (did(code)) {
+            runnable.run();
+        }
+    }
+}