slight code fixes