Small refactoring, and add comments