From 176026c58d86981945a5463401f0732d6a29ff97 Mon Sep 17 00:00:00 2001
From: Dmitriy Pleshevskiy <dmitriy@pleshevski.ru>
Date: Wed, 17 Apr 2024 14:55:12 +0300
Subject: [PATCH] add local and func

---
 lib.nix      | 58 ++++++++++++++++++++++++++++++++--------------------
 lib.test.nix | 18 +++++++++++++++-
 2 files changed, 53 insertions(+), 23 deletions(-)

diff --git a/lib.nix b/lib.nix
index 8e639ec..7bb9ac5 100644
--- a/lib.nix
+++ b/lib.nix
@@ -49,6 +49,27 @@ let
       ++ [ (raw "") ]
     );
   pipe = join ".";
+  spaceBetween = join " ";
+
+  op = operation: left: right: concat [ (raw "(") (join " ${operation} " [ left right ]) (raw ")") ];
+
+  add = op "+";
+  sub = op "-";
+  mul = op "*";
+  div = op "/";
+  mod = op "%";
+  exp = op "^";
+
+  eq = op "==";
+  ne = op "~=";
+  gt = op ">";
+  lt = op "<";
+  gte = op ">=";
+  lte = op "<=";
+
+  and = op "and";
+  or = op "or";
+  not = expr: spaceBetween [ (raw "not") expr ];
 
   isRaw = expr: getType expr == "raw";
   raw = expr:
@@ -72,28 +93,21 @@ let
     );
   require = name: call "require" [ name ];
 
+  local = expr: spaceBetween [ (raw "local") expr ];
   set = variable: value: join " = " [ (raw variable) value ];
-  lset = variable: value: join " " [ (raw "local") (set variable value) ];
-
-  op = operation: left: right: concat [ (raw "(") (join " ${operation} " [ left right ]) (raw ")") ];
-
-  add = op "+";
-  sub = op "-";
-  mul = op "*";
-  div = op "/";
-  mod = op "%";
-  exp = op "^";
-
-  eq = op "==";
-  ne = op "~=";
-  gt = op ">";
-  lt = op "<";
-  gte = op ">=";
-  lte = op "<=";
-
-  and = op "and";
-  or = op "or";
-  not = expr: join " " [ (raw "not") expr ];
+  func = name: params: body:
+    (concatLines
+      ([
+        (concat [
+          (spaceBetween (map raw [ "function" name ]))
+          (raw "(")
+          (join ", " (map raw params))
+          (raw ")")
+        ])
+      ]
+      ++ body
+      ++ [ (raw "end") ])
+    );
 
   isLuaNil = expr: getType expr == "nil";
   LuaNil = { _type = "nil"; };
@@ -148,7 +162,7 @@ in
 
   inherit LuaNil;
   inherit raw join concat concatLines pipe;
-  inherit namedField call require set lset;
+  inherit namedField call require local set func;
 
   inherit op;
   inherit eq ne gt lt gte lte;
diff --git a/lib.test.nix b/lib.test.nix
index 97c128a..046b027 100644
--- a/lib.test.nix
+++ b/lib.test.nix
@@ -145,7 +145,7 @@ with nix2lua; pkgs.lib.runTests {
     expected = "parser_config.d2 = {  }";
   };
   "test returns lua with setting value to the local variable" = {
-    expr = toLua (lset "parser_config.d2" { });
+    expr = toLua (local (set "parser_config.d2" { }));
     expected = "local parser_config.d2 = {  }";
   };
   "test returns all operations" = {
@@ -162,4 +162,20 @@ with nix2lua; pkgs.lib.runTests {
       (10 > 5)
     '';
   };
+  "test returns defined 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))
+      (not 10)
+      (gt 10 5)
+    ]);
+    expected = ''
+      function hello(a, b)
+      (((a + 2) * (2 - 1)) == 3)
+      not (((b + 2) * (2 - 1)) == 3)
+      not 10
+      (10 > 5)
+      end
+    '';
+  };
 }