Some code improvements