[keep-tree-state] Tree trait refactoring