Make some code improvements