Fix cycling for tree
Refactor code