Combine "else if" in formatter. (#2395)

Fixes #2393

executor/src/witgen/jit/compiler.rs

@@ -288,13 +288,17 @@ fn format_effects_inner<T: FieldElement>(
     effects: &[Effect<T, Variable>],
     is_top_level: bool,
 ) -> String {
-    indent(
-        effects
-            .iter()
-            .map(|effect| format_effect(effect, is_top_level))
-            .join("\n"),
-        1,
-    )
+    indent(format_effects_inner_unindented(effects, is_top_level), 1)
+}
+
+fn format_effects_inner_unindented<T: FieldElement>(
+    effects: &[Effect<T, Variable>],
+    is_top_level: bool,
+) -> String {
+    effects
+        .iter()
+        .map(|effect| format_effect(effect, is_top_level))
+        .join("\n")
 }
 
 fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>, is_top_level: bool) -> String {
@@ -369,12 +373,22 @@ fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>, is_top_level: bo
             } else {
                 "".to_string()
             };
-            format!(
-                "{var_decls}if {} {{\n{}\n}} else {{\n{}\n}}",
-                format_condition(condition),
-                format_effects_inner(first, false),
-                format_effects_inner(second, false)
-            )
+
+            if matches!(second[..], [Effect::Branch(..)]) {
+                format!(
+                    "{var_decls}if {} {{\n{}\n}} else if {}",
+                    format_condition(condition),
+                    format_effects_inner(first, false),
+                    format_effects_inner_unindented(second, false)
+                )
+            } else {
+                format!(
+                    "{var_decls}if {} {{\n{}\n}} else {{\n{}\n}}",
+                    format_condition(condition),
+                    format_effects_inner(first, false),
+                    format_effects_inner(second, false)
+                )
+            }
         }
     }
 }
@@ -998,19 +1012,29 @@ extern \"C\" fn witgen(
     fn branches_codegen() {
         let x = param(0);
         let y = param(1);
+        let z = param(2);
         let branch_effect = Effect::Branch(
             BranchCondition {
                 variable: x.clone(),
                 condition: RangeConstraint::from_range(7.into(), 20.into()),
             },
             vec![assignment(&y, symbol(&x) + number(1))],
-            vec![assignment(&y, symbol(&x) + number(2))],
+            vec![Effect::Branch(
+                BranchCondition {
+                    variable: z.clone(),
+                    condition: RangeConstraint::from_range(7.into(), 20.into()),
+                },
+                vec![assignment(&y, symbol(&x) + number(2))],
+                vec![assignment(&y, symbol(&x) + number(3))],
+            )],
         );
         let expectation = "    let p_1;
     if 7 <= IntType::from(p_0) && IntType::from(p_0) <= 20 {
         p_1 = (p_0 + FieldElement::from(1));
-    } else {
+    } else if if 7 <= IntType::from(p_2) && IntType::from(p_2) <= 20 {
         p_1 = (p_0 + FieldElement::from(2));
+    } else {
+        p_1 = (p_0 + FieldElement::from(3));
     }";
         assert_eq!(format_effects(&[branch_effect]), expectation);
     }

executor/src/witgen/jit/effect.rs

@@ -145,11 +145,22 @@ pub fn format_code<T: FieldElement>(effects: &[Effect<T, Variable>]) -> String {
                 panic!("Range constraints should not be part of the code.")
             }
             Effect::Branch(condition, first, second) => {
-                let first = indent(format_code(first), 1);
-                let second = indent(format_code(second), 1);
+                let first = format_code(first);
+                let second_str = format_code(second);
                 let condition = format_condition(condition);
 
-                format!("if ({condition}) {{\n{first}\n}} else {{\n{second}\n}}")
+                if matches!(second[..], [Effect::Branch(..)]) {
+                    format!(
+                        "if ({condition}) {{\n{}\n}} else {second_str}",
+                        indent(first, 1),
+                    )
+                } else {
+                    format!(
+                        "if ({condition}) {{\n{}\n}} else {{\n{}\n}}",
+                        indent(first, 1),
+                        indent(second_str, 1)
+                    )
+                }
             }
         })
         .join("\n")
@@ -168,3 +179,92 @@ fn format_condition<T: FieldElement>(
         Ordering::Greater => format!("{variable} <= {min} || {variable} >= {max}"),
     }
 }
+
+#[cfg(test)]
+mod test {
+    use powdr_number::GoldilocksField;
+
+    use crate::witgen::jit::variable::Cell;
+
+    use pretty_assertions::assert_eq;
+
+    use super::*;
+    type T = GoldilocksField;
+
+    fn var(id: u64) -> Variable {
+        Variable::WitnessCell(Cell {
+            column_name: format!("v{id}"),
+            id,
+            row_offset: 0,
+        })
+    }
+
+    #[test]
+    fn combine_if_else() {
+        let effects = vec![
+            Effect::Assignment(var(0), SymbolicExpression::from(T::from(1))),
+            Effect::Branch(
+                BranchCondition {
+                    variable: var(0),
+                    condition: RangeConstraint::from_range(T::from(1), T::from(2)),
+                },
+                vec![Effect::Assignment(
+                    var(1),
+                    SymbolicExpression::from(T::from(2)),
+                )],
+                vec![Effect::Branch(
+                    BranchCondition {
+                        variable: var(1),
+                        condition: RangeConstraint::from_range(T::from(5), T::from(6)),
+                    },
+                    vec![Effect::Branch(
+                        BranchCondition {
+                            variable: var(2),
+                            condition: RangeConstraint::from_range(T::from(5), T::from(6)),
+                        },
+                        vec![Effect::Assignment(
+                            var(8),
+                            SymbolicExpression::from(T::from(3)),
+                        )],
+                        vec![Effect::Assignment(
+                            var(9),
+                            SymbolicExpression::from(T::from(4)),
+                        )],
+                    )],
+                    vec![Effect::Branch(
+                        BranchCondition {
+                            variable: var(3),
+                            condition: RangeConstraint::from_range(T::from(5), T::from(6)),
+                        },
+                        vec![Effect::Assignment(
+                            var(21),
+                            SymbolicExpression::from(T::from(3)),
+                        )],
+                        vec![Effect::Assignment(
+                            var(22),
+                            SymbolicExpression::from(T::from(4)),
+                        )],
+                    )],
+                )],
+            ),
+        ];
+        let code = format_code(&effects);
+        assert_eq!(
+            code,
+            "v0[0] = 1;
+if (1 <= v0[0] && v0[0] <= 2) {
+    v1[0] = 2;
+} else if (5 <= v1[0] && v1[0] <= 6) {
+    if (5 <= v2[0] && v2[0] <= 6) {
+        v8[0] = 3;
+    } else {
+        v9[0] = 4;
+    }
+} else if (5 <= v3[0] && v3[0] <= 6) {
+    v21[0] = 3;
+} else {
+    v22[0] = 4;
+}"
+        );
+    }
+}