Fix TreeWalk to correctly handle missing subtrees