Small code fixes