add some additional checks and docstrings for fn*

1 files changed, 71 insertions, 50 deletions

diff --git a/macros/fn.fnl b/macros/fn.fnl
index 5be8abe..7d12008 100644
--- a/macros/fn.fnl
+++ b/macros/fn.fnl
@@ -1,38 +1,88 @@
 (local _unpack (or table.unpack unpack))
+(local insert table.insert)
 
 (fn string? [x]
   (= (type x) "string"))
 
 (fn has-amp? [args]
+  "Check if arglist has `&' and return its position of `false'.
+Performs additional checks for `&' usage in arglist."
+
   (var res false)
   (each [i s (ipairs args)]
-    (when (= (tostring s) "&")
-      (set res i)))
+    (if (= (tostring s) "&")
+        (if res (assert-compile false "only one `&' can be specified in arglist." args)
+            (set res i))
+        (and res (> i (+ res 1)))
+        (assert-compile false "only one `more' arg can be supplied after `&' in arglist." args)))
   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)))
+  "Forms three values, representing data needed to create dispatcher:
+
+- the lengs of arglist;
+- the body of the function we generate;
+- position of `&' in the arglist. "
+  (assert-compile (sequence? args) "fn*: expected parameters table.
+* Try adding function parameters as a list of identifiers in brackets." args)
+  (values (length args)
+          (list 'let [args ['...]] (_unpack body))
+          (has-amp? args)))
+
+(fn arity-dispatcher [len fixed amp-body name]
+  "Forms an `if' expression with all fixed arities first, then `&'
+arity, if present, and default error message as last arity.
 
-(fn arity-dispatcher [size fixed amp-body name]
-  (let [bodies []]
+`len' is a symbol, that represens the length of the current argumen
+list, and is computed at runtime.
+
+`fixed' is a table of arities with fixed amount of arguments. These are put in this `if' as:
+`(= len fixed-len)', where `fixed-len' is the length of current arity arglist, computed with `gen-arity'.
+
+`amp-body' stores size of fixed part of arglist, that is, everything up until `&'"
+  (let [bodies '(if)]
     (each [i body (pairs (doto fixed))]
-      (table.insert bodies (list '= size i))
-      (table.insert bodies body))
+      (insert bodies (list '= len i))
+      (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))))
+        (insert bodies (list '>= len (- i 1)))
+        (insert bodies body)))
+    (insert bodies (list 'error
+                         (.. "wrong argument amount"
+                             (if name (.. " for "  name) "")) 3))
+    bodies))
+
+(fn single-arity-body [args fname]
+  (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))))
 
+(fn multi-arity-body [args fname]
+  (let [bodies {}
+        amp-bodies {}]
+    (each [_ arity (ipairs args)]
+      (let [(n body amp) (gen-arity arity)]
+        (if amp
+            (do (insert amp-bodies amp)
+                (insert amp-bodies body)
+                (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))))
 
 (fn fn* [name doc? ...]
   (assert-compile (not (string? name)) "fn* expects symbol as function name" name)
@@ -42,37 +92,8 @@ Try wrapping arguments in square brackets." args)
                  (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)))
+        body (if (sequence? x) (single-arity-body args fname)
+                 (list? x) (multi-arity-body args fname)
                  (assert-compile false "fn* expects vector as its arguments" x))]
     (if (sym? name)
         `(fn ,name [...] ,docstring ,body)