Refactor the add state function.