Refactoring a bit of logic improvement