Handle refactors (#5)

* Handle refactors Separate generator interfaces * Fix addTo generation and cleanup deleted fields * Reuse addElement * Remove ancestor check
junkfactory · Aug 5, 2024 · 836d9bf · 836d9bf
1 parent 97b236f
commit 836d9bf
Show file tree

Hide file tree

Showing 11 changed files with 130 additions and 79 deletions.
diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/AbstractGenerator.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/AbstractGenerator.java
@@ -24,28 +24,34 @@ protected AbstractGenerator(GeneratorFactory generatorFactory, GeneratorParams g
         this.generatorParams = generatorParams;
     }
 
+    protected PsiElement addElement(PsiElement target, PsiElement element, PsiElement after) {
+        if (after != null) {
+            return target.addAfter(element, after);
+        }
+        return target.add(element);
+    }
+
     protected PsiElement addMethod(@NotNull final PsiClass target, @Nullable final PsiElement after,
                                    @NotNull final PsiMethod newMethod, final boolean replace) {
         var existingMethod = target.findMethodBySignature(newMethod, false);
         if (existingMethod == null && newMethod.isConstructor()) {
-            for (final PsiMethod constructor : target.getConstructors()) {
-                if (Utils.areParameterListsEqual(constructor.getParameterList(),
-                        newMethod.getParameterList())) {
-                    existingMethod = constructor;
-                    break;
-                }
-            }
+            existingMethod = findConstructor(target, newMethod);
         }
         if (existingMethod == null) {
-            if (after != null) {
-                return target.addAfter(newMethod, after);
-            } else {
-                return target.add(newMethod);
-            }
+            return addElement(target, newMethod, after);
         } else if (replace) {
             existingMethod.replace(newMethod);
         }
         return existingMethod;
     }
 
+    private PsiMethod findConstructor(PsiClass target, PsiMethod newMethod) {
+        for (var constructor : target.getConstructors()) {
+            if (Utils.areParameterListsEqual(constructor.getParameterList(), newMethod.getParameterList())) {
+                return constructor;
+            }
+        }
+        return null;
+    }
+
 }
diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderClassGenerator.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderClassGenerator.java
@@ -7,16 +7,12 @@
 class BuilderClassGenerator extends AbstractGenerator {
 
     private final BuilderClassParams builderClassParams;
-    private final Runnable fieldsGenerator;
-    private final Runnable methodsGenerator;
 
     BuilderClassGenerator(GeneratorFactory generatorFactory,
                           GeneratorParams generatorParams,
                           BuilderClassParams builderClassParams) {
         super(generatorFactory, generatorParams);
         this.builderClassParams = builderClassParams;
-        this.fieldsGenerator = generatorFactory.createBuilderFieldsGenerator(generatorParams, builderClassParams);
-        this.methodsGenerator = generatorFactory.createBuilderMethodsGenerator(generatorParams, builderClassParams);
     }
 
     @Override
@@ -26,7 +22,11 @@ public void run() {
         var builderConstructor = generateBuilderConstructor();
         addMethod(builderClass, null, builderConstructor, false);
 
+        var fieldsGenerator = generatorFactory.createBuilderFieldsGenerator(generatorParams, builderClassParams);
         fieldsGenerator.run();
+
+        var methodsGenerator = generatorFactory.createBuilderMethodsGenerator(generatorParams,
+                builderClassParams, fieldsGenerator);
         methodsGenerator.run();
     }
 

diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderFieldsGenerator.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderFieldsGenerator.java
@@ -3,11 +3,17 @@
 import com.intellij.codeInsight.generation.PsiFieldMember;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiField;
 import org.jetbrains.annotations.Nullable;
 
-class BuilderFieldsGenerator extends AbstractGenerator {
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Optional;
+
+class BuilderFieldsGenerator extends AbstractGenerator implements FieldsGenerator {
 
     private final BuilderClassParams builderClassParams;
+    private final List<PsiField> fields = new LinkedList<>();
 
     BuilderFieldsGenerator(GeneratorFactory generatorFactory,
                            GeneratorParams generatorParams,
@@ -16,37 +22,56 @@ class BuilderFieldsGenerator extends AbstractGenerator {
         this.builderClassParams = builderClassParams;
     }
 
+    @Override
+    public List<PsiField> getFields() {
+        return fields;
+    }
+
     @Override
     public void run() {
-        PsiElement lastAddedField = null;
+        PsiField lastAddedField = null;
         for (var fieldMember : generatorParams.psi().selectedFields()) {
-            lastAddedField = findOrCreateField(builderClassParams.builderClass(), fieldMember, lastAddedField);
+            lastAddedField = createOrUpdateField(builderClassParams.builderClass(), fieldMember, lastAddedField);
+            fields.add(lastAddedField);
         }
+        cleanupFields(builderClassParams.builderClass());
     }
 
-    private PsiElement findOrCreateField(final PsiClass builderClass, final PsiFieldMember member,
+    private void cleanupFields(PsiClass builderClass) {
+        for (var field : builderClass.getFields()) {
+            if (!fields.contains(field)) {
+                deleteFieldAndMethodIfExists(builderClass, field);
+            }
+        }
+    }
+
+    private PsiField createOrUpdateField(final PsiClass builderClass, final PsiFieldMember member,
                                          @Nullable final PsiElement last) {
         var psiFactory = generatorParams.psi().factory();
         var field = member.getElement();
         var fieldName = field.getName();
         var fieldType = field.getType();
         var existingField = builderClass.findFieldByName(fieldName, false);
         if (existingField == null || Utils.areTypesPresentableNotEqual(existingField.getType(), fieldType)) {
-            if (existingField != null) {
-                existingField.delete();
-            }
+            deleteFieldAndMethodIfExists(builderClass, existingField);
             var newField = psiFactory.createField(fieldName, fieldType);
             newField.setInitializer(field.getInitializer());
             if (!builderClassParams.targetClass().isRecord()) {
                 field.setInitializer(null);
             }
-            if (last != null) {
-                return builderClass.addAfter(newField, last);
-            } else {
-                return builderClass.add(newField);
-            }
+            existingField = (PsiField) addElement(builderClass, newField, last);
         }
         return existingField;
     }
 
+    private void deleteFieldAndMethodIfExists(PsiClass builderClass, PsiField field) {
+        if (null == field) {
+            return;
+        }
+        Optional.ofNullable(field.getCopyableUserData(UserDataKey.METHOD_REF))
+                .map(m -> builderClass.findMethodBySignature(m, false))
+                .ifPresent(PsiElement::delete);
+        field.delete();
+    }
+
 }
diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderMethodsGenerator.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/BuilderMethodsGenerator.java
@@ -1,35 +1,41 @@
 package com.github.junkfactory.innerbuilder.generators;
 
 import com.github.junkfactory.innerbuilder.ui.JavaInnerBuilderOption;
-import com.intellij.codeInsight.generation.PsiFieldMember;
 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;
 import com.intellij.psi.PsiModifier;
 import com.intellij.psi.PsiStatement;
 import com.intellij.psi.util.PsiUtil;
 
 import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-class BuilderMethodsGenerator extends AbstractGenerator {
+class BuilderMethodsGenerator extends AbstractGenerator implements MethodsGenerator {
 
     private final BuilderClassParams builderClassParams;
+    private final FieldsGenerator fieldsGenerator;
 
     BuilderMethodsGenerator(GeneratorFactory generatorFactory,
                             GeneratorParams generatorParams,
-                            BuilderClassParams builderClassParams) {
+                            BuilderClassParams builderClassParams,
+                            FieldsGenerator fieldsGenerator) {
         super(generatorFactory, generatorParams);
         this.builderClassParams = builderClassParams;
+        this.fieldsGenerator = fieldsGenerator;
     }
 
     @Override
     public void run() {
         var builderClass = builderClassParams.builderClass();
         PsiElement lastAddedElement = null;
-        for (var member : generatorParams.psi().selectedFields()) {
-            var setterMethod = generateFieldMethod(member);
+        for (var field : fieldsGenerator.getFields()) {
+            var setterMethod = generateFieldMethod(field);
+            field.putCopyableUserData(UserDataKey.METHOD_REF, setterMethod);
             lastAddedElement = addMethod(builderClass, lastAddedElement, setterMethod, false);
         }
 
@@ -40,7 +46,7 @@ public void run() {
         }
 
         var buildMethod = generateBuildMethod();
-        addMethod(builderClass, null, buildMethod, builderClassParams.builderClass().isRecord());
+        addMethod(builderClass, null, buildMethod, builderClassParams.targetClass().isRecord());
     }
 
     private PsiMethod generateValidateMethod() {
@@ -51,16 +57,28 @@ private PsiMethod generateValidateMethod() {
         return validateMethod;
     }
 
-    private PsiMethod generateFieldMethod(PsiFieldMember member) {
-        var addMethod = Utils.findFieldAddMethod(builderClassParams.builderClass(), member);
+    private PsiMethod generateFieldMethod(PsiField field) {
+        var addMethod = field.hasInitializer() ? findAddMethod(field) : null;
         if (null != addMethod) {
-            return generateAddToCollection(member, addMethod);
+            return generateAddToCollection(field, addMethod);
         }
-        return generateBuilderSetter(member);
+        return generateBuilderSetter(field);
     }
 
-    private PsiMethod generateAddToCollection(PsiFieldMember member, PsiMethod fieldAddMethod) {
-        var field = member.getElement();
+    private PsiMethod findAddMethod(PsiField field) {
+        var fieldClass = PsiUtil.resolveClassInClassTypeOnly(field.getType());
+        var methods = Optional.ofNullable(fieldClass)
+                .map(PsiClass::getAllMethods)
+                .orElseGet(() -> new PsiMethod[0]);
+        for (var method : methods) {
+            if (method.getName().equals("add") && method.getParameterList().getParametersCount() == 1) {
+                return method;
+            }
+        }
+        return null;
+    }
+
+    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 = PsiUtil.resolveGenericsClassInType(field.getType())
@@ -84,9 +102,7 @@ private PsiMethod generateAddToCollection(PsiFieldMember member, PsiMethod field
         return addMethod;
     }
 
-    private PsiMethod generateBuilderSetter(final PsiFieldMember member) {
-
-        var field = member.getElement();
+    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();

diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/FieldCollector.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/FieldCollector.java
@@ -5,7 +5,6 @@
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiClass;
-import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiField;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiModifier;
@@ -45,17 +44,16 @@ public static List<PsiFieldMember> collectFields(final PsiFile file, final Edito
 
         PsiClass classToExtractFieldsFrom = clazz;
         while (classToExtractFieldsFrom != null) {
-            var classFieldMembers = collectFieldsInClass(element, clazz, classToExtractFieldsFrom);
+            var classFieldMembers = collectFieldsInClass(clazz, classToExtractFieldsFrom);
             allFields.addAll(0, classFieldMembers);
             classToExtractFieldsFrom = classToExtractFieldsFrom.getSuperClass();
         }
 
         return allFields;
     }
 
-    private static List<PsiFieldMember> collectFieldsInClass(final PsiElement element,
-                                                             final PsiClass accessObjectClass,
-                                                             final PsiClass classToExtractFieldsFrom) {
+    private static List<PsiFieldMember> collectFieldsInClass(PsiClass accessObjectClass,
+                                                             PsiClass classToExtractFieldsFrom) {
         if (AbstractGenerator.BUILDER_CLASS_NAME.equals(classToExtractFieldsFrom.getName()) ||
                 OBJECT_CLASS_NAME.equals(classToExtractFieldsFrom.getName())) {
             return List.of();
@@ -64,7 +62,6 @@ private static List<PsiFieldMember> collectFieldsInClass(final PsiElement elemen
         return Arrays.stream(classToExtractFieldsFrom.getFields())
                 .filter(field -> helper.isAccessible(field, classToExtractFieldsFrom, accessObjectClass) ||
                         hasSetter(classToExtractFieldsFrom, field.getName()))
-                .filter(field -> !PsiTreeUtil.isAncestor(field, element, false))
                 .filter(field -> !field.hasModifierProperty(PsiModifier.STATIC))
                 .filter(field -> hasLowerCaseChar(field.getName()))
                 .filter(field -> Objects.nonNull(field.getContainingClass()))

diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/FieldsGenerator.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/FieldsGenerator.java
@@ -0,0 +1,9 @@
+package com.github.junkfactory.innerbuilder.generators;
+
+import com.intellij.psi.PsiField;
+
+import java.util.List;
+
+public interface FieldsGenerator extends Runnable {
+    List<PsiField> getFields();
+}
diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/GeneratorFactory.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/GeneratorFactory.java
@@ -11,9 +11,10 @@ static GeneratorFactory create() {
     Runnable createBuilderClassGenerator(GeneratorParams generatorParams,
                                          BuilderClassParams builderClassParams);
 
-    Runnable createBuilderFieldsGenerator(GeneratorParams generatorParams,
-                                          BuilderClassParams builderClassParams);
+    FieldsGenerator createBuilderFieldsGenerator(GeneratorParams generatorParams,
+                                                 BuilderClassParams builderClassParams);
 
-    Runnable createBuilderMethodsGenerator(GeneratorParams generatorParams,
-                                           BuilderClassParams builderClassParams);
+    MethodsGenerator createBuilderMethodsGenerator(GeneratorParams generatorParams,
+                                                   BuilderClassParams builderClassParams,
+                                                   FieldsGenerator fieldsGenerator);
 }
diff --git a/...ain/java/com/github/junkfactory/innerbuilder/generators/InnerBuilderGeneratorFactory.java b/...ain/java/com/github/junkfactory/innerbuilder/generators/InnerBuilderGeneratorFactory.java
@@ -13,14 +13,15 @@ public Runnable createBuilderClassGenerator(GeneratorParams generatorParams,
     }
 
     @Override
-    public Runnable createBuilderFieldsGenerator(GeneratorParams generatorParams,
-                                                 BuilderClassParams builderClassParams) {
+    public FieldsGenerator createBuilderFieldsGenerator(GeneratorParams generatorParams,
+                                                        BuilderClassParams builderClassParams) {
         return new BuilderFieldsGenerator(this, generatorParams, builderClassParams);
     }
 
     @Override
-    public Runnable createBuilderMethodsGenerator(GeneratorParams generatorParams,
-                                                  BuilderClassParams builderClassParams) {
-        return new BuilderMethodsGenerator(this, generatorParams, builderClassParams);
+    public MethodsGenerator createBuilderMethodsGenerator(GeneratorParams generatorParams,
+                                                          BuilderClassParams builderClassParams,
+                                                          FieldsGenerator fieldsGenerator) {
+        return new BuilderMethodsGenerator(this, generatorParams, builderClassParams, fieldsGenerator);
     }
 }
diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/MethodsGenerator.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/MethodsGenerator.java
@@ -0,0 +1,4 @@
+package com.github.junkfactory.innerbuilder.generators;
+
+public interface MethodsGenerator extends Runnable {
+}
diff --git a/src/main/java/com/github/junkfactory/innerbuilder/generators/UserDataKey.java b/src/main/java/com/github/junkfactory/innerbuilder/generators/UserDataKey.java
@@ -0,0 +1,13 @@
+package com.github.junkfactory.innerbuilder.generators;
+
+import com.intellij.openapi.util.Key;
+import com.intellij.psi.PsiMethod;
+
+final class UserDataKey {
+
+    private UserDataKey() {
+    }
+
+    static final Key<PsiMethod> METHOD_REF = Key.create("METHOD_REF");
+
+}