From ede4ec6c0a504158861795a5dbb99ac094d44636 Mon Sep 17 00:00:00 2001
From: Wojtek Mach <wojtek@wojtekmach.pl>
Date: Tue, 13 Jul 2021 16:05:21 +0200
Subject: [PATCH] Change type_data/2 to take module_state

---
 lib/ex_doc/language.ex        |  2 ++
 lib/ex_doc/language/elixir.ex | 25 +++++++++++++++++++++++--
 lib/ex_doc/language/erlang.ex |  9 ++++++++-
 lib/ex_doc/retriever.ex       | 22 +++++-----------------
 4 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/lib/ex_doc/language.ex b/lib/ex_doc/language.ex
index e64824d92..33bcee621 100644
--- a/lib/ex_doc/language.ex
+++ b/lib/ex_doc/language.ex
@@ -94,6 +94,8 @@ defmodule ExDoc.Language do
   """
   @callback type_data(entry :: tuple(), spec :: term()) :: data
             when data: %{
+                   type: :type | :opaque,
+                   line: non_neg_integer(),
                    spec: spec_ast(),
                    signature_fallback: (() -> String.t()) | nil
                  }
diff --git a/lib/ex_doc/language/elixir.ex b/lib/ex_doc/language/elixir.ex
index 1e6c13dfb..6988dfc9b 100644
--- a/lib/ex_doc/language/elixir.ex
+++ b/lib/ex_doc/language/elixir.ex
@@ -82,11 +82,15 @@ defmodule ExDoc.Language.Elixir do
   end
 
   @impl true
-  def type_data(entry, spec) do
-    {{kind, _name, arity}, _anno, _signature, _doc, _metadata} = entry
+  def type_data(entry, module_state) do
+    {{kind, name, arity}, _anno, _signature, _doc, _metadata} = entry
+
+    %{type: type, spec: spec, line: line} = type_from_module_state(name, arity, module_state)
     spec = spec |> Code.Typespec.type_to_quoted() |> process_type_ast(kind)
 
     %{
+      type: type,
+      line: line,
       spec: spec,
       signature_fallback: fn ->
         get_typespec_signature(spec, arity)
@@ -94,6 +98,23 @@ defmodule ExDoc.Language.Elixir do
     }
   end
 
+  @doc false
+  def type_from_module_state(name, arity, module_state) do
+    Enum.find_value(module_state.abst_code, fn
+      {:attribute, anno, type, {^name, _, args} = spec} ->
+        if type in [:opaque, :type] and length(args) == arity do
+          %{
+            type: type,
+            spec: spec,
+            line: anno_line(anno)
+          }
+        end
+
+      _ ->
+        nil
+    end)
+  end
+
   @impl true
   def autolink_doc(ast, opts) do
     config = struct!(Autolink, opts)
diff --git a/lib/ex_doc/language/erlang.ex b/lib/ex_doc/language/erlang.ex
index a4832a9be..e7c23e1fb 100644
--- a/lib/ex_doc/language/erlang.ex
+++ b/lib/ex_doc/language/erlang.ex
@@ -61,8 +61,15 @@ defmodule ExDoc.Language.Erlang do
   end
 
   @impl true
-  def type_data(_entry, spec) do
+  def type_data(entry, module_state) do
+    {{_kind, name, arity}, _anno, _signature, _doc, _metadata} = entry
+
+    %{type: type, spec: spec, line: line} =
+      ExDoc.Language.Elixir.type_from_module_state(name, arity, module_state)
+
     %{
+      type: type,
+      line: line,
       spec: {:attribute, 0, :type, spec},
       signature_fallback: fn -> nil end
     }
diff --git a/lib/ex_doc/retriever.ex b/lib/ex_doc/retriever.ex
index 783b00aac..62ded0561 100644
--- a/lib/ex_doc/retriever.ex
+++ b/lib/ex_doc/retriever.ex
@@ -406,36 +406,24 @@ defmodule ExDoc.Retriever do
     doc_line = anno_line(anno)
     annotations = annotations_from_metadata(metadata)
 
-    {:attribute, anno, type, spec} =
-      Enum.find(module_state.abst_code, fn
-        {:attribute, _, type, {^name, _, args}} ->
-          type in [:opaque, :type] and length(args) == arity
-
-        _ ->
-          false
-      end)
-
-    line = anno_line(anno)
-    type_data = module_state.language.type_data(type_entry, spec)
-    spec = type_data.spec
-
+    type_data = module_state.language.type_data(type_entry, module_state)
     signature = signature(signature) || type_data.signature_fallback.()
 
-    annotations = if type == :opaque, do: ["opaque" | annotations], else: annotations
+    annotations = if type_data.type == :opaque, do: ["opaque" | annotations], else: annotations
     doc_ast = doc_ast(content_type, doc, file: source.path)
 
     %ExDoc.TypeNode{
       id: "#{name}/#{arity}",
       name: name,
       arity: arity,
-      type: type,
-      spec: spec,
+      type: type_data.type,
+      spec: type_data.spec,
       deprecated: metadata[:deprecated],
       doc: doc_ast,
       doc_line: doc_line,
       signature: signature,
       source_path: source.path,
-      source_url: source_link(source, line),
+      source_url: source_link(source, type_data.line),
       annotations: annotations
     }
   end