Some code refactoring
Added some missing comments