Control flow fixes