Bug fix in flattening if/then/else