Code cleanup as suggested in pull request #3