Minor refactoring
More code refactoring