Tweak the code to make it more reusable.