More code improvements