platform/core-api/src/com/intellij/psi/search/PsiElementProcessor.java

/*
 * Copyright 2000-2013 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.
 */
package com.intellij.psi.search;

import com.intellij.psi.PsiElement;
import com.intellij.psi.util.PsiElementFilter;
import com.intellij.psi.util.PsiUtilCore;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @see com.intellij.psi.util.PsiTreeUtil#processElements(com.intellij.psi.PsiElement, PsiElementProcessor)
 */
public interface PsiElementProcessor<T extends PsiElement> {
  /**
   * Processes a PsiElement
   *
   * @param element currently processed element.
   * @return false to stop processing.
   */
  boolean execute(@NotNull T element);

  class CollectElements<T extends PsiElement> implements PsiElementProcessor<T> {
    private final Collection<T> myCollection;

    public CollectElements() {
      this(new ArrayList<T>());
    }

    public CollectElements(@NotNull Collection<T> collection) {
      myCollection = Collections.synchronizedCollection(collection);
    }

    @NotNull
    public PsiElement[] toArray() {
      return PsiUtilCore.toPsiElementArray(myCollection);
    }

    @NotNull
    public Collection<T> getCollection() {
      return myCollection;
    }

    @NotNull
    public T[] toArray(T[] array) {
      return myCollection.toArray(array);
    }

    @Override
    public boolean execute(@NotNull T element) {
      myCollection.add(element);
      return true;
    }
  }

  class CollectFilteredElements<T extends PsiElement> extends CollectElements<T> {
    private final PsiElementFilter myFilter;

    public CollectFilteredElements(@NotNull PsiElementFilter filter, @NotNull Collection<T> collection) {
      super(collection);
      myFilter = filter;
    }

    public CollectFilteredElements(@NotNull PsiElementFilter filter) {
      myFilter = filter;
    }

    @Override
    public boolean execute(@NotNull T element) {
      return !myFilter.isAccepted(element) || super.execute(element);
    }
  }

  class CollectElementsWithLimit<T extends PsiElement> extends CollectElements<T>{
    private final AtomicInteger myCount = new AtomicInteger(0);
    private volatile boolean myOverflow = false;
    private final int myLimit;

    public CollectElementsWithLimit(int limit) {
      myLimit = limit;
    }

    public CollectElementsWithLimit(int limit, @NotNull Collection<T> collection) {
      super(collection);
      myLimit = limit;
    }

    @Override
    public boolean execute(@NotNull T element) {
      if (myCount.get() == myLimit){
        myOverflow = true;
        return false;
      }
      myCount.incrementAndGet();
      return super.execute(element);
    }

    public boolean isOverflow() {
      return myOverflow;
    }
  }

  class FindElement<T extends PsiElement> implements PsiElementProcessor<T> {
    private volatile T myFoundElement = null;

    public boolean isFound() {
      return myFoundElement != null;
    }

    @Nullable
    public T getFoundElement() {
      return myFoundElement;
    }

    public boolean setFound(T element) {
      myFoundElement = element;
      return false;
    }

    @Override
    public boolean execute(@NotNull T element) {
      return setFound(element);
    }
  }

  class FindFilteredElement<T extends PsiElement> extends FindElement<T> {
    private final PsiElementFilter myFilter;

    public FindFilteredElement(@NotNull PsiElementFilter filter) {
      myFilter = filter;
    }

    @Override
    public boolean execute(@NotNull T element) {
      return !myFilter.isAccepted(element) || super.execute(element);
    }
  }
}