Fixing some bugs which just popped up. (I should really add more unit tests)