Add a sanity check and do minor refactoring