Rework indentation after assignment

This patch reworks the indentation after assignment. For "blocklike"
(try/if/triple strings) right hand sides the indentation is disabled.
Superfluous newlines are removed. For example

```julia
x =
if a
    x
else
    y
end
```

will become

```julia
x = if a
    x
else
    y
end
```

A valid alternative is

```julia
x = if a
        x
    else
        y
    end
```

This highlights the start/end of the right hand side of the assignment
in a better way, but it doesn't look very nice and this style is never
really seen in the wild. For triple strings I think it is OK though,
e.g.

```julia
x = """
    foo
    bar
    """
```

vs the current (and seen more in the wild)

```julia
x = """
foo
bar
"""
```

Closes #39.

src/chisels.jl

@@ -310,7 +310,12 @@ end
 
 function is_assignment(node::Node)
     return JuliaSyntax.is_prec_assignment(node)
-    # return !is_leaf(node) && JuliaSyntax.is_prec_assignment(node)
+end
+
+# Assignment node but exclude loop variable assignment
+function is_variable_assignment(ctx, node::Node)
+    return !is_leaf(node) && is_assignment(node) &&
+        !(ctx.lineage_kinds[end] in KSet"for generator cartesian_iterator filter")
 end
 
 function unwrap_to_call_or_tuple(x)
@@ -553,6 +558,18 @@ function is_triple_string(node)
         JuliaSyntax.has_flags(node, JuliaSyntax.TRIPLE_STRING_FLAG)
 end
 
+function is_triple_string_macro(node)
+    if kind(node) === K"macrocall"
+        kids = verified_kids(node)
+        if length(kids) >= 2 &&
+                kind(kids[1]) in KSet"StringMacroName CmdMacroName core_@cmd" &&
+                is_triple_string(kids[2])
+            return true
+        end
+    end
+    return false
+end
+
 ##########################
 # Utilities for IOBuffer #
 ##########################

src/runestone.jl

@@ -2074,9 +2074,7 @@ function indent_call(ctx::Context, node::Node)
 end
 
 
-# TODO: I feel like this function can be removed. The use in `indent_assignment` can be
-# replaced with `continue_all_newlines` but the use in `indent_op_call` might have to be
-# tweaked slightly.
+# TODO: I feel like this function can be removed. It is only used in `indent_op_call`
 function indent_newlines_between_indices(
         ctx::Context, node::Node, open_idx::Int, close_idx::Int;
         indent_closing_token::Bool = false,
@@ -2532,19 +2530,110 @@ function indent_ternary(ctx::Context, node::Node)
 end
 
 function indent_iterator(ctx::Context, node::Node)
-    @assert kind(node) === K"cartesian_iterator"
+    @assert kind(node) in KSet"cartesian_iterator generator"
     return continue_all_newlines(ctx, node)
 end
 
 function indent_assignment(ctx::Context, node::Node)
+    @assert !is_leaf(node)
+    @assert is_variable_assignment(ctx, node)
     kids = verified_kids(node)
-    # Also catches for loop specifications (but at this point we have normalized `=` and `∈`
-    # to `in`).
-    op_idx = findfirst(x -> is_assignment(x) || kind(x) === K"in", kids)::Int
-    last_non_ws_idx = findlast(!JuliaSyntax.is_whitespace, kids)::Int
-    return indent_newlines_between_indices(
-        ctx, node, op_idx, last_non_ws_idx; indent_closing_token = true,
-    )
+    lhsidx = findfirst(!JuliaSyntax.is_whitespace, kids)::Int
+    eqidx = findnext(!JuliaSyntax.is_whitespace, kids, lhsidx + 1)::Int
+    if !(kind(node) === kind(kids[eqidx]))
+        @info "Indent assignment failure $node" ctx.lineage_kinds[end]
+    end
+    @assert kind(node) === kind(kids[eqidx])
+    @assert length(kids) > eqidx
+    rhsidx = findnext(!JuliaSyntax.is_whitespace, kids, eqidx + 1)::Int
+    r = (eqidx + 1):(rhsidx-1)
+    length(r) == 0 && return nothing
+    if length(r) == 1 && kind(kids[r[1]]) === K"Whitespace"
+        return nothing
+    end
+    if !any(i -> kind(kids[i]) === K"NewlineWs", r)
+        return nothing
+    end
+    rhs = kids[rhsidx]
+    # These right hand sides have more "inertia" towards indentation. This is so that we
+    # will end up with
+    # ```
+    # x = if cond
+    # end
+    # ```
+    # instead of
+    # ```
+    # x = if cond
+    #     end
+    # ```
+    # (Although perhaps the latter is actually more readable?). If possible we remove
+    # any newlines inbetween the `=` and the rhs because
+    # ```
+    # x =
+    # if cond
+    # end
+    # ```
+    # doesn't look good.
+    blocklike = kind(rhs) in KSet"if try function" ||
+        is_triple_string(rhs) || is_triple_string_macro(rhs)
+    has_comments = any(i -> kind(kids[i]) === K"Comment", r)
+    if has_comments
+        # TODO: Perhaps delete newlines in between comments or something
+        blocklike && return nothing
+        # Continue all newlines between the `=` and the rhs
+        kids′ = kids
+        changed = false
+        for i in r
+            kid = kids[i]
+            if kind(kid) === K"NewlineWs" && !has_tag(kid, TAG_LINE_CONT)
+                kid′ = add_tag(kid, TAG_LINE_CONT)
+                changed = true
+            else
+                kid′ = kid
+            end
+            if changed
+                if kids′ === kids
+                    kids′ = kids[1:(i - 1)]
+                end
+                push!(kids′, kid′)
+            end
+        end
+        if changed
+            @assert kids !== kids′
+            for i in rhsidx:length(kids)
+                push!(kids′, kids[i])
+            end
+            return make_node(node, kids′)
+        else
+            return nothing
+        end
+    end
+    if !blocklike && length(r) == 1 && has_tag(kids[r[1]], TAG_LINE_CONT)
+        @assert kind(kids[r[1]]) === K"NewlineWs"
+        return nothing
+    end
+    @assert all(i -> kind(kids[i]) === K"NewlineWs", r)
+    kids′ = kids[1:eqidx]
+    for i in 1:eqidx
+        accept_node!(ctx, kids[i])
+    end
+    for i in r
+        if i == r[end] && !blocklike
+            kid = kids[i]
+            @assert kind(kid) === K"NewlineWs"
+            if !has_tag(kid, TAG_LINE_CONT)
+                kid = add_tag(kid, TAG_LINE_CONT)
+            end
+            push!(kids′, kid)
+        else
+            replace_bytes!(ctx, "", span(kids[i]))
+        end
+    end
+    # TODO: Insert a whitespace node here to skip another recursion.
+    for i in rhsidx:length(kids)
+        push!(kids′, kids[i])
+    end
+    return make_node(node, kids′)
 end
 
 function indent_paren_block(ctx::Context, node::Node)
@@ -2730,13 +2819,13 @@ function insert_delete_mark_newlines(ctx::Context, node::Node)
         return indent_short_circuit(ctx, node)
     elseif kind(node) in KSet"using import export public"
         return indent_using_import_export_public(ctx, node)
-    elseif is_assignment(node)
+    elseif is_variable_assignment(ctx, node)
         return indent_assignment(ctx, node)
     elseif kind(node) === K"parameters"
         return indent_parameters(ctx, node)
     elseif kind(node) === K"?"
         return indent_ternary(ctx, node)
-    elseif kind(node) === K"cartesian_iterator"
+    elseif kind(node) in KSet"cartesian_iterator generator"
         return indent_iterator(ctx, node)
     elseif kind(node) === K"try"
         return indent_try(ctx, node)

test/runtests.jl

@@ -745,8 +745,24 @@ end
             end
         end
         # assignment
-        for op in ("=", "+=")
-            @test format_string("a $(op)\n$(sp)b") == "a $(op)\n    b"
+        for nl in ("\n", "\n\n")
+            # Regular line continuation of newlines between `=` and rhs
+            for op in ("=", "+=", ".=", ".+=")
+                @test format_string("a $(op)$(nl)b") == "a $(op)\n    b"
+                @test format_string("a $(op)$(nl)# comment$(nl)b") ==
+                    "a $(op)$(nl)    # comment$(nl)    b"
+            end
+            @test format_string("f(a) =$(nl)b") == "f(a) =\n    b"
+            # Blocklike RHS
+            for thing in (
+                    "if c\n    x\nend", "try\n    x\ncatch\n    y\nend",
+                    "\"\"\"\nfoo\n\"\"\"", "r\"\"\"\nfoo\n\"\"\"",
+                    "```\nfoo\n```", "r```\nfoo\n```",
+                )
+                @test format_string("a =$(nl)$(thing)") == "a = $(thing)"
+                @test format_string("a =$(nl)# comment$(nl)$(thing)") ==
+                    "a =$(nl)# comment$(nl)$(thing)"
+            end
         end
         # using/import
         for verb in ("using", "import")