From f3796d29b4a90df7f207bd41d9ddb347d9c1f125 Mon Sep 17 00:00:00 2001
From: Dmitriy Pleshevskiy <dmitriy@pleshevski.ru>
Date: Sat, 20 Apr 2024 17:36:34 +0300
Subject: [PATCH] improve function api

---
 lib.nix      | 20 ++++++++++++--------
 lib.test.nix |  6 +++++-
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/lib.nix b/lib.nix
index ad36379..fc0886d 100644
--- a/lib.nix
+++ b/lib.nix
@@ -75,8 +75,8 @@ let
 
   call = fnName: args:
     concat [ (validFuncName fnName) (wrapParen (join ", " args)) ];
-  call0 = call [ ];
-  call1 = arg: call [ arg ];
+  call0 = fnName: call fnName [ ];
+  call1 = fnName: arg: call fnName [ arg ];
   require = name: call1 "require" name;
 
   kw_and = raw "and";
@@ -119,20 +119,24 @@ let
 
   local = expr: spaceBetween [ kw_local expr ];
   set = variable: value: join " = " [ (raw variable) value ];
-  func' = fnName: params: body:
-    (spaceBetween
+  func = fnName: params: body:
+    let
+      f =
+        if isList body then concatLines
+        else if isAttrs body then spaceBetween
+        else error "Value ${toString body} is not a valid function body";
+    in
+    (f
       ([
         (concat [
           (spaceBetween [ kw_function (validFuncName fnName) ])
           (wrapParen (join ", " (map raw params)))
         ])
       ]
-      ++ body
+      ++ (if isList body then body else if isAttrs body then [ body ] else [ ])
       ++ [ kw_end ])
     );
-  func = fnName: params: body:
-    func' fnName params (concatLines body);
-  lambda = func' "";
+  lambda = func "";
   lambda0 = lambda [ ];
 
   return = expr: spaceBetween ([ kw_return expr ]);
diff --git a/lib.test.nix b/lib.test.nix
index 26465b0..57b95de 100644
--- a/lib.test.nix
+++ b/lib.test.nix
@@ -168,7 +168,7 @@ with nix2lua; pkgs.lib.runTests {
     );
     expected = "name = ((a >= 1) and (a <= 10))";
   };
-  "test returns defined function" = {
+  "test returns defined multiline function" = {
     expr = toLua (func "hello" [ "a" "b" ] [
       (eq (mul (add (var "a") 2) (sub 2 1)) 3)
       (not (eq (mul (add (var "b") 2) (sub 2 1)) 3))
@@ -184,6 +184,10 @@ with nix2lua; pkgs.lib.runTests {
       end
     '';
   };
+  "test returns defined singleline function" = {
+    expr = toLua (func "hello" [ "a" ] (call1 "h.world" (var "a")));
+    expected = "function hello(a) h.world(a) end";
+  };
   "test returns if statement" = {
     expr = toLua (if' (eq 10 10) [
       (call "print" [ 10 ])