Add my refactoring attempt to end of code samples