Code refactor

Minor code refactoring.