Repository: domluna/JuliaFormatter.jl
Pull request #132 (2019-11-24 00:19:41)

Unnesting
_unnesting_ can occur when the RHS an assignment or function definition
is dedented, which may allow a previously nested expression to fit in the margin at full
length.

Motivating example:

```julia
let X = LinearAlgebra.Symmetric{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where {T},
    Y = Union{
	LinearAlgebra.Hermitian{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
	LinearAlgebra.Symmetric{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
    }

    @test X <: Y
end
```

The above fits in a margin of 92 and the arguments inside of

```julia
    ...
    Y = Union{
	LinearAlgebra.Hermitian{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
	LinearAlgebra.Symmetric{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
    }
    ...
```

fit in a margin of 90.

Previously `Y = Union{...` would initially be nested to

```julia
Y =
    Union{...
```

The formatter then notices it could fit the caller on the initial line
and the arguments would be dedented.

```julia
Y = Union{
    ...
}
```

However, there was no recursive check to see if nodes could be unnested
given the additional space. Now there is.

This also adds indexing methods from Base to make things a little
easier.


Patch:
diff --git a/src/nest.jl b/src/nest.jl
index f8139eaf..26f902fe 100644
--- a/src/nest.jl
+++ b/src/nest.jl
@@ -5,7 +5,7 @@
 # This is done by replacing `PLACEHOLDER` nodes with `NEWLINE` 
 # nodes and updating the PTree's indent.
 #
-# `extra_width` provides additional width to consider from
+# `extra_margin` provides additional width to consider from
 # the top-level node.
 #
 # Example:
@@ -14,24 +14,41 @@
 #
 # the length of " op" will be considered when nesting arg1
 
-function skip_indent(x::PTree)
-    if x.typ === CSTParser.LITERAL && x.val == ""
+# unnest, converts newlines to whitespace
+unnest!(pt::PTree, ind::Int) = pt[ind] = Whitespace(pt[ind].len)
+function unnest!(pt::PTree, nl_inds::Vector{Int})
+    for (i, ind) in enumerate(nl_inds)
+        unnest!(pt, ind)
+        i == length(nl_inds) || continue
+        ind2 = ind - 1
+        if pt[ind2].typ === TRAILINGCOMMA
+            pt[ind2].val = ""
+            pt[ind2].len = 0
+        elseif pt[ind-1].typ === TRAILINGSEMICOLON
+            pt[ind2].val = ";"
+            pt[ind2].len = 1
+        end
+    end
+end
+
+function skip_indent(pt::PTree)
+    if pt.typ === CSTParser.LITERAL && pt.val == ""
         return true
-    elseif x.typ === NEWLINE || x.typ === NOTCODE
+    elseif pt.typ === NEWLINE || pt.typ === NOTCODE
         return true
     end
     false
 end
 
-function walk(f, x::PTree, s::State)
-    f(x, s)
-    is_leaf(x) && return
-    for (i, n) in enumerate(x.nodes)
-        if n.typ === NEWLINE && i < length(x.nodes)
-            if is_closer(x.nodes[i+1])
-                s.line_offset = x.nodes[i+1].indent
-            elseif !skip_indent(x.nodes[i+1])
-                s.line_offset = x.indent
+function walk(f, pt::PTree, s::State)
+    f(pt, s)
+    is_leaf(pt) && return
+    for (i, n) in enumerate(pt.nodes)
+        if n.typ === NEWLINE && i < length(pt.nodes)
+            if is_closer(pt[i+1])
+                s.line_offset = pt[i+1].indent
+            elseif !skip_indent(pt[i+1])
+                s.line_offset = pt.indent
             end
         else
             walk(f, n, s)
@@ -40,7 +57,7 @@ function walk(f, x::PTree, s::State)
 end
 
 function reset_line_offset!(x::PTree, s::State)
-    !is_leaf(x) && return
+    is_leaf(x) || return
     s.line_offset += length(x)
 end
 
@@ -52,7 +69,30 @@ function add_indent!(x::PTree, s::State, indent)
     s.line_offset = lo
 end
 
-function nest!(nodes::Vector{PTree}, s::State, indent; extra_width = 0)
+function dedent!(x::PTree, s::State)
+    if is_leaf(x)
+        s.line_offset += length(x)
+        is_closer(x) && (x.indent -= s.indent_size)
+        return
+    end
+    x.typ === CSTParser.ConditionalOpCall && return
+    x.typ === CSTParser.Comparison && return
+    x.typ === CSTParser.ChainOpCall && return
+    x.typ === CSTParser.BinaryOpCall && return
+    x.typ === CSTParser.StringH && return
+
+    # dedent
+    x.indent -= s.indent_size
+
+    nl_inds = findall(n -> n.typ === NEWLINE && !n.force_nest, x.nodes)
+    length(nl_inds) > 0 || return
+    margin = s.line_offset + x.extra_margin + length(x)
+    # @info "" x.typ margin s.line_offset x.extra_margin length(x) length(nl_inds)
+    margin <= s.margin || return
+    unnest!(x, nl_inds)
+end
+
+function nest!(nodes::Vector{PTree}, s::State, indent::Int; extra_margin = 0)
     for (i, n) in enumerate(nodes)
         if n.typ === NEWLINE && nodes[i+1].typ === CSTParser.Block
             s.line_offset = nodes[i+1].indent
@@ -61,94 +101,97 @@ function nest!(nodes::Vector{PTree}, s::State, indent; extra_width = 0)
         elseif n.typ === NEWLINE
             s.line_offset = indent
         else
-            nest!(n, s, extra_width = extra_width)
+            n.extra_margin = extra_margin
+            nest!(n, s)
         end
     end
 end
 
-function nest!(x::PTree, s::State; extra_width = 0)
+function nest!(x::PTree, s::State)
     if is_leaf(x)
         s.line_offset += length(x)
         return
     end
 
     if x.typ === CSTParser.Import
-        n_import!(x, s, extra_width = extra_width)
+        n_import!(x, s)
     elseif x.typ === CSTParser.Export
-        n_import!(x, s, extra_width = extra_width)
+        n_import!(x, s)
     elseif x.typ === CSTParser.Using
-        n_import!(x, s, extra_width = extra_width)
+        n_import!(x, s)
     elseif x.typ === CSTParser.ImportAll
-        n_import!(x, s, extra_width = extra_width)
+        n_import!(x, s)
     elseif x.typ === CSTParser.WhereOpCall
-        n_wherecall!(x, s, extra_width = extra_width)
+        n_wherecall!(x, s)
     elseif x.typ === CSTParser.ConditionalOpCall
-        n_condcall!(x, s, extra_width = extra_width)
+        n_condcall!(x, s)
     elseif x.typ === CSTParser.BinaryOpCall
-        n_binarycall!(x, s, extra_width = extra_width)
+        n_binarycall!(x, s)
     elseif x.typ === CSTParser.Curly
-        n_call!(x, s, extra_width = extra_width)
+        n_call!(x, s)
     elseif x.typ === CSTParser.Call
-        n_call!(x, s, extra_width = extra_width)
+        n_call!(x, s)
     elseif x.typ === CSTParser.MacroCall
-        n_call!(x, s, extra_width = extra_width)
+        n_call!(x, s)
     elseif x.typ === CSTParser.Ref
-        n_call!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.ChainOpCall
-        n_chainopcall!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.Comparison
-        n_comparison!(x, s, extra_width = extra_width)
+        n_call!(x, s)
+    elseif x.typ === CSTParser.TypedVcat
+        n_call!(x, s)
     elseif x.typ === CSTParser.TupleH
-        n_tuple!(x, s, extra_width = extra_width)
+        n_tuple!(x, s)
     elseif x.typ === CSTParser.Vect
-        n_tuple!(x, s, extra_width = extra_width)
+        n_tuple!(x, s)
     elseif x.typ === CSTParser.Vcat
-        n_tuple!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.Block
-        n_block!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.For
-        n_for!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.Let
-        n_for!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.TypedVcat
-        n_call!(x, s, extra_width = extra_width)
+        n_tuple!(x, s)
     elseif x.typ === CSTParser.Parameters
-        n_tuple!(x, s, extra_width = extra_width)
+        n_tuple!(x, s)
     elseif x.typ === CSTParser.Braces
-        n_tuple!(x, s, extra_width = extra_width)
+        n_tuple!(x, s)
     elseif x.typ === CSTParser.InvisBrackets
-        n_tuple!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.UnaryOpCall && x.nodes[2].typ === CSTParser.OPERATOR
-        nest!(x.nodes[1], s, extra_width = extra_width + length(x.nodes[2]))
-        nest!(x.nodes[2], s)
+        n_tuple!(x, s)
+    elseif x.typ === CSTParser.Comprehension
+        n_tuple!(x, s)
     elseif x.typ === CSTParser.Do
-        n_do!(x, s, extra_width = extra_width)
+        n_do!(x, s)
     elseif x.typ === CSTParser.Generator
-        n_gen!(x, s, extra_width = extra_width)
+        n_gen!(x, s)
     elseif x.typ === CSTParser.Filter
-        n_gen!(x, s, extra_width = extra_width)
-    elseif x.typ === CSTParser.Comprehension
-        n_tuple!(x, s, extra_width = extra_width)
+        n_gen!(x, s)
+    elseif x.typ === CSTParser.Block
+        n_block!(x, s)
+    elseif x.typ === CSTParser.ChainOpCall
+        n_block!(x, s)
+    elseif x.typ === CSTParser.Comparison
+        n_block!(x, s)
+    elseif x.typ === CSTParser.For
+        n_for!(x, s)
+    elseif x.typ === CSTParser.Let
+        n_for!(x, s)
+    elseif x.typ === CSTParser.UnaryOpCall && x[2].typ === CSTParser.OPERATOR
+        x[1].extra_margin = x.extra_margin + length(x[2])
+        nest!(x[1], s)
+        nest!(x[2], s)
     else
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
     end
 end
 
-function n_do!(x, s; extra_width = 0)
-    ew = sum(length.(x.nodes[2:3]))
+function n_do!(x, s)
+    extra_margin = sum(length.(x[2:3]))
     # make sure there are nodes after "do"
-    if x.nodes[4].typ === WHITESPACE
-        ew += length(x.nodes[4])
-        ew += length(x.nodes[5])
+    if x[4].typ === WHITESPACE
+        extra_margin += length(x[4])
+        extra_margin += length(x[5])
     end
-    nest!(x.nodes[1], s, extra_width = extra_width + ew)
-    nest!(x.nodes[2:end], s, x.indent, extra_width = extra_width)
+    x[1].extra_margin = x.extra_margin + extra_margin
+    nest!(x[1], s)
+    nest!(x[2:end], s, x.indent, extra_margin = x.extra_margin)
 end
 
 
 # Import,Using,Export,ImportAll
-function n_import!(x, s; extra_width = 0)
-    line_margin = s.line_offset + length(x) + extra_width
+function n_import!(x, s)
+    line_margin = s.line_offset + length(x) + x.extra_margin
     idx = findfirst(n -> n.typ === PLACEHOLDER, x.nodes)
     if idx !== nothing && (line_margin > s.margin || x.force_nest)
         sidx = findfirst(is_colon, x.nodes)
@@ -158,33 +201,33 @@ function n_import!(x, s; extra_width = 0)
             # for the WHITESPACE node after the colon
             sidx += 1
         end
-        x.indent = s.line_offset + sum(length.(x.nodes[1:sidx]))
+        x.indent = s.line_offset + sum(length.(x[1:sidx]))
         s.line_offset = x.indent
 
         for (i, n) in enumerate(x.nodes)
             if n.typ === NEWLINE
                 s.line_offset = x.indent
             elseif n.typ === PLACEHOLDER
-                x.nodes[i] = Newline()
+                x[i] = Newline(length = n.len)
                 s.line_offset = x.indent
             else
                 nest!(n, s)
             end
         end
     else
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
     end
 end
 
-function n_tuple!(x, s; extra_width = 0)
-    line_margin = s.line_offset + length(x) + extra_width
+function n_tuple!(x, s)
+    line_margin = s.line_offset + length(x) + x.extra_margin
     idx = findlast(n -> n.typ === PLACEHOLDER, x.nodes)
-    # @info "ENTERING" idx x.typ s.line_offset length(x) extra_width
-    opener = length(x.nodes) > 0 ? is_opener(x.nodes[1]) : false
+    # @info "ENTERING" idx x.typ s.line_offset length(x) x.extra_margin
+    opener = length(x.nodes) > 0 ? is_opener(x[1]) : false
     if idx !== nothing && (line_margin > s.margin || x.force_nest)
         line_offset = s.line_offset
         if opener
-            x.nodes[end].indent = x.indent
+            x[end].indent = x.indent
         end
         if x.typ !== CSTParser.Parameters && !(x.typ === CSTParser.TupleH && !opener)
             x.indent += s.indent_size
@@ -195,43 +238,45 @@ function n_tuple!(x, s; extra_width = 0)
             if n.typ === NEWLINE
                 s.line_offset = x.indent
             elseif n.typ === PLACEHOLDER
-                x.nodes[i] = Newline()
+                x[i] = Newline(length = n.len)
                 s.line_offset = x.indent
             elseif n.typ === TRAILINGCOMMA
-                x.nodes[i].val = ","
-                x.nodes[i].len = 1
-                nest!(x.nodes[i], s)
+                n.val = ","
+                n.len = 1
+                nest!(n, s)
             elseif n.typ === TRAILINGSEMICOLON
-                x.nodes[i].val = ""
-                x.nodes[i].len = 0
-                nest!(x.nodes[i], s)
+                n.val = ""
+                n.len = 0
+                nest!(n, s)
             elseif opener && (i == 1 || i == length(x.nodes))
                 nest!(n, s)
             else
-                diff = x.indent - x.nodes[i].indent
+                diff = x.indent - x[i].indent
                 add_indent!(n, s, diff)
-                nest!(n, s, extra_width = 1)
+                n.extra_margin = 1
+                nest!(n, s)
             end
         end
 
         if opener
-            s.line_offset = x.nodes[end].indent + 1
+            s.line_offset = x[end].indent + 1
         end
-        # @info "EXITING" x.typ s.line_offset x.indent x.nodes[end].indent
+        # @info "EXITING" x.typ s.line_offset x.indent x[end].indent
     else
-        opener && (extra_width += 1)
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        extra_margin = x.extra_margin
+        opener && (extra_margin += 1)
+        nest!(x.nodes, s, x.indent, extra_margin = extra_margin)
     end
 end
 
 
-function n_call!(x, s; extra_width = 0)
-    line_margin = s.line_offset + length(x) + extra_width
+function n_call!(x, s)
+    line_margin = s.line_offset + length(x) + x.extra_margin
     idx = findlast(n -> n.typ === PLACEHOLDER, x.nodes)
-    # @info "ENTERING" x.typ s.line_offset length(x) extra_width s.margin x.nodes[1].val
+    # @info "ENTERING" x.typ s.line_offset length(x) x.extra_margin s.margin x[1].val
     if idx !== nothing && (line_margin > s.margin || x.force_nest)
         line_offset = s.line_offset
-        x.nodes[end].indent = x.indent
+        x[end].indent = x.indent
         x.indent += s.indent_size
 
         # @info "DURING" x.typ x.indent s.line_offset
@@ -239,96 +284,99 @@ function n_call!(x, s; extra_width = 0)
             if n.typ === NEWLINE
                 s.line_offset = x.indent
             elseif n.typ === PLACEHOLDER
-                x.nodes[i] = Newline()
+                x[i] = Newline(length = n.len)
                 s.line_offset = x.indent
             elseif n.typ === TRAILINGCOMMA
-                x.nodes[i].val = ","
-                x.nodes[i].len = 1
-                nest!(x.nodes[i], s)
+                n.val = ","
+                n.len = 1
+                nest!(n, s)
             elseif n.typ === TRAILINGSEMICOLON
-                x.nodes[i].val = ""
-                x.nodes[i].len = 0
-                nest!(x.nodes[i], s)
+                n.val = ""
+                n.len = 0
+                nest!(n, s)
             elseif i == 1 || i == length(x.nodes)
                 n.typ === CSTParser.Parameters && (n.force_nest = true)
-                nest!(n, s, extra_width = 1)
+                n.extra_margin = 1
+                nest!(n, s)
             else
-                diff = x.indent - x.nodes[i].indent
+                diff = x.indent - x[i].indent
                 add_indent!(n, s, diff)
                 n.typ === CSTParser.Parameters && (n.force_nest = true)
-                nest!(n, s, extra_width = 1)
+                n.extra_margin = 1
+                nest!(n, s)
             end
         end
 
-        s.line_offset = x.nodes[end].indent + 1
+        s.line_offset = x[end].indent + 1
     else
-        is_closer(x.nodes[end]) && (extra_width += 1)
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        extra_margin = x.extra_margin
+        is_closer(x[end]) && (extra_margin += 1)
+        nest!(x.nodes, s, x.indent, extra_margin = extra_margin)
     end
 end
 
 # "A where B"
-function n_wherecall!(x, s; extra_width = 0)
-    line_margin = s.line_offset + length(x) + extra_width
-    # @info "" s.line_offset x.typ line_margin extra_width length(x) x.force_nest line_margin > s.margin
+function n_wherecall!(x, s)
+    line_margin = s.line_offset + length(x) + x.extra_margin
+    # @info "" s.line_offset x.typ line_margin x.extra_margin length(x) x.force_nest line_margin > s.margin
     if line_margin > s.margin || x.force_nest
         line_offset = s.line_offset
         # after "A where "
         idx = findfirst(n -> n.typ === PLACEHOLDER, x.nodes)
-        Blen = sum(length.(x.nodes[idx+1:end]))
+        Blen = sum(length.(x[idx+1:end]))
 
         # A, +7 is the length of " where "
-        ew = Blen + 7 + extra_width
-        nest!(x.nodes[1], s, extra_width = ew)
+        x[1].extra_margin = Blen + 7 + x.extra_margin
+        nest!(x[1], s)
 
-        for (i, n) in enumerate(x.nodes[2:idx-1])
+        for (i, n) in enumerate(x[2:idx-1])
             nest!(n, s)
         end
 
         # "B"
 
-        has_braces = is_closer(x.nodes[end])
+        has_braces = is_closer(x[end])
         if has_braces
-            x.nodes[end].indent = x.indent
+            x[end].indent = x.indent
         end
 
-        over = (s.line_offset + Blen + extra_width > s.margin) || x.force_nest
+        over = (s.line_offset + Blen + x.extra_margin > s.margin) || x.force_nest
         x.indent += s.indent_size
 
-        last_typ::Union{CSTParser.Head,Nothing} = nothing
-        for (i, n) in enumerate(x.nodes[idx+1:end])
+        for (i, n) in enumerate(x[idx+1:end])
             if n.typ === NEWLINE
                 s.line_offset = x.indent
             elseif is_opener(n)
                 if x.indent - s.line_offset > 1
                     x.indent = s.line_offset + 1
-                    x.nodes[end].indent = s.line_offset
+                    x[end].indent = s.line_offset
                 end
                 nest!(n, s)
             elseif n.typ === PLACEHOLDER && over
-                x.nodes[i+idx] = Newline()
+                x[i+idx] = Newline(length = n.len)
                 s.line_offset = x.indent
             elseif n.typ === TRAILINGCOMMA && over
-                x.nodes[i+idx].val = ","
-                x.nodes[i+idx].len = 1
-                nest!(x.nodes[i+idx], s)
+                n.val = ","
+                n.len = 1
+                nest!(n, s)
             elseif has_braces
-                nest!(n, s, extra_width = 1 + extra_width)
+                n.extra_margin = 1 + x.extra_margin
+                nest!(n, s)
             else
-                nest!(n, s, extra_width = extra_width)
+                n.extra_margin = x.extra_margin
+                nest!(n, s)
             end
-            (!is_leaf(n) || n.typ === CSTParser.IDENTIFIER) && (last_typ = n.typ)
         end
 
         s.line_offset = line_offset
         walk(reset_line_offset!, x, s)
     else
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
     end
 end
 
-function n_condcall!(x, s; extra_width = 0)
-    line_margin = s.line_offset + length(x) + extra_width
+function n_condcall!(x, s)
+    line_margin = s.line_offset + length(x) + x.extra_margin
     # @info "ENTERING" x.typ s.line_offset line_margin x.force_nest
     if line_margin > s.margin || x.force_nest
         line_offset = s.line_offset
@@ -336,24 +384,25 @@ function n_condcall!(x, s; extra_width = 0)
         phs = reverse(findall(n -> n.typ === PLACEHOLDER, x.nodes))
         for (i, idx) in enumerate(phs)
             if i == 1
-                x.nodes[idx] = Newline()
+                x[idx] = Newline(length = x[idx].len)
             else
                 nidx = phs[i-1]
-                l1 = sum(length.(x.nodes[1:idx-1]))
-                l2 = sum(length.(x.nodes[idx:nidx-1]))
+                l1 = sum(length.(x[1:idx-1]))
+                l2 = sum(length.(x[idx:nidx-1]))
                 width = line_offset + l1 + l2
                 if x.force_nest || width > s.margin
-                    x.nodes[idx] = Newline()
+                    x[idx] = Newline(length = x[idx].len)
                 end
             end
         end
 
         for (i, n) in enumerate(x.nodes)
             if i == length(x.nodes)
-                nest!(n, s, extra_width = extra_width)
-            elseif x.nodes[i+1].typ === WHITESPACE
-                ew = length(x.nodes[i+1]) + length(x.nodes[i+2])
-                nest!(n, s, extra_width = ew)
+                n.extra_margin = x.extra_margin
+                nest!(n, s)
+            elseif x[i+1].typ === WHITESPACE
+                n.extra_margin = length(x[i+1]) + length(x[i+2])
+                nest!(n, s)
             elseif n.typ === NEWLINE
                 s.line_offset = x.indent
             else
@@ -363,56 +412,32 @@ function n_condcall!(x, s; extra_width = 0)
 
         s.line_offset = line_offset
         for (i, n) in enumerate(x.nodes)
-            if n.typ === NEWLINE && !is_comment(x.nodes[i+1]) && !is_comment(x.nodes[i-1])
+            if n.typ === NEWLINE && !is_comment(x[i+1]) && !is_comment(x[i-1])
                 # +1 for newline to whitespace conversion
                 width = s.line_offset + 1
                 if i == length(x.nodes) - 1
-                    width += sum(length.(x.nodes[i+1:end])) + extra_width
+                    width += sum(length.(x[i+1:end])) + x.extra_margin
                 else
-                    width += sum(length.(x.nodes[i+1:i+3]))
+                    width += sum(length.(x[i+1:i+3]))
                 end
                 # @debug "" s.line_offset l  s.margin
                 if width <= s.margin
-                    x.nodes[i] = Whitespace(1)
+                    x[i] = Whitespace(1)
                 else
                     s.line_offset = x.indent
                 end
             end
-            s.line_offset += length(x.nodes[i])
+            s.line_offset += length(x[i])
         end
 
         s.line_offset = line_offset
         walk(reset_line_offset!, x, s)
     else
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
     end
 end
-n_comparison!(x, s; extra_width = 0) = n_block!(x, s, extra_width = extra_width)
-n_chainopcall!(x, s; extra_width = 0) = n_block!(x, s, extra_width = extra_width)
 
-function dedent!(x::PTree, s::State, pindent::Int, line_margin::Int)
-    is_leaf(x) && return
-    x.typ === CSTParser.ConditionalOpCall && return
-    x.typ === CSTParser.Comparison && return
-    x.typ === CSTParser.ChainOpCall && return
-    x.typ === CSTParser.BinaryOpCall && return
 
-    if is_iterable(x)
-        indent = x.indent
-        args = get_args(x)
-        while length(args) == 1 && is_iterable(args[1])
-            indent = args[1].indent
-            args = get_args(args[1])
-        end
-        diff = min(line_margin, pindent + s.indent_size) - indent
-        # @info "" pindent diff x.indent line_margin indent
-        # @info "" pindent diff x.indent x.typ
-        add_indent!(x, s, diff)
-        x.nodes[end].indent = pindent
-    else
-        add_indent!(x, s, -s.indent_size)
-    end
-end
 
 # arg1 op arg2
 #
@@ -420,44 +445,45 @@ end
 #
 # arg1 op
 # arg2
-function n_binarycall!(x, s; extra_width = 0)
+function n_binarycall!(x, s)
     # If there's no placeholder the binary call is not nestable
     idxs = findall(n -> n.typ === PLACEHOLDER, x.nodes)
-    line_margin = s.line_offset + length(x) + extra_width
-    # @info "ENTERING" x.typ extra_width s.line_offset length(x) idxs x.ref[][2]
+    line_margin = s.line_offset + length(x) + x.extra_margin
+    # @info "ENTERING" x.typ x.extra_margin s.line_offset length(x) idxs x.ref[][2]
     if length(idxs) == 2 && (line_margin > s.margin || x.force_nest)
         line_offset = s.line_offset
         i1 = idxs[1]
         i2 = idxs[2]
-        x.nodes[i1] = Newline()
+        x[i1] = Newline(length = x[i1].len)
         cst = x.ref[]
 
         has_eq = CSTParser.defines_function(cst) || nest_assignment(cst)
         if has_eq
             s.line_offset = x.indent + s.indent_size
-            x.nodes[i2] = Whitespace(s.indent_size)
-            add_indent!(x.nodes[end], s, s.indent_size)
+            x[i2] = Whitespace(s.indent_size)
+            add_indent!(x[end], s, s.indent_size)
         else
             x.indent = s.line_offset
         end
 
         # arg2
-        nest!(x.nodes[end], s, extra_width = extra_width)
+        x[end].extra_margin = x.extra_margin
+        nest!(x[end], s)
 
-        # arg1 op
+        # "arg1 op" arg2
         s.line_offset = line_offset
-        # extra_width for " op"
-        ew = length(x.nodes[2]) + length(x.nodes[3])
-        # @debug "DURING" ew s.line_offset x.typ
-        nest!(x.nodes[1], s, extra_width = ew)
-        for n in x.nodes[2:i1]
+
+        # extra margin for " op"
+        x[1].extra_margin = length(x[2]) + length(x[3])
+        nest!(x[1], s)
+        for n in x[2:i1]
             nest!(n, s)
         end
 
         # Undo nest if possible
         if !x.force_nest
-            arg2 = x.nodes[end]
-            arg2.typ === CSTParser.Block && (arg2 = arg2.nodes[1])
+            arg2 = x[end]
+            arg2.typ === CSTParser.Block && (arg2 = arg2[1])
             cst = arg2.ref[]
 
             if (
@@ -465,33 +491,27 @@ function n_binarycall!(x, s; extra_width = 0)
                     !(is_lazy_op(cst) && !has_eq) && cst[2].kind !== Tokens.IN
                 )
             ) || arg2.typ === CSTParser.UnaryOpCall || is_block(cst)
-                line_margin = s.line_offset + 1 + length(x.nodes[end])
+                line_margin = s.line_offset + 1 + length(x[end])
             else
                 rw, _ = length_to(x, [NEWLINE], start = i2 + 1)
-                line_margin = s.line_offset + 1 + rw
+                line_margin = s.line_offset + rw
             end
-            # @info "" can_unnest arg2.typ has_eq
 
-            if line_margin + extra_width <= s.margin
-                x.nodes[i1] = Whitespace(1)
+            # @info "" arg2.typ has_eq s.line_offset line_margin x.extra_margin length(x[end])
+
+            if line_margin + x.extra_margin <= s.margin
+                x[i1] = Whitespace(1)
                 if has_eq
-                    x.nodes[i2] = Placeholder(0)
-                    # recursive dedent
-                    if arg2.typ === CSTParser.BinaryOpCall
-                        dedent!(arg2.nodes[1], s, x.indent, line_margin)
-                        dedent!(arg2.nodes[end], s, x.indent, line_margin)
-                    else
-                        dedent!(arg2, s, x.indent, line_margin)
-                    end
+                    x[i2] = Whitespace(0)
+                    # s.line_offset = line_offset
+                    walk(dedent!, arg2, s)
                 end
             end
         end
 
         # @info "before reset" line_offset s.line_offset
-
         s.line_offset = line_offset
         walk(reset_line_offset!, x, s)
-
         # @info "after reset" line_offset s.line_offset
     else
         # Handles the case of a function def defined
@@ -520,52 +540,55 @@ function n_binarycall!(x, s; extra_width = 0)
         idx = findfirst(n -> n.typ === CSTParser.WhereOpCall, x.nodes)
         return_width = 0
         if idx !== nothing && idx > 1
-            return_width = length(x.nodes[idx].nodes[1]) + length(x.nodes[2])
+            return_width = length(x[idx][1]) + length(x[2])
         elseif idx === nothing
             return_width, _ = length_to(x, [NEWLINE], start = 2)
         end
 
         # @info "" return_width
-        # @info "" s.line_offset return_width length(x.nodes[1])
+        # @info "" s.line_offset return_width length(x[1])
 
         for (i, n) in enumerate(x.nodes)
             if n.typ === NEWLINE
                 s.line_offset = x.indent
             elseif i == 1
-                nest!(n, s, extra_width = return_width)
+                n.extra_margin = return_width
+                nest!(n, s)
             elseif i == idx
-                nest!(n, s, extra_width = extra_width)
+                n.extra_margin = x.extra_margin
+                nest!(n, s)
             else
-                nest!(n, s, extra_width = extra_width)
+                n.extra_margin = x.extra_margin
+                nest!(n, s)
             end
         end
     end
 end
 
-function n_for!(x, s; extra_width = 0)
+function n_for!(x, s)
     block_idx = findfirst(n -> !is_leaf(n), x.nodes)
     if block_idx === nothing
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
         return
     end
-    ph_idx = findfirst(n -> n.typ === PLACEHOLDER, x.nodes[block_idx].nodes)
-    nest!(x.nodes, s, x.indent, extra_width = extra_width)
+    ph_idx = findfirst(n -> n.typ === PLACEHOLDER, x[block_idx].nodes)
+    nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
 
     # return if the argument block was nested
-    ph_idx !== nothing && x.nodes[3].nodes[ph_idx].typ === NEWLINE && return
+    ph_idx !== nothing && x[3][ph_idx].typ === NEWLINE && return
 
     idx = 5
-    n = x.nodes[idx]
+    n = x[idx]
     if n.typ === NOTCODE && n.startline == n.endline
         res = get(s.doc.comments, n.startline, (0, ""))
         res == (0, "") && deleteat!(x.nodes, idx - 1)
     end
 end
 
-function n_block!(x, s; extra_width = 0, custom_indent = 0)
-    line_margin = s.line_offset + length(x) + extra_width
+function n_block!(x, s; custom_indent = 0)
+    line_margin = s.line_offset + length(x) + x.extra_margin
     idx = findfirst(n -> n.typ === PLACEHOLDER, x.nodes)
-    # @info "ENTERING" idx x.typ s.line_offset length(x) extra_width
+    # @info "ENTERING" idx x.typ s.line_offset length(x) x.extra_margin
     if idx !== nothing && (line_margin > s.margin || x.force_nest)
         line_offset = s.line_offset
 
@@ -576,28 +599,28 @@ function n_block!(x, s; extra_width = 0, custom_indent = 0)
             if n.typ === NEWLINE
                 s.line_offset = x.indent
             elseif n.typ === PLACEHOLDER
-                x.nodes[i] = Newline()
+                x[i] = Newline(length = n.len)
                 s.line_offset = x.indent
             elseif n.typ === TRAILINGCOMMA
-                x.nodes[i].val = ","
-                x.nodes[i].len = 1
-                nest!(x.nodes[i], s)
+                n.val = ","
+                n.len = 1
+                nest!(n, s)
             elseif n.typ === TRAILINGSEMICOLON
-                x.nodes[i].val = ""
-                x.nodes[i].len = 0
-                nest!(x.nodes[i], s)
+                n.val = ""
+                n.len = 0
+                nest!(n, s)
             else
-                diff = x.indent - x.nodes[i].indent
+                diff = x.indent - x[i].indent
                 add_indent!(n, s, diff)
-                nest!(n, s, extra_width = 1)
+                n.extra_margin = 1
+                nest!(n, s)
             end
         end
 
-        # @info "EXITING" x.typ s.line_offset x.indent x.nodes[end].indent
+        # @info "EXITING" x.typ s.line_offset x.indent x[end].indent
     else
-        nest!(x.nodes, s, x.indent, extra_width = extra_width)
+        nest!(x.nodes, s, x.indent, extra_margin = x.extra_margin)
     end
 end
 
-n_gen!(x, s; extra_width = 0) =
-    n_block!(x, s, extra_width = extra_width, custom_indent = x.indent)
+n_gen!(x, s) = n_block!(x, s, custom_indent = x.indent)
diff --git a/src/pretty.jl b/src/pretty.jl
index 015a3676..fe0715dc 100644
--- a/src/pretty.jl
+++ b/src/pretty.jl
@@ -22,40 +22,51 @@ mutable struct PTree
     nodes::Union{Nothing,Vector{PTree}}
     ref::Union{Nothing,Ref{CSTParser.EXPR}}
     force_nest::Bool
+    extra_margin::Int
 end
 
 PTree(x::CSTParser.EXPR, indent::Integer) =
-    PTree(x.typ, -1, -1, indent, 0, nothing, PTree[], Ref(x), false)
+    PTree(x.typ, -1, -1, indent, 0, nothing, PTree[], Ref(x), false, 0)
 
 function PTree(x::CSTParser.EXPR, startline::Integer, endline::Integer, val::AbstractString)
-    PTree(x.typ, startline, endline, 0, length(val), val, nothing, Ref(x), false)
+    PTree(x.typ, startline, endline, 0, length(val), val, nothing, Ref(x), false, 0)
 end
 
 function PTree(x::CSTParser.Head, startline::Integer, endline::Integer, val::AbstractString)
-    PTree(x, startline, endline, 0, length(val), val, nothing, nothing, false)
+    PTree(x, startline, endline, 0, length(val), val, nothing, nothing, false, 0)
 end
 
-Newline() = PTree(NEWLINE, -1, -1, 0, 0, "\n", nothing, nothing, false)
-Semicolon() = PTree(SEMICOLON, -1, -1, 0, 1, ";", nothing, nothing, false)
-TrailingComma() = PTree(TRAILINGCOMMA, -1, -1, 0, 0, "", nothing, nothing, false)
-TrailingSemicolon() = PTree(TRAILINGSEMICOLON, -1, -1, 0, 1, ";", nothing, nothing, false)
-Whitespace(n) = PTree(WHITESPACE, -1, -1, 0, n, " "^n, nothing, nothing, false)
-Placeholder(n) = PTree(PLACEHOLDER, -1, -1, 0, n, " "^n, nothing, nothing, false)
+function Base.setindex!(pt::PTree, node::PTree, ind::Int)
+    # pt.len = pt.len - pt.nodes[ind].len + node.len
+    pt.nodes[ind] = node
+end
+Base.getindex(pt::PTree, inds...) = pt.nodes[inds...]
+Base.lastindex(pt::PTree) = length(pt.nodes)
+
+
+Newline(; length = 0, force_nest = false) =
+    PTree(NEWLINE, -1, -1, 0, length, "\n", nothing, nothing, force_nest, 0)
+Semicolon() = PTree(SEMICOLON, -1, -1, 0, 1, ";", nothing, nothing, false, 0)
+TrailingComma() = PTree(TRAILINGCOMMA, -1, -1, 0, 0, "", nothing, nothing, false, 0)
+TrailingSemicolon() =
+    PTree(TRAILINGSEMICOLON, -1, -1, 0, 1, ";", nothing, nothing, false, 0)
+Whitespace(n) = PTree(WHITESPACE, -1, -1, 0, n, " "^n, nothing, nothing, false, 0)
+Placeholder(n) = PTree(PLACEHOLDER, -1, -1, 0, n, " "^n, nothing, nothing, false, 0)
 Notcode(startline, endline) =
-    PTree(NOTCODE, startline, endline, 0, 0, "", nothing, nothing, false)
-InlineComment(line) = PTree(INLINECOMMENT, line, line, 0, 0, "", nothing, nothing, false)
+    PTree(NOTCODE, startline, endline, 0, 0, "", nothing, nothing, false, 0)
+InlineComment(line) = PTree(INLINECOMMENT, line, line, 0, 0, "", nothing, nothing, false, 0)
 
-Base.length(x::PTree) = x.len
+Base.length(pt::PTree) = pt.len
 
-is_leaf(x::PTree) = x.nodes === nothing
-empty_start(x::PTree) = x.startline == 1 && x.endline == 1 && x.val == ""
+is_leaf(pt::PTree) = pt.nodes === nothing
+empty_start(pt::PTree) = pt.startline == 1 && pt.endline == 1 && pt.val == ""
 
 is_punc(x) = CSTParser.ispunctuation(x)
 is_end(x) = x.typ === CSTParser.KEYWORD && x.val == "end"
 is_colon(x) = x.typ === CSTParser.OPERATOR && x.val == ":"
-is_comma(x::PTree) =
-    (x.typ === CSTParser.PUNCTUATION && x.val == ",") || x.typ === TRAILINGCOMMA
-is_comment(x::PTree) = x.typ === INLINECOMMENT || x.typ === NOTCODE
+is_comma(pt::PTree) =
+    (pt.typ === CSTParser.PUNCTUATION && pt.val == ",") || pt.typ === TRAILINGCOMMA
+is_comment(pt::PTree) = pt.typ === INLINECOMMENT || pt.typ === NOTCODE
 
 is_colon_op(x) =
     (
@@ -68,8 +79,8 @@ is_lazy_op(x) =
     )
 
 # f a function which returns a bool
-function parent_is(x, f; ignore_typs = [])
-    p = x.parent
+function parent_is(cst::CSTParser.EXPR, f; ignore_typs = [])
+    p = cst.parent
     p === nothing && return false
     while p !== nothing && p.typ in ignore_typs
         p = p.parent
@@ -225,17 +236,17 @@ function add_node!(t::PTree, n::PTree, s::State; join_lines = false, max_padding
             hs = hascomment(s.doc, current_line)
             hs && add_node!(t, InlineComment(current_line), s)
             if nt !== PLACEHOLDER
-                add_node!(t, Newline(), s)
+                add_node!(t, Newline(force_nest = true), s)
             elseif hs && nt === PLACEHOLDER
                 # swap PLACEHOLDER (will be NEWLINE) with INLINECOMMENT node
                 idx = length(t.nodes)
                 t.nodes[idx-1], t.nodes[idx] = t.nodes[idx], t.nodes[idx-1]
             end
             add_node!(t, Notcode(notcode_startline, notcode_endline), s)
-            add_node!(t, Newline(), s)
+            add_node!(t, Newline(force_nest = true), s)
         elseif !join_lines
             hascomment(s.doc, current_line) && add_node!(t, InlineComment(current_line), s)
-            add_node!(t, Newline(), s)
+            add_node!(t, Newline(force_nest = true), s)
         elseif nt === PLACEHOLDER &&
                current_line != n.startline && hascomment(s.doc, current_line)
             t.force_nest = true
@@ -572,7 +583,7 @@ function p_literal(x, s)
 
     # @debug "" lines x.val loc loc[2] sidx
 
-    t = PTree(CSTParser.StringH, -1, -1, loc[2] - 1, 0, nothing, PTree[], Ref(x), false)
+    t = PTree(CSTParser.StringH, -1, -1, loc[2] - 1, 0, nothing, PTree[], Ref(x), false, 0)
     for (i, l) in enumerate(lines)
         ln = startline + i - 1
         l = i == 1 ? l : l[sidx:end]
@@ -586,6 +597,7 @@ function p_literal(x, s)
             nothing,
             nothing,
             false,
+            0,
         )
         add_node!(t, tt, s)
     end
@@ -631,6 +643,7 @@ function p_stringh(x, s)
             nothing,
             nothing,
             false,
+            0,
         )
         add_node!(t, tt, s)
     end
diff --git a/src/print.jl b/src/print.jl
index b4f399d6..f82a8e2b 100644
--- a/src/print.jl
+++ b/src/print.jl
@@ -89,7 +89,7 @@ function print_stringh(io::IOBuffer, x::PTree, s::State)
     # The new indent for the string is index of when a character in
     # the multiline string is FIRST encountered in the source file - the above difference
     # +1 since the character is 1 space after the indent
-    x.indent = max(x.nodes[1].indent + 1 - diff, 0)
+    x.indent = max(x[1].indent + 1 - diff, 0)
     print_tree(io, x, s)
 end
 


Test patch:
diff --git a/test/files/reverse.jl b/test/files/reverse.jl
index f3503396..ccc00598 100644
--- a/test/files/reverse.jl
+++ b/test/files/reverse.jl
@@ -296,10 +296,7 @@ function adjoint(pr::Primal)
             if haskey(pr.pullbacks, v)
                 g = push!(
                     rb,
-                    stmt(
-                        Expr(:call, alpha(pr.pullbacks[v]), grad(v)),
-                        line = b[v].line,
-                    ),
+                    stmt(Expr(:call, alpha(pr.pullbacks[v]), grad(v)), line = b[v].line),
                 )
                 for (i, x) in enumerate(ex.args)
                     x isa Variable || continue
@@ -337,7 +334,7 @@ function adjoint(pr::Primal)
             Δ = push!(
                 rb,
                 pr.varargs == nothing ? xcall(Zygote, :tuple, gs...) :
-                xcall(Zygote, :tuple_va, Val(pr.varargs), gs...),
+                    xcall(Zygote, :tuple_va, Val(pr.varargs), gs...),
             )
             branches(rb)[1].args[1] = Δ
         end
diff --git a/test/runtests.jl b/test/runtests.jl
index 7b419867..2255a627 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -779,7 +779,7 @@ end
             )
         end"""
         @test fmt(str, 4, 32) == str
-        @test_broken fmt(str, 4, 27) == str
+        @test fmt(str, 4, 27) == str
         str = """
         begin
             foo() = (
@@ -792,18 +792,6 @@ end
         end"""
         @test fmt(str, 4, 26) == str
 
-        str = """
-        begin
-                  foo() = (
-                           one,
-                           x -> (
-                                     true,
-                                     false,
-                           ),
-                  )
-        end"""
-        @test fmt(str, 10, 42) == str
-
         str = """
         ignored_f(f) = f in (
             GlobalRef(Base, :not_int),
@@ -3569,4 +3557,31 @@ some_function(
         @test fmt(str_, 4, 1) == str
     end
 
+    @testset "unnest" begin
+        str = """
+        let X = LinearAlgebra.Symmetric{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where {T},
+            Y = Union{
+                LinearAlgebra.Hermitian{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
+                LinearAlgebra.Symmetric{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
+            }
+
+            @test X <: Y
+        end"""
+        @test fmt(str, 4, 92) == str
+
+        str = """
+        let X = LinearAlgebra.Symmetric{
+                T,
+                S,
+            } where {S<:(AbstractArray{U,2} where {U<:T})} where {T},
+            Y = Union{
+                LinearAlgebra.Hermitian{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
+                LinearAlgebra.Symmetric{T,S} where {S<:(AbstractArray{U,2} where {U<:T})} where T,
+            }
+
+            @test X <: Y
+        end"""
+        @test fmt(str, 4, 90) == str
+    end
+
 end