Improve code.
Merge pull request #9 from YetiOr/master

Improve code.