From 9ba6bb749df7d760aeba9c7a3be555c80fd65916 Mon Sep 17 00:00:00 2001
From: Kyrylo Simonov <xi@resolvent.net>
Date: Fri, 14 Jun 2024 14:06:44 -0500
Subject: [PATCH] Support for docstrings in @funsql notation

---
 docs/src/test/nodes.md | 20 ++++++++++++++++++--
 src/nodes.jl           | 13 +++++++++++--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/docs/src/test/nodes.md b/docs/src/test/nodes.md
index 6d4a6974..fcba928e 100644
--- a/docs/src/test/nodes.md
+++ b/docs/src/test/nodes.md
@@ -247,10 +247,10 @@ A single `@funsql` macro can wrap multiple definitions.
     @funsql begin
         SNOMED(codes...) = concept_by_code("SNOMED", $(codes...))
 
-        MYOCARDIAL_INFARCTION() = SNOMED("22298006")
+        `MYOCARDIAL INFARCTION`() = SNOMED("22298006")
     end
 
-    display(@funsql MYOCARDIAL_INFARCTION())
+    display(@funsql `MYOCARDIAL INFARCTION`())
     #=>
     let q1 = From(:concept),
         q2 = q1 |>
@@ -260,6 +260,22 @@ A single `@funsql` macro can wrap multiple definitions.
     end
     =#
 
+A query function may have a docstring.
+
+    @funsql begin
+        "SNOMED concept set with the given `codes`"
+        SNOMED
+
+        "Visit concept set with the given `codes`"
+        Visit(codes...) = concept_by_code("Visit", $(codes...))
+    end
+
+    @doc funsql_SNOMED
+    #-> SNOMED concept set with the given `codes`
+
+    @doc funsql_Visit
+    #-> Visit concept set with the given `codes`
+
 An ill-formed `@funsql` query triggers an error.
 
     @funsql for p in person; end
diff --git a/src/nodes.jl b/src/nodes.jl
index 23167580..58d955a8 100644
--- a/src/nodes.jl
+++ b/src/nodes.jl
@@ -513,7 +513,16 @@ function transliterate(@nospecialize(ex), ctx::TransliterateContext)
         if @dissect(ex, Expr(:($), arg))
             # $(...)
             return esc(arg)
-        elseif @dissect(ex, Expr(:(=), Expr(:call, name::Symbol, args...), body))
+        elseif @dissect(ex, Expr(:macrocall, ref := GlobalRef(Core, Symbol("@doc")), ln::LineNumberNode, doc, arg))
+            # "..." ...
+            if @dissect(arg, name::Symbol || Expr(:macrocall, GlobalRef(Core, Symbol("@cmd")), ::LineNumberNode, name::String))
+                arg = Symbol("funsql_$name")
+            else
+                ctx = TransliterateContext(ctx, src = ln)
+                arg = transliterate(arg, ctx)
+            end
+            return Expr(:macrocall, ref, ln, doc, arg)
+        elseif @dissect(ex, Expr(:(=), Expr(:call, name::Symbol || Expr(:macrocall, GlobalRef(Core, Symbol("@cmd")), ::LineNumberNode, name::String), args...), body))
             # name(args...) = body
             ctx = TransliterateContext(ctx, decl = true)
             trs = Any[transliterate(arg, ctx) for arg in args]
@@ -623,7 +632,7 @@ function transliterate(@nospecialize(ex), ctx::TransliterateContext)
             return :($(esc(Symbol("funsql_$name")))($(trs...)))
         elseif @dissect(ex, Expr(:block, args...))
             # begin; args...; end
-            if all(@dissect(arg, ::LineNumberNode || Expr(:(=), _...))
+            if all(@dissect(arg, ::LineNumberNode || Expr(:(=), _...) || Expr(:macrocall, GlobalRef(Core, Symbol("@doc")), _...))
                    for arg in args)
                 trs = Any[]
                 for arg in args