CheckBoxTree.java

package jetbrains.mps.build.mps.pluginSolution.plugin;

/*Generated by MPS */

import javax.swing.JPanel;
import javax.swing.JTree;
import java.util.Set;
import jetbrains.mps.internal.collections.runtime.SetSequence;
import java.util.HashSet;
import java.awt.BorderLayout;
import com.intellij.ui.treeStructure.Tree;
import javax.swing.tree.TreeSelectionModel;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import com.intellij.ui.ScrollPaneFactory;
import javax.swing.tree.TreePath;

public class CheckBoxTree<N extends NodeData> extends JPanel {
  private final JTree myTree;
  private final Set<N> mySelectedItems = SetSequence.fromSet(new HashSet<N>());
  public CheckBoxTree(CheckBoxNode node) {
    super(new BorderLayout());
    this.myTree = new Tree(node);
    this.myTree.setCellRenderer(new CheckBoxCellRenderrer());
    this.myTree.getSelectionModel().setSelectionMode(TreeSelectionModel.SINGLE_TREE_SELECTION);
    this.myTree.addMouseListener(new MouseAdapter() {
      @Override
      public void mousePressed(MouseEvent e) {
        CheckBoxTree.this.mouseParessed(e.getX(), e.getY());
      }
    });
    this.add(ScrollPaneFactory.createScrollPane(myTree), BorderLayout.CENTER);
  }
  public void mouseParessed(int x, int y) {
    int row = this.myTree.getRowForLocation(x, y);
    if (row == 0) {
      this.myTree.revalidate();
      this.myTree.repaint();
    }
    TreePath path = this.myTree.getPathForRow(row);
    if (path == null) {
      return;
    }
    Object lastPathComponent = path.getLastPathComponent();
    CheckBoxNode<N> cbNode = (CheckBoxNode) lastPathComponent;
    boolean isChecked = this.isChecked(cbNode);
    this.checkNodeRecursively(cbNode, !(isChecked));
    if (isChecked) {
      this.uncheckParents(cbNode);
    }
    this.repaint();
  }
  public void checkNodeRecursively(CheckBoxNode<N> checkBoxNode, boolean check) {
    this.checkNode(checkBoxNode, check);
    int childCount = checkBoxNode.getChildCount();
    for (int i = 0; i < childCount; i++) {
      this.checkNodeRecursively((CheckBoxNode<N>) checkBoxNode.getChildAt(i), check);
    }
  }
  public boolean isChecked(CheckBoxNode<N> cbNode) {
    return SetSequence.fromSet(CheckBoxTree.this.mySelectedItems).contains(cbNode.getData());
  }
  public void uncheckParents(CheckBoxNode<N> cbNode) {
    CheckBoxNode<N> parent = (CheckBoxNode<N>) cbNode.getParent();
    if (parent == null) {
      return;
    }
    this.checkNode(parent, false);
    this.uncheckParents(parent);
  }
  private void checkNode(CheckBoxNode<N> checkBoxNode, boolean check) {
    if (check) {
      SetSequence.fromSet(CheckBoxTree.this.mySelectedItems).addElement(checkBoxNode.getData());
    } else {
      SetSequence.fromSet(CheckBoxTree.this.mySelectedItems).removeElement(checkBoxNode.getData());
    }
    checkBoxNode.setChecked(check);
  }
  public Set<N> getSelectedItems() {
    return SetSequence.fromSetWithValues(new HashSet<N>(), this.mySelectedItems);
  }
}