Fix simplify_tree

Needs more testing.