Fix add return statement if last statement is an ifStatement