From 2abf9db43100663893b65d7905e908ad0b16a383 Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Mon, 21 Oct 2024 14:15:14 +0200
Subject: [PATCH] Fix indentation of documented modules in files with more
 toplevel code

---
 src/runestone.jl | 29 +++++++++++++++++++++++++++--
 test/runtests.jl |  3 +++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/runestone.jl b/src/runestone.jl
index e343894..6a42d87 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -2936,9 +2936,31 @@ function indent_comparison(ctx::Context, node::Node)
     return continue_all_newlines(ctx, node)
 end
 
+# Indent a nested documented module
+function indent_doc_module(ctx::Context, node::Node; do_indent::Bool)
+    @assert kind(node) === K"doc"
+    kids = verified_kids(node)
+    mod_idx = findfirst(x -> kind(x) === K"module", kids)::Int
+    pos = position(ctx.fmt_io)
+    for i in 1:(mod_idx - 1)
+        accept_node!(ctx, kids[i])
+    end
+    mod′ = indent_module(ctx, kids[mod_idx]; do_indent = do_indent)
+    seek(ctx.fmt_io, pos)
+    if mod′ === nothing
+        return nothing
+    end
+    kids′ = copy(kids)
+    kids′[mod_idx] = mod′
+    return make_node(node, kids′)
+end
+
 # Indent a nested module
 function indent_module(ctx::Context, node::Node; do_indent::Bool = true)
-    @assert kind(node) === K"module"
+    @assert kind(node) in KSet"module doc"
+    if kind(node) === K"doc"
+        return indent_doc_module(ctx, node; do_indent = do_indent)
+    end
     kids = verified_kids(node)
     any_kid_changed = false
     pos = position(ctx.fmt_io)
@@ -2998,7 +3020,10 @@ end
 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)
+    mod_idx = findfirst(kids) do x
+        kind(x) === K"module" ||
+            (kind(x) === K"doc" && findfirst(y -> kind(y) === K"module", verified_kids(x)) !== nothing)
+    end
     if mod_idx === nothing
         # No module here
         return nothing
diff --git a/test/runtests.jl b/test/runtests.jl
index 1a4fece..f66b755 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -673,6 +673,9 @@ end
             # 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"
+            # toplevel documented module with more things
+            @test format_string("\"doc\"\n$(b)module A\n$(sp)x\nend\nf") ==
+                "\"doc\"\n$(b)module A\n    x\nend\nf"
             # var"" as module name
             @test format_string("$(b)module var\"A\"\n$(sp)x\n$(sp)end\nf") ==
                 "$(b)module var\"A\"\n    x\nend\nf"