diff --git a/pom.xml b/pom.xml
index d3a38df..bb64491 100644
--- a/pom.xml
+++ b/pom.xml
@@ -81,6 +81,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>com.google.code.gson</groupId>
+            <artifactId>gson</artifactId>
+            <version>2.10.1</version>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
 
     <build>
diff --git a/src/main/java/xyz/zhouxy/plusone/commons/util/TreeBuilder.java b/src/main/java/xyz/zhouxy/plusone/commons/util/TreeBuilder.java
index 71790d4..6be4bfe 100644
--- a/src/main/java/xyz/zhouxy/plusone/commons/util/TreeBuilder.java
+++ b/src/main/java/xyz/zhouxy/plusone/commons/util/TreeBuilder.java
@@ -11,32 +11,29 @@ import java.util.stream.Collectors;
 import xyz.zhouxy.plusone.commons.collection.CollectionTools;
 
 public class TreeBuilder<T, TSubTree extends T, TIdentity> {
-    private final Collection<T> nodes;
     private final Function<T, TIdentity> identityGetter;
     private final Function<T, Optional<TIdentity>> parentIdentityGetter;
-    private final BiConsumer<TSubTree, T> addChildrenMethod;
+    private final BiConsumer<TSubTree, T> addChildMethod;
 
-    public TreeBuilder(Collection<T> nodes, Function<T, TIdentity> identityGetter,
-            Function<T, Optional<TIdentity>> parentIdentityGetter, BiConsumer<TSubTree, T> addChildren) {
-        this.nodes = nodes;
+    public TreeBuilder(Function<T, TIdentity> identityGetter, Function<T, Optional<TIdentity>> parentIdentityGetter,
+            BiConsumer<TSubTree, T> addChild) {
         this.identityGetter = identityGetter;
         this.parentIdentityGetter = parentIdentityGetter;
-        this.addChildrenMethod = addChildren;
+        this.addChildMethod = addChild;
     }
 
-    public List<T> buildTree() {
+    public List<T> buildTree(Collection<T> nodes) {
         Map<TIdentity, T> identityNodeMap = CollectionTools.toHashMap(nodes, identityGetter);
-        List<T> result = this.nodes.stream()
+        List<T> result = nodes.stream()
                 .filter(node -> !this.parentIdentityGetter.apply(node).isPresent())
                 .collect(Collectors.toList());
-        for (T node : this.nodes) {
-            Optional<TIdentity> parentIdentity = parentIdentityGetter.apply(node);
-            if (parentIdentity.isPresent() && identityNodeMap.containsKey(parentIdentity.get())) {
-                @SuppressWarnings("all")
-                TSubTree parentNode = (TSubTree) identityNodeMap.get(parentIdentity.get());
-                addChildrenMethod.accept(parentNode, node);
+        nodes.forEach(node -> parentIdentityGetter.apply(node).ifPresent(parentIdentity -> {
+            if (identityNodeMap.containsKey(parentIdentity)) {
+                @SuppressWarnings("unchecked")
+                TSubTree parentNode = (TSubTree) identityNodeMap.get(parentIdentity);
+                addChildMethod.accept(parentNode, node);
             }
-        }
+        }));
         return result;
     }
 }
diff --git a/src/test/java/xyz/zhouxy/plusone/commons/util/TreeBuilderTests.java b/src/test/java/xyz/zhouxy/plusone/commons/util/TreeBuilderTests.java
index e3a4d43..14b3828 100644
--- a/src/test/java/xyz/zhouxy/plusone/commons/util/TreeBuilderTests.java
+++ b/src/test/java/xyz/zhouxy/plusone/commons/util/TreeBuilderTests.java
@@ -8,12 +8,17 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import com.google.common.collect.Lists;
+import com.google.gson.Gson;
 
 import lombok.ToString;
 
 class TreeBuilderTests {
 
     private static final Logger log = LoggerFactory.getLogger(TreeBuilderTests.class);
+    private final TreeBuilder<Menu, MenuList, String> treeBuilder = new TreeBuilder<>(
+            Menu::getMenuCode,
+            Menu::getParentMenuCode,
+            MenuList::addChild);
 
     @Test
     void testBuildTree() {
@@ -31,13 +36,10 @@ class TreeBuilderTests {
                 MenuItem.of("C", "C2", "二级菜单C2", "/c/c2"),
                 MenuItem.of("C", "C3", "二级菜单C3", "/c/c3")
         );
-        List<Menu> menuTree = new TreeBuilder<>(
-                menus,
-                Menu::getMenuCode,
-                Menu::getParentMenuCode,
-                MenuList::addChild)
-                .buildTree();
-        log.info("menuTree: {}", menuTree);
+
+        List<Menu> menuTree = treeBuilder.buildTree(menus);
+        log.info("menuTree: {}", new Gson().toJson(menuTree));
+
     }
 }