extend defun

3 files changed, 24 insertions, 12 deletions

diff --git a/src/builtins.c b/src/builtins.c
index f1c7d10..efad674 100644
--- a/src/builtins.c
+++ b/src/builtins.c
@@ -75,7 +75,7 @@ SExpRef builtin_foldl(Interp *interp, SExpRef args) {
     }
     SExpRef fn = CAR(args), init = CADR(args), lst = CADDR(args);
     SExpRef ret = init;
-    if (VALTYPE(fn) != kUserFuncSExp && VALTYPE(fn) != kFuncSExp) {
+    if (!CALLABLE(fn)) {
         return new_error(interp, "foldl: type error.\n");
     }
     if (!lisp_check_list(interp, lst)) {
diff --git a/src/interp.h b/src/interp.h
index 1d3a2b4..4e43cd8 100644
--- a/src/interp.h
+++ b/src/interp.h
@@ -60,6 +60,7 @@ SExpRef Interp_load_file(Interp *interp, const char *filename);
     || REF((_x))->type == kBreakSignal \
     || REF((_x))->type == kContinueSignal)
 #define VALTYPE(_x) (REF((_x))->type)
+#define CALLABLE(_x) (VALTYPE(_x) == kFuncSExp || VALTYPE(_x) == kUserFuncSExp)
 #define NIL (interp->nil)
 #define CAR(_x) (lisp_car(interp, (_x)))
 #define CDR(_x) (lisp_cdr(interp, (_x)))
diff --git a/src/primitives.c b/src/primitives.c
index 0eadbb7..73812b8 100644
--- a/src/primitives.c
+++ b/src/primitives.c
@@ -321,17 +321,28 @@ error:
 SExpRef primitive_defun(Interp *interp, SExpRef args, bool istail) {
     SExpRef name, param, body, function;
 
-    if (LENGTH(args) < 3) goto error;
-    if (CAR(interp->stack).idx != interp->top_level.idx) {
-        return new_error(interp, "defun: functions can only be defined in top level.\n");
-    }
-    name = CAR(args);
-    if (VALTYPE(name) != kSymbolSExp) goto error;
-    param = CADR(args);
-    body = CDDR(args);
-    function = new_lambda(interp, param, body, interp->top_level);
-    lisp_defun(interp, name, function);
-    return name;
+    if (LENGTH(args) == 2) {
+        if (CAR(interp->stack).idx != interp->top_level.idx) {
+            return new_error(interp, "defun: functions can only be defined in top level.\n");
+        }
+        name = CAR(args);
+        if (VALTYPE(name) != kSymbolSExp) goto error;
+        function = EVAL(CADR(args)); 
+        if (!CALLABLE(function)) goto error;
+        lisp_defun(interp, name, function);
+        return name;
+    } else if (LENGTH(args) >= 3) {
+        if (CAR(interp->stack).idx != interp->top_level.idx) {
+            return new_error(interp, "defun: functions can only be defined in top level.\n");
+        }
+        name = CAR(args);
+        if (VALTYPE(name) != kSymbolSExp) goto error;
+        param = CADR(args);
+        body = CDDR(args);
+        function = new_lambda(interp, param, body, interp->top_level);
+        lisp_defun(interp, name, function);
+        return name;
+    } else goto error;
 error:
     return new_error(interp, "defun: syntax error.\n");
 }