Fix lint warning, and cleanup the code