From b117caed9303d1d9936674ffdc37182f9248b53d Mon Sep 17 00:00:00 2001
From: Lukasz Samson <lukaszsamson@gmail.com>
Date: Thu, 19 Sep 2024 22:13:22 +0200
Subject: [PATCH] better type inference in map guards

---
 lib/elixir_sense/core/type_inference/guard.ex | 37 ++++++++++++++--
 .../core/type_inference/guard_test.exs        | 44 ++++++++++++++++++-
 2 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/lib/elixir_sense/core/type_inference/guard.ex b/lib/elixir_sense/core/type_inference/guard.ex
index 9f6c4441..9867045c 100644
--- a/lib/elixir_sense/core/type_inference/guard.ex
+++ b/lib/elixir_sense/core/type_inference/guard.ex
@@ -77,11 +77,19 @@ defmodule ElixirSense.Core.TypeInference.Guard do
     end)
   end
 
+  # {{:., _, [target, key]}, _, []}
+  def type_information_from_guards({{:., _, [target, key]}, _, []}) when is_atom(key) do
+    case extract_var_type(target, {:map, [{key, {:atom, true}}], []}) do
+      nil -> %{}
+      {var, type} -> %{var => type}
+    end
+  end
+
   # Standalone variable: func my_func(x) when x
   def type_information_from_guards({var, meta, context}) when is_atom(var) and is_atom(context) do
     case Keyword.fetch(meta, :version) do
       {:ok, version} ->
-        %{{var, version} => :boolean}
+        %{{var, version} => {:atom, true}}
 
       _ ->
         %{}
@@ -94,8 +102,7 @@ defmodule ElixirSense.Core.TypeInference.Guard do
         {{:., _dot_meta, [:erlang, fun]}, _call_meta, params}, acc
         when is_atom(fun) and is_list(params) ->
           with {type, binding} <- guard_predicate_type(fun, params),
-               {var, meta, context} when is_atom(var) and is_atom(context) <- binding,
-               {:ok, version} <- Keyword.fetch(meta, :version) do
+               {{var, version}, type} <- extract_var_type(binding, type) do
             # If we found the predicate type, we can prematurely exit traversing the subtree
             {nil, Map.put(acc, {var, version}, type)}
           else
@@ -115,6 +122,22 @@ defmodule ElixirSense.Core.TypeInference.Guard do
     acc
   end
 
+  defp extract_var_type({var, meta, context}, type) when is_atom(var) and is_atom(context) do
+    case Keyword.fetch(meta, :version) do
+      {:ok, version} ->
+        {{var, version}, type}
+
+      _ ->
+        nil
+    end
+  end
+
+  defp extract_var_type({{:., _, [target, key]}, _, []}, type) when is_atom(key) do
+    extract_var_type(target, {:map, [{key, type}], []})
+  end
+
+  defp extract_var_type(_, _), do: nil
+
   # TODO div and rem only work on first arg
   defp guard_predicate_type(p, [first | _])
        when p in [
@@ -221,11 +244,19 @@ defmodule ElixirSense.Core.TypeInference.Guard do
     {type_of(value), lhs}
   end
 
+  defp guard_predicate_type(p, [{{:., _, _}, _, []} = lhs, value]) when p in [:==, :===] do
+    {type_of(value), lhs}
+  end
+
   defp guard_predicate_type(p, [value, {variable, _, context} = rhs])
        when p in [:==, :===] and is_atom(variable) and is_atom(context) do
     guard_predicate_type(p, [rhs, value])
   end
 
+  defp guard_predicate_type(p, [value, {{:., _, _}, _, []} = rhs]) when p in [:==, :===] do
+    guard_predicate_type(p, [rhs, value])
+  end
+
   defp guard_predicate_type(:is_map, [first | _]), do: {{:map, [], nil}, first}
   defp guard_predicate_type(:is_non_struct_map, [first | _]), do: {{:map, [], nil}, first}
   defp guard_predicate_type(:map_size, [first | _]), do: {{:map, [], nil}, first}
diff --git a/test/elixir_sense/core/type_inference/guard_test.exs b/test/elixir_sense/core/type_inference/guard_test.exs
index 2aaefd01..a89535ea 100644
--- a/test/elixir_sense/core/type_inference/guard_test.exs
+++ b/test/elixir_sense/core/type_inference/guard_test.exs
@@ -60,7 +60,7 @@ defmodule ElixirSense.Core.TypeInference.GuardTest do
     test "infers type from naked var" do
       guard_expr = quote(do: x) |> expand()
       result = Guard.type_information_from_guards(guard_expr)
-      assert result == %{{:x, 0} => :boolean}
+      assert result == %{{:x, 0} => {:atom, true}}
     end
 
     # 1. Simple guards
@@ -270,7 +270,7 @@ defmodule ElixirSense.Core.TypeInference.GuardTest do
       guard = quote(do: is_number(x) or is_atom(x) or (is_nil(x) or x)) |> expand()
 
       result = Guard.type_information_from_guards(guard)
-      assert result == %{{:x, 0} => {:union, [:number, :atom, {:atom, nil}, :boolean]}}
+      assert result == %{{:x, 0} => {:union, [:number, :atom, {:atom, nil}, {:atom, true}]}}
     end
 
     test "handles nested when" do
@@ -280,4 +280,44 @@ defmodule ElixirSense.Core.TypeInference.GuardTest do
       assert result == %{{:x, 0} => {:union, [:number, :binary]}}
     end
   end
+
+  describe "guard on map field" do
+    test "naked" do
+      guard = quote(do: x.foo) |> expand()
+
+      result = Guard.type_information_from_guards(guard)
+      assert result == %{{:x, 0} => {:map, [{:foo, {:atom, true}}], []}}
+    end
+
+    test "naked nested" do
+      guard = quote(do: x.foo.bar) |> expand()
+
+      result = Guard.type_information_from_guards(guard)
+      assert result == %{{:x, 0} => {:map, [{:foo, {:map, [{:bar, {:atom, true}}], []}}], []}}
+    end
+
+    test "simple" do
+      guard = quote(do: is_atom(x.foo)) |> expand()
+
+      result = Guard.type_information_from_guards(guard)
+      assert result == %{{:x, 0} => {:map, [{:foo, :atom}], []}}
+    end
+
+    test "nested" do
+      guard = quote(do: is_atom(x.foo.bar.baz)) |> expand()
+
+      result = Guard.type_information_from_guards(guard)
+
+      assert result == %{
+               {:x, 0} => {:map, [{:foo, {:map, [{:bar, {:map, [{:baz, :atom}], []}}], []}}], []}
+             }
+    end
+
+    test "with operator" do
+      guard = quote(do: x.foo == 1) |> expand()
+
+      result = Guard.type_information_from_guards(guard)
+      assert result == %{{:x, 0} => {:map, [{:foo, {:integer, 1}}], []}}
+    end
+  end
 end