/*
 * 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.
 */
package com.intellij.ui.tree;

import com.intellij.openapi.actionSystem.DataContext;
import com.intellij.openapi.actionSystem.PlatformDataKeys;
import com.intellij.openapi.ide.SmartSelectProvider;
import org.jetbrains.annotations.Nullable;

import javax.swing.*;
import javax.swing.tree.TreePath;
import java.awt.*;

/**
 * @author Konstantin Bulenkov
 */
public class TreeSmartSelectProvider implements SmartSelectProvider<JTree> {
  @Override
  public void increaseSelection(JTree tree) {
    TreePath path = tree.getLeadSelectionPath();
    if (path == null) return;
    increaseSelection(path, tree);
  }

  private static void increaseSelection(TreePath path, JTree tree) {
    TreePath parentPath = path.getParentPath();

    if (parentPath == null) {
      return;
    }

    boolean madeSelection = false;

    for (int i = 0; i < tree.getRowCount(); i++) {
      TreePath row = tree.getPathForRow(i);
      if (parentPath.isDescendant(row) && !row.equals(parentPath)) {
        if (!tree.isRowSelected(i)) {
          madeSelection = true;
          addSelection(tree, row);
        }
      }
    }

    if (!madeSelection) {
      if (tree.isRowSelected(tree.getRowForPath(parentPath))) {
        increaseSelection(parentPath, tree);
      } else {
        addSelection(tree, parentPath);
      }
    }
  }

  private static void addSelection(JTree tree, TreePath path) {
    TreePath leadPath = tree.getLeadSelectionPath();
    tree.getSelectionModel().addSelectionPath(path);
    tree.setLeadSelectionPath(leadPath);
  }

  @Override
  public void decreaseSelection(JTree tree) {
    //todo[kb]
  }

  @Nullable
  @Override
  public JTree getSource(DataContext context) {
    Component component = PlatformDataKeys.CONTEXT_COMPONENT.getData(context);
    return component instanceof JTree ? (JTree)component : null;
  }
}

@@ -909,16 +909,44 @@
 a.awt.*;
+%0Aimport java.util.ArrayList;
 %0A%0A/**%0A *
@@ -2222,18 +2222,1544 @@
     
-//todo%5Bkb%5D
+TreePath%5B%5D paths = tree.getSelectionPaths();%0A    TreePath leadSelection = tree.getLeadSelectionPath();%0A%0A    if (paths == null %7C%7C paths.length %3C 2) return;%0A    Object%5B%5D selected = leadSelection.getPath();%0A    if (selected.length %3C 2) return;%0A%0A    int i = 0;%0A    ArrayList%3CTreePath%3E toRemove = new ArrayList%3C%3E();%0A    while (i %3C selected.length) %7B%0A      for (TreePath path : paths) %7B%0A        if (!hasCommonStart(path, leadSelection, i + 1)) %7B%0A          toRemove.add(path);%0A        %7D%0A      %7D%0A%0A      if (!toRemove.isEmpty()) %7B%0A        if (toRemove.size() == paths.length - 1 && tree.isPathSelected(leadSelection.getParentPath())) %7B%0A          tree.removeSelectionPath(leadSelection.getParentPath());%0A        %7D else %7B%0A          for (TreePath path : toRemove) %7B%0A            tree.removeSelectionPath(path);%0A          %7D%0A        %7D%0A%0A        tree.setLeadSelectionPath(leadSelection);%0A        return;%0A      %7D else %7B%0A        i++;%0A      %7D%0A    %7D%0A  %7D%0A%0A  private static boolean hasCommonStart(TreePath path, TreePath leadSelection, int commonStartLength) %7B%0A    Object%5B%5D pathObjects = path.getPath();%0A    if (pathObjects.length %3C commonStartLength) return false;%0A    Object%5B%5D leadSelectionObjects = leadSelection.getPath();%0A    for (int i = 0; i %3C pathObjects.length && i %3C leadSelectionObjects.length && i %3C commonStartLength; i++) %7B%0A      if (pathObjects%5Bi%5D == leadSelectionObjects%5Bi%5D) %7B%0A        continue;%0A      %7D%0A%0A      if (pathObjects%5Bi%5D == null %7C%7C !pathObjects%5Bi%5D.equals(leadSelectionObjects%5Bi%5D)) %7B%0A        return false;%0A      %7D%0A    %7D%0A    return true;
 %0A  %7D

support for Smart Unselect in trees