diff --git a/examples/simple_app/lib/simple_app.ex b/examples/simple_app/lib/simple_app.ex
index 6c3603fc..a196bf11 100644
--- a/examples/simple_app/lib/simple_app.ex
+++ b/examples/simple_app/lib/simple_app.ex
@@ -35,5 +35,4 @@ defmodule SimpleApp do
     l = fn -> :ok end
     l.()
   end
-
 end
diff --git a/examples/simple_app/lib/simple_app/box.ex b/examples/simple_app/lib/simple_app/box.ex
index a24a5831..728443e3 100644
--- a/examples/simple_app/lib/simple_app/box.ex
+++ b/examples/simple_app/lib/simple_app/box.ex
@@ -1,5 +1,4 @@
 defmodule SimpleApp.Box do
-
   @spec get_l() :: {:ok, integer()} | :error
   def get_l, do: {:ok, 5}
 
diff --git a/lib/gradient.ex b/lib/gradient.ex
index 367dc619..13008752 100644
--- a/lib/gradient.ex
+++ b/lib/gradient.ex
@@ -9,7 +9,7 @@ defmodule Gradient do
 
   alias Gradient.ElixirFileUtils
   alias Gradient.ElixirFmt
-  alias Gradient.SpecifyErlAst
+  alias Gradient.AstSpecifier
 
   require Logger
 
@@ -23,7 +23,7 @@ defmodule Gradient do
       forms =
         forms
         |> put_code_path(opts)
-        |> SpecifyErlAst.specify()
+        |> AstSpecifier.specify()
 
       case :gradualizer.type_check_forms(forms, opts) do
         [] ->
diff --git a/lib/gradient/specify_erl_ast.ex b/lib/gradient/ast_specifier.ex
similarity index 74%
rename from lib/gradient/specify_erl_ast.ex
rename to lib/gradient/ast_specifier.ex
index 90628520..f79e648a 100644
--- a/lib/gradient/specify_erl_ast.ex
+++ b/lib/gradient/ast_specifier.ex
@@ -1,13 +1,15 @@
-defmodule Gradient.SpecifyErlAst do
+defmodule Gradient.AstSpecifier do
   @moduledoc """
-  Module adds missing line information to the Erlang abstract code produced 
-  from Elixir AST.
+  Module adds missing location information to the Erlang abstract code produced
+  from Elixir AST. Moreover it can be used to catch some ast pattern and replace 
+  it to forms that cannot be produced from Elixir directly.
 
   FIXME Optimize tokens searching. Find out why some tokens are dropped 
 
   NOTE Mapper implements:
   - function [x]
   - fun [x] 
+  - fun @spec [x]
   - clause [x] 
   - case [x]
   - block [X] 
@@ -42,22 +44,23 @@ defmodule Gradient.SpecifyErlAst do
   - guards [X]
   """
 
-  import Gradient.Utils
+  import Gradient.Tokens
 
   require Logger
 
-  @type token :: tuple()
-  @type tokens :: [tuple()]
-  @type form ::
-          :erl_parse.abstract_clause()
-          | :erl_parse.abstract_expr()
-          | :erl_parse.abstract_form()
-          | :erl_parse.abstract_type()
-  @type forms :: [form()]
-  @type options :: keyword()
+  alias Gradient.Types
 
-  @doc """
+  @type token :: Types.token()
+  @type tokens :: Types.tokens()
+  @type form :: Types.form()
+  @type forms :: Types.forms()
+  @type options :: Types.options()
+
+  # Api
 
+  @doc """
+  Read and tokenize code file. Than run mappers on the given AST (with obtained tokens) 
+  to specify missing locations or replace some parts that match pattern.
   """
   @spec specify(nonempty_list(:erl_parse.abstract_form())) :: [:erl_parse.abstract_form()]
   def specify(forms) do
@@ -65,7 +68,7 @@ defmodule Gradient.SpecifyErlAst do
          path <- to_string(path),
          {:ok, code} <- File.read(path),
          {:ok, tokens} <- :elixir.string_to_tokens(String.to_charlist(code), line, line, path, []) do
-      add_missing_loc_literals(forms, tokens)
+      run_mappers(forms, tokens)
     else
       error ->
         IO.puts("Error occured when specifying forms : #{inspect(error)}")
@@ -74,14 +77,14 @@ defmodule Gradient.SpecifyErlAst do
   end
 
   @doc """
-  Function takes forms and traverse them to add missing location for literals. 
-  Firstly the parent location is set, then it is matched 
-  with tokens to precise the literal line.
+  Function takes forms and traverse them in order to specify location or modify
+  forms matching the pattern. The tokens are required to obtain the missing location 
+  as precise as possible.
   """
-  @spec add_missing_loc_literals([:erl_parse.abstract_form()], tokens()) :: [
+  @spec run_mappers([:erl_parse.abstract_form()], tokens()) :: [
           :erl_parse.abstract_form()
         ]
-  def add_missing_loc_literals(forms, tokens) do
+  def run_mappers(forms, tokens) do
     opts = [end_line: -1]
 
     {forms, _} =
@@ -92,6 +95,8 @@ defmodule Gradient.SpecifyErlAst do
     forms
   end
 
+  # Mappers
+
   @doc """
   Map over the forms using mapper and attach a context i.e. end line. 
   """
@@ -106,7 +111,7 @@ defmodule Gradient.SpecifyErlAst do
   end
 
   @doc """
-    Fold over the forms using mapper and attach a context i.e. end line.
+  Fold over the forms using mapper and attach a context i.e. end line.
   """
   @spec context_mapper_fold(forms(), tokens(), options()) :: {forms(), tokens()}
   def context_mapper_fold(forms, tokens, opts, mapper \\ &mapper/3)
@@ -119,50 +124,26 @@ defmodule Gradient.SpecifyErlAst do
     {[form | forms], res_tokens}
   end
 
-  def set_form_end_line(opts, form, forms) do
-    case Enum.find(forms, fn f ->
-           anno = elem(f, 1)
-
-           # Maybe should try to go deeper when generated and try to obtain 
-           # the line from the first child. It should work for sure for clauses, 
-           # but it has to be in the right order (e.g. if clauses are reversed)
-           :erl_anno.line(anno) > 0 and not :erl_anno.generated(anno)
-         end) do
-      nil ->
-        opts
-
-      next_form ->
-        current_line = :erl_anno.line(elem(form, 1))
-        next_line = :erl_anno.line(elem(next_form, 1))
+  @doc """
+  The main mapper function traverses AST and specifies missing locations
+  or replaces parts that match the pattern. 
+  """
+  @spec mapper(form(), [token()], options()) :: {form(), [token()]}
+  def mapper(form, tokens, opts)
 
-        if current_line == next_line do
-          Keyword.put(opts, :end_line, next_line + 1)
-        else
-          Keyword.put(opts, :end_line, next_line)
-        end
-    end
-  end
+  def mapper({:attribute, anno, :spec, {name_arity, specs}}, tokens, opts) do
+    new_specs = context_mapper_map(specs, [], opts, &spec_mapper/3)
 
-  def prepare_forms_order(forms) do
-    forms
-    |> Enum.sort(fn l, r -> elem(l, 0) == elem(r, 0) and elem(l, 1) > elem(r, 1) end)
-    |> Enum.reverse()
-  end
-
-  @spec pass_tokens(any(), tokens()) :: {any(), tokens()}
-  defp pass_tokens(form, tokens) do
-    {form, tokens}
+    {:attribute, anno, :spec, {name_arity, new_specs}}
+    |> pass_tokens(tokens)
   end
 
-  @spec mapper(form(), [token()], options()) :: {form(), [token()]}
-  defp mapper(form, tokens, opts)
-
-  defp mapper({:function, _line, :__info__, _arity, _children} = form, tokens, _opts) do
+  def mapper({:function, _line, :__info__, _arity, _children} = form, tokens, _opts) do
     # skip analysis for __info__ functions
     pass_tokens(form, tokens)
   end
 
-  defp mapper({:function, anno, name, arity, clauses}, tokens, opts) do
+  def mapper({:function, anno, name, arity, clauses}, tokens, opts) do
     # anno has line
     {clauses, tokens} = context_mapper_fold(clauses, tokens, opts)
 
@@ -170,7 +151,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:fun, anno, {:clauses, clauses}}, tokens, opts) do
+  def mapper({:fun, anno, {:clauses, clauses}}, tokens, opts) do
     # anno has line
     {clauses, tokens} = context_mapper_fold(clauses, tokens, opts)
 
@@ -178,7 +159,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:case, anno, condition, clauses}, tokens, opts) do
+  def mapper({:case, anno, condition, clauses}, tokens, opts) do
     # anno has line
     # NOTE In Elixir `if`, `case` and `cond` statements are represented 
     # as a `case` in abstract code.
@@ -205,7 +186,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:clause, anno, args, guards, children}, tokens, opts) do
+  def mapper({:clause, anno, args, guards, children}, tokens, opts) do
     # anno has line
     # FIXME Handle generated clauses. Right now the literals inherit lines 
     # from the parents without checking them with tokens 
@@ -237,7 +218,7 @@ defmodule Gradient.SpecifyErlAst do
     end
   end
 
-  defp mapper({:block, anno, body}, tokens, opts) do
+  def mapper({:block, anno, body}, tokens, opts) do
     # TODO check if anno has line
     {:ok, _line, anno, opts, _} = get_line(anno, opts)
 
@@ -247,7 +228,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:match, anno, left, right}, tokens, opts) do
+  def mapper({:match, anno, left, right}, tokens, opts) do
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
     {left, tokens} = mapper(left, tokens, opts)
@@ -257,35 +238,35 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:map, anno, pairs}, tokens, opts) do
+  def mapper({:map, anno, pairs}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
-    {pairs, tokens} = context_mapper_fold(pairs, tokens, opts, &map_element/3)
+    {pairs, tokens} = context_mapper_fold(pairs, tokens, opts, &map_element_mapper/3)
 
     {:map, anno, pairs}
     |> pass_tokens(tokens)
   end
 
   # update map pattern
-  defp mapper({:map, anno, map, pairs}, tokens, opts) do
+  def mapper({:map, anno, map, pairs}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
     {map, tokens} = mapper(map, tokens, opts)
-    {pairs, tokens} = context_mapper_fold(pairs, tokens, opts, &map_element/3)
+    {pairs, tokens} = context_mapper_fold(pairs, tokens, opts, &map_element_mapper/3)
 
     {:map, anno, map, pairs}
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:cons, anno, value, more} = cons, tokens, opts) do
+  def mapper({:cons, anno, value, more} = cons, tokens, opts) do
     # anno could be 0
     {:ok, line, anno, opts, _} = get_line(anno, opts)
 
     tokens = drop_tokens_to_line(tokens, line)
 
-    case get_list_from_tokens(tokens, opts) do
+    case get_list(tokens, opts) do
       {:list, tokens} ->
         cons_mapper(cons, tokens, opts)
 
@@ -303,13 +284,13 @@ defmodule Gradient.SpecifyErlAst do
     end
   end
 
-  defp mapper({:tuple, anno, elements}, tokens, opts) do
+  def mapper({:tuple, anno, elements}, tokens, opts) do
     # anno could be 0
     {:ok, line, anno, opts, has_line?} = get_line(anno, opts)
 
     tokens
     |> drop_tokens_to_line(line)
-    |> get_tuple_from_tokens(opts)
+    |> get_tuple(opts)
     |> case do
       {:tuple, tokens} ->
         {anno, opts} = update_line_from_tokens(tokens, anno, opts, has_line?)
@@ -327,7 +308,7 @@ defmodule Gradient.SpecifyErlAst do
     end
   end
 
-  defp mapper({:receive, anno, clauses}, tokens, opts) do
+  def mapper({:receive, anno, clauses}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
@@ -338,7 +319,7 @@ defmodule Gradient.SpecifyErlAst do
   end
 
   # receive with timeout
-  defp mapper({:receive, anno, clauses, after_val, after_block}, tokens, opts) do
+  def mapper({:receive, anno, clauses, after_val, after_block}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
@@ -350,7 +331,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:try, anno, body, else_block, catchers, after_block}, tokens, opts) do
+  def mapper({:try, anno, body, else_block, catchers, after_block}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
@@ -366,19 +347,19 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper(
-         {:call, anno, {:atom, _, name_atom} = name,
-          [expr, {:bin, _, [{:bin_element, _, {:string, _, _} = val, :default, :default}]}]},
-         tokens,
-         _opts
-       )
-       when name_atom in [:"::", :":::"] do
+  def mapper(
+        {:call, anno, {:atom, _, name_atom} = name,
+         [expr, {:bin, _, [{:bin_element, _, {:string, _, _} = val, :default, :default}]}]},
+        tokens,
+        _opts
+      )
+      when name_atom in [:"::", :":::"] do
     # unwrap string from binary for correct type annotation matching
     {:call, anno, name, [expr, val]}
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:call, anno, name, args}, tokens, opts) do
+  def mapper({:call, anno, name, args}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
@@ -390,7 +371,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:op, anno, op, left, right}, tokens, opts) do
+  def mapper({:op, anno, op, left, right}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
@@ -401,7 +382,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:op, anno, op, right}, tokens, opts) do
+  def mapper({:op, anno, op, right}, tokens, opts) do
     # anno has correct line
     {:ok, _, anno, opts, _} = get_line(anno, opts)
 
@@ -411,7 +392,7 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  defp mapper({:bin, anno, elements}, tokens, opts) do
+  def mapper({:bin, anno, elements}, tokens, opts) do
     # anno could be 0
     {:ok, line, anno, opts, _} = get_line(anno, opts)
 
@@ -423,16 +404,16 @@ defmodule Gradient.SpecifyErlAst do
 
       _ ->
         {bin_tokens, other_tokens} = cut_tokens_to_bin(tokens, line)
-        bin_tokens = flat_tokens(bin_tokens)
-        {elements, _} = context_mapper_fold(elements, bin_tokens, opts, &bin_element/3)
+        bin_tokens = flatten_tokens(bin_tokens)
+        {elements, _} = context_mapper_fold(elements, bin_tokens, opts, &bin_element_mapper/3)
 
         {:bin, anno, elements}
         |> pass_tokens(other_tokens)
     end
   end
 
-  defp mapper({type, 0, value}, tokens, opts)
-       when type in [:atom, :char, :float, :integer, :string, :bin] do
+  def mapper({type, 0, value}, tokens, opts)
+      when type in [:atom, :char, :float, :integer, :string, :bin] do
     # TODO check what happend for :string
     {:ok, line} = Keyword.fetch(opts, :line)
 
@@ -440,31 +421,86 @@ defmodule Gradient.SpecifyErlAst do
     |> specify_line(tokens, opts)
   end
 
-  defp mapper(skip, tokens, _opts)
-       when elem(skip, 0) in [
-              :fun,
-              :attribute,
-              :var,
-              nil,
-              :atom,
-              :char,
-              :float,
-              :integer,
-              :string,
-              :bin
-            ] do
+  def mapper(skip, tokens, _opts)
+      when elem(skip, 0) in [
+             :fun,
+             :attribute,
+             :var,
+             nil,
+             :atom,
+             :char,
+             :float,
+             :integer,
+             :string,
+             :bin
+           ] do
     # NOTE fun - I skipped here checking &name/arity or &module.name/arity
     # skip forms that don't need analysis and do not display warning
     pass_tokens(skip, tokens)
   end
 
-  defp mapper(form, tokens, _opts) do
+  def mapper(form, tokens, _opts) do
     Logger.warn("Not found mapper for #{inspect(form)}")
     pass_tokens(form, tokens)
   end
 
   @doc """
-  Adds missing line to the module literal
+  Adds missing location to the function specification.
+  """
+  @spec spec_mapper(form(), tokens(), options()) :: {form(), tokens()}
+  def spec_mapper({:type, anno, :tuple, :any}, tokens, _opts) do
+    {:type, anno, :tuple, :any}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:type, anno, :map, :any}, tokens, _opts) do
+    {:type, anno, :map, :any}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:type, anno, :any}, tokens, _opts) do
+    {:type, anno, :any}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:type, anno, type_name, args}, tokens, opts) do
+    {:ok, _line, anno, opts, _} = get_line(anno, opts)
+    new_args = context_mapper_map(args, tokens, opts, &spec_mapper/3)
+
+    {:type, anno, type_name, new_args}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:remote_type, anno, [mod, type, args]}, tokens, opts) do
+    {:ok, _line, anno, opts, _} = get_line(anno, opts)
+    {new_mod, _} = spec_mapper(mod, tokens, opts)
+    {new_type, _} = spec_mapper(type, tokens, opts)
+    new_args = context_mapper_map(args, tokens, opts, &spec_mapper/3)
+
+    {:remote_type, anno, [new_mod, new_type, new_args]}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:user_type, anno, name, args}, tokens, opts) do
+    new_args = context_mapper_map(args, tokens, opts, &spec_mapper/3)
+
+    {:user_type, anno, name, new_args}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:ann_type, anno, attrs}, tokens, opts) do
+    new_attrs = context_mapper_map(attrs, tokens, opts, &spec_mapper/3)
+
+    {:ann_type, anno, new_attrs}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper(type, tokens, opts) do
+    mapper(type, tokens, opts)
+  end
+
+  @doc """
+  Adds missing location to the module literal
   """
   def remote_mapper({:remote, line, {:atom, 0, mod}, fun}) do
     {:remote, line, {:atom, line, mod}, fun}
@@ -490,7 +526,11 @@ defmodule Gradient.SpecifyErlAst do
     end)
   end
 
-  def map_element({field, anno, key, value}, tokens, opts)
+  @doc """
+  Run mapper on map value and key.
+  """
+  @spec map_element_mapper(tuple(), tokens(), options()) :: {tuple(), tokens()}
+  def map_element_mapper({field, anno, key, value}, tokens, opts)
       when field in [:map_field_assoc, :map_field_exact] do
     line = :erl_anno.line(anno)
     opts = Keyword.put(opts, :line, line)
@@ -502,7 +542,11 @@ defmodule Gradient.SpecifyErlAst do
     |> pass_tokens(tokens)
   end
 
-  def bin_element({:bin_element, anno, value, size, tsl}, tokens, opts) do
+  @doc """
+  Run mapper on bin element value.
+  """
+  @spec bin_element_mapper(tuple(), tokens(), options()) :: {tuple(), tokens()}
+  def bin_element_mapper({:bin_element, anno, value, size, tsl}, tokens, opts) do
     {:ok, _line, anno, opts, _} = get_line(anno, opts)
 
     {value, tokens} = mapper(value, tokens, opts)
@@ -515,7 +559,6 @@ defmodule Gradient.SpecifyErlAst do
   Iterate over the list in abstract code format and runs mapper on each element 
   """
   @spec cons_mapper(form(), [token()], options()) :: {form(), tokens()}
-
   def cons_mapper({:cons, anno, value, tail}, tokens, opts) do
     {:ok, _, anno, opts, has_line?} = get_line(anno, opts)
 
@@ -532,77 +575,12 @@ defmodule Gradient.SpecifyErlAst do
   def cons_mapper(other, tokens, opts), do: mapper(other, tokens, opts)
 
   @doc """
-  Drop tokens to the first conditional occurance. Returns type of the encountered conditional and following tokens.
+  Update form anno with location taken from the corresponding token, if found.
+  Otherwise return form unchanged.
   """
-  @spec get_conditional([token()], integer(), options()) ::
-          {:case, [token()]}
-          | {:cond, [token()]}
-          | {:unless, [token()]}
-          | {:if, [token()]}
-          | {:with, [token()]}
-          | :undefined
-  def get_conditional(tokens, line, opts) do
-    conditionals = [:if, :unless, :cond, :case, :with]
-    {:ok, limit_line} = Keyword.fetch(opts, :end_line)
-
-    drop_tokens_while(tokens, limit_line, fn
-      {:do_identifier, _, c} -> c not in conditionals
-      {:paren_identifier, _, c} -> c not in conditionals
-      {:identifier, _, c} -> c not in conditionals
-      _ -> true
-    end)
-    |> case do
-      [token | _] = tokens when elem(elem(token, 1), 0) == line -> {elem(token, 2), tokens}
-      _ -> :undefined
-    end
-  end
-
-  @spec get_list_from_tokens([token()], options()) ::
-          {:list, [token()]} | {:keyword, [token()]} | {:charlist, [token()]} | :undefined
-  def get_list_from_tokens(tokens, opts) do
-    tokens = flat_tokens(tokens)
-    {:ok, limit_line} = Keyword.fetch(opts, :end_line)
-
-    res =
-      drop_tokens_while(tokens, limit_line, fn
-        {:"[", _} -> false
-        {:list_string, _, _} -> false
-        {:kw_identifier, _, id} when id not in [:do] -> false
-        _ -> true
-      end)
-
-    case res do
-      [{:"[", _} | _] = list -> {:list, list}
-      [{:list_string, _, _} | _] = list -> {:charlist, list}
-      [{:kw_identifier, _, _} | _] = list -> {:keyword, list}
-      _ -> :undefined
-    end
-  end
-
-  @spec get_tuple_from_tokens(tokens, options()) ::
-          {:tuple, tokens()} | :undefined
-  def get_tuple_from_tokens(tokens, opts) do
-    {:ok, limit_line} = Keyword.fetch(opts, :end_line)
-
-    res =
-      drop_tokens_while(tokens, limit_line, fn
-        {:"{", _} -> false
-        {:kw_identifier, _, _} -> false
-        _ -> true
-      end)
-
-    case res do
-      [{:"{", _} | _] = tuple -> {:tuple, tuple}
-      [{:kw_identifier, _, _} | _] = tuple -> {:tuple, tuple}
-      _ -> :undefined
-    end
-  end
-
   @spec specify_line(form(), [token()], options()) :: {form(), [token()]}
-  # def specify_line(form, []), do: raise("ehh -- #{inspect form}")
   def specify_line(form, tokens, opts) do
     if not :erl_anno.generated(elem(form, 1)) do
-      # Logger.debug("#{inspect(form)} --- #{inspect(tokens, limit: :infinity)}")
       {:ok, end_line} = Keyword.fetch(opts, :end_line)
 
       res = drop_tokens_while(tokens, end_line, &(!match_token_to_form(&1, form)))
@@ -612,7 +590,6 @@ defmodule Gradient.SpecifyErlAst do
           {take_loc_from_token(token, form), tokens}
 
         [] ->
-          # Logger.info("Not found - #{inspect(form)}")
           {form, tokens}
       end
     else
@@ -620,6 +597,8 @@ defmodule Gradient.SpecifyErlAst do
     end
   end
 
+  # Private Helpers
+
   @spec match_token_to_form(token(), form()) :: boolean()
   defp match_token_to_form({:int, {l1, _, v1}, _}, {:integer, l2, v2}) do
     l2 = :erl_anno.line(l2)
@@ -756,7 +735,7 @@ defmodule Gradient.SpecifyErlAst do
     {anno, opts}
   end
 
-  def get_line(anno, opts) do
+  defp get_line(anno, opts) do
     case :erl_anno.line(anno) do
       0 ->
         case Keyword.fetch(opts, :line) do
@@ -773,4 +752,40 @@ defmodule Gradient.SpecifyErlAst do
         {:ok, line, anno, opts, true}
     end
   end
+
+  @spec prepare_forms_order(forms()) :: forms()
+  defp prepare_forms_order(forms) do
+    forms
+    |> Enum.sort(fn l, r -> elem(l, 0) == elem(r, 0) and elem(l, 1) > elem(r, 1) end)
+    |> Enum.reverse()
+  end
+
+  defp set_form_end_line(opts, form, forms) do
+    case Enum.find(forms, fn f ->
+           anno = elem(f, 1)
+
+           # Maybe should try to go deeper when generated and try to obtain 
+           # the line from the first child. It should work for sure for clauses, 
+           # but it has to be in the right order (e.g. if clauses are reversed)
+           :erl_anno.line(anno) > 0 and not :erl_anno.generated(anno)
+         end) do
+      nil ->
+        opts
+
+      next_form ->
+        current_line = :erl_anno.line(elem(form, 1))
+        next_line = :erl_anno.line(elem(next_form, 1))
+
+        if current_line == next_line do
+          Keyword.put(opts, :end_line, next_line + 1)
+        else
+          Keyword.put(opts, :end_line, next_line)
+        end
+    end
+  end
+
+  @spec pass_tokens(any(), tokens()) :: {any(), tokens()}
+  defp pass_tokens(form, tokens) do
+    {form, tokens}
+  end
 end
diff --git a/lib/gradient/elixir_fmt.ex b/lib/gradient/elixir_fmt.ex
index 8ef70625..2616e6ea 100644
--- a/lib/gradient/elixir_fmt.ex
+++ b/lib/gradient/elixir_fmt.ex
@@ -5,6 +5,7 @@ defmodule Gradient.ElixirFmt do
   @behaviour Gradient.Fmt
 
   alias :gradualizer_fmt, as: FmtLib
+  alias Gradient.ElixirType
 
   def print_errors(errors, opts) do
     for {file, e} <- errors do
@@ -23,7 +24,12 @@ defmodule Gradient.ElixirFmt do
       _ -> :io.format("~s: ", [file])
     end
 
-    :io.put_chars(format_type_error(error, opts))
+    :io.put_chars(format_error(error, opts))
+  end
+
+  def format_error(error, opts) do
+    opts = Keyword.put(opts, :fmt_type_fun, &ElixirType.pretty_print/1)
+    format_type_error(error, opts)
   end
 
   @impl Gradient.Fmt
@@ -32,6 +38,79 @@ defmodule Gradient.ElixirFmt do
     format_expr_type_error(expression, actual_type, expected_type, opts)
   end
 
+  def format_type_error({:call_undef, anno, module, func, arity}, opts) do
+    :io_lib.format(
+      "~sCall to undefined function ~p:~p/~p~s~n",
+      [
+        format_location(anno, :brief, opts),
+        module,
+        func,
+        arity,
+        format_location(anno, :verbose, opts)
+      ]
+    )
+  end
+
+  def format_type_error({:undef, :record, anno, {module, recName}}, opts) do
+    :io_lib.format(
+      "~sUndefined record ~p:~p~s~n",
+      [
+        format_location(anno, :brief, opts),
+        module,
+        recName,
+        format_location(anno, :verbose, opts)
+      ]
+    )
+  end
+
+  def format_type_error({:undef, :record, anno, recName}, opts) do
+    :io_lib.format(
+      "~sUndefined record ~p~s~n",
+      [format_location(anno, :brief, opts), recName, format_location(anno, :verbose, opts)]
+    )
+  end
+
+  def format_type_error({:undef, :record_field, fieldName}, opts) do
+    :io_lib.format(
+      "~sUndefined record field ~s~s~n",
+      [
+        format_location(fieldName, :brief, opts),
+        pp_expr(fieldName, opts),
+        format_location(fieldName, :verbose, opts)
+      ]
+    )
+  end
+
+  def format_type_error({:undef, :user_type, anno, {name, arity}}, opts) do
+    :io_lib.format(
+      "~sUndefined type ~p/~p~s~n",
+      [format_location(anno, :brief, opts), name, arity, format_location(anno, :verbose, opts)]
+    )
+  end
+
+  def format_type_error({:undef, type, anno, {module, name, arity}}, opts)
+      when type in [:user_type, :remote_type] do
+    type =
+      case type do
+        :user_type -> "type"
+        :remote_type -> "remote type"
+      end
+
+    module = "#{inspect(module)}"
+
+    :io_lib.format(
+      "~sUndefined ~s ~s:~p/~p~s~n",
+      [
+        format_location(anno, :brief, opts),
+        type,
+        module,
+        name,
+        arity,
+        format_location(anno, :verbose, opts)
+      ]
+    )
+  end
+
   def format_type_error(error, opts) do
     :gradualizer_fmt.format_type_error(error, opts) ++ '\n'
   end
@@ -68,9 +147,9 @@ defmodule Gradient.ElixirFmt do
     IO.ANSI.blue() <> "#{inspect(expression)}" <> IO.ANSI.reset()
   end
 
-  def pp_type(expression, _opts) do
-    pp = expression |> :typelib.pp_type() |> to_string()
-    IO.ANSI.yellow() <> pp <> IO.ANSI.reset()
+  def pp_type(type, _opts) do
+    pp = ElixirType.pretty_print(type)
+    IO.ANSI.cyan() <> pp <> IO.ANSI.reset()
   end
 
   def try_highlight_in_context(expression, opts) do
diff --git a/lib/gradient/elixir_type.ex b/lib/gradient/elixir_type.ex
new file mode 100644
index 00000000..96ad3fa0
--- /dev/null
+++ b/lib/gradient/elixir_type.ex
@@ -0,0 +1,135 @@
+defmodule Gradient.ElixirType do
+  @moduledoc """
+  Module to format types.
+
+  TODO records
+  FIXME add tests
+  """
+
+  @doc """
+  Take type and prepare a pretty string representation.
+  """
+  @spec pretty_print(tuple()) :: String.t()
+  def pretty_print({:remote_type, _, [{:atom, _, mod}, {:atom, _, type}, args]}) do
+    args_str = Enum.map(args, &pretty_print(&1)) |> Enum.join(", ")
+    type_str = Atom.to_string(type)
+    mod_str = parse_module(mod)
+    mod_str <> type_str <> "(#{args_str})"
+  end
+
+  def pretty_print({:user_type, _, type, args}) do
+    args_str = Enum.map(args, &pretty_print(&1)) |> Enum.join(", ")
+    type_str = Atom.to_string(type)
+    "#{type_str}(#{args_str})"
+  end
+
+  def pretty_print({:ann_type, _, [var_name, var_type]}) do
+    pretty_print(var_name) <> pretty_print(var_type)
+  end
+
+  def pretty_print({:type, _, :map, :any}) do
+    "map()"
+  end
+
+  def pretty_print({:type, _, :map, assocs}) do
+    assocs_str = Enum.map(assocs, &association_type(&1)) |> Enum.join(", ")
+    "%{" <> assocs_str <> "}"
+  end
+
+  def pretty_print({:op, _, op, type}) do
+    Atom.to_string(op) <> pretty_print(type)
+  end
+
+  def pretty_print({:op, _, op, left_type, right_type}) do
+    operator = " " <> Atom.to_string(op) <> " "
+    pretty_print(left_type) <> operator <> pretty_print(right_type)
+  end
+
+  def pretty_print({:type, _, :fun, [{:type, _, :product, arg_types}, res_type]}) do
+    args = Enum.map(arg_types, &pretty_print(&1)) |> Enum.join(", ")
+    res = pretty_print(res_type)
+    "(" <> args <> " -> " <> res <> ")"
+  end
+
+  def pretty_print({:type, _, :fun, [{:type, _, :any}, res_type]}) do
+    res = pretty_print(res_type)
+    "(... -> " <> res <> ")"
+  end
+
+  def pretty_print({:type, _, :tuple, :any}) do
+    "tuple()"
+  end
+
+  def pretty_print({:type, _, :tuple, elements}) do
+    elements_str = Enum.map(elements, &pretty_print(&1)) |> Enum.join(", ")
+    "{" <> elements_str <> "}"
+  end
+
+  def pretty_print({:atom, _, nil}) do
+    "nil"
+  end
+
+  def pretty_print({:atom, _, val}) when val in [true, false] do
+    Atom.to_string(val)
+  end
+
+  def pretty_print({:atom, _, val}) do
+    ":" <> Atom.to_string(val)
+  end
+
+  def pretty_print({:integer, _, val}) do
+    Integer.to_string(val)
+  end
+
+  def pretty_print({:type, _, :binary, _}) do
+    "binary()"
+  end
+
+  def pretty_print({:type, _, nil, []}) do
+    "[]"
+  end
+
+  def pretty_print({:type, _, :union, [{:atom, _, true}, {:atom, _, false}]}) do
+    "boolean()"
+  end
+
+  def pretty_print({:type, _, :union, args}) do
+    args |> Enum.map(&pretty_print/1) |> Enum.join(" | ")
+  end
+
+  def pretty_print({:type, _, type, args}) do
+    args_str = Enum.map(args, &pretty_print(&1)) |> Enum.join(", ")
+    Atom.to_string(type) <> "(#{args_str})"
+  end
+
+  def pretty_print(type) do
+    "#{inspect(type)}"
+  end
+
+  ######
+  ### Private
+  ######
+
+  @spec parse_module(atom()) :: String.t()
+  defp parse_module(:elixir), do: ""
+
+  defp parse_module(mod) do
+    case Atom.to_string(mod) do
+      "Elixir." <> mod_str -> mod_str <> "."
+      mod -> mod <> "."
+    end
+  end
+
+  @spec association_type(tuple()) :: String.t()
+  defp association_type({:type, _, :map_field_assoc, [key, value]}) do
+    key_str = pretty_print(key)
+    value_str = pretty_print(value)
+    "optional(#{key_str}) => #{value_str}"
+  end
+
+  defp association_type({:type, _, :map_field_exact, [key, value]}) do
+    key_str = pretty_print(key)
+    value_str = pretty_print(value)
+    "required(#{key_str}) => #{value_str}"
+  end
+end
diff --git a/lib/gradient/tokens.ex b/lib/gradient/tokens.ex
new file mode 100644
index 00000000..42f99ce5
--- /dev/null
+++ b/lib/gradient/tokens.ex
@@ -0,0 +1,201 @@
+defmodule Gradient.Tokens do
+  @moduledoc """
+  Functions useful for token management.
+  """
+  alias Gradient.Types, as: T
+
+  @typedoc "Type of conditional with following tokens"
+  @type conditional_t() ::
+          {:case, T.tokens()}
+          | {:cond, T.tokens()}
+          | {:unless, T.tokens()}
+          | {:if, T.tokens()}
+          | {:with, T.tokens()}
+          | :undefined
+
+  @doc """
+  Drop tokens to the first conditional occurrence. Returns type of the encountered
+  conditional and the following tokens.
+  """
+  @spec get_conditional(T.tokens(), integer(), T.options()) :: conditional_t()
+  def get_conditional(tokens, line, opts) do
+    conditionals = [:if, :unless, :cond, :case, :with]
+    {:ok, limit_line} = Keyword.fetch(opts, :end_line)
+
+    drop_tokens_while(tokens, limit_line, fn
+      {:do_identifier, _, c} -> c not in conditionals
+      {:paren_identifier, _, c} -> c not in conditionals
+      {:identifier, _, c} -> c not in conditionals
+      _ -> true
+    end)
+    |> case do
+      [token | _] = tokens when elem(elem(token, 1), 0) == line -> {elem(token, 2), tokens}
+      _ -> :undefined
+    end
+  end
+
+  @doc """
+  Drop tokens to the first list occurrence. Returns type of the encountered
+  list and the following tokens.
+  """
+  @spec get_list(T.tokens(), T.options()) ::
+          {:list, T.tokens()} | {:keyword, T.tokens()} | {:charlist, T.tokens()} | :undefined
+  def get_list(tokens, opts) do
+    tokens = flatten_tokens(tokens)
+    {:ok, limit_line} = Keyword.fetch(opts, :end_line)
+
+    res =
+      drop_tokens_while(tokens, limit_line, fn
+        {:"[", _} -> false
+        {:list_string, _, _} -> false
+        {:kw_identifier, _, id} when id not in [:do] -> false
+        _ -> true
+      end)
+
+    case res do
+      [{:"[", _} | _] = list -> {:list, list}
+      [{:list_string, _, _} | _] = list -> {:charlist, list}
+      [{:kw_identifier, _, _} | _] = list -> {:keyword, list}
+      _ -> :undefined
+    end
+  end
+
+  @doc """
+  Drop tokens to the first tuple occurrence. Returns type of the encountered 
+  list and the following tokens.
+  """
+  @spec get_tuple(T.tokens(), T.options()) ::
+          {:tuple, T.tokens()} | :undefined
+  def get_tuple(tokens, opts) do
+    {:ok, limit_line} = Keyword.fetch(opts, :end_line)
+
+    res =
+      drop_tokens_while(tokens, limit_line, fn
+        {:"{", _} -> false
+        {:kw_identifier, _, _} -> false
+        _ -> true
+      end)
+
+    case res do
+      [{:"{", _} | _] = tuple -> {:tuple, tuple}
+      [{:kw_identifier, _, _} | _] = tuple -> {:tuple, tuple}
+      _ -> :undefined
+    end
+  end
+
+  @doc """
+  Drop tokens till the matcher returns false or the token's line exceeds the limit.
+  """
+  @spec drop_tokens_while(T.tokens(), integer(), (T.token() -> boolean())) :: T.tokens()
+  def drop_tokens_while(tokens, limit_line \\ -1, matcher)
+  def drop_tokens_while([], _, _), do: []
+
+  def drop_tokens_while([token | tokens] = all, limit_line, matcher) do
+    line = get_line_from_token(token)
+
+    limit_passed = limit_line < 0 or line < limit_line
+
+    cond do
+      matcher.(token) and limit_passed ->
+        drop_tokens_while(tokens, limit_line, matcher)
+
+      not limit_passed ->
+        []
+
+      true ->
+        all
+    end
+  end
+
+  @doc """
+  Drop tokens while the token's line is lower than the given location.
+  """
+  @spec drop_tokens_to_line(T.tokens(), integer()) :: T.tokens()
+  def drop_tokens_to_line(tokens, line) do
+    Enum.drop_while(tokens, fn t ->
+      elem(elem(t, 1), 0) < line
+    end)
+  end
+
+  @doc """
+  Get line from token.
+  """
+  @spec get_line_from_token(T.token()) :: integer()
+  def get_line_from_token(token), do: elem(elem(token, 1), 0)
+
+  def get_line_from_form(form) do
+    form
+    |> elem(1)
+    |> :erl_anno.line()
+  end
+
+  @doc """
+  Drop the tokens to binary occurrence and then collect all belonging tokens. 
+  Return tuple where the first element is a list of tokens making up the binary, and the second 
+  element is a list of tokens after the binary.
+  """
+  @spec cut_tokens_to_bin(T.tokens(), integer()) :: {T.tokens(), T.tokens()}
+  def cut_tokens_to_bin(tokens, line) do
+    tokens = drop_tokens_to_line(tokens, line)
+
+    drop_tokens_while(tokens, fn
+      {:"<<", _} -> false
+      {:bin_string, _, _} -> false
+      _ -> true
+    end)
+    |> case do
+      [{:"<<", _} | _] = ts -> cut_bottom(ts, 0)
+      [{:bin_string, _, _} = t | ts] -> {[t], ts}
+      [] -> {[], tokens}
+    end
+  end
+
+  @doc """
+  Flatten the tokens, mostly binaries or string interpolation.
+  """
+  @spec flatten_tokens(T.tokens()) :: T.tokens()
+  def flatten_tokens(tokens) do
+    Enum.map(tokens, &flatten_token/1)
+    |> Enum.concat()
+  end
+
+  # Private
+
+  defp flatten_token(token) do
+    case token do
+      {:bin_string, _, [s]} = t when is_binary(s) ->
+        [t]
+
+      {:bin_string, _, ts} ->
+        flatten_tokens(ts)
+
+      {{_, _, nil}, {_, _, nil}, ts} ->
+        flatten_tokens(ts)
+
+      str when is_binary(str) ->
+        [{:str, {0, 0, nil}, str}]
+
+      _otherwise ->
+        [token]
+    end
+  end
+
+  defp cut_bottom([{:"<<", _} = t | ts], deep) do
+    {ts, cut_ts} = cut_bottom(ts, deep + 1)
+    {[t | ts], cut_ts}
+  end
+
+  defp cut_bottom([{:">>", _} = t | ts], deep) do
+    if deep - 1 > 0 do
+      {ts, cut_ts} = cut_bottom(ts, deep - 1)
+      {[t | ts], cut_ts}
+    else
+      {[t], ts}
+    end
+  end
+
+  defp cut_bottom([t | ts], deep) do
+    {ts, cut_ts} = cut_bottom(ts, deep)
+    {[t | ts], cut_ts}
+  end
+end
diff --git a/lib/gradient/type_annotation.ex b/lib/gradient/type_annotation.ex
index bd8b4e85..9028adc8 100644
--- a/lib/gradient/type_annotation.ex
+++ b/lib/gradient/type_annotation.ex
@@ -1,41 +1,41 @@
 defmodule Gradient.TypeAnnotation do
-
   defmacro annotate_type(expr, type),
-    do: annotate(:'::', expr, type)
+    do: annotate(:"::", expr, type)
 
   defmacro assert_type(expr, type),
-    do: annotate(:':::', expr, type)
+    do: annotate(:":::", expr, type)
 
   defp annotate(type_op, expr, type) do
     erlang_type = elixir_type_to_erlang(type)
-    #IO.inspect(erlang_type, label: "erlang type")
+    # IO.inspect(erlang_type, label: "erlang type")
     {type_op, [], [expr, Macro.to_string(erlang_type)]}
-    #|> IO.inspect(label: "annotation node")
+    # |> IO.inspect(label: "annotation node")
   end
 
   defp elixir_type_to_erlang(type) do
     case type do
       {{:., _, [{:__aliases__, _, path}, name]}, _, [] = _args} ->
-        #unquote({:{}, [], [:string, 0, '\'Elixir.Fake\':t()']})
-        #{:string, 0, Macro.escape("'Elixir.#{Enum.join(path, ".")}':#{name}()" |> to_charlist())}
+        # unquote({:{}, [], [:string, 0, '\'Elixir.Fake\':t()']})
+        # {:string, 0, Macro.escape("'Elixir.#{Enum.join(path, ".")}':#{name}()" |> to_charlist())}
         "'Elixir.#{Enum.join(path, ".")}':#{name}()"
+
       _ when is_atom(type) ->
         Atom.to_string(type)
+
       other ->
-        #unquote({:{}, [], [:string, 0, '\'Elixir.Fake\':t()']})
+        # unquote({:{}, [], [:string, 0, '\'Elixir.Fake\':t()']})
         other
     end
   end
-  
+
   defmacro __using__(_) do
     quote [] do
       import Gradient.TypeAnnotation
       require Gradient.TypeAnnotation
 
-      @compile {:inline, '::': 2, ':::': 2}
-      def unquote(:'::')(expr, _type), do: expr
-      def unquote(:':::')(expr, _type), do: expr
-
+      @compile {:inline, "::": 2, ":::": 2}
+      def unquote(:"::")(expr, _type), do: expr
+      def unquote(:":::")(expr, _type), do: expr
     end
   end
 end
diff --git a/lib/gradient/typed_server.ex b/lib/gradient/typed_server.ex
index 31c32e44..055fe65e 100644
--- a/lib/gradient/typed_server.ex
+++ b/lib/gradient/typed_server.ex
@@ -1,5 +1,4 @@
 defmodule Gradient.TypedServer do
-
   ## TODO: docs, docs, docs!
 
   defmacro __using__(_) do
diff --git a/lib/gradient/typed_server/compile_hooks.ex b/lib/gradient/typed_server/compile_hooks.ex
index 2334629d..0cddbe1b 100644
--- a/lib/gradient/typed_server/compile_hooks.ex
+++ b/lib/gradient/typed_server/compile_hooks.ex
@@ -3,9 +3,10 @@ defmodule Gradient.TypedServer.CompileHooks do
 
   defmacro __before_compile__(env) do
     response_types = Module.get_attribute(env.module, :response_types)
-    #IO.inspect(response_types, label: "response types")
+    # IO.inspect(response_types, label: "response types")
     for {request_tag, response_type} <- response_types do
-      name = Macro.escape(:'call_#{request_tag}')
+      name = Macro.escape(:"call_#{request_tag}")
+
       quote do
         @spec unquote(name)(t(), any) :: unquote(response_type)
         defp unquote(name)(pid, arg) do
@@ -17,17 +18,21 @@ defmodule Gradient.TypedServer.CompileHooks do
 
   def __on_definition__(env, _kind, name, args, _guards, body) do
     if name == :handle do
-      #IO.inspect({name, env}, limit: :infinity)
+      # IO.inspect({name, env}, limit: :infinity)
       IO.inspect({env.module, Module.get_attribute(env.module, :spec)})
     end
+
     request_handler = Module.get_attribute(env.module, :request_handler, :handle)
+
     case request_handler do
       ^name ->
         response_type = find_response_type(env, body)
+
         if response_type != nil do
           {request_tag, _} = Enum.at(args, 0)
           Module.put_attribute(env.module, :response_types, {request_tag, response_type})
         end
+
       _ ->
         :ok
     end
@@ -49,10 +54,13 @@ defmodule Gradient.TypedServer.CompileHooks do
         case Macro.expand(path, env) do
           Gradient.TypedServer ->
             get_response_type_from_typed_call(env, Macro.decompose_call(reply_call))
+
           _other ->
             :ok
         end
+
         reply_call
+
       not_a_call ->
         not_a_call
     end
diff --git a/lib/gradient/types.ex b/lib/gradient/types.ex
new file mode 100644
index 00000000..d37595da
--- /dev/null
+++ b/lib/gradient/types.ex
@@ -0,0 +1,11 @@
+defmodule Gradient.Types do
+  @type token :: tuple()
+  @type tokens :: [tuple()]
+  @type form ::
+          :erl_parse.abstract_clause()
+          | :erl_parse.abstract_expr()
+          | :erl_parse.abstract_form()
+          | :erl_parse.abstract_type()
+  @type forms :: [form()]
+  @type options :: keyword()
+end
diff --git a/lib/gradient/utils.ex b/lib/gradient/utils.ex
deleted file mode 100644
index a1e1761a..00000000
--- a/lib/gradient/utils.ex
+++ /dev/null
@@ -1,100 +0,0 @@
-defmodule Gradient.Utils do
-  @moduledoc """
-  Utility functions that helps obtain info about location from data 
-  """
-
-  @doc """
-  Drop tokens till the matcher returns false or the token's line exceeds the limit.
-  """
-  def drop_tokens_while(tokens, limit_line \\ -1, matcher)
-  def drop_tokens_while([], _, _), do: []
-
-  def drop_tokens_while([token | tokens] = all, limit_line, matcher) do
-    line = get_line_from_token(token)
-
-    limit_passed = limit_line < 0 or line < limit_line
-
-    cond do
-      matcher.(token) and limit_passed ->
-        drop_tokens_while(tokens, limit_line, matcher)
-
-      not limit_passed ->
-        []
-
-      true ->
-        all
-    end
-  end
-
-  def drop_tokens_to_line(tokens, line) do
-    Enum.drop_while(tokens, fn t ->
-      elem(elem(t, 1), 0) < line
-    end)
-  end
-
-  def get_line_from_token(token), do: elem(elem(token, 1), 0)
-
-  def get_line_from_form(form) do
-    form
-    |> elem(1)
-    |> :erl_anno.line()
-  end
-
-  def cut_tokens_to_bin(tokens, line) do
-    tokens = drop_tokens_to_line(tokens, line)
-
-    drop_tokens_while(tokens, fn
-      {:"<<", _} -> false
-      {:bin_string, _, _} -> false
-      _ -> true
-    end)
-    |> case do
-      [{:"<<", _} | _] = ts -> cut_bottom(ts, 0)
-      [{:bin_string, _, _} = t | ts] -> {[t], ts}
-      [] -> {[], tokens}
-    end
-  end
-
-  defp cut_bottom([{:"<<", _} = t | ts], deep) do
-    {ts, cut_ts} = cut_bottom(ts, deep + 1)
-    {[t | ts], cut_ts}
-  end
-
-  defp cut_bottom([{:">>", _} = t | ts], deep) do
-    if deep - 1 > 0 do
-      {ts, cut_ts} = cut_bottom(ts, deep - 1)
-      {[t | ts], cut_ts}
-    else
-      {[t], ts}
-    end
-  end
-
-  defp cut_bottom([t | ts], deep) do
-    {ts, cut_ts} = cut_bottom(ts, deep)
-    {[t | ts], cut_ts}
-  end
-
-  def flat_tokens(tokens) do
-    Enum.map(tokens, &flat_token/1)
-    |> Enum.concat()
-  end
-
-  def flat_token(token) do
-    case token do
-      {:bin_string, _, [s]} = t when is_binary(s) ->
-        [t]
-
-      {:bin_string, _, ts} ->
-        flat_tokens(ts)
-
-      {{_, _, nil}, {_, _, nil}, ts} ->
-        flat_tokens(ts)
-
-      str when is_binary(str) ->
-        [{:str, {0, 0, nil}, str}]
-
-      _otherwise ->
-        [token]
-    end
-  end
-end
diff --git a/lib/mix/tasks/gradient.ex b/lib/mix/tasks/gradient.ex
index 7ccf6a6d..850be6ee 100644
--- a/lib/mix/tasks/gradient.ex
+++ b/lib/mix/tasks/gradient.ex
@@ -7,7 +7,7 @@ defmodule Mix.Tasks.Gradient do
     Mix.Tasks.Compile.run([])
 
     files = get_beams_paths()
-    #IO.puts("Found files:\n #{Enum.join(Enum.concat(Map.values(files)), "\n ")}")
+    # IO.puts("Found files:\n #{Enum.join(Enum.concat(Map.values(files)), "\n ")}")
 
     Application.ensure_all_started(:gradualizer)
 
diff --git a/mix.exs b/mix.exs
index 549e49e8..875226c2 100644
--- a/mix.exs
+++ b/mix.exs
@@ -6,6 +6,7 @@ defmodule Gradient.MixProject do
       app: :gradient,
       version: "0.1.0",
       elixir: "~> 1.12",
+      elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,
       deps: deps(),
       aliases: aliases(),
@@ -13,6 +14,10 @@ defmodule Gradient.MixProject do
     ]
   end
 
+  # Specifies which paths to compile per environment.
+  defp elixirc_paths(:test), do: ["lib", "test/support"]
+  defp elixirc_paths(_), do: ["lib"]
+
   # Run "mix help compile.app" to learn about applications.
   def application do
     [
@@ -23,8 +28,8 @@ defmodule Gradient.MixProject do
   # Run "mix help deps" to learn about dependencies.
   def deps do
     [
-      # {:gradualizer, path: "../Gradualizer", manager: :rebar3}
       {:gradualizer, github: "josefs/Gradualizer", ref: "master", manager: :rebar3},
+      # {:gradualizer, path: "../Gradualizer/", manager: :rebar3},
       {:dialyxir, "~> 1.0", only: [:dev], runtime: false}
     ]
   end
diff --git a/mix.lock b/mix.lock
index 4066e3c8..1055a1ce 100644
--- a/mix.lock
+++ b/mix.lock
@@ -1,6 +1,5 @@
 %{
-  "astranaut": {:hex, :astranaut, "0.9.0", "76f94d98bb1dfadf3b53fbfa3e25d2738582aaf9607fb85a6d8b6e4b92adfbdc", [:rebar3], [], "hexpm", "cf88a1ae5702005bb005facd8702ab9fcca4645d1a06f5446171bf2bb9ca52bb"},
   "dialyxir": {:hex, :dialyxir, "1.1.0", "c5aab0d6e71e5522e77beff7ba9e08f8e02bad90dfbeffae60eaf0cb47e29488", [:mix], [{:erlex, ">= 0.2.6", [hex: :erlex, repo: "hexpm", optional: false]}], "hexpm", "07ea8e49c45f15264ebe6d5b93799d4dd56a44036cf42d0ad9c960bc266c0b9a"},
   "erlex": {:hex, :erlex, "0.2.6", "c7987d15e899c7a2f34f5420d2a2ea0d659682c06ac607572df55a43753aa12e", [:mix], [], "hexpm", "2ed2e25711feb44d52b17d2780eabf998452f6efda104877a3881c2f8c0c0c75"},
-  "gradualizer": {:git, "https://github.com/josefs/Gradualizer.git", "051040400924dfcfe59f7a8936beb34b43278191", [ref: "master"]},
+  "gradualizer": {:git, "https://github.com/josefs/Gradualizer.git", "051040400924dfcfe59f7a8936beb34b43278191", [ref: "master"]}
 }
diff --git a/test/examples/Elixir.Typespec.beam b/test/examples/Elixir.Typespec.beam
new file mode 100644
index 00000000..85aec163
Binary files /dev/null and b/test/examples/Elixir.Typespec.beam differ
diff --git a/test/examples/type/Elixir.WrongRet.beam b/test/examples/type/Elixir.WrongRet.beam
new file mode 100644
index 00000000..cf73ca46
Binary files /dev/null and b/test/examples/type/Elixir.WrongRet.beam differ
diff --git a/test/examples/type/wrong_ret.ex b/test/examples/type/wrong_ret.ex
new file mode 100644
index 00000000..9171431c
--- /dev/null
+++ b/test/examples/type/wrong_ret.ex
@@ -0,0 +1,52 @@
+defmodule WrongRet do
+  @spec ret_wrong_atom() :: atom()
+  def ret_wrong_atom, do: 1
+
+  @spec ret_wrong_atom2() :: atom()
+  def ret_wrong_atom2, do: {:ok, []}
+
+  @spec ret_wrong_atom3() :: atom()
+  def ret_wrong_atom3, do: %{a: 1}
+
+  # @spec ret_wrong_atom4() :: atom()
+  # def ret_wrong_atom4, do: false
+
+  @spec ret_wrong_integer() :: integer()
+  def ret_wrong_integer, do: 1.0
+
+  @spec ret_wrong_integer2() :: integer()
+  def ret_wrong_integer2, do: :ok
+
+  @spec ret_wrong_integer3() :: integer()
+  def ret_wrong_integer3, do: true
+
+  @spec ret_wrong_integer4() :: integer()
+  def ret_wrong_integer4, do: [1, 2, 3]
+
+  @spec ret_out_of_range_int() :: 1..10
+  def ret_out_of_range_int, do: 12
+
+  @spec ret_wrong_boolean() :: boolean()
+  def ret_wrong_boolean, do: :ok
+
+  @spec ret_wrong_boolean2() :: boolean()
+  def ret_wrong_boolean2, do: "1234"
+
+  @spec ret_wrong_boolean3() :: boolean()
+  def ret_wrong_boolean3, do: 1
+
+  @spec ret_wrong_boolean4() :: boolean()
+  def ret_wrong_boolean4, do: [a: 1, b: 2]
+
+  @spec ret_wrong_keyword() :: keyword()
+  def ret_wrong_keyword, do: [1, 2, 3]
+
+  @spec ret_wrong_tuple() :: tuple()
+  def ret_wrong_tuple, do: %{a: 1, b: 2}
+
+  @spec ret_wrong_map() :: map()
+  def ret_wrong_map, do: {:a, 1, 2}
+
+  @spec ret_wrong_fun() :: (... -> atom())
+  def ret_wrong_fun, do: fn -> 12 end
+end
diff --git a/test/examples/typespec.ex b/test/examples/typespec.ex
new file mode 100644
index 00000000..d7aafd23
--- /dev/null
+++ b/test/examples/typespec.ex
@@ -0,0 +1,18 @@
+defmodule Typespec do
+  @type mylist(t) :: [t]
+
+  @spec missing_type() :: Unknown.atom()
+  def missing_type, do: :ok
+
+  @spec missing_type_arg() :: mylist(Unknown.atom())
+  def missing_type_arg, do: [:ok]
+
+  @spec named_type(name :: Unknown.atom()) :: atom()
+  def named_type(name), do: name
+
+  @spec atoms_type(:ok | :error) :: :ok | :error
+  def atoms_type(name), do: name
+
+  @spec atoms_type2(:ok | :error) :: Unknown.atom(:ok | :error)
+  def atoms_type2(name), do: name
+end
diff --git a/test/gradient/specify_erl_ast_test.exs b/test/gradient/ast_specifier_test.exs
similarity index 88%
rename from test/gradient/specify_erl_ast_test.exs
rename to test/gradient/ast_specifier_test.exs
index f5867e92..7a178b65 100644
--- a/test/gradient/specify_erl_ast_test.exs
+++ b/test/gradient/ast_specifier_test.exs
@@ -1,71 +1,19 @@
-defmodule Gradient.SpecifyErlAstTest do
+defmodule Gradient.AstSpecifierTest do
   use ExUnit.Case
-  doctest Gradient.SpecifyErlAst
+  doctest Gradient.AstSpecifier
 
-  alias Gradient.SpecifyErlAst
+  alias Gradient.AstSpecifier
 
-  import Gradient.Utils
-
-  @examples_path "test/examples"
+  import Gradient.TestHelpers
 
   setup_all state do
     {:ok, state}
   end
 
-  describe "get_conditional/1" do
-    test "case" do
-      {tokens, _ast} = load("/conditional/Elixir.Conditional.Case.beam", "/conditional/case.ex")
-      tokens = drop_tokens_to_line(tokens, 2)
-      opts = [end_line: -1]
-      assert {:case, _} = SpecifyErlAst.get_conditional(tokens, 4, opts)
-
-      tokens = drop_tokens_to_line(tokens, 9)
-      assert {:case, _} = SpecifyErlAst.get_conditional(tokens, 10, opts)
-    end
-
-    test "if" do
-      {tokens, _ast} = load("/conditional/Elixir.Conditional.If.beam", "/conditional/if.ex")
-      tokens = drop_tokens_to_line(tokens, 2)
-      opts = [end_line: -1]
-      assert {:if, _} = SpecifyErlAst.get_conditional(tokens, 4, opts)
-
-      tokens = drop_tokens_to_line(tokens, 12)
-      assert {:if, _} = SpecifyErlAst.get_conditional(tokens, 13, opts)
-    end
-
-    test "unless" do
-      {tokens, _ast} =
-        load("/conditional/Elixir.Conditional.Unless.beam", "/conditional/unless.ex")
-
-      tokens = drop_tokens_to_line(tokens, 2)
-      opts = [end_line: -1]
-      assert {:unless, _} = SpecifyErlAst.get_conditional(tokens, 3, opts)
-    end
-
-    test "cond" do
-      {tokens, _ast} = load("/conditional/Elixir.Conditional.Cond.beam", "/conditional/cond.ex")
-
-      tokens = drop_tokens_to_line(tokens, 2)
-      opts = [end_line: -1]
-      assert {:cond, _} = SpecifyErlAst.get_conditional(tokens, 4, opts)
-
-      tokens = drop_tokens_to_line(tokens, 10)
-      assert {:cond, _} = SpecifyErlAst.get_conditional(tokens, 13, opts)
-    end
-
-    test "with" do
-      {tokens, _ast} = load("/conditional/Elixir.Conditional.With.beam", "/conditional/with.ex")
-
-      tokens = drop_tokens_to_line(tokens, 6)
-      opts = [end_line: -1]
-      assert {:with, _} = SpecifyErlAst.get_conditional(tokens, 7, opts)
-    end
-  end
-
-  describe "add_missing_loc_literals/2" do
+  describe "run_mappers/2" do
     test "messy test on simple_app" do
       {tokens, ast} = example_data()
-      new_ast = SpecifyErlAst.add_missing_loc_literals(ast, tokens)
+      new_ast = AstSpecifier.run_mappers(ast, tokens)
 
       assert is_list(new_ast)
     end
@@ -73,7 +21,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "integer" do
       {tokens, ast} = load("/basic/Elixir.Basic.Int.beam", "/basic/int.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 2, :int, 0, [{:clause, 2, [], [], [{:integer, 2, 1}]}]} = inline
 
@@ -83,7 +31,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "float" do
       {tokens, ast} = load("/basic/Elixir.Basic.Float.beam", "/basic/float.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
       assert {:function, 2, :float, 0, [{:clause, 2, [], [], [{:float, 2, 0.12}]}]} = inline
 
       assert {:function, 4, :float_block, 0, [{:clause, 4, [], [], [{:float, 5, 0.12}]}]} = block
@@ -92,7 +40,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "atom" do
       {tokens, ast} = load("/basic/Elixir.Basic.Atom.beam", "/basic/atom.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 2, :atom, 0, [{:clause, 2, [], [], [{:atom, 2, :ok}]}]} = inline
 
@@ -102,7 +50,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "char" do
       {tokens, ast} = load("/basic/Elixir.Basic.Char.beam", "/basic/char.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 2, :char, 0, [{:clause, 2, [], [], [{:integer, 2, 99}]}]} = inline
 
@@ -112,7 +60,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "charlist" do
       {tokens, ast} = load("/basic/Elixir.Basic.Charlist.beam", "/basic/charlist.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       # TODO propagate location to each charlist element
       assert {:function, 2, :charlist, 0,
@@ -137,7 +85,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "string" do
       {tokens, ast} = load("/basic/Elixir.Basic.String.beam", "/basic/string.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 2, :string, 0,
               [
@@ -156,7 +104,7 @@ defmodule Gradient.SpecifyErlAstTest do
       {tokens, ast} = load("/Elixir.Tuple.beam", "/tuple.ex")
 
       [tuple_in_str2, tuple_in_str, tuple_in_list, _list_in_tuple, tuple | _] =
-        SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+        AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       # FIXME
       assert {:function, 18, :tuple_in_str2, 0,
@@ -264,7 +212,7 @@ defmodule Gradient.SpecifyErlAstTest do
       {tokens, ast} = load("/basic/Elixir.Basic.Binary.beam", "/basic/binary.ex")
 
       [complex2, complex, bin_block, bin | _] =
-        SpecifyErlAst.add_missing_loc_literals(ast, tokens)
+        AstSpecifier.run_mappers(ast, tokens)
         |> Enum.reverse()
 
       assert {:function, 13, :complex2, 0,
@@ -333,7 +281,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "case conditional" do
       {tokens, ast} = load("/conditional/Elixir.Conditional.Case.beam", "/conditional/case.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 2, :case_, 0,
               [
@@ -363,8 +311,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "if conditional" do
       {tokens, ast} = load("/conditional/Elixir.Conditional.If.beam", "/conditional/if.ex")
 
-      [block, inline, if_ | _] =
-        SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline, if_ | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 12, :if_block, 0,
               [
@@ -413,7 +360,7 @@ defmodule Gradient.SpecifyErlAstTest do
       {tokens, ast} =
         load("/conditional/Elixir.Conditional.Unless.beam", "/conditional/unless.ex")
 
-      [block | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {
                :function,
@@ -438,7 +385,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "cond conditional" do
       {tokens, ast} = load("/conditional/Elixir.Conditional.Cond.beam", "/conditional/cond.ex")
 
-      [block, inline | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block, inline | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 2, :cond_, 1,
               [
@@ -509,7 +456,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "with conditional" do
       {tokens, ast} = load("/conditional/Elixir.Conditional.With.beam", "/conditional/with.ex")
 
-      [block | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      [block | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
       assert {:function, 6, :test_with, 0,
               [
@@ -543,7 +490,7 @@ defmodule Gradient.SpecifyErlAstTest do
       beam_file = "/Elixir.Basic.beam"
       {tokens, ast} = load(beam_file, ex_file)
 
-      specified_ast = SpecifyErlAst.add_missing_loc_literals(ast, tokens)
+      specified_ast = AstSpecifier.run_mappers(ast, tokens)
       IO.inspect(specified_ast)
       assert is_list(specified_ast)
     end
@@ -554,10 +501,10 @@ defmodule Gradient.SpecifyErlAstTest do
     opts = [end_line: -1]
 
     assert {{:integer, 21, 12}, tokens} =
-             SpecifyErlAst.specify_line({:integer, 21, 12}, tokens, opts)
+             AstSpecifier.specify_line({:integer, 21, 12}, tokens, opts)
 
     assert {{:integer, 22, 12}, _tokens} =
-             SpecifyErlAst.specify_line({:integer, 20, 12}, tokens, opts)
+             AstSpecifier.specify_line({:integer, 20, 12}, tokens, opts)
   end
 
   test "cons_to_charlist/1" do
@@ -565,17 +512,7 @@ defmodule Gradient.SpecifyErlAstTest do
       {:cons, 0, {:integer, 0, 49},
        {:cons, 0, {:integer, 0, 48}, {:cons, 0, {:integer, 0, 48}, {nil, 0}}}}
 
-    assert '100' == SpecifyErlAst.cons_to_charlist(cons)
-  end
-
-  test "get_list_from_tokens" do
-    tokens = example_string_tokens()
-    ts = drop_tokens_to_line(tokens, 4)
-    opts = [end_line: -1]
-    assert {:charlist, _} = SpecifyErlAst.get_list_from_tokens(ts, opts)
-
-    ts = drop_tokens_to_line(ts, 6)
-    assert {:list, _} = SpecifyErlAst.get_list_from_tokens(ts, opts)
+    assert '100' == AstSpecifier.cons_to_charlist(cons)
   end
 
   describe "test that prints result" do
@@ -583,7 +520,7 @@ defmodule Gradient.SpecifyErlAstTest do
     test "specify/1" do
       {_tokens, forms} = example_data()
 
-      SpecifyErlAst.specify(forms)
+      AstSpecifier.specify(forms)
       |> IO.inspect()
     end
 
@@ -597,7 +534,7 @@ defmodule Gradient.SpecifyErlAstTest do
   test "function call" do
     {tokens, ast} = load("/Elixir.Call.beam", "/call.ex")
 
-    [call, _ | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+    [call, _ | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 5, :call, 0,
             [
@@ -617,7 +554,7 @@ defmodule Gradient.SpecifyErlAstTest do
   test "pipe" do
     {tokens, ast} = load("/Elixir.Pipe.beam", "/pipe_op.ex")
 
-    [block | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+    [block | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 2, :pipe, 0,
             [
@@ -649,8 +586,7 @@ defmodule Gradient.SpecifyErlAstTest do
   test "guards" do
     {tokens, ast} = load("/conditional/Elixir.Conditional.Guard.beam", "/conditional/guards.ex")
 
-    [guarded_case, guarded_fun | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+    [guarded_case, guarded_fun | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 3, :guarded_fun, 1,
             [
@@ -691,7 +627,7 @@ defmodule Gradient.SpecifyErlAstTest do
     {tokens, ast} = load("/Elixir.RangeEx.beam", "/range.ex")
 
     [to_list, match_range, rev_range_step, range_step, range | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 18, :to_list, 0,
             [
@@ -771,7 +707,7 @@ defmodule Gradient.SpecifyErlAstTest do
   test "list comprehension" do
     {tokens, ast} = load("/Elixir.ListComprehension.beam", "/list_comprehension.ex")
 
-    [block | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+    [block | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 2, :lc, 0,
             [
@@ -819,8 +755,7 @@ defmodule Gradient.SpecifyErlAstTest do
   test "list" do
     {tokens, ast} = load("/Elixir.ListEx.beam", "/list.ex")
 
-    [ht2, ht, list, _wrap | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+    [ht2, ht, list, _wrap | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 5, :list, 0,
             [
@@ -861,7 +796,7 @@ defmodule Gradient.SpecifyErlAstTest do
     {tokens, ast} = load("/Elixir.Try.beam", "/try.ex")
 
     [body_after, try_after, try_else, try_rescue | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 2, :try_rescue, 0,
             [
@@ -1042,7 +977,7 @@ defmodule Gradient.SpecifyErlAstTest do
     {tokens, ast} = load("/Elixir.MapEx.beam", "/map.ex")
 
     [pattern_matching_str, pattern_matching, test_map_str, test_map, empty_map | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 2, :empty_map, 0, [{:clause, 2, [], [], [{:map, 3, []}]}]} = empty_map
 
@@ -1106,7 +1041,7 @@ defmodule Gradient.SpecifyErlAstTest do
     {tokens, ast} = load("/struct/Elixir.StructEx.beam", "/struct/struct.ex")
 
     [get2, get, update, empty, struct | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 8, :update, 0,
             [
@@ -1223,7 +1158,7 @@ defmodule Gradient.SpecifyErlAstTest do
     {tokens, ast} = load("/record/Elixir.RecordEx.beam", "/record/record.ex")
 
     [update, init, empty, macro3, macro2, macro1 | _] =
-      SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+      AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 7, :empty, 0,
             [
@@ -1304,7 +1239,7 @@ defmodule Gradient.SpecifyErlAstTest do
   test "receive" do
     {tokens, ast} = load("/Elixir.Receive.beam", "/receive.ex")
 
-    [recv, recv2 | _] = SpecifyErlAst.add_missing_loc_literals(ast, tokens) |> Enum.reverse()
+    [recv, recv2 | _] = AstSpecifier.run_mappers(ast, tokens) |> Enum.reverse()
 
     assert {:function, 2, :recv2, 0,
             [
@@ -1351,61 +1286,83 @@ defmodule Gradient.SpecifyErlAstTest do
             ]} = recv
   end
 
-  @spec load(String.t(), String.t()) :: {list(), list()}
-  def load(beam_file, ex_file) do
-    beam_file = String.to_charlist(@examples_path <> beam_file)
-    ex_file = @examples_path <> ex_file
-
-    code =
-      File.read!(ex_file)
-      |> String.to_charlist()
-
-    {:ok, tokens} =
-      code
-      |> :elixir.string_to_tokens(1, 1, ex_file, [])
-
-    {:ok, {_, [abstract_code: {:raw_abstract_v1, ast}]}} =
-      :beam_lib.chunks(beam_file, [:abstract_code])
-
-    ast = replace_file_path(ast, ex_file)
-    {tokens, ast}
-  end
-
-  def example_data() do
-    beam_path = (@examples_path <> "/Elixir.SimpleApp.beam") |> String.to_charlist()
-    file_path = @examples_path <> "/simple_app.ex"
-
-    code =
-      File.read!(file_path)
-      |> String.to_charlist()
-
-    {:ok, tokens} =
-      code
-      |> :elixir.string_to_tokens(1, 1, file_path, [])
-
-    {:ok, {SimpleApp, [abstract_code: {:raw_abstract_v1, ast}]}} =
-      :beam_lib.chunks(beam_path, [:abstract_code])
-
-    ast = replace_file_path(ast, file_path)
-    {tokens, ast}
+  test "typespec" do
+    {tokens, ast} = load("/Elixir.Typespec.beam", "/typespec.ex")
+
+    [atoms_type2, atoms_type, named_type, missing_type_arg, missing_type | _] =
+      AstSpecifier.run_mappers(ast, tokens)
+      |> filter_specs()
+      |> Enum.reverse()
+
+    assert {:attribute, 17, :spec,
+            {{:atoms_type2, 1},
+             [
+               {:type, 17, :fun,
+                [
+                  {:type, 17, :product,
+                   [{:type, 17, :union, [{:atom, 17, :ok}, {:atom, 17, :error}]}]},
+                  {:remote_type, 17,
+                   [
+                     {:atom, 17, Unknown},
+                     {:atom, 17, :atom},
+                     [{:type, 17, :union, [{:atom, 17, :ok}, {:atom, 17, :error}]}]
+                   ]}
+                ]}
+             ]}} = atoms_type2
+
+    assert {:attribute, 14, :spec,
+            {{:atoms_type, 1},
+             [
+               {:type, 14, :fun,
+                [
+                  {:type, 14, :product,
+                   [{:type, 14, :union, [{:atom, 14, :ok}, {:atom, 14, :error}]}]},
+                  {:type, 14, :union, [{:atom, 14, :ok}, {:atom, 14, :error}]}
+                ]}
+             ]}} = atoms_type
+
+    assert {:attribute, 11, :spec,
+            {{:named_type, 1},
+             [
+               {:type, 11, :fun,
+                [
+                  {:type, 11, :product,
+                   [
+                     {:ann_type, 11,
+                      [
+                        {:var, 11, :name},
+                        {:remote_type, 11, [{:atom, 11, Unknown}, {:atom, 11, :atom}, []]}
+                      ]}
+                   ]},
+                  {:type, 11, :atom, []}
+                ]}
+             ]}} = named_type
+
+    assert {:attribute, 8, :spec,
+            {{:missing_type_arg, 0},
+             [
+               {:type, 8, :fun,
+                [
+                  {:type, 8, :product, []},
+                  {:user_type, 8, :mylist,
+                   [{:remote_type, 8, [{:atom, 8, Unknown}, {:atom, 8, :atom}, []]}]}
+                ]}
+             ]}} = missing_type_arg
+
+    assert {:attribute, 5, :spec,
+            {{:missing_type, 0},
+             [
+               {:type, 5, :fun,
+                [
+                  {:type, 5, :product, []},
+                  {:remote_type, 5, [{:atom, 5, Unknown}, {:atom, 5, :atom}, []]}
+                ]}
+             ]}} = missing_type
   end
 
-  def example_string_tokens() do
-    file_path = @examples_path <> "/string_example.ex"
-
-    code =
-      File.read!(file_path)
-      |> String.to_charlist()
-
-    {:ok, tokens} =
-      code
-      |> :elixir.string_to_tokens(1, 1, file_path, [])
-
-    tokens
-  end
+  # Helpers
 
-  def replace_file_path([_ | forms], path) do
-    path = String.to_charlist(path)
-    [{:attribute, 1, :file, {path, 1}} | forms]
+  def filter_specs(ast) do
+    Enum.filter(ast, &match?({:attribute, _, :spec, _}, &1))
   end
 end
diff --git a/test/gradient/elixir_fmt_test.exs b/test/gradient/elixir_fmt_test.exs
index be7e7afb..16fcbb51 100644
--- a/test/gradient/elixir_fmt_test.exs
+++ b/test/gradient/elixir_fmt_test.exs
@@ -3,9 +3,15 @@ defmodule Gradient.ElixirFmtTest do
   doctest Gradient.ElixirFmt
 
   alias Gradient.ElixirFmt
+  import Gradient.TestHelpers
+  alias Gradient.AstSpecifier
 
   @example_module_path "test/examples/simple_app.ex"
 
+  setup_all config do
+    load_wrong_ret_error_examples(config)
+  end
+
   test "try_highlight_in_context/2" do
     opts = [forms: basic_erlang_forms()]
     expression = {:integer, 31, 12}
@@ -18,6 +24,113 @@ defmodule Gradient.ElixirFmtTest do
     assert res == expected
   end
 
+  describe "types format" do
+    test "return integer() instead atom()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_atom)
+
+      assert String.contains?(msg, "atom()")
+      assert String.contains?(msg, "1")
+    end
+
+    test "return tuple() instead atom()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_atom2)
+
+      assert String.contains?(msg, "atom()")
+      assert String.contains?(msg, "{:ok, []}")
+    end
+
+    test "return map() instead atom()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_atom3)
+
+      assert String.contains?(msg, "atom()")
+      assert String.contains?(msg, "%{required(:a) => 1}")
+    end
+
+    test "return float() instead integer()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_integer)
+
+      assert String.contains?(msg, "integer()")
+      assert String.contains?(msg, "1.0")
+    end
+
+    test "return atom() instead integer()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_integer2)
+
+      assert String.contains?(msg, "integer()")
+      assert String.contains?(msg, ":ok")
+    end
+
+    test "return boolean() instead integer()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_integer3)
+
+      assert String.contains?(msg, "integer()")
+      assert String.contains?(msg, "true")
+    end
+
+    test "return list() instead integer()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_integer4)
+
+      assert String.contains?(msg, "integer()")
+      assert String.contains?(msg, "nonempty_list()")
+    end
+
+    test "return integer() out of the range()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_out_of_range_int)
+
+      assert String.contains?(msg, "range(1, 10)")
+      assert String.contains?(msg, "12")
+    end
+
+    test "return atom() instead boolean()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_boolean)
+
+      assert String.contains?(msg, "boolean()")
+      assert String.contains?(msg, ":ok")
+    end
+
+    test "return binary() instead boolean()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_boolean2)
+
+      assert String.contains?(msg, "boolean()")
+      assert String.contains?(msg, "binary()")
+    end
+
+    test "return integer() instead boolean()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_boolean3)
+
+      assert String.contains?(msg, "boolean()")
+      assert String.contains?(msg, "1")
+    end
+
+    test "return keyword() instead boolean()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_boolean4)
+
+      assert String.contains?(msg, "boolean()")
+      assert String.contains?(msg, "nonempty_list()")
+    end
+
+    test "return list() instead keyword()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_keyword)
+
+      assert String.contains?(msg, "{atom(), any()}")
+      assert String.contains?(msg, "1")
+    end
+
+    test "return tuple() instead map()", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_map)
+
+      assert String.contains?(msg, "map()")
+      assert String.contains?(msg, "{:a, 1, 2}")
+    end
+
+    test "return lambda with wrong returned type", %{wrong_ret_errors: errors} do
+      msg = format_error_to_binary(errors.ret_wrong_fun)
+
+      assert String.contains?(msg, "atom()")
+      assert String.contains?(msg, "12")
+    end
+  end
+
   @tag :skip
   test "format_expr_type_error/4" do
     opts = [forms: basic_erlang_forms()]
@@ -29,7 +142,52 @@ defmodule Gradient.ElixirFmtTest do
     IO.puts(res)
   end
 
-  def basic_erlang_forms() do
+  # Helpers
+
+  defp basic_erlang_forms() do
     [{:attribute, 1, :file, {@example_module_path, 1}}]
   end
+
+  defp type_check_file(ast, opts) do
+    forms = AstSpecifier.specify(ast)
+    opts = Keyword.put(opts, :return_errors, true)
+    opts = Keyword.put(opts, :forms, forms)
+
+    errors =
+      forms
+      |> :gradualizer.type_check_forms(opts)
+      |> Enum.map(&elem(&1, 1))
+
+    {errors, forms}
+  end
+
+  defp format_error_to_binary(error, opts \\ []) do
+    error
+    |> ElixirFmt.format_error(opts)
+    |> :erlang.iolist_to_binary()
+  end
+
+  @spec load_wrong_ret_error_examples(map()) :: map()
+  defp load_wrong_ret_error_examples(config) do
+    {_tokens, ast} = load("/type/Elixir.WrongRet.beam", "/type/wrong_ret.ex")
+
+    {errors, forms} = type_check_file(ast, [])
+    names = get_function_names_from_ast(forms)
+
+    errors_map =
+      Enum.zip(names, errors)
+      |> Map.new()
+
+    Map.put(config, :wrong_ret_errors, errors_map)
+  end
+
+  @spec get_function_names_from_ast([tuple()]) :: [atom()]
+  def get_function_names_from_ast(ast) do
+    ast
+    |> Enum.filter(fn
+      {:function, _, name, _, _} -> name != :__info__
+      _ -> false
+    end)
+    |> Enum.map(&elem(&1, 2))
+  end
 end
diff --git a/test/gradient/tokens_test.exs b/test/gradient/tokens_test.exs
new file mode 100644
index 00000000..93ab5733
--- /dev/null
+++ b/test/gradient/tokens_test.exs
@@ -0,0 +1,97 @@
+defmodule Gradient.TokensTest do
+  use ExUnit.Case
+  doctest Gradient.Tokens
+
+  alias Gradient.Tokens
+
+  import Gradient.TestHelpers
+
+  test "drop_tokens_while" do
+    tokens = example_tokens()
+
+    matcher = fn
+      {:atom, _, :ok} -> false
+      _ -> true
+    end
+
+    assert [] =
+             Tokens.drop_tokens_while(
+               tokens,
+               5,
+               matcher
+             )
+
+    refute [] ==
+             Tokens.drop_tokens_while(
+               tokens,
+               6,
+               matcher
+             )
+
+    refute [] ==
+             Tokens.drop_tokens_while(
+               tokens,
+               matcher
+             )
+  end
+
+  test "get_list_from_tokens" do
+    tokens = example_string_tokens()
+    ts = Tokens.drop_tokens_to_line(tokens, 4)
+    opts = [end_line: -1]
+    assert {:charlist, _} = Tokens.get_list(ts, opts)
+
+    ts = Tokens.drop_tokens_to_line(ts, 6)
+    assert {:list, _} = Tokens.get_list(ts, opts)
+  end
+
+  describe "get_conditional/1" do
+    test "case" do
+      {tokens, _ast} = load("/conditional/Elixir.Conditional.Case.beam", "/conditional/case.ex")
+      tokens = Tokens.drop_tokens_to_line(tokens, 2)
+      opts = [end_line: -1]
+      assert {:case, _} = Tokens.get_conditional(tokens, 4, opts)
+
+      tokens = Tokens.drop_tokens_to_line(tokens, 9)
+      assert {:case, _} = Tokens.get_conditional(tokens, 10, opts)
+    end
+
+    test "if" do
+      {tokens, _ast} = load("/conditional/Elixir.Conditional.If.beam", "/conditional/if.ex")
+      tokens = Tokens.drop_tokens_to_line(tokens, 2)
+      opts = [end_line: -1]
+      assert {:if, _} = Tokens.get_conditional(tokens, 4, opts)
+
+      tokens = Tokens.drop_tokens_to_line(tokens, 12)
+      assert {:if, _} = Tokens.get_conditional(tokens, 13, opts)
+    end
+
+    test "unless" do
+      {tokens, _ast} =
+        load("/conditional/Elixir.Conditional.Unless.beam", "/conditional/unless.ex")
+
+      tokens = Tokens.drop_tokens_to_line(tokens, 2)
+      opts = [end_line: -1]
+      assert {:unless, _} = Tokens.get_conditional(tokens, 3, opts)
+    end
+
+    test "cond" do
+      {tokens, _ast} = load("/conditional/Elixir.Conditional.Cond.beam", "/conditional/cond.ex")
+
+      tokens = Tokens.drop_tokens_to_line(tokens, 2)
+      opts = [end_line: -1]
+      assert {:cond, _} = Tokens.get_conditional(tokens, 4, opts)
+
+      tokens = Tokens.drop_tokens_to_line(tokens, 10)
+      assert {:cond, _} = Tokens.get_conditional(tokens, 13, opts)
+    end
+
+    test "with" do
+      {tokens, _ast} = load("/conditional/Elixir.Conditional.With.beam", "/conditional/with.ex")
+
+      tokens = Tokens.drop_tokens_to_line(tokens, 6)
+      opts = [end_line: -1]
+      assert {:with, _} = Tokens.get_conditional(tokens, 7, opts)
+    end
+  end
+end
diff --git a/test/gradient/utils_test.exs b/test/gradient/utils_test.exs
deleted file mode 100644
index 2815f37a..00000000
--- a/test/gradient/utils_test.exs
+++ /dev/null
@@ -1,50 +0,0 @@
-defmodule Gradient.UtilsTest do
-  use ExUnit.Case
-
-  alias Gradient.Utils
-
-  @examples_path "test/examples"
-
-  test "drop_tokens_while" do
-    tokens = example_tokens()
-
-    matcher = fn
-      {:atom, _, :ok} -> false
-      _ -> true
-    end
-
-    assert [] =
-             Utils.drop_tokens_while(
-               tokens,
-               5,
-               matcher
-             )
-
-    refute [] ==
-             Utils.drop_tokens_while(
-               tokens,
-               6,
-               matcher
-             )
-
-    refute [] ==
-             Utils.drop_tokens_while(
-               tokens,
-               matcher
-             )
-  end
-
-  def example_tokens() do
-    file_path = @examples_path <> "/conditional/cond.ex"
-
-    code =
-      File.read!(file_path)
-      |> String.to_charlist()
-
-    {:ok, tokens} =
-      code
-      |> :elixir.string_to_tokens(1, 1, file_path, [])
-
-    tokens
-  end
-end
diff --git a/test/support/helpers.ex b/test/support/helpers.ex
new file mode 100644
index 00000000..b8763286
--- /dev/null
+++ b/test/support/helpers.ex
@@ -0,0 +1,80 @@
+defmodule Gradient.TestHelpers do
+  alias Gradient.Types, as: T
+
+  @examples_path "test/examples"
+
+  @spec load(String.t(), String.t()) :: {T.tokens(), T.forms()}
+  def load(beam_file, ex_file) do
+    beam_file = String.to_charlist(@examples_path <> beam_file)
+    ex_file = @examples_path <> ex_file
+
+    code =
+      File.read!(ex_file)
+      |> String.to_charlist()
+
+    {:ok, tokens} =
+      code
+      |> :elixir.string_to_tokens(1, 1, ex_file, [])
+
+    {:ok, {_, [abstract_code: {:raw_abstract_v1, ast}]}} =
+      :beam_lib.chunks(beam_file, [:abstract_code])
+
+    ast = replace_file_path(ast, ex_file)
+    {tokens, ast}
+  end
+
+  @spec example_data() :: {T.tokens(), T.forms()}
+  def example_data() do
+    beam_path = (@examples_path <> "/Elixir.SimpleApp.beam") |> String.to_charlist()
+    file_path = @examples_path <> "/simple_app.ex"
+
+    code =
+      File.read!(file_path)
+      |> String.to_charlist()
+
+    {:ok, tokens} =
+      code
+      |> :elixir.string_to_tokens(1, 1, file_path, [])
+
+    {:ok, {SimpleApp, [abstract_code: {:raw_abstract_v1, ast}]}} =
+      :beam_lib.chunks(beam_path, [:abstract_code])
+
+    ast = replace_file_path(ast, file_path)
+    {tokens, ast}
+  end
+
+  @spec example_tokens() :: T.tokens()
+  def example_tokens() do
+    file_path = @examples_path <> "/conditional/cond.ex"
+
+    code =
+      File.read!(file_path)
+      |> String.to_charlist()
+
+    {:ok, tokens} =
+      code
+      |> :elixir.string_to_tokens(1, 1, file_path, [])
+
+    tokens
+  end
+
+  @spec example_string_tokens() :: T.tokens()
+  def example_string_tokens() do
+    file_path = @examples_path <> "/string_example.ex"
+
+    code =
+      File.read!(file_path)
+      |> String.to_charlist()
+
+    {:ok, tokens} =
+      code
+      |> :elixir.string_to_tokens(1, 1, file_path, [])
+
+    tokens
+  end
+
+  defp replace_file_path([_ | forms], path) do
+    path = String.to_charlist(path)
+    [{:attribute, 1, :file, {path, 1}} | forms]
+  end
+end