support for multi-arity functions in raw mode

lib/zig/_nif.ex

@@ -100,8 +100,8 @@ defmodule Zig.Nif do
     }
   end
 
-  def arity(%{raw: nil, signature: %{arity: arity}}), do: arity
-  def arity(%{raw: t, params: arity}) when not is_nil(t), do: arity
+  def arities(%{raw: nil, signature: %{arity: arity}}), do: [arity]
+  def arities(%{raw: t, params: arities}) when not is_nil(t), do: arities
 
   def render_elixir(%{concurrency: concurrency} = nif) do
     doc =
@@ -116,11 +116,19 @@ defmodule Zig.Nif do
         false ->
           quote do
           end
-
         _ ->
-          quote do
-            @spec unquote(render_elixir_spec(nif))
-          end
+          nif.spec
+          |> List.wrap
+          |> Enum.map(fn spec ->
+
+          end)
+
+        {_, list} when is_list(list) ->
+          Enum.map(list, fn spec ->
+            quote do 
+              @spec unquote(spec)
+            end
+          end)
       end
 
     functions = concurrency.render_elixir(nif)
@@ -132,17 +140,19 @@ defmodule Zig.Nif do
     end
   end
 
-  def render_elixir_spec(%{raw: t, params: arity} = nif) when not is_nil(t) do
-    param_spec = case arity do
-      0 -> []
-      arity -> Enum.map(0..arity - 1, fn _ -> {:term, [], []} end)
-    end
-
-    return_spec = Type.render_elixir_spec(nif.return.type, :return, nif.return)
-
-    quote do
-      unquote(nif.name)(unquote_splicing(param_spec)) :: unquote(return_spec)
-    end
+  def render_elixir_spec(%{raw: t, params: arities} = nif) when not is_nil(t) do
+    Enum.map(arities, fn arity ->
+      param_spec = case arity do
+        0 -> []
+        arity -> Enum.map(0..arity - 1, fn _ -> {:term, [], []} end)
+      end
+
+      return_spec = Type.render_elixir_spec(nif.return.type, :return, nif.return)
+
+      quote do
+        unquote(nif.name)(unquote_splicing(param_spec)) :: unquote(return_spec)
+      end
+    end)
   end
 
   def render_elixir_spec(nif) do
@@ -183,10 +193,10 @@ defmodule Zig.Nif do
 
   def table_entries(nif) do
     nif.concurrency.table_entries(nif)
-    |> Enum.map(fn
+    |> Enum.flat_map(fn
       {function, fptr, concurrency} ->
         flags = Map.fetch!(@flags, concurrency)
-        ~s(.{.name="#{function}", .arity=#{arity(nif)}, .fptr=#{fptr}, .flags=#{flags}})
+        Enum.map(arities(nif), &~s(.{.name="#{function}", .arity=#{&1}, .fptr=#{fptr}, .flags=#{flags}}))
     end)
   end
 

lib/zig/nif/_basic.ex

@@ -57,13 +57,14 @@ defmodule Zig.Nif.Basic do
   end
 
   def render_elixir(%{raw: raw} = nif) when not is_nil(raw) do
-    unused_params = Nif.elixir_parameters(nif.params, false)
-
-    quote do
-      unquote(style(nif))(unquote(nif.name)(unquote_splicing(unused_params))) do
-        :erlang.nif_error(unquote(error_text(nif, nif.params)))
+    Enum.map(nif.params, fn arity ->
+      unused_params = Nif.elixir_parameters(arity, false)
+      quote do
+        unquote(style(nif))(unquote(nif.name)(unquote_splicing(unused_params))) do
+          :erlang.nif_error(unquote(error_text(nif, nif.params)))
+        end
       end
-    end
+    end)
   end
 
   def render_elixir(%{signature: %{arity: arity}} = nif) do

lib/zig/sema.ex

@@ -144,12 +144,18 @@ defmodule Zig.Sema do
          opts
        )
        when t in ~w[term erl_nif_term]a do
-    case Keyword.fetch(opts, :arity) do
-      {:ok, arity} when arity in 0..63 ->
-        %{nif | signature: sema, raw: t, params: arity, return: Return.new(t)}
+    arities = case Keyword.fetch(opts, :arity) do
+      {:ok, arity} when arity in 0..63 -> arities(arity)
+      {:ok, {:.., _, _} = range} -> arities(range)
+      {:ok, list} when is_list(list) -> Enum.flat_map(list, &arities/1)
     end
+
+    %{nif | signature: sema, raw: t, params: arities, return: Return.new(t)}
   end
 
+  defp arities(integer) when is_integer(integer), do: [integer]
+  defp arities({:.., _, [start, finish]}), do: Enum.to_list(start..finish)
+
   defp apply_from_sema(nif, sema, opts) do
     %{
       nif

test/integration/raw/basic_raw_test.exs

@@ -14,20 +14,20 @@ defmodule ZiglerTest.Raw.BasicRawTest do
   const beam = @import("beam");
   const e = @import("erl_nif");
 
-  pub fn raw_beam_term(env: beam.env, count: c_int, list: [*]const beam.term) beam.term {
-    return beam.make(.{.count = count, .item = list[0]}, .{.env = env});
+  pub fn raw_beam_term(env: beam.env, arity: c_int, params: [*]const beam.term) beam.term {
+    return beam.make(.{.arity = arity, .item = params[0]}, .{.env = env});
   }
 
-  pub fn raw_erl_nif_term(env: beam.env, count: c_int, list: [*]const e.ErlNifTerm) e.ErlNifTerm {
-    return beam.make(.{.count = count, .item = beam.term{.v = list[0]}}, .{.env = env}).v;
+  pub fn raw_erl_nif_term(env: beam.env, arity: c_int, params: [*]const e.ErlNifTerm) e.ErlNifTerm {
+    return beam.make(.{.arity = arity, .item = beam.term{.v = params[0]}}, .{.env = env}).v;
   }
   """
 
   test "raw call with beam.term" do
-    assert %{count: 1, item: 1.0} = raw_beam_term(1.0)
+    assert %{arity: 1, item: 1.0} = raw_beam_term(1.0)
   end
 
   test "raw call with erl_nif_term" do
-    assert %{count: 1, item: 1.0} = raw_erl_nif_term(1.0)
+    assert %{arity: 1, item: 1.0} = raw_erl_nif_term(1.0)
   end
 end

test/integration/raw/multi_arity_test.exs

@@ -0,0 +1,58 @@
+defmodule ZiglerTest.Raw.MultiArityTest do
+  use ZiglerTest.IntegrationCase, async: true
+
+  @moduletag :raw
+
+  use Zig,
+    otp_app: :zigler,
+    nifs: [
+      range_arity: [arity: 1..3],
+      list_arity: [arity: [1, 3..4]]
+    ]
+
+  ~Z"""
+  const beam = @import("beam");
+  const e = @import("erl_nif");
+
+  pub fn range_arity(env: beam.env, arity: c_int, params: [*]const beam.term) beam.term {
+    return beam.make(params[0..@intCast(arity)], .{.env = env});
+  }
+
+  pub fn list_arity(env: beam.env, arity: c_int, params: [*]const beam.term) beam.term {
+    return beam.make(params[0..@intCast(arity)], .{.env = env});
+  }
+  """
+
+  test "raw call with range arity" do
+    for arity <- 1..3 do
+      foos = List.duplicate(:foo, arity)
+      assert foos == apply(__MODULE__, :range_arity, foos)
+    end
+
+    assert_raise UndefinedFunctionError, fn ->
+      apply(__MODULE__, :range_arity, [])
+    end
+
+    assert_raise UndefinedFunctionError, fn ->
+      apply(__MODULE__, :range_arity, [:foo, :foo, :foo, :foo])
+    end
+  end
+
+  test "raw call with list of arity mixtures" do
+    for arity <- 3..4 do
+      assert List.duplicate(:foo, arity) == apply(__MODULE__, :list_arity, List.duplicate(:foo, arity))
+    end
+
+    assert_raise UndefinedFunctionError, fn ->
+      apply(__MODULE__, :list_arity, [])
+    end
+
+    assert_raise UndefinedFunctionError, fn ->
+      apply(__MODULE__, :list_arity, [:foo, :foo])
+    end
+
+    assert_raise UndefinedFunctionError, fn ->
+      apply(__MODULE__, :list_arity, [:foo, :foo, :foo, :foo, :foo])
+    end
+  end
+end