Best fit arguments with YAS style

src/nest_utils.jl

@@ -212,19 +212,47 @@ function find_all_segment_splits(n::Int, k::Int)
 end
 
 """
-Finds the optimal placeholders to turn into a newlines such that the length of the arguments on each line is as close as possible while following margin constraints.
+Finds the optimal placeholders to turn into a newlines such that the length of
+the arguments on each line is as close as possible while following margin constraints.
+
+first_segment_offset: The offset of the first line of the segment.
+segment_offset: The offset of the rest of the segments.
+max_margin: The maximum margin allowed.
+
+The offsets are used with paired with segment lengths to determine if the margin is exceeded.
+The offsets for the first segment and the rest are typically the same value.
+
+A scenario in which they are not is YASStyle, where the first argument cannot be nested
+and so to determine the "best fit" the `first_segment_offset` is the length up until
+the the first placeholder. Example:
+
+```
+function foo(arg1, arg2, ....)
+                  ^(First Placeholder)
+    body
+end
+```
+
+length of "function foo(arg1," is the `first_segment_offset`.
 """
 function find_optimal_nest_placeholders(
     fst::FST,
-    start_line_offset::Int,
+    first_segment_offset::Int,
+    segment_offset::Int,
     max_margin::Int,
 )::Vector{Int}
     placeholder_inds = findall(n -> n.typ === PLACEHOLDER, fst.nodes)
+
+    # For 1 or fewer placeholders, all are optimal
     if length(placeholder_inds) <= 1
         return placeholder_inds
     end
-    newline_inds = findall(n -> n.typ === NEWLINE, fst.nodes)
 
+    # Split `placeholder_inds` at newlines.
+    # The first entry of `placeholder_groups` will contain all placeholders before the first
+    # newline, the second entry will contain all placeholders between the first and second
+    # newline, and so on.
+    newline_inds = findall(n -> n.typ === NEWLINE, fst.nodes)
     placeholder_groups = Vector{Int}[]
     i = 1
     current_group = Int[]
@@ -241,12 +269,14 @@ function find_optimal_nest_placeholders(
 
     # @info "groups" placeholder_groups
 
+    # Pass individual placeholder groups to the function below
     optimal_placeholders = Int[]
     for (i, g) in enumerate(placeholder_groups)
         optinds = find_optimal_nest_placeholders(
             fst,
             g,
-            start_line_offset,
+            i == 1 ? first_segment_offset : segment_offset,
+            segment_offset,
             max_margin,
             last_group = i == length(placeholder_groups),
         )
@@ -260,21 +290,23 @@ end
 function find_optimal_nest_placeholders(
     fst::FST,
     placeholder_inds::Vector{Int},
-    initial_offset::Int,
+    first_segment_offset::Int,
+    segment_offset::Int,
     max_margin::Int;
     last_group::Bool = false,
 )::Vector{Int}
     # Function to calculate the length of a segment
     segment_length =
         (start_idx::Int, end_idx::Int) -> begin
-            if placeholder_inds[end] == end_idx && last_group
+            if end_idx == placeholder_inds[end] && last_group
                 sum(length.(fst[start_idx:end])) + fst.extra_margin
             else
                 sum(length.(fst[start_idx:end_idx-1]))
             end
         end
 
     n = length(placeholder_inds)
+    # @info "" placeholder_inds
     dp = fill(0, n - 1, n - 1)
 
     # Initialize the lengths of segments with single placeholders
@@ -285,7 +317,6 @@ function find_optimal_nest_placeholders(
         end
     end
 
-    # @info "" dp placeholder_inds
 
     N = size(dp, 1)
 
@@ -325,18 +356,19 @@ function find_optimal_nest_placeholders(
         return best_split
     end
 
-    fst_line_offset = fst.indent
+    # @info "" offset dp max_margin
     # Calculate best splits for each number of segments s
     segments = Tuple{Int,Int}[]
     for s in 1:N
         segments = find_best_segments(s)
         fits = true
         for (i, s) in enumerate(segments)
-            if i == 1
-                fits &= fst_line_offset + dp[first(s), last(s)] <= max_margin
+            val = if i == 1
+                first_segment_offset + dp[first(s), last(s)]
             else
-                fits &= fst_line_offset + dp[first(s), last(s)] <= max_margin
+                segment_offset + dp[first(s), last(s)]
             end
+            fits &= val <= max_margin
         end
         fits && break
     end

src/styles/sciml/nest.jl

@@ -97,8 +97,7 @@ function _n_tuple!(ss::SciMLStyle, fst::FST, s::State)
         false
     end
 
-    optimal_placeholders =
-        find_optimal_nest_placeholders(fst, start_line_offset, s.opts.margin)
+    optimal_placeholders = find_optimal_nest_placeholders(fst, fst.indent, fst.indent, s.opts.margin)
 
     for i in optimal_placeholders
         fst[i] = Newline(length = fst[i].len)

src/styles/yas/nest.jl

@@ -22,6 +22,19 @@ function n_call!(ys::YASStyle, fst::FST, s::State)
 
     f = n -> n.typ === PLACEHOLDER || n.typ === NEWLINE
 
+    start_line_offset = s.line_offset
+    indent_offset = start_line_offset + sum(length.(fst[1:2]))
+    first_ph_idx = findfirst(n -> n.typ === PLACEHOLDER, fst.nodes::Vector)
+    first_line_indent_offset = start_line_offset
+    if first_ph_idx !== nothing
+        first_line_indent_offset += sum(length.(fst[1:first_ph_idx]))
+    end
+    optimal_placeholders = find_optimal_nest_placeholders(fst, first_line_indent_offset, indent_offset, s.opts.margin)
+
+    for i in optimal_placeholders
+        fst[i] = Newline(length = fst[i].len)
+    end
+
     nodes = fst.nodes::Vector
     for (i, n) in enumerate(nodes)
         if i == 3
@@ -52,6 +65,11 @@ function n_call!(ys::YASStyle, fst::FST, s::State)
             nest!(style, n, s)
         end
     end
+
+    s.line_offset = start_line_offset
+    walk(unnest!(style; dedent = false), fst, s)
+    s.line_offset = start_line_offset
+    walk(increment_line_offset!, fst, s)
 end
 n_curly!(ys::YASStyle, fst::FST, s::State) = n_call!(ys, fst, s)
 n_ref!(ys::YASStyle, fst::FST, s::State) = n_call!(ys, fst, s)

test/yas_style.jl

@@ -278,13 +278,13 @@
         var = fcall(arg1, arg2, arg3, # comment
                     arg4, arg5)"""
         @test yasfmt(str_, 4, 80, join_lines_based_on_source = false) == str
-        @test yasfmt(str_, 4, 29, join_lines_based_on_source = false) == str
+        @test yasfmt(str_, 4, 30, join_lines_based_on_source = false) == str
 
         str = """
         var = fcall(arg1, arg2,
                     arg3, # comment
                     arg4, arg5)"""
-        @test yasfmt(str_, 4, 28, join_lines_based_on_source = false) == str
+        @test yasfmt(str_, 4, 29, join_lines_based_on_source = false) == str
         @test yasfmt(str_, 4, 23, join_lines_based_on_source = false) == str
 
         str = """
@@ -446,7 +446,8 @@
 
     @testset "issue 321 - exponential inline comments !!!" begin
         str = """
-        scaled_ticks, mini, maxi = optimize_ticks(scale_func(lmin), scale_func(lmax); k_min=4, # minimum number of ticks
+        scaled_ticks, mini, maxi = optimize_ticks(scale_func(lmin), scale_func(lmax);
+                                                  k_min=4, # minimum number of ticks
                                                   k_max=8)"""
         @test yasfmt(str, 4, 92, whitespace_in_kwargs = false) == str
     end
@@ -501,7 +502,8 @@
         @test yasfmt(str_, 4, 22) == str_
 
         str = """
-        T[10 20; 30 40;
+        T[10 20;
+          30 40;
           50 60;
           10
           10]"""
@@ -772,4 +774,46 @@
         @test format_text(str, YASStyle(), variable_call_indent = ["SVector", "test"]) ==
               formatted_str2
     end
+
+    @testset "optimal nesting" begin
+        str = """
+        function foo(arg1, arg2, arg3, arg4, arg5)
+            return body
+        end
+        """
+
+        @test format_text(str, YASStyle(), margin = 42) == str
+
+        fstr = """
+        function foo(arg1, arg2, arg3,
+                     arg4, arg5)
+            return body
+        end
+        """
+        @test format_text(str, YASStyle(), margin = 41) == fstr
+        # should be 30? might be a unnesting off by 1 error
+        @test format_text(str, YASStyle(), margin = 31) == fstr
+
+        fstr = """
+        function foo(arg1, arg2,
+                     arg3,
+                     arg4, arg5)
+            return body
+        end
+        """
+        @test format_text(str, YASStyle(), margin = 30) == fstr
+        @test format_text(str, YASStyle(), margin = 24) == fstr
+
+        fstr = """
+        function foo(arg1,
+                     arg2,
+                     arg3,
+                     arg4,
+                     arg5)
+            return body
+        end
+        """
+        @test format_text(str, YASStyle(), margin = 23) == fstr
+        @test format_text(str, YASStyle(), margin = 10) == fstr
+    end
 end