Small code improvements