[minor] Clean up of the code.