More changes from code refactoring