Small code refactor & comments edit