refactor: split if then else

src/alejandra_engine/src/builder.rs

@@ -288,10 +288,10 @@ fn format_wider(
 
 pub fn fits_in_single_line(
     build_ctx: &crate::builder::BuildCtx,
-    node: rnix::SyntaxElement,
+    element: rnix::SyntaxElement,
 ) -> bool {
     let line = build_ctx.pos_new.line;
-    let maybe_green_node = build(node, true, build_ctx.path.clone(), false);
+    let maybe_green_node = build(element, true, build_ctx.path.clone(), false);
 
     match maybe_green_node {
         Some(_) => build_ctx.pos_new.line == line,

src/alejandra_engine/src/children.rs

@@ -134,29 +134,6 @@ impl Children {
         })
     }
 
-    // #[deprecated]
-    pub fn drain_comments<F: FnMut(String)>(&mut self, mut callback: F) {
-        while let Some(child) = self.peek_next() {
-            match child.element.kind() {
-                rnix::SyntaxKind::TOKEN_COMMENT => {
-                    callback(
-                        child
-                            .element
-                            .as_token()
-                            .as_ref()
-                            .unwrap()
-                            .text()
-                            .to_string(),
-                    );
-                    self.move_next();
-                }
-                _ => {
-                    break;
-                }
-            }
-        }
-    }
-
     pub fn drain_trivia<F: FnMut(Trivia)>(&mut self, mut callback: F) {
         while let Some(child) = self.peek_next() {
             match child.element.kind() {

src/alejandra_engine/src/parsers/if_else.rs

@@ -0,0 +1,84 @@
+use std::collections::LinkedList;
+
+#[derive(Debug, Default)]
+pub struct IfElse {
+    pub comments_before_if_expr:   LinkedList<String>,
+    pub if_expr:                   Option<rnix::SyntaxElement>,
+    pub comments_after_if_expr:    LinkedList<String>,
+    pub comments_before_then_expr: LinkedList<String>,
+    pub then_expr:                 Option<rnix::SyntaxElement>,
+    pub comments_after_then_expr:  LinkedList<String>,
+    pub comments_before_else_expr: LinkedList<String>,
+    pub else_expr:                 Option<rnix::SyntaxElement>,
+}
+
+pub fn parse(
+    build_ctx: &crate::builder::BuildCtx,
+    node: &rnix::SyntaxNode,
+) -> IfElse {
+    let mut if_else = IfElse::default();
+
+    let mut children = crate::children::Children::new_with_configuration(
+        build_ctx, node, true,
+    );
+
+    // if
+    children.get_next().unwrap();
+
+    // /**/
+    children.drain_trivia(|element| match element {
+        crate::children::Trivia::Comment(text) => {
+            if_else.comments_before_if_expr.push_back(text);
+        }
+        crate::children::Trivia::Whitespace(_) => {}
+    });
+
+    // expr
+    if_else.if_expr = Some(children.get_next().unwrap().element);
+
+    // /**/
+    children.drain_trivia(|element| match element {
+        crate::children::Trivia::Comment(text) => {
+            if_else.comments_after_if_expr.push_back(text);
+        }
+        crate::children::Trivia::Whitespace(_) => {}
+    });
+
+    // then
+    children.get_next().unwrap();
+
+    // /**/
+    children.drain_trivia(|element| match element {
+        crate::children::Trivia::Comment(text) => {
+            if_else.comments_before_then_expr.push_back(text);
+        }
+        crate::children::Trivia::Whitespace(_) => {}
+    });
+
+    // expr
+    if_else.then_expr = Some(children.get_next().unwrap().element);
+
+    // /**/
+    children.drain_trivia(|element| match element {
+        crate::children::Trivia::Comment(text) => {
+            if_else.comments_after_then_expr.push_back(text);
+        }
+        crate::children::Trivia::Whitespace(_) => {}
+    });
+
+    // else
+    children.get_next().unwrap();
+
+    // /**/
+    children.drain_trivia(|element| match element {
+        crate::children::Trivia::Comment(text) => {
+            if_else.comments_before_else_expr.push_back(text);
+        }
+        crate::children::Trivia::Whitespace(_) => {}
+    });
+
+    // expr
+    if_else.else_expr = Some(children.get_next().unwrap().element);
+
+    if_else
+}

src/alejandra_engine/src/parsers/mod.rs

@@ -1 +1,2 @@
+pub mod if_else;
 pub mod pattern;

src/alejandra_engine/src/rules/if_else.rs

@@ -4,67 +4,142 @@ pub fn rule(
 ) -> std::collections::LinkedList<crate::builder::Step> {
     let mut steps = std::collections::LinkedList::new();
 
-    let mut children = crate::children::Children::new(build_ctx, node);
+    let if_else = crate::parsers::if_else::parse(build_ctx, node);
 
-    for branch in ["if", "then", "else"] {
+    // if
-        // if/then/else
+    steps.push_back(crate::builder::Step::Token(
-        let child = children.get_next().unwrap();
+        rnix::SyntaxKind::TOKEN_IF,
-        steps.push_back(crate::builder::Step::Format(child.element));
+        "if".to_string(),
+    ));
 
-        if let rnix::SyntaxKind::TOKEN_COMMENT =
+    if if_else.comments_before_if_expr.is_empty() {
-            children.peek_next().unwrap().element.kind()
+        // expr
-        {
+        let element = if_else.if_expr.unwrap();
-            steps.push_back(crate::builder::Step::Indent);
+        if crate::builder::fits_in_single_line(build_ctx, element.clone()) {
-
+            steps.push_back(crate::builder::Step::Whitespace);
-            // /**/
+            steps.push_back(crate::builder::Step::FormatWider(element));
-            children.drain_comments(|text| {
-                steps.push_back(crate::builder::Step::NewLine);
-                steps.push_back(crate::builder::Step::Pad);
-                steps.push_back(crate::builder::Step::Comment(text));
-            });
-
-            // expr
-            let child = children.get_next().unwrap();
-            steps.push_back(crate::builder::Step::NewLine);
-            steps.push_back(crate::builder::Step::Pad);
-            steps.push_back(crate::builder::Step::FormatWider(child.element));
-            steps.push_back(crate::builder::Step::Dedent);
         } else {
-            let child = children.get_next().unwrap();
+            steps.push_back(crate::builder::Step::Indent);
-
-            // expr
-            if crate::builder::fits_in_single_line(
-                build_ctx,
-                child.element.clone(),
-            ) || (branch == "else"
-                && child.element.kind() == rnix::SyntaxKind::NODE_IF_ELSE)
-            {
-                steps.push_back(crate::builder::Step::Whitespace);
-                steps.push_back(crate::builder::Step::FormatWider(
-                    child.element,
-                ));
-            } else {
-                steps.push_back(crate::builder::Step::Indent);
-                steps.push_back(crate::builder::Step::NewLine);
-                steps.push_back(crate::builder::Step::Pad);
-                steps.push_back(crate::builder::Step::FormatWider(
-                    child.element,
-                ));
-                steps.push_back(crate::builder::Step::Dedent);
-            }
-        }
-
-        if branch != "else" {
             steps.push_back(crate::builder::Step::NewLine);
             steps.push_back(crate::builder::Step::Pad);
-
+            steps.push_back(crate::builder::Step::FormatWider(element));
-            // /**/
+            steps.push_back(crate::builder::Step::Dedent);
-            children.drain_comments(|text| {
-                steps.push_back(crate::builder::Step::Comment(text));
-                steps.push_back(crate::builder::Step::NewLine);
-                steps.push_back(crate::builder::Step::Pad);
-            });
         }
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+    } else {
+        // /**/
+        steps.push_back(crate::builder::Step::Indent);
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+        for text in if_else.comments_before_if_expr {
+            steps.push_back(crate::builder::Step::Comment(text));
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+        }
+        // expr
+        steps.push_back(crate::builder::Step::FormatWider(
+            if_else.if_expr.unwrap(),
+        ));
+        steps.push_back(crate::builder::Step::Dedent);
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+    }
+
+    // /**/
+    if !if_else.comments_after_if_expr.is_empty() {
+        for text in if_else.comments_after_if_expr {
+            steps.push_back(crate::builder::Step::Comment(text));
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+        }
+    }
+
+    // then
+    steps.push_back(crate::builder::Step::Token(
+        rnix::SyntaxKind::TOKEN_THEN,
+        "then".to_string(),
+    ));
+
+    if if_else.comments_before_then_expr.is_empty() {
+        // expr
+        let element = if_else.then_expr.unwrap();
+        if crate::builder::fits_in_single_line(build_ctx, element.clone()) {
+            steps.push_back(crate::builder::Step::Whitespace);
+            steps.push_back(crate::builder::Step::FormatWider(element));
+        } else {
+            steps.push_back(crate::builder::Step::Indent);
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+            steps.push_back(crate::builder::Step::FormatWider(element));
+            steps.push_back(crate::builder::Step::Dedent);
+        }
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+    } else {
+        // /**/
+        steps.push_back(crate::builder::Step::Indent);
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+        for text in if_else.comments_before_then_expr {
+            steps.push_back(crate::builder::Step::Comment(text));
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+        }
+        // expr
+        steps.push_back(crate::builder::Step::FormatWider(
+            if_else.then_expr.unwrap(),
+        ));
+        steps.push_back(crate::builder::Step::Dedent);
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+    }
+
+    // /**/
+    if !if_else.comments_after_then_expr.is_empty() {
+        for text in if_else.comments_after_then_expr {
+            steps.push_back(crate::builder::Step::Comment(text));
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+        }
+    }
+
+    // else
+    steps.push_back(crate::builder::Step::Token(
+        rnix::SyntaxKind::TOKEN_ELSE,
+        "else".to_string(),
+    ));
+
+    if if_else.comments_before_else_expr.is_empty() {
+        // expr
+        let element = if_else.else_expr.unwrap();
+        if matches!(element.kind(), rnix::SyntaxKind::NODE_IF_ELSE)
+            || crate::builder::fits_in_single_line(build_ctx, element.clone())
+        {
+            steps.push_back(crate::builder::Step::Whitespace);
+            steps.push_back(crate::builder::Step::FormatWider(element));
+        } else {
+            steps.push_back(crate::builder::Step::Indent);
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+            steps.push_back(crate::builder::Step::FormatWider(element));
+            steps.push_back(crate::builder::Step::Dedent);
+        }
+    } else {
+        // /**/
+        steps.push_back(crate::builder::Step::Indent);
+        steps.push_back(crate::builder::Step::NewLine);
+        steps.push_back(crate::builder::Step::Pad);
+        for text in if_else.comments_before_else_expr {
+            steps.push_back(crate::builder::Step::Comment(text));
+            steps.push_back(crate::builder::Step::NewLine);
+            steps.push_back(crate::builder::Step::Pad);
+        }
+        // expr
+        steps.push_back(crate::builder::Step::FormatWider(
+            if_else.else_expr.unwrap(),
+        ));
+        steps.push_back(crate::builder::Step::Dedent);
     }
 
     steps