Small improvements to tree generation