Using states instead of stages