Refactor the logic for editing.