Refactor diff logic - #6