all class inheritors optimisation (do not try to analyze the base class concurrently) and proper handling of PCE (do not lose the element currently being analyzed)

java/java-indexing-impl/src/com/intellij/psi/impl/search/JavaClassInheritorsSearcher.java

@@ -105,25 +105,19 @@ public class JavaClassInheritorsSearcher extends QueryExecutorBase<PsiClass, Cla
 
   @NotNull
   private static Iterable<PsiClass> getOrComputeSubClasses(@NotNull Project project, @NotNull PsiClass baseClass) {
-    ConcurrentMap<PsiClass, Iterable<PsiClass>> CACHE = HighlightingCaches.getInstance(project).ALL_SUB_CLASSES;
-    Iterable<PsiClass> cached = CACHE.get(baseClass);
+    ConcurrentMap<PsiClass, Iterable<PsiClass>> map = HighlightingCaches.getInstance(project).ALL_SUB_CLASSES;
+    Iterable<PsiClass> cached = map.get(baseClass);
     if (cached == null) {
-      cached = computeAllSubClasses(project, baseClass); // it's almost empty now, no big deal
+      // returns lazy collection of subclasses. Each call to next() leads to calculation of next batch of subclasses.
+      Iterable<PsiClass> computed = new AllSubClassesLazyCollection(project, baseClass); // it's empty now, no big deal
       // for non-physical elements ignore the cache completely because non-physical elements created so often/unpredictably so I can't figure out when to clear caches in this case
-      if (ApplicationManager.getApplication().runReadAction((Computable<Boolean>)baseClass::isPhysical)) {
-        // make sure concurrent calls of this method always return the same collection to avoid expensive duplicate work
-        cached = ConcurrencyUtil.cacheOrGet(CACHE, baseClass, cached);
-      }
+      boolean isPhysical = ApplicationManager.getApplication().runReadAction((Computable<Boolean>)baseClass::isPhysical);
+      // make sure concurrent calls of this method always return the same collection to avoid expensive duplicate work
+      cached = isPhysical ? ConcurrencyUtil.cacheOrGet(map, baseClass, computed) : computed;
     }
     return cached;
   }
 
-  @NotNull
-  // returns lazy collection of subclasses. Each call to next() leads to calculation of next batch of subclasses.
-  private static Iterable<PsiClass> computeAllSubClasses(@NotNull Project project, @NotNull PsiClass baseClass) {
-    return new AllSubClassesLazyCollection(project, baseClass);
-  }
-
   private static boolean processLocalScope(@NotNull final Project project,
                                            @NotNull final ClassInheritorsSearch.SearchParameters parameters,
                                            @NotNull LocalSearchScope searchScope,
@@ -198,33 +192,36 @@ public class JavaClassInheritorsSearcher extends QueryExecutorBase<PsiClass, Cla
     // Already computed subclasses are stored in this collection.
     // There are two iterators maintained for this collection:
     // - 'candidatesToFindSubclassesIterator' points to the next element for which direct inheritors haven't been searched yet.
-    // - 'subClassIterator' created in AllSubClassesLazyCollection.iterator() maintains state of the AllSubClassesLazyCollection iterator in a lazy fashion. If more elements requested for this iterator, the findNextSubclasses() is called which tries to populate 'subClasses' with more inheritors.
-    private final HashSetQueue<PsiAnchor> subClasses = new HashSetQueue<>();
+    // - 'subClassIterator' created in AllSubClassesLazyCollection.iterator() maintains state of the AllSubClassesLazyCollection iterator in a lazy fashion. If more elements requested for this iterator, the processMoreSubclasses() is called which tries to populate 'subClasses' with more inheritors.
+    private final HashSetQueue<PsiAnchor> subClasses = new HashSetQueue<>(); // guarded by lock
     private final Object lock = new Object();
     private final GlobalSearchScope projectScope;
     private final Semaphore currentlyProcessingClasses = new Semaphore();
-    private final PsiClass myBaseClass;
+
+    private final HashSetQueue.ResettableIterator<PsiAnchor> candidatesToFindSubclassesIterator = subClasses.iterator(); // guarded by lock
 
     AllSubClassesLazyCollection(@NotNull Project project, @NotNull PsiClass baseClass) {
-      myBaseClass = baseClass;
       projectScope = GlobalSearchScope.allScope(project);
-      // populate with at least one subclass
-      findNextSubclasses();
+      subClasses.add(ApplicationManager.getApplication().runReadAction((Computable<PsiAnchor>)() -> PsiAnchor.create(baseClass)));
     }
 
     @NotNull
     @Override
     public Iterator<PsiClass> iterator() {
       return new Iterator<PsiClass>() {
-        private final Iterator<PsiAnchor> subClassIterator = subClasses.iterator();
-
+        private final Iterator<PsiAnchor> subClassIterator = subClasses.iterator(); // guarded by lock
+        {
+          synchronized (lock) {
+            subClassIterator.next(); //skip the baseClass which stored in the subClasses first element
+          }
+        }
         @Override
         public boolean hasNext() {
           synchronized (lock) {
             if (subClassIterator.hasNext()) return true;
           }
 
-          findNextSubclasses();
+          processMoreSubclasses();
 
           synchronized (lock) {
             return subClassIterator.hasNext();
@@ -242,32 +239,29 @@ public class JavaClassInheritorsSearcher extends QueryExecutorBase<PsiClass, Cla
       };
     }
 
-    private Iterator<PsiAnchor> candidatesToFindSubclassesIterator; // guarded by lock
-
-    private void findNextSubclasses() {
+    // polls 'subClasses' for more sub classes and call DirectClassInheritorsSearch for them
+    private void processMoreSubclasses() {
       while (true) {
         ProgressManager.checkCanceled();
 
         PsiAnchor next;
+        Object markedPosition;
         synchronized (lock) {
-          if (candidatesToFindSubclassesIterator == null) {
-            candidatesToFindSubclassesIterator = subClasses.iterator();
-            next = null;
+          if (!candidatesToFindSubclassesIterator.hasNext()) {
+            // no candidates left, exit
+            // but first, wait for other threads to process their candidates
+            break;
           }
-          else {
-            if (!candidatesToFindSubclassesIterator.hasNext()) {
-              // no candidates left, exit
-              // but first, wait for other threads to process their candidates
-              break;
-            }
-            next = candidatesToFindSubclassesIterator.next();
-          }
-          currentlyProcessingClasses.down(); // tell other threads we are going to process this candidate
+          // in case of a sudden PCE thrown in DirectClassInheritorsSearch below, we have to re-analyze this candidate later
+          // to do that, we mark this current position in subClasses list and revert to this position when PCE throws.
+          // Note that in time of reverting, other threads might have advanced this iterator for a number of candidates forward but it's ok if we re-analyze them all.
+          markedPosition = candidatesToFindSubclassesIterator.markPosition();
+          next = candidatesToFindSubclassesIterator.next();
         }
-
+        currentlyProcessingClasses.down(); // tell other threads we are going to process something
         boolean added;
         try {
-          PsiClass candidate = next == null ? myBaseClass : ReadAction.compute(() -> (PsiClass)next.retrieve());
+          PsiClass candidate = ReadAction.compute(() -> (PsiClass)next.retrieve());
 
           if (candidate == null || candidate instanceof PsiAnonymousClass || isFinal(candidate)) {
             added = false;
@@ -282,11 +276,19 @@ public class JavaClassInheritorsSearcher extends QueryExecutorBase<PsiClass, Cla
             });
           }
         }
+        catch (Exception e) {
+          synchronized (lock) {
+            // revert back to the position before analyzing "next" candidate
+            // in case of multiple passes through the "next" candidate, we might try to add the same classes several times but nothing bad will happen thanks to the "subClasses" being a Set
+            candidatesToFindSubclassesIterator.resetPosition(markedPosition);
+          }
+          throw e;
+        }
         finally {
           currentlyProcessingClasses.up();
         }
         if (added) {
-          // just allow the iterator to move forward; more elements will be added on the next call to .next()
+          // just allow the iterator to move forward; more elements will be added on the subsequent call to .next()
           return;
         }
       }

platform/util/src/com/intellij/util/containers/HashSetQueue.java

@@ -153,9 +153,10 @@ public class HashSetQueue<T> extends AbstractCollection<T> implements Queue<T> {
 
   @NotNull
   @Override
-  public Iterator<T> iterator() {
-    return new Iterator<T>() {
+  public ResettableIterator<T> iterator() {
+    return new ResettableIterator<T>() {
       private QueueEntry<T> cursor = TOMB;
+      private long count;
       @Override
       public boolean hasNext() {
         return cursor.next != TOMB;
@@ -164,6 +165,7 @@ public class HashSetQueue<T> extends AbstractCollection<T> implements Queue<T> {
       @Override
       public T next() {
         cursor = cursor.next;
+        count++;
         return cursor.t;
       }
 
@@ -172,6 +174,40 @@ public class HashSetQueue<T> extends AbstractCollection<T> implements Queue<T> {
         if (cursor == TOMB) throw new NoSuchElementException();
         HashSetQueue.this.remove(cursor.t);
       }
+
+      @Override
+      public Object markPosition() {
+        return new IteratorPosition<T>(cursor, count);
+      }
+
+      @Override
+      public boolean resetPosition(Object p) {
+        @SuppressWarnings("unchecked")
+        IteratorPosition<T> requested = (IteratorPosition<T>)p;
+
+        if (requested.count <= count) {
+          cursor = requested.cursor;
+          count = requested.count;
+          return true;
+        }
+        return false;
+      }
     };
   }
+
+  private static class IteratorPosition<T> {
+    private final QueueEntry<T> cursor;
+    private final long count;
+
+    IteratorPosition(@NotNull QueueEntry<T> cursor, long count) {
+      this.cursor = cursor;
+      this.count = count;
+    }
+  }
+
+  public interface ResettableIterator<T> extends Iterator<T> {
+    Object markPosition();
+    // returns true if reset successfully, false if failed (e.g. the requested position is ahead of current)
+    boolean resetPosition(Object pos);
+  }
 }

platform/util/testSrc/com/intellij/util/containers/HashSetQueueTest.java

@@ -25,7 +25,7 @@ import java.util.Iterator;
 
 public class HashSetQueueTest extends TestCase {
   private final Assertion CHECK = new Assertion();
-  private final HashSetQueue<String> myQueue = new HashSetQueue<String>();
+  private final HashSetQueue<String> myQueue = new HashSetQueue<>();
 
   public void testEmpty() {
     assertEquals(0, myQueue.size());
@@ -45,7 +45,7 @@ public class HashSetQueueTest extends TestCase {
     assertEquals(1, myQueue.size());
     myQueue.add("3");
     assertEquals(2, myQueue.size());
-    CHECK.compareAll(new Object[]{"2", "3"}, new ArrayList<String>(myQueue));
+    CHECK.compareAll(new Object[]{"2", "3"}, new ArrayList<>(myQueue));
     assertEquals("2", myQueue.poll());
     assertEquals("3", myQueue.poll());
     testEmpty();
@@ -128,4 +128,28 @@ public class HashSetQueueTest extends TestCase {
     assertEquals("2", iterator.next());
     assertFalse(iterator.hasNext());
   }
+
+  public void testResettableIterator() {
+    assertTrue(myQueue.add("1"));
+    HashSetQueue.ResettableIterator<String> iterator = myQueue.iterator();
+    Object position = iterator.markPosition();
+    assertTrue(iterator.hasNext());
+    assertEquals("1", iterator.next());
+    assertFalse(iterator.hasNext());
+    Object pos2 = iterator.markPosition();
+
+    boolean reset = iterator.resetPosition(position);
+    assertTrue(reset);
+
+    assertTrue(iterator.hasNext());
+    assertEquals("1", iterator.next());
+    assertFalse(iterator.hasNext());
+
+    reset = iterator.resetPosition(position);
+    assertTrue(reset);
+    assertTrue(iterator.hasNext());
+
+    boolean reset2 = iterator.resetPosition(pos2);
+    assertFalse(reset2);
+  }
 }