diff --git a/lib/lua.nix b/lib/lua.nix
index 3d8d227..257a8ab 100644
--- a/lib/lua.nix
+++ b/lib/lua.nix
@@ -1,66 +1,79 @@
 {lib, ...}: let
   inherit (lib.strings) removeSuffix;
-  inherit (builtins) match;
-  handleAst = data: let
-    ast = data.__ast;
-  in
-    if ast == "raw"
-    then data.data
-    else if ast == "return"
-    # TODO nix2lua can return null
-    then "return ${nix2lua data.data}"
-    else if ast == "fn"
-    then "function ${data.name}()\n${nix2lua data.data}\nend"
-    else abort ''Unknown ast type ${ast.__ast}'';
+  inherit (builtins) match isNull typeOf concatStringsSep attrNames concatMap;
+
+  commaJoin = concatStringsSep ", ";
+  wrapNotNull' = val: value:
+    if isNull val
+    then []
+    else [value];
+  wrapNotNull = val: wrapNotNull' val val;
+
+  toLua = {
+    ast = {
+      raw = {data, ...}: data;
+      return = {data, ...}: "return ${nix2lua data}";
+      fn = {
+        data,
+        name,
+        args,
+        ...
+      }: "function ${name}(${commaJoin args})\n${nix2lua data}\nend";
+    };
+
+    type = rec {
+      null = _: "nil";
+      string = data: ''"${data}"'';
+      path = string;
+      lambda = data: builtins.trace "Skipping function" null;
+      int = data: toString data;
+
+      bool = data:
+        if data
+        then "true"
+        else "false";
+
+      ast = data: let
+        astType = data.__ast;
+      in
+        if toLua.ast ? ${astType}
+        then toLua.ast.${astType} data
+        else abort ''Unknown ast type ${astType}'';
+
+      list = data: let
+        nix2luaList = val: wrapNotNull (nix2lua val);
+        listContent = commaJoin (concatMap nix2luaList data);
+      in "{ ${listContent} }";
+
+      set = data: let
+        mkKeyValue = key: let
+          value = data.${key};
+          luaKey =
+            if isNull (match "[a-zA-Z_][a-zA-Z_0-9]+" key)
+            then ''["${key}"]''
+            else key;
+          luaValue = nix2lua value;
+        in
+          wrapNotNull' luaValue ''
+            ${luaKey} = ${luaValue},
+          '';
+        attrsContent = concatMap mkKeyValue (attrNames data);
+      in ''
+        {
+          ${concatStringsSep "" attrsContent}}
+      '';
+    };
+  };
+
+  newAst = type: set: set // {__ast = type;};
 
   nix2lua = data: let
-    inherit (builtins) isInt isBool isNull isString isList isPath isAttrs isFunction typeOf concatStringsSep attrNames concatMap;
     type = typeOf data;
-    condValue2list = val: value:
-      if isNull val
-      then []
-      else [value];
-    value2list = val: condValue2list val val;
   in
     if data ? __ast
-    then handleAst data
-    else if isInt data
-    then toString data
-    else if isBool data
-    then
-      if data
-      then "true"
-      else "false"
-    else if isString data || isPath data
-    then ''"${data}"''
-    else if isNull data
-    then "nil"
-    else if isFunction data
-    then (builtins.trace "Skipping function" null)
-    else if isList data
-    then let
-      nix2luaList = val: value2list (nix2lua val);
-      listContent = concatStringsSep ", " (concatMap nix2luaList data);
-    in "{ ${listContent} }"
-    else if isAttrs data
-    then let
-      mkKeyValue = key: let
-        value = data.${key};
-        luaKey =
-          if isNull (match "[a-zA-Z_][a-zA-Z_0-9]+" key)
-          then ''["${key}"]''
-          else key;
-        luaValue = nix2lua value;
-      in
-        # TODO Handle indentifier keys
-        condValue2list luaValue ''
-          ${luaKey} = ${luaValue},
-        '';
-      attrsContent = concatMap mkKeyValue (attrNames data);
-    in ''
-      {
-        ${concatStringsSep "" attrsContent}}
-    ''
+    then toLua.type.ast data
+    else if toLua.type ? ${type}
+    then toLua.type.${type} data
     else abort ''Type "${type}"'';
 in rec {
   inherit nix2lua;
@@ -75,21 +88,16 @@ in rec {
     (${wrapFunction (removeSuffix "\n" lua)})();
     -- end ${section}
   '';
+  raw = data: newAst "raw" {inherit data;};
 
-  raw = data: {
-    inherit data;
-    __ast = "raw";
-  };
+  functionWithArgs = name: args: data:
+    newAst "fn" {
+      inherit name data args;
+    };
 
-  function = name: data: {
-    inherit name data;
-    __ast = "fn";
-  };
+  function = name: data: functionWithArgs name [] data;
 
   lambda = function "";
 
-  return = data: {
-    inherit data;
-    __ast = "return";
-  };
+  return = data: newAst "return" {inherit data;};
 }