Code cleanup :heart: