More code tweaks