Repository: fika-lang/fika (elixir)

Issue:
Add support for union types
A union type is a set consisting of multiple non-overlapping types. With union types, we can have function signatures of the type:

```elixir
fn foo : {:ok, Int} | {:error, String} do
  if true do
    {:ok, 123}
  else
    {:error, "Something went wrong"}
  end
end
```

We'll need to understand the concept of subtypes before we start working on this. [This wiki](https://en.wikipedia.org/wiki/Subtyping) is a good introduction.

In this example:
![image](https://user-images.githubusercontent.com/1707078/95010225-33699680-0645-11eb-9064-032975e96198.png)
duck is a subtype of bird (which is its supertype), which means it can be substituted wherever a bird is expected. Similarly, a union type is a supertype of each member of the set.

Changes required:

1. Allow if-else branches to return different types. The inferred type of the if-else expression will be a union of the two types. This can be represented using a new "Union" struct that can be defined [here](https://github.com/fika-lang/fika/blob/master/lib/fika/types.ex) like the FunctionRef struct. The to_string function of the struct will join the types separated by " | ".
2. When [inferring the return type of a function](https://github.com/fika-lang/fika/blob/master/lib/fika/type_checker.ex#L52), we'll now expand union types and add all possible function types in [add_function_type](https://github.com/fika-lang/fika/blob/master/lib/fika/type_checker.ex#L402). Example:

```elixir
# A function that takes x which can be either :foo or :bar and returns either :ok or :error
fn foo(x: :foo | :bar) : :ok | :error do
  # ...
end
```

The type checker will add the following function types to its "function_types" map in env:
1. module.foo(:foo) => :ok 
2. module.foo(:foo) => :error
3. module.foo(:bar) => :ok
4. module.foo(:bar) => :error 

We can optimize the storing and checking of function types later, but this should do for now.

Please let me know if I've missed something here!

Pull request description:
feat: add union types
New PR because the old one got too outdated with the new structs

closes #24

Patch:
diff --git a/lib/fika/parser/function_def.ex b/lib/fika/parser/function_def.ex
index f368be7..c1bba02 100644
--- a/lib/fika/parser/function_def.ex
+++ b/lib/fika/parser/function_def.ex
@@ -24,11 +24,11 @@ defmodule Fika.Parser.FunctionDef do
 
   args =
     arg
-    |> optional(
+    |> repeat(
       allow_space
       |> ignore(string(","))
       |> concat(allow_space)
-      |> parsec(:args)
+      |> concat(arg)
     )
 
   arg_parens =
diff --git a/lib/fika/parser/helper.ex b/lib/fika/parser/helper.ex
index cb88002..f0e2fc5 100644
--- a/lib/fika/parser/helper.ex
+++ b/lib/fika/parser/helper.ex
@@ -67,6 +67,10 @@ defmodule Fika.Parser.Helper do
     %T.List{type: inner_type}
   end
 
+  def do_to_ast({inner_types, _line}, :tuple_type) do
+    %T.Tuple{elements: inner_types}
+  end
+
   def do_to_ast({[{:atom, _, atom}], _line}, :atom_type) do
     atom
   end
@@ -78,6 +82,10 @@ defmodule Fika.Parser.Helper do
     %T.FunctionRef{arg_types: arg_types, return_type: return_type}
   end
 
+  def do_to_ast({types, _line}, :union_type) do
+    T.Union.new(types)
+  end
+
   def do_to_ast({[{_, _, type}], line}, :type) when is_struct(type) do
     {:type, line, type}
   end
diff --git a/lib/fika/parser/non_literal_exps.ex b/lib/fika/parser/non_literal_exps.ex
index 4fffe97..2359adc 100644
--- a/lib/fika/parser/non_literal_exps.ex
+++ b/lib/fika/parser/non_literal_exps.ex
@@ -49,6 +49,7 @@ defmodule Fika.Parser.NonLiteralExps do
     ])
     |> label("function call")
 
+  # TO-DO: accept nested if-else expressions
   exp_if_else =
     ignore(string("if"))
     |> concat(require_space)
diff --git a/lib/fika/parser/types.ex b/lib/fika/parser/types.ex
index 2efca55..a901873 100644
--- a/lib/fika/parser/types.ex
+++ b/lib/fika/parser/types.ex
@@ -72,11 +72,11 @@ defmodule Fika.Parser.Types do
     |> label("map type")
 
   tuple_type =
-    string("{")
+    ignore(string("{"))
     |> concat(type_args)
-    |> string("}")
-    |> reduce({Enum, :join, []})
+    |> ignore(string("}"))
     |> label("tuple type")
+    |> Helper.to_ast(:tuple_type)
 
   list_type =
     ignore(string("List("))
@@ -105,7 +105,7 @@ defmodule Fika.Parser.Types do
     |> label("nothing")
     |> reduce({Helper, :to_atom, []})
 
-  type =
+  base_type =
     choice([
       string_type,
       int_type,
@@ -119,6 +119,20 @@ defmodule Fika.Parser.Types do
       tuple_type
     ])
 
+  union_type =
+    base_type
+    |> times(
+      allow_space
+      |> ignore(string("|"))
+      |> concat(allow_space)
+      |> concat(base_type),
+      min: 1
+    )
+    |> label("union type")
+    |> Helper.to_ast(:union_type)
+
+  type = choice([union_type, base_type])
+
   parse_type =
     type
     |> Helper.to_ast(:type)
diff --git a/lib/fika/type_checker.ex b/lib/fika/type_checker.ex
index 6814eea..9b0c069 100644
--- a/lib/fika/type_checker.ex
+++ b/lib/fika/type_checker.ex
@@ -38,8 +38,8 @@ defmodule Fika.TypeChecker do
       {:ok, ^expected_type, _env} = result ->
         result
 
-      {:ok, other_type, _} ->
-        {:error, "Expected type: #{expected_type}, got: #{other_type}"}
+      {:ok, inferred_type, _env} ->
+        {:error, "Expected type: #{expected_type}, got: #{inferred_type}"}
 
       error ->
         error
@@ -301,19 +301,13 @@ defmodule Fika.TypeChecker do
   end
 
   defp infer_if_else_blocks(env, if_block, else_block) do
-    {_, if_type_val, env} = infer_block(env, if_block)
-    {_, else_type_val, _} = type = infer_block(env, else_block)
-
-    unless if_type_val == else_type_val do
-      Logger.debug("if and else block have different types")
-
-      error =
-        "Expected if and else blocks to have same return type. " <>
-          "Got #{if_type_val} and #{else_type_val}"
-
-      {:error, error}
-    else
-      type
+    with {:ok, if_type_val, if_env} <- infer_block(env, if_block),
+         {:ok, else_type_val, else_env} <- infer_block(if_env, else_block) do
+      if if_type_val == else_type_val do
+        {:ok, if_type_val, else_env}
+      else
+        {:ok, T.Union.new([if_type_val, else_type_val]), else_env}
+      end
     end
   end
 
diff --git a/lib/fika/types/union.ex b/lib/fika/types/union.ex
new file mode 100644
index 0000000..49735fc
--- /dev/null
+++ b/lib/fika/types/union.ex
@@ -0,0 +1,30 @@
+defmodule Fika.Types.Union do
+  defstruct types: MapSet.new()
+
+  @type t :: %__MODULE__{types: MapSet.t()}
+
+  @spec new(types :: Enumerable.t()) :: t()
+  def new(types) do
+    %__MODULE__{types: flatten_types(types)}
+  end
+
+  @spec flatten_types(types :: Enumerable.t()) :: MapSet.t()
+  @doc false
+  def flatten_types(types) do
+    types
+    |> Enum.flat_map(fn
+      %__MODULE__{types: t} ->
+        flatten_types(t)
+
+      t ->
+        [t]
+    end)
+    |> MapSet.new()
+  end
+
+  defimpl String.Chars, for: Fika.Types.Union do
+    def to_string(%{types: types}) do
+      Enum.join(types, " | ")
+    end
+  end
+end


Test patch:
diff --git a/test/fika/parser_test.exs b/test/fika/parser_test.exs
index d7a69b8..979e413 100644
--- a/test/fika/parser_test.exs
+++ b/test/fika/parser_test.exs
@@ -501,6 +501,45 @@ defmodule Fika.ParserTest do
                  ]}}
              ]
     end
+
+    test "with union types" do
+      code = """
+      fn foo(a: Int | Float, b: Nothing) : :ok | {:error, String} do
+        if true do
+          :ok
+        else
+          {:error, "not ok"}
+        end
+      end
+      """
+
+      arg_union = MapSet.new([:Int, :Float])
+
+      return_union =
+        MapSet.new([
+          :ok,
+          %T.Tuple{elements: [:error, :String]}
+        ])
+
+      assert {:function, [position: {7, 120, 123}],
+              {:foo,
+               [
+                 {{:identifier, {1, 0, 8}, :a}, {:type, {1, 0, 21}, %T.Union{types: ^arg_union}}},
+                 {{:identifier, {1, 0, 24}, :b}, {:type, {1, 0, 33}, :Nothing}}
+               ],
+               {:type, {1, 0, 59},
+                %T.Union{
+                  types: ^return_union
+                }},
+               [
+                 {{:if, {6, 114, 119}}, {:boolean, {2, 63, 72}, true},
+                  [{:atom, {3, 76, 83}, :ok}],
+                  [
+                    {:tuple, {5, 91, 113},
+                     [{:atom, {5, 91, 102}, :error}, {:string, {5, 91, 112}, ["not ok"]}]}
+                  ]}
+               ]}} = TestParser.function_def!(code)
+    end
   end
 
   describe "if-else expression" do
diff --git a/test/fika/type_checker_test.exs b/test/fika/type_checker_test.exs
index 0881d62..c4ead9a 100644
--- a/test/fika/type_checker_test.exs
+++ b/test/fika/type_checker_test.exs
@@ -439,7 +439,7 @@ defmodule Fika.TypeCheckerTest do
              } = TypeChecker.infer_exp(env, ast)
     end
 
-    test "completes when if and else blocks have same return type" do
+    test "completes when if and else blocks have same return types" do
       str = """
       if true do
         "done"
@@ -454,7 +454,7 @@ defmodule Fika.TypeCheckerTest do
       assert {:ok, :String, _env} = TypeChecker.infer_exp(env, ast)
     end
 
-    test "error when if and else blocks have different return types" do
+    test "completes when if and else blocks have different return types" do
       str = """
       if false do
         "done"
@@ -466,10 +466,8 @@ defmodule Fika.TypeCheckerTest do
       ast = TestParser.expression!(str)
       env = Env.init_module_env(Env.init(), "test", ast)
 
-      assert {
-               :error,
-               "Expected if and else blocks to have same return type. Got String and Int"
-             } = TypeChecker.infer_exp(env, ast)
+      types = MapSet.new([:String, :Int])
+      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.infer_exp(env, ast)
     end
 
     test "with multiple expressions in blocks" do
@@ -508,6 +506,50 @@ defmodule Fika.TypeCheckerTest do
       assert {:ok, :Bool, _} = TypeChecker.infer(function, env)
     end
 
+    test "when function returns is expected to return an union type and has if-else clause" do
+      str = """
+      fn foo(x: String, y: Int) : :ok | :error do
+        f = &test2.bar(String, Int)
+        if f.(x, y) do
+          :ok
+        else
+          :error
+        end
+      end
+      """
+
+      {:module, _, [function]} = ast = Fika.Parser.parse_module(str, "test1")
+
+      env =
+        Env.init()
+        |> Env.init_module_env("test", ast)
+        |> Env.add_function_type("test2.bar(String, Int)", :Bool)
+
+      types = MapSet.new([:ok, :error])
+      assert {:ok, %T.Union{types: ^types}, _} = TypeChecker.infer(function, env)
+    end
+
+    test "when function accepts union types and calls a function ref" do
+      str = """
+      fn foo(x: String, y: Int) : :ok | :error do
+        f = &test2.bar(String, Int)
+        f.(x, y)
+      end
+      """
+
+      {:module, _, [function]} = ast = Fika.Parser.parse_module(str, "test1")
+
+      types = MapSet.new([:error, :ok])
+
+      env =
+        Env.init()
+        |> Env.init_module_env("test", ast)
+        |> Env.add_function_type("test2.bar(String, Int)", %T.Union{types: types})
+
+      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.infer(function, env)
+      assert {:ok, %T.Union{types: ^types}, _env} = TypeChecker.check(function, env)
+    end
+
     test "identifier is not a reference" do
       str = """
       fn foo do
diff --git a/test/fika/types/union_type_test.exs b/test/fika/types/union_type_test.exs
new file mode 100644
index 0000000..28e0eef
--- /dev/null
+++ b/test/fika/types/union_type_test.exs
@@ -0,0 +1,32 @@
+defmodule Fika.Types.UnionTypeTest do
+  use ExUnit.Case, async: true
+
+  alias Fika.Types.Union
+
+  describe "new/1" do
+    test "handles a flat list" do
+      expected = MapSet.new([:a, :b, :c])
+      assert %Union{types: expected} == Union.new([:a, :b, :c])
+    end
+  end
+
+  describe "flatten_types/1" do
+    test "handles a list which contains nested Unions" do
+      expected = MapSet.new([:a, :b, :c, :d])
+
+      assert expected == Union.flatten_types([%Union{types: [:a, :b]}, :c, :d])
+      assert expected == Union.flatten_types([:a, %Union{types: [:b, :c]}, :d])
+      assert expected == Union.flatten_types([:a, :b, %Union{types: [:c, :d]}])
+
+      assert expected ==
+               Union.flatten_types([
+                 %Union{
+                   types: [
+                     %Union{types: [:a, :b]},
+                     %Union{types: [:c, :d]}
+                   ]
+                 }
+               ])
+    end
+  end
+end