基于hutool树构建工具(TreeUtil)封装的搜索树工具类

单调树结构搜索实现

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Filter;
import cn.hutool.core.lang.func.Func1;
import cn.hutool.core.lang.func.LambdaUtil;
import cn.hutool.core.lang.tree.Tree;
import lombok.Getter;
import lombok.RequiredArgsConstructor;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 支持多参数检索，保留树层级结构，即匹配到子节点时父节点也保留
 * 
 * 下面是伪代码
 * <pre>{@code
 * List<Tree<Long>> source = new ArrayList<>();
 * List<Tree<Long>> result = TreeSearcher.of(source)
 * .applyEq("key1", "value1")
 * .applyEq("key2", 2)
 * .applyLike("key3", "value3")
 * .getResult();
 * }
 * </pre>
 *
 * @author Alone
 */
public class TreeSearcher<T, R> {

    private final List<Tree<T>> source;
    private final Map<String, SearchItem> searchMap;

    private TreeSearcher(List<Tree<T>> source) {
        this.source = source;
        this.searchMap = new HashMap<>(4);
    }

    @SuppressWarnings("unused")
    public static <T, R> TreeSearcher<T, R> of(List<Tree<T>> source, Class<R> targetClass) {
        return new TreeSearcher<>(source);
    }
    public static <T, R> TreeSearcher<T, R> of(List<Tree<T>> source) {
        return new TreeSearcher<>(source);
    }

    public TreeSearcher<T, R> applyEq(boolean condition, Func1<R, ?> key, Number data) {
        return apply(condition, LambdaUtil.getFieldName(key), data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyEq(Func1<R, ?> key, Number data) {
        return apply(true, LambdaUtil.getFieldName(key), data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyEq(Func1<R, ?> key, Object data) {
        return apply(true, LambdaUtil.getFieldName(key), data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyLike(Func1<R, ?> key, Object data) {
        return apply(true, LambdaUtil.getFieldName(key), data, SearchModel.LIKE);
    }

    public TreeSearcher<T, R> applyEq(boolean condition, Func1<R, ?> key, Object data) {
        return apply(condition, LambdaUtil.getFieldName(key), data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyLike(boolean condition, Func1<R, ?> key, Object data) {
        return apply(condition, LambdaUtil.getFieldName(key), data, SearchModel.LIKE);
    }

    public TreeSearcher<T, R> applyEq(boolean condition, String key, Number data) {
        return apply(condition, key, data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyEq(String key, Number data) {
        return apply(true, key, data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyEq(String key, Object data) {
        return apply(true, key, data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyLike(String key, Object data) {
        return apply(true, key, data, SearchModel.LIKE);
    }

    public TreeSearcher<T, R> applyEq(boolean condition, String key, Object data) {
        return apply(condition, key, data, SearchModel.EQ);
    }

    public TreeSearcher<T, R> applyLike(boolean condition, String key, Object data) {
        return apply(condition, key, data, SearchModel.LIKE);
    }

    private TreeSearcher<T, R> apply(boolean condition, String key, Object data, SearchModel model) {
        if (condition) {
            BizAssert.isNotBlank(key, "key不能空");
            SearchItem searchItem = new SearchItem(data, model);
            searchMap.put(key, searchItem);
        }
        return this;
    }

    public List<Tree<T>> getResult() {
        if (CollUtil.isEmpty(this.source)) {
            return Collections.emptyList();
        }
        if (CollUtil.isEmpty(searchMap)) {
            return this.source;
        }
        return search();
    }

    /**
     * hutool自带的搜索逻辑
     */
    public List<Tree<T>> getResult2() {
        if (CollUtil.isEmpty(this.source)) {
            return Collections.emptyList();
        }
        if (CollUtil.isEmpty(searchMap)) {
            return this.source;
        }
        return search2();
    }

    private List<Tree<T>> search2() {
        return this.source.stream()
                .map(tree -> tree.filter(match))
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    private List<Tree<T>> search() {
        return this.source.stream()
                .map(item -> doSearch(item, false))
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    private Tree<T> doSearch(Tree<T> tree, boolean pMatches) {
        if (pMatches) {
            return tree;
        }
        boolean currentMatches = match(tree);
        if (tree.hasChild()) {
            List<Tree<T>> children = tree.getChildren().stream()
                    .map(item -> doSearch(item, currentMatches))
                    .filter(Objects::nonNull)
                    .collect(Collectors.toList());
            tree.setChildren(children);
            if (CollUtil.isNotEmpty(children)) {
                return tree;
            }
        }
        if (currentMatches) {
            return tree;
        } else {
            return null;
        }
    }

    private final Filter<Tree<T>> match = this::match;
    private boolean match(Tree<T> tree) {
        return searchMap.keySet().stream().allMatch(key -> {
            SearchItem searchItem = searchMap.get(key);
            Object val = tree.get(key);
            if (Objects.isNull(val) && Objects.isNull(searchItem.getData())) {
                return false;
            }
            boolean res = false;
            if (val instanceof Number) {
                BizAssert.isTrue(SearchModel.EQ.equals(searchItem.getModel()), "数字类型值只支持等值比较");
                BizAssert.isTrue(searchItem.getData() instanceof Number, "数字类型值只支持匹配数字");
                res = val == searchItem.getData();
            } else {
                String valStr = val.toString();
                if (SearchModel.EQ.equals(searchItem.getModel())) {
                    res = valStr.trim().equals(searchItem.getData());
                }
                if (SearchModel.LIKE.equals(searchItem.getModel())) {
                    String regex = ".*" + searchItem.getData() + ".*";
                    res = valStr.matches(regex);
                }
            }
            return res;
        });
    }

    @SuppressWarnings("all")
    private enum SearchModel {
        EQ,
        LIKE
    }

    @Getter
    @RequiredArgsConstructor
    private static class SearchItem {
        private final Object data;
        private final SearchModel model;
    }

}

异构树结构搜索实现

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Pair;
import cn.hutool.core.lang.tree.Tree;
import com.google.common.base.Function;
import com.google.common.base.Predicate;
import com.google.common.base.Predicates;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;

/**
 * 异构多子树搜索实现，搜索逻辑与TreeSearcher一致
 *
 * @author Alone
 * @see TreeSearcher
 */
@SuppressWarnings({"rawtypes", "unchecked"})
abstract class MultiTreeSearcherBase<T, Children extends MultiTreeSearcherBase<T, Children>> {

    private final List<Tree<T>> source;

    private final Map<
            Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>>,
            Predicate<Tree>
            > matcherMap;

    /**
     * 主树的孩子节点的 getter setter
     */
    protected final Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> defaultKey;

    @SuppressWarnings("unchecked")
    protected final Children typedThis = (Children) this;

    protected MultiTreeSearcherBase(List<Tree<T>> source) {
        this.source = source;
        this.matcherMap = new HashMap<>(1);
        this.defaultKey = Pair.of(Tree::getChildren, Tree::setChildren);
    }

    protected Children eq(Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> key, Function<Tree, ?> valueFunc, Object value) {
        and(key, tree -> valueFunc.apply(tree).equals(value));
        return typedThis;
    }

    protected Children like(Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> key, Function<Tree, String> valueFunc, String value) {
        String regex = ".*" + value + ".*";
        and(key, tree -> valueFunc.apply(tree).matches(regex));
        return typedThis;
    }

    private void and(Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> key, Predicate<Tree>... predicate) {
        Predicate<Tree> matcher = matcherMap.get(key);
        if (Objects.isNull(matcher)) {
            matcher = Predicates.and(predicate);
        } else {
            matcher = Predicates.and(matcher, Predicates.and(predicate));
        }
        matcherMap.put(key, matcher);
    }

    public List<Tree<T>> getResult() {
        if (CollUtil.isEmpty(this.source)) {
            return Collections.emptyList();
        }
        if (CollUtil.isEmpty(matcherMap.values())) {
            return this.source;
        }
        return search();
    }

    private List<Tree<T>> search() {
        return this.source.stream()
                .map(item -> {
                    Tree<T> res = null;
                    for (Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> pair : matcherMap.keySet()) {
                        Tree tree = searchTree(item, pair, matcherMap.get(pair), false);
                        if (Objects.nonNull(tree)) {
                            res = tree;
                        }
                    }
                    return res;
                })
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    private Tree searchTree(Tree tree, Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> pair, Predicate<Tree> matcher, boolean pMatches) {
        if (pMatches) {
            return tree;
        }
        boolean currentMatches = matcher.test(tree);
        List<Tree> child = pair.getKey().apply(tree);
        if (CollUtil.isNotEmpty(child)) {
            child = child.stream()
                    .map(item -> searchTree(item, defaultKey, matcher, currentMatches))
                    .filter(Objects::nonNull)
                    .collect(Collectors.toList());
            pair.getValue().accept(tree, child);
            if (CollUtil.isNotEmpty(child)) {
                return tree;
            }
        }
        if (currentMatches) {
            return tree;
        } else {
            return null;
        }
    }
}

import cn.hutool.core.lang.Pair;
import cn.hutool.core.lang.func.Func1;
import cn.hutool.core.lang.func.LambdaUtil;
import cn.hutool.core.lang.tree.Tree;
import com.google.common.base.Function;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiConsumer;

/**
 * 基于MultiTreeSearcherBase封装的便捷api
 *
 * @author Alone
 * @see MultiTreeSearcherBase
 */
@SuppressWarnings("all")
public class MultiTreeSearcher<T, R> extends MultiTreeSearcherBase<T, MultiTreeSearcher<T, R>> {

    public static <T, R> MultiTreeSearcher<T, R> of(List<Tree<T>> source) {
        return new MultiTreeSearcher<>(source);
    }

    public static <T, R> MultiTreeSearcher<T, R> of(List<Tree<T>> source, Class<R> clazz) {
        return new MultiTreeSearcher<>(source);
    }

    private final Function<Func1<R, ?>, Function<Tree, String>> strValueFunc = key -> {
        String fieldName = LambdaUtil.getFieldName(key);
        return tree -> tree.get(fieldName).toString();
    };

    private final Function<Func1<R, ?>, Function<Tree, ?>> valueFunc = key -> {
        String fieldName = LambdaUtil.getFieldName(key);
        return tree -> tree.get(fieldName);
    };

    private final Map<String, Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>>> matcherMapKeyCache;

    private MultiTreeSearcher(List<Tree<T>> source) {
        super(source);
        matcherMapKeyCache = new HashMap<>(1);
    }

    public MultiTreeSearcher<T, R> eq(String key, Object value) {
        return conditionEq(true, defaultKey, tree -> tree.get(key), value);
    }

    public MultiTreeSearcher<T, R> like(String key, String value) {
        return conditionLike(true, defaultKey, tree -> tree.get(key).toString(), value);
    }

    public MultiTreeSearcher<T, R> eq(Func1<R, ?> key, Object value) {
        return conditionEq(true, defaultKey, valueFunc.apply(key), value);
    }

    public MultiTreeSearcher<T, R> like(Func1<R, String> key, String value) {
        return conditionLike(true, defaultKey, strValueFunc.apply(key), value);
    }

    public MultiTreeSearcher<T, R> eqSubTree(String treeKey, String key, Object value) {
        return conditionEq(true, getMatcherMapKey(treeKey), tree -> tree.get(key), value);
    }

    public MultiTreeSearcher<T, R> likeSubTree(String treeKey, String key, String value) {
        return conditionLike(true, getMatcherMapKey(treeKey), tree -> tree.get(key).toString(), value);
    }

    public MultiTreeSearcher<T, R> eq(boolean condition, String key, Object value) {
        return conditionEq(condition, defaultKey, tree -> tree.get(key), value);
    }

    public MultiTreeSearcher<T, R> like(boolean condition, String key, String value) {
        return conditionLike(condition, defaultKey, tree -> tree.get(key).toString(), value);
    }

    public MultiTreeSearcher<T, R> eq(boolean condition, Func1<R, ?> key, Object value) {
        return conditionEq(condition, defaultKey, valueFunc.apply(key), value);
    }

    public MultiTreeSearcher<T, R> like(boolean condition, Func1<R, String> key, String value) {
        return conditionLike(condition, defaultKey, strValueFunc.apply(key), value);
    }

    public MultiTreeSearcher<T, R> eqSubTree(boolean condition, String treeKey, String key, Object value) {
        return conditionEq(condition, getMatcherMapKey(treeKey), tree -> tree.get(key), value);
    }

    public MultiTreeSearcher<T, R> likeSubTree(boolean condition, String treeKey, String key, String value) {
        return conditionLike(condition, getMatcherMapKey(treeKey), tree -> tree.get(key).toString(), value);
    }


    private MultiTreeSearcher<T, R> conditionEq(boolean condition, Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> key, Function<Tree, ?> valueFunc, Object value) {
        if (condition) {
            return super.eq(key, valueFunc, value);
        }
        return typedThis;
    }

    private MultiTreeSearcher<T, R> conditionLike(boolean condition, Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> key, Function<Tree, String> valueFunc, String value) {
        if (condition) {
            return super.like(key, valueFunc, value);
        }
        return typedThis;
    }

    private final Function<String, Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>>>
            matcherMapKeyFunc = key -> Pair.of(tree -> (List<Tree>) tree.get(key), (tree, list) -> tree.putExtra(key, list));

    private synchronized Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> getMatcherMapKey(String key) {
        Pair<Function<Tree, List<Tree>>, BiConsumer<Tree, List<Tree>>> res = matcherMapKeyCache.get(key);
        if (Objects.isNull(res)) {
            res = matcherMapKeyFunc.apply(key);
            matcherMapKeyCache.put(key, res);
        }
        return res;
    }

}