Fix SciMLStyle regressions

src/styles/default/nest.jl

@@ -951,6 +951,14 @@ function n_binaryopcall!(
         # Undo nest if possible
         if can_nest(fst) && !no_unnest(rhs) && !src_diff_line
             line_margin = s.line_offset
+            _, rhs_has_newline = length_to(rhs, (NEWLINE,))
+            if op_kind(fst) === K"=>" &&
+               rhs_has_newline &&
+               fst[1].indent > 0 &&
+               fst[1].line_offset <= fst[1].indent + 1 &&
+               fst[1].indent < line_offset
+                line_margin -= line_offset - fst[1].indent
+            end
 
             # replace IN with all of precedence level 6
             if (rhs.typ === Binary && !(op_kind(rhs) in KSet"in ::")) ||
@@ -978,7 +986,15 @@ function n_binaryopcall!(
                 fst[i1] = Whitespace(1)
                 if indent_nest || style isa YASStyle
                     fst[i2] = Whitespace(0)
-                    walk(unnest!(style; dedent = true), rhs, s)
+                    unnest_style =
+                        if op_kind(fst) === K"=>" &&
+                           s.opts.yas_style_nesting &&
+                           !(style isa YASStyle)
+                            YASStyle(style)
+                        else
+                            style
+                        end
+                    walk(unnest!(unnest_style; dedent = true), rhs, s)
                 end
             end
         end

src/styles/sciml/nest.jl

@@ -44,6 +44,195 @@ end
 #     nest!(style, fst[end], s, lineage)
 # end
 
+function _is_for_tuple_binding(
+    fst::FST,
+    s::State,
+    lineage::Vector{Tuple{FNode,Union{Nothing,Metadata}}},
+)
+    length(lineage) >= 2 &&
+        lineage[end-1][1] === CartesianIterator &&
+        op_kind(fst) in KSet"in ∈" &&
+        fst[1].typ === TupleN &&
+        s.line_offset + length(fst[1]) <= s.opts.margin
+end
+
+function _nearest_binary_is_assignment(
+    lineage::Vector{Tuple{FNode,Union{Nothing,Metadata}}},
+)
+    i = findlast(x -> x[1] === Binary, lineage)
+    i === nothing && return false
+
+    metadata = lineage[i][2]
+    return !isnothing(metadata) && metadata.is_assignment
+end
+
+function _align_tuple_comments!(fst::FST)
+    for n in fst.nodes::Vector
+        n.typ === NOTCODE && (n.indent = fst.indent)
+    end
+end
+
+function _is_dict_call(fst::FST)
+    fst.typ === Call &&
+        length(fst.nodes::Vector) > 0 &&
+        fst[1].typ === IDENTIFIER &&
+        fst[1].val == "Dict"
+end
+
+function _is_tuple_pair(fst::FST)
+    fst.typ === Binary && op_kind(fst) === K"=>" && fst[end].typ === TupleN
+end
+
+function _align_dict_tuple_pair_arrows!(fst::FST)
+    pair_inds = findall(_is_tuple_pair, fst.nodes::Vector)
+    length(pair_inds) > 1 || return
+
+    align_len = maximum(node_align_length(fst[i][1]) for i in pair_inds)
+    for i in pair_inds
+        pair = fst[i]
+        ws = align_len - node_align_length(pair[1]) + 1
+        align_binaryopcall!(pair, ws)
+    end
+end
+
+function _tuple_rhs_line_fits(rhs::FST, indent::Int, s::State)
+    nodes = rhs.nodes::Vector
+    idx = findfirst(
+        n ->
+            !(n.typ in (PUNCTUATION, WHITESPACE, PLACEHOLDER, NEWLINE, TRAILINGCOMMA)) && !is_closer(n),
+        nodes,
+    )
+    isnothing(idx) && return true
+
+    width, _ = length_to(rhs, (NEWLINE, PLACEHOLDER); start = idx)
+    return indent + width + rhs.extra_margin <= s.opts.margin
+end
+
+function _align_pair_tuple_rhs!(
+    fst::FST,
+    s::State;
+    align_to_pair::Bool = true,
+    fallback_indent::Int = fst.indent + s.opts.indent,
+)
+    rhs = fst[end]
+    rhs.typ === TupleN || return
+
+    desired_indent = if !align_to_pair
+        fallback_indent
+    elseif length(rhs.nodes::Vector) > 1 && rhs[2].typ === NEWLINE
+        rhs.indent - 1
+    else
+        fst.indent + node_align_length(fst[1:(end-1)]) + 1
+    end
+    add_indent!(rhs, s, desired_indent - rhs.indent)
+    _align_tuple_comments!(rhs)
+end
+
+function _unnest_pair_tuple_body!(style::AbstractStyle, rhs::FST, s::State)
+    lo = s.line_offset
+    s.line_offset = rhs.indent
+    for n in rhs.nodes::Vector
+        n.typ === NEWLINE && (s.line_offset = rhs.indent; continue)
+        walk(unnest!(style; dedent = false), n, s)
+    end
+    s.line_offset = lo
+end
+
+function _preserve_pair_tuple_newlines!(fst::FST; closer_indent::Int = fst.indent)
+    rhs = fst[end]
+    nodes = rhs.nodes::Vector
+    length(nodes) > 2 || return
+
+    nodes[2].typ !== NEWLINE && (nodes[2] = Newline(; length = nodes[2].len))
+    if is_closer(nodes[end]) && nodes[end-1].typ !== NEWLINE
+        nodes[end-1] = Newline(; length = nodes[end-1].len)
+        nodes[end].indent = closer_indent
+    end
+end
+
+function _preserve_multiline_closer!(fst::FST)
+    nodes = fst.nodes::Vector
+    length(nodes) > 2 && is_closer(nodes[end]) || return
+
+    nodes[end-1].typ !== NEWLINE && (nodes[end-1] = Newline(; length = nodes[end-1].len))
+end
+
+function _unnest_short_binary_lines!(fst::FST, s::State)
+    is_leaf(fst) && return
+
+    if fst.typ === Binary && !contains_comment(fst) && fst.line_offset >= 0
+        nl_inds = findall(n -> n.typ === NEWLINE, fst.nodes::Vector)
+        if length(nl_inds) > 0 &&
+           fst.line_offset + fst.extra_margin + node_align_length(fst) <= s.opts.margin
+            nl_to_ws!(fst, nl_inds)
+        end
+    end
+
+    for n in fst.nodes::Vector
+        _unnest_short_binary_lines!(n, s)
+    end
+end
+
+function n_binaryopcall!(
+    ss::SciMLStyle,
+    fst::FST,
+    s::State,
+    lineage::Vector{Tuple{FNode,Union{Nothing,Metadata}}};
+    indent::Int = -1,
+)
+    if op_kind(fst) === K"=>" &&
+       fst[end].typ === TupleN &&
+       fst[1].endline == fst[end].startline
+        style = getstyle(ss)
+        rhs_style = YASStyle(style)
+        nodes = fst.nodes::Vector
+        oplen = sum(length.(fst[2:end]))
+        line_offset = s.line_offset
+        fallback_indent = line_offset + s.opts.indent
+        nested = false
+        for (i, n) in enumerate(nodes)
+            if n.typ === NEWLINE
+                s.line_offset = fst.indent
+            elseif i == 1
+                n.extra_margin = oplen + fst.extra_margin
+                nested |= nest!(style, n, s, lineage)
+            elseif i == length(nodes)
+                n.extra_margin = fst.extra_margin
+                align_to_pair = _tuple_rhs_line_fits(n, s.line_offset, s)
+                n.indent = align_to_pair ? s.line_offset : fallback_indent
+                nested |= nest!(rhs_style, n, s, lineage)
+                _align_pair_tuple_rhs!(fst, s; align_to_pair, fallback_indent)
+                if align_to_pair
+                    lo = s.line_offset
+                    walk(unnest!(rhs_style; dedent = false), n, s)
+                    s.line_offset = lo
+                    _unnest_short_binary_lines!(n, s)
+                else
+                    _unnest_pair_tuple_body!(rhs_style, n, s)
+                    _preserve_pair_tuple_newlines!(fst; closer_indent = line_offset)
+                end
+                _align_tuple_comments!(n)
+            else
+                nested |= nest!(style, n, s, lineage)
+            end
+        end
+        return nested
+    end
+
+    if _is_for_tuple_binding(fst, s, lineage)
+        lhs = fst[1]
+        nest_behavior = lhs.nest_behavior
+        lhs.nest_behavior = NeverNest
+        try
+            return n_binaryopcall!(DefaultStyle(getstyle(ss)), fst, s, lineage; indent)
+        finally
+            lhs.nest_behavior = nest_behavior
+        end
+    end
+
+    n_binaryopcall!(DefaultStyle(getstyle(ss)), fst, s, lineage; indent)
+end
+
 function n_functiondef!(
     ss::SciMLStyle,
     fst::FST,
@@ -96,15 +285,56 @@ function n_macro!(
     n_functiondef!(ss, fst, s, lineage)
 end
 
+function _is_multiline_typed_ref(fst::FST)
+    fst.typ === RefN &&
+        length(fst.nodes::Vector) > 1 &&
+        fst[1].typ === Curly &&
+        any(n -> n.typ === NEWLINE, fst.nodes::Vector)
+end
+
+function _has_multiline_do_args(fst::FST)
+    length(fst.nodes::Vector) >= 5 &&
+        fst[4].typ === WHITESPACE &&
+        !is_leaf(fst[5]) &&
+        any(n -> n.typ === NEWLINE, fst[5].nodes::Vector)
+end
+
+function n_do!(
+    ss::SciMLStyle,
+    fst::FST,
+    s::State,
+    lineage::Vector{Tuple{FNode,Union{Nothing,Metadata}}},
+)
+    style = getstyle(ss)
+    extra_margin = sum(length.(fst[2:3]))
+    if fst[4].typ === WHITESPACE
+        extra_margin += length(fst[4])
+        if !_has_multiline_do_args(fst)
+            extra_margin += length(fst[5])
+        end
+    end
+    fst[1].extra_margin = fst.extra_margin + extra_margin
+
+    nested = false
+    nested |= nest!(style, fst[1], s, lineage)
+    nested |=
+        nest!(style, fst[2:end], s, fst.indent, lineage; extra_margin = fst.extra_margin)
+    return nested
+end
+
 function _n_tuple!(
     ss::SciMLStyle,
     fst::FST,
     s::State,
     lineage::Vector{Tuple{FNode,Union{Nothing,Metadata}}},
 )
     style = getstyle(ss)
-    line_margin = s.line_offset + length(fst) + fst.extra_margin
     nodes = fst.nodes::Vector
+    _is_dict_call(fst) &&
+        fst.startline != fst.endline &&
+        _align_dict_tuple_pair_arrows!(fst)
+
+    line_margin = s.line_offset + length(fst) + fst.extra_margin
     has_closer = is_closer(fst[end])
     start_line_offset = s.line_offset
 
@@ -224,6 +454,7 @@ function _n_tuple!(
 
     s.line_offset = start_line_offset
     walk(unnest!(style; dedent = false), fst, s)
+    _is_dict_call(fst) && fst.startline != fst.endline && _preserve_multiline_closer!(fst)
     s.line_offset = start_line_offset
     walk(increment_line_offset!, fst, s)
 
@@ -237,40 +468,21 @@ function n_ref!(
     s::State,
     lineage::Vector{Tuple{FNode,Union{Nothing,Metadata}}},
 )
-    # Check if this RefN is the LHS of an assignment
-    # Look through the lineage to see if we have a Binary assignment parent
-    # and this RefN comes before any other Binary operators
-    is_lhs_of_assignment = false
-
-    if length(lineage) >= 2
-        # Check if we have a Binary assignment in the lineage
-        for i in length(lineage):-1:1
-            if lineage[i][1] === Binary &&
-               !isnothing(lineage[i][2]) &&
-               lineage[i][2].is_assignment
-                # Check if there are any other Binary nodes between us and the assignment
-                has_intermediate_binary = false
-                for j in (i+1):length(lineage)
-                    if lineage[j][1] === Binary
-                        has_intermediate_binary = true
-                        break
-                    end
-                end
-
-                if !has_intermediate_binary
-                    is_lhs_of_assignment = true
-                end
-                break
+    if _nearest_binary_is_assignment(lineage) && fst.extra_margin > 0
+        # Don't break the LHS of an assignment
+        # Format children but keep them on the same line
+        nodes = fst.nodes::Vector{FST}
+        for (i, n) in enumerate(nodes)
+            if n.typ === NEWLINE &&
+               (i == 1 || !is_comment(nodes[i-1])) &&
+               (i == length(nodes) || !is_comment(nodes[i+1]))
+                nodes[i] = Whitespace(n.len)
             end
         end
-    end
 
-    if is_lhs_of_assignment
-        # Don't break the LHS of an assignment
-        # Format children but keep them on the same line
         lo = s.line_offset
         nested = false
-        for (i, n) in enumerate(fst.nodes)
+        for n in nodes
             nested |= nest!(ss, n, s, lineage)
             if n.typ !== NEWLINE  # Prevent any newlines
                 s.line_offset += length(n)
@@ -280,6 +492,10 @@ function n_ref!(
         return nested
     end
 
+    if _is_multiline_typed_ref(fst)
+        return n_ref!(YASStyle(getstyle(ss)), fst, s, lineage)
+    end
+
     # Otherwise use the default behavior
     if s.opts.yas_style_nesting
         return n_ref!(YASStyle(getstyle(ss)), fst, s, lineage)