1 files changed, 84 insertions, 0 deletions
diff --git a/macros/fn.fnl b/macros/fn.fnl
new file mode 100644
index 0000000..5be8abe
--- /dev/null
+++ b/macros/fn.fnl
@@ -0,0 +1,84 @@
+(local _unpack (or table.unpack unpack))
+
+(fn string? [x]
+  (= (type x) "string"))
+
+(fn has-amp? [args]
+  (var res false)
+  (each [i s (ipairs args)]
+    (when (= (tostring s) "&")
+      (set res i)))
+  res)
+
+(fn gen-arity [[args & body]]
+  (assert-compile (sequence? args) "fn* expects vector as it's parameter list.
+Try wrapping arguments in square brackets." args)
+  (list (length args)
+        (list 'let [args ['...]] (_unpack body))
+        (has-amp? args)))
+
+(fn arity-dispatcher [size fixed amp-body name]
+  (let [bodies []]
+    (each [i body (pairs (doto fixed))]
+      (table.insert bodies (list '= size i))
+      (table.insert bodies body))
+    (when amp-body
+      (let [[i body] amp-body]
+        (table.insert bodies (list '>= size (- i 1)))
+        (table.insert bodies body)))
+    (table.insert
+     bodies
+     (list 'error
+           (.. "wrong argument amount"
+               (if name (.. " for "  name) "")) 3))
+    (list 'if (_unpack bodies))))
+
+
+(fn fn* [name doc? ...]
+  (assert-compile (not (string? name)) "fn* expects symbol as function name" name)
+  (let [docstring (if (string? doc?) doc? nil)
+        fname (if (sym? name) (tostring name))
+        args (if (sym? name)
+                 (if (string? doc?) [...] [doc? ...])
+                 [name doc? ...])
+        [x] args
+        body (if (sequence? x)
+                 ;; Single-arity function
+                 (let [[args & body] args
+                       [arity body amp] (gen-arity [args (_unpack body)])]
+                   `(let [len# (length [...])]
+                      ,(arity-dispatcher
+                        'len#
+                        (if amp {} {arity body})
+                        (if amp [amp body])
+                        fname)))
+                 ;; Multi-arity function
+                 (list? x)
+                 (let [bodies {}
+                       amp-bodies {}]
+                   (each [_ arity (ipairs args)]
+                     (let [[n body amp] (gen-arity arity)]
+                       (if amp
+                           (do (table.insert amp-bodies amp)
+                               (table.insert amp-bodies body)
+                               (table.insert amp-bodies arity))
+                           (tset bodies n body))))
+                   (assert-compile (<= (length amp-bodies) 3)
+                                   "fn* must have only one arity with &:"
+                                   (. amp-bodies (length amp-bodies)))
+                   `(let [len# (length [...])]
+                      ,(arity-dispatcher
+                        'len#
+                        bodies
+                        (if (~= (next amp-bodies) nil)
+                            amp-bodies)
+                        fname)))
+                 (assert-compile false "fn* expects vector as its arguments" x))]
+    (if (sym? name)
+        `(fn ,name [...] ,docstring ,body)
+        `(fn [...] ,docstring ,body))))
+
+{: fn*}
+
+;; (import-macros {: fn*} :fn)
+;; (fn* f ([a] a) ([a b] (+ a b)))