Auto wrap functions passed to Lua.set! to receive %Lua{} (#53)

lib/lua.ex

@@ -165,8 +165,40 @@ defmodule Lua do
       iex> lua = Lua.set!(Lua.new(), [:sum], fn args -> [Enum.sum(args)] end)
       iex> {[10], _lua} = Lua.eval!(lua, "return sum(1, 2, 3, 4)")
 
+
+  Functions can also take a second argument for the state of Lua
+
+      iex> lua =
+      ...>   Lua.set!(Lua.new(), [:set_count], fn args, state ->
+      ...>     {[], Lua.set!(state, :count, Enum.count(args))}
+      ...>   end)
+      iex> {[3], _} = Lua.eval!(lua, "set_count(1, 2, 3); return count")
+
   """
-  def set!(%__MODULE__{state: state}, keys, value) do
+  def set!(%__MODULE__{} = lua, keys, value) do
+    value =
+      case value do
+        func when is_function(func, 1) ->
+          func
+
+        func when is_function(func, 2) ->
+          fn args, state ->
+            case func.(args, wrap(state)) do
+              {value, %__MODULE__{} = lua} -> {List.wrap(value), lua.state}
+              value -> {List.wrap(value), state}
+            end
+          end
+
+        value ->
+          value
+      end
+
+    do_set(lua.state, keys, value)
+  end
+
+  defp do_set(state, keys, value) do
+    keys = List.wrap(keys)
+
     {_keys, state} =
       Enum.reduce_while(keys, {[], state}, fn key, {keys, state} ->
         keys = keys ++ [key]
@@ -484,8 +516,8 @@ defmodule Lua do
     |> Enum.reduce(lua, fn {name, with_state?, variadic?}, lua ->
       arities = Map.get(funcs, name)
 
-      Lua.set!(
-        lua,
+      do_set(
+        lua.state,
         List.wrap(scope) ++ [name],
         if variadic? do
           wrap_variadic_function(module, name, with_state?)
@@ -502,7 +534,7 @@ defmodule Lua do
 
     if with_state? do
       fn args, state ->
-        execute_function(module, function_name, [args, wrap(state)])
+        execute_function_with_state(module, function_name, [args, wrap(state)], wrap(state))
       end
     else
       fn args ->
@@ -518,7 +550,7 @@ defmodule Lua do
     if with_state? do
       fn args, state ->
         if (length(args) + 1) in arities do
-          execute_function(module, function_name, args ++ [wrap(state)])
+          execute_function_with_state(module, function_name, args ++ [wrap(state)], wrap(state))
         else
           arities = Enum.map(arities, &(&1 - 1))
 
@@ -546,6 +578,7 @@ defmodule Lua do
     # Luerl mandates lists as return values; this function ensures all results conform.
     case apply(module, function_name, args) do
       {ret, %Lua{state: state}} -> {ret, state}
+      [{_, _} | _rest] = table -> [table]
       other -> List.wrap(other)
     end
   catch
@@ -554,6 +587,19 @@ defmodule Lua do
        "Value thrown during function '#{function_name}' execution: #{inspect(thrown_value)}"}
   end
 
+  defp execute_function_with_state(module, function_name, args, lua) do
+    # Luerl mandates lists as return values; this function ensures all results conform.
+    case apply(module, function_name, args) do
+      {ret, %Lua{state: state}} -> {ret, state}
+      [{_, _} | _rest] = table -> {[table], lua.state}
+      other -> {List.wrap(other), lua.state}
+    end
+  catch
+    thrown_value ->
+      {:error,
+       "Value thrown during function '#{function_name}' execution: #{inspect(thrown_value)}"}
+  end
+
   defp ensure_scope!(lua, []) do
     lua
   end

test/lua_test.exs

@@ -179,6 +179,46 @@ defmodule LuaTest do
       assert {[10], _} = Lua.eval!(lua, ~LUA"return sum(1, 2, 3, 4)"c)
     end
 
+    test "it can register functions with two arguments that receive state" do
+      foo = 22
+      my_table = %{"a" => 1, "b" => 2}
+
+      lua =
+        Lua.new()
+        |> Lua.set!([:foo], foo)
+        |> Lua.set!([:my_table], my_table)
+        |> Lua.set!([:get_foo1], fn _args, state ->
+          # Just the value
+          Lua.get!(state, [:foo])
+        end)
+        |> Lua.set!([:get_foo2], fn _args, state ->
+          # Value and state, not wrapped in a list
+          {Lua.get!(state, [:foo]), state}
+        end)
+        |> Lua.set!([:get_foo3], fn _args, state ->
+          # Value and state, wrapped in a list
+          {[Lua.get!(state, [:foo])], state}
+        end)
+        |> Lua.set!([:get_my_table1], fn _args, state ->
+          {Lua.get!(state, [:my_table]), state}
+        end)
+        |> Lua.set!([:get_my_table2], fn _args, state ->
+          {[Lua.get!(state, [:my_table])], state}
+        end)
+
+      assert {[^foo], %Lua{}} = Lua.call_function!(lua, [:get_foo1], [])
+      assert {[^foo], %Lua{}} = Lua.call_function!(lua, [:get_foo2], [])
+      assert {[^foo], %Lua{}} = Lua.call_function!(lua, [:get_foo3], [])
+
+      # Unwrapped table
+      assert {[table], %Lua{} = lua} = Lua.call_function!(lua, [:get_my_table1], [])
+      assert lua |> Lua.decode!(table) |> Lua.Table.as_map() == my_table
+
+      # Wrapped table
+      assert {[table], %Lua{} = lua} = Lua.call_function!(lua, [:get_my_table2], [])
+      assert lua |> Lua.decode!(table) |> Lua.Table.as_map() == my_table
+    end
+
     test "it can evaluate chunks" do
       assert %Lua.Chunk{} = chunk = ~LUA[return 2 + 2]c
 
@@ -322,6 +362,34 @@ defmodule LuaTest do
                """)
     end
 
+    test "you can return single values from the state variant of deflua" do
+      defmodule SingleValueState do
+        use Lua.API, scope: "single"
+
+        deflua foo(value), _state do
+          value
+        end
+
+        deflua bar(value) do
+          value
+        end
+      end
+
+      lua = Lua.load_api(Lua.new(), SingleValueState)
+      assert {[22], _lua} = Lua.eval!(lua, "return single.foo(22)")
+      assert {[], _lua} = Lua.eval!(lua, "return single.foo(nil)")
+
+      # There is ambiguity for empty tables, we treat as an empty return value
+      assert {[], _lua} = Lua.eval!(lua, "return single.foo({})")
+
+      assert {[[{"a", 1}]], _lua} = Lua.eval!(lua, "return single.foo({ a = 1 })")
+
+      assert {[22], _lua} = Lua.eval!(lua, "return single.bar(22)")
+      assert {[], _lua} = Lua.eval!(lua, "return single.bar(nil)")
+      assert {[], _lua} = Lua.eval!(lua, "return single.bar({})")
+      assert {[[{"a", 1}]], _lua} = Lua.eval!(lua, "return single.bar({ a = 1 })")
+    end
+
     test "calling non-functions raises" do
       {_, lua} =
         Lua.eval!("""
@@ -653,16 +721,6 @@ defmodule LuaTest do
       end
     end
 
-    defmodule GlobalVar do
-      use Lua.API, scope: "gv"
-
-      deflua get(name), state do
-        table = Lua.get!(state, [name])
-
-        {[table], state}
-      end
-    end
-
     setup do
       {:ok, lua: Lua.new()}
     end
@@ -721,6 +779,16 @@ defmodule LuaTest do
     end
 
     test "it can handle tref values", %{lua: lua} do
+      defmodule GlobalVar do
+        use Lua.API, scope: "gv"
+
+        deflua get(name), state do
+          table = Lua.get!(state, [name])
+
+          {[table], state}
+        end
+      end
+
       lua = Lua.load_api(lua, GlobalVar)
 
       assert {[[{"a", 1}]], _lua} =