From 30959320f7c486ed5170852a2989aeb56af85d1a Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 18:32:30 -0300
Subject: [PATCH 01/10] feat: start bringing features from other branch

---
 .../compiler/code_server/dependency_graph.ex  |  38 ++++++
 lib/fika/compiler/function_signature.ex       |   2 +
 lib/fika/compiler/parser/helper.ex            |   8 ++
 lib/fika/compiler/parser/types.ex             |   8 ++
 lib/fika/compiler/type_checker/types/loop.ex  |  25 ++++
 test/fika/compiler/module_compiler_test.exs   | 110 ++++++++++++++++++
 test/fika/compiler/parser_test.exs            |  52 +++++++++
 7 files changed, 243 insertions(+)
 create mode 100644 lib/fika/compiler/code_server/dependency_graph.ex
 create mode 100644 lib/fika/compiler/type_checker/types/loop.ex

diff --git a/lib/fika/compiler/code_server/dependency_graph.ex b/lib/fika/compiler/code_server/dependency_graph.ex
new file mode 100644
index 0000000..1c6b3d7
--- /dev/null
+++ b/lib/fika/compiler/code_server/dependency_graph.ex
@@ -0,0 +1,38 @@
+defmodule CodeServer.FunctionDependencies do
+  alias Fika.Compiler.FunctionSignature
+
+  def new_graph, do: :digraph.new()
+
+  @spec set_function_dependency(
+          graph :: :digraph.graph(),
+          source_function :: FunctionSignature.t(),
+          sink_function :: FunctionSignature.t()
+        ) :: :digraph.graph()
+  def set_function_dependency(graph, source_function, sink_function) do
+    graph
+    |> add_vertex(source_function)
+    |> add_vertex(sink_function)
+    |> add_edge(source_function, sink_function)
+    |> check_cycle(source_function)
+  end
+
+  def check_cycle(graph, function) do
+    if :digraph.get_cycle(graph, function) do
+      {:error, :cycle_encountered}
+    else
+      :ok
+    end
+  end
+
+  defp add_vertex(graph, v) do
+    :digraph.add_vertex(graph, v, v)
+    graph
+  end
+
+  defp add_edge(graph, source, target) do
+    edge = {source, target}
+
+    :digraph.add_edge(graph, edge, source, target, edge)
+    graph
+  end
+end
diff --git a/lib/fika/compiler/function_signature.ex b/lib/fika/compiler/function_signature.ex
index 78752e7..5732af6 100644
--- a/lib/fika/compiler/function_signature.ex
+++ b/lib/fika/compiler/function_signature.ex
@@ -1,6 +1,8 @@
 defmodule Fika.Compiler.FunctionSignature do
   defstruct [:module, :function, :return, types: []]
 
+  @type t :: %__MODULE__{}
+
   alias Fika.Compiler.TypeChecker.Types, as: T
 
   defimpl String.Chars, for: __MODULE__ do
diff --git a/lib/fika/compiler/parser/helper.ex b/lib/fika/compiler/parser/helper.ex
index 08b5ceb..95464fe 100644
--- a/lib/fika/compiler/parser/helper.ex
+++ b/lib/fika/compiler/parser/helper.ex
@@ -125,6 +125,14 @@ defmodule Fika.Compiler.Parser.Helper do
     %T.Effect{type: inner_type}
   end
 
+  def do_to_ast({[{_, _, inner_type}], _line}, :loop_type) when is_struct(inner_type) do
+    %T.Loop{type: inner_type}
+  end
+
+  def do_to_ast({[inner_type], _line}, :loop_type) do
+    %T.Loop{type: inner_type}
+  end
+
   def do_to_ast({inner_types, _line}, :tuple_type) do
     %T.Tuple{elements: inner_types}
   end
diff --git a/lib/fika/compiler/parser/types.ex b/lib/fika/compiler/parser/types.ex
index 08c6a7f..78c9d0e 100644
--- a/lib/fika/compiler/parser/types.ex
+++ b/lib/fika/compiler/parser/types.ex
@@ -114,6 +114,13 @@ defmodule Fika.Compiler.Parser.Types do
     identifier_str
     |> label("type variable")
 
+  loop_type =
+    ignore(string("Loop("))
+    |> concat(parsec(:type))
+    |> ignore(string(")"))
+    |> label("loop type")
+    |> Helper.to_ast(:loop_type)
+
   base_type =
     choice([
       string_type,
@@ -124,6 +131,7 @@ defmodule Fika.Compiler.Parser.Types do
       function_type,
       list_type,
       effect_type,
+      loop_type,
       record_type,
       map_type,
       tuple_type,
diff --git a/lib/fika/compiler/type_checker/types/loop.ex b/lib/fika/compiler/type_checker/types/loop.ex
new file mode 100644
index 0000000..c63cbfc
--- /dev/null
+++ b/lib/fika/compiler/type_checker/types/loop.ex
@@ -0,0 +1,25 @@
+defmodule Fika.Compiler.TypeChecker.Types.Loop do
+  @moduledoc """
+  Defines when a function may not terminate due to recursion
+  """
+  defstruct [:type]
+
+  alias Fika.Compiler.TypeChecker.Types, as: T
+
+  def new(), do: %__MODULE__{}
+  def new([]), do: %__MODULE__{}
+  def new([type]), do: %__MODULE__{type: type}
+  def new(type), do: %__MODULE__{type: type}
+
+  def is_loop(%__MODULE__{}), do: true
+  def is_loop(_), do: false
+
+  def is_empty_loop(%__MODULE__{type: nil}), do: true
+  def is_empty_loop(_), do: false
+
+  defimpl String.Chars, for: T.Loop do
+    def to_string(%{type: type}) do
+      "Loop(#{type})"
+    end
+  end
+end
diff --git a/test/fika/compiler/module_compiler_test.exs b/test/fika/compiler/module_compiler_test.exs
index 7e4e385..d28ba06 100644
--- a/test/fika/compiler/module_compiler_test.exs
+++ b/test/fika/compiler/module_compiler_test.exs
@@ -90,4 +90,114 @@ defmodule Fika.Compiler.ModuleCompilerTest do
 
     File.rm!(temp_file)
   end
+
+  test "does not hang when there is local recursion with both direct and indirect cycles" do
+    module = Path.join(System.tmp_dir!(), "foo") |> String.to_atom()
+
+    temp_file = "#{module}.fi"
+
+    str = """
+    fn f : Loop(nil) do
+      g()
+    end
+    fn g : Loop(nil) do
+      h()
+    end
+    fn h : Loop(nil) do
+      f()
+    end
+    fn a : Loop(nil) do
+      c()
+      b()
+    end
+    fn b : Loop(nil) do
+      c()
+      a()
+    end
+    fn c : Int do
+      1 + 1
+    end
+    """
+
+    File.write!(temp_file, str)
+
+    assert {:error, "Type check error"} == ModuleCompiler.compile(module)
+
+    File.rm!(temp_file)
+  end
+
+  test "does not hang when there is intermodule recursion" do
+    foo_module = Path.join(System.tmp_dir!(), "foo") |> String.to_atom()
+    foo_temp_file = "#{foo_module}.fi"
+
+    bar_module = Path.join(System.tmp_dir!(), "bar") |> String.to_atom()
+    bar_temp_file = "#{bar_module}.fi"
+
+    foo_str = """
+    use bar
+
+    fn f : Loop(nil) do
+      bar.g()
+    end
+    """
+
+    bar_str = """
+    use foo
+
+    fn g : Loop(nil) do
+      foo.f()
+    end
+    """
+
+    File.write!(foo_temp_file, foo_str)
+    File.write!(bar_temp_file, bar_str)
+
+    File.cd!(System.tmp_dir!(), fn ->
+      assert {:ok, _, _, _} = ModuleCompiler.compile("foo")
+    end)
+
+    File.rm!(foo_temp_file)
+    File.rm!(bar_temp_file)
+  end
+
+  test "infer return type when the recursion is not at the top level" do
+    tmp_dir = System.tmp_dir!()
+    module = Path.join(tmp_dir, "foo") |> String.to_atom()
+
+    temp_file = "#{module}.fi"
+
+    str = """
+    fn factorial(x: Int) : Int do
+      do_factorial(x, 0)
+    end
+    fn do_factorial(x: Int, acc: Int) : Loop(Int) do
+      if x <= 1 do
+        acc
+      else
+        do_factorial(x - 1, acc * x)
+      end
+    end
+    fn top_level(x: Int) : Int do
+      second_level_a(x)
+    end
+    fn second_level_a(x: Int) : Loop(Int) do
+      if x > 1 do
+        second_level_b(x - 1)
+      else
+        1
+      end
+    end
+    fn second_level_b(x: Int) : Loop(Int) do
+      second_level_a(x - 1)
+    end
+    """
+
+    File.write!(temp_file, str)
+
+    File.cd!(tmp_dir, fn ->
+      assert {:ok, "foo", "foo.fi", _binary} = ModuleCompiler.compile("foo")
+    end)
+
+    File.rm!(temp_file)
+  end
 end
diff --git a/test/fika/compiler/parser_test.exs b/test/fika/compiler/parser_test.exs
index 02b3342..20315b7 100644
--- a/test/fika/compiler/parser_test.exs
+++ b/test/fika/compiler/parser_test.exs
@@ -382,6 +382,58 @@ defmodule Fika.Compiler.ParserTest do
              } ==
                TestParser.function_def!(str)
     end
+
+    test "parses loop type" do
+      str = """
+      fn foo(x: Int) : Loop(Int) do
+        if x <= 1 do
+          1
+        else
+          foo(x - 1)
+        end
+      end
+      """
+
+      assert {
+               :function,
+               [position: {7, 79, 82}],
+               {
+                 :foo,
+                 [
+                   {
+                     {:identifier, {1, 0, 8}, :x},
+                     {:type, {1, 0, 13}, :Int}
+                   }
+                 ],
+                 {
+                   :type,
+                   {1, 0, 26},
+                   %Fika.Compiler.TypeChecker.Types.Loop{type: :Int}
+                 },
+                 [
+                   {
+                     {:if, {6, 73, 78}},
+                     {
+                       :call,
+                       {:<=, {2, 30, 41}},
+                       [{:identifier, {2, 30, 36}, :x}, {:integer, {2, 30, 41}, 1}],
+                       "fika/kernel"
+                     },
+                     [{:integer, {3, 45, 50}, 1}],
+                     [
+                       {:call, {:foo, {5, 58, 72}},
+                        [
+                          {:call, {:-, {5, 58, 71}},
+                           [{:identifier, {5, 58, 67}, :x}, {:integer, {5, 58, 71}, 1}],
+                           "fika/kernel"}
+                        ], nil}
+                     ]
+                   }
+                 ]
+               }
+             } ==
+               TestParser.function_def!(str)
+    end
   end
 
   describe "if-else expression" do

From 07fb031d1e4f5e4811f95073c701bc9af5a5d31f Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 18:53:18 -0300
Subject: [PATCH 02/10] feat: make first-pass loop inference work

---
 lib/fika/compiler/code_server.ex              | 59 ++++++++++++++++++-
 ...ency_graph.ex => function_dependencies.ex} |  2 +-
 lib/fika/compiler/erl_translate.ex            |  2 +-
 lib/fika/compiler/type_checker.ex             | 37 ++++++++----
 .../type_checker/parallel_type_checker.ex     |  3 +-
 test/fika/compiler/module_compiler_test.exs   | 26 ++++----
 6 files changed, 101 insertions(+), 28 deletions(-)
 rename lib/fika/compiler/code_server/{dependency_graph.ex => function_dependencies.ex} (93%)

diff --git a/lib/fika/compiler/code_server.ex b/lib/fika/compiler/code_server.ex
index 5f6d1ca..e551fc9 100644
--- a/lib/fika/compiler/code_server.ex
+++ b/lib/fika/compiler/code_server.ex
@@ -65,6 +65,23 @@ defmodule Fika.Compiler.CodeServer do
     GenServer.call(__MODULE__, {:write_binaries, dest})
   end
 
+  @doc false
+  def get_dependency_graph do
+    GenServer.call(__MODULE__, :get_dependency_graph)
+  end
+
+  @spec set_function_dependency(source :: String.t() | nil, target :: String.t() | nil) ::
+          :ok | {:error, :cycle_encountered}
+  def set_function_dependency(source, target) do
+    GenServer.call(__MODULE__, {:set_function_dependency, source, target})
+  end
+
+  @spec check_cycle(source :: String.t() | nil) ::
+          :ok | {:error, :cycle_encountered}
+  def check_cycle(node) do
+    GenServer.call(__MODULE__, {:check_cycle, node})
+  end
+
   def init(_) do
     state = init_state()
 
@@ -123,6 +140,15 @@ defmodule Fika.Compiler.CodeServer do
     {:reply, :ok, init_state()}
   end
 
+  def handle_call(:get_dependency_graph, _from, %{function_dependencies: graph} = state) do
+    vertices = graph |> :digraph.vertices() |> Enum.sort()
+    edges = graph |> :digraph.edges() |> Enum.sort()
+
+    deps = %{vertices: vertices, edges: edges}
+
+    {:reply, deps, state}
+  end
+
   def handle_call({:get_type, signature}, from, state) do
     module = signature.module
     function = signature.function
@@ -150,6 +176,36 @@ defmodule Fika.Compiler.CodeServer do
     {:noreply, state}
   end
 
+  def handle_call({:set_function_dependency, source, target}, _from, state)
+      when is_nil(source) or is_nil(target) do
+    {:reply, :ok, state}
+  end
+
+  def handle_call(
+        {:set_function_dependency, source, target},
+        _from,
+        %{function_dependencies: graph} = state
+      ) do
+    response = __MODULE__.FunctionDependencies.set_function_dependency(graph, source, target)
+
+    {:reply, response, state}
+  end
+
+  def handle_call({:check_cycle, node}, _from, state)
+      when is_nil(node) do
+    {:reply, :ok, state}
+  end
+
+  def handle_call(
+        {:check_cycle, node},
+        _from,
+        %{function_dependencies: graph} = state
+      ) do
+    response = __MODULE__.FunctionDependencies.check_cycle(graph, node)
+
+    {:reply, response, state}
+  end
+
   def handle_call(:list_all_types, _from, state) do
     {:reply, state.public_functions, state}
   end
@@ -280,7 +336,8 @@ defmodule Fika.Compiler.CodeServer do
       parent_pid: nil,
       dev_token: Application.get_env(:fika, :dev_token),
       remote_endpoint: Application.get_env(:fika, :remote_endpoint, "https://fikaapp.com/code"),
-      binaries: []
+      binaries: [],
+      function_dependencies: __MODULE__.FunctionDependencies.new_graph()
     }
     |> put_default_types(DefaultTypes.kernel())
     |> put_default_types(DefaultTypes.io())
diff --git a/lib/fika/compiler/code_server/dependency_graph.ex b/lib/fika/compiler/code_server/function_dependencies.ex
similarity index 93%
rename from lib/fika/compiler/code_server/dependency_graph.ex
rename to lib/fika/compiler/code_server/function_dependencies.ex
index 1c6b3d7..9567571 100644
--- a/lib/fika/compiler/code_server/dependency_graph.ex
+++ b/lib/fika/compiler/code_server/function_dependencies.ex
@@ -1,4 +1,4 @@
-defmodule CodeServer.FunctionDependencies do
+defmodule Fika.Compiler.CodeServer.FunctionDependencies do
   alias Fika.Compiler.FunctionSignature
 
   def new_graph, do: :digraph.new()
diff --git a/lib/fika/compiler/erl_translate.ex b/lib/fika/compiler/erl_translate.ex
index 7327f97..dba38a9 100644
--- a/lib/fika/compiler/erl_translate.ex
+++ b/lib/fika/compiler/erl_translate.ex
@@ -20,7 +20,7 @@ defmodule Fika.Compiler.ErlTranslate do
     translate_exp(exp)
   end
 
-  def erl_module_name(module_name_str) do
+  def erl_module_name(module_name_str) when is_binary(module_name_str) do
     module_name_str
     |> String.replace("/", ".")
     |> String.to_atom()
diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index 72fad52..5424ab2 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -516,17 +516,32 @@ defmodule Fika.Compiler.TypeChecker do
     %FunctionSignature{module: module, function: to_string(function_name), types: arg_types}
   end
 
-  # Local function
-  defp get_type(nil, signature, env) do
-    if pid = env[:type_checker_pid] do
-      ParallelTypeChecker.get_result(pid, signature)
-    else
-      SequentialTypeChecker.get_result(signature, env)
-    end
-  end
+  defp get_type(module, target_signature, env) do
+    is_local_call = is_nil(module) or module == env[:module_name]
+
+    current_signature = env[:current_signature]
+
+    pid = env[:type_checker_pid]
+
+    function_dependency =
+      if current_signature do
+        CodeServer.set_function_dependency(current_signature, target_signature)
+      else
+        :ok
+      end
+
+    case {is_local_call, function_dependency, pid} do
+      {true, :ok, pid} when is_pid(pid) ->
+        ParallelTypeChecker.get_result(pid, target_signature)
 
-  # Remote function
-  defp get_type(_module, signature, _env) do
-    CodeServer.get_type(signature)
+      {true, :ok, nil} ->
+        SequentialTypeChecker.get_result(target_signature, env)
+
+      {_, {:error, :cycle_encountered}, _} ->
+        {:ok, T.Loop.new()}
+
+      {false, _, _} ->
+        CodeServer.get_type(target_signature)
+    end
   end
 end
diff --git a/lib/fika/compiler/type_checker/parallel_type_checker.ex b/lib/fika/compiler/type_checker/parallel_type_checker.ex
index f40dbd1..eb8e578 100644
--- a/lib/fika/compiler/type_checker/parallel_type_checker.ex
+++ b/lib/fika/compiler/type_checker/parallel_type_checker.ex
@@ -58,7 +58,8 @@ defmodule Fika.Compiler.TypeChecker.ParallelTypeChecker do
         result =
           TypeChecker.check(function, %{
             type_checker_pid: pid,
-            module: state.module_name
+            module: state.module_name,
+            current_signature: signature
           })
 
         __MODULE__.post_result(pid, signature, result)
diff --git a/test/fika/compiler/module_compiler_test.exs b/test/fika/compiler/module_compiler_test.exs
index d28ba06..0c58b1e 100644
--- a/test/fika/compiler/module_compiler_test.exs
+++ b/test/fika/compiler/module_compiler_test.exs
@@ -92,25 +92,25 @@ defmodule Fika.Compiler.ModuleCompilerTest do
   end
 
   test "does not hang when there is local recursion with both direct and indirect cycles" do
-    module = Path.join(System.tmp_dir!(), "foo") |> String.to_atom()
+    module = Path.join(System.tmp_dir!(), "foo")
 
     temp_file = "#{module}.fi"
 
     str = """
-    fn f : Loop(nil) do
+    fn f : Loop(:nil) do
       g()
     end
-    fn g : Loop(nil) do
+    fn g : Loop(:nil) do
       h()
     end
-    fn h : Loop(nil) do
+    fn h : Loop(:nil) do
       f()
     end
-    fn a : Loop(nil) do
+    fn a : Loop(:nil) | Int do
       c()
       b()
     end
-    fn b : Loop(nil) do
+    fn b : Loop(:nil) | Int do
       c()
       a()
     end
@@ -136,7 +136,7 @@ defmodule Fika.Compiler.ModuleCompilerTest do
     foo_str = """
     use bar
 
-    fn f : Loop(nil) do
+    fn f : Loop(:nil) do
       bar.g()
     end
     """
@@ -144,7 +144,7 @@ defmodule Fika.Compiler.ModuleCompilerTest do
     bar_str = """
     use foo
 
-    fn g : Loop(nil) do
+    fn g : Loop(:nil) do
       foo.f()
     end
     """
@@ -167,27 +167,27 @@ defmodule Fika.Compiler.ModuleCompilerTest do
     temp_file = "#{module}.fi"
 
     str = """
-    fn factorial(x: Int) : Int do
+    fn factorial(x: Int) : Int | Loop(:nil) do
       do_factorial(x, 0)
     end
-    fn do_factorial(x: Int, acc: Int) : Loop(Int) do
+    fn do_factorial(x: Int, acc: Int) : Int | Loop(:nil) do
       if x <= 1 do
         acc
       else
         do_factorial(x - 1, acc * x)
       end
     end
-    fn top_level(x: Int) : Int do
+    fn top_level(x: Int) : Int | Loop(:nil) do
       second_level_a(x)
     end
-    fn second_level_a(x: Int) : Loop(Int) do
+    fn second_level_a(x: Int) : Int | Loop(:nil) do
       if x > 1 do
         second_level_b(x - 1)
       else
         1
       end
     end
-    fn second_level_b(x: Int) : Loop(Int) do
+    fn second_level_b(x: Int) : Int | Loop(:nil) do
       second_level_a(x - 1)
     end
     """

From 60ea42fb85b71e5f355e77b81abf3fb1cc64b156 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 19:30:30 -0300
Subject: [PATCH 03/10] feat: introduce TypeChecker.Env struct for better
 manipulation of env

---
 lib/fika/compiler/module_compiler.ex          |   3 +-
 lib/fika/compiler/type_checker.ex             | 146 ++++++++++++------
 lib/fika/compiler/type_checker/match.ex       |   4 +-
 .../type_checker/parallel_type_checker.ex     |   2 +-
 test/fika/compiler/type_checker_test.exs      | 118 +++++++-------
 5 files changed, 163 insertions(+), 110 deletions(-)

diff --git a/lib/fika/compiler/module_compiler.ex b/lib/fika/compiler/module_compiler.ex
index fdced84..9b1f03c 100644
--- a/lib/fika/compiler/module_compiler.ex
+++ b/lib/fika/compiler/module_compiler.ex
@@ -4,6 +4,7 @@ defmodule Fika.Compiler.ModuleCompiler do
   alias Fika.Compiler.{
     Parser,
     TypeChecker.ParallelTypeChecker,
+    TypeChecker.Env,
     ErlTranslate,
     CodeServer
   }
@@ -43,7 +44,7 @@ defmodule Fika.Compiler.ModuleCompiler do
   end
 
   defp init(module_atom) do
-    %{
+    %Env{
       # Full name of the current module as atom
       module_name: module_atom,
       # Path of module file
diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index 5424ab2..70bbf72 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -14,9 +14,31 @@ defmodule Fika.Compiler.TypeChecker do
 
   require Logger
 
+  defmodule Env do
+    defstruct [
+      :ast,
+      :latest_called_function,
+      :type_checker_pid,
+      :module,
+      :module_name,
+      :file,
+      :current_signature,
+      has_effect: false,
+      scope: %{}
+    ]
+
+    def add_variable_to_scope(env, variable, type) do
+      %{env | scope: put_in(env.scope, [variable], type)}
+    end
+
+    def update_scope(env, variable, update_fn) do
+      %{env | scope: update_in(env.scope, [variable], update_fn)}
+    end
+  end
+
   # Given the AST of a function definition, this function checks if the return
   # type is indeed the type that's inferred from the body of the function.
-  def check({:function, _line, {_, _, return_type, _}} = function, env) do
+  def check({:function, _line, {_, _, return_type, _}} = function, %Env{} = env) do
     {:type, _line, expected_type} = return_type
 
     case infer(function, env) do
@@ -31,11 +53,31 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
+  defp unwrap_loop(env, t, %T.Loop{type: t} = loop_type) do
+    # The function can be a top-level function which depends on a loop
+    # In this case, we can unwrap the loop
+
+    case CodeServer.check_cycle(env.current_signature) do
+      :ok ->
+        {:ok, t}
+
+      _ ->
+        {:error, "Expected type: #{t}, got: #{loop_type}"}
+    end
+  end
+
+  defp unwrap_loop(_env, t, t), do: {:ok, t}
+
+  defp unwrap_loop(_env, expected_type, inferred_type),
+    do: {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
+
   # Given the AST of a function definition, this function infers the
   # return type of the body of the function.
-  def infer({:function, _line, {name, args, _type, exprs}}, env) do
+  def infer({:function, _line, {name, args, _type, exprs}}, %Env{} = env) do
     Logger.debug("Inferring type of function: #{name}")
 
+    Logger.debug("First call to #{inspect(env.current_signature)}")
+
     env =
       env
       |> Map.put(:scope, %{})
@@ -45,7 +87,7 @@ defmodule Fika.Compiler.TypeChecker do
     case infer_block(env, exprs) do
       {:ok, type, env} ->
         type =
-          if env[:has_effect] do
+          if env.has_effect do
             %T.Effect{type: type}
           else
             type
@@ -58,37 +100,37 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
-  def infer_block(env, []) do
+  def infer_block(%Env{} = env, []) do
     Logger.debug("Block is empty.")
     {:ok, nil, env}
   end
 
-  def infer_block(env, [exp]) do
-    infer_exp(env, exp)
+  def infer_block(%Env{} = env, [exp]) do
+    infer_exp(%Env{} = env, exp)
   end
 
-  def infer_block(env, [exp | exp_list]) do
-    case infer_exp(env, exp) do
-      {:ok, _type, env} -> infer_block(env, exp_list)
+  def infer_block(%Env{} = env, [exp | exp_list]) do
+    case infer_exp(%Env{} = env, exp) do
+      {:ok, _type, env} -> infer_block(%Env{} = env, exp_list)
       error -> error
     end
   end
 
   # Integer literals
-  def infer_exp(env, {:integer, _line, integer}) do
+  def infer_exp(%Env{} = env, {:integer, _line, integer}) do
     Logger.debug("Integer #{integer} found. Type: Int")
     {:ok, :Int, env}
   end
 
   # Booleans
-  def infer_exp(env, {:boolean, _line, boolean}) do
+  def infer_exp(%Env{} = env, {:boolean, _line, boolean}) do
     Logger.debug("Boolean #{boolean} found. Type: Bool")
     {:ok, :Bool, env}
   end
 
   # Variables
-  def infer_exp(env, {:identifier, _line, name}) do
-    type = get_in(env, [:scope, name])
+  def infer_exp(%Env{} = env, {:identifier, _line, name}) do
+    type = env.scope[name]
 
     if type do
       Logger.debug("Variable type found from scope: #{name}:#{type}")
@@ -100,13 +142,13 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # External function calls
-  def infer_exp(env, {:ext_call, _line, {m, f, _, type}}) do
+  def infer_exp(%Env{} = env, {:ext_call, _line, {m, f, _, type}}) do
     Logger.debug("Return type of ext function #{m}.#{f} specified as #{type}")
     {:ok, type, env}
   end
 
   # Function calls
-  def infer_exp(env, {:call, {name, _line}, args, module}) do
+  def infer_exp(%Env{} = env, {:call, {name, _line}, args, module}) do
     exp = %{args: args, name: name}
     # module_name = module || Env.current_module(env)
     Logger.debug("Inferring type of function: #{name}")
@@ -115,8 +157,8 @@ defmodule Fika.Compiler.TypeChecker do
 
   # Function calls using reference
   # exp has to be a function ref type
-  def infer_exp(env, {:call, {exp, _line}, args}) do
-    case infer_exp(env, exp) do
+  def infer_exp(%Env{} = env, {:call, {exp, _line}, args}) do
+    case infer_exp(%Env{} = env, exp) do
       {:ok, %T.FunctionRef{arg_types: arg_types, return_type: type}, env} ->
         case do_infer_args_without_name(env, args) do
           {:ok, ^arg_types, env} ->
@@ -140,11 +182,11 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # =
-  def infer_exp(env, {{:=, _}, {:identifier, _line, left}, right}) do
-    case infer_exp(env, right) do
+  def infer_exp(%Env{} = env, {{:=, _}, {:identifier, _line, left}, right}) do
+    case infer_exp(%Env{} = env, right) do
       {:ok, type, env} ->
         Logger.debug("Adding variable to scope: #{left}:#{type}")
-        env = put_in(env, [:scope, left], type)
+        env = Env.add_variable_to_scope(env, left, type)
         {:ok, type, env}
 
       error ->
@@ -153,7 +195,7 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # String
-  def infer_exp(env, {:string, _line, string_parts}) do
+  def infer_exp(%Env{} = env, {:string, _line, string_parts}) do
     Enum.reduce_while(string_parts, {:ok, nil, env}, fn
       string, {:ok, _acc_type, acc_env} when is_binary(string) ->
         Logger.debug("String #{string} found. Type: String")
@@ -179,12 +221,12 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # List
-  def infer_exp(env, {:list, _, exps}) do
+  def infer_exp(%Env{} = env, {:list, _, exps}) do
     infer_list_exps(env, exps)
   end
 
   # Tuple
-  def infer_exp(env, {:tuple, _, exps}) do
+  def infer_exp(%Env{} = env, {:tuple, _, exps}) do
     case do_infer_tuple_exps(exps, env) do
       {:ok, exp_types, env} ->
         {:ok, %T.Tuple{elements: exp_types}, env}
@@ -195,7 +237,7 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # Record
-  def infer_exp(env, {:record, _, name, key_values}) do
+  def infer_exp(%Env{} = env, {:record, _, name, key_values}) do
     if name do
       # Lookup type of name, ensure it matches.
       Logger.error("Not implemented")
@@ -212,18 +254,18 @@ defmodule Fika.Compiler.TypeChecker do
 
   # Map
   # TODO: refactor this when union types are available
-  def infer_exp(env, {:map, _, [kv | rest_kvs]}) do
+  def infer_exp(%Env{} = env, {:map, _, [kv | rest_kvs]}) do
     {key, value} = kv
 
-    with {:ok, key_type, env} <- infer_exp(env, key),
-         {:ok, value_type, env} <- infer_exp(env, value) do
+    with {:ok, key_type, env} <- infer_exp(%Env{} = env, key),
+         {:ok, value_type, env} <- infer_exp(%Env{} = env, value) do
       map_type = %T.Map{key_type: key_type, value_type: value_type}
 
       Enum.reduce_while(rest_kvs, {:ok, map_type, env}, fn {k, v}, {:ok, type, env} ->
         %{key_type: key_type, value_type: value_type} = type
 
-        with {:key, {:ok, ^key_type, env}} <- {:key, infer_exp(env, k)},
-             {:value, {:ok, ^value_type, env}} <- {:value, infer_exp(env, v)} do
+        with {:key, {:ok, ^key_type, env}} <- {:key, infer_exp(%Env{} = env, k)},
+             {:value, {:ok, ^value_type, env}} <- {:value, infer_exp(%Env{} = env, v)} do
           {:cont, {:ok, type, env}}
         else
           {:key, {:ok, diff_type, _}} ->
@@ -243,10 +285,10 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # Function ref
-  def infer_exp(env, {:function_ref, _, {module, function_name, arg_types}}) do
+  def infer_exp(%Env{} = env, {:function_ref, _, {module, function_name, arg_types}}) do
     Logger.debug("Inferring type of function: #{function_name}")
 
-    signature = get_function_signature(module || env[:module], function_name, arg_types)
+    signature = get_function_signature(module || env.module, function_name, arg_types)
 
     case get_type(module, signature, env) do
       {:ok, type} ->
@@ -259,13 +301,13 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # Atom value
-  def infer_exp(env, {:atom, _line, atom}) do
+  def infer_exp(%Env{} = env, {:atom, _line, atom}) do
     Logger.debug("Atom value found. Type: #{atom}")
     {:ok, atom, env}
   end
 
   # if-else expression
-  def infer_exp(env, {{:if, _line}, condition, if_block, else_block}) do
+  def infer_exp(%Env{} = env, {{:if, _line}, condition, if_block, else_block}) do
     Logger.debug("Inferring an if-else expression")
 
     case infer_if_else_condition(env, condition) do
@@ -275,19 +317,19 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # case expression
-  def infer_exp(env, {{:case, _line}, exp, clauses}) do
+  def infer_exp(%Env{} = env, {{:case, _line}, exp, clauses}) do
     Logger.debug("Inferring a case expression")
 
     # Check the type of exp.
     # For each clause, ensure all of the patterns return {:ok, env}
-    with {:ok, rhs_type, env} <- infer_exp(env, exp),
+    with {:ok, rhs_type, env} <- infer_exp(%Env{} = env, exp),
          {:ok, type, env} <- infer_case_clauses(env, rhs_type, clauses) do
       {:ok, type, env}
     end
   end
 
   # anonymous function
-  def infer_exp(env, {:anonymous_function, _line, args, exps}) do
+  def infer_exp(%Env{} = env, {:anonymous_function, _line, args, exps}) do
     Logger.debug("Inferring type of anonymous function")
 
     env =
@@ -295,7 +337,7 @@ defmodule Fika.Compiler.TypeChecker do
       |> Map.put(:scope, %{})
       |> add_args_to_scope(args)
 
-    case infer_block(env, exps) do
+    case infer_block(%Env{} = env, exps) do
       {:ok, return_type, _env} ->
         arg_types = Enum.map(args, fn {_, {:type, _, type}} -> type end)
         type = %T.FunctionRef{arg_types: arg_types, return_type: return_type}
@@ -312,7 +354,7 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   def init_env(ast) do
-    %{ast: ast, scope: %{}}
+    %Env{ast: ast, scope: %{}}
   end
 
   # TODO: made it work, now make it pretty.
@@ -324,7 +366,7 @@ defmodule Fika.Compiler.TypeChecker do
                                                               {env, types, unmatched} ->
         case Match.match_case(env, pattern, unmatched) do
           {:ok, env, unmatched} ->
-            case infer_block(env, block) do
+            case infer_block(%Env{} = env, block) do
               {:ok, type, env} -> {:cont, {env, [type | types], unmatched}}
               error -> {:halt, error}
             end
@@ -377,7 +419,7 @@ defmodule Fika.Compiler.TypeChecker do
 
   defp do_infer_key_values(key_values, env) do
     Enum.reduce_while(key_values, {:ok, [], env}, fn {k, v}, {:ok, acc, env} ->
-      case infer_exp(env, v) do
+      case infer_exp(%Env{} = env, v) do
         {:ok, type, env} ->
           {:identifier, _, key} = k
           {:cont, {:ok, [{key, type} | acc], env}}
@@ -390,7 +432,7 @@ defmodule Fika.Compiler.TypeChecker do
 
   defp do_infer_tuple_exps(exps, env) do
     Enum.reduce_while(exps, {:ok, [], env}, fn exp, {:ok, acc, env} ->
-      case infer_exp(env, exp) do
+      case infer_exp(%Env{} = env, exp) do
         {:ok, exp_type, env} ->
           {:cont, {:ok, [exp_type | acc], env}}
 
@@ -410,7 +452,7 @@ defmodule Fika.Compiler.TypeChecker do
   def infer_args(env, exp, module) do
     case do_infer_args(env, exp) do
       {:ok, type_acc, env} ->
-        signature = get_function_signature(module || env[:module], exp.name, type_acc)
+        signature = get_function_signature(module || env.module, exp.name, type_acc)
 
         case get_type(module, signature, env) do
           {:ok, %T.Effect{type: type}} ->
@@ -434,14 +476,14 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   defp infer_list_exps(env, [exp]) do
-    case infer_exp(env, exp) do
+    case infer_exp(%Env{} = env, exp) do
       {:ok, type, env} -> {:ok, %T.List{type: type}, env}
       error -> error
     end
   end
 
   defp infer_list_exps(env, [exp | rest]) do
-    {:ok, type, env} = infer_exp(env, exp)
+    {:ok, type, env} = infer_exp(%Env{} = env, exp)
 
     Enum.reduce_while(rest, {:ok, %T.List{type: type}, env}, fn exp, {:ok, acc_type, acc_env} ->
       case infer_exp(acc_env, exp) do
@@ -464,7 +506,7 @@ defmodule Fika.Compiler.TypeChecker do
 
   defp do_infer_args(env, exp) do
     Enum.reduce_while(exp.args, {:ok, [], env}, fn arg, {:ok, type_acc, env} ->
-      case infer_exp(env, arg) do
+      case infer_exp(%Env{} = env, arg) do
         {:ok, type, env} ->
           Logger.debug("Argument of #{exp.name} is type: #{type}")
           {:cont, {:ok, [type | type_acc], env}}
@@ -486,7 +528,7 @@ defmodule Fika.Compiler.TypeChecker do
   # TODO: This can be merged with do_infer_args
   defp do_infer_args_without_name(env, args) do
     Enum.reduce_while(args, {:ok, [], env}, fn arg, {:ok, acc, env} ->
-      case infer_exp(env, arg) do
+      case infer_exp(%Env{} = env, arg) do
         {:ok, type, env} ->
           Logger.debug("Argument is type: #{type}")
           {:cont, {:ok, [type | acc], env}}
@@ -508,20 +550,24 @@ defmodule Fika.Compiler.TypeChecker do
   defp add_args_to_scope(env, args) do
     Enum.reduce(args, env, fn {{:identifier, _, name}, {:type, _, type}}, env ->
       Logger.debug("Adding arg type to scope: #{name}:#{type}")
-      put_in(env, [:scope, name], type)
+      Env.add_variable_to_scope(env, name, type)
     end)
   end
 
+  defp set_latest_call(env, signature) do
+    Map.put(env, :latest_called_function, signature)
+  end
+
   defp get_function_signature(module, function_name, arg_types) do
     %FunctionSignature{module: module, function: to_string(function_name), types: arg_types}
   end
 
   defp get_type(module, target_signature, env) do
-    is_local_call = is_nil(module) or module == env[:module_name]
+    is_local_call = is_nil(module) or module == env.module_name
 
-    current_signature = env[:current_signature]
+    current_signature = env.current_signature
 
-    pid = env[:type_checker_pid]
+    pid = env.type_checker_pid
 
     function_dependency =
       if current_signature do
diff --git a/lib/fika/compiler/type_checker/match.ex b/lib/fika/compiler/type_checker/match.ex
index 88c3d16..5be4c4c 100644
--- a/lib/fika/compiler/type_checker/match.ex
+++ b/lib/fika/compiler/type_checker/match.ex
@@ -15,6 +15,8 @@ defmodule Fika.Compiler.TypeChecker.Match do
 
   """
 
+  alias Fika.Compiler.TypeChecker.Env
+
   # Returns:
   # {:ok, env, unmatched_types} | :error
   def match_case(env, lhs_ast, rhs_types) when is_list(rhs_types) do
@@ -126,7 +128,7 @@ defmodule Fika.Compiler.TypeChecker.Match do
 
   defp do_match_case(env, {:identifier, _, name}, rhs) do
     env =
-      update_in(env, [:scope, name], fn
+      Env.update_scope(env, name, fn
         nil -> rhs
         %T.Union{types: types} -> T.Union.new([rhs | T.Union.to_list(types)])
         type -> T.Union.new([rhs, type])
diff --git a/lib/fika/compiler/type_checker/parallel_type_checker.ex b/lib/fika/compiler/type_checker/parallel_type_checker.ex
index eb8e578..fae31dd 100644
--- a/lib/fika/compiler/type_checker/parallel_type_checker.ex
+++ b/lib/fika/compiler/type_checker/parallel_type_checker.ex
@@ -56,7 +56,7 @@ defmodule Fika.Compiler.TypeChecker.ParallelTypeChecker do
     Enum.each(state.unchecked_functions, fn {signature, function} ->
       Task.start_link(fn ->
         result =
-          TypeChecker.check(function, %{
+          TypeChecker.check(function, %TypeChecker.Env{
             type_checker_pid: pid,
             module: state.module_name,
             current_signature: signature
diff --git a/test/fika/compiler/type_checker_test.exs b/test/fika/compiler/type_checker_test.exs
index 170b57a..aea72b9 100644
--- a/test/fika/compiler/type_checker_test.exs
+++ b/test/fika/compiler/type_checker_test.exs
@@ -8,6 +8,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
     FunctionSignature
   }
 
+  alias Fika.Compiler.TypeChecker.Env
   alias Fika.Compiler.TypeChecker.Types, as: T
 
   setup do
@@ -19,7 +20,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:ok, :Int, _} = TypeChecker.infer_exp(%{}, ast)
+    assert {:ok, :Int, _} = TypeChecker.infer_exp(%Env{}, ast)
   end
 
   test "infer type of atom expressions" do
@@ -27,7 +28,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     {:atom, {1, 0, 2}, :a} = ast = TestParser.expression!(str)
 
-    assert {:ok, :a, _} = TypeChecker.infer_exp(%{}, ast)
+    assert {:ok, :a, _} = TypeChecker.infer_exp(%Env{}, ast)
   end
 
   test "infer type for list of atom expressions" do
@@ -35,7 +36,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:ok, %T.List{type: :a}, _} = TypeChecker.infer_exp(%{}, ast)
+    assert {:ok, %T.List{type: :a}, _} = TypeChecker.infer_exp(%Env{}, ast)
   end
 
   test "infer type of arithmetic expressions" do
@@ -43,7 +44,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:ok, :Int, _} = TypeChecker.infer_exp(%{}, ast)
+    assert {:ok, :Int, _} = TypeChecker.infer_exp(%Env{}, ast)
   end
 
   describe "logical operators" do
@@ -51,23 +52,23 @@ defmodule Fika.Compiler.TypeCheckerTest do
       # and
       str = "true & false"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
 
       # or
       str = "true | false"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
 
       # negation
       str = "!true"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "infer type of logical expressions when using atoms" do
       str = "true & :false"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -76,7 +77,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:error, "Unknown variable: foo"} = TypeChecker.infer_exp(%{}, ast)
+    assert {:error, "Unknown variable: foo"} = TypeChecker.infer_exp(%Env{}, ast)
   end
 
   test "infer ext function's return type" do
@@ -107,7 +108,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:ok, :Int} = TypeChecker.infer(function, %{})
+    assert {:ok, :Int} = TypeChecker.infer(function, %Env{})
   end
 
   test "check returns error when return type is not the inferred type" do
@@ -121,7 +122,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:error, "Expected type: Float, got: Int"} = TypeChecker.check(function, %{})
+    assert {:error, "Expected type: Float, got: Int"} = TypeChecker.check(function, %Env{})
   end
 
   test "checks tuple return type for function" do
@@ -135,7 +136,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:error, "Expected type: {Float}, got: {Int}"} = TypeChecker.check(function, %{})
+    assert {:error, "Expected type: {Float}, got: {Int}"} = TypeChecker.check(function, %Env{})
   end
 
   test "infer return type of another function in the module" do
@@ -196,8 +197,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:ok, :Float} = TypeChecker.infer(function, %{})
-    assert {:ok, :Float} = TypeChecker.check(function, %{})
+    assert {:ok, :Float} = TypeChecker.infer(function, %Env{})
+    assert {:ok, :Float} = TypeChecker.check(function, %Env{})
   end
 
   test "infer function calls considering union types" do
@@ -216,7 +217,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:ok, :ok} = TypeChecker.infer(function, %{})
+    assert {:ok, :ok} = TypeChecker.infer(function, %Env{})
   end
 
   describe "string" do
@@ -225,7 +226,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :String, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "parses string inside interpolation" do
@@ -235,7 +236,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :String, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "returns error when unknown variable in string interpolation" do
@@ -245,7 +246,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%{}, ast)
+      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "parses known variable in string interpolation" do
@@ -284,7 +285,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       assert {
                :error,
                "Expression used in string interpolation expected to be String, got Int"
-             } = TypeChecker.infer_exp(%{}, ast)
+             } = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -294,7 +295,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: :Int}, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.List{type: :Int}, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "list of integers and floats" do
@@ -303,7 +304,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:error, "Elements of list have different types. Expected: Int, got: Float"} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "list of floats inferred from fn calls" do
@@ -311,7 +312,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: :Float}, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.List{type: :Float}, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "List of strings" do
@@ -319,7 +320,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: :String}, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.List{type: :String}, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "List of list of integers" do
@@ -327,7 +328,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: %T.List{type: :Int}}, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.List{type: %T.List{type: :Int}}, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "empty list" do
@@ -335,7 +336,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: nil}, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.List{type: nil}, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -345,7 +346,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.Tuple{elements: [:Int, :Int, :Int]}, _env} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.Tuple{elements: [:Int, :Int, :Int]}, _env} =
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "tuple of integers and floats" do
@@ -354,7 +356,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:ok, %T.Tuple{elements: [:Int, :Float, :Int]}, _env} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "tuple of floats inferred from fn calls" do
@@ -362,7 +364,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.Tuple{elements: [:Float, :Float]}, _env} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.Tuple{elements: [:Float, :Float]}, _env} =
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "tuple of strings" do
@@ -370,7 +373,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.Tuple{elements: [:String, :String]}, _env} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.Tuple{elements: [:String, :String]}, _env} =
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "tuple of tuple of mixed types" do
@@ -384,7 +388,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
                   %T.Tuple{elements: [:Int, :Float]},
                   %T.Tuple{elements: [:String, :Bool]}
                 ]
-              }, _env} = TypeChecker.infer_exp(%{}, ast)
+              }, _env} = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -395,7 +399,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:ok, %T.Record{fields: [bar: :String, foo: :Int]}, _} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "error" do
@@ -403,7 +407,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%{}, ast)
+      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -414,7 +418,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:ok, %T.Map{key_type: :String, value_type: :Int}, _} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "type check for map with mixed type" do
@@ -423,14 +427,14 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:error, "Expected map key of type Int, but got String"} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
 
       str = ~s({"foo" => [1, 2], "bar" => 345})
 
       ast = TestParser.expression!(str)
 
       assert {:error, "Expected map value of type List(Int), but got Int"} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -449,7 +453,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
               %T.FunctionRef{
                 arg_types: [:Int, :Int],
                 return_type: :Int
-              }, _} = TypeChecker.infer_exp(%{}, ast)
+              }, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "without args" do
@@ -463,7 +467,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       CodeServer.set_type(signature("bar", "sum", []), {:ok, :Int})
 
       assert {:ok, %T.FunctionRef{arg_types: [], return_type: :Int}, _} =
-               TypeChecker.infer_exp(%{}, ast)
+               TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -472,14 +476,14 @@ defmodule Fika.Compiler.TypeCheckerTest do
       str = "true"
       ast = TestParser.expression!(str)
 
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "false" do
       str = "false"
       ast = TestParser.expression!(str)
 
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -498,7 +502,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       assert {
                :error,
                "Wrong type for if condition. Expected: Bool, Got: String"
-             } = TypeChecker.infer_exp(%{}, ast)
+             } = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "completes when if and else blocks have same return types" do
@@ -512,7 +516,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, _env} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :String, _env} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "completes when if and else blocks have different return types" do
@@ -526,7 +530,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
       types = MapSet.new([:String, :Int])
-      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "with multiple expressions in blocks" do
@@ -631,7 +635,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, :Bool} = TypeChecker.infer(function, %{})
+      assert {:ok, :Bool} = TypeChecker.infer(function, %Env{})
     end
 
     test "when function returns is expected to return a union type and has if-else clause" do
@@ -653,7 +657,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       types = MapSet.new([:ok, :error])
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.Union{types: ^types}} = TypeChecker.infer(function, %{})
+      assert {:ok, %T.Union{types: ^types}} = TypeChecker.infer(function, %Env{})
     end
 
     test "when function accepts union types and calls a function ref" do
@@ -694,7 +698,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       [function] = ast[:function_defs]
 
       assert {:error, "Expected a function reference, but got type: Int"} =
-               TypeChecker.infer(function, %{})
+               TypeChecker.infer(function, %Env{})
     end
 
     test "function ref when given wrong types" do
@@ -714,7 +718,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       error =
         "Expected function reference to be called with arguments (String, Int), but it was called with arguments (Int)"
 
-      assert {:error, ^error} = TypeChecker.infer(function, %{})
+      assert {:error, ^error} = TypeChecker.infer(function, %Env{})
     end
   end
 
@@ -724,7 +728,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, %{has_effect: true}} = TypeChecker.infer_exp(%{}, ast)
+      assert {:ok, :String, %{has_effect: true}} = TypeChecker.infer_exp(%Env{}, ast)
     end
 
     test "functions with effects inside them become effectful" do
@@ -739,8 +743,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.Effect{type: :String}} = TypeChecker.infer(function, %{})
-      assert {:ok, %T.Effect{type: :String}} = TypeChecker.check(function, %{})
+      assert {:ok, %T.Effect{type: :String}} = TypeChecker.infer(function, %Env{})
+      assert {:ok, %T.Effect{type: :String}} = TypeChecker.check(function, %Env{})
     end
 
     test "function with multiple effects" do
@@ -758,8 +762,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.infer(function, %{})
-      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.check(function, %{})
+      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.infer(function, %Env{})
+      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.check(function, %Env{})
     end
 
     test "effectful function ref call" do
@@ -797,7 +801,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
               %T.FunctionRef{
                 arg_types: [:Int, :Int],
                 return_type: :Int
-              }, _} = TypeChecker.infer_exp(%{}, ast)
+              }, _} = TypeChecker.infer_exp(%Env{}, ast)
     end
   end
 
@@ -813,7 +817,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.List{type: "b"}} = TypeChecker.infer(function, %{})
+      assert {:ok, %T.List{type: "b"}} = TypeChecker.infer(function, %Env{})
     end
 
     test "can't call incompatible functions on type variables" do
@@ -828,7 +832,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       [function] = ast[:function_defs]
 
       assert {:error, "Function fika/kernel.+(a, Int) does not exist"} =
-               TypeChecker.infer(function, %{})
+               TypeChecker.infer(function, %Env{})
     end
 
     test "infers return types of functions with type variables" do
@@ -846,7 +850,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [_foo, bar] = ast[:function_defs]
 
-      assert {:ok, :Int} = TypeChecker.infer(bar, %{ast: ast})
+      assert {:ok, :Int} = TypeChecker.infer(bar, %Env{ast: ast})
     end
 
     test "infers return types of functions when type variables are passed as args" do
@@ -867,7 +871,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [_foo, bar] = ast[:function_defs]
 
-      assert TypeChecker.infer(bar, %{ast: ast}) == {:ok, T.Union.new([:String, "z"])}
+      assert TypeChecker.infer(bar, %Env{ast: ast}) == {:ok, T.Union.new([:String, "z"])}
     end
 
     test "type variables used in function ref calls" do
@@ -892,7 +896,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [_foo, _bar, baz] = ast[:function_defs]
 
-      assert TypeChecker.infer(baz, %{ast: ast}) == {:ok, T.Union.new([:String, "c", :Int])}
+      assert TypeChecker.infer(baz, %Env{ast: ast}) == {:ok, T.Union.new([:String, "c", :Int])}
     end
   end
 

From 96139f2e2fb41111bcc9e64ffb974ce05d311001 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 19:57:27 -0300
Subject: [PATCH 04/10] feat: unwrap loops from inside unions

---
 lib/fika/compiler/type_checker.ex             | 45 +++++++++++++++----
 lib/fika/compiler/type_checker/types/loop.ex  |  5 +--
 lib/fika/compiler/type_checker/types/union.ex |  9 +++-
 test/fika/compiler/module_compiler_test.exs   | 10 ++---
 4 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index 70bbf72..8d64056 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -46,14 +46,41 @@ defmodule Fika.Compiler.TypeChecker do
         result
 
       {:ok, inferred_type} ->
-        {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
+        unwrap_type(env, expected_type, inferred_type, unwrap_union: true)
 
       error ->
         error
     end
   end
 
-  defp unwrap_loop(env, t, %T.Loop{type: t} = loop_type) do
+  defp unwrap_type(%Env{} = env, t, %T.Union{types: union_types}, unwrap_union: true) do
+    case Enum.split_with(union_types, &match?(%T.Loop{}, &1)) do
+      {[], union_types} ->
+        unwrap_type(env, t, %T.Union{types: union_types}, unwrap_union: false)
+
+      {loops, union_types} ->
+        left_union =
+          loops |> Enum.reject(&T.Loop.is_empty_loop/1) |> Enum.map(& &1.type) |> T.Union.new()
+
+        if Enum.empty?(union_types) do
+          unwrap_type(
+            env,
+            t,
+            %T.Loop{is_empty_loop: false, type: T.Union.new(left_union)},
+            unwrap_union: false
+          )
+        else
+          unwrap_type(
+            env,
+            t,
+            %T.Loop{is_empty_loop: false, type: T.Union.new([left_union, union_types])},
+            unwrap_union: false
+          )
+        end
+    end
+  end
+
+  defp unwrap_type(%Env{} = env, t, %T.Loop{type: t} = loop_type, _) do
     # The function can be a top-level function which depends on a loop
     # In this case, we can unwrap the loop
 
@@ -66,10 +93,14 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
-  defp unwrap_loop(_env, t, t), do: {:ok, t}
+  defp unwrap_type(_env, %T.Loop{type: t}, %T.Loop{type: t}, _),
+    do: {:ok, %T.Loop{type: t, is_empty_loop: false}}
+
+  defp unwrap_type(_env, t, t, _), do: {:ok, t}
 
-  defp unwrap_loop(_env, expected_type, inferred_type),
-    do: {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
+  defp unwrap_type(_env, expected_type, inferred_type, _) do
+    {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
+  end
 
   # Given the AST of a function definition, this function infers the
   # return type of the body of the function.
@@ -554,10 +585,6 @@ defmodule Fika.Compiler.TypeChecker do
     end)
   end
 
-  defp set_latest_call(env, signature) do
-    Map.put(env, :latest_called_function, signature)
-  end
-
   defp get_function_signature(module, function_name, arg_types) do
     %FunctionSignature{module: module, function: to_string(function_name), types: arg_types}
   end
diff --git a/lib/fika/compiler/type_checker/types/loop.ex b/lib/fika/compiler/type_checker/types/loop.ex
index c63cbfc..361ae83 100644
--- a/lib/fika/compiler/type_checker/types/loop.ex
+++ b/lib/fika/compiler/type_checker/types/loop.ex
@@ -2,7 +2,7 @@ defmodule Fika.Compiler.TypeChecker.Types.Loop do
   @moduledoc """
   Defines when a function may not terminate due to recursion
   """
-  defstruct [:type]
+  defstruct [:type, is_empty_loop: true]
 
   alias Fika.Compiler.TypeChecker.Types, as: T
 
@@ -14,8 +14,7 @@ defmodule Fika.Compiler.TypeChecker.Types.Loop do
   def is_loop(%__MODULE__{}), do: true
   def is_loop(_), do: false
 
-  def is_empty_loop(%__MODULE__{type: nil}), do: true
-  def is_empty_loop(_), do: false
+  def is_empty_loop(%__MODULE__{is_empty_loop: is_empty_loop}), do: is_empty_loop == true
 
   defimpl String.Chars, for: T.Loop do
     def to_string(%{type: type}) do
diff --git a/lib/fika/compiler/type_checker/types/union.ex b/lib/fika/compiler/type_checker/types/union.ex
index db60f88..4b91551 100644
--- a/lib/fika/compiler/type_checker/types/union.ex
+++ b/lib/fika/compiler/type_checker/types/union.ex
@@ -5,7 +5,13 @@ defmodule Fika.Compiler.TypeChecker.Types.Union do
 
   @spec new(types :: Enumerable.t()) :: t()
   def new(types) do
-    %__MODULE__{types: flatten_types(types)}
+    union = %__MODULE__{types: flatten_types(types)}
+
+    if MapSet.size(union.types) == 1 do
+      Enum.at(union.types, 0)
+    else
+      union
+    end
   end
 
   @spec flatten_types(types :: Enumerable.t()) :: MapSet.t()
@@ -19,6 +25,7 @@ defmodule Fika.Compiler.TypeChecker.Types.Union do
       t ->
         [t]
     end)
+    |> List.flatten()
     |> MapSet.new()
   end
 
diff --git a/test/fika/compiler/module_compiler_test.exs b/test/fika/compiler/module_compiler_test.exs
index 0c58b1e..3723dc9 100644
--- a/test/fika/compiler/module_compiler_test.exs
+++ b/test/fika/compiler/module_compiler_test.exs
@@ -167,27 +167,27 @@ defmodule Fika.Compiler.ModuleCompilerTest do
     temp_file = "#{module}.fi"
 
     str = """
-    fn factorial(x: Int) : Int | Loop(:nil) do
+    fn factorial(x: Int) : Loop(Int) do
       do_factorial(x, 0)
     end
-    fn do_factorial(x: Int, acc: Int) : Int | Loop(:nil) do
+    fn do_factorial(x: Int, acc: Int) : Loop(Int) do
       if x <= 1 do
         acc
       else
         do_factorial(x - 1, acc * x)
       end
     end
-    fn top_level(x: Int) : Int | Loop(:nil) do
+    fn top_level(x: Int) : Loop(Int) do
       second_level_a(x)
     end
-    fn second_level_a(x: Int) : Int | Loop(:nil) do
+    fn second_level_a(x: Int) : Loop(Int) do
       if x > 1 do
         second_level_b(x - 1)
       else
         1
       end
     end
-    fn second_level_b(x: Int) : Int | Loop(:nil) do
+    fn second_level_b(x: Int) : Loop(Int) do
       second_level_a(x - 1)
     end
     """

From f9f2d9a4352b337f7d930af1b599e8d9a0f78b64 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 19:59:41 -0300
Subject: [PATCH 05/10] feat: ensure top level functions are not inferred as
 loops

---
 test/fika/compiler/module_compiler_test.exs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/test/fika/compiler/module_compiler_test.exs b/test/fika/compiler/module_compiler_test.exs
index 3723dc9..7e4cc60 100644
--- a/test/fika/compiler/module_compiler_test.exs
+++ b/test/fika/compiler/module_compiler_test.exs
@@ -167,9 +167,10 @@ defmodule Fika.Compiler.ModuleCompilerTest do
     temp_file = "#{module}.fi"
 
     str = """
-    fn factorial(x: Int) : Loop(Int) do
+    fn factorial(x: Int) : Int do
       do_factorial(x, 0)
     end
+
     fn do_factorial(x: Int, acc: Int) : Loop(Int) do
       if x <= 1 do
         acc
@@ -177,7 +178,8 @@ defmodule Fika.Compiler.ModuleCompilerTest do
         do_factorial(x - 1, acc * x)
       end
     end
-    fn top_level(x: Int) : Loop(Int) do
+
+    fn top_level(x: Int) : Int do
       second_level_a(x)
     end
     fn second_level_a(x: Int) : Loop(Int) do
@@ -196,6 +198,8 @@ defmodule Fika.Compiler.ModuleCompilerTest do
 
     File.cd!(tmp_dir, fn ->
       assert {:ok, "foo", "foo.fi", _binary} = ModuleCompiler.compile("foo")
+
+      Fika.Compiler.CodeServer.list_all_types() |> IO.inspect()
     end)
 
     File.rm!(temp_file)

From eeb495aca66ef610f53cfefab80e9ceb8eac0c67 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 20:52:01 -0300
Subject: [PATCH 06/10] wip: make type checker tests work for inference as well
 as checking

---
 config/test.exs                              |   2 +
 lib/fika/compiler/type_checker.ex            | 152 ++++++++++---------
 lib/fika/compiler/type_checker/types/loop.ex |   3 +
 test/fika/compiler/module_compiler_test.exs  |   3 +-
 test/fika/compiler/type_checker_test.exs     |  75 +++++++++
 test/test_helper.exs                         |   1 -
 6 files changed, 166 insertions(+), 70 deletions(-)

diff --git a/config/test.exs b/config/test.exs
index 29ad4ac..5508554 100644
--- a/config/test.exs
+++ b/config/test.exs
@@ -2,3 +2,5 @@ import Config
 
 config :fika,
   disable_web_server: true
+
+config :logger, level: :debug
diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index 8d64056..a028665 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -34,6 +34,36 @@ defmodule Fika.Compiler.TypeChecker do
     def update_scope(env, variable, update_fn) do
       %{env | scope: update_in(env.scope, [variable], update_fn)}
     end
+
+    def reset_scope_and_set_signature(
+          %__MODULE__{} = env,
+          {:function, _, {_name, args, _type, _}} = ast
+        ) do
+      current_signature = Fika.Compiler.TypeChecker.function_ast_signature(env.module, ast)
+
+      %{
+        env
+        | scope: %{},
+          has_effect: false,
+          current_signature: current_signature,
+          latest_called_function: env.current_signature
+      }
+      |> add_args_to_scope(args)
+    end
+
+    def reset_scope_and_set_signature(env, {:anonymous_function, _line, args, _}) do
+      %{env | scope: %{}}
+      |> add_args_to_scope(args)
+    end
+
+    def reset_scope_and_set_signature(env, _ast), do: env
+
+    defp add_args_to_scope(env, args) do
+      Enum.reduce(args, env, fn {{:identifier, _, name}, {:type, _, type}}, env ->
+        Logger.debug("Adding arg type to scope: #{name}:#{type}")
+        Env.add_variable_to_scope(env, name, type)
+      end)
+    end
   end
 
   # Given the AST of a function definition, this function checks if the return
@@ -46,85 +76,81 @@ defmodule Fika.Compiler.TypeChecker do
         result
 
       {:ok, inferred_type} ->
-        unwrap_type(env, expected_type, inferred_type, unwrap_union: true)
+        unwrap_type(env, expected_type, inferred_type)
 
       error ->
         error
     end
   end
 
-  defp unwrap_type(%Env{} = env, t, %T.Union{types: union_types}, unwrap_union: true) do
+  defp unwrap_type(
+         %Env{current_signature: current_signature},
+         wrapped_expected_type,
+         wrapped_inferred_type
+       ) do
+    is_top_level_function = CodeServer.check_cycle(current_signature) == :ok
+
+    expected_type = do_unwrap_type(wrapped_expected_type)
+    inferred_type = do_unwrap_type(wrapped_inferred_type)
+
+    cond do
+      is_top_level_function and T.Loop.is_loop(expected_type) ->
+        {:error, "Top level function cannot be a loop"}
+
+      is_top_level_function and T.Loop.is_loop(inferred_type) and
+          inferred_type.type == expected_type ->
+        {:ok, expected_type}
+
+      T.Loop.equals?(expected_type, inferred_type) ->
+        {:ok, %{expected_type | is_empty_loop: false}}
+
+      match?(^expected_type, inferred_type) ->
+        {:ok, expected_type}
+
+      true ->
+        {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
+    end
+  end
+
+  defp do_unwrap_type(%T.Union{types: union_types}) do
     case Enum.split_with(union_types, &match?(%T.Loop{}, &1)) do
       {[], union_types} ->
-        unwrap_type(env, t, %T.Union{types: union_types}, unwrap_union: false)
+        %T.Union{types: union_types}
 
       {loops, union_types} ->
         left_union =
           loops |> Enum.reject(&T.Loop.is_empty_loop/1) |> Enum.map(& &1.type) |> T.Union.new()
 
         if Enum.empty?(union_types) do
-          unwrap_type(
-            env,
-            t,
-            %T.Loop{is_empty_loop: false, type: T.Union.new(left_union)},
-            unwrap_union: false
-          )
+          %T.Loop{is_empty_loop: false, type: T.Union.new(left_union)}
         else
-          unwrap_type(
-            env,
-            t,
-            %T.Loop{is_empty_loop: false, type: T.Union.new([left_union, union_types])},
-            unwrap_union: false
-          )
+          %T.Loop{is_empty_loop: false, type: T.Union.new([left_union, union_types])}
         end
     end
   end
 
-  defp unwrap_type(%Env{} = env, t, %T.Loop{type: t} = loop_type, _) do
-    # The function can be a top-level function which depends on a loop
-    # In this case, we can unwrap the loop
-
-    case CodeServer.check_cycle(env.current_signature) do
-      :ok ->
-        {:ok, t}
-
-      _ ->
-        {:error, "Expected type: #{t}, got: #{loop_type}"}
-    end
-  end
-
-  defp unwrap_type(_env, %T.Loop{type: t}, %T.Loop{type: t}, _),
-    do: {:ok, %T.Loop{type: t, is_empty_loop: false}}
-
-  defp unwrap_type(_env, t, t, _), do: {:ok, t}
-
-  defp unwrap_type(_env, expected_type, inferred_type, _) do
-    {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
-  end
+  defp do_unwrap_type(t), do: t
 
   # Given the AST of a function definition, this function infers the
   # return type of the body of the function.
-  def infer({:function, _line, {name, args, _type, exprs}}, %Env{} = env) do
-    Logger.debug("Inferring type of function: #{name}")
-
-    Logger.debug("First call to #{inspect(env.current_signature)}")
-
-    env =
-      env
-      |> Map.put(:scope, %{})
-      |> Map.put(:has_effect, false)
-      |> add_args_to_scope(args)
-
-    case infer_block(env, exprs) do
-      {:ok, type, env} ->
-        type =
-          if env.has_effect do
-            %T.Effect{type: type}
-          else
-            type
-          end
+  def infer({:function, _line, {_name, _args, _type, exprs}} = ast, %Env{} = env) do
+    env = Env.reset_scope_and_set_signature(env, ast)
+    Logger.debug("Inferring type of function: #{env.current_signature}")
+
+    with :ok <-
+           CodeServer.set_function_dependency(env.latest_called_function, env.current_signature),
+         {:ok, type, env} <- infer_block(env, exprs) do
+      type =
+        if env.has_effect do
+          %T.Effect{type: type}
+        else
+          type
+        end
 
-        {:ok, type}
+      {:ok, type}
+    else
+      {:error, :cycle_encountered} ->
+        {:ok, T.Loop.new()}
 
       error ->
         error
@@ -360,13 +386,10 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # anonymous function
-  def infer_exp(%Env{} = env, {:anonymous_function, _line, args, exps}) do
+  def infer_exp(%Env{} = env, {:anonymous_function, _line, args, exps} = ast) do
     Logger.debug("Inferring type of anonymous function")
 
-    env =
-      env
-      |> Map.put(:scope, %{})
-      |> add_args_to_scope(args)
+    env = Env.reset_scope_and_set_signature(env, ast)
 
     case infer_block(%Env{} = env, exps) do
       {:ok, return_type, _env} ->
@@ -578,13 +601,6 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
-  defp add_args_to_scope(env, args) do
-    Enum.reduce(args, env, fn {{:identifier, _, name}, {:type, _, type}}, env ->
-      Logger.debug("Adding arg type to scope: #{name}:#{type}")
-      Env.add_variable_to_scope(env, name, type)
-    end)
-  end
-
   defp get_function_signature(module, function_name, arg_types) do
     %FunctionSignature{module: module, function: to_string(function_name), types: arg_types}
   end
diff --git a/lib/fika/compiler/type_checker/types/loop.ex b/lib/fika/compiler/type_checker/types/loop.ex
index 361ae83..be95f2c 100644
--- a/lib/fika/compiler/type_checker/types/loop.ex
+++ b/lib/fika/compiler/type_checker/types/loop.ex
@@ -14,6 +14,9 @@ defmodule Fika.Compiler.TypeChecker.Types.Loop do
   def is_loop(%__MODULE__{}), do: true
   def is_loop(_), do: false
 
+  def equals?(%__MODULE__{type: t}, %__MODULE__{type: t}), do: true
+  def equals?(_, _), do: false
+
   def is_empty_loop(%__MODULE__{is_empty_loop: is_empty_loop}), do: is_empty_loop == true
 
   defimpl String.Chars, for: T.Loop do
diff --git a/test/fika/compiler/module_compiler_test.exs b/test/fika/compiler/module_compiler_test.exs
index 7e4cc60..a723de7 100644
--- a/test/fika/compiler/module_compiler_test.exs
+++ b/test/fika/compiler/module_compiler_test.exs
@@ -154,6 +154,7 @@ defmodule Fika.Compiler.ModuleCompilerTest do
 
     File.cd!(System.tmp_dir!(), fn ->
       assert {:ok, _, _, _} = ModuleCompiler.compile("foo")
+      assert {:ok, _, _, _} = ModuleCompiler.compile("bar")
     end)
 
     File.rm!(foo_temp_file)
@@ -199,7 +200,7 @@ defmodule Fika.Compiler.ModuleCompilerTest do
     File.cd!(tmp_dir, fn ->
       assert {:ok, "foo", "foo.fi", _binary} = ModuleCompiler.compile("foo")
 
-      Fika.Compiler.CodeServer.list_all_types() |> IO.inspect()
+      Fika.Compiler.CodeServer.list_all_types()
     end)
 
     File.rm!(temp_file)
diff --git a/test/fika/compiler/type_checker_test.exs b/test/fika/compiler/type_checker_test.exs
index aea72b9..3f003a2 100644
--- a/test/fika/compiler/type_checker_test.exs
+++ b/test/fika/compiler/type_checker_test.exs
@@ -158,6 +158,81 @@ defmodule Fika.Compiler.TypeCheckerTest do
     assert {:ok, :Int} = TypeChecker.infer(foo, env)
   end
 
+  test "infer return type when there is simple recursion" do
+    str = """
+    fn factorial(x: Int, acc: Int) : Loop(Int) do
+      if x <= 1 do
+        acc
+      else
+        factorial(x - 1, acc * x)
+      end
+    end
+    """
+
+    {:ok, ast} = Parser.parse_module(str)
+    [factorial] = ast[:function_defs]
+
+    env = TypeChecker.init_env(ast)
+
+    assert {:ok, %T.Loop{type: :Int}} = TypeChecker.infer(factorial, env)
+  end
+
+  test "infer return type when there is intra-module recursion" do
+    str = """
+      fn foo(a: Int) : :nil do
+        bar(a)
+      end
+      fn bar(a: Int) : :nil do
+        if a <= 1 do
+          foobar()
+        else
+          baz(a - 1)
+        end
+      end
+      fn baz(a: Int) : :nil do
+        if a == 1 do
+          1
+        else
+          foo(a - 1)
+        end
+      end
+      fn foobar : :nil do
+        :nil
+      end
+    """
+
+    {:ok, ast} = Parser.parse_module(str)
+    [foo, bar, baz, foobar] = ast[:function_defs]
+
+    env = TypeChecker.init_env(ast)
+
+    assert {:ok, nil} = TypeChecker.infer(foobar, env)
+
+    types = MapSet.new([:Int, nil])
+    assert {:ok, %T.Loop{type: %T.Union{types: ^types}}} = TypeChecker.infer(foo, env)
+    assert {:ok, %T.Loop{type: %T.Union{types: ^types}}} = TypeChecker.infer(bar, env)
+    assert {:ok, %T.Loop{type: %T.Union{types: ^types}}} = TypeChecker.infer(baz, env)
+  end
+
+  test "infer empty-function recursion" do
+    str = """
+    fn foo : Loop(:nil) do
+      bar()
+    end
+    fn bar : Loop(:nil) do
+      foo()
+    end
+    """
+
+    {:ok, ast} = Parser.parse_module(str) |> IO.inspect(label: "parsed")
+    [foo, bar] = ast[:function_defs]
+
+    env = TypeChecker.init_env(ast) |> IO.inspect(label: "env")
+
+    assert {:ok, %T.Loop{type: nil}} = TypeChecker.infer(foo, env) |> IO.inspect(label: "tc foo")
+    assert {:ok, %T.Loop{type: nil}} = TypeChecker.infer(bar, env) |> IO.inspect(label: "tc bar")
+  end
+
   test "infer calls with multiple args" do
     str = """
     fn foo(a: Int, b: String) : Int do
diff --git a/test/test_helper.exs b/test/test_helper.exs
index e2cf047..869559e 100644
--- a/test/test_helper.exs
+++ b/test/test_helper.exs
@@ -1,2 +1 @@
-Logger.configure(level: :error)
 ExUnit.start()

From c96607432ce7fe06b77d2dfed1df8b4bf6681879 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sat, 6 Mar 2021 23:25:10 -0300
Subject: [PATCH 07/10] fix cycle-finding

---
 config/test.exs                                     |  2 +-
 .../compiler/code_server/function_dependencies.ex   |  8 +++++---
 lib/fika/compiler/type_checker.ex                   | 13 +++++++++----
 test/fika/compiler/type_checker_test.exs            |  8 ++++----
 4 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/config/test.exs b/config/test.exs
index 5508554..cd2fbfe 100644
--- a/config/test.exs
+++ b/config/test.exs
@@ -3,4 +3,4 @@ import Config
 config :fika,
   disable_web_server: true
 
-config :logger, level: :debug
+config :logger, level: :error
diff --git a/lib/fika/compiler/code_server/function_dependencies.ex b/lib/fika/compiler/code_server/function_dependencies.ex
index 9567571..17959d0 100644
--- a/lib/fika/compiler/code_server/function_dependencies.ex
+++ b/lib/fika/compiler/code_server/function_dependencies.ex
@@ -13,11 +13,13 @@ defmodule Fika.Compiler.CodeServer.FunctionDependencies do
     |> add_vertex(source_function)
     |> add_vertex(sink_function)
     |> add_edge(source_function, sink_function)
-    |> check_cycle(source_function)
+    |> check_cycle(source_function, sink_function)
   end
 
-  def check_cycle(graph, function) do
-    if :digraph.get_cycle(graph, function) do
+  def check_cycle(graph, source_function, sink_function) do
+    vertices = :digraph.get_cycle(graph, source_function)
+
+    if vertices && sink_function in vertices do
       {:error, :cycle_encountered}
     else
       :ok
diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index a028665..3dae9b2 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -24,6 +24,7 @@ defmodule Fika.Compiler.TypeChecker do
       :file,
       :current_signature,
       has_effect: false,
+      first_pass: false,
       scope: %{}
     ]
 
@@ -46,7 +47,8 @@ defmodule Fika.Compiler.TypeChecker do
         | scope: %{},
           has_effect: false,
           current_signature: current_signature,
-          latest_called_function: env.current_signature
+          latest_called_function: if(env.first_pass, do: env.current_signature),
+          first_pass: false
       }
       |> add_args_to_scope(args)
     end
@@ -115,7 +117,7 @@ defmodule Fika.Compiler.TypeChecker do
   defp do_unwrap_type(%T.Union{types: union_types}) do
     case Enum.split_with(union_types, &match?(%T.Loop{}, &1)) do
       {[], union_types} ->
-        %T.Union{types: union_types}
+        T.Union.new(union_types)
 
       {loops, union_types} ->
         left_union =
@@ -129,6 +131,7 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
+  defp do_unwrap_type(%T.Effect{type: t}), do: %T.Effect{type: do_unwrap_type(t)}
   defp do_unwrap_type(t), do: t
 
   # Given the AST of a function definition, this function infers the
@@ -139,7 +142,8 @@ defmodule Fika.Compiler.TypeChecker do
 
     with :ok <-
            CodeServer.set_function_dependency(env.latest_called_function, env.current_signature),
-         {:ok, type, env} <- infer_block(env, exprs) do
+         {:ok, type, env} <-
+           infer_block(env, exprs) do
       type =
         if env.has_effect do
           %T.Effect{type: type}
@@ -147,7 +151,7 @@ defmodule Fika.Compiler.TypeChecker do
           type
         end
 
-      {:ok, type}
+      {:ok, do_unwrap_type(type)}
     else
       {:error, :cycle_encountered} ->
         {:ok, T.Loop.new()}
@@ -512,6 +516,7 @@ defmodule Fika.Compiler.TypeChecker do
           {:ok, %T.Effect{type: type}} ->
             env = Map.put(env, :has_effect, true)
             {:ok, type, env}
+            1
 
           {:ok, type} ->
             {:ok, type, env}
diff --git a/test/fika/compiler/type_checker_test.exs b/test/fika/compiler/type_checker_test.exs
index 3f003a2..611f5fc 100644
--- a/test/fika/compiler/type_checker_test.exs
+++ b/test/fika/compiler/type_checker_test.exs
@@ -224,13 +224,13 @@ defmodule Fika.Compiler.TypeCheckerTest do
     end
     """
 
-    {:ok, ast} = Parser.parse_module(str) |> IO.inspect(label: "parsed")
+    {:ok, ast} = Parser.parse_module(str)
     [foo, bar] = ast[:function_defs]
 
-    env = TypeChecker.init_env(ast) |> IO.inspect(label: "env")
+    env = TypeChecker.init_env(ast)
 
-    assert {:ok, %T.Loop{type: nil}} = TypeChecker.infer(foo, env) |> IO.inspect(label: "tc foo")
-    assert {:ok, %T.Loop{type: nil}} = TypeChecker.infer(bar, env) |> IO.inspect(label: "tc bar")
+    assert {:ok, %T.Loop{type: nil}} = TypeChecker.infer(foo, env)
+    assert {:ok, %T.Loop{type: nil}} = TypeChecker.infer(bar, env)
   end
 
   test "infer calls with multiple args" do

From 4ca3b1d16929662208adedb92fb84975f33c9260 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Sun, 7 Mar 2021 00:12:44 -0300
Subject: [PATCH 08/10] fix: apply some fixes

---
 lib/fika/compiler/code_server.ex              | 21 ++++++++++++
 .../code_server/function_dependencies.ex      | 28 ++++++++++++----
 lib/fika/compiler/type_checker.ex             | 32 +++++++++++++++----
 test/fika/compiler/type_checker_test.exs      |  6 ++--
 4 files changed, 72 insertions(+), 15 deletions(-)

diff --git a/lib/fika/compiler/code_server.ex b/lib/fika/compiler/code_server.ex
index e551fc9..58ee6ed 100644
--- a/lib/fika/compiler/code_server.ex
+++ b/lib/fika/compiler/code_server.ex
@@ -82,6 +82,12 @@ defmodule Fika.Compiler.CodeServer do
     GenServer.call(__MODULE__, {:check_cycle, node})
   end
 
+  @spec get_cycle(source :: String.t() | nil) ::
+          {:ok, list()} | {:error, :no_cycle_found}
+  def get_cycle(node) do
+    GenServer.call(__MODULE__, {:get_cycle, node})
+  end
+
   def init(_) do
     state = init_state()
 
@@ -191,6 +197,21 @@ defmodule Fika.Compiler.CodeServer do
     {:reply, response, state}
   end
 
+  def handle_call({:get_cycle, node}, _from, state)
+      when is_nil(node) do
+    {:reply, {:error, :no_cycle_found}, state}
+  end
+
+  def handle_call(
+        {:get_cycle, node},
+        _from,
+        %{function_dependencies: graph} = state
+      ) do
+    response = __MODULE__.FunctionDependencies.get_cycle(graph, node)
+
+    {:reply, response, state}
+  end
+
   def handle_call({:check_cycle, node}, _from, state)
       when is_nil(node) do
     {:reply, :ok, state}
diff --git a/lib/fika/compiler/code_server/function_dependencies.ex b/lib/fika/compiler/code_server/function_dependencies.ex
index 17959d0..b3349d8 100644
--- a/lib/fika/compiler/code_server/function_dependencies.ex
+++ b/lib/fika/compiler/code_server/function_dependencies.ex
@@ -7,7 +7,7 @@ defmodule Fika.Compiler.CodeServer.FunctionDependencies do
           graph :: :digraph.graph(),
           source_function :: FunctionSignature.t(),
           sink_function :: FunctionSignature.t()
-        ) :: :digraph.graph()
+        ) :: :ok | {:error, :cycle_encountered}
   def set_function_dependency(graph, source_function, sink_function) do
     graph
     |> add_vertex(source_function)
@@ -16,13 +16,29 @@ defmodule Fika.Compiler.CodeServer.FunctionDependencies do
     |> check_cycle(source_function, sink_function)
   end
 
-  def check_cycle(graph, source_function, sink_function) do
+  def get_cycle(graph, function) do
+    if vertices = :digraph.get_cycle(graph, function) do
+      {:ok, vertices}
+    else
+      {:error, :no_cycle_found}
+    end
+  end
+
+  def check_cycle(graph, source_function, sink_function \\ nil) do
     vertices = :digraph.get_cycle(graph, source_function)
 
-    if vertices && sink_function in vertices do
-      {:error, :cycle_encountered}
-    else
-      :ok
+    cond do
+      !vertices ->
+        :ok
+
+      is_nil(sink_function) ->
+        {:error, :cycle_encountered}
+
+      sink_function in vertices ->
+        {:error, :cycle_encountered}
+
+      true ->
+        :ok
     end
   end
 
diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index 3dae9b2..fd29beb 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -90,10 +90,23 @@ defmodule Fika.Compiler.TypeChecker do
          wrapped_expected_type,
          wrapped_inferred_type
        ) do
-    is_top_level_function = CodeServer.check_cycle(current_signature) == :ok
+    {is_top_level_function, cycle} =
+      case CodeServer.get_cycle(current_signature) do
+        {:ok, cycle} ->
+          {false, cycle}
+
+        _ ->
+          {true, []}
+      end
 
     expected_type = do_unwrap_type(wrapped_expected_type)
-    inferred_type = do_unwrap_type(wrapped_inferred_type)
+
+    types =
+      cycle
+      |> Enum.map(&CodeServer.get_type/1)
+      |> Keyword.get_values(:ok)
+
+    inferred_type = [wrapped_inferred_type | types] |> T.Union.new() |> do_unwrap_type()
 
     cond do
       is_top_level_function and T.Loop.is_loop(expected_type) ->
@@ -124,7 +137,7 @@ defmodule Fika.Compiler.TypeChecker do
           loops |> Enum.reject(&T.Loop.is_empty_loop/1) |> Enum.map(& &1.type) |> T.Union.new()
 
         if Enum.empty?(union_types) do
-          %T.Loop{is_empty_loop: false, type: T.Union.new(left_union)}
+          %T.Loop{is_empty_loop: false, type: T.Union.new([left_union])}
         else
           %T.Loop{is_empty_loop: false, type: T.Union.new([left_union, union_types])}
         end
@@ -142,6 +155,8 @@ defmodule Fika.Compiler.TypeChecker do
 
     with :ok <-
            CodeServer.set_function_dependency(env.latest_called_function, env.current_signature),
+         {:cycle, {:error, :no_cycle_found}} <-
+           {:cycle, CodeServer.get_cycle(env.current_signature)},
          {:ok, type, env} <-
            infer_block(env, exprs) do
       type =
@@ -153,8 +168,14 @@ defmodule Fika.Compiler.TypeChecker do
 
       {:ok, do_unwrap_type(type)}
     else
-      {:error, :cycle_encountered} ->
-        {:ok, T.Loop.new()}
+      {:cycle, {:ok, cycle}} ->
+        types =
+          cycle
+          |> Enum.map(&CodeServer.get_type/1)
+          |> Keyword.get_values(:ok)
+
+        result = [T.Loop.new() | types] |> T.Union.new() |> do_unwrap_type()
+        {:ok, result}
 
       error ->
         error
@@ -516,7 +537,6 @@ defmodule Fika.Compiler.TypeChecker do
           {:ok, %T.Effect{type: type}} ->
             env = Map.put(env, :has_effect, true)
             {:ok, type, env}
-            1
 
           {:ok, type} ->
             {:ok, type, env}
diff --git a/test/fika/compiler/type_checker_test.exs b/test/fika/compiler/type_checker_test.exs
index 611f5fc..d619308 100644
--- a/test/fika/compiler/type_checker_test.exs
+++ b/test/fika/compiler/type_checker_test.exs
@@ -179,17 +179,17 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
   test "infer return type when there is intra-module recursion" do
     str = """
-      fn foo(a: Int) : :nil do
+      fn foo(a: Int) : Int | :nil do
         bar(a)
       end
-      fn bar(a: Int) : :nil do
+      fn bar(a: Int) : Loop(Int | :nil) do
         if a <= 1 do
           foobar()
         else
           baz(a - 1)
         end
       end
-      fn baz(a: Int) : :nil do
+      fn baz(a: Int) : Loop(Int | :nil) do
         if a == 1 do
           1
         else

From 736fc70e1145292455f4557ff208c95557ce25cd Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Mon, 8 Mar 2021 00:35:53 -0300
Subject: [PATCH 09/10] chore: revert env struct

---
 lib/fika/compiler/module_compiler.ex          |   3 +-
 lib/fika/compiler/type_checker.ex             | 205 +++++++++---------
 lib/fika/compiler/type_checker/match.ex       |   2 +-
 .../type_checker/parallel_type_checker.ex     |  13 +-
 test/fika/compiler/type_checker_test.exs      | 117 +++++-----
 5 files changed, 166 insertions(+), 174 deletions(-)

diff --git a/lib/fika/compiler/module_compiler.ex b/lib/fika/compiler/module_compiler.ex
index 9b1f03c..fdced84 100644
--- a/lib/fika/compiler/module_compiler.ex
+++ b/lib/fika/compiler/module_compiler.ex
@@ -4,7 +4,6 @@ defmodule Fika.Compiler.ModuleCompiler do
   alias Fika.Compiler.{
     Parser,
     TypeChecker.ParallelTypeChecker,
-    TypeChecker.Env,
     ErlTranslate,
     CodeServer
   }
@@ -44,7 +43,7 @@ defmodule Fika.Compiler.ModuleCompiler do
   end
 
   defp init(module_atom) do
-    %Env{
+    %{
       # Full name of the current module as atom
       module_name: module_atom,
       # Path of module file
diff --git a/lib/fika/compiler/type_checker.ex b/lib/fika/compiler/type_checker.ex
index fd29beb..fdbe2d7 100644
--- a/lib/fika/compiler/type_checker.ex
+++ b/lib/fika/compiler/type_checker.ex
@@ -14,63 +14,9 @@ defmodule Fika.Compiler.TypeChecker do
 
   require Logger
 
-  defmodule Env do
-    defstruct [
-      :ast,
-      :latest_called_function,
-      :type_checker_pid,
-      :module,
-      :module_name,
-      :file,
-      :current_signature,
-      has_effect: false,
-      first_pass: false,
-      scope: %{}
-    ]
-
-    def add_variable_to_scope(env, variable, type) do
-      %{env | scope: put_in(env.scope, [variable], type)}
-    end
-
-    def update_scope(env, variable, update_fn) do
-      %{env | scope: update_in(env.scope, [variable], update_fn)}
-    end
-
-    def reset_scope_and_set_signature(
-          %__MODULE__{} = env,
-          {:function, _, {_name, args, _type, _}} = ast
-        ) do
-      current_signature = Fika.Compiler.TypeChecker.function_ast_signature(env.module, ast)
-
-      %{
-        env
-        | scope: %{},
-          has_effect: false,
-          current_signature: current_signature,
-          latest_called_function: if(env.first_pass, do: env.current_signature),
-          first_pass: false
-      }
-      |> add_args_to_scope(args)
-    end
-
-    def reset_scope_and_set_signature(env, {:anonymous_function, _line, args, _}) do
-      %{env | scope: %{}}
-      |> add_args_to_scope(args)
-    end
-
-    def reset_scope_and_set_signature(env, _ast), do: env
-
-    defp add_args_to_scope(env, args) do
-      Enum.reduce(args, env, fn {{:identifier, _, name}, {:type, _, type}}, env ->
-        Logger.debug("Adding arg type to scope: #{name}:#{type}")
-        Env.add_variable_to_scope(env, name, type)
-      end)
-    end
-  end
-
   # Given the AST of a function definition, this function checks if the return
   # type is indeed the type that's inferred from the body of the function.
-  def check({:function, _line, {_, _, return_type, _}} = function, %Env{} = env) do
+  def check({:function, _line, {_, _, return_type, _}} = function, env) do
     {:type, _line, expected_type} = return_type
 
     case infer(function, env) do
@@ -86,10 +32,12 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   defp unwrap_type(
-         %Env{current_signature: current_signature},
+         env,
          wrapped_expected_type,
          wrapped_inferred_type
        ) do
+    current_signature = env[:current_signature]
+
     {is_top_level_function, cycle} =
       case CodeServer.get_cycle(current_signature) do
         {:ok, cycle} ->
@@ -149,8 +97,8 @@ defmodule Fika.Compiler.TypeChecker do
 
   # Given the AST of a function definition, this function infers the
   # return type of the body of the function.
-  def infer({:function, _line, {_name, _args, _type, exprs}} = ast, %Env{} = env) do
-    env = Env.reset_scope_and_set_signature(env, ast)
+  def infer({:function, _line, {_name, _args, _type, exprs}} = ast, env) do
+    env = reset_scope_and_set_signature(env, ast)
     Logger.debug("Inferring type of function: #{env.current_signature}")
 
     with :ok <-
@@ -182,37 +130,37 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
-  def infer_block(%Env{} = env, []) do
+  def infer_block(env, []) do
     Logger.debug("Block is empty.")
     {:ok, nil, env}
   end
 
-  def infer_block(%Env{} = env, [exp]) do
-    infer_exp(%Env{} = env, exp)
+  def infer_block(env, [exp]) do
+    infer_exp(env, exp)
   end
 
-  def infer_block(%Env{} = env, [exp | exp_list]) do
-    case infer_exp(%Env{} = env, exp) do
-      {:ok, _type, env} -> infer_block(%Env{} = env, exp_list)
+  def infer_block(env, [exp | exp_list]) do
+    case infer_exp(env, exp) do
+      {:ok, _type, env} -> infer_block(env, exp_list)
       error -> error
     end
   end
 
   # Integer literals
-  def infer_exp(%Env{} = env, {:integer, _line, integer}) do
+  def infer_exp(env, {:integer, _line, integer}) do
     Logger.debug("Integer #{integer} found. Type: Int")
     {:ok, :Int, env}
   end
 
   # Booleans
-  def infer_exp(%Env{} = env, {:boolean, _line, boolean}) do
+  def infer_exp(env, {:boolean, _line, boolean}) do
     Logger.debug("Boolean #{boolean} found. Type: Bool")
     {:ok, :Bool, env}
   end
 
   # Variables
-  def infer_exp(%Env{} = env, {:identifier, _line, name}) do
-    type = env.scope[name]
+  def infer_exp(env, {:identifier, _line, name}) do
+    type = get_in(env, [:scope, name])
 
     if type do
       Logger.debug("Variable type found from scope: #{name}:#{type}")
@@ -224,13 +172,13 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # External function calls
-  def infer_exp(%Env{} = env, {:ext_call, _line, {m, f, _, type}}) do
+  def infer_exp(env, {:ext_call, _line, {m, f, _, type}}) do
     Logger.debug("Return type of ext function #{m}.#{f} specified as #{type}")
     {:ok, type, env}
   end
 
   # Function calls
-  def infer_exp(%Env{} = env, {:call, {name, _line}, args, module}) do
+  def infer_exp(env, {:call, {name, _line}, args, module}) do
     exp = %{args: args, name: name}
     # module_name = module || Env.current_module(env)
     Logger.debug("Inferring type of function: #{name}")
@@ -239,8 +187,8 @@ defmodule Fika.Compiler.TypeChecker do
 
   # Function calls using reference
   # exp has to be a function ref type
-  def infer_exp(%Env{} = env, {:call, {exp, _line}, args}) do
-    case infer_exp(%Env{} = env, exp) do
+  def infer_exp(env, {:call, {exp, _line}, args}) do
+    case infer_exp(env, exp) do
       {:ok, %T.FunctionRef{arg_types: arg_types, return_type: type}, env} ->
         case do_infer_args_without_name(env, args) do
           {:ok, ^arg_types, env} ->
@@ -264,11 +212,11 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # =
-  def infer_exp(%Env{} = env, {{:=, _}, {:identifier, _line, left}, right}) do
-    case infer_exp(%Env{} = env, right) do
+  def infer_exp(env, {{:=, _}, {:identifier, _line, left}, right}) do
+    case infer_exp(env, right) do
       {:ok, type, env} ->
         Logger.debug("Adding variable to scope: #{left}:#{type}")
-        env = Env.add_variable_to_scope(env, left, type)
+        env = add_variable_to_scope(env, left, type)
         {:ok, type, env}
 
       error ->
@@ -277,7 +225,7 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # String
-  def infer_exp(%Env{} = env, {:string, _line, string_parts}) do
+  def infer_exp(env, {:string, _line, string_parts}) do
     Enum.reduce_while(string_parts, {:ok, nil, env}, fn
       string, {:ok, _acc_type, acc_env} when is_binary(string) ->
         Logger.debug("String #{string} found. Type: String")
@@ -303,12 +251,12 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # List
-  def infer_exp(%Env{} = env, {:list, _, exps}) do
+  def infer_exp(env, {:list, _, exps}) do
     infer_list_exps(env, exps)
   end
 
   # Tuple
-  def infer_exp(%Env{} = env, {:tuple, _, exps}) do
+  def infer_exp(env, {:tuple, _, exps}) do
     case do_infer_tuple_exps(exps, env) do
       {:ok, exp_types, env} ->
         {:ok, %T.Tuple{elements: exp_types}, env}
@@ -319,7 +267,7 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # Record
-  def infer_exp(%Env{} = env, {:record, _, name, key_values}) do
+  def infer_exp(env, {:record, _, name, key_values}) do
     if name do
       # Lookup type of name, ensure it matches.
       Logger.error("Not implemented")
@@ -336,18 +284,18 @@ defmodule Fika.Compiler.TypeChecker do
 
   # Map
   # TODO: refactor this when union types are available
-  def infer_exp(%Env{} = env, {:map, _, [kv | rest_kvs]}) do
+  def infer_exp(env, {:map, _, [kv | rest_kvs]}) do
     {key, value} = kv
 
-    with {:ok, key_type, env} <- infer_exp(%Env{} = env, key),
-         {:ok, value_type, env} <- infer_exp(%Env{} = env, value) do
+    with {:ok, key_type, env} <- infer_exp(env, key),
+         {:ok, value_type, env} <- infer_exp(env, value) do
       map_type = %T.Map{key_type: key_type, value_type: value_type}
 
       Enum.reduce_while(rest_kvs, {:ok, map_type, env}, fn {k, v}, {:ok, type, env} ->
         %{key_type: key_type, value_type: value_type} = type
 
-        with {:key, {:ok, ^key_type, env}} <- {:key, infer_exp(%Env{} = env, k)},
-             {:value, {:ok, ^value_type, env}} <- {:value, infer_exp(%Env{} = env, v)} do
+        with {:key, {:ok, ^key_type, env}} <- {:key, infer_exp(env, k)},
+             {:value, {:ok, ^value_type, env}} <- {:value, infer_exp(env, v)} do
           {:cont, {:ok, type, env}}
         else
           {:key, {:ok, diff_type, _}} ->
@@ -367,10 +315,10 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # Function ref
-  def infer_exp(%Env{} = env, {:function_ref, _, {module, function_name, arg_types}}) do
+  def infer_exp(env, {:function_ref, _, {module, function_name, arg_types}}) do
     Logger.debug("Inferring type of function: #{function_name}")
 
-    signature = get_function_signature(module || env.module, function_name, arg_types)
+    signature = get_function_signature(module || env[:module], function_name, arg_types)
 
     case get_type(module, signature, env) do
       {:ok, type} ->
@@ -383,13 +331,13 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # Atom value
-  def infer_exp(%Env{} = env, {:atom, _line, atom}) do
+  def infer_exp(env, {:atom, _line, atom}) do
     Logger.debug("Atom value found. Type: #{atom}")
     {:ok, atom, env}
   end
 
   # if-else expression
-  def infer_exp(%Env{} = env, {{:if, _line}, condition, if_block, else_block}) do
+  def infer_exp(env, {{:if, _line}, condition, if_block, else_block}) do
     Logger.debug("Inferring an if-else expression")
 
     case infer_if_else_condition(env, condition) do
@@ -399,24 +347,24 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   # case expression
-  def infer_exp(%Env{} = env, {{:case, _line}, exp, clauses}) do
+  def infer_exp(env, {{:case, _line}, exp, clauses}) do
     Logger.debug("Inferring a case expression")
 
     # Check the type of exp.
     # For each clause, ensure all of the patterns return {:ok, env}
-    with {:ok, rhs_type, env} <- infer_exp(%Env{} = env, exp),
+    with {:ok, rhs_type, env} <- infer_exp(env, exp),
          {:ok, type, env} <- infer_case_clauses(env, rhs_type, clauses) do
       {:ok, type, env}
     end
   end
 
   # anonymous function
-  def infer_exp(%Env{} = env, {:anonymous_function, _line, args, exps} = ast) do
+  def infer_exp(env, {:anonymous_function, _line, args, exps} = ast) do
     Logger.debug("Inferring type of anonymous function")
 
-    env = Env.reset_scope_and_set_signature(env, ast)
+    env = reset_scope_and_set_signature(env, ast)
 
-    case infer_block(%Env{} = env, exps) do
+    case infer_block(env, exps) do
       {:ok, return_type, _env} ->
         arg_types = Enum.map(args, fn {_, {:type, _, type}} -> type end)
         type = %T.FunctionRef{arg_types: arg_types, return_type: return_type}
@@ -432,8 +380,20 @@ defmodule Fika.Compiler.TypeChecker do
     get_function_signature(module, name, arg_types)
   end
 
-  def init_env(ast) do
-    %Env{ast: ast, scope: %{}}
+  def init_env(ast, args \\ %{}) do
+    %{
+      ast: ast,
+      latest_called_function: nil,
+      type_checker_pid: nil,
+      module: nil,
+      module_name: nil,
+      file: nil,
+      current_signature: nil,
+      has_effect: false,
+      first_pass: false,
+      scope: %{}
+    }
+    |> Map.merge(args)
   end
 
   # TODO: made it work, now make it pretty.
@@ -445,7 +405,7 @@ defmodule Fika.Compiler.TypeChecker do
                                                               {env, types, unmatched} ->
         case Match.match_case(env, pattern, unmatched) do
           {:ok, env, unmatched} ->
-            case infer_block(%Env{} = env, block) do
+            case infer_block(env, block) do
               {:ok, type, env} -> {:cont, {env, [type | types], unmatched}}
               error -> {:halt, error}
             end
@@ -498,7 +458,7 @@ defmodule Fika.Compiler.TypeChecker do
 
   defp do_infer_key_values(key_values, env) do
     Enum.reduce_while(key_values, {:ok, [], env}, fn {k, v}, {:ok, acc, env} ->
-      case infer_exp(%Env{} = env, v) do
+      case infer_exp(env, v) do
         {:ok, type, env} ->
           {:identifier, _, key} = k
           {:cont, {:ok, [{key, type} | acc], env}}
@@ -511,7 +471,7 @@ defmodule Fika.Compiler.TypeChecker do
 
   defp do_infer_tuple_exps(exps, env) do
     Enum.reduce_while(exps, {:ok, [], env}, fn exp, {:ok, acc, env} ->
-      case infer_exp(%Env{} = env, exp) do
+      case infer_exp(env, exp) do
         {:ok, exp_type, env} ->
           {:cont, {:ok, [exp_type | acc], env}}
 
@@ -531,7 +491,7 @@ defmodule Fika.Compiler.TypeChecker do
   def infer_args(env, exp, module) do
     case do_infer_args(env, exp) do
       {:ok, type_acc, env} ->
-        signature = get_function_signature(module || env.module, exp.name, type_acc)
+        signature = get_function_signature(module || env[:module], exp.name, type_acc)
 
         case get_type(module, signature, env) do
           {:ok, %T.Effect{type: type}} ->
@@ -550,19 +510,52 @@ defmodule Fika.Compiler.TypeChecker do
     end
   end
 
+  defp reset_scope_and_set_signature(env, {:function, _, {_name, args, _type, _}} = ast) do
+    current_signature = Fika.Compiler.TypeChecker.function_ast_signature(env[:module], ast)
+
+    env
+    |> Map.merge(%{
+      scope: %{},
+      has_effect: false,
+      current_signature: current_signature,
+      latest_called_function: if(env[:first_pass], do: env[:current_signature]),
+      first_pass: false
+    })
+    |> add_args_to_scope(args)
+  end
+
+  defp reset_scope_and_set_signature(env, {:anonymous_function, _line, args, _}) do
+    env
+    |> Map.put(:scope, %{})
+    |> add_args_to_scope(args)
+  end
+
+  defp reset_scope_and_set_signature(env, _ast), do: env
+
+  defp add_variable_to_scope(env, variable, type) do
+    put_in(env, [:scope, variable], type)
+  end
+
+  defp add_args_to_scope(env, args) do
+    Enum.reduce(args, env, fn {{:identifier, _, name}, {:type, _, type}}, env ->
+      Logger.debug("Adding arg type to scope: #{name}:#{type}")
+      add_variable_to_scope(env, name, type)
+    end)
+  end
+
   defp infer_list_exps(env, []) do
     {:ok, %T.List{}, env}
   end
 
   defp infer_list_exps(env, [exp]) do
-    case infer_exp(%Env{} = env, exp) do
+    case infer_exp(env, exp) do
       {:ok, type, env} -> {:ok, %T.List{type: type}, env}
       error -> error
     end
   end
 
   defp infer_list_exps(env, [exp | rest]) do
-    {:ok, type, env} = infer_exp(%Env{} = env, exp)
+    {:ok, type, env} = infer_exp(env, exp)
 
     Enum.reduce_while(rest, {:ok, %T.List{type: type}, env}, fn exp, {:ok, acc_type, acc_env} ->
       case infer_exp(acc_env, exp) do
@@ -585,7 +578,7 @@ defmodule Fika.Compiler.TypeChecker do
 
   defp do_infer_args(env, exp) do
     Enum.reduce_while(exp.args, {:ok, [], env}, fn arg, {:ok, type_acc, env} ->
-      case infer_exp(%Env{} = env, arg) do
+      case infer_exp(env, arg) do
         {:ok, type, env} ->
           Logger.debug("Argument of #{exp.name} is type: #{type}")
           {:cont, {:ok, [type | type_acc], env}}
@@ -607,7 +600,7 @@ defmodule Fika.Compiler.TypeChecker do
   # TODO: This can be merged with do_infer_args
   defp do_infer_args_without_name(env, args) do
     Enum.reduce_while(args, {:ok, [], env}, fn arg, {:ok, acc, env} ->
-      case infer_exp(%Env{} = env, arg) do
+      case infer_exp(env, arg) do
         {:ok, type, env} ->
           Logger.debug("Argument is type: #{type}")
           {:cont, {:ok, [type | acc], env}}
@@ -631,11 +624,11 @@ defmodule Fika.Compiler.TypeChecker do
   end
 
   defp get_type(module, target_signature, env) do
-    is_local_call = is_nil(module) or module == env.module_name
+    is_local_call = is_nil(module) or module == env[:module_name]
 
-    current_signature = env.current_signature
+    current_signature = env[:current_signature]
 
-    pid = env.type_checker_pid
+    pid = env[:type_checker_pid]
 
     function_dependency =
       if current_signature do
diff --git a/lib/fika/compiler/type_checker/match.ex b/lib/fika/compiler/type_checker/match.ex
index 5be4c4c..17c71a1 100644
--- a/lib/fika/compiler/type_checker/match.ex
+++ b/lib/fika/compiler/type_checker/match.ex
@@ -128,7 +128,7 @@ defmodule Fika.Compiler.TypeChecker.Match do
 
   defp do_match_case(env, {:identifier, _, name}, rhs) do
     env =
-      Env.update_scope(env, name, fn
+      update_in(env, [:scope, name], fn
         nil -> rhs
         %T.Union{types: types} -> T.Union.new([rhs | T.Union.to_list(types)])
         type -> T.Union.new([rhs, type])
diff --git a/lib/fika/compiler/type_checker/parallel_type_checker.ex b/lib/fika/compiler/type_checker/parallel_type_checker.ex
index fae31dd..aa2738c 100644
--- a/lib/fika/compiler/type_checker/parallel_type_checker.ex
+++ b/lib/fika/compiler/type_checker/parallel_type_checker.ex
@@ -56,11 +56,14 @@ defmodule Fika.Compiler.TypeChecker.ParallelTypeChecker do
     Enum.each(state.unchecked_functions, fn {signature, function} ->
       Task.start_link(fn ->
         result =
-          TypeChecker.check(function, %TypeChecker.Env{
-            type_checker_pid: pid,
-            module: state.module_name,
-            current_signature: signature
-          })
+          TypeChecker.check(
+            function,
+            TypeChecker.init_env(nil, %{
+              type_checker_pid: pid,
+              module: state.module_name,
+              current_signature: signature
+            })
+          )
 
         __MODULE__.post_result(pid, signature, result)
       end)
diff --git a/test/fika/compiler/type_checker_test.exs b/test/fika/compiler/type_checker_test.exs
index d619308..84838a6 100644
--- a/test/fika/compiler/type_checker_test.exs
+++ b/test/fika/compiler/type_checker_test.exs
@@ -20,7 +20,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:ok, :Int, _} = TypeChecker.infer_exp(%Env{}, ast)
+    assert {:ok, :Int, _} = TypeChecker.infer_exp(%{}, ast)
   end
 
   test "infer type of atom expressions" do
@@ -28,7 +28,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     {:atom, {1, 0, 2}, :a} = ast = TestParser.expression!(str)
 
-    assert {:ok, :a, _} = TypeChecker.infer_exp(%Env{}, ast)
+    assert {:ok, :a, _} = TypeChecker.infer_exp(%{}, ast)
   end
 
   test "infer type for list of atom expressions" do
@@ -36,7 +36,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:ok, %T.List{type: :a}, _} = TypeChecker.infer_exp(%Env{}, ast)
+    assert {:ok, %T.List{type: :a}, _} = TypeChecker.infer_exp(%{}, ast)
   end
 
   test "infer type of arithmetic expressions" do
@@ -44,7 +44,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:ok, :Int, _} = TypeChecker.infer_exp(%Env{}, ast)
+    assert {:ok, :Int, _} = TypeChecker.infer_exp(%{}, ast)
   end
 
   describe "logical operators" do
@@ -52,23 +52,23 @@ defmodule Fika.Compiler.TypeCheckerTest do
       # and
       str = "true & false"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
 
       # or
       str = "true | false"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
 
       # negation
       str = "!true"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "infer type of logical expressions when using atoms" do
       str = "true & :false"
       ast = TestParser.expression!(str)
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -77,7 +77,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     ast = TestParser.expression!(str)
 
-    assert {:error, "Unknown variable: foo"} = TypeChecker.infer_exp(%Env{}, ast)
+    assert {:error, "Unknown variable: foo"} = TypeChecker.infer_exp(%{}, ast)
   end
 
   test "infer ext function's return type" do
@@ -108,7 +108,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:ok, :Int} = TypeChecker.infer(function, %Env{})
+    assert {:ok, :Int} = TypeChecker.infer(function, %{})
   end
 
   test "check returns error when return type is not the inferred type" do
@@ -122,7 +122,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:error, "Expected type: Float, got: Int"} = TypeChecker.check(function, %Env{})
+    assert {:error, "Expected type: Float, got: Int"} = TypeChecker.check(function, %{})
   end
 
   test "checks tuple return type for function" do
@@ -136,7 +136,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:error, "Expected type: {Float}, got: {Int}"} = TypeChecker.check(function, %Env{})
+    assert {:error, "Expected type: {Float}, got: {Int}"} = TypeChecker.check(function, %{})
   end
 
   test "infer return type of another function in the module" do
@@ -272,8 +272,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:ok, :Float} = TypeChecker.infer(function, %Env{})
-    assert {:ok, :Float} = TypeChecker.check(function, %Env{})
+    assert {:ok, :Float} = TypeChecker.infer(function, %{})
+    assert {:ok, :Float} = TypeChecker.check(function, %{})
   end
 
   test "infer function calls considering union types" do
@@ -292,7 +292,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
     [function] = ast[:function_defs]
 
-    assert {:ok, :ok} = TypeChecker.infer(function, %Env{})
+    assert {:ok, :ok} = TypeChecker.infer(function, %{})
   end
 
   describe "string" do
@@ -301,7 +301,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :String, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "parses string inside interpolation" do
@@ -311,7 +311,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :String, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "returns error when unknown variable in string interpolation" do
@@ -321,7 +321,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "parses known variable in string interpolation" do
@@ -360,7 +360,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       assert {
                :error,
                "Expression used in string interpolation expected to be String, got Int"
-             } = TypeChecker.infer_exp(%Env{}, ast)
+             } = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -370,7 +370,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: :Int}, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.List{type: :Int}, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "list of integers and floats" do
@@ -379,7 +379,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:error, "Elements of list have different types. Expected: Int, got: Float"} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
     end
 
     test "list of floats inferred from fn calls" do
@@ -387,7 +387,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: :Float}, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.List{type: :Float}, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "List of strings" do
@@ -395,7 +395,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: :String}, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.List{type: :String}, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "List of list of integers" do
@@ -403,7 +403,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: %T.List{type: :Int}}, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.List{type: %T.List{type: :Int}}, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "empty list" do
@@ -411,7 +411,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.List{type: nil}, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.List{type: nil}, _} = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -421,8 +421,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.Tuple{elements: [:Int, :Int, :Int]}, _env} =
-               TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.Tuple{elements: [:Int, :Int, :Int]}, _env} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "tuple of integers and floats" do
@@ -431,7 +430,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:ok, %T.Tuple{elements: [:Int, :Float, :Int]}, _env} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
     end
 
     test "tuple of floats inferred from fn calls" do
@@ -439,8 +438,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.Tuple{elements: [:Float, :Float]}, _env} =
-               TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.Tuple{elements: [:Float, :Float]}, _env} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "tuple of strings" do
@@ -448,8 +446,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, %T.Tuple{elements: [:String, :String]}, _env} =
-               TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.Tuple{elements: [:String, :String]}, _env} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "tuple of tuple of mixed types" do
@@ -463,7 +460,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
                   %T.Tuple{elements: [:Int, :Float]},
                   %T.Tuple{elements: [:String, :Bool]}
                 ]
-              }, _env} = TypeChecker.infer_exp(%Env{}, ast)
+              }, _env} = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -474,7 +471,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:ok, %T.Record{fields: [bar: :String, foo: :Int]}, _} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
     end
 
     test "error" do
@@ -482,7 +479,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:error, "Unknown variable: x"} = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -493,7 +490,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:ok, %T.Map{key_type: :String, value_type: :Int}, _} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
     end
 
     test "type check for map with mixed type" do
@@ -502,14 +499,14 @@ defmodule Fika.Compiler.TypeCheckerTest do
       ast = TestParser.expression!(str)
 
       assert {:error, "Expected map key of type Int, but got String"} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
 
       str = ~s({"foo" => [1, 2], "bar" => 345})
 
       ast = TestParser.expression!(str)
 
       assert {:error, "Expected map value of type List(Int), but got Int"} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -528,7 +525,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
               %T.FunctionRef{
                 arg_types: [:Int, :Int],
                 return_type: :Int
-              }, _} = TypeChecker.infer_exp(%Env{}, ast)
+              }, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "without args" do
@@ -542,7 +539,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       CodeServer.set_type(signature("bar", "sum", []), {:ok, :Int})
 
       assert {:ok, %T.FunctionRef{arg_types: [], return_type: :Int}, _} =
-               TypeChecker.infer_exp(%Env{}, ast)
+               TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -551,14 +548,14 @@ defmodule Fika.Compiler.TypeCheckerTest do
       str = "true"
       ast = TestParser.expression!(str)
 
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "false" do
       str = "false"
       ast = TestParser.expression!(str)
 
-      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :Bool, _} = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -577,7 +574,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       assert {
                :error,
                "Wrong type for if condition. Expected: Bool, Got: String"
-             } = TypeChecker.infer_exp(%Env{}, ast)
+             } = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "completes when if and else blocks have same return types" do
@@ -591,7 +588,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, _env} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :String, _env} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "completes when if and else blocks have different return types" do
@@ -605,7 +602,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
       types = MapSet.new([:String, :Int])
-      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "with multiple expressions in blocks" do
@@ -710,7 +707,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, :Bool} = TypeChecker.infer(function, %Env{})
+      assert {:ok, :Bool} = TypeChecker.infer(function, %{})
     end
 
     test "when function returns is expected to return a union type and has if-else clause" do
@@ -732,7 +729,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       types = MapSet.new([:ok, :error])
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.Union{types: ^types}} = TypeChecker.infer(function, %Env{})
+      assert {:ok, %T.Union{types: ^types}} = TypeChecker.infer(function, %{})
     end
 
     test "when function accepts union types and calls a function ref" do
@@ -773,7 +770,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       [function] = ast[:function_defs]
 
       assert {:error, "Expected a function reference, but got type: Int"} =
-               TypeChecker.infer(function, %Env{})
+               TypeChecker.infer(function, %{})
     end
 
     test "function ref when given wrong types" do
@@ -793,7 +790,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       error =
         "Expected function reference to be called with arguments (String, Int), but it was called with arguments (Int)"
 
-      assert {:error, ^error} = TypeChecker.infer(function, %Env{})
+      assert {:error, ^error} = TypeChecker.infer(function, %{})
     end
   end
 
@@ -803,7 +800,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       ast = TestParser.expression!(str)
 
-      assert {:ok, :String, %{has_effect: true}} = TypeChecker.infer_exp(%Env{}, ast)
+      assert {:ok, :String, %{has_effect: true}} = TypeChecker.infer_exp(%{}, ast)
     end
 
     test "functions with effects inside them become effectful" do
@@ -818,8 +815,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.Effect{type: :String}} = TypeChecker.infer(function, %Env{})
-      assert {:ok, %T.Effect{type: :String}} = TypeChecker.check(function, %Env{})
+      assert {:ok, %T.Effect{type: :String}} = TypeChecker.infer(function, %{})
+      assert {:ok, %T.Effect{type: :String}} = TypeChecker.check(function, %{})
     end
 
     test "function with multiple effects" do
@@ -837,8 +834,8 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.infer(function, %Env{})
-      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.check(function, %Env{})
+      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.infer(function, %{})
+      assert {:ok, %T.Effect{type: :Int}} = TypeChecker.check(function, %{})
     end
 
     test "effectful function ref call" do
@@ -876,7 +873,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
               %T.FunctionRef{
                 arg_types: [:Int, :Int],
                 return_type: :Int
-              }, _} = TypeChecker.infer_exp(%Env{}, ast)
+              }, _} = TypeChecker.infer_exp(%{}, ast)
     end
   end
 
@@ -892,7 +889,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [function] = ast[:function_defs]
 
-      assert {:ok, %T.List{type: "b"}} = TypeChecker.infer(function, %Env{})
+      assert {:ok, %T.List{type: "b"}} = TypeChecker.infer(function, %{})
     end
 
     test "can't call incompatible functions on type variables" do
@@ -907,7 +904,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
       [function] = ast[:function_defs]
 
       assert {:error, "Function fika/kernel.+(a, Int) does not exist"} =
-               TypeChecker.infer(function, %Env{})
+               TypeChecker.infer(function, %{})
     end
 
     test "infers return types of functions with type variables" do
@@ -925,7 +922,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [_foo, bar] = ast[:function_defs]
 
-      assert {:ok, :Int} = TypeChecker.infer(bar, %Env{ast: ast})
+      assert {:ok, :Int} = TypeChecker.infer(bar, %{ast: ast})
     end
 
     test "infers return types of functions when type variables are passed as args" do
@@ -946,7 +943,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [_foo, bar] = ast[:function_defs]
 
-      assert TypeChecker.infer(bar, %Env{ast: ast}) == {:ok, T.Union.new([:String, "z"])}
+      assert TypeChecker.infer(bar, %{ast: ast}) == {:ok, T.Union.new([:String, "z"])}
     end
 
     test "type variables used in function ref calls" do
@@ -971,7 +968,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
       [_foo, _bar, baz] = ast[:function_defs]
 
-      assert TypeChecker.infer(baz, %Env{ast: ast}) == {:ok, T.Union.new([:String, "c", :Int])}
+      assert TypeChecker.infer(baz, %{ast: ast}) == {:ok, T.Union.new([:String, "c", :Int])}
     end
   end
 

From 53efffa413218dff4ae5df594c93190e25de9175 Mon Sep 17 00:00:00 2001
From: Paulo Valente <polvalente@gmail.com>
Date: Mon, 8 Mar 2021 01:07:20 -0300
Subject: [PATCH 10/10] fix: use correct type in test module definition

---
 test/fika/compiler/type_checker_test.exs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/fika/compiler/type_checker_test.exs b/test/fika/compiler/type_checker_test.exs
index 84838a6..ae7f6ad 100644
--- a/test/fika/compiler/type_checker_test.exs
+++ b/test/fika/compiler/type_checker_test.exs
@@ -179,7 +179,7 @@ defmodule Fika.Compiler.TypeCheckerTest do
 
   test "infer return type when there is intra-module recursion" do
     str = """
-      fn foo(a: Int) : Int | :nil do
+      fn foo(a: Int) : Loop(Int | :nil) do
         bar(a)
       end
       fn bar(a: Int) : Loop(Int | :nil) do