Small changes in tree generator