Merge pull request #275 from raghavgautam/fixes

Small refactoring and fixes