Some code tweaks