Merge pull request #7 from SummerSun/master

Update code according to Noah's comments