Improve handling of stage 1.