small refactoring and improvement