Fixed some crashes after refactoring