From 18ad14f571cf296a1b87c9b5da559807da324be1 Mon Sep 17 00:00:00 2001
From: Blake Kostner <git@btkostner.io>
Date: Sun, 20 Oct 2024 04:01:44 -0700
Subject: [PATCH] Add protobuf code comments as Elixir module documentation
 (#352)

* Add protobuf code comments as moduledoc

* remove debug docs

* Remove grpc dep

* Refactor test to avoid dependency on grpc library

* Ensure moduledoc is set even if modules don't have comments

* Fix warning during code generation with multi-line comments

---------

Co-authored-by: v0idpwn <v0idpwn@gmail.com>
---
 lib/protobuf/protoc/context.ex                |  23 ++-
 lib/protobuf/protoc/generator.ex              |  20 +-
 lib/protobuf/protoc/generator/comment.ex      |  58 ++++++
 lib/protobuf/protoc/generator/enum.ex         |   2 +
 lib/protobuf/protoc/generator/extension.ex    |  24 ++-
 lib/protobuf/protoc/generator/message.ex      |  30 ++-
 lib/protobuf/protoc/generator/service.ex      |   2 +
 lib/protobuf/protoc/generator/util.ex         |  13 ++
 mix.exs                                       |  10 +-
 priv/templates/enum.ex.eex                    |   9 +-
 priv/templates/extension.ex.eex               |   9 +
 priv/templates/message.ex.eex                 |   9 +-
 priv/templates/service.ex.eex                 |   9 +-
 test/protobuf/protoc/cli_integration_test.exs | 195 ++++++++++++++----
 test/protobuf/protoc/generator/enum_test.exs  |  32 ++-
 .../protoc/generator/message_test.exs         |  16 +-
 .../protoc/generator/service_test.exs         |  42 +++-
 test/protobuf/protoc/proto/service.proto      |  13 ++
 test/protobuf/protoc/proto/test.proto         |   6 +
 test/test_helper.exs                          |   6 +
 20 files changed, 443 insertions(+), 85 deletions(-)
 create mode 100644 lib/protobuf/protoc/generator/comment.ex
 create mode 100644 test/protobuf/protoc/proto/service.proto

diff --git a/lib/protobuf/protoc/context.ex b/lib/protobuf/protoc/context.ex
index 6c996ecc..282b8e4a 100644
--- a/lib/protobuf/protoc/context.ex
+++ b/lib/protobuf/protoc/context.ex
@@ -8,6 +8,10 @@ defmodule Protobuf.Protoc.Context do
 
             ### All files scope
 
+            # All parsed comments from the source file (mapping from path to comment)
+            # %{"1.4.2" => "this is a comment", "1.5.2.4.2" => "more comment\ndetails"}
+            comments: %{},
+
             # Mapping from file name to (mapping from type name to metadata, like elixir type name)
             # %{"example.proto" => %{".example.FooMsg" => %{type_name: "Example.FooMsg"}}}
             global_type_mapping: %{},
@@ -42,7 +46,11 @@ defmodule Protobuf.Protoc.Context do
             include_docs?: false,
 
             # Elixirpb.FileOptions
-            custom_file_options: %{}
+            custom_file_options: %{},
+
+            # Current comment path. The locations are encoded into with a joining "."
+            # character. E.g. "4.2.3.0"
+            current_comment_path: ""
 
   @spec custom_file_options_from_file_desc(t(), Google.Protobuf.FileDescriptorProto.t()) :: t()
   def custom_file_options_from_file_desc(ctx, desc)
@@ -68,4 +76,17 @@ defmodule Protobuf.Protoc.Context do
         module_prefix: Map.get(custom_file_opts, :module_prefix)
     }
   end
+
+  @doc """
+  Appends a comment path to the current comment path.
+
+  ## Examples
+
+      iex> append_comment_path(%{current_comment_path: "4"}, "2.4")
+      %{current_comment_path: "4.2.4"}
+
+  """
+  def append_comment_path(ctx, path) do
+    %{ctx | current_comment_path: String.trim(ctx.current_comment_path <> "." <> path, ".")}
+  end
 end
diff --git a/lib/protobuf/protoc/generator.ex b/lib/protobuf/protoc/generator.ex
index 6011f1a1..fa5434c9 100644
--- a/lib/protobuf/protoc/generator.ex
+++ b/lib/protobuf/protoc/generator.ex
@@ -36,13 +36,20 @@ defmodule Protobuf.Protoc.Generator do
     ctx =
       %Context{
         ctx
-        | syntax: syntax(desc.syntax),
+        | comments: Protobuf.Protoc.Generator.Comment.parse(desc),
+          syntax: syntax(desc.syntax),
           package: desc.package,
           dep_type_mapping: get_dep_type_mapping(ctx, desc.dependency, desc.name)
       }
       |> Protobuf.Protoc.Context.custom_file_options_from_file_desc(desc)
 
-    enum_defmodules = Enum.map(desc.enum_type, &Generator.Enum.generate(ctx, &1))
+    enum_defmodules =
+      desc.enum_type
+      |> Enum.with_index()
+      |> Enum.map(fn {enum, index} ->
+        {Context.append_comment_path(ctx, "5.#{index}"), enum}
+      end)
+      |> Enum.map(fn {ctx, enum} -> Generator.Enum.generate(ctx, enum) end)
 
     {nested_enum_defmodules, message_defmodules} =
       Generator.Message.generate_list(ctx, desc.message_type)
@@ -51,7 +58,14 @@ defmodule Protobuf.Protoc.Generator do
 
     service_defmodules =
       if "grpc" in ctx.plugins do
-        Enum.map(desc.service, &Generator.Service.generate(ctx, &1))
+        desc.service
+        |> Enum.with_index()
+        |> Enum.map(fn {service, index} ->
+          Generator.Service.generate(
+            Context.append_comment_path(ctx, "6.#{index}"),
+            service
+          )
+        end)
       else
         []
       end
diff --git a/lib/protobuf/protoc/generator/comment.ex b/lib/protobuf/protoc/generator/comment.ex
new file mode 100644
index 00000000..44489e19
--- /dev/null
+++ b/lib/protobuf/protoc/generator/comment.ex
@@ -0,0 +1,58 @@
+defmodule Protobuf.Protoc.Generator.Comment do
+  @moduledoc false
+
+  alias Protobuf.Protoc.Context
+
+  @doc """
+  Parses comment information from `Google.Protobuf.FileDescriptorProto`
+  into a map with path keys.
+  """
+  @spec parse(Google.Protobuf.FileDescriptorProto.t()) :: %{optional(String.t()) => String.t()}
+  def parse(file_descriptor_proto) do
+    file_descriptor_proto
+    |> get_locations()
+    |> Enum.reject(&empty_comment?/1)
+    |> Map.new(fn location ->
+      {Enum.join(location.path, "."), format_comment(location)}
+    end)
+  end
+
+  defp get_locations(%{source_code_info: %{location: value}}) when is_list(value),
+    do: value
+
+  defp get_locations(_value), do: []
+
+  defp empty_comment?(%{leading_comments: value}) when not is_nil(value) and value != "",
+    do: false
+
+  defp empty_comment?(%{trailing_comments: value}) when not is_nil(value) and value != "",
+    do: false
+
+  defp empty_comment?(%{leading_detached_comments: value}), do: Enum.empty?(value)
+
+  defp format_comment(location) do
+    [location.leading_comments, location.trailing_comments | location.leading_detached_comments]
+    |> Enum.reject(&is_nil/1)
+    |> Enum.map(&String.replace(&1, ~r/^\s*\*/, "", global: true))
+    |> Enum.join("\n\n")
+    |> String.replace(~r/\n{3,}/, "\n")
+    |> String.trim()
+  end
+
+  @doc """
+  Finds a comment via the context. Returns an empty string if the
+  comment is not found or if `include_docs?` is set to false.
+  """
+  @spec get(Context.t()) :: String.t()
+  def get(%{include_docs?: false}), do: ""
+
+  def get(%{comments: comments, current_comment_path: path}),
+    do: get(comments, path)
+
+  @doc """
+  Finds a comment via a map of comments and a path. Returns an
+  empty string if the comment is not found
+  """
+  @spec get(%{optional(String.t()) => String.t()}, String.t()) :: String.t()
+  def get(comments, path), do: Map.get(comments, path, "")
+end
diff --git a/lib/protobuf/protoc/generator/enum.ex b/lib/protobuf/protoc/generator/enum.ex
index 210aafaf..bda2cbef 100644
--- a/lib/protobuf/protoc/generator/enum.ex
+++ b/lib/protobuf/protoc/generator/enum.ex
@@ -2,6 +2,7 @@ defmodule Protobuf.Protoc.Generator.Enum do
   @moduledoc false
 
   alias Protobuf.Protoc.Context
+  alias Protobuf.Protoc.Generator.Comment
   alias Protobuf.Protoc.Generator.Util
 
   require EEx
@@ -34,6 +35,7 @@ defmodule Protobuf.Protoc.Generator.Enum do
 
     content =
       enum_template(
+        comment: Comment.get(ctx),
         module: msg_name,
         use_options: use_options,
         fields: desc.value,
diff --git a/lib/protobuf/protoc/generator/extension.ex b/lib/protobuf/protoc/generator/extension.ex
index 9c57a74a..69415208 100644
--- a/lib/protobuf/protoc/generator/extension.ex
+++ b/lib/protobuf/protoc/generator/extension.ex
@@ -3,6 +3,7 @@ defmodule Protobuf.Protoc.Generator.Extension do
 
   alias Google.Protobuf.{DescriptorProto, FieldDescriptorProto, FileDescriptorProto}
   alias Protobuf.Protoc.Context
+  alias Protobuf.Protoc.Generator.Comment
   alias Protobuf.Protoc.Generator.Util
 
   require EEx
@@ -29,7 +30,13 @@ defmodule Protobuf.Protoc.Generator.Extension do
 
     module_contents =
       Util.format(
-        extension_template(use_options: use_options, module: mod_name, extends: extensions)
+        extension_template(
+          comment: Comment.get(ctx),
+          use_options: use_options,
+          module: mod_name,
+          extends: extensions,
+          module_doc?: ctx.include_docs?
+        )
       )
 
     {mod_name, module_contents}
@@ -75,10 +82,15 @@ defmodule Protobuf.Protoc.Generator.Extension do
   end
 
   defp get_extensions_from_messages(%Context{} = ctx, use_options, descs) do
-    Enum.flat_map(descs, fn %DescriptorProto{} = desc ->
-      generate_module(ctx, use_options, desc) ++
+    descs
+    |> Enum.with_index()
+    |> Enum.flat_map(fn {desc, index} ->
+      generate_module(Context.append_comment_path(ctx, "7.#{index}"), use_options, desc) ++
         get_extensions_from_messages(
-          %Context{ctx | namespace: ctx.namespace ++ [Macro.camelize(desc.name)]},
+          %Context{
+            Context.append_comment_path(ctx, "6.#{index}")
+            | namespace: ctx.namespace ++ [Macro.camelize(desc.name)]
+          },
           use_options,
           desc.nested_type
         )
@@ -96,9 +108,11 @@ defmodule Protobuf.Protoc.Generator.Extension do
     module_contents =
       Util.format(
         extension_template(
+          comment: Comment.get(ctx),
           module: module_name,
           use_options: use_options,
-          extends: Enum.map(desc.extension, &generate_extend_dsl(ctx, &1, _ns = ""))
+          extends: Enum.map(desc.extension, &generate_extend_dsl(ctx, &1, _ns = "")),
+          module_doc?: ctx.include_docs?
         )
       )
 
diff --git a/lib/protobuf/protoc/generator/message.ex b/lib/protobuf/protoc/generator/message.ex
index c858b3e6..82c1dde5 100644
--- a/lib/protobuf/protoc/generator/message.ex
+++ b/lib/protobuf/protoc/generator/message.ex
@@ -4,6 +4,7 @@ defmodule Protobuf.Protoc.Generator.Message do
   alias Google.Protobuf.{DescriptorProto, FieldDescriptorProto}
 
   alias Protobuf.Protoc.Context
+  alias Protobuf.Protoc.Generator.Comment
   alias Protobuf.Protoc.Generator.Util
   alias Protobuf.Protoc.Generator.Enum, as: EnumGenerator
 
@@ -21,7 +22,10 @@ defmodule Protobuf.Protoc.Generator.Message do
            messages :: [{mod_name :: String.t(), contents :: String.t()}]}
   def generate_list(%Context{} = ctx, descs) when is_list(descs) do
     descs
-    |> Enum.map(fn desc -> generate(ctx, desc) end)
+    |> Enum.with_index()
+    |> Enum.map(fn {desc, index} ->
+      generate(Context.append_comment_path(ctx, "4.#{index}"), desc)
+    end)
     |> Enum.unzip()
   end
 
@@ -46,6 +50,7 @@ defmodule Protobuf.Protoc.Generator.Message do
       {msg_name,
        Util.format(
          message_template(
+           comment: Comment.get(ctx),
            module: msg_name,
            use_options: msg_opts_str(ctx, desc.options),
            oneofs: desc.oneof_decl,
@@ -61,11 +66,19 @@ defmodule Protobuf.Protoc.Generator.Message do
   end
 
   defp gen_nested_msgs(ctx, desc) do
-    Enum.map(desc.nested_type, fn msg_desc -> generate(ctx, msg_desc) end)
+    desc.nested_type
+    |> Enum.with_index()
+    |> Enum.map(fn {msg_desc, index} ->
+      generate(Context.append_comment_path(ctx, "3.#{index}"), msg_desc)
+    end)
   end
 
   defp gen_nested_enums(ctx, desc) do
-    Enum.map(desc.enum_type, fn enum_desc -> EnumGenerator.generate(ctx, enum_desc) end)
+    desc.enum_type
+    |> Enum.with_index()
+    |> Enum.map(fn {enum_desc, index} ->
+      EnumGenerator.generate(Context.append_comment_path(ctx, "4.#{index}"), enum_desc)
+    end)
   end
 
   defp gen_fields(syntax, fields) do
@@ -103,7 +116,15 @@ defmodule Protobuf.Protoc.Generator.Message do
     oneofs = get_real_oneofs(desc.oneof_decl)
 
     nested_maps = nested_maps(ctx, desc)
-    for field <- desc.field, do: get_field(ctx, field, nested_maps, oneofs)
+
+    for {field, index} <- Enum.with_index(desc.field) do
+      get_field(
+        Context.append_comment_path(ctx, "2.#{index}"),
+        field,
+        nested_maps,
+        oneofs
+      )
+    end
   end
 
   # Public and used by extensions.
@@ -137,6 +158,7 @@ defmodule Protobuf.Protoc.Generator.Message do
 
     %{
       name: field_desc.name,
+      comment: Comment.get(ctx),
       number: field_desc.number,
       label: label_name(field_desc.label),
       type: type,
diff --git a/lib/protobuf/protoc/generator/service.ex b/lib/protobuf/protoc/generator/service.ex
index 08aaf424..71775233 100644
--- a/lib/protobuf/protoc/generator/service.ex
+++ b/lib/protobuf/protoc/generator/service.ex
@@ -2,6 +2,7 @@ defmodule Protobuf.Protoc.Generator.Service do
   @moduledoc false
 
   alias Protobuf.Protoc.Context
+  alias Protobuf.Protoc.Generator.Comment
   alias Protobuf.Protoc.Generator.Util
 
   require EEx
@@ -31,6 +32,7 @@ defmodule Protobuf.Protoc.Generator.Service do
     {mod_name,
      Util.format(
        service_template(
+         comment: Comment.get(ctx),
          module: mod_name,
          service_name: name,
          package: ctx.package,
diff --git a/lib/protobuf/protoc/generator/util.ex b/lib/protobuf/protoc/generator/util.ex
index e3f23343..70a6ca50 100644
--- a/lib/protobuf/protoc/generator/util.ex
+++ b/lib/protobuf/protoc/generator/util.ex
@@ -86,6 +86,19 @@ defmodule Protobuf.Protoc.Generator.Util do
     |> IO.iodata_to_binary()
   end
 
+  @spec pad_comment(String.t(), non_neg_integer()) :: String.t()
+  def pad_comment(comment, size) do
+    padding = String.duplicate(" ", size)
+
+    comment
+    |> String.split("\n")
+    |> Enum.map(fn line ->
+      trimmed = String.trim_leading(line, " ")
+      padding <> trimmed
+    end)
+    |> Enum.join("\n")
+  end
+
   @spec version() :: String.t()
   def version do
     {:ok, value} = :application.get_key(:protobuf, :vsn)
diff --git a/mix.exs b/mix.exs
index 1c7e01bf..9eb58f39 100644
--- a/mix.exs
+++ b/mix.exs
@@ -149,15 +149,15 @@ defmodule Protobuf.Mixfile do
     proto_src = path_in_protobuf_source(["src"])
 
     protoc!(
-      "-I #{proto_src} -I src -I test/protobuf/protoc/proto",
+      "-I #{proto_src} -I src -I test/protobuf/protoc/proto --elixir_opt=include_docs=true",
       "./generated",
       ["test/protobuf/protoc/proto/extension.proto"]
     )
 
     protoc!(
-      "-I test/protobuf/protoc/proto --elixir_opt=package_prefix=my",
+      "-I test/protobuf/protoc/proto --elixir_opt=package_prefix=my,include_docs=true",
       "./generated",
-      ["test/protobuf/protoc/proto/test.proto"]
+      ["test/protobuf/protoc/proto/test.proto", "test/protobuf/protoc/proto/service.proto"]
     )
 
     protoc!(
@@ -168,7 +168,7 @@ defmodule Protobuf.Mixfile do
     )
 
     protoc!(
-      "-I test/protobuf/protoc/proto",
+      "-I test/protobuf/protoc/proto --elixir_opt=include_docs=true",
       "./generated",
       ["test/protobuf/protoc/proto/no_package.proto"]
     )
@@ -194,7 +194,7 @@ defmodule Protobuf.Mixfile do
       google/protobuf/test_messages_proto3.proto
     )
 
-    protoc!("-I \"#{proto_root}\"", "./generated", files)
+    protoc!("-I \"#{proto_root}\" --elixir_opt=include_docs=true", "./generated", files)
   end
 
   defp gen_conformance_protos(_args) do
diff --git a/priv/templates/enum.ex.eex b/priv/templates/enum.ex.eex
index 19460b4d..f0ffdbee 100644
--- a/priv/templates/enum.ex.eex
+++ b/priv/templates/enum.ex.eex
@@ -1,7 +1,14 @@
 defmodule <%= @module %> do
-  <%= unless @module_doc? do %>
+  <%= if @module_doc? do %>
+  <%= if @comment != "" do %>
+  @moduledoc """
+<%= Protobuf.Protoc.Generator.Util.pad_comment(@comment, 2) %>
+  """
+  <% end %>
+  <% else %>
   @moduledoc false
   <% end %>
+
   use Protobuf, <%= @use_options %>
 
   <%= if @descriptor_fun_body do %>
diff --git a/priv/templates/extension.ex.eex b/priv/templates/extension.ex.eex
index 554fb2c1..7d4c25d5 100644
--- a/priv/templates/extension.ex.eex
+++ b/priv/templates/extension.ex.eex
@@ -1,5 +1,14 @@
 defmodule <%= @module %> do
+  <%= if @module_doc? do %>
+  <%= if @comment != "" do %>
+  @moduledoc """
+<%= Protobuf.Protoc.Generator.Util.pad_comment(@comment, 2) %>
+  """
+  <% end %>
+  <% else %>
   @moduledoc false
+  <% end %>
+
   use Protobuf, <%= @use_options %>
 
   <% if @extends == [], do: raise("Fuck! #{@module}") %>
diff --git a/priv/templates/message.ex.eex b/priv/templates/message.ex.eex
index b3b83971..cbe1b16f 100644
--- a/priv/templates/message.ex.eex
+++ b/priv/templates/message.ex.eex
@@ -1,7 +1,14 @@
 defmodule <%= @module %> do
-  <%= unless @module_doc? do %>
+  <%= if @module_doc? do %>
+  <%= if @comment != "" do %>
+  @moduledoc """
+<%= Protobuf.Protoc.Generator.Util.pad_comment(@comment, 2) %>
+  """
+  <% end %>
+  <% else %>
   @moduledoc false
   <% end %>
+
   use Protobuf<%= @use_options %>
 
   <%= if @descriptor_fun_body do %>
diff --git a/priv/templates/service.ex.eex b/priv/templates/service.ex.eex
index a42800dc..1cf8e682 100644
--- a/priv/templates/service.ex.eex
+++ b/priv/templates/service.ex.eex
@@ -1,7 +1,14 @@
 defmodule <%= @module %>.Service do
-  <%= unless @module_doc? do %>
+  <%= if @module_doc? do %>
+  <%= if @comment != "" do %>
+  @moduledoc """
+<%= Protobuf.Protoc.Generator.Util.pad_comment(@comment, 2) %>
+  """
+  <% end %>
+  <% else %>
   @moduledoc false
   <% end %>
+
   use GRPC.Service, name: <%= inspect(@service_name) %>, protoc_gen_elixir_version: "<%= @version %>"
 
   <%= if @descriptor_fun_body do %>
diff --git a/test/protobuf/protoc/cli_integration_test.exs b/test/protobuf/protoc/cli_integration_test.exs
index a2608be7..ee669c87 100644
--- a/test/protobuf/protoc/cli_integration_test.exs
+++ b/test/protobuf/protoc/cli_integration_test.exs
@@ -72,53 +72,52 @@ defmodule Protobuf.Protoc.CLIIntegrationTest do
       assert descriptor.name == "User"
     end
 
-    if System.version() < "1.14" do
-      test "include_docs option", %{tmp_dir: tmp_dir, proto_path: proto_path} do
-        protoc!([
-          "--proto_path=#{tmp_dir}",
-          "--elixir_out=#{tmp_dir}",
-          "--elixir_opt=include_docs=true",
-          "--plugin=./protoc-gen-elixir",
-          proto_path
-        ])
-
-        modules_and_docs = get_docs_and_clean_modules_on_exit("#{tmp_dir}/user.pb.ex")
-
-        assert [{Foo.User, docs}] = modules_and_docs
-        assert {:docs_v1, _, :elixir, _, module_doc, _, _} = docs
-        assert module_doc != :hidden
-      end
-
-      test "hides docs when include_docs is not true", %{tmp_dir: tmp_dir, proto_path: proto_path} do
-        protoc!([
-          "--proto_path=#{tmp_dir}",
-          "--elixir_out=#{tmp_dir}",
-          "--plugin=./protoc-gen-elixir",
-          proto_path
-        ])
-
-        modules_and_docs = get_docs_and_clean_modules_on_exit("#{tmp_dir}/user.pb.ex")
-
-        assert [{Foo.User, docs}] = modules_and_docs
-        assert {:docs_v1, _, :elixir, _, :hidden, _, _} = docs
-      end
-
-      defp get_docs_and_clean_modules_on_exit(path) do
-        modules_and_docs =
-          path
-          |> Code.compile_file()
-          |> Enum.map(fn {mod, bytecode} ->
-            {mod, fetch_docs_from_bytecode(bytecode)}
-          end)
-
-        on_exit(fn ->
-          modules_and_docs
-          |> Enum.map(fn {mod, _bytecode} -> mod end)
-          |> TestHelpers.purge_modules()
+    test "include_docs option", %{tmp_dir: tmp_dir, proto_path: proto_path} do
+      protoc!([
+        "--proto_path=#{tmp_dir}",
+        "--elixir_out=#{tmp_dir}",
+        "--elixir_opt=include_docs=true",
+        "--plugin=./protoc-gen-elixir",
+        proto_path
+      ])
+
+      modules_and_docs = get_docs_and_clean_modules_on_exit("#{tmp_dir}/user.pb.ex")
+      assert [{Foo.User, docs}] = modules_and_docs
+      assert {:docs_v1, _, :elixir, _, module_doc, _, _} = docs
+      assert module_doc != :hidden
+    end
+
+    test "hides docs when include_docs is not true", %{tmp_dir: tmp_dir, proto_path: proto_path} do
+      protoc!([
+        "--proto_path=#{tmp_dir}",
+        "--elixir_out=#{tmp_dir}",
+        "--plugin=./protoc-gen-elixir",
+        proto_path
+      ])
+
+      modules_and_docs = get_docs_and_clean_modules_on_exit("#{tmp_dir}/user.pb.ex")
+
+      assert [{Foo.User, docs}] = modules_and_docs
+      assert {:docs_v1, _, :elixir, _, :hidden, _, _} = docs
+    end
+
+    defp get_docs_and_clean_modules_on_exit(path) do
+      Code.put_compiler_option(:docs, true)
+
+      modules_and_docs =
+        path
+        |> Code.compile_file()
+        |> Enum.map(fn {mod, bytecode} ->
+          {mod, fetch_docs_from_bytecode(bytecode)}
         end)
 
+      on_exit(fn ->
         modules_and_docs
-      end
+        |> Enum.map(fn {mod, _docs} -> mod end)
+        |> TestHelpers.purge_modules()
+      end)
+
+      modules_and_docs
     end
 
     test "package_prefix mypkg", %{tmp_dir: tmp_dir, proto_path: proto_path} do
@@ -278,6 +277,114 @@ defmodule Protobuf.Protoc.CLIIntegrationTest do
     assert base_mod.get_extension(message, Bugs.Ext2.PbExtension, :first_name) == "Ext2 FN"
   end
 
+  test "generates documentation tags from comments", %{tmp_dir: tmp_dir} do
+    proto_path_one = Path.join(tmp_dir, "one.proto")
+    proto_path_two = Path.join(tmp_dir, "two.proto")
+
+    File.write!(proto_path_one, """
+    syntax = "proto3";
+
+    package tests;
+
+    // This is the comment for module One
+    message One {
+      // Comment about optional name field
+      optional string name = 1;
+
+      // Here is a nested message for One
+      message NestedOne {
+        optional string name = 1;
+      }
+
+      // Field referencing nested message
+      NestedOne nested_one = 2;
+    }
+    """)
+
+    File.write!(proto_path_two, """
+    syntax = "proto3";
+
+    import "one.proto";
+
+    package tests;
+
+    // This enum represents days of the week.
+    //
+    // It is a multi line description.
+    enum TestEnum {
+      // Monday is the first day
+      MONDAY = 0;
+      // Tuesday the second
+      TUESDAY = 1;
+    }
+
+    // This is a message that might be sent somewhere.
+    message Request {
+      // An enum of colors
+      enum Color {
+        RED = 0;
+        GREEN = 1; // My favorite color!
+        BLUE = 2;
+      }
+      //  optional imp.ImportedMessage imported_message = 2;
+      optional Color hue = 3;  // no default
+      optional One one = 4;
+
+      // This is a map field. It will generate map[int32]string.
+      map<int32, string> name_mapping = 14;
+    }
+    """)
+
+    protoc!([
+      "--proto_path=#{tmp_dir}",
+      "--elixir_out=#{tmp_dir}",
+      "--plugin=./protoc-gen-elixir",
+      "--elixir_opt=include_docs=true",
+      proto_path_one,
+      proto_path_two
+    ])
+
+    elixir_one = File.read!("#{tmp_dir}/one.pb.ex")
+    elixir_two = File.read!("#{tmp_dir}/two.pb.ex")
+
+    assert elixir_one =~ """
+           defmodule Tests.One do
+             @moduledoc \"\"\"
+             This is the comment for module One
+             \"\"\"
+           """
+
+    assert elixir_one =~ """
+           defmodule Tests.One.NestedOne do
+             @moduledoc \"\"\"
+             Here is a nested message for One
+             \"\"\"
+           """
+
+    assert elixir_two =~ """
+           defmodule Tests.TestEnum do
+             @moduledoc \"\"\"
+             This enum represents days of the week.
+
+             It is a multi line description.
+             \"\"\"
+           """
+
+    assert elixir_two =~ """
+           defmodule Tests.Request do
+             @moduledoc \"\"\"
+             This is a message that might be sent somewhere.
+             \"\"\"
+           """
+
+    assert elixir_two =~ """
+           defmodule Tests.Request.Color do
+             @moduledoc \"\"\"
+             An enum of colors
+             \"\"\"
+           """
+  end
+
   @tag :skip
   test "extensions defined and used in the same protoc call", %{tmp_dir: tmp_dir} do
     proto_path = Path.join(tmp_dir, "extensions.proto")
diff --git a/test/protobuf/protoc/generator/enum_test.exs b/test/protobuf/protoc/generator/enum_test.exs
index 653cb75f..0a839197 100644
--- a/test/protobuf/protoc/generator/enum_test.exs
+++ b/test/protobuf/protoc/generator/enum_test.exs
@@ -111,17 +111,33 @@ defmodule Protobuf.Protoc.Generator.EnumTest do
   end
 
   describe "generate/2 include_docs" do
-    test "does not include `@moduledoc false` when flag is true" do
-      ctx = %Context{include_docs?: true}
-      desc = %Google.Protobuf.EnumDescriptorProto{name: "valueType"}
-
-      {_module, msg} = Generator.generate(ctx, desc)
-
-      refute msg =~ "@moduledoc\n"
+    test "includes enum comment for `@moduledoc` when flag is true" do
+      test_pb = TestHelpers.read_generated_file("test.pb.ex")
+
+      assert test_pb =~ """
+             defmodule My.Test.Days do
+               @moduledoc \"\"\"
+               This enum represents days of the week.
+               \"\"\"
+             """
+
+      assert test_pb =~ """
+             defmodule My.Test.HatType do
+               @moduledoc \"\"\"
+               This enum represents different kinds of hats.
+               \"\"\"
+             """
+
+      assert test_pb =~ """
+             defmodule My.Test.Request.Color do
+               @moduledoc \"\"\"
+               This enum represents three different colors.
+               \"\"\"
+             """
     end
 
     test "includes `@moduledoc false` by default" do
-      ctx = %Context{}
+      ctx = %Context{include_docs?: false}
       desc = %Google.Protobuf.EnumDescriptorProto{name: "valueType"}
 
       {_module, msg} = Generator.generate(ctx, desc)
diff --git a/test/protobuf/protoc/generator/message_test.exs b/test/protobuf/protoc/generator/message_test.exs
index 7ce24c25..3bef634c 100644
--- a/test/protobuf/protoc/generator/message_test.exs
+++ b/test/protobuf/protoc/generator/message_test.exs
@@ -795,17 +795,21 @@ defmodule Protobuf.Protoc.Generator.MessageTest do
   end
 
   describe "generate/2 include_docs" do
-    test "does not include `@moduledoc false` when flag is true" do
-      ctx = %Context{include_docs?: true}
-      desc = %Google.Protobuf.DescriptorProto{name: "Foo"}
+    test "includes message comment for `@moduledoc` when flag is true" do
+      test_pb = TestHelpers.read_generated_file("test.pb.ex")
 
-      {[], [{_mod, msg}]} = Generator.generate(ctx, desc)
+      assert test_pb =~ """
+             defmodule My.Test.Request do
+               @moduledoc \"\"\"
+               This is a message that might be sent somewhere.
 
-      refute msg =~ "@moduledoc\n"
+               Here is another line for a documentation example.
+               \"\"\"
+             """
     end
 
     test "includes `@moduledoc false` by default" do
-      ctx = %Context{}
+      ctx = %Context{include_docs?: false}
       desc = %Google.Protobuf.DescriptorProto{name: "Foo"}
 
       {[], [{_mod, msg}]} = Generator.generate(ctx, desc)
diff --git a/test/protobuf/protoc/generator/service_test.exs b/test/protobuf/protoc/generator/service_test.exs
index f626ad75..8a7a881e 100644
--- a/test/protobuf/protoc/generator/service_test.exs
+++ b/test/protobuf/protoc/generator/service_test.exs
@@ -64,17 +64,47 @@ defmodule Protobuf.Protoc.Generator.ServiceTest do
   end
 
   describe "generate/2 include_docs" do
-    test "does not include `@moduledoc false` when flag is true" do
-      ctx = %Context{include_docs?: true}
-      desc = %Google.Protobuf.ServiceDescriptorProto{name: "ServiceFoo"}
+    test "includes service comment for `@moduledoc` when flag is true" do
+      ctx = %Context{
+        package: "foo",
+        include_docs?: true,
+        comments: %{
+          "" =>
+            "An example test service that has\n" <>
+              "a test method. It expects a Request\n" <>
+              "and returns a Reply."
+        },
+        dep_type_mapping: %{
+          ".foo.Input0" => %{type_name: "Foo.Input0"},
+          ".foo.Output0" => %{type_name: "Foo.Output0"}
+        },
+        module_prefix: "Foo"
+      }
 
-      {_module, msg} = Generator.generate(ctx, desc)
+      desc = %Google.Protobuf.ServiceDescriptorProto{
+        name: "ServiceFoo",
+        method: [
+          %Google.Protobuf.MethodDescriptorProto{
+            name: "MethodA",
+            input_type: ".foo.Input0",
+            output_type: ".foo.Output0"
+          }
+        ]
+      }
+
+      assert {"Foo.ServiceFoo", msg} = Generator.generate(ctx, desc)
 
-      refute msg =~ "@moduledoc"
+      assert msg =~ """
+               @moduledoc \"\"\"
+               An example test service that has
+               a test method. It expects a Request
+               and returns a Reply.
+               \"\"\"
+             """
     end
 
     test "includes `@moduledoc false` by default" do
-      ctx = %Context{}
+      ctx = %Context{include_docs?: false}
       desc = %Google.Protobuf.ServiceDescriptorProto{name: "ServiceFoo"}
 
       {_module, msg} = Generator.generate(ctx, desc)
diff --git a/test/protobuf/protoc/proto/service.proto b/test/protobuf/protoc/proto/service.proto
new file mode 100644
index 00000000..250a0d7d
--- /dev/null
+++ b/test/protobuf/protoc/proto/service.proto
@@ -0,0 +1,13 @@
+syntax = "proto2";
+
+// This package holds interesting messages.
+package test;  // dotted package name
+
+import "test.proto";
+
+// An example test service that has
+// a test method. It expects a Request
+// and returns a Reply.
+service TestService {
+  rpc test (Request) returns (Reply);
+}
diff --git a/test/protobuf/protoc/proto/test.proto b/test/protobuf/protoc/proto/test.proto
index 997e286f..11b27ec7 100644
--- a/test/protobuf/protoc/proto/test.proto
+++ b/test/protobuf/protoc/proto/test.proto
@@ -3,6 +3,7 @@ syntax = "proto2";
 // This package holds interesting messages.
 package test;  // dotted package name
 
+// This enum represents different kinds of hats.
 enum HatType {
   // deliberately skipping 0
   FEDORA = 1;
@@ -19,7 +20,10 @@ enum Days {
 }
 
 // This is a message that might be sent somewhere.
+//
+// Here is another line for a documentation example.
 message Request {
+  // This enum represents three different colors.
   enum Color {
     RED = 0;
     GREEN = 1;
@@ -71,10 +75,12 @@ message OtherBase {
 }
 
 message ReplyExtensions {
+  // Extends Reply
   extend Reply {
     optional double time = 101;
     optional ReplyExtensions carrot = 105;
   }
+  // Yet another base message
   extend OtherBase {
     optional ReplyExtensions donut = 101;
   }
diff --git a/test/test_helper.exs b/test/test_helper.exs
index 8e60cbd1..0379e2da 100644
--- a/test/test_helper.exs
+++ b/test/test_helper.exs
@@ -28,6 +28,12 @@ defmodule Protobuf.TestHelpers do
     Map.put(context, :tmp_dir, tmp_dir_name)
   end
 
+  def read_generated_file(relative_path) do
+    [__DIR__, "../generated", relative_path]
+    |> Path.join()
+    |> File.read!()
+  end
+
   def get_type_spec_as_string(module, bytecode, type)
       when is_atom(module) and is_binary(bytecode) and is_atom(type) do
     # This code is taken from Code.Typespec in Elixir (v1.13 in particular).