Type check size in binaries

Author: josevalim

lib/elixir/lib/module/types/expr.ex

@@ -27,7 +27,8 @@ defmodule Module.Types.Expr do
   @atom_true atom([true])
   @exception open_map(__struct__: atom(), __exception__: @atom_true)
 
-  defp of_expr(expr, expected_expr, stack, context) do
+  # of_expr/4 is public as it is called recursively from Of.binary
+  def of_expr(expr, expected_expr, stack, context) do
     with {:ok, actual, context} <- of_expr(expr, stack, context) do
       Of.intersect(actual, expected_expr, stack, context)
     end
@@ -81,7 +82,7 @@ defmodule Module.Types.Expr do
 
   # <<...>>>
   def of_expr({:<<>>, _meta, args}, stack, context) do
-    case Of.binary(args, :expr, stack, context, &of_expr/4) do
+    case Of.binary(args, :expr, stack, context) do
       {:ok, context} -> {:ok, binary(), context}
       # It is safe to discard errors from binary inside expressions
       {:error, context} -> {:ok, binary(), context}

lib/elixir/lib/module/types/of.ex

@@ -206,63 +206,94 @@ defmodule Module.Types.Of do
   In the stack, we add nodes such as <<expr>>, <<..., expr>>, etc,
   based on the position of the expression within the binary.
   """
-  def binary([], _kind, _stack, context, _of_fun) do
+  def binary([], _kind, _stack, context) do
     {:ok, context}
   end
 
-  def binary([head], kind, stack, context, of_fun) do
-    binary_segment(head, kind, [head], stack, context, of_fun)
+  def binary([head], kind, stack, context) do
+    binary_segment(head, kind, [head], stack, context)
   end
 
-  def binary([head | tail], kind, stack, context, of_fun) do
-    case binary_segment(head, kind, [head, @suffix], stack, context, of_fun) do
-      {:ok, context} -> binary_many(tail, kind, stack, context, of_fun)
-      {:error, reason} -> {:error, reason}
+  def binary([head | tail], kind, stack, context) do
+    case binary_segment(head, kind, [head, @suffix], stack, context) do
+      {:ok, context} -> binary_many(tail, kind, stack, context)
+      {:error, context} -> {:error, context}
     end
   end
 
-  defp binary_many([last], kind, stack, context, of_fun) do
-    binary_segment(last, kind, [@prefix, last], stack, context, of_fun)
+  defp binary_many([last], kind, stack, context) do
+    binary_segment(last, kind, [@prefix, last], stack, context)
   end
 
-  defp binary_many([head | tail], kind, stack, context, of_fun) do
-    case binary_segment(head, kind, [@prefix, head, @suffix], stack, context, of_fun) do
-      {:ok, context} -> binary_many(tail, kind, stack, context, of_fun)
-      {:error, reason} -> {:error, reason}
+  defp binary_many([head | tail], kind, stack, context) do
+    case binary_segment(head, kind, [@prefix, head, @suffix], stack, context) do
+      {:ok, context} -> binary_many(tail, kind, stack, context)
+      {:error, context} -> {:error, context}
     end
   end
 
   # If the segment is a literal, the compiler has already checked its validity,
   # so we just skip it.
-  defp binary_segment({:"::", _meta, [left, _right]}, _kind, _args, _stack, context, _of_fun)
+  defp binary_segment({:"::", _meta, [left, _right]}, _kind, _args, _stack, context)
        when is_binary(left) or is_number(left) do
     {:ok, context}
   end
 
-  defp binary_segment({:"::", meta, [left, right]}, kind, args, stack, context, of_fun) do
-    expected_type = specifier_info(kind, right)
+  defp binary_segment({:"::", meta, [left, right]}, kind, args, stack, context) do
+    expected_type = specifier_type(kind, right)
     expr = {:<<>>, meta, args}
 
-    with {:ok, _type, context} <- of_fun.(left, {expected_type, expr}, stack, context) do
-      {:ok, context}
+    result =
+      case kind do
+        :pattern -> Module.Types.Pattern.of_pattern(left, {expected_type, expr}, stack, context)
+        :guard -> Module.Types.Pattern.of_guard(left, {expected_type, expr}, stack, context)
+        :expr -> Module.Types.Expr.of_expr(left, {expected_type, expr}, stack, context)
+      end
+
+    with {:ok, _type, context} <- result do
+      {:ok, specifier_size(kind, right, stack, context)}
+    end
+  end
+
+  defp specifier_type(kind, {:-, _, [left, _right]}), do: specifier_type(kind, left)
+  defp specifier_type(:pattern, {:utf8, _, _}), do: @integer
+  defp specifier_type(:pattern, {:utf16, _, _}), do: @integer
+  defp specifier_type(:pattern, {:utf32, _, _}), do: @integer
+  defp specifier_type(:pattern, {:float, _, _}), do: @float
+  defp specifier_type(_kind, {:float, _, _}), do: @integer_or_float
+  defp specifier_type(_kind, {:utf8, _, _}), do: @integer_or_binary
+  defp specifier_type(_kind, {:utf16, _, _}), do: @integer_or_binary
+  defp specifier_type(_kind, {:utf32, _, _}), do: @integer_or_binary
+  defp specifier_type(_kind, {:integer, _, _}), do: @integer
+  defp specifier_type(_kind, {:bits, _, _}), do: @binary
+  defp specifier_type(_kind, {:bitstring, _, _}), do: @binary
+  defp specifier_type(_kind, {:bytes, _, _}), do: @binary
+  defp specifier_type(_kind, {:binary, _, _}), do: @binary
+  defp specifier_type(_kind, _specifier), do: @integer
+
+  defp specifier_size(kind, {:-, _, [left, right]}, stack, context) do
+    specifier_size(kind, right, stack, specifier_size(kind, left, stack, context))
+  end
+
+  defp specifier_size(:expr, {:size, _, [arg]} = expr, stack, context)
+       when not is_integer(arg) do
+    case Module.Types.Expr.of_expr(arg, {integer(), expr}, stack, context) do
+      {:ok, _, context} -> context
+      {:error, context} -> context
     end
   end
 
-  defp specifier_info(kind, {:-, _, [left, _right]}), do: specifier_info(kind, left)
-  defp specifier_info(:expr, {:float, _, _}), do: @integer_or_float
-  defp specifier_info(:expr, {:utf8, _, _}), do: @integer_or_binary
-  defp specifier_info(:expr, {:utf16, _, _}), do: @integer_or_binary
-  defp specifier_info(:expr, {:utf32, _, _}), do: @integer_or_binary
-  defp specifier_info(:pattern, {:utf8, _, _}), do: @integer
-  defp specifier_info(:pattern, {:utf16, _, _}), do: @integer
-  defp specifier_info(:pattern, {:utf32, _, _}), do: @integer
-  defp specifier_info(:pattern, {:float, _, _}), do: @float
-  defp specifier_info(_kind, {:integer, _, _}), do: @integer
-  defp specifier_info(_kind, {:bits, _, _}), do: @binary
-  defp specifier_info(_kind, {:bitstring, _, _}), do: @binary
-  defp specifier_info(_kind, {:bytes, _, _}), do: @binary
-  defp specifier_info(_kind, {:binary, _, _}), do: @binary
-  defp specifier_info(_kind, _specifier), do: @integer
+  defp specifier_size(_pattern_or_guard, {:size, _, [arg]} = expr, stack, context)
+       when not is_integer(arg) do
+    case Module.Types.Pattern.of_guard(arg, {integer(), expr}, stack, context) do
+      {:ok, _, context} -> context
+      {:error, context} -> context
+    end
+  end
+
+  defp specifier_size(_kind, _expr, _stack, context) do
+    context
+  end
 
   ## Apply
 

lib/elixir/lib/module/types/pattern.ex

@@ -30,21 +30,22 @@ defmodule Module.Types.Pattern do
   end
 
   ## Patterns
+  # of_pattern is public as it is called recursively from Of.binary
 
-  defp of_pattern(expr, stack, context) do
-    # TODO: Remove the hardcoding of dynamic
-    # TODO: Possibly remove this function
+  # TODO: Remove the hardcoding of dynamic
+  # TODO: Remove this function
+  def of_pattern(expr, stack, context) do
     of_pattern(expr, {dynamic(), expr}, stack, context)
   end
 
   # ^var
-  defp of_pattern({:^, _meta, [var]}, expected_expr, stack, context) do
+  def of_pattern({:^, _meta, [var]}, expected_expr, stack, context) do
     Of.intersect(Of.var(var, context), expected_expr, stack, context)
   end
 
   # left = right
   # TODO: Track variables and handle nesting
-  defp of_pattern({:=, _meta, [left_expr, right_expr]}, {expected, expr}, stack, context) do
+  def of_pattern({:=, _meta, [left_expr, right_expr]}, {expected, expr}, stack, context) do
     case {is_var(left_expr), is_var(right_expr)} do
       {true, false} ->
         with {:ok, type, context} <- of_pattern(right_expr, {expected, expr}, stack, context) do
@@ -64,13 +65,13 @@ defmodule Module.Types.Pattern do
   end
 
   # %var{...} and %^var{...}
-  defp of_pattern(
-         {:%, _meta, [struct_var, {:%{}, _meta2, args}]} = expr,
-         expected_expr,
-         stack,
-         context
-       )
-       when not is_atom(struct_var) do
+  def of_pattern(
+        {:%, _meta, [struct_var, {:%{}, _meta2, args}]} = expr,
+        expected_expr,
+        stack,
+        context
+      )
+      when not is_atom(struct_var) do
     with {:ok, struct_type, context} <-
            of_pattern(struct_var, {atom(), expr}, %{stack | refine: false}, context),
          {:ok, map_type, context} <-
@@ -83,35 +84,35 @@ defmodule Module.Types.Pattern do
   end
 
   # %Struct{...}
-  defp of_pattern({:%, _meta, [module, {:%{}, _, args}]} = expr, expected_expr, stack, context)
-       when is_atom(module) do
+  def of_pattern({:%, _meta, [module, {:%{}, _, args}]} = expr, expected_expr, stack, context)
+      when is_atom(module) do
     with {:ok, actual, context} <-
            Of.struct(expr, module, args, :merge_defaults, stack, context, &of_pattern/3) do
       Of.intersect(actual, expected_expr, stack, context)
     end
   end
 
   # %{...}
-  defp of_pattern({:%{}, _meta, args}, expected_expr, stack, context) do
+  def of_pattern({:%{}, _meta, args}, expected_expr, stack, context) do
     of_open_map(args, [], expected_expr, stack, context)
   end
 
   # <<...>>>
-  defp of_pattern({:<<>>, _meta, args}, _expected_expr, stack, context) do
-    case Of.binary(args, :pattern, stack, context, &of_pattern/4) do
+  def of_pattern({:<<>>, _meta, args}, _expected_expr, stack, context) do
+    case Of.binary(args, :pattern, stack, context) do
       {:ok, context} -> {:ok, binary(), context}
       {:error, context} -> {:error, context}
     end
   end
 
   # _
-  defp of_pattern({:_, _meta, _var_context}, {expected, _expr}, _stack, context) do
+  def of_pattern({:_, _meta, _var_context}, {expected, _expr}, _stack, context) do
     {:ok, expected, context}
   end
 
   # var
-  defp of_pattern({name, meta, ctx} = var, {expected, expr}, stack, context)
-       when is_atom(name) and is_atom(ctx) do
+  def of_pattern({name, meta, ctx} = var, {expected, expr}, stack, context)
+      when is_atom(name) and is_atom(ctx) do
     case stack do
       %{refine: true} ->
         Of.refine_var(var, expected, expr, stack, context)
@@ -129,7 +130,7 @@ defmodule Module.Types.Pattern do
     end
   end
 
-  defp of_pattern(expr, expected_expr, stack, context) do
+  def of_pattern(expr, expected_expr, stack, context) do
     of_shared(expr, expected_expr, stack, context, &of_pattern/4)
   end
 
@@ -151,11 +152,11 @@ defmodule Module.Types.Pattern do
     end
   end
 
-  @doc """
-  Refines the type variables in the typing context using type check guards
-  such as `is_integer/1`.
-  """
+  ## Guards
+  # of_guard is public as it is called recursively from Of.binary
 
+  # TODO: Remove the hardcoding of dynamic
+  # TODO: Remove this function
   def of_guard(expr, stack, context) do
     of_guard(expr, {dynamic(), expr}, stack, context)
   end
@@ -173,7 +174,7 @@ defmodule Module.Types.Pattern do
 
   # <<>>
   def of_guard({:<<>>, _meta, args}, _expected_expr, stack, context) do
-    case Of.binary(args, :expr, stack, context, &of_guard/4) do
+    case Of.binary(args, :guard, stack, context) do
       {:ok, context} -> {:ok, binary(), context}
       # It is safe to discard errors from binary inside expressions
       {:error, context} -> {:ok, binary(), context}
@@ -197,9 +198,15 @@ defmodule Module.Types.Pattern do
     end
   end
 
+  # ^var
+  # Happens from inside size(^...) and map keys
+  def of_guard({:^, _meta, [var]}, expected_expr, stack, context) do
+    Of.intersect(Of.var(var, context), expected_expr, stack, context)
+  end
+
   # var
-  def of_guard(var, _expected_expr, _stack, context) when is_var(var) do
-    {:ok, Of.var(var, context), context}
+  def of_guard(var, expected_expr, stack, context) when is_var(var) do
+    Of.intersect(Of.var(var, context), expected_expr, stack, context)
   end
 
   def of_guard(expr, expected_expr, stack, context) do

lib/elixir/test/elixir/module/types/expr_test.exs

@@ -199,6 +199,34 @@ defmodule Module.Types.ExprTest do
                 #{hints(:inferred_bitstring_spec)}
                 """}
     end
+
+    test "size ok" do
+      assert typecheck!([<<x, y>>, z], <<z::size(x - y)>>) == binary()
+    end
+
+    test "size error" do
+      assert typewarn!([<<x::binary>>, y], <<y::size(x)>>) ==
+               {binary(),
+                ~l"""
+                incompatible types in expression:
+
+                    size(x)
+
+                expected type:
+
+                    integer()
+
+                but got type:
+
+                    binary()
+
+                where "x" was given the type:
+
+                    # type: binary()
+                    # from: types_test.ex:LINE-2
+                    <<x::binary>>
+                """}
+    end
   end
 
   describe "tuples" do
@@ -544,7 +572,7 @@ defmodule Module.Types.ExprTest do
                 where "x" was given the type:
 
                     # type: integer()
-                    # from: types_test.ex:533
+                    # from: types_test.ex:LINE-2
                     <<x>>
 
                 #{hints(:inferred_bitstring_spec)}

lib/elixir/test/elixir/module/types/pattern_test.exs

@@ -131,5 +131,33 @@ defmodule Module.Types.PatternTest do
              #{hints(:inferred_bitstring_spec)}
              """
     end
+
+    test "size ok" do
+      assert typecheck!([<<x, y, _::size(x - y)>>], :ok) == atom([:ok])
+    end
+
+    test "size error" do
+      assert typewarn!([<<x::float, _::size(x)>>], :ok) ==
+               {atom([:ok]),
+                ~l"""
+                incompatible types in expression:
+
+                    size(x)
+
+                expected type:
+
+                    integer()
+
+                but got type:
+
+                    float()
+
+                where "x" was given the type:
+
+                    # type: float()
+                    # from: types_test.ex:LINE-2
+                    <<x::float, ...>>
+                """}
+    end
   end
 end