Updated to new tree API