Fix undo problem
Minor refactor