API to receive information about batches of completion items

platform/lang-api/src/com/intellij/codeInsight/completion/BatchConsumer.java

@@ -0,0 +1,15 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.codeInsight.completion;
+
+import com.intellij.util.Consumer;
+
+/**
+ * Pass an implementation of this interface to {@link CompletionService#performCompletion(CompletionParameters, Consumer)} to receive
+ * updates when each batch of completion items generated by a completion contributor is done.
+ *
+ * @author yole
+ */
+public interface BatchConsumer<T> extends Consumer<T> {
+  default void startBatch() {}
+  default void endBatch() {}
+}

platform/lang-api/src/com/intellij/codeInsight/completion/CompletionResultSet.java

@@ -1,18 +1,4 @@
-/*
- * Copyright 2000-2014 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.codeInsight.completion;
 
 import com.intellij.codeInsight.lookup.LookupElement;
@@ -71,6 +57,18 @@ public abstract class CompletionResultSet implements Consumer<LookupElement> {
     myConsumer.consume(result);
   }
 
+  public void startBatch() {
+    if (myConsumer instanceof BatchConsumer) {
+      ((BatchConsumer)myConsumer).startBatch();
+    }
+  }
+
+  public void endBatch() {
+    if (myConsumer instanceof BatchConsumer) {
+      ((BatchConsumer)myConsumer).endBatch();
+    }
+  }
+
   /**
    * Adds all elements from the given collection that match the prefix for further processing. The elements are processed in batch,
    * so that they'll appear in lookup all together.<p/>
@@ -82,6 +80,7 @@ public abstract class CompletionResultSet implements Consumer<LookupElement> {
    * instead of {@link #addElement(LookupElement)} helps to avoid that.
    */
   public void addAllElements(@NotNull final Iterable<? extends LookupElement> elements) {
+    startBatch();
     int seldomCounter = 0;
     for (LookupElement element : elements) {
       seldomCounter++;
@@ -90,6 +89,7 @@ public abstract class CompletionResultSet implements Consumer<LookupElement> {
         ProgressManager.checkCanceled();
       }
     }
+    endBatch();
   }
 
   @Contract(value="", pure=true)
@@ -146,7 +146,22 @@ public abstract class CompletionResultSet implements Consumer<LookupElement> {
     if (stop) {
       stopHere();
     }
-    myCompletionService.getVariantsFromContributors(parameters, myContributor, consumer);
+    myCompletionService.getVariantsFromContributors(parameters, myContributor, new BatchConsumer<CompletionResult>() {
+      @Override
+      public void startBatch() {
+        CompletionResultSet.this.startBatch();
+      }
+
+      @Override
+      public void endBatch() {
+        CompletionResultSet.this.endBatch();
+      }
+
+      @Override
+      public void consume(CompletionResult result) {
+        consumer.consume(result);
+      }
+    });
   }
 
   /**

platform/lang-api/src/com/intellij/codeInsight/completion/CompletionService.java

@@ -111,14 +111,32 @@ public abstract class CompletionService {
   /**
    * The main method that is invoked to collect all the completion variants
    * @param parameters Parameters specifying current completion environment
-   * @param consumer This consumer will directly add lookup elements to the lookup
+   * @param consumer The consumer of the completion variants. Pass an instance of {@link BatchConsumer} if you need to receive information
+   *                 about item batches generated by each completion contributor.
    */
   public void performCompletion(CompletionParameters parameters, Consumer<? super CompletionResult> consumer) {
     final Set<LookupElement> lookupSet = ContainerUtil.newConcurrentSet();
 
-    getVariantsFromContributors(parameters, null, result -> {
-      if (lookupSet.add(result.getLookupElement())) {
-        consumer.consume(result);
+    getVariantsFromContributors(parameters, null, new BatchConsumer<CompletionResult>() {
+      @Override
+      public void startBatch() {
+        if (consumer instanceof BatchConsumer) {
+          ((BatchConsumer<? super CompletionResult>)consumer).startBatch();
+        }
+      }
+
+      @Override
+      public void endBatch() {
+        if (consumer instanceof BatchConsumer) {
+          ((BatchConsumer<? super CompletionResult>)consumer).endBatch();
+        }
+      }
+
+      @Override
+      public void consume(CompletionResult result) {
+        if (lookupSet.add(result.getLookupElement())) {
+          consumer.consume(result);
+        }
       }
     });
   }

platform/lang-impl/src/com/intellij/codeInsight/template/impl/LiveTemplateCompletionContributor.java

@@ -1,18 +1,4 @@
-/*
- * Copyright 2000-2017 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.codeInsight.template.impl;
 
 import com.intellij.codeInsight.completion.*;