Fixed an issue with my change to the state update order.