Fix bugs in control flow optimisations