implement Clojure's defn

4 files changed, 128 insertions, 42 deletions

diff --git a/.dir-locals.el b/.dir-locals.el
new file mode 100644
index 0000000..3c181f6
--- /dev/null
+++ b/.dir-locals.el
@@ -0,0 +1,4 @@
+;;; Directory Local Variables
+;;; For more information see (info "(emacs) Directory Variables")
+
+((fennel-mode . ((eval . (put 'fn* 'fennel-indent-function 'defun)))))
diff --git a/README.org b/README.org
index ae75172..e681f26 100644
--- a/README.org
+++ b/README.org
@@ -1,3 +1,4 @@
+
 * Fennel Cljlib
 Library for [[https://fennel-lang.org/][Fennel]] language that adds a lot of functions from [[https://clojure.org/][Clojure]] standard library.
 This is not a one to one port of Clojure =core=, because many Clojure functions require certain guarantees, like immutability of the underlying data structures.
@@ -9,6 +10,45 @@ Goals of this project are:
 - Be close to the platform, e.g. implement functions in a way that is efficient to use in Lua VM,
 - Be well documented library, with good test coverage.
 
+** Macros
+*** =fn*=
+Clojure's =defn= equivalent.
+Returns a function of fixed arity by doing runtime dispatch, based on argument amount.
+Capable of producing multi-arity functions:
+
+#+begin_src fennel
+  (fn* square "square number" [x] (^ x 2))
+
+  (square 9) ;; 81
+  (square 1 2) ;; error
+
+  (fn* range
+    "Returns increasing sequence  of numbers from `lower' to `upper'.
+  If `lower' is not provided, sequence starts from zero.
+  Accepts optional `step'"
+    ([upper] (range 0 upper 1))
+    ([lower upper] (range lower upper 1))
+    ([lower upper step]
+     (let [res []]
+       (for [i lower upper step]
+         (table.insert res i))
+       res)))
+
+  (range 10)
+  ;; [0 1 2 3 4 5 6 7 8 9 10]
+  (range -10 0)
+  ;; [-10 -9 -8 -7 -6 -5 -4 -3 -2 -1 0]
+  (range 0 1 0.2)
+  ;; [0.0 0.2 0.4 0.6 0.8 1.0]
+
+  ;; both variants support up to one arity with & more:
+  (fn* list [& xs] xs)
+
+  (list 1 2 3)
+  ;; [1 2 3]
+#+end_src
+
+See =core.fnl= for more examples.
 ** Functions
 Here are some important functions from the library.
 Full set can be examined by requiring the module.
diff --git a/fn.fnl b/fn.fnl
deleted file mode 100644
index e937878..0000000
--- a/fn.fnl
+++ /dev/null
@@ -1,42 +0,0 @@
-(fn string? [x]
-  (= (type x) "string"))
-
-(fn has-amp? [args]
-  (var res false)
-  (each [_ s (ipairs args)]
-    (when (= (tostring s) "&")
-      (set res true)))
-  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)
-  (let [arg-length (if (has-amp? args) (sym "_") (length args))
-        body (list 'let [args [(sym "...")]] (unpack body))]
-    (list arg-length body)))
-
-(fn fn* [name doc? ...]
-  (assert-compile (not (string? name)) "fn* expects symbol as function name" name)
-  (let [docstring (if (string? doc?) doc? nil)
-        args (if (sym? name)
-                 (if (string? doc?) [...] [doc? ...])
-                 [name doc? ...])
-        [x & xs] args]
-    (if (sequence? x)
-        ;; Ordinary function
-        (let [[args & body] args]
-          (if (sym? name)
-              `(fn ,name ,args ,docstring ,(unpack body))
-              `(fn ,args ,docstring ,(unpack body))))
-        ;; Multi-arity function
-        (list? x)
-        (let [bodies []]
-          (each [_ arity (ipairs args)]
-            (let [[arity body] (gen-arity arity)]
-              (table.insert bodies arity)
-              (table.insert bodies body)))
-          `(fn ,name [...] ,docstring (match (length [...]) ,(unpack bodies)))))))
-
-{: fn*}
-
-;; (import-macros {: fn*} :fn) (macrodebug (fn* f ([a] a)))
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)))