calculated correct generic type for toBuilder

#134
(TODO: clean up and refactor code)

src/main/java/de/plushnikov/intellij/plugin/processor/handler/BuilderHandler.java

@@ -2,11 +2,13 @@
 
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiAnnotation;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiClassType;
 import com.intellij.psi.PsiCodeBlock;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiElementFactory;
 import com.intellij.psi.PsiField;
 import com.intellij.psi.PsiMethod;
 import com.intellij.psi.PsiModifier;
@@ -18,6 +20,7 @@
 import com.intellij.psi.PsiTypeParameter;
 import com.intellij.psi.PsiTypeParameterListOwner;
 import com.intellij.psi.PsiVariable;
+import com.intellij.psi.impl.source.PsiClassReferenceType;
 import com.intellij.psi.util.PsiTypesUtil;
 import de.plushnikov.intellij.plugin.problem.ProblemBuilder;
 import de.plushnikov.intellij.plugin.processor.ShouldGenerateFullCodeBlock;
@@ -305,20 +308,55 @@ public void createBuilderMethodIfNecessary(@NotNull Collection<? super PsiElemen
 
   public void createToBuilderMethodIfNecessary(@NotNull Collection<? super PsiElement> target, @NotNull PsiClass containingClass, @Nullable PsiMethod psiMethod, @NotNull PsiClass builderPsiClass, @NotNull PsiAnnotation psiAnnotation) {
     if (PsiAnnotationUtil.getBooleanAnnotationValue(psiAnnotation, TO_BUILDER_ANNOTATION_KEY, false)) {
-      final PsiType psiTypeWithGenerics = PsiClassUtil.getTypeWithGenerics(builderPsiClass);
+
+      final PsiType psiTypeWithGenerics;
+      if (null != psiMethod) {
+        psiTypeWithGenerics = calculateResultType(psiMethod, builderPsiClass, containingClass);
+      } else {
+        psiTypeWithGenerics = PsiClassUtil.getTypeWithGenerics(builderPsiClass);
+      }
 
       final LombokLightMethodBuilder method = new LombokLightMethodBuilder(containingClass.getManager(), TO_BUILDER_METHOD_NAME)
         .withMethodReturnType(psiTypeWithGenerics)
         .withContainingClass(containingClass)
         .withNavigationElement(psiAnnotation)
         .withModifier(PsiModifier.PUBLIC);
-      addTypeParameters(builderPsiClass, psiMethod, method);
 
       method.withBody(createToBuilderMethodCodeBlock(containingClass, psiMethod, psiTypeWithGenerics));
       target.add(method);
     }
   }
 
+  private PsiType calculateResultType(@NotNull PsiMethod psiMethod, PsiClass builderPsiClass, PsiClass psiClass) {
+    final Collection<PsiParameter> builderParameters = getBuilderParameters(psiMethod, Collections.<PsiField>emptySet());
+
+    final Collection<PsiType> types = new ArrayList<>();
+    for (PsiVariable psiVariable : builderParameters) {
+
+      final PsiAnnotation obtainViaAnnotation = PsiAnnotationSearchUtil.findAnnotation(psiVariable, BUILDER_OBTAIN_VIA_ANNOTATION);
+      if (null != obtainViaAnnotation) {
+        final String viaFieldName = PsiAnnotationUtil.getStringAnnotationValue(obtainViaAnnotation, BUILDER_OBTAIN_VIA_FIELD);
+
+        final PsiField fieldByName = psiClass.findFieldByName(viaFieldName, false);
+        if (fieldByName != null) {
+          psiVariable = fieldByName;
+        }
+      }
+
+
+      final PsiType psiVariableType = psiVariable.getType();
+
+      if (psiVariableType instanceof PsiClassReferenceType) {
+        final PsiClass resolvedPsiVariableClass = ((PsiClassReferenceType) psiVariableType).resolve();
+        if (resolvedPsiVariableClass instanceof PsiTypeParameter) {
+          types.add(psiVariableType);
+        }
+      }
+    }
+    final PsiElementFactory factory = JavaPsiFacade.getElementFactory(psiClass.getProject());
+    return factory.createType(builderPsiClass, types.toArray(new PsiType[types.size()]));
+  }
+
   @NotNull
   private PsiCodeBlock createBuilderMethodCodeBlock(@NotNull PsiClass containingClass, @NotNull PsiType psiTypeWithGenerics) {
     final String blockText;
@@ -347,10 +385,6 @@ private PsiCodeBlock createToBuilderMethodCodeBlock(@NotNull PsiClass psiClass,
 
       final StringBuilder methodCalls = new StringBuilder();
       for (PsiVariable psiVariable : builderVariables) {
-        //TODO
-//        psiVariable.getType()
-//        JavaPsiFacade.getElementFactory(project).createType(genericClass, psiTypes)
-
         methodCalls.append('.');
         methodCalls.append(accessorsInfo.removePrefix(psiVariable.getName()));
         methodCalls.append('(');

test-manual/src/main/java/de/plushnikov/builder/BuilderExperimentalExample.java

@@ -1,6 +1,6 @@
 package de.plushnikov.builder;
 
-import lombok.Builder;
+import lombok.experimental.Builder;
 
 @Builder
 public class BuilderExperimentalExample {

test-manual/src/main/java/de/plushnikov/builder/tobuilder/BuilderWithToBuilderOnClass.java

@@ -1,6 +1,6 @@
 package de.plushnikov.builder.tobuilder;
 
-import java.util.Arrays;
+import java.util.Collections;
 
 @lombok.Builder(toBuilder = true)
 @lombok.experimental.Accessors(prefix = "m")
@@ -18,16 +18,16 @@ public static <K> K rrr(BuilderWithToBuilderOnClass<K> x) {
   }
 
   public static void main(String[] args) {
-    BuilderWithToBuilderOnClass<String> bean = new BuilderWithToBuilderOnClass<String>("mOneParam", "mTwoParam", "fooParam", Arrays.asList("barsParam1"));
+    BuilderWithToBuilderOnClass<String> bean = new BuilderWithToBuilderOnClass<String>("mOneParam", "mTwoParam", "fooParam", Collections.singletonList("barsParam1"));
     bean.mOne = "mOne";
     bean.mTwo = "mTwo";
     bean.foo = "foo";
 
-    BuilderWithToBuilderOnClass.BuilderWithToBuilderOnClassBuilder x = bean.toBuilder();
+    BuilderWithToBuilderOnClass.BuilderWithToBuilderOnClassBuilder<String> x = bean.toBuilder();
     System.out.println(x);
 
     x.one("builderOne");
     x.bar("builderBar");
     System.out.println(x);
   }
-}
+}

test-manual/src/main/java/de/plushnikov/builder/tobuilder/BuilderWithToBuilderOnConstructor.java

@@ -19,7 +19,7 @@ public static void main(String[] args) {
     bean.mTwo = "mTwo";
     bean.foo = "foo";
 
-    BuilderWithToBuilderOnConstructor.BuilderWithToBuilderOnConstructorBuilder x = bean.toBuilder();
+    BuilderWithToBuilderOnConstructor.BuilderWithToBuilderOnConstructorBuilder<String> x = bean.toBuilder();
     System.out.println(x);
 
     x.mOne("builderOne");

test-manual/src/main/java/de/plushnikov/builder/tobuilder/BuilderWithToBuilderOnMethod.java

@@ -21,7 +21,7 @@ public static void main(String[] args) {
     bean.bar = "bar";
     bean.foo = "foo";
 
-    BuilderWithToBuilderOnMethod.BuilderWithToBuilderOnMethodBuilder x = bean.toBuilder();
+    BuilderWithToBuilderOnMethod.BuilderWithToBuilderOnMethodBuilder<String> x = bean.toBuilder();
     System.out.println(x);
 
     x.mOne("builderOne");

test-manual/src/main/java/de/plushnikov/builder/tobuilder/SimpleBuilderWithToBuilderOnClass.java

@@ -1,6 +1,6 @@
 package de.plushnikov.builder.tobuilder;
 
-import java.util.Arrays;
+import java.util.Collections;
 
 @lombok.Builder(toBuilder = true)
 public class SimpleBuilderWithToBuilderOnClass {
@@ -17,7 +17,7 @@ public static String rrr(SimpleBuilderWithToBuilderOnClass x) {
   }
 
   public static void main(String[] args) {
-    SimpleBuilderWithToBuilderOnClass bean = new SimpleBuilderWithToBuilderOnClass("mOneParam", "mTwoParam", "fooParam", Arrays.asList("barsParam1"));
+    SimpleBuilderWithToBuilderOnClass bean = new SimpleBuilderWithToBuilderOnClass("mOneParam", "mTwoParam", "fooParam", Collections.singletonList("barsParam1"));
     bean.one = "one";
     bean.two = "two";
     bean.foo = "foo";
@@ -29,4 +29,4 @@ public static void main(String[] args) {
     x.bar("builderBar");
     System.out.println(x);
   }
-}
+}

testData/after/builder/BuilderWithToBuilder.java

@@ -93,15 +93,15 @@ public BuilderWithToBuilderBuilder<T> clearBars() {
 		public BuilderWithToBuilder<T> build() {
 			java.util.List<T> bars;
 			switch (this.bars == null ? 0 : this.bars.size()) {
-			case 0: 
+			case 0:
 				bars = java.util.Collections.emptyList();
 				break;
 
-			case 1: 
+			case 1:
 				bars = java.util.Collections.singletonList(this.bars.get(0));
 				break;
 
-			default: 
+			default:
 				bars = java.util.Collections.unmodifiableList(new java.util.ArrayList<T>(this.bars));
 			}
 			return new BuilderWithToBuilder<T>(one, two, foo, bars);
@@ -241,7 +241,7 @@ public StaticWithToBuilderBuilder<Z> bar(final Z bar) {
 
 		@java.lang.SuppressWarnings("all")
 		@javax.annotation.Generated("lombok")
-		public StaticWithToBuilder build() {
+		public StaticWithToBuilder<Z, String> build() {
 			return StaticWithToBuilder.<Z>test(mOne, bar);
 		}
 

testData/after/builder/BuilderWithToBuilderOnMethod.java

@@ -43,7 +43,7 @@ public BuilderWithToBuilderOnMethodBuilder<Z> bar(final Z bar) {
 
 		@java.lang.SuppressWarnings("all")
 		@javax.annotation.Generated("lombok")
-		public BuilderWithToBuilderOnMethod build() {
+		public BuilderWithToBuilderOnMethod<Z, String> build() {
 			return BuilderWithToBuilderOnMethod.<Z>test(one, bar);
 		}