Pull request comments and minor refactoring