Implement type inferences from guard expressions

lib/elixir_sense/core/guard.ex

@@ -0,0 +1,138 @@
+defmodule ElixirSense.Core.Guard do
+  @moduledoc """
+  This module is responsible for infer type information from guard expressions
+  """
+
+  import ElixirSense.Core.State
+
+  alias ElixirSense.Core.MetadataBuilder
+
+  # 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
+  def 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 -> {:intersection, [v1, v2]} end)
+  end
+
+  def type_information_from_guards({:or, _, [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
+        {{:union, types_1}, {:union, types_2}} -> {:union, types_1 ++ types_2}
+        {{:union, types}, _} -> {:union, types ++ [v2]}
+        {_, {:union, types}} -> {:union, [v1 | types]}
+        _ -> {:union, [v1, v2]}
+      end
+    end)
+  end
+
+  def 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 ->
+          case guard_predicate_type(guard_predicate, params, state) do
+            {type, binding} ->
+              {var, _, nil} = binding
+              # If we found the predicate type, we can prematurely exit traversing the subtree
+              {[], [{var, type} | acc]}
+
+            nil ->
+              {node, acc}
+          end
+
+        node, acc ->
+          {node, acc}
+      end)
+
+    acc
+  end
+
+  defp guard_predicate_type(p, params, _)
+       when p in [:is_number, :is_float, :is_integer, :round, :trunc, :div, :rem, :abs],
+       do: {:number, hd(params)}
+
+  defp guard_predicate_type(p, params, _) when p in [:is_binary, :binary_part],
+    do: {:binary, hd(params)}
+
+  defp guard_predicate_type(p, params, _) when p in [:is_bitstring, :bit_size, :byte_size],
+    do: {:bitstring, hd(params)}
+
+  defp guard_predicate_type(p, params, _) when p in [:is_list, :length], do: {:list, hd(params)}
+
+  defp guard_predicate_type(p, params, _) when p in [:hd, :tl],
+    do: {{:list, :boolean}, hd(params)}
+
+  # when hd(x) == 1
+  # when tl(x) <= 2
+  defp guard_predicate_type(p, [{guard, _, guard_params}, rhs], _)
+       when p in [:==, :===, :>=, :>, :<=, :<] and guard in [:hd, :tl] do
+    rhs_type =
+      cond do
+        is_number(rhs) -> :number
+        is_binary(rhs) -> :binary
+        is_bitstring(rhs) -> :bitstring
+        is_atom(rhs) -> :atom
+        is_boolean(rhs) -> :boolean
+        true -> nil
+      end
+
+    rhs_type = if rhs_type, do: {:list, rhs_type}, else: :list
+
+    {rhs_type, hd(guard_params)}
+  end
+
+  defp guard_predicate_type(p, params, _) when p in [:is_tuple, :elem],
+    do: {:tuple, hd(params)}
+
+  # when tuple_size(x) == 1
+  # when tuple_size(x) == 2
+  defp guard_predicate_type(p, [{:tuple_size, _, guard_params}, size], _)
+       when p in [:==, :===] do
+    type =
+      if is_integer(size) do
+        {:tuple, size, []}
+      else
+        :tuple
+      end
+
+    {type, hd(guard_params)}
+  end
+
+  defp guard_predicate_type(:is_map, params, _), do: {{:map, [], nil}, hd(params)}
+  defp guard_predicate_type(:map_size, params, _), do: {{:map, [], nil}, hd(params)}
+
+  defp guard_predicate_type(:is_map_key, [var, key], state) do
+    type =
+      case MetadataBuilder.get_binding_type(state, key) do
+        {:atom, key} -> {:map, [{key, nil}], nil}
+        nil when is_binary(key) -> {:map, [{key, nil}], nil}
+        _ -> {:map, [], nil}
+      end
+
+    {type, var}
+  end
+
+  defp guard_predicate_type(:is_atom, params, _), do: {:atom, hd(params)}
+  defp guard_predicate_type(:is_boolean, params, _), do: {:boolean, hd(params)}
+
+  defp guard_predicate_type(:is_struct, [var, {:__aliases__, _, list}], state) do
+    type = {:struct, [], {:atom, expand_alias(state, list)}, nil}
+    {type, var}
+  end
+
+  defp guard_predicate_type(:is_struct, params, _), do: {{:struct, [], nil, nil}, hd(params)}
+  defp guard_predicate_type(_, _, _), do: nil
+end

lib/elixir_sense/core/metadata_builder.ex

@@ -13,6 +13,7 @@ defmodule ElixirSense.Core.MetadataBuilder do
   alias ElixirSense.Core.State
   alias ElixirSense.Core.State.VarInfo
   alias ElixirSense.Core.TypeInfo
+  alias ElixirSense.Core.Guard
 
   @scope_keywords [:for, :fn, :with]
   @block_keywords [:do, :else, :rescue, :catch, :after]
@@ -219,6 +220,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 +561,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 +1669,18 @@ defmodule ElixirSense.Core.MetadataBuilder do
     {ast, {vars, match_context}}
   end
 
+  def infer_type_from_guards(guard_ast, vars, state) do
+    type_info = Guard.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
+
   # struct or struct update
   def get_binding_type(
         state,

test/elixir_sense/core/metadata_builder_test.exs

@@ -1477,6 +1477,114 @@ 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, :number}} = 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, 1, []}} = 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, [], nil}} = var_with_guards("map_size(x) == 1")
+
+      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, [], nil, nil}} = 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: {:intersection, [:number, :boolean]}} =
+               var_with_guards("is_number(x) and x >= 1")
+
+      assert %VarInfo{
+               name: :x,
+               type: {:intersection, [{:map, [a: nil], nil}, {:map, [b: nil], nil}]}
+             } = var_with_guards("is_map_key(x, :a) and is_map_key(x, :b)")
+    end
+
+    test "or combination predicate guards can be merge into union type" do
+      assert %VarInfo{name: :x, type: {:union, [:number, :atom]}} =
+               var_with_guards("is_number(x) or is_atom(x)")
+
+      assert %VarInfo{name: :x, type: {:union, [:number, :atom, :binary]}} =
+               var_with_guards("is_number(x) or is_atom(x) or is_binary(x)")
+    end
+  end
+
   test "aliases" do
     state =
       """