fix descriptor recursion

Author: zhihuizhang17

.gitignore

@@ -31,3 +31,4 @@ priv/plts
 # Ignore cursor and memory-bank
 .cursor/
 memory-bank/
+custom_modes/

lib/grpc_reflection/service/builder.ex

@@ -2,149 +2,243 @@ defmodule GrpcReflection.Service.Builder do
   @moduledoc false
 
   alias Google.Protobuf.FileDescriptorProto
-  alias GrpcReflection.Service.State
-  alias GrpcReflection.Service.Builder.Util
+  alias GrpcReflection.Service.Builder.Acc
   alias GrpcReflection.Service.Builder.Extensions
+  alias GrpcReflection.Service.Builder.Util
+  alias GrpcReflection.Service.State
 
+  @spec build_reflection_tree(any()) ::
+          {:error, <<_::216>>} | {:ok, GrpcReflection.Service.State.t()}
   def build_reflection_tree(services) do
     with :ok <- Util.validate_services(services) do
-      tree =
-        Enum.reduce(services, State.new(services), fn service, state ->
-          new_state = process_service(service)
-          State.merge(state, new_state)
+      acc =
+        Enum.reduce(services, %Acc{services: services}, fn service, acc ->
+          process_service(acc, service)
         end)
 
-      {:ok, tree}
+      {:ok, finalize(acc)}
     end
   end
 
-  defp process_service(service) do
-    service_name = service.__meta__(:name)
-    service_response = build_response(service_name, service)
+  defp finalize(%Acc{} = acc) do
+    files = build_files(acc)
+
+    symbols =
+      acc.symbol_info
+      |> Enum.reduce(%{}, fn {symbol, info}, map ->
+        Map.put(map, symbol, Map.fetch!(files, info.file))
+      end)
+      |> then(fn base ->
+        Enum.reduce(acc.aliases, base, fn {alias_symbol, target_symbol}, map ->
+          payload = Map.fetch!(map, target_symbol)
+          Map.put(map, alias_symbol, payload)
+        end)
+      end)
+
+    State.new(acc.services)
+    |> State.add_files(Map.merge(files, acc.extension_files))
+    |> State.add_symbols(symbols)
+    |> State.add_extensions(acc.extensions)
+  end
+
+  defp build_files(%Acc{} = acc) do
+    acc.symbol_info
+    |> Enum.group_by(fn {_symbol, info} -> info.file end)
+    |> Enum.reduce(%{}, fn {file_name, entries}, files ->
+      descriptors = Enum.map(entries, fn {_symbol, info} -> info end)
+
+      syntax = descriptors |> List.first() |> Map.fetch!(:syntax)
+      package = descriptors |> List.first() |> Map.fetch!(:package)
+
+      dependencies =
+        descriptors
+        |> Enum.flat_map(& &1.deps)
+        |> Enum.map(&resolve_dependency_file(&1, acc))
+        |> Enum.reject(&is_nil/1)
+        |> Enum.reject(&(&1 == file_name))
+        |> Enum.uniq()
+
+      {messages, enums, services} =
+        Enum.reduce(descriptors, {[], [], []}, fn %{descriptor: descriptor},
+                                                  {messages, enums, services} ->
+          case descriptor do
+            %Google.Protobuf.DescriptorProto{} ->
+              {[descriptor | messages], enums, services}
+
+            %Google.Protobuf.EnumDescriptorProto{} ->
+              {messages, [descriptor | enums], services}
+
+            %Google.Protobuf.ServiceDescriptorProto{} ->
+              {messages, enums, [descriptor | services]}
+          end
+        end)
+        |> then(fn {messages, enums, services} ->
+          {Enum.reverse(messages), Enum.reverse(enums), Enum.reverse(services)}
+        end)
+
+      file_proto = %FileDescriptorProto{
+        name: file_name,
+        package: package,
+        dependency: dependencies,
+        syntax: syntax,
+        message_type: messages,
+        enum_type: enums,
+        service: services
+      }
 
-    State.new()
-    |> State.add_symbols(%{service_name => service_response})
-    |> State.add_files(%{(service_name <> ".proto") => service_response})
-    |> trace_service_refs(service)
+      Map.put(files, file_name, %{file_descriptor_proto: [FileDescriptorProto.encode(file_proto)]})
+    end)
   end
 
-  defp trace_service_refs(state, module) do
-    service_name = module.__meta__(:name)
-    methods = get_descriptor(module).method
+  defp process_service(%Acc{} = acc, service) do
+    service_name = service.__meta__(:name)
+    {acc, _} = register_symbol(acc, service_name, service, :service)
+
+    methods = get_descriptor(service).method
 
-    module.__rpc_calls__()
-    |> Enum.reduce(state, fn call, state ->
-      {function, {request, _}, {response, _}, _} = call
+    Enum.reduce(service.__rpc_calls__(), acc, fn call, acc ->
+      {function, {req, _}, {resp, _}, _} = call
 
       %{input_type: req_symbol, output_type: resp_symbol} =
         Enum.find(methods, fn method -> method.name == to_string(function) end)
 
-      call_symbol = service_name <> "." <> to_string(function)
-      call_response = build_response(service_name, module)
       req_symbol = Util.trim_symbol(req_symbol)
-      req_response = build_response(req_symbol, request)
       resp_symbol = Util.trim_symbol(resp_symbol)
-      resp_response = build_response(resp_symbol, response)
-
-      state
-      |> Extensions.add_extensions(service_name, module)
-      |> State.add_symbols(%{
-        call_symbol => call_response,
-        req_symbol => req_response,
-        resp_symbol => resp_response
-      })
-      |> State.add_files(%{
-        (req_symbol <> ".proto") => req_response,
-        (resp_symbol <> ".proto") => resp_response
-      })
-      |> Extensions.add_extensions(req_symbol, request)
-      |> Extensions.add_extensions(resp_symbol, response)
-      |> trace_message_refs(req_symbol, request)
-      |> trace_message_refs(resp_symbol, response)
+
+      method_symbol = service_name <> "." <> to_string(function)
+
+      acc
+      |> register_alias(method_symbol, service_name)
+      |> Extensions.add_extensions(service_name, service)
+      |> process_message(req_symbol, req)
+      |> process_message(resp_symbol, resp)
     end)
   end
 
-  defp trace_message_refs(state, parent_symbol, module) do
-    case module.descriptor() do
-      %{field: fields} ->
-        trace_message_fields(state, parent_symbol, module, fields)
+  defp process_message(%Acc{} = acc, nil, _module, _root_symbol), do: acc
 
-      _ ->
-        state
-    end
-  end
+  defp process_message(%Acc{} = acc, symbol, module, root_symbol) do
+    symbol = Util.trim_symbol(symbol)
+    root_symbol = root_symbol || symbol
 
-  defp trace_message_fields(state, parent_symbol, module, fields) do
-    # nested types arent a "separate file", they return their parents' response
-    nested_types = Util.get_nested_types(parent_symbol, module.descriptor())
-
-    module.__message_props__().field_props
-    |> Map.values()
-    |> Enum.map(fn %{name: name, type: type} ->
-      %{
-        mod:
-          case type do
-            {_, mod} -> mod
-            mod -> mod
-          end,
-        symbol: Enum.find(fields, fn f -> f.name == name end).type_name
-      }
-    end)
-    |> Enum.reject(fn %{symbol: s} -> s == nil end)
-    |> Enum.reduce(state, fn %{mod: mod, symbol: symbol}, state ->
-      symbol = Util.trim_symbol(symbol)
+    {acc, already_processed} =
+      if root_symbol == symbol do
+        register_symbol(acc, symbol, module, :message)
+      else
+        acc = register_alias(acc, symbol, root_symbol)
 
-      response =
-        if symbol in nested_types do
-          build_response(parent_symbol, module)
+        if MapSet.member?(acc.visited, symbol) do
+          {acc, true}
         else
-          build_response(symbol, mod)
+          {acc, false}
         end
+      end
 
-      state
-      |> Extensions.add_extensions(symbol, mod)
-      |> State.add_symbols(%{symbol => response})
-      |> State.add_files(%{(symbol <> ".proto") => response})
-      |> trace_message_refs(symbol, mod)
-    end)
+    if already_processed do
+      acc
+    else
+      acc = %{acc | visited: MapSet.put(acc.visited, symbol)}
+      acc = Extensions.add_extensions(acc, symbol, module)
+
+      case module.descriptor() do
+        %{field: fields} = descriptor ->
+          nested_symbols = Util.get_nested_types(symbol, descriptor)
+
+          module.__message_props__().field_props
+          |> Map.values()
+          |> Enum.map(fn %{name: name, type: type} ->
+            %{
+              mod:
+                case type do
+                  {_, mod} -> mod
+                  mod -> mod
+                end,
+              symbol: Enum.find(fields, fn f -> f.name == name end).type_name
+            }
+          end)
+          |> Enum.reject(fn %{symbol: s} -> is_nil(s) end)
+          |> Enum.reduce(acc, fn %{mod: mod, symbol: child_symbol}, acc ->
+            child_symbol = Util.trim_symbol(child_symbol)
+
+            if child_symbol in nested_symbols do
+              process_message(acc, child_symbol, mod, root_symbol)
+            else
+              process_message(acc, child_symbol, mod)
+            end
+          end)
+
+        _ ->
+          acc
+      end
+    end
   end
 
-  defp build_response(symbol, module) do
-    # we build our own file responses, so unwrap any present
-    descriptor = get_descriptor(module)
-
-    dependencies =
-      descriptor
-      |> Util.types_from_descriptor()
-      |> Enum.uniq()
-      |> Kernel.--(Util.get_nested_types(symbol, descriptor))
-      |> Enum.map(fn name ->
-        Util.trim_symbol(name) <> ".proto"
-      end)
-
-    syntax = Util.get_syntax(module)
+  defp process_message(%Acc{} = acc, symbol, module) do
+    process_message(acc, symbol, module, nil)
+  end
 
-    response_stub =
-      %FileDescriptorProto{
-        name: symbol <> ".proto",
+  defp register_symbol(%Acc{} = acc, symbol, module, kind) do
+    symbol = Util.trim_symbol(symbol)
+
+    if Map.has_key?(acc.symbol_info, symbol) do
+      {acc, true}
+    else
+      descriptor = get_descriptor(module)
+
+      info = %{
+        descriptor: descriptor,
+        deps:
+          descriptor
+          |> Util.types_from_descriptor()
+          |> Enum.map(&Util.trim_symbol/1)
+          |> Enum.uniq(),
+        file: Util.proto_filename(module),
+        syntax: Util.get_syntax(module),
         package: Util.get_package(symbol),
-        dependency: dependencies,
-        syntax: syntax
+        kind: kind
       }
 
-    unencoded_payload =
-      case descriptor = descriptor do
-        %Google.Protobuf.DescriptorProto{} -> %{response_stub | message_type: [descriptor]}
-        %Google.Protobuf.ServiceDescriptorProto{} -> %{response_stub | service: [descriptor]}
-        %Google.Protobuf.EnumDescriptorProto{} -> %{response_stub | enum_type: [descriptor]}
-      end
+      {%{
+         acc
+         | symbol_info: Map.put(acc.symbol_info, symbol, info),
+           visited: MapSet.put(acc.visited, symbol)
+       }, false}
+    end
+  end
+
+  defp register_alias(%Acc{} = acc, alias_symbol, target_symbol) do
+    alias_symbol = Util.trim_symbol(alias_symbol)
+    target_symbol = Util.trim_symbol(target_symbol)
+
+    cond do
+      alias_symbol == target_symbol -> acc
+      Map.get(acc.aliases, alias_symbol) == target_symbol -> acc
+      true -> %{acc | aliases: Map.put(acc.aliases, alias_symbol, target_symbol)}
+    end
+  end
 
-    %{file_descriptor_proto: [FileDescriptorProto.encode(unencoded_payload)]}
+  defp resolve_dependency_file(nil, _acc), do: nil
+
+  defp resolve_dependency_file(symbol, %Acc{} = acc) do
+    symbol = Util.trim_symbol(symbol)
+
+    cond do
+      info = Map.get(acc.symbol_info, symbol) ->
+        info.file
+
+      target = Map.get(acc.aliases, symbol) ->
+        acc.symbol_info
+        |> Map.get(target)
+        |> case do
+          nil -> symbol <> ".proto"
+          info -> info.file
+        end
+
+      true ->
+        symbol <> ".proto"
+    end
   end
 
-  # protoc with the elixir generator and protobuf.generate slightly differ for how they
-  # generate descriptors.  Use this to potentially unwrap the service proto when dealing
-  # with descriptors that could come from a service module.
   defp get_descriptor(module) do
     case module.descriptor() do
       %FileDescriptorProto{service: [proto]} -> proto

lib/grpc_reflection/service/builder/acc.ex

@@ -0,0 +1,10 @@
+defmodule GrpcReflection.Service.Builder.Acc do
+  @moduledoc false
+
+  defstruct services: [],
+            symbol_info: %{},
+            aliases: %{},
+            visited: MapSet.new(),
+            extension_files: %{},
+            extensions: %{}
+end