IDEA-CR-36670: refactoring after review

platform/platform-api/src/com/intellij/util/ui/tree/TreeUtil.java

@@ -121,22 +121,14 @@ public final class TreeUtil {
     return null;
   }
 
+  /**
+   * @param tree a tree, which selection is processed
+   * @param type a {@code Class} object to filter selected user objects
+   * @return a list of user objects of the specified type retrieved from all selected paths
+   */
   @NotNull
-  public static <T> List<T> collectSelectedObjectsOfType(@NotNull JTree tree, @NotNull Class<T> clazz) {
-    final TreePath[] selections = tree.getSelectionPaths();
-    if (selections != null) {
-      final ArrayList<T> result = new ArrayList<>();
-      for (TreePath selection : selections) {
-        final DefaultMutableTreeNode node = (DefaultMutableTreeNode)selection.getLastPathComponent();
-        final Object userObject = node.getUserObject();
-        if (clazz.isInstance(userObject)) {
-          //noinspection unchecked
-          result.add((T)userObject);
-        }
-      }
-      return result;
-    }
-    return Collections.emptyList();
+  public static <T> List<T> collectSelectedObjectsOfType(@NotNull JTree tree, @NotNull Class<T> type) {
+    return collectSelectedObjects(tree, path -> getLastUserObject(type, path));
   }
 
   /**
@@ -933,8 +925,7 @@ public final class TreeUtil {
   }
 
   /**
-   * @param tree   a tree, which selection is processed
-   * @param mapper a function to convert a selected tree path to a corresponding object
+   * @param tree a tree, which selection is processed
    * @return a list of all selected paths
    */
   @NotNull
@@ -943,8 +934,7 @@ public final class TreeUtil {
   }
 
   /**
-   * @param tree   a tree, which selection is processed
-   * @param mapper a function to convert a selected tree path to a corresponding object
+   * @param tree a tree, which selection is processed
    * @return a list of user objects which correspond to all selected paths
    */
   @NotNull

platform/platform-tests/testSrc/com/intellij/ui/tree/TreeUtilVisitTest.java

@@ -790,6 +790,29 @@ public final class TreeUtilVisitTest {
                     consumer);
   }
 
+  @Test
+  public void testCollectSelectedObjectsOfType() {
+    TreeTest.test(() -> node(Boolean.TRUE, node(101), node(1.1f)), test
+      -> test.assertTree("+true\n", ()
+      -> TreeUtil.expandAll(test.getTree(), ()
+      -> test.assertTree("-true\n 101\n 1.1\n", ()
+      -> {
+      TreeUtil.visitVisibleRows(test.getTree(), path -> {
+        test.getTree().addSelectionPath(path);
+        return TreeVisitor.Action.CONTINUE;
+      });
+      test.assertTree("-[true]\n [101]\n [1.1]\n", true, () -> {
+        Assert.assertEquals(3, TreeUtil.collectSelectedObjectsOfType(test.getTree(), Object.class).size());
+        Assert.assertEquals(2, TreeUtil.collectSelectedObjectsOfType(test.getTree(), Number.class).size());
+        Assert.assertEquals(1, TreeUtil.collectSelectedObjectsOfType(test.getTree(), Boolean.class).size());
+        Assert.assertEquals(1, TreeUtil.collectSelectedObjectsOfType(test.getTree(), Integer.class).size());
+        Assert.assertEquals(1, TreeUtil.collectSelectedObjectsOfType(test.getTree(), Float.class).size());
+        Assert.assertEquals(0, TreeUtil.collectSelectedObjectsOfType(test.getTree(), String.class).size());
+        test.done();
+      });
+    }))));
+  }
+
   @Test
   public void testSelectFirstEmpty() {
     testSelectFirst(() -> null, true, "");