/*
 * 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.
 */
package com.jetbrains.python.psi;

import com.google.common.collect.FluentIterable;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiNamedElement;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.extenstions.PsiElementExtKt;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * JQuery-like tool that makes PSI navigation easier.
 * You just "drive" your query filtering results, and no need to check result for null.
 * There are 3 axes: ancestors, descendants and siblings. Each axis has N elements which you can filter or map.
 *
 * @param <T> type of current elements
 * @author Ilya.Kazakevich
 */
public final class PsiQuery<T extends PsiElement> {

  @NotNull
  private final List<T> myElements;


  public PsiQuery(@NotNull final T elementToStart) {
    this(Collections.singletonList(elementToStart));
  }


  @SuppressWarnings("MethodCanBeVariableArityMethod") // To prevent type safety
  public PsiQuery(@NotNull final T[] elementsToStart) {
    this(Arrays.asList(elementsToStart));
  }

  public PsiQuery(@NotNull final List<T> elementsToStart) {
    myElements = Collections.unmodifiableList(elementsToStart);
  }

  /**
   * Adds elements to query
   */
  @NotNull
  public PsiQuery<T> addElements(@NotNull final PsiQuery<? extends T> newElements) {
    final List<T> newList = new ArrayList<>(myElements);
    newList.addAll(newElements.myElements);
    return new PsiQuery<>(newList);
  }


  @NotNull
  public static <T extends PsiElement> PsiQuery<T> create(@NotNull final T elementToStart) {
    return new PsiQuery<>(elementToStart);
  }

  @NotNull
  public static <T extends PsiElement> PsiQuery<T> createFromQueries(@NotNull final List<PsiQuery<? extends T>> queriesWithElements) {
    final Set<T> result = new LinkedHashSet<>();
    queriesWithElements.forEach(o -> result.addAll(o.getElements()));
    return new PsiQuery<>(new ArrayList<>(result));
  }

  @NotNull
  public <R extends T> PsiQuery<R> filter(@NotNull final PsiFilter<R> filter) {
    return new PsiQuery<>(filter.filter(myElements));
  }

  @NotNull
  public <R extends T> PsiQuery<R> filter(@NotNull final Class<R> filterToType) {
    return filter(new PsiFilter<>(filterToType));
  }

  @NotNull
  public PsiQuery<T> filter(@NotNull final Predicate<T> filter) {
    return new PsiQuery<>(asStream().filter(filter).collect(Collectors.toList()));
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> map(@NotNull final Function<T, R> map) {
    return new PsiQuery<>(asStream().map(map).collect(Collectors.toList()));
  }

  @NotNull
  public <R extends PsiElement, F_R extends T> PsiQuery<R> map(@NotNull final PsiFilter<F_R> preFilter,
                                                               @NotNull final Function<F_R, R> map) {
    return filter(preFilter).map(map);
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> descendants(@NotNull final Class<R> type) {
    return descendants(new PsiFilter<>(type));
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> ancestors(@NotNull final Class<R> type) {
    return ancestors(new PsiFilter<>(type));
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> siblings(@NotNull final Class<R> type) {
    return siblings(new PsiFilter<>(type));
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> descendants(@NotNull final PsiFilter<R> filter) {
    return getQueryWithProducer(o -> PsiTreeUtil.findChildrenOfType(o, filter.myClass), filter);
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> ancestors(@NotNull final PsiFilter<R> filter) {
    return getQueryWithProducer(o -> PsiElementExtKt.getAncestors(o, o.getContainingFile(), filter.myClass), filter);
  }

  @NotNull
  public <R extends PsiElement> PsiQuery<R> siblings(@NotNull final PsiFilter<R> filter) {
    return getQueryWithProducer(o -> PsiTreeUtil.findChildrenOfType(o.getParent(), filter.myClass), filter);
  }

  @NotNull
  private <R extends PsiElement> PsiQuery<R> getQueryWithProducer(@NotNull final Function<T, Collection<R>> elementsProducer,
                                                                  @NotNull final PsiFilter<R> filter) {
    final Set<R> result = new LinkedHashSet<>(); // Set to get rid of duplicates but preserve order
    asStream()
      .map(o -> elementsProducer.apply(o)
        .stream().filter(o2 -> !o2.equals(o)) // Filter out same element in case of siblings
        .collect(Collectors.toList()))
      .forEach(result::addAll);
    return new PsiQuery<>(new ArrayList<>(filter.filter(result)));
  }


  @NotNull
  public Stream<T> asStream() {
    return myElements.stream();
  }

  @NotNull
  public List<T> getElements() {
    return asStream().collect(Collectors.toList());
  }

  @Nullable
  public T getFirstElement() {
    final List<T> elements = getElements();
    return (elements.isEmpty() ? null : elements.get(0));
  }

  @Nullable
  public T getLastElement() {
    final List<T> elements = getElements();
    return (elements.isEmpty() ? null : elements.get(elements.size() - 1));
  }

  public boolean exists() {
    return !getElements().isEmpty();
  }

  public static class PsiFilter<T extends PsiElement> {
    public static final PsiFilter<PsiElement> ANY = new PsiFilter<>(PsiElement.class);
    @NotNull
    private final Class<T> myClass;
    @NotNull
    private final Predicate<T> myPredicate;

    /**
     * Include current element in result set, or not (false by default)
     */

    public PsiFilter(@NotNull final Class<T> aClass, @NotNull final Predicate<T> predicate) {
      myClass = aClass;
      myPredicate = predicate;
    }

    public PsiFilter(@NotNull final Class<T> aClass) {
      this(aClass, o -> true);
    }

    @NotNull
    List<T> filter(@NotNull final Collection<?> list) {
      // IDEA-165420
      //noinspection Guava
      return FluentIterable.from(list).filter(myClass).toList().stream().filter(myPredicate).collect(Collectors.toList());
    }
  }

  /**
   * Filter by name for {@link PsiNamedElement}
   */
  public static class PsiNameFilter<T extends PsiNamedElement> extends PsiFilter<T> {

    @NotNull
    public static PsiNameFilter<PsiNamedElement> create(@NotNull final String name) {
      return new PsiNameFilter<>(PsiNamedElement.class, name);
    }

    public PsiNameFilter(@NotNull final Class<T> type, @NotNull final String name) {
      super(type, o -> name.equals(o.getName()));
    }
  }
}

@@ -4626,36 +4626,35 @@
  -%3E PsiTreeUtil.
-find
+get
 ChildrenOfType(o
@@ -4651,16 +4651,22 @@
 enOfType
+AsList
 (o.getPa

Use one level (getChildrenOfTypeAsList) instead of recursive (findChildrenOfType)