Implement type inferences from guard expressions

lib/elixir_sense/core/metadata_builder.ex

@@ -219,6 +219,11 @@ defmodule ElixirSense.Core.MetadataBuilder do
       |> find_vars(params)
       |> merge_same_name_vars()
 
+    vars =
+      if options[:guards],
+        do: infer_type_from_guards(options[:guards], vars, state),
+        else: vars
+
     {position, end_position} = extract_range(meta)
 
     options = Keyword.put(options, :generated, state.generated)
@@ -555,7 +560,7 @@ defmodule ElixirSense.Core.MetadataBuilder do
        )
        when def_name in @defs do
     ast_without_params = {def_name, meta, [{name, add_no_call(meta2), []}, guards, body]}
-    pre_func(ast_without_params, state, meta, name, params)
+    pre_func(ast_without_params, state, meta, name, params, guards: guards)
   end
 
   defp pre(
@@ -1663,6 +1668,98 @@ defmodule ElixirSense.Core.MetadataBuilder do
     {ast, {vars, match_context}}
   end
 
+  def infer_type_from_guards(guard_ast, vars, state) do
+    type_info = type_information_from_guards(guard_ast, state)
+
+    Enum.reduce(type_info, vars, fn {var, type}, acc ->
+      index = Enum.find_index(acc, &(&1.name == var))
+
+      if index,
+        do: List.update_at(acc, index, &Map.put(&1, :type, type)),
+        else: acc
+    end)
+  end
+
+  # A guard expression can be in either these form:
+  #        :and                          :or
+  #      /     \            or          /   \              or      guard_expr
+  # guard_expr  guard_expr        guard_expr  guard_expr
+  #
+  # type information from :and subtrees are mergeable
+  # type information from :or subtrees are discarded
+  defp type_information_from_guards({:and, _, [guard_l, guard_r]}, state) do
+    left = type_information_from_guards(guard_l, state)
+    right = type_information_from_guards(guard_r, state)
+
+    Keyword.merge(left, right, fn _k, v1, v2 ->
+      case {v1, v2} do
+        # func my_func(x) when is_map_key(x, :a) and is_map_key(x, :b)
+        {{:map, fields1, _}, {:map, fields2, _}} ->
+          {:map, Enum.uniq_by(fields1 ++ fields2, &elem(&1, 0)), nil}
+
+        # In case we can't merge, just pick one
+        _ ->
+          v1
+      end
+    end)
+  end
+
+  defp type_information_from_guards({:or, _, [_guard_l, _guard_r]}, _state), do: []
+
+  defp type_information_from_guards(guard_ast, state) do
+    {_, acc} =
+      Macro.prewalk(guard_ast, [], fn
+        # Standalone variable: func my_func(x) when x
+        {var, _, nil} = node, acc ->
+          {node, [{var, :boolean} | acc]}
+
+        {guard_predicate, _, params} = node, acc ->
+          if type = guard_predicate_type(guard_predicate, params, state) do
+            [{var, _, nil} | _] = params
+            # If we found the predicate type, we can prematurely exit traversing the subtree
+            {[], [{var, type} | acc]}
+          else
+            {node, acc}
+          end
+
+        node, acc ->
+          {node, acc}
+      end)
+
+    acc
+  end
+
+  defp guard_predicate_type(p, _, _)
+       when p in [:is_number, :is_float, :is_integer, :round, :trunc, :div, :rem, :abs],
+       do: :number
+
+  defp guard_predicate_type(p, _, _) when p in [:is_binary, :binary_part], do: :binary
+
+  defp guard_predicate_type(p, _, _) when p in [:is_bitstring, :bit_size, :byte_size],
+    do: :bitstring
+
+  defp guard_predicate_type(p, _, _) when p in [:is_list, :hd, :tl, :length], do: :list
+  defp guard_predicate_type(p, _, _) when p in [:is_tuple, :tuple_size, :elem], do: :tuple
+  defp guard_predicate_type(:is_map, _, _), do: {:map, [], nil}
+
+  defp guard_predicate_type(:is_map_key, [_, key], state) do
+    case get_binding_type(state, key) do
+      {:atom, key} -> {:map, [{key, nil}], nil}
+      nil when is_binary(key) -> {:map, [{key, nil}], nil}
+      _ -> {:map, [], nil}
+    end
+  end
+
+  defp guard_predicate_type(:is_atom, _, _), do: :atom
+  defp guard_predicate_type(:is_boolean, _, _), do: :boolean
+
+  defp guard_predicate_type(:is_struct, [_, {:__aliases__, _, list}], state) do
+    {:struct, [], {:atom, expand_alias(state, list)}, nil}
+  end
+
+  defp guard_predicate_type(:is_struct, _, _), do: :struct
+  defp guard_predicate_type(_, _, _), do: nil
+
   # struct or struct update
   def get_binding_type(
         state,

test/elixir_sense/core/metadata_builder_test.exs

@@ -1477,6 +1477,106 @@ defmodule ElixirSense.Core.MetadataBuilderTest do
            ] = state |> get_line_vars(4)
   end
 
+  describe "infer vars type information from guards" do
+    defp var_with_guards(guard) do
+      """
+      defmodule MyModule do
+        def func(x) when #{guard} do
+          x
+        end
+      end
+      """
+      |> string_to_state()
+      |> get_line_vars(3)
+      |> hd()
+    end
+
+    test "number guards" do
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("is_number(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("is_float(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("is_integer(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("round(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("trunc(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("div(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("rem(x)")
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("abs(x)")
+    end
+
+    test "binary guards" do
+      assert %VarInfo{name: :x, type: :binary} = var_with_guards("is_binary(x)")
+      assert %VarInfo{name: :x, type: :binary} = var_with_guards(~s/binary_part(x, 0, 1) == "a"/)
+    end
+
+    test "bitstring guards" do
+      assert %VarInfo{name: :x, type: :bitstring} = var_with_guards("is_bitstring(x)")
+      assert %VarInfo{name: :x, type: :bitstring} = var_with_guards("bit_size(x) == 1")
+      assert %VarInfo{name: :x, type: :bitstring} = var_with_guards("byte_size(x) == 1")
+    end
+
+    test "list guards" do
+      assert %VarInfo{name: :x, type: :list} = var_with_guards("is_list(x)")
+      assert %VarInfo{name: :x, type: :list} = var_with_guards("hd(x) == 1")
+      assert %VarInfo{name: :x, type: :list} = var_with_guards("tl(x) == 1")
+      assert %VarInfo{name: :x, type: :list} = var_with_guards("length(x) == 1")
+    end
+
+    test "tuple guards" do
+      assert %VarInfo{name: :x, type: :tuple} = var_with_guards("is_tuple(x)")
+      assert %VarInfo{name: :x, type: :tuple} = var_with_guards("tuple_size(x) == 1")
+      assert %VarInfo{name: :x, type: :tuple} = var_with_guards("elem(x, 0) == 1")
+    end
+
+    test "atom guards" do
+      assert %VarInfo{name: :x, type: :atom} = var_with_guards("is_atom(x)")
+    end
+
+    test "boolean guards" do
+      assert %VarInfo{name: :x, type: :boolean} = var_with_guards("is_boolean(x)")
+    end
+
+    test "map guards" do
+      assert %VarInfo{name: :x, type: {:map, [], nil}} = var_with_guards("is_map(x)")
+
+      assert %VarInfo{name: :x, type: {:map, [a: nil], nil}} =
+               var_with_guards("is_map_key(x, :a)")
+
+      assert %VarInfo{name: :x, type: {:map, [{"a", nil}], nil}} =
+               var_with_guards(~s/is_map_key(x, "a")/)
+    end
+
+    test "struct guards" do
+      assert %VarInfo{name: :x, type: :struct} = var_with_guards("is_struct(x)")
+
+      assert %VarInfo{name: :x, type: {:struct, [], {:atom, URI}, nil}} =
+               var_with_guards("is_struct(x, URI)")
+
+      assert %VarInfo{name: :x, type: {:struct, [], {:atom, URI}, nil}} =
+               """
+               defmodule MyModule do
+                 alias URI, as: MyURI
+
+                 def func(x) when is_struct(x, MyURI) do
+                   x
+                 end
+               end
+               """
+               |> string_to_state()
+               |> get_line_vars(5)
+               |> hd()
+    end
+
+    test "and combination predicate guards can be merge" do
+      assert %VarInfo{name: :x, type: :number} = var_with_guards("is_number(x) and x >= 1")
+
+      assert %VarInfo{name: :x, type: {:map, [a: nil, b: nil], nil}} =
+               var_with_guards("is_map_key(x, :a) and is_map_key(x, :b)")
+    end
+
+    test "or combination predicate guards can not be used" do
+      assert %VarInfo{name: :x, type: nil} = var_with_guards("is_number(x) or is_atom(x)")
+    end
+  end
+
   test "aliases" do
     state =
       """