fix formatting and merging if chains in attributes (#2697)

* fix formatting and merging if chains in attributes

* fix autoformat if attribute chains

* Fix IfAttributeValue ToTokens implementation

---------

Co-authored-by: Jonathan Kelley <jkelleyrtp@gmail.com>

examples/rsx_usage.rs

@@ -85,6 +85,9 @@ fn app() -> Element {
                 class: "{asd}",
                 // if statements can be used to conditionally render attributes
                 class: if formatting.contains("form") { "{asd}" },
+                // longer if chains also work
+                class: if formatting.contains("form") { "{asd}" } else if formatting.contains("my other form") { "{asd}" },
+                class: if formatting.contains("form") { "{asd}" } else if formatting.contains("my other form") { "{asd}" } else { "{asd}" },
                 div {
                     class: {
                         const WORD: &str = "expressions";

packages/autofmt/src/rsx_block.rs

@@ -226,14 +226,8 @@ impl Writer<'_> {
 
     fn write_attribute_value(&mut self, value: &AttributeValue) -> Result {
         match value {
-            AttributeValue::AttrOptionalExpr { condition, value } => {
-                write!(
-                    self.out,
-                    "if {condition} {{ ",
-                    condition = unparse_expr(condition),
-                )?;
-                self.write_attribute_value(value)?;
-                write!(self.out, " }}")?;
+            AttributeValue::IfExpr(if_chain) => {
+                self.write_attribute_if_chain(if_chain)?;
             }
             AttributeValue::AttrLiteral(value) => {
                 write!(self.out, "{value}")?;
@@ -258,6 +252,26 @@ impl Writer<'_> {
         Ok(())
     }
 
+    fn write_attribute_if_chain(&mut self, if_chain: &IfAttributeValue) -> Result {
+        write!(self.out, "if {} {{ ", unparse_expr(&if_chain.condition))?;
+        self.write_attribute_value(&if_chain.then_value)?;
+        write!(self.out, " }}")?;
+        match if_chain.else_value.as_deref() {
+            Some(AttributeValue::IfExpr(else_if_chain)) => {
+                write!(self.out, "else ")?;
+                self.write_attribute_if_chain(else_if_chain)?;
+            }
+            Some(other) => {
+                write!(self.out, "else {{")?;
+                self.write_attribute_value(other)?;
+                write!(self.out, " }}")?;
+            }
+            None => {}
+        }
+
+        Ok(())
+    }
+
     fn write_mulitiline_tokens(&mut self, out: String) -> Result {
         let mut lines = out.split('\n').peekable();
         let first = lines.next().unwrap();

packages/autofmt/src/writer.rs

@@ -249,10 +249,18 @@ impl<'a> Writer<'a> {
 
     pub(crate) fn attr_value_len(&mut self, value: &ElementAttrValue) -> usize {
         match value {
-            ElementAttrValue::AttrOptionalExpr { condition, value } => {
-                let condition_len = self.retrieve_formatted_expr(condition).len();
-                let value_len = self.attr_value_len(value);
-                condition_len + value_len + 6
+            ElementAttrValue::IfExpr(if_chain) => {
+                let condition_len = self.retrieve_formatted_expr(&if_chain.condition).len();
+                let value_len = self.attr_value_len(&if_chain.then_value);
+                let if_len = 2;
+                let brace_len = 2;
+                let space_len = 2;
+                let else_len = if_chain
+                    .else_value
+                    .as_ref()
+                    .map(|else_value| self.attr_value_len(else_value) + 1)
+                    .unwrap_or_default();
+                condition_len + value_len + if_len + brace_len + space_len + else_len
             }
             ElementAttrValue::AttrLiteral(lit) => lit.to_string().len(),
             ElementAttrValue::Shorthand(expr) => expr.span().line_length(),

packages/rsx/src/attribute.rs

@@ -23,8 +23,9 @@ use std::fmt::Display;
 use syn::{
     ext::IdentExt,
     parse::{Parse, ParseStream},
+    parse_quote,
     spanned::Spanned,
-    Expr, ExprClosure, ExprIf, Ident, Lit, LitBool, LitFloat, LitInt, LitStr, Token,
+    Block, Expr, ExprClosure, ExprIf, Ident, Lit, LitBool, LitFloat, LitInt, LitStr, Token,
 };
 
 #[cfg(feature = "hot_reload")]
@@ -241,7 +242,7 @@ impl Attribute {
                 AttributeValue::AttrLiteral(_)
                 | AttributeValue::AttrExpr(_)
                 | AttributeValue::Shorthand(_)
-                | AttributeValue::AttrOptionalExpr { .. }
+                | AttributeValue::IfExpr { .. }
                     if is_not_event =>
                 {
                     let name = &self.name;
@@ -439,20 +440,17 @@ pub enum AttributeValue {
     /// A series of tokens that represent an event handler
     ///
     /// We use a special type here so we can get autocomplete in the closure using partial expansion.
-    /// We also do some extra wrapping for improved type hinting since rust sometimes as trouble with
+    /// We also do some extra wrapping for improved type hinting since rust sometimes has trouble with
     /// generics and closures.
     EventTokens(PartialClosure),
 
-    /// Unterminated expression - full expressions are handled by AttrExpr
+    /// Conditional expression
     ///
-    /// attribute: if bool { "value" }
+    /// attribute: if bool { "value" } else if bool { "other value" } else { "default value" }
     ///
     /// Currently these don't get hotreloading super powers, but they could, depending on how far
     /// we want to go with it
-    AttrOptionalExpr {
-        condition: Expr,
-        value: Box<AttributeValue>,
-    },
+    IfExpr(IfAttributeValue),
 
     /// attribute: some_expr
     /// attribute: {some_expr} ?
@@ -463,45 +461,7 @@ impl Parse for AttributeValue {
     fn parse(content: ParseStream) -> syn::Result<Self> {
         // Attempt to parse the unterminated if statement
         if content.peek(Token![if]) {
-            let if_expr = content.parse::<ExprIf>()?;
-
-            if is_if_chain_terminated(&if_expr) {
-                return Ok(AttributeValue::AttrExpr(
-                    syn::parse2(if_expr.to_token_stream()).unwrap(),
-                ));
-            }
-
-            let stmts = &if_expr.then_branch.stmts;
-
-            if stmts.len() != 1 {
-                return Err(syn::Error::new(
-                    if_expr.then_branch.span(),
-                    "Expected a single statement in the if block",
-                ));
-            }
-
-            // either an ifmt or an expr in the block
-            let stmt = &stmts[0];
-
-            // Either it's a valid ifmt or an expression
-            let value = match stmt {
-                syn::Stmt::Expr(exp, None) => {
-                    // Try parsing the statement as an IfmtInput by passing it through tokens
-                    let value: Result<HotLiteral, syn::Error> = syn::parse2(quote! { #exp });
-                    match value {
-                        Ok(res) => Box::new(AttributeValue::AttrLiteral(res)),
-                        Err(_) => Box::new(AttributeValue::AttrExpr(
-                            syn::parse2(if_expr.to_token_stream()).unwrap(),
-                        )),
-                    }
-                }
-                _ => return Err(syn::Error::new(stmt.span(), "Expected an expression")),
-            };
-
-            return Ok(AttributeValue::AttrOptionalExpr {
-                condition: *if_expr.cond,
-                value,
-            });
+            return Ok(Self::IfExpr(content.parse::<IfAttributeValue>()?));
         }
 
         // Use the move and/or bars as an indicator that we have an event handler
@@ -534,15 +494,7 @@ impl ToTokens for AttributeValue {
         match self {
             Self::Shorthand(ident) => ident.to_tokens(tokens),
             Self::AttrLiteral(ifmt) => ifmt.to_tokens(tokens),
-            Self::AttrOptionalExpr { condition, value } => tokens.append_all(quote! {
-                {
-                    if #condition {
-                        Some(#value)
-                    } else {
-                        None
-                    }
-                }
-            }),
+            Self::IfExpr(if_expr) => if_expr.to_tokens(tokens),
             Self::AttrExpr(expr) => expr.to_tokens(tokens),
             Self::EventTokens(closure) => closure.to_tokens(tokens),
         }
@@ -554,13 +506,209 @@ impl AttributeValue {
         match self {
             Self::Shorthand(ident) => ident.span(),
             Self::AttrLiteral(ifmt) => ifmt.span(),
-            Self::AttrOptionalExpr { value, .. } => value.span(),
+            Self::IfExpr(if_expr) => if_expr.span(),
             Self::AttrExpr(expr) => expr.span(),
             Self::EventTokens(closure) => closure.span(),
         }
     }
 }
 
+/// A if else chain attribute value
+#[derive(PartialEq, Eq, Clone, Debug, Hash)]
+pub struct IfAttributeValue {
+    pub condition: Expr,
+    pub then_value: Box<AttributeValue>,
+    pub else_value: Option<Box<AttributeValue>>,
+}
+
+impl IfAttributeValue {
+    /// Convert the if expression to an expression that returns a string. If the unterminated case is hit, it returns an empty string
+    pub(crate) fn quote_as_string(&self, diagnostics: &mut Diagnostics) -> Expr {
+        let mut expression = quote! {};
+        let mut current_if_value = self;
+
+        let mut non_string_diagnostic = |span: proc_macro2::Span| -> Expr {
+            Element::add_merging_non_string_diagnostic(diagnostics, span);
+            parse_quote! { ::std::string::String::new() }
+        };
+
+        loop {
+            let AttributeValue::AttrLiteral(lit) = current_if_value.then_value.as_ref() else {
+                return non_string_diagnostic(current_if_value.span());
+            };
+
+            let HotLiteralType::Fmted(new) = &lit.value else {
+                return non_string_diagnostic(current_if_value.span());
+            };
+
+            let condition = &current_if_value.condition;
+            expression.extend(quote! {
+                if #condition {
+                    #new.to_string()
+                } else
+            });
+            match current_if_value.else_value.as_deref() {
+                // If the else value is another if expression, then we need to continue the loop
+                Some(AttributeValue::IfExpr(else_value)) => {
+                    current_if_value = else_value;
+                }
+                // If the else value is a literal, then we need to append it to the expression and break
+                Some(AttributeValue::AttrLiteral(lit)) => {
+                    if let HotLiteralType::Fmted(new) = &lit.value {
+                        expression.extend(quote! { { #new.to_string() } });
+                        break;
+                    } else {
+                        return non_string_diagnostic(current_if_value.span());
+                    }
+                }
+                // If it is the end of the if expression without an else, then we need to append the default value and break
+                None => {
+                    expression.extend(quote! { { ::std::string::String::new() } });
+                    break;
+                }
+                _ => {
+                    return non_string_diagnostic(current_if_value.else_value.span());
+                }
+            }
+        }
+
+        parse_quote! {
+            {
+                #expression
+            }
+        }
+    }
+
+    fn span(&self) -> proc_macro2::Span {
+        self.then_value.span()
+    }
+
+    fn is_terminated(&self) -> bool {
+        match &self.else_value {
+            Some(attribute) => match attribute.as_ref() {
+                AttributeValue::IfExpr(if_expr) => if_expr.is_terminated(),
+                _ => true,
+            },
+            None => false,
+        }
+    }
+
+    fn parse_attribute_value_from_block(block: &Block) -> syn::Result<Box<AttributeValue>> {
+        let stmts = &block.stmts;
+
+        if stmts.len() != 1 {
+            return Err(syn::Error::new(
+                block.span(),
+                "Expected a single statement in the if block",
+            ));
+        }
+
+        // either an ifmt or an expr in the block
+        let stmt = &stmts[0];
+
+        // Either it's a valid ifmt or an expression
+        match stmt {
+            syn::Stmt::Expr(exp, None) => {
+                // Try parsing the statement as an IfmtInput by passing it through tokens
+                let value: Result<HotLiteral, syn::Error> = syn::parse2(quote! { #exp });
+                Ok(match value {
+                    Ok(res) => Box::new(AttributeValue::AttrLiteral(res)),
+                    Err(_) => Box::new(AttributeValue::AttrExpr(PartialExpr::from_expr(exp))),
+                })
+            }
+            _ => Err(syn::Error::new(stmt.span(), "Expected an expression")),
+        }
+    }
+
+    fn to_tokens_with_terminated(&self, tokens: &mut TokenStream2, terminated: bool) {
+        let IfAttributeValue {
+            condition,
+            then_value,
+            else_value,
+        } = self;
+        let then_value = if terminated {
+            quote! { #then_value }
+        }
+        // Otherwise we need to return an Option and a None if the else value is None
+        else {
+            quote! { Some(#then_value) }
+        };
+
+        let else_value = match else_value.as_deref() {
+            Some(AttributeValue::IfExpr(else_value)) => {
+                let mut tokens = TokenStream2::new();
+                else_value.to_tokens_with_terminated(&mut tokens, terminated);
+                tokens
+            }
+            Some(other) => {
+                if terminated {
+                    quote! { #other }
+                } else {
+                    quote! { Some(#other) }
+                }
+            }
+            None => quote! { None },
+        };
+
+        tokens.append_all(quote! {
+            {
+                if #condition {
+                    #then_value
+                } else {
+                    #else_value
+                }
+            }
+        });
+    }
+}
+
+impl Parse for IfAttributeValue {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let if_expr = input.parse::<ExprIf>()?;
+
+        let stmts = &if_expr.then_branch.stmts;
+
+        if stmts.len() != 1 {
+            return Err(syn::Error::new(
+                if_expr.then_branch.span(),
+                "Expected a single statement in the if block",
+            ));
+        }
+
+        // Parse the then branch into a single attribute value
+        let then_value = Self::parse_attribute_value_from_block(&if_expr.then_branch)?;
+
+        // If there's an else branch, parse it as a single attribute value or an if expression
+        let else_value = match if_expr.else_branch.as_ref() {
+            Some((_, else_branch)) => {
+                // The else branch if either a block or another if expression
+                let attribute_value = match else_branch.as_ref() {
+                    // If it is a block, then the else is terminated
+                    Expr::Block(block) => Self::parse_attribute_value_from_block(&block.block)?,
+                    // Otherwise try to parse it as an if expression
+                    _ => Box::new(syn::parse2(quote! { #else_branch })?),
+                };
+                Some(attribute_value)
+            }
+            None => None,
+        };
+
+        Ok(Self {
+            condition: *if_expr.cond,
+            then_value,
+            else_value,
+        })
+    }
+}
+
+impl ToTokens for IfAttributeValue {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        // If the if expression is terminated, we can just return the then value
+        let terminated = self.is_terminated();
+        self.to_tokens_with_terminated(tokens, terminated)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -581,22 +729,33 @@ mod tests {
         let _parsed: Attribute = parse2(quote! { "custom": false, }).unwrap();
         let _parsed: Attribute = parse2(quote! { name: false, }).unwrap();
 
-        // with expressions
-        let _parsed: Attribute = parse2(quote! { name: if true { "value" } }).unwrap();
-        let _parsed: Attribute =
+        // with if chains
+        let parsed: Attribute = parse2(quote! { name: if true { "value" } }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::IfExpr(_)));
+        let parsed: Attribute =
             parse2(quote! { name: if true { "value" } else { "other" } }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::IfExpr(_)));
+        let parsed: Attribute =
+            parse2(quote! { name: if true { "value" } else if false { "other" } }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::IfExpr(_)));
 
         // with shorthand
         let _parsed: Attribute = parse2(quote! { name }).unwrap();
         let _parsed: Attribute = parse2(quote! { name, }).unwrap();
 
         // Events - make sure they get partial expansion
-        let _parsed: Attribute = parse2(quote! { onclick: |e| {} }).unwrap();
-        let _parsed: Attribute = parse2(quote! { onclick: |e| { "value" } }).unwrap();
-        let _parsed: Attribute = parse2(quote! { onclick: |e| { value. } }).unwrap();
-        let _parsed: Attribute = parse2(quote! { onclick: move |e| { value. } }).unwrap();
-        let _parsed: Attribute = parse2(quote! { onclick: move |e| value }).unwrap();
-        let _parsed: Attribute = parse2(quote! { onclick: |e| value, }).unwrap();
+        let parsed: Attribute = parse2(quote! { onclick: |e| {} }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::EventTokens(_)));
+        let parsed: Attribute = parse2(quote! { onclick: |e| { "value" } }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::EventTokens(_)));
+        let parsed: Attribute = parse2(quote! { onclick: |e| { value. } }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::EventTokens(_)));
+        let parsed: Attribute = parse2(quote! { onclick: move |e| { value. } }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::EventTokens(_)));
+        let parsed: Attribute = parse2(quote! { onclick: move |e| value }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::EventTokens(_)));
+        let parsed: Attribute = parse2(quote! { onclick: |e| value, }).unwrap();
+        assert!(matches!(parsed.value, AttributeValue::EventTokens(_)));
     }
 
     #[test]

packages/rsx/src/element.rs

@@ -211,6 +211,12 @@ impl ToTokens for Element {
 }
 
 impl Element {
+    pub(crate) fn add_merging_non_string_diagnostic(diagnostics: &mut Diagnostics, span: Span) {
+        diagnostics.push(span.error("Cannot merge non-fmt literals").help(
+            "Only formatted strings can be merged together. If you want to merge literals, you can use a format string.",
+        ));
+    }
+
     /// Collapses ifmt attributes into a single dynamic attribute using a space or `;` as a delimiter
     ///
     /// ```ignore,
@@ -271,22 +277,16 @@ impl Element {
                     }
                 }
 
-                // Merge `if cond { "abc" }` into the output
-                if let AttributeValue::AttrOptionalExpr { condition, value } = &matching_attr.value
-                {
-                    if let AttributeValue::AttrLiteral(lit) = value.as_ref() {
-                        if let HotLiteralType::Fmted(new) = &lit.value {
-                            out.push_condition(condition.clone(), new.clone());
-                            continue;
-                        }
-                    }
+                // Merge `if cond { "abc" } else if ...` into the output
+                if let AttributeValue::IfExpr(value) = &matching_attr.value {
+                    out.push_expr(value.quote_as_string(&mut self.diagnostics));
+                    continue;
                 }
 
-                // unwind in case there's a test or two that cares about this weird state
-                _ = out.segments.pop();
-                self.diagnostics.push(matching_attr.span().error("Cannot merge non-fmt literals").help(
-                    "Only formatted strings can be merged together. If you want to merge literals, you can use a format string.",
-                ));
+                Self::add_merging_non_string_diagnostic(
+                    &mut self.diagnostics,
+                    matching_attr.span(),
+                );
             }
 
             let out_lit = HotLiteral {
@@ -533,11 +533,13 @@ fn merges_attributes() {
     let input = quote::quote! {
         div {
             class: "hello world",
-            class: if count > 3 { "abc {def}" }
+            class: if count > 3 { "abc {def}" },
+            class: if count < 50 { "small" } else { "big" }
         }
     };
 
     let parsed: Element = syn::parse2(input).unwrap();
+    assert_eq!(parsed.diagnostics.len(), 0);
     assert_eq!(parsed.merged_attributes.len(), 1);
     assert_eq!(
         parsed.merged_attributes[0].name.to_string(),

packages/rsx/src/ifchain.rs

@@ -7,7 +7,7 @@ use quote::quote;
 use quote::{ToTokens, TokenStreamExt};
 use syn::{
     parse::{Parse, ParseStream},
-    Expr, ExprIf, Result, Token,
+    Expr, Result, Token,
 };
 
 use crate::TemplateBody;
@@ -138,21 +138,6 @@ impl ToTokens for IfChain {
     }
 }
 
-pub(crate) fn is_if_chain_terminated(chain: &ExprIf) -> bool {
-    let mut current = chain;
-    loop {
-        if let Some((_, else_block)) = &current.else_branch {
-            if let Expr::If(else_if) = else_block.as_ref() {
-                current = else_if;
-            } else {
-                return true;
-            }
-        } else {
-            return false;
-        }
-    }
-}
-
 #[test]
 fn parses_if_chain() {
     let input = quote! {

packages/rsx/src/ifmt.rs

@@ -20,6 +20,12 @@ pub struct IfmtInput {
     pub segments: Vec<Segment>,
 }
 
+impl Default for IfmtInput {
+    fn default() -> Self {
+        Self::new(Span::call_site())
+    }
+}
+
 impl IfmtInput {
     pub fn new(span: Span) -> Self {
         Self {
@@ -45,22 +51,6 @@ impl IfmtInput {
         self.segments.extend(other.segments);
     }
 
-    pub fn push_condition(&mut self, condition: Expr, contents: IfmtInput) {
-        let desugared = quote! {
-            {
-                let _cond = if #condition { #contents.to_string() } else { String::new() };
-                _cond
-            }
-        };
-
-        let parsed = syn::parse2::<Expr>(desugared).unwrap();
-
-        self.segments.push(Segment::Formatted(FormattedSegment {
-            format_args: String::new(),
-            segment: FormattedSegmentType::Expr(Box::new(parsed)),
-        }));
-    }
-
     pub fn push_expr(&mut self, expr: Expr) {
         self.segments.push(Segment::Formatted(FormattedSegment {
             format_args: String::new(),
@@ -481,18 +471,6 @@ mod tests {
         assert!(input.is_static());
     }
 
-    #[test]
-    fn pushing_conditional() {
-        let mut input = syn::parse2::<IfmtInput>(quote! { "hello " }).unwrap();
-
-        input.push_condition(
-            parse_quote! { true },
-            syn::parse2::<IfmtInput>(quote! { "world" }).unwrap(),
-        );
-        println!("{}", input.to_token_stream().pretty_unparse());
-        dbg!(input.segments);
-    }
-
     #[test]
     fn fmt_segments() {
         let left = syn::parse2::<IfmtInput>(quote! { "thing {abc}" }).unwrap();

packages/rsx/src/scoring.rs

@@ -1,6 +1,6 @@
 #![cfg(feature = "hot_reload")]
 
-use crate::{Attribute, AttributeValue, BodyNode, HotLiteralType, IfmtInput};
+use crate::{Attribute, AttributeValue, BodyNode, HotLiteralType, IfAttributeValue, IfmtInput};
 
 /// Take two nodes and return their similarity score
 ///
@@ -114,17 +114,24 @@ fn score_attr_value(old_attr: &AttributeValue, new_attr: &AttributeValue) -> usi
         }
 
         (
-            AttrOptionalExpr {
+            IfExpr(IfAttributeValue {
                 condition: cond_a,
-                value: value_a,
-            },
-            AttrOptionalExpr {
+                then_value: value_a,
+                else_value: else_value_a,
+            }),
+            IfExpr(IfAttributeValue {
                 condition: cond_b,
-                value: value_b,
-            },
+                then_value: value_b,
+                else_value: else_value_b,
+            }),
         ) if cond_a == cond_b => {
             // If the condition is the same, we can hotreload it
             score_attr_value(value_a, value_b)
+                + match (else_value_a, else_value_b) {
+                    (Some(a), Some(b)) => score_attr_value(a, b),
+                    (None, None) => 0,
+                    _ => usize::MAX,
+                }
         }
 
         // todo: we should try and score recursively if we can - templates need to propagate up their