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 ca4d49e..bfeb301 100644
--- a/src/main/java/xyz/zhouxy/plusone/commons/util/TreeBuilder.java
+++ b/src/main/java/xyz/zhouxy/plusone/commons/util/TreeBuilder.java
@@ -8,14 +8,14 @@ import java.util.function.BiConsumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-public class TreeBuilder<T, TIdentity> {
+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<T, T> addChildrenMethod;
+    private final BiConsumer<TSubTree, T> addChildrenMethod;
 
     public TreeBuilder(Collection<T> nodes, Function<T, TIdentity> identityGetter,
-            Function<T, Optional<TIdentity>> parentIdentityGetter, BiConsumer<T, T> addChildren) {
+            Function<T, Optional<TIdentity>> parentIdentityGetter, BiConsumer<TSubTree, T> addChildren) {
         this.nodes = nodes;
         this.identityGetter = identityGetter;
         this.parentIdentityGetter = parentIdentityGetter;
@@ -30,7 +30,8 @@ public class TreeBuilder<T, TIdentity> {
         for (T node : this.nodes) {
             Optional<TIdentity> parentIdentity = parentIdentityGetter.apply(node);
             if (parentIdentity.isPresent() && identityNodeMap.containsKey(parentIdentity.get())) {
-                T parentNode = identityNodeMap.get(parentIdentity.get());
+                @SuppressWarnings("all")
+                TSubTree parentNode = (TSubTree) identityNodeMap.get(parentIdentity.get());
                 addChildrenMethod.accept(parentNode, node);
             }
         }
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 dd2c063..e3a4d43 100644
--- a/src/test/java/xyz/zhouxy/plusone/commons/util/TreeBuilderTests.java
+++ b/src/test/java/xyz/zhouxy/plusone/commons/util/TreeBuilderTests.java
@@ -35,7 +35,7 @@ class TreeBuilderTests {
                 menus,
                 Menu::getMenuCode,
                 Menu::getParentMenuCode,
-                Menu::addChild)
+                MenuList::addChild)
                 .buildTree();
         log.info("menuTree: {}", menuTree);
     }
@@ -64,8 +64,6 @@ abstract class Menu {
     public String getTitle() {
         return title;
     }
-
-    public abstract void addChild(Menu child);
 }
 
 @ToString(callSuper = true)
@@ -89,12 +87,6 @@ class MenuItem extends Menu {
     public String getUrl() {
         return url;
     }
-
-    @Override
-    @Deprecated
-    public void addChild(Menu child) {
-        throw new UnsupportedOperationException("Unimplemented method 'addChild'");
-    }
 }
 
 @ToString(callSuper = true)
@@ -126,7 +118,6 @@ class MenuList extends Menu {
         return instance;
     }
 
-    @Override
     public void addChild(Menu child) {
         if (this.children == null) {
             this.children = Lists.newArrayList();