Resolve types from transitive import public deps

Closes #351.

Author: whatyouhide

lib/protobuf/protoc/cli.ex

@@ -158,14 +158,56 @@ defmodule Protobuf.Protoc.CLI do
           Context.t()
   def find_types(%Context{} = ctx, descs, files_to_generate)
       when is_list(descs) and is_list(files_to_generate) do
-    global_type_mapping =
+    base_mapping =
       Map.new(descs, fn %Google.Protobuf.FileDescriptorProto{name: filename} = desc ->
         {filename, find_types_in_proto(ctx, desc, files_to_generate)}
       end)
 
+    # `import public` re-exports types from the imported file. A file that
+    # depends on `shared.proto` should see types from anything `shared.proto`
+    # publicly imports (transitively). We fold those types into each file's
+    # mapping here so `get_dep_type_mapping/3` keeps working unchanged.
+    public_deps_by_file =
+      Map.new(descs, fn %Google.Protobuf.FileDescriptorProto{} = desc ->
+        deps = desc.dependency
+        public_deps = Enum.map(desc.public_dependency, &Enum.at(deps, &1))
+        {desc.name, public_deps}
+      end)
+
+    global_type_mapping =
+      Map.new(base_mapping, fn {filename, mapping} ->
+        transitive =
+          collect_public_deps([filename], public_deps_by_file, MapSet.new([filename]))
+          |> MapSet.delete(filename)
+
+        merged =
+          Enum.reduce(transitive, mapping, fn other, acc ->
+            Map.merge(acc, Map.get(base_mapping, other, %{}))
+          end)
+
+        {filename, merged}
+      end)
+
     %{ctx | global_type_mapping: global_type_mapping}
   end
 
+  defp collect_public_deps([], _public_deps_by_file, visited), do: visited
+
+  defp collect_public_deps([filename | rest], public_deps_by_file, visited) do
+    deps = Map.get(public_deps_by_file, filename, [])
+
+    {new_deps, visited} =
+      Enum.reduce(deps, {[], visited}, fn dep, {to_visit, visited} ->
+        if MapSet.member?(visited, dep) do
+          {to_visit, visited}
+        else
+          {[dep | to_visit], MapSet.put(visited, dep)}
+        end
+      end)
+
+    collect_public_deps(new_deps ++ rest, public_deps_by_file, visited)
+  end
+
   defp find_types_in_proto(
          %Context{} = ctx,
          %Google.Protobuf.FileDescriptorProto{} = desc,

test/protobuf/protoc/cli_test.exs

@@ -226,5 +226,47 @@ defmodule Protobuf.Protoc.CLITest do
                }
              }
     end
+
+    test "transitively folds in types from `import public` dependencies (#351)" do
+      descs = [
+        %FileDescriptorProto{
+          name: "arm_status.proto",
+          package: "mypackage.arm_status",
+          enum_type: [%EnumDescriptorProto{name: "ArmStatus"}]
+        },
+        %FileDescriptorProto{
+          name: "leg_status.proto",
+          package: "mypackage.leg_status",
+          enum_type: [%EnumDescriptorProto{name: "LegStatus"}]
+        },
+        %FileDescriptorProto{
+          name: "shared.proto",
+          package: "mypackage",
+          dependency: ["arm_status.proto", "leg_status.proto"],
+          public_dependency: [0, 1]
+        },
+        %FileDescriptorProto{
+          name: "person.proto",
+          package: "mypackage.person",
+          dependency: ["shared.proto"],
+          message_type: [%DescriptorProto{name: "Person"}]
+        }
+      ]
+
+      mapping = find_types(%Context{}, descs, Enum.map(descs, & &1.name)).global_type_mapping
+
+      # `shared.proto` re-exports the enums from its public deps, so looking up
+      # by `shared.proto` includes them.
+      assert mapping["shared.proto"] == %{
+               ".mypackage.arm_status.ArmStatus" => %{type_name: "Mypackage.ArmStatus.ArmStatus"},
+               ".mypackage.leg_status.LegStatus" => %{type_name: "Mypackage.LegStatus.LegStatus"}
+             }
+
+      # `person.proto` only directly depends on `shared.proto` but the publicly
+      # re-exported types are still reachable.
+      assert mapping["person.proto"] == %{
+               ".mypackage.person.Person" => %{type_name: "Mypackage.Person.Person"}
+             }
+    end
   end
 end