Implementation of LombokReferenceSearcher and LombokMethodReferenceSearcher (fix for #685)

src/main/java/de/plushnikov/intellij/plugin/extension/LombokMethodReferenceSearcher.java

@@ -0,0 +1,141 @@
+package de.plushnikov.intellij.plugin.extension;
+
+import com.intellij.openapi.application.QueryExecutorBase;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiMethod;
+import com.intellij.psi.PsiNewExpression;
+import com.intellij.psi.PsiReference;
+import com.intellij.psi.PsiReferenceBase;
+import com.intellij.psi.search.searches.MethodReferencesSearch;
+import com.intellij.util.Processor;
+import de.plushnikov.intellij.plugin.psi.LombokLightMethodBuilder;
+import lombok.Builder;
+import lombok.With;
+import lombok.experimental.Wither;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class LombokMethodReferenceSearcher extends QueryExecutorBase<PsiReference, MethodReferencesSearch.SearchParameters> {
+
+  public LombokMethodReferenceSearcher(boolean requireReadAction) {
+    super(requireReadAction);
+    System.out.println("LombokMethodReferenceSearcher(boolean)");
+  }
+
+  public LombokMethodReferenceSearcher() {
+    System.out.println("LombokMethodReferenceSearcher()");
+  }
+
+  @Override
+  public void processQuery(@NotNull MethodReferencesSearch.SearchParameters queryParameters, @NotNull Processor<? super PsiReference> consumer) {
+    PsiMethod queryMethod = queryParameters.getMethod();
+    LombokLightMethodBuilder lombokMethod = LombokLightMethodBuilder.getLightMethodBuilder(queryMethod);
+    PsiMethod actualMethod = lombokMethod == null ? queryMethod : lombokMethod;
+
+    // Each Lombok-generated PsiMethod exists twice.
+    // - Once as an instance of LombokLightMethodBuilder which is properly wired by to its containing class.
+    // - Once as a degenerated instance of PsiMethodImpl which is built by LombokLightMethodBuilder and is not properly wired to its containing class
+
+    if (lombokMethod != null && queryMethod != lombokMethod) {
+      // We replace the search for the degenerated PsiMethodImpl by a search for the rich LombokLightMethodBuilder
+      // (but only the query is not already the LombokLightMethodBuilder, otherwise we'll cause an infinite loop)
+      MethodReferencesSearch.SearchParameters newParameters = new MethodReferencesSearch.SearchParameters(
+        lombokMethod,
+        queryParameters.getScopeDeterminedByUser(),
+        queryParameters.isStrictSignatureSearch(),
+        queryParameters.getOptimizer());
+      MethodReferencesSearch.search(newParameters).forEach(consumer);
+    }
+
+
+    if (shouldLookForConstructorReferences(actualMethod)) {
+      findConstructorReferencesInClassMethods(actualMethod.getContainingClass(), actualMethod, consumer);
+      findConstructorReferencesInInnerClasses(actualMethod.getContainingClass(), actualMethod, consumer);
+    }
+  }
+
+  /**
+   * If the method is a constructor, we must also look into class methods (it might be used in @With methods),
+   * and inner classes (it might be used in @Builder methods). Note that @With and @Builder will use physical
+   * constructors if available, so we have to look for constructor references even if the constructor is not
+   * generated by lombok.
+   * <p>
+   * However, in order to avoid searching deep into method implementations for *each* Constructor reference search and
+   * save a bit of CPU, we will only search for constructor references if the containing class is annotated with
+   *
+   * @Builder, @Wither or @With
+   */
+  private boolean shouldLookForConstructorReferences(PsiMethod constructor) {
+    if (!constructor.isConstructor()) {
+      return false;
+    }
+
+    PsiClass containingClass = constructor.getContainingClass();
+    if (containingClass == null) {
+      return false;
+    }
+
+    return containingClass.getAnnotation(Builder.class.getName()) != null
+      || containingClass.getAnnotation(With.class.getName()) != null
+      || containingClass.getAnnotation(Wither.class.getName()) != null;
+  }
+
+  private void findConstructorReferencesInInnerClasses(PsiClass containingClass, PsiMethod constructor, Processor<? super PsiReference> consumer) {
+    for (PsiClass clazz : containingClass.getInnerClasses()) {
+      findConstructorReferencesInClassMethods(clazz, constructor, consumer);
+    }
+  }
+
+  private boolean findConstructorReferencesInClassMethods(PsiClass clazz, PsiMethod constructor, Processor<? super PsiReference> consumer) {
+    for (PsiMethod method : clazz.getMethods()) {
+      if (method instanceof LombokLightMethodBuilder) {
+        // only look in the methods we have generated ourselves. Other methods are already indexed and the references they contain can
+        // already be found by the native ReferenceSearch
+        if (!reportConstructorReferencesInElement(method, constructor, consumer)) {
+          return false;
+        }
+      }
+    }
+    return true;
+  }
+
+  private boolean reportConstructorReferencesInElement(PsiElement haystack, PsiMethod needle, Processor<? super PsiReference> consumer) {
+    if (haystack instanceof PsiNewExpression) {
+      PsiNewExpression newExpression = (PsiNewExpression) haystack;
+      PsiMethod resolvedConstructor = newExpression.resolveConstructor();
+      if (resolvedConstructor == needle) {
+        consumer.process(new LombokConstructorReference<>(
+          newExpression,
+          new TextRange(3, 4),
+          resolvedConstructor
+        ));
+      }
+    }
+
+    PsiElement[] children = haystack.getChildren();
+    for (PsiElement child : children) {
+      if (!reportConstructorReferencesInElement(child, needle, consumer)) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+
+  private static class LombokConstructorReference<T extends PsiElement> extends PsiReferenceBase<T> {
+    private final PsiMethod resolved;
+
+    public LombokConstructorReference(@NotNull T element, TextRange rangeInElement, PsiMethod resolved) {
+      super(element, rangeInElement);
+      this.resolved = resolved;
+    }
+
+    @Nullable
+    @Override
+    public PsiElement resolve() {
+      return resolved;
+    }
+  }
+}

src/main/java/de/plushnikov/intellij/plugin/extension/LombokReferenceSearcher.java

@@ -1,22 +1,16 @@
 package de.plushnikov.intellij.plugin.extension;
 
 import com.intellij.openapi.application.QueryExecutorBase;
-import com.intellij.openapi.project.DumbService;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiField;
+import com.intellij.psi.PsiMethod;
 import com.intellij.psi.PsiReference;
-import com.intellij.psi.search.SearchRequestCollector;
-import com.intellij.psi.search.UsageSearchContext;
 import com.intellij.psi.search.searches.ReferencesSearch;
 import com.intellij.util.Processor;
-import de.plushnikov.intellij.plugin.psi.LombokLightFieldBuilder;
 import de.plushnikov.intellij.plugin.psi.LombokLightMethodBuilder;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.Arrays;
-import java.util.Objects;
-
 /**
  * Annas example: org.jetbrains.plugins.javaFX.fxml.refs.JavaFxControllerFieldSearcher
  * Alternative Implementation for LombokFieldFindUsagesHandlerFactory
@@ -28,46 +22,48 @@ public LombokReferenceSearcher() {
   }
 
   @Override
-  public void processQuery(@NotNull ReferencesSearch.SearchParameters queryParameters, @NotNull Processor consumer) {
+  public void processQuery(@NotNull ReferencesSearch.SearchParameters queryParameters, @NotNull Processor<? super PsiReference> consumer) {
     PsiElement refElement = queryParameters.getElementToSearch();
 
     if (refElement instanceof PsiField) {
-      DumbService.getInstance(queryParameters.getProject()).runReadActionInSmartMode(() ->
-        processPsiField((PsiField) refElement, queryParameters.getOptimizer()));
+      searchInContainingClass((PsiField) refElement, consumer);
     }
   }
 
-  private void processPsiField(final PsiField refPsiField, final SearchRequestCollector collector) {
+  private void searchInContainingClass(final PsiField refPsiField, Processor<? super PsiReference> consumer) {
     final PsiClass containingClass = refPsiField.getContainingClass();
     if (null != containingClass) {
-      processClassMethods(refPsiField, collector, containingClass);
-
-      final PsiClass[] innerClasses = containingClass.getInnerClasses();
-      Arrays.stream(innerClasses)
-        .forEach(psiClass -> processClassMethods(refPsiField, collector, psiClass));
+      boolean mayContinueSearching = searchInClassMethods(containingClass, refPsiField, consumer);
 
-      Arrays.stream(innerClasses)
-        .forEach(psiClass -> processClassFields(refPsiField, collector, psiClass));
+      // TODO : look in generated inner classes methods (like Builders)
     }
   }
 
-  private void processClassMethods(PsiField refPsiField, SearchRequestCollector collector, PsiClass containingClass) {
-    Arrays.stream(containingClass.getMethods())
-      .filter(LombokLightMethodBuilder.class::isInstance)
-      .filter(psiMethod -> psiMethod.getNavigationElement() == refPsiField)
-      .forEach(psiMethod -> {
-        collector.searchWord(psiMethod.getName(), psiMethod.getUseScope(), UsageSearchContext.IN_CODE, true, psiMethod);
-      });
+  private boolean searchInClassMethods(PsiClass containingClass, PsiElement element, Processor<? super PsiReference> consumer) {
+    for(PsiMethod method : containingClass.getMethods()){
+      if(method instanceof LombokLightMethodBuilder){
+        // only look in the methods we have generated ourselves. Other methods are already indexed and the references they contain can
+        // already be found by the native ReferenceSearch
+        if(!reportReferencesInElement(method, element, consumer)){
+          return false;
+        }
+      }
+    }
+    return true;
   }
 
-  private void processClassFields(PsiField refPsiField, SearchRequestCollector collector, PsiClass containingClass) {
-    Arrays.stream(containingClass.getFields())
-      .filter(LombokLightFieldBuilder.class::isInstance)
-      .filter(psiField -> psiField.getNavigationElement() == refPsiField)
-      .filter(psiField -> Objects.nonNull(psiField.getName()))
-      .forEach(psiField -> {
-        collector.searchWord(psiField.getName(), psiField.getUseScope(), UsageSearchContext.IN_CODE, true, psiField);
-      });
-  }
+  private boolean reportReferencesInElement(PsiElement haystack, PsiElement needle, Processor<? super PsiReference> consumer) {
+    PsiReference ref = haystack.getReference();
+    if (ref != null && ref.isReferenceTo(needle)) {
+      return consumer.process(ref);
+    }
+    PsiElement[] children = haystack.getChildren();
+    for (PsiElement child : children) {
+      if(!reportReferencesInElement(child, needle, consumer)){
+        return false;
+      }
+    }
 
+    return true;
+  }
 }

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

@@ -255,7 +255,7 @@ public String renderSuperBuilderConstruction() {
   }
 
   public String renderBuildCall() {
-    return fieldInBuilderName;
+    return builderClass.getName() + ".this." + fieldInBuilderName;
   }
 
   public CharSequence renderToBuilderCall() {

src/main/java/de/plushnikov/intellij/plugin/processor/handler/singular/AbstractSingularHandler.java

@@ -39,8 +39,7 @@ public Collection<PsiField> renderBuilderFields(@NotNull BuilderInfo info) {
     return Collections.singleton(
       new LombokLightFieldBuilder(info.getManager(), info.getFieldName(), builderFieldType)
         .withContainingClass(info.getBuilderClass())
-        .withModifier(PsiModifier.PRIVATE)
-        .withNavigationElement(info.getVariable()));
+        .withModifier(PsiModifier.PRIVATE));
   }
 
   @NotNull

src/main/java/de/plushnikov/intellij/plugin/processor/handler/singular/NonSingularHandler.java

@@ -8,6 +8,7 @@
 import de.plushnikov.intellij.plugin.psi.LombokLightFieldBuilder;
 import de.plushnikov.intellij.plugin.psi.LombokLightMethodBuilder;
 import de.plushnikov.intellij.plugin.util.PsiMethodUtil;
+import lombok.NonNull;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -24,13 +25,12 @@ public Collection<PsiField> renderBuilderFields(@NotNull BuilderInfo info) {
     return Collections.singleton(
       new LombokLightFieldBuilder(info.getManager(), info.getFieldName(), info.getFieldType())
         .withContainingClass(info.getBuilderClass())
-        .withModifier(PsiModifier.PRIVATE)
-        .withNavigationElement(info.getVariable()));
+        .withModifier(PsiModifier.PRIVATE));
   }
 
   @Override
   public Collection<PsiMethod> renderBuilderMethod(@NotNull BuilderInfo info) {
-    final String blockText = getAllMethodBody(info.getFieldName(), info.getBuilderChainResult());
+    final String blockText = getAllMethodBody(info);
     final LombokLightMethodBuilder methodBuilder = new LombokLightMethodBuilder(info.getManager(), info.getFieldName())
       .withContainingClass(info.getBuilderClass())
       .withMethodReturnType(info.getBuilderType())
@@ -51,8 +51,8 @@ public String createSingularName(PsiAnnotation singularAnnotation, String psiFie
     return psiFieldName;
   }
 
-  private String getAllMethodBody(@NotNull String psiFieldName, @NotNull String builderChainResult) {
-    final String codeBlockTemplate = "this.{0} = {0};\nreturn {1};";
-    return MessageFormat.format(codeBlockTemplate, psiFieldName, builderChainResult);
+  private String getAllMethodBody(@NonNull BuilderInfo info) {
+    final String codeBlockTemplate = "{2}.this.{0} = {0};\nreturn {1};";
+    return MessageFormat.format(codeBlockTemplate, info.getFieldName(), info.getBuilderChainResult(), info.getBuilderClass().getName());
   }
 }

src/main/java/de/plushnikov/intellij/plugin/processor/handler/singular/SingularGuavaMapHandler.java

@@ -33,8 +33,7 @@ public Collection<PsiField> renderBuilderFields(@NotNull BuilderInfo info) {
     return Collections.singleton(
       new LombokLightFieldBuilder(info.getManager(), info.getFieldName(), builderFieldKeyType)
         .withContainingClass(info.getBuilderClass())
-        .withModifier(PsiModifier.PRIVATE)
-        .withNavigationElement(info.getVariable()));
+        .withModifier(PsiModifier.PRIVATE));
   }
 
   @NotNull

src/main/java/de/plushnikov/intellij/plugin/processor/handler/singular/SingularGuavaTableHandler.java

@@ -35,8 +35,7 @@ public Collection<PsiField> renderBuilderFields(@NotNull BuilderInfo info) {
     return Collections.singleton(
       new LombokLightFieldBuilder(info.getManager(), info.getFieldName(), builderFieldKeyType)
         .withContainingClass(info.getBuilderClass())
-        .withModifier(PsiModifier.PRIVATE)
-        .withNavigationElement(info.getVariable()));
+        .withModifier(PsiModifier.PRIVATE));
   }
 
   @NotNull

src/main/java/de/plushnikov/intellij/plugin/processor/handler/singular/SingularMapHandler.java

@@ -48,12 +48,10 @@ public Collection<PsiField> renderBuilderFields(@NotNull BuilderInfo info) {
     return Arrays.asList(
       new LombokLightFieldBuilder(info.getManager(), info.getFieldName() + LOMBOK_KEY, builderFieldKeyType)
         .withContainingClass(info.getBuilderClass())
-        .withModifier(PsiModifier.PRIVATE)
-        .withNavigationElement(info.getVariable()),
+        .withModifier(PsiModifier.PRIVATE),
       new LombokLightFieldBuilder(info.getManager(), info.getFieldName() + LOMBOK_VALUE, builderFieldValueType)
         .withContainingClass(info.getBuilderClass())
-        .withModifier(PsiModifier.PRIVATE)
-        .withNavigationElement(info.getVariable()));
+        .withModifier(PsiModifier.PRIVATE));
   }
 
   @NotNull

src/main/java/de/plushnikov/intellij/plugin/psi/LombokLightFieldBuilder.java

@@ -4,13 +4,16 @@
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiIdentifier;
 import com.intellij.psi.PsiManager;
 import com.intellij.psi.PsiModifier;
 import com.intellij.psi.PsiType;
 import com.intellij.psi.impl.CheckUtil;
 import com.intellij.psi.impl.light.LightFieldBuilder;
 import com.intellij.psi.impl.light.LightModifierList;
+import com.intellij.psi.impl.source.DummyHolder;
+import com.intellij.psi.impl.source.DummyHolderFactory;
 import com.intellij.util.IncorrectOperationException;
 import de.plushnikov.intellij.plugin.icon.LombokIcons;
 import org.jetbrains.annotations.NonNls;
@@ -26,12 +29,14 @@ public class LombokLightFieldBuilder extends LightFieldBuilder {
   private String myName;
   private final LombokLightIdentifier myNameIdentifier;
   private final LombokLightModifierList myModifierList;
+  private final DummyHolder myHolder;
 
   public LombokLightFieldBuilder(@NotNull PsiManager manager, @NotNull String name, @NotNull PsiType type) {
     super(manager, name, type);
     myName = name;
     myNameIdentifier = new LombokLightIdentifier(manager, name);
     myModifierList = new LombokLightModifierList(manager, JavaLanguage.INSTANCE, Collections.emptyList());
+    myHolder = DummyHolderFactory.createHolder(manager, getContext());
     setBaseIcon(LombokIcons.FIELD_ICON);
   }
 
@@ -77,6 +82,11 @@ public LombokLightFieldBuilder withNavigationElement(PsiElement navigationElemen
     return this;
   }
 
+  @Override
+  public PsiFile getContainingFile() {
+    return myHolder;
+  }
+
   @NotNull
   @Override
   public String getName() {