Indent nested modules, fixes #5.

src/Runic.jl

@@ -123,6 +123,7 @@ mutable struct Context
     prev_sibling::Union{Node, Nothing}
     next_sibling::Union{Node, Nothing}
     # parent::Union{Node, Nothing}
+    lineage_kinds::Vector{JuliaSyntax.Kind}
 end
 
 function Context(
@@ -151,9 +152,10 @@ function Context(
     indent_level = 0
     call_depth = 0
     prev_sibling = next_sibling = nothing
+    lineage_kinds = JuliaSyntax.Kind[]
     return Context(
         src_str, src_tree, src_io, fmt_io, fmt_tree, quiet, verbose, assert, debug, check,
-        diff, call_depth, indent_level, prev_sibling, next_sibling,
+        diff, call_depth, indent_level, prev_sibling, next_sibling, lineage_kinds,
     )
 end
 
@@ -200,6 +202,7 @@ function format_node_with_kids!(ctx::Context, node::Node)
     next_sibling = ctx.next_sibling
     ctx.prev_sibling = nothing
     ctx.next_sibling = nothing
+    push!(ctx.lineage_kinds, kind(node))
 
     # The new node parts. `kids′` aliases `kids` and only copied below if any of the
     # nodes change ("copy-on-write").
@@ -255,6 +258,7 @@ function format_node_with_kids!(ctx::Context, node::Node)
     # Reset the siblings
     ctx.prev_sibling = prev_sibling
     ctx.next_sibling = next_sibling
+    pop!(ctx.lineage_kinds)
     ctx.call_depth -= 1
     # Return a new node if any of the kids changed
     if any_kid_changed
@@ -501,8 +505,8 @@ function format_tree!(ctx::Context)
             root′ = root′′
         end
         # The root node must only change once.
-        if (itr += 1) == 2
-            error("root node modified more than once")
+        if (itr += 1) > 100
+            error("root node modified more than 100 times?")
         end
     end
     # Truncate the output at the root span

src/runestone.jl

@@ -1152,6 +1152,69 @@ function indent_comparison(ctx::Context, node::Node)
     return continue_all_newlines(ctx, node)
 end
 
+# Indent a nested module
+function indent_module(ctx::Context, node::Node)
+    kids = verified_kids(node)
+    any_kid_changed = false
+    # First node is the module keyword
+    mod_idx = 1
+    mod_node = kids[mod_idx]
+    @assert is_leaf(mod_node) && kind(mod_node) in KSet"module baremodule"
+    if !has_tag(mod_node, TAG_INDENT)
+        kids[mod_idx] = add_tag(mod_node, TAG_INDENT)
+        any_kid_changed = true
+    end
+    # Second node is the space between keyword and name
+    # TODO: Make sure there is just a single space
+    space_idx = 2
+    space_node = kids[space_idx]
+    @assert is_leaf(space_node) && kind(space_node) === K"Whitespace"
+    # Third node is the module identifier
+    id_idx = 3
+    id_node = kids[id_idx]
+    @assert kind(id_node) === K"Identifier"
+    # Fourth node is the module body block.
+    block_idx = 4
+    block_node′ = indent_block(ctx, kids[block_idx])
+    if block_node′ !== nothing
+        kids[block_idx] = block_node′
+        any_kid_changed = true
+    end
+    # Fifth node is the closing end keyword
+    end_idx = 5
+    end_node = kids[end_idx]
+    @assert is_leaf(end_node) && kind(end_node) === K"end"
+    if !has_tag(end_node, TAG_DEDENT)
+        kids[end_idx] = add_tag(end_node, TAG_DEDENT)
+        any_kid_changed = true
+    end
+    @assert verified_kids(node) === kids
+    return any_kid_changed ? node : nothing
+end
+
+# The only thing at top level that we need to indent are modules which don't occupy the full
+# top level expression, for example a file with an inner module followed by some code.
+function indent_toplevel(ctx::Context, node::Node)
+    @assert kind(node) === K"toplevel"
+    kids = verified_kids(node)
+    mod_idx = findfirst(x -> kind(x) === K"module", kids)
+    if mod_idx === nothing || count(!JuliaSyntax.is_whitespace, kids) == 1
+        # No module or module that is the only top level expression
+        return nothing
+    end
+    any_kid_changed = false
+    while mod_idx !== nothing
+        mod_node = kids[mod_idx]
+        mod_node′ = indent_module(ctx, mod_node)
+        if mod_node′ !== nothing
+            kids[mod_idx] = mod_node′
+            any_kid_changed = true
+        end
+        mod_idx = findnext(x -> kind(x) === K"module", kids, mod_idx + 1)
+    end
+    return any_kid_changed ? node : nothing
+end
+
 function insert_delete_mark_newlines(ctx::Context, node::Node)
     if is_leaf(node)
         return nothing
@@ -1201,6 +1264,10 @@ function insert_delete_mark_newlines(ctx::Context, node::Node)
         return indent_array_row(ctx, node)
     elseif kind(node) === K"comparison"
         return indent_comparison(ctx, node)
+    elseif kind(node) === K"toplevel"
+        return indent_toplevel(ctx, node)
+    elseif kind(node) === K"module" && findlast(x -> x === K"module", ctx.lineage_kinds) !== nothing
+        return indent_module(ctx, node)
     end
     return nothing
 end

test/runtests.jl

@@ -388,6 +388,31 @@ end
         # do-end
         @test format_string("open() do\n$(sp)a\n$(sp)end") == "open() do\n    a\nend"
         @test format_string("open() do io\n$(sp)a\n$(sp)end") == "open() do io\n    a\nend"
+        # module-end, baremodule-end
+        for b in ("", "bare")
+            # Just a module
+            @test format_string("$(b)module A\n$(sp)x\n$(sp)end") == "$(b)module A\nx\nend"
+            # Comment before
+            @test format_string("# c\n$(b)module A\n$(sp)x\n$(sp)end") ==
+                "# c\n$(b)module A\nx\nend"
+            # Docstring before
+            @test format_string("\"doc\"\n$(b)module A\n$(sp)x\n$(sp)end") ==
+                "\"doc\"\n$(b)module A\nx\nend"
+            # code before
+            @test format_string("f\n$(b)module A\n$(sp)x\n$(sp)end") ==
+                "f\n$(b)module A\n    x\nend"
+            @test format_string("f\n$(b)module A\n$(sp)x\n$(sp)end\n$(b)module B\n$(sp)x\n$(sp)end") ==
+                "f\n$(b)module A\n    x\nend\n$(b)module B\n    x\nend"
+            # code after
+            @test format_string("$(b)module A\n$(sp)x\n$(sp)end\nf") ==
+                "$(b)module A\n    x\nend\nf"
+            # nested modules
+            @test format_string("$(b)module A\n$(sp)$(b)module B\n$(sp)x\n$(sp)end\n$(sp)end") ==
+                "$(b)module A\n$(b)module B\n    x\nend\nend"
+            # nested documented modules
+            @test format_string("\"doc\"\n$(b)module A\n\"doc\"\n$(b)module B\n$(sp)x\n$(sp)end\n$(sp)end") ==
+                "\"doc\"\n$(b)module A\n\"doc\"\n$(b)module B\n    x\nend\nend"
+        end
     end
 end