1;; This module contains all the special forms; all built in Fennel constructs
2;; which cannot be implemented as macros. It also contains some core compiler
3;; functionality which is kept in this module for circularity reasons.
4
5(local utils (require :fennel.utils))
6(local view (require :fennel.view))
7(local parser (require :fennel.parser))
8(local compiler (require :fennel.compiler))
9(local unpack (or table.unpack _G.unpack))
10
11(local SPECIALS compiler.scopes.global.specials)
12
13(fn wrap-env [env]
14  "Convert a fennel environment table to a Lua environment table.
15This means automatically unmangling globals when getting a value,
16and mangling values when setting a value. This means the original env
17will see its values updated as expected, regardless of mangling rules."
18  (setmetatable []
19                {:__index (fn [_ key]
20                            (if (= (type key) :string)
21                                (. env (compiler.global-unmangling key))
22                                (. env key)))
23                 :__newindex (fn [_ key value]
24                               (if (= (type key) :string)
25                                   (tset env (compiler.global-unmangling key)
26                                         value)
27                                   (tset env key value)))
28                 ;; manually in 5.1
29                 :__pairs (fn []
30                            (fn putenv [k v]
31                              (values (if (= (type k) :string)
32                                          (compiler.global-unmangling k)
33                                          k)
34                                      v))
35
36                            (values next (utils.kvmap env putenv) nil))}))
37
38(fn current-global-names [?env]
39  ;; if there's a metatable on ?env, we need to make sure it's one that has a
40  ;; __pairs metamethod, otherwise we give up entirely on globals checking.
41  (let [mt (match (getmetatable ?env)
42             ;; newer lua versions know about __pairs natively but not 5.1
43             {:__pairs mtpairs} (collect [k v (mtpairs ?env)] (values k v))
44             nil (or ?env _G))]
45    (and mt (utils.kvmap mt compiler.global-unmangling))))
46
47(fn load-code [code ?env ?filename]
48  "Load Lua code with an environment in all recent Lua versions"
49  (let [env (or ?env (rawget _G :_ENV) _G)]
50    (if (and (rawget _G :setfenv) (rawget _G :loadstring))
51        (let [f (assert (_G.loadstring code ?filename))]
52          (doto f (setfenv env)))
53        (assert (load code ?filename :t env)))))
54
55(fn doc* [tgt name]
56  "Return a docstring for tgt."
57  (if (not tgt)
58      (.. name " not found")
59      (let [docstring (-> (: compiler.metadata :get tgt :fnl/docstring)
60                          (or "#<undocumented>")
61                          (: :gsub "\n$" "")
62                          (: :gsub "\n" "\n  "))
63            mt (getmetatable tgt)]
64        (if (or (= (type tgt) :function)
65                (and (= (type mt) :table) (= (type (. mt :__call)) :function)))
66            (let [arglist (table.concat (or (: compiler.metadata :get tgt
67                                               :fnl/arglist)
68                                            ["#<unknown-arguments>"])
69                                        " ")]
70              (string.format "(%s%s%s)\n  %s" name
71                             (if (> (length arglist) 0) " " "") arglist
72                             docstring))
73            (string.format "%s\n  %s" name docstring)))))
74
75;; TODO: replace this with using the special fn's own docstring
76(fn doc-special [name arglist docstring body-form?]
77  "Add a docstring to a special form."
78  (tset compiler.metadata (. SPECIALS name)
79        {:fnl/arglist arglist :fnl/docstring docstring :fnl/body-form? body-form?}))
80
81(fn compile-do [ast scope parent ?start]
82  "Compile a list of forms for side effects."
83  (let [start (or ?start 2)
84        len (length ast)
85        sub-scope (compiler.make-scope scope)]
86    (for [i start len]
87      (compiler.compile1 (. ast i) sub-scope parent {:nval 0}))))
88
89(fn SPECIALS.do [ast scope parent opts ?start ?chunk ?sub-scope ?pre-syms]
90  "Implements a do statement, starting at the 'start'-th element.
91By default, start is 2."
92  (let [start (or ?start 2)
93        sub-scope (or ?sub-scope (compiler.make-scope scope))
94        chunk (or ?chunk [])
95        len (length ast)
96        retexprs {:returned true}]
97    (fn compile-body [outer-target outer-tail outer-retexprs]
98      (if (< len start)
99          ;; In the unlikely event we do a do with no arguments
100          (compiler.compile1 nil sub-scope chunk
101                             {:tail outer-tail :target outer-target})
102          ;; There will be side-effects
103          (for [i start len]
104            (let [subopts {:nval (or (and (not= i len) 0) opts.nval)
105                           :tail (or (and (= i len) outer-tail) nil)
106                           :target (or (and (= i len) outer-target) nil)}
107                  _ (utils.propagate-options opts subopts)
108                  subexprs (compiler.compile1 (. ast i) sub-scope chunk subopts)]
109              (when (not= i len)
110                (compiler.keep-side-effects subexprs parent nil (. ast i))))))
111      (compiler.emit parent chunk ast)
112      (compiler.emit parent :end ast)
113      (utils.hook :do ast sub-scope)
114      (or outer-retexprs retexprs))
115
116    ;; See if we need special handling to get the return values of the do block
117    (if (or opts.target (= opts.nval 0) opts.tail)
118        (do
119          (compiler.emit parent :do ast)
120          (compile-body opts.target opts.tail))
121        opts.nval
122        ;; generate a local target
123        (let [syms []]
124          (for [i 1 opts.nval]
125            (let [s (or (and ?pre-syms (. ?pre-syms i)) (compiler.gensym scope))]
126              (tset syms i s)
127              (tset retexprs i (utils.expr s :sym))))
128          (let [outer-target (table.concat syms ", ")]
129            (compiler.emit parent (string.format "local %s" outer-target) ast)
130            (compiler.emit parent :do ast)
131            (compile-body outer-target opts.tail)))
132        ;; we will use an IIFE for the do
133        (let [fname (compiler.gensym scope)
134              fargs (if scope.vararg "..." "")]
135          (compiler.emit parent
136                         (string.format "local function %s(%s)" fname fargs) ast)
137          (compile-body nil true
138                        (utils.expr (.. fname "(" fargs ")") :statement))))))
139
140(doc-special :do ["..."] "Evaluate multiple forms; return last value." true)
141
142(fn SPECIALS.values [ast scope parent]
143  "Unlike most expressions and specials, 'values' resolves with multiple
144values, one for each argument, allowing multiple return values. The last
145expression can return multiple arguments as well, allowing for more than
146the number of expected arguments."
147  (let [len (length ast)
148        exprs []]
149    (for [i 2 len]
150      (let [subexprs (compiler.compile1 (. ast i) scope parent
151                                        {:nval (and (not= i len) 1)})]
152        (table.insert exprs (. subexprs 1))
153        (when (= i len)
154          (for [j 2 (length subexprs)]
155            (table.insert exprs (. subexprs j))))))
156    exprs))
157
158(doc-special :values ["..."]
159             "Return multiple values from a function. Must be in tail position.")
160
161(fn deep-tostring [x key?]
162  "Tostring for literal tables created with {} or [].
163Recursively transforms tables into one-line string representation.
164Main purpose to print function argument list in docstring."
165  (if (utils.sequence? x)
166      (.. "[" (table.concat (icollect [_ v (ipairs x)]
167                              (deep-tostring v))
168                            " ") "]")
169      (utils.table? x)
170      (.. "{" (table.concat (icollect [k v (pairs x)]
171                              (.. (deep-tostring k true) " "
172                                  (deep-tostring v)))
173                            " ") "}")
174      (and key? (= (type x) :string) (x:find "^[-%w?\\^_!$%&*+./@:|<=>]+$"))
175      (.. ":" x)
176      (= (type x) :string)
177      (-> (string.format "%q" x)
178          (: :gsub "\\\"" "\\\\\"")
179          (: :gsub "\"" "\\\""))
180      (tostring x)))
181
182(fn set-fn-metadata [arg-list docstring parent fn-name]
183  (when utils.root.options.useMetadata
184    (let [args (utils.map arg-list #(: "\"%s\"" :format (deep-tostring $)))
185          meta-fields ["\"fnl/arglist\"" (.. "{" (table.concat args ", ") "}")]]
186      (when docstring
187        (table.insert meta-fields "\"fnl/docstring\"")
188        (table.insert meta-fields (.. "\""
189                                      (-> docstring
190                                          (: :gsub "%s+$" "")
191                                          (: :gsub "\\" "\\\\")
192                                          (: :gsub "\n" "\\n")
193                                          (: :gsub "\"" "\\\""))
194                                      "\"")))
195      (let [meta-str (: "require(\"%s\").metadata" :format
196                        (or utils.root.options.moduleName :fennel))]
197        (compiler.emit parent
198                       (: "pcall(function() %s:setall(%s, %s) end)" :format
199                          meta-str fn-name (table.concat meta-fields ", ")))))))
200
201(fn get-fn-name [ast scope fn-name multi]
202  (if (and fn-name (not= (. fn-name 1) :nil))
203      (values (if (not multi)
204                  (compiler.declare-local fn-name [] scope ast)
205                  (. (compiler.symbol-to-expression fn-name scope) 1))
206              (not multi) 3)
207      (values nil true 2)))
208
209(fn compile-named-fn [ast f-scope f-chunk parent index fn-name local?
210                      arg-name-list arg-list docstring]
211  (for [i (+ index 1) (length ast)]
212    (compiler.compile1 (. ast i) f-scope f-chunk
213                       {:nval (or (and (not= i (length ast)) 0) nil)
214                        :tail (= i (length ast))}))
215  (compiler.emit parent
216                 (string.format (if local? "local function %s(%s)"
217                                    "%s = function(%s)")
218                                fn-name (table.concat arg-name-list ", "))
219                 ast)
220  (compiler.emit parent f-chunk ast)
221  (compiler.emit parent :end ast)
222  (set-fn-metadata arg-list docstring parent fn-name)
223  (utils.hook :fn ast f-scope)
224  (utils.expr fn-name :sym))
225
226(fn compile-anonymous-fn [ast f-scope f-chunk parent index
227                          arg-name-list arg-list docstring scope]
228  ;; TODO: eventually compile this to an actual function value instead of
229  ;; binding it to a local and using the symbol. the difficulty here is that
230  ;; a function is a chunk with many lines, and the current representation of
231  ;; an expr can only be a string, making it difficult to pass around without
232  ;; losing line numbering information.
233  (let [fn-name (compiler.gensym scope)]
234    (compile-named-fn ast f-scope f-chunk parent index fn-name true
235                      arg-name-list arg-list docstring)))
236
237(fn SPECIALS.fn [ast scope parent]
238  (let [f-scope (doto (compiler.make-scope scope)
239                  (tset :vararg false))
240        f-chunk []
241        fn-sym (utils.sym? (. ast 2))
242        multi (and fn-sym (utils.multi-sym? (. fn-sym 1)))
243        (fn-name local? index) (get-fn-name ast scope fn-sym multi)
244        arg-list (compiler.assert (utils.table? (. ast index))
245                                  "expected parameters table" ast)]
246    (compiler.assert (or (not multi) (not multi.multi-sym-method-call))
247                     (.. "unexpected multi symbol " (tostring fn-name)) fn-sym)
248
249    (fn get-arg-name [arg]
250      (if (utils.varg? arg)
251          (do
252            (compiler.assert (= arg (. arg-list (length arg-list)))
253                             "expected vararg as last parameter" ast)
254            (set f-scope.vararg true)
255            "...")
256          (and (utils.sym? arg) (not= (tostring arg) :nil)
257               (not (utils.multi-sym? (tostring arg))))
258          (compiler.declare-local arg [] f-scope ast)
259          (utils.table? arg)
260          (let [raw (utils.sym (compiler.gensym scope))
261                declared (compiler.declare-local raw [] f-scope ast)]
262            (compiler.destructure arg raw ast f-scope f-chunk
263                                  {:declaration true
264                                   :nomulti true
265                                   :symtype :arg})
266            declared)
267          (compiler.assert false
268                           (: "expected symbol for function parameter: %s"
269                              :format (tostring arg))
270                           (. ast index))))
271
272    (let [arg-name-list (utils.map arg-list get-arg-name)
273          (index docstring) (if (and (= (type (. ast (+ index 1))) :string)
274                                     (< (+ index 1) (length ast)))
275                                (values (+ index 1) (. ast (+ index 1)))
276                                (values index nil))]
277      (if fn-name
278          (compile-named-fn ast f-scope f-chunk parent index fn-name local?
279                            arg-name-list arg-list docstring)
280          (compile-anonymous-fn ast f-scope f-chunk parent index
281                                arg-name-list arg-list docstring scope)))))
282
283(doc-special :fn [:name? :args :docstring? "..."]
284             "Function syntax. May optionally include a name and docstring.
285If a name is provided, the function will be bound in the current scope.
286When called with the wrong number of args, excess args will be discarded
287and lacking args will be nil, use lambda for arity-checked functions." true)
288
289;; FORBIDDEN KNOWLEDGE:
290;; (lua "print('hello!')") -> prints hello, evaluates to nil
291;; (lua "print 'hello!'" "10") -> prints hello, evaluates to the number 10
292;; (lua nil "{1,2,3}") -> Evaluates to a table literal
293(fn SPECIALS.lua [ast _ parent]
294  (compiler.assert (or (= (length ast) 2) (= (length ast) 3))
295                   "expected 1 or 2 arguments" ast)
296  (when (not= :nil (-?> (utils.sym? (. ast 2)) tostring))
297    (table.insert parent {: ast :leaf (tostring (. ast 2))}))
298  (when (not= :nil (-?> (utils.sym? (. ast 3)) tostring))
299    (tostring (. ast 3))))
300
301(fn dot [ast scope parent]
302  "Table lookup; equivalent to tbl[] in Lua."
303  (compiler.assert (< 1 (length ast)) "expected table argument" ast)
304  (let [len (length ast)
305        [lhs] (compiler.compile1 (. ast 2) scope parent {:nval 1})]
306    (if (= len 2)
307        (tostring lhs)
308        (let [indices []]
309          (for [i 3 len]
310            (let [index (. ast i)]
311              (if (and (= (type index) :string)
312                       (utils.valid-lua-identifier? index))
313                  (table.insert indices (.. "." index))
314                  (let [[index] (compiler.compile1 index scope parent {:nval 1})]
315                    (table.insert indices (.. "[" (tostring index) "]"))))))
316          ;; Extra parens are needed unless the target is a table literal
317          (if (or (: (tostring lhs) :find "[{\"0-9]") (= :nil (tostring lhs)))
318              (.. "(" (tostring lhs) ")" (table.concat indices))
319              (.. (tostring lhs) (table.concat indices)))))))
320
321(tset SPECIALS "." dot)
322
323(doc-special "." [:tbl :key1 "..."]
324             "Look up key1 in tbl table. If more args are provided, do a nested lookup.")
325
326(fn SPECIALS.global [ast scope parent]
327  (compiler.assert (= (length ast) 3) "expected name and value" ast)
328  (compiler.destructure (. ast 2) (. ast 3) ast scope parent
329                        {:forceglobal true :nomulti true :symtype :global})
330  nil)
331
332(doc-special :global [:name :val] "Set name as a global with val.")
333
334(fn SPECIALS.set [ast scope parent]
335  (compiler.assert (= (length ast) 3) "expected name and value" ast)
336  (compiler.destructure (. ast 2) (. ast 3) ast scope parent
337                        {:noundef true :symtype :set})
338  nil)
339
340(doc-special :set [:name :val]
341             "Set a local variable to a new value. Only works on locals using var.")
342
343(fn set-forcibly!* [ast scope parent]
344  (compiler.assert (= (length ast) 3) "expected name and value" ast)
345  (compiler.destructure (. ast 2) (. ast 3) ast scope parent
346                        {:forceset true :symtype :set})
347  nil)
348
349(tset SPECIALS :set-forcibly! set-forcibly!*)
350
351(fn local* [ast scope parent]
352  (compiler.assert (= (length ast) 3) "expected name and value" ast)
353  (compiler.destructure (. ast 2) (. ast 3) ast scope parent
354                        {:declaration true :nomulti true :symtype :local})
355  nil)
356
357(tset SPECIALS :local local*)
358
359(doc-special :local [:name :val] "Introduce new top-level immutable local.")
360
361(fn SPECIALS.var [ast scope parent]
362  (compiler.assert (= (length ast) 3) "expected name and value" ast)
363  (compiler.destructure (. ast 2) (. ast 3) ast scope parent
364                        {:declaration true
365                         :isvar true
366                         :nomulti true
367                         :symtype :var})
368  nil)
369
370(doc-special :var [:name :val] "Introduce new mutable local.")
371
372(fn kv? [t] (. (icollect [k (pairs t)] (if (not (= :number (type k))) k)) 1))
373
374(fn SPECIALS.let [ast scope parent opts]
375  (let [bindings (. ast 2)
376        pre-syms []]
377    (compiler.assert (and (utils.table? bindings) (not (kv? bindings)))
378                     "expected binding sequence" bindings)
379    (compiler.assert (= (% (length bindings) 2) 0)
380                     "expected even number of name/value bindings" (. ast 2))
381    (compiler.assert (>= (length ast) 3) "expected body expression" (. ast 1))
382    ;; we have to gensym the binding for the let body's return value before
383    ;; compiling the binding vector, otherwise there's a possibility to conflict
384    (for [_ 1 (or opts.nval 0)]
385      (table.insert pre-syms (compiler.gensym scope)))
386    (let [sub-scope (compiler.make-scope scope)
387          sub-chunk []]
388      (for [i 1 (length bindings) 2]
389        (compiler.destructure (. bindings i) (. bindings (+ i 1)) ast sub-scope
390                              sub-chunk
391                              {:declaration true :nomulti true :symtype :let}))
392      (SPECIALS.do ast scope parent opts 3 sub-chunk sub-scope pre-syms))))
393
394(doc-special :let ["[name1 val1 ... nameN valN]" "..."]
395             "Introduces a new scope in which a given set of local bindings are used."
396             true)
397
398(fn get-prev-line [parent]
399  (if (= :table (type parent))
400      (get-prev-line (or parent.leaf (. parent (length parent))))
401      (or parent "")))
402
403(fn disambiguate? [rootstr parent]
404  (or (rootstr:match "^{")
405      (match (get-prev-line parent)
406        prev-line (prev-line:match "%)$"))))
407
408(fn SPECIALS.tset [ast scope parent]
409  "For setting items in a table."
410  (compiler.assert (> (length ast) 3)
411                   "expected table, key, and value arguments" ast)
412  (let [root (. (compiler.compile1 (. ast 2) scope parent {:nval 1}) 1)
413        keys []]
414    (for [i 3 (- (length ast) 1)]
415      (let [[key] (compiler.compile1 (. ast i) scope parent {:nval 1})]
416        (table.insert keys (tostring key))))
417    (let [value (. (compiler.compile1 (. ast (length ast)) scope parent
418                                      {:nval 1}) 1)
419          rootstr (tostring root)
420          fmtstr (if (disambiguate? rootstr parent)
421                     "do end (%s)[%s] = %s"
422                     "%s[%s] = %s")]
423      (compiler.emit parent
424                     (: fmtstr :format rootstr (table.concat keys "][")
425                        (tostring value)) ast))))
426
427(doc-special :tset [:tbl :key1 "..." :keyN :val]
428             "Set the value of a table field. Can take additional keys to set
429nested values, but all parents must contain an existing table.")
430
431(fn calculate-target [scope opts]
432  (if (not (or opts.tail opts.target opts.nval))
433      (values :iife true nil)
434      (and opts.nval (not= opts.nval 0) (not opts.target))
435      (let [accum []
436            target-exprs []]
437        ;; We need to create a target
438        (for [i 1 opts.nval]
439          (let [s (compiler.gensym scope)]
440            (tset accum i s)
441            (tset target-exprs i (utils.expr s :sym))))
442        (values :target opts.tail (table.concat accum ", ") target-exprs))
443      (values :none opts.tail opts.target)))
444
445;; TODO: refactor; too long!
446(fn if* [ast scope parent opts]
447  (compiler.assert (< 2 (length ast)) "expected condition and body" ast)
448  (let [do-scope (compiler.make-scope scope)
449        branches []
450        (wrapper inner-tail inner-target target-exprs) (calculate-target scope
451                                                                         opts)
452        body-opts {:nval opts.nval :tail inner-tail :target inner-target}]
453    (fn compile-body [i]
454      (let [chunk []
455            cscope (compiler.make-scope do-scope)]
456        (compiler.keep-side-effects (compiler.compile1 (. ast i) cscope chunk
457                                                       body-opts)
458                                    chunk nil (. ast i))
459        {: chunk :scope cscope}))
460
461    ;; Implicit else becomes nil
462    (when (= 1 (% (length ast) 2))
463      (table.insert ast (utils.sym :nil)))
464
465    (for [i 2 (- (length ast) 1) 2]
466      (let [condchunk []
467            res (compiler.compile1 (. ast i) do-scope condchunk {:nval 1})
468            cond (. res 1)
469            branch (compile-body (+ i 1))]
470        (set branch.cond cond)
471        (set branch.condchunk condchunk)
472        (set branch.nested (and (not= i 2) (= (next condchunk nil) nil)))
473        (table.insert branches branch)))
474    ;; Emit code
475    (let [else-branch (compile-body (length ast))
476          s (compiler.gensym scope)
477          buffer []]
478      (var last-buffer buffer)
479      (for [i 1 (length branches)]
480        (let [branch (. branches i)
481              fstr (if (not branch.nested) "if %s then" "elseif %s then")
482              cond (tostring branch.cond)
483              cond-line (: fstr :format cond)]
484          (if branch.nested
485              (compiler.emit last-buffer branch.condchunk ast)
486              (each [_ v (ipairs branch.condchunk)]
487                (compiler.emit last-buffer v ast)))
488          (compiler.emit last-buffer cond-line ast)
489          (compiler.emit last-buffer branch.chunk ast)
490          (if (= i (length branches))
491              (do
492                (compiler.emit last-buffer :else ast)
493                (compiler.emit last-buffer else-branch.chunk ast)
494                (compiler.emit last-buffer :end ast))
495              (not (. (. branches (+ i 1)) :nested))
496              (let [next-buffer []]
497                (compiler.emit last-buffer :else ast)
498                (compiler.emit last-buffer next-buffer ast)
499                (compiler.emit last-buffer :end ast)
500                (set last-buffer next-buffer)))))
501      ;; Emit if
502      (if (= wrapper :iife)
503          (let [iifeargs (or (and scope.vararg "...") "")]
504            (compiler.emit parent
505                           (: "local function %s(%s)" :format (tostring s)
506                              iifeargs) ast)
507            (compiler.emit parent buffer ast)
508            (compiler.emit parent :end ast)
509            (utils.expr (: "%s(%s)" :format (tostring s) iifeargs) :statement))
510          (= wrapper :none) ; Splice result right into code
511          (do
512            (for [i 1 (length buffer)]
513              (compiler.emit parent (. buffer i) ast))
514            {:returned true})
515          ;; wrapper is target
516          (do
517            (compiler.emit parent (: "local %s" :format inner-target) ast)
518            (for [i 1 (length buffer)]
519              (compiler.emit parent (. buffer i) ast))
520            target-exprs)))))
521
522(tset SPECIALS :if if*)
523
524(doc-special :if [:cond1 :body1 "..." :condN :bodyN]
525             "Conditional form.
526Takes any number of condition/body pairs and evaluates the first body where
527the condition evaluates to truthy. Similar to cond in other lisps.")
528
529(fn remove-until-condition [bindings]
530  (when (= :until (. bindings (- (length bindings) 1)))
531    (table.remove bindings (- (length bindings) 1))
532    (table.remove bindings)))
533
534(fn compile-until [condition scope chunk]
535  (when condition
536    (let [[condition-lua] (compiler.compile1 condition scope chunk {:nval 1})]
537      (compiler.emit chunk (: "if %s then break end" :format
538                              (tostring condition-lua))
539                     (utils.expr condition :expression)))))
540
541(fn SPECIALS.each [ast scope parent]
542  (compiler.assert (>= (length ast) 3) "expected body expression" (. ast 1))
543  (let [binding (compiler.assert (utils.table? (. ast 2))
544                                 "expected binding table" ast)
545        _ (compiler.assert (<= 2 (length binding))
546                           "expected binding and iterator" binding)
547        until-condition (remove-until-condition binding)
548        iter (table.remove binding (length binding))
549        ; last item is iterator call
550        destructures []
551        new-manglings []
552        sub-scope (compiler.make-scope scope)]
553    (fn destructure-binding [v]
554      (compiler.assert (not= :string (type v))
555                       (.. "unexpected iterator clause " (tostring v)) binding)
556      (if (utils.sym? v)
557          (compiler.declare-local v [] sub-scope ast new-manglings)
558          (let [raw (utils.sym (compiler.gensym sub-scope))]
559            (tset destructures raw v)
560            (compiler.declare-local raw [] sub-scope ast))))
561
562    (let [bind-vars (utils.map binding destructure-binding)
563          vals (compiler.compile1 iter scope parent)
564          val-names (utils.map vals tostring)
565          chunk []]
566      (compiler.emit parent
567                     (: "for %s in %s do" :format (table.concat bind-vars ", ")
568                        (table.concat val-names ", ")) ast)
569      (each [raw args (utils.stablepairs destructures)]
570        (compiler.destructure args raw ast sub-scope chunk
571                              {:declaration true :nomulti true :symtype :each}))
572      (compiler.apply-manglings sub-scope new-manglings ast)
573      (compile-until until-condition sub-scope chunk)
574      (compile-do ast sub-scope chunk 3)
575      (compiler.emit parent chunk ast)
576      (compiler.emit parent :end ast))))
577
578(doc-special :each ["[key value (iterator)]" "..."]
579             "Runs the body once for each set of values provided by the given iterator.
580Most commonly used with ipairs for sequential tables or pairs for  undefined
581order, but can be used with any iterator." true)
582
583(fn while* [ast scope parent]
584  (let [len1 (length parent)
585        condition (. (compiler.compile1 (. ast 2) scope parent {:nval 1}) 1)
586        len2 (length parent)
587        sub-chunk []]
588    (if (not= len1 len2)
589        ;; compound condition; move new compilation to subchunk
590        (do
591          (for [i (+ len1 1) len2]
592            (table.insert sub-chunk (. parent i))
593            (tset parent i nil))
594          (compiler.emit parent "while true do" ast)
595          (compiler.emit sub-chunk
596                         (: "if not %s then break end" :format (. condition 1))
597                         ast))
598        ;; simple condition
599        (compiler.emit parent (.. "while " (tostring condition) " do") ast))
600    (compile-do ast (compiler.make-scope scope) sub-chunk 3)
601    (compiler.emit parent sub-chunk ast)
602    (compiler.emit parent :end ast)))
603
604(tset SPECIALS :while while*)
605
606(doc-special :while [:condition "..."]
607             "The classic while loop. Evaluates body until a condition is non-truthy."
608             true)
609
610(fn for* [ast scope parent]
611  (let [ranges (compiler.assert (utils.table? (. ast 2))
612                                "expected binding table" ast)
613        until-condition (remove-until-condition (. ast 2))
614        binding-sym (table.remove (. ast 2) 1)
615        sub-scope (compiler.make-scope scope)
616        range-args []
617        chunk []]
618    (compiler.assert (utils.sym? binding-sym)
619                     (: "unable to bind %s %s" :format (type binding-sym)
620                        (tostring binding-sym)) (. ast 2))
621    (compiler.assert (>= (length ast) 3) "expected body expression" (. ast 1))
622    (compiler.assert (<= (length ranges) 3) "unexpected arguments" (. ranges 4))
623    (for [i 1 (math.min (length ranges) 3)]
624      (tset range-args i (tostring (. (compiler.compile1 (. ranges i) scope
625                                                         parent {:nval 1}) 1))))
626    (compiler.emit parent
627                   (: "for %s = %s do" :format
628                      (compiler.declare-local binding-sym [] sub-scope ast)
629                      (table.concat range-args ", ")) ast)
630    (compile-until until-condition sub-scope chunk)
631    (compile-do ast sub-scope chunk 3)
632    (compiler.emit parent chunk ast)
633    (compiler.emit parent :end ast)))
634
635(tset SPECIALS :for for*)
636
637(doc-special :for ["[index start stop step?]" "..."]
638             "Numeric loop construct.
639Evaluates body once for each value between start and stop (inclusive)." true)
640
641(fn native-method-call [ast _scope _parent target args]
642  "Prefer native Lua method calls when method name is a valid Lua identifier."
643  (let [[_ _ method-string] ast
644        call-string (if (or (= target.type :literal)
645                            (= target.type :varg)
646                            (= target.type :expression))
647                        "(%s):%s(%s)" "%s:%s(%s)")]
648    (utils.expr (string.format call-string (tostring target) method-string
649                               (table.concat args ", "))
650                :statement)))
651
652(fn nonnative-method-call [ast scope parent target args]
653  "When we don't have to protect against double-evaluation, it's not so bad."
654  (let [method-string (tostring (. (compiler.compile1 (. ast 3) scope parent
655                                                      {:nval 1})
656                                   1))
657        args [(tostring target) (unpack args)]]
658    (utils.expr (string.format "%s[%s](%s)" (tostring target) method-string
659                               (table.concat args ", "))
660                :statement)))
661
662(fn double-eval-protected-method-call [ast scope parent target args]
663  "When double-evaluation is a concern, we have to wrap an IIFE."
664  (let [method-string (tostring (. (compiler.compile1 (. ast 3) scope parent
665                                                      {:nval 1})
666                                   1))
667        call "(function(tgt, m, ...) return tgt[m](tgt, ...) end)(%s, %s)"]
668    (table.insert args 1 method-string)
669    (utils.expr (string.format call (tostring target) (table.concat args ", "))
670                :statement)))
671
672(fn method-call [ast scope parent]
673  (compiler.assert (< 2 (length ast)) "expected at least 2 arguments" ast)
674  (let [[target] (compiler.compile1 (. ast 2) scope parent {:nval 1})
675        args []]
676    (for [i 4 (length ast)]
677      (let [subexprs (compiler.compile1 (. ast i) scope parent
678                                        {:nval (if (not= i (length ast)) 1)})]
679        (utils.map subexprs tostring args)))
680    (if (and (= (type (. ast 3)) :string)
681             (utils.valid-lua-identifier? (. ast 3)))
682        (native-method-call ast scope parent target args)
683        (= target.type :sym)
684        (nonnative-method-call ast scope parent target args)
685        ;; When the target is an expression, we can't use the naive
686        ;; nonnative-method-call approach, because it will cause the target
687        ;; to be evaluated twice. This is fine if it's a symbol but if it's
688        ;; the result of a function call, that function could have side-effects.
689        ;; See test-short-circuit in test/misc.fnl for an example of the problem.
690        (double-eval-protected-method-call ast scope parent target args))))
691
692(tset SPECIALS ":" method-call)
693
694(doc-special ":" [:tbl :method-name "..."] "Call the named method on tbl with the provided args.
695Method name doesn't have to be known at compile-time; if it is, use
696(tbl:method-name ...) instead.")
697
698(fn SPECIALS.comment [ast _ parent]
699  (let [els []]
700    (for [i 2 (length ast)]
701      (table.insert els (view (. ast i) {:one-line? true})))
702    (compiler.emit parent (.. "--[[ " (table.concat els " ") " ]]--") ast)))
703
704(doc-special :comment ["..."] "Comment which will be emitted in Lua output." true)
705
706(fn hashfn-max-used [f-scope i max]
707  (let [max (if (. f-scope.symmeta (.. "$" i) :used) i max)]
708    (if (< i 9)
709        (hashfn-max-used f-scope (+ i 1) max)
710        max)))
711
712(fn SPECIALS.hashfn [ast scope parent]
713  (compiler.assert (= (length ast) 2) "expected one argument" ast)
714  (let [f-scope (doto (compiler.make-scope scope)
715                  (tset :vararg false)
716                  (tset :hashfn true))
717        f-chunk []
718        name (compiler.gensym scope)
719        symbol (utils.sym name)
720        args []]
721    (compiler.declare-local symbol [] scope ast)
722    (for [i 1 9]
723      (tset args i (compiler.declare-local (utils.sym (.. "$" i)) [] f-scope
724                                           ast)))
725    ;; recursively walk the AST, transforming $... into ...
726
727    (fn walker [idx node parent-node]
728      (if (and (utils.sym? node) (= (tostring node) "$..."))
729          (do
730            (tset parent-node idx (utils.varg))
731            (set f-scope.vararg true))
732          (or (utils.list? node) (utils.table? node))))
733
734    (utils.walk-tree (. ast 2) walker)
735    ;; compile body
736    (compiler.compile1 (. ast 2) f-scope f-chunk {:tail true})
737    (let [max-used (hashfn-max-used f-scope 1 0)]
738      (when f-scope.vararg
739        (compiler.assert (= max-used 0)
740                         "$ and $... in hashfn are mutually exclusive" ast))
741      (let [arg-str (if f-scope.vararg
742                        (tostring (utils.varg))
743                        (table.concat args ", " 1 max-used))]
744        (compiler.emit parent
745                       (string.format "local function %s(%s)" name arg-str) ast)
746        (compiler.emit parent f-chunk ast)
747        (compiler.emit parent :end ast)
748        (utils.expr name :sym)))))
749
750(doc-special :hashfn ["..."]
751             "Function literal shorthand; args are either $... OR $1, $2, etc.")
752
753(fn arithmetic-special [name zero-arity unary-prefix ast scope parent]
754  (let [len (length ast) operands []
755        padded-op (.. " " name " ")]
756    (for [i 2 len]
757      (let [subexprs (compiler.compile1 (. ast i) scope parent)]
758        (if (= i len)
759            ;; last arg gets all its exprs but everyone else only gets one
760            (utils.map subexprs tostring operands)
761            (table.insert operands (tostring (. subexprs 1))))))
762    (match (length operands)
763      0 (utils.expr (doto zero-arity
764                      (compiler.assert "Expected more than 0 arguments" ast))
765                    :literal)
766      1 (if unary-prefix
767            (.. "(" unary-prefix padded-op (. operands 1) ")")
768            (. operands 1))
769      _ (.. "(" (table.concat operands padded-op) ")"))))
770
771(fn define-arithmetic-special [name zero-arity unary-prefix ?lua-name]
772  (tset SPECIALS name (partial arithmetic-special (or ?lua-name name) zero-arity
773                               unary-prefix))
774  (doc-special name [:a :b "..."]
775               "Arithmetic operator; works the same as Lua but accepts more arguments."))
776
777(define-arithmetic-special "+" :0)
778(define-arithmetic-special ".." "''")
779(define-arithmetic-special "^")
780(define-arithmetic-special "-" nil "")
781(define-arithmetic-special "*" :1)
782(define-arithmetic-special "%")
783(define-arithmetic-special "/" nil :1)
784(define-arithmetic-special "//" nil :1)
785
786(fn SPECIALS.or [ast scope parent]
787  ;; and/or have nval=nil in order to trigger IIFE so they can short-circuit
788  (arithmetic-special :or :false nil ast scope parent))
789
790(fn SPECIALS.and [ast scope parent]
791  (arithmetic-special :and :true nil ast scope parent))
792
793(doc-special :and [:a :b "..."]
794             "Boolean operator; works the same as Lua but accepts more arguments.")
795
796(doc-special :or [:a :b "..."]
797             "Boolean operator; works the same as Lua but accepts more arguments.")
798
799(fn bitop-special [native-name lib-name zero-arity unary-prefix ast scope parent]
800  (if (= (length ast) 1)
801      (compiler.assert zero-arity "Expected more than 0 arguments." ast)
802      (let [len (length ast)
803            operands []
804            padded-native-name (.. " " native-name " ")
805            prefixed-lib-name (.. "bit." lib-name)]
806        (for [i 2 len]
807          (let [subexprs (compiler.compile1 (. ast i) scope parent
808                                            {:nval (if (not= i len) 1)})]
809            (utils.map subexprs tostring operands)))
810        (if (= (length operands) 1)
811            (if utils.root.options.useBitLib
812                (.. prefixed-lib-name "(" unary-prefix ", " (. operands 1) ")")
813                (.. "(" unary-prefix padded-native-name (. operands 1) ")"))
814            (if utils.root.options.useBitLib
815                (.. prefixed-lib-name "(" (table.concat operands ", ") ")")
816                (.. "(" (table.concat operands padded-native-name) ")"))))))
817
818(fn define-bitop-special [name zero-arity unary-prefix native]
819  (tset SPECIALS name (partial bitop-special native name zero-arity unary-prefix)))
820
821(define-bitop-special :lshift nil :1 "<<")
822(define-bitop-special :rshift nil :1 ">>")
823(define-bitop-special :band :0 :0 "&")
824(define-bitop-special :bor :0 :0 "|")
825(define-bitop-special :bxor :0 :0 "~")
826
827(doc-special :lshift [:x :n]
828             "Bitwise logical left shift of x by n bits.
829Only works in Lua 5.3+ or LuaJIT with the --use-bit-lib flag.")
830
831(doc-special :rshift [:x :n]
832             "Bitwise logical right shift of x by n bits.
833Only works in Lua 5.3+ or LuaJIT with the --use-bit-lib flag.")
834
835(doc-special :band [:x1 :x2 "..."] "Bitwise AND of any number of arguments.
836Only works in Lua 5.3+ or LuaJIT with the --use-bit-lib flag.")
837
838(doc-special :bor [:x1 :x2 "..."] "Bitwise OR of any number of arguments.
839Only works in Lua 5.3+ or LuaJIT with the --use-bit-lib flag.")
840
841(doc-special :bxor [:x1 :x2 "..."] "Bitwise XOR of any number of arguments.
842Only works in Lua 5.3+ or LuaJIT with the --use-bit-lib flag.")
843
844(doc-special ".." [:a :b "..."]
845             "String concatenation operator; works the same as Lua but accepts more arguments.")
846
847(fn native-comparator [op [_ lhs-ast rhs-ast] scope parent]
848  "Naively compile a binary comparison to Lua."
849  (let [[lhs] (compiler.compile1 lhs-ast scope parent {:nval 1})
850        [rhs] (compiler.compile1 rhs-ast scope parent {:nval 1})]
851    (string.format "(%s %s %s)" (tostring lhs) op (tostring rhs))))
852
853(fn double-eval-protected-comparator [op chain-op ast scope parent]
854  "Compile a multi-arity comparison to a binary Lua comparison."
855  (let [arglist []
856        comparisons []
857        vals []
858        chain (string.format " %s " (or chain-op :and))]
859    (for [i 2 (length ast)]
860      (table.insert arglist (tostring (compiler.gensym scope)))
861      (table.insert vals (tostring (. (compiler.compile1 (. ast i) scope parent
862                                                         {:nval 1})
863                                      1))))
864    (for [i 1 (- (length arglist) 1)]
865      (table.insert comparisons
866                    (string.format "(%s %s %s)" (. arglist i) op
867                                   (. arglist (+ i 1)))))
868    ;; The function call here introduces some overhead, but it is the only way
869    ;; to compile this safely while preventing both double-evaluation of
870    ;; side-effecting values and early evaluation of values which should never
871    ;; happen in the case of a short-circuited call. See test-short-circuit in
872    ;; test/misc.fnl for an example of the problem.
873    (string.format "(function(%s) return %s end)(%s)"
874                   (table.concat arglist ",") (table.concat comparisons chain)
875                   (table.concat vals ","))))
876
877(fn define-comparator-special [name ?lua-op ?chain-op]
878  (let [op (or ?lua-op name)]
879    (fn opfn [ast scope parent]
880      (compiler.assert (< 2 (length ast)) "expected at least two arguments" ast)
881      (if (= 3 (length ast))
882          (native-comparator op ast scope parent)
883          (double-eval-protected-comparator op ?chain-op ast scope parent)))
884
885    (tset SPECIALS name opfn))
886  (doc-special name [:a :b "..."]
887               "Comparison operator; works the same as Lua but accepts more arguments."))
888
889(define-comparator-special ">")
890(define-comparator-special "<")
891(define-comparator-special ">=")
892(define-comparator-special "<=")
893(define-comparator-special "=" "==")
894(define-comparator-special :not= "~=" :or)
895
896(fn define-unary-special [op ?realop]
897  (fn opfn [ast scope parent]
898    (compiler.assert (= (length ast) 2) "expected one argument" ast)
899    (let [tail (compiler.compile1 (. ast 2) scope parent {:nval 1})]
900      (.. (or ?realop op) (tostring (. tail 1)))))
901
902  (tset SPECIALS op opfn))
903
904(define-unary-special :not "not ")
905(doc-special :not [:x] "Logical operator; works the same as Lua.")
906(define-unary-special :bnot "~")
907(doc-special :bnot [:x] "Bitwise negation; only works in Lua 5.3+ or LuaJIT with the --use-bit-lib flag.")
908(define-unary-special :length "#")
909(doc-special :length [:x] "Returns the length of a table or string.")
910
911;; backwards-compatibility aliases
912(tset SPECIALS "~=" (. SPECIALS :not=))
913(tset SPECIALS "#" (. SPECIALS :length))
914
915(fn SPECIALS.quote [ast scope parent]
916  (compiler.assert (= (length ast) 2) "expected one argument" ast)
917  (var (runtime this-scope) (values true scope))
918  (while this-scope
919    (set this-scope this-scope.parent)
920    (when (= this-scope compiler.scopes.compiler)
921      (set runtime false)))
922  (compiler.do-quote (. ast 2) scope parent runtime))
923
924(doc-special :quote [:x]
925             "Quasiquote the following form. Only works in macro/compiler scope.")
926
927;; This is the compile-env equivalent of package.loaded. It's used by
928;; require-macros and import-macros, but also by require when used from within
929;; default compiler scope.
930(local macro-loaded {})
931
932(fn safe-getmetatable [tbl]
933  (let [mt (getmetatable tbl)]
934    ;; we can't let the string metatable leak
935    (assert (not= mt (getmetatable "")) "Illegal metatable access!")
936    mt))
937
938;; Circularity
939(var safe-require nil)
940
941(fn safe-compiler-env []
942  {:table (utils.copy table)
943   :math (utils.copy math)
944   :string (utils.copy string)
945   : pairs : ipairs : select : tostring : tonumber :bit (rawget _G :bit)
946   : pcall : xpcall : next : print : type : assert : error
947   : setmetatable :getmetatable safe-getmetatable :require safe-require
948   :rawlen (rawget _G :rawlen) : rawget : rawset : rawequal : _VERSION})
949
950(fn combined-mt-pairs [env]
951  (let [combined {}
952        {: __index} (getmetatable env)]
953    (when (= :table (type __index))
954      (each [k v (pairs __index)]
955        (tset combined k v)))
956    (each [k v (values next env nil)]
957      (tset combined k v))
958    (values next combined nil)))
959
960(fn make-compiler-env [ast scope parent ?opts]
961  (let [provided (match (or ?opts utils.root.options)
962                   {:compiler-env :strict} (safe-compiler-env)
963                   {: compilerEnv} compilerEnv
964                   {: compiler-env} compiler-env
965                   _ (safe-compiler-env false))
966        env {:_AST ast
967             :_CHUNK parent
968             :_IS_COMPILER true
969             :_SCOPE scope
970             :_SPECIALS compiler.scopes.global.specials
971             :_VARARG (utils.varg)
972             : macro-loaded
973             : unpack
974             :assert-compile compiler.assert
975             : view
976             :version utils.version
977             :metadata compiler.metadata
978             ;; AST functions
979             :list utils.list :list? utils.list? :table? utils.table?
980             :sequence utils.sequence :sequence? utils.sequence?
981             :sym utils.sym :sym? utils.sym? :multi-sym? utils.multi-sym?
982             :comment utils.comment :comment? utils.comment? :varg? utils.varg?
983             ;; scoping functions
984             :gensym (fn [base]
985                       (utils.sym (compiler.gensym (or compiler.scopes.macro
986                                                       scope)
987                                                   base)))
988             :get-scope (fn []
989                          compiler.scopes.macro)
990             :in-scope? (fn [symbol]
991                          (compiler.assert compiler.scopes.macro
992                                           "must call from macro" ast)
993                          (. compiler.scopes.macro.manglings
994                             (tostring symbol)))
995             :macroexpand (fn [form]
996                            (compiler.assert compiler.scopes.macro
997                                             "must call from macro" ast)
998                            (compiler.macroexpand form
999                                                  compiler.scopes.macro))}]
1000    (set env._G env)
1001    (setmetatable env
1002                  {:__index provided
1003                   :__newindex provided
1004                   :__pairs combined-mt-pairs})))
1005
1006;; search-module uses package.config to process package.path (windows compat)
1007(local [dirsep pathsep pathmark]
1008       (icollect [c (string.gmatch (or package.config "") "([^\n]+)")] c))
1009(local pkg-config {:dirsep (or dirsep "/")
1010                   :pathmark (or pathmark ";")
1011                   :pathsep (or pathsep "?")})
1012
1013(fn escapepat [str]
1014  "Escape a string for safe use in a Lua pattern."
1015  (string.gsub str "[^%w]" "%%%1"))
1016
1017(fn search-module [modulename ?pathstring]
1018  (let [pathsepesc (escapepat pkg-config.pathsep)
1019        pattern (: "([^%s]*)%s" :format pathsepesc pathsepesc)
1020        no-dot-module (modulename:gsub "%." pkg-config.dirsep)
1021        fullpath (.. (or ?pathstring utils.fennel-module.path)
1022                     pkg-config.pathsep)]
1023    (fn try-path [path]
1024      (let [filename (: path :gsub (escapepat pkg-config.pathmark)
1025                        no-dot-module)
1026            filename2 (: path :gsub (escapepat pkg-config.pathmark) modulename)]
1027        (match (or (io.open filename) (io.open filename2))
1028          file (do
1029                 (file:close)
1030                 filename))))
1031
1032    (fn find-in-path [start]
1033      (match (fullpath:match pattern start)
1034        path (or (try-path path) (find-in-path (+ start (length path) 1)))))
1035
1036    (find-in-path 1)))
1037
1038(fn make-searcher [?options]
1039  "This will allow regular `require` to work with Fennel:
1040table.insert(package.loaders, fennel.searcher)"
1041  (fn [module-name]
1042    (let [opts (utils.copy utils.root.options)]
1043      (each [k v (pairs (or ?options {}))]
1044        (tset opts k v))
1045      (set opts.module-name module-name)
1046      (match (search-module module-name)
1047        filename (values (partial utils.fennel-module.dofile filename opts)
1048                         filename)))))
1049
1050(fn fennel-macro-searcher [module-name]
1051  (let [opts (doto (utils.copy utils.root.options)
1052               (tset :env :_COMPILER)
1053               (tset :requireAsInclude false)
1054               (tset :allowedGlobals nil))]
1055    (match (search-module module-name utils.fennel-module.macro-path)
1056      filename (values (partial utils.fennel-module.dofile filename opts)
1057                       filename))))
1058
1059(fn lua-macro-searcher [module-name]
1060  (match (search-module module-name package.path)
1061    filename (let [code (with-open [f (io.open filename)] (assert (f:read :*a)))
1062                   chunk (load-code code (make-compiler-env) filename)]
1063               (values chunk filename))))
1064
1065(local macro-searchers [fennel-macro-searcher lua-macro-searcher])
1066
1067(fn search-macro-module [modname n]
1068  (match (. macro-searchers n)
1069    f (match (f modname)
1070        (loader ?filename) (values loader ?filename)
1071        _ (search-macro-module modname (+ n 1)))))
1072
1073(fn metadata-only-fennel [modname]
1074  "Let limited Fennel module thru just for purposes of compiling docstrings."
1075  (if (or (= modname :fennel.macros)
1076          (and package package.loaded
1077               (= :table (type (. package.loaded modname)))
1078               (= (. package.loaded modname :metadata) compiler.metadata)))
1079      {:metadata compiler.metadata}))
1080
1081(set safe-require (fn [modname]
1082                    "This is a replacement for require for use in macro contexts.
1083It ensures that compile-scoped modules are loaded differently from regular
1084modules in the compiler environment."
1085                    (or (. macro-loaded modname) (metadata-only-fennel modname)
1086                        (let [(loader filename) (search-macro-module modname 1)]
1087                          (compiler.assert loader (.. modname " module not found."))
1088                          (tset macro-loaded modname (loader modname filename))
1089                          (. macro-loaded modname)))))
1090
1091(fn add-macros [macros* ast scope]
1092  (compiler.assert (utils.table? macros*) "expected macros to be table" ast)
1093  (each [k v (pairs macros*)]
1094    (compiler.assert (= (type v) :function)
1095                     "expected each macro to be function" ast)
1096    (tset scope.macros k v)))
1097
1098(fn resolve-module-name [{: filename 2 second} _scope _parent opts]
1099  ;; Compile module path to resolve real module name.  Allows using
1100  ;; (.. ... :.foo.bar) expressions and self-contained
1101  ;; statement-expressions in `require`, `include`, `require-macros`,
1102  ;; and `import-macros`.
1103  (let [filename (or filename (and (utils.table? second) second.filename))
1104        module-name utils.root.options.module-name
1105        modexpr (compiler.compile second opts)
1106        modname-chunk (load-code modexpr)]
1107    (modname-chunk module-name filename)))
1108
1109(fn SPECIALS.require-macros [ast scope parent ?real-ast]
1110  (compiler.assert (= (length ast) 2) "Expected one module name argument"
1111                   (or ?real-ast ast)) ; real-ast comes from import-macros
1112  (let [modname (resolve-module-name ast scope parent {})]
1113    (compiler.assert (= :string (type modname))
1114                     "module name must compile to string" (or ?real-ast ast))
1115    (when (not (. macro-loaded modname))
1116      (let [(loader filename) (search-macro-module modname 1)]
1117        (compiler.assert loader (.. modname " module not found.") ast)
1118        (tset macro-loaded modname (loader modname filename))))
1119    ;; if we're called from import-macros, return the modname, else add them
1120    ;; to scope directly
1121    (if (= :import-macros (tostring (. ast 1)))
1122        (. macro-loaded modname)
1123        (add-macros (. macro-loaded modname) ast scope parent))))
1124
1125(doc-special :require-macros [:macro-module-name]
1126             "Load given module and use its contents as macro definitions in current scope.
1127Macro module should return a table of macro functions with string keys.
1128Consider using import-macros instead as it is more flexible.")
1129
1130(fn emit-included-fennel [src path opts sub-chunk]
1131  "Emit Fennel code in src into sub-chunk."
1132  (let [subscope (compiler.make-scope utils.root.scope.parent)
1133        forms []]
1134    (when utils.root.options.requireAsInclude
1135      (set subscope.specials.require compiler.require-include))
1136    ;; parse Fennel src into table of exprs to know which expr is the tail
1137    (each [_ val (parser.parser (parser.string-stream src) path)]
1138      (table.insert forms val))
1139    ;; Compile the forms into sub-chunk; compiler.compile1 is necessary
1140    ;; for all nested includes to be emitted in the same root chunk
1141    ;; in the top-level module.
1142    (for [i 1 (length forms)]
1143      (let [subopts (if (= i (length forms)) {:tail true} {:nval 0})]
1144        (utils.propagate-options opts subopts)
1145        (compiler.compile1 (. forms i) subscope sub-chunk subopts)))))
1146
1147(fn include-path [ast opts path mod fennel?]
1148  "Helper function for include once we have determined the path to use."
1149  (tset utils.root.scope.includes mod :fnl/loading)
1150  (let [src (with-open [f (assert (io.open path))]
1151              (: (f:read :*all) :gsub "[\r\n]*$" ""))
1152        ;; splice in source and memoize it in compiler AND package.preload
1153        ;; so we can include it again without duplication, even in runtime
1154        ret (utils.expr (.. "require(\"" mod "\")") :statement)
1155        target (: "package.preload[%q]" :format mod)
1156        preload-str (.. target " = " target " or function(...)")
1157        (temp-chunk sub-chunk) (values [] [])]
1158    (compiler.emit temp-chunk preload-str ast)
1159    (compiler.emit temp-chunk sub-chunk)
1160    (compiler.emit temp-chunk :end ast)
1161    ;; Splice temp-chunk to begining of root chunk
1162    (each [i v (ipairs temp-chunk)]
1163      (table.insert utils.root.chunk i v))
1164    ;; For fennel source, compile sub-chunk AFTER splicing into start of
1165    ;; root chunk.
1166    (if fennel?
1167        (emit-included-fennel src path opts sub-chunk)
1168        ;; For Lua source, simply emit src into the loaders's body
1169        (compiler.emit sub-chunk src ast))
1170    ;; Put in cache and return
1171    (tset utils.root.scope.includes mod ret)
1172    ret))
1173
1174(fn include-circular-fallback [mod modexpr fallback ast]
1175  "If a circular include is detected, fall back to require if possible."
1176  (when (= (. utils.root.scope.includes mod) :fnl/loading) ; circular include
1177    (compiler.assert fallback "circular include detected" ast)
1178    (fallback modexpr)))
1179
1180(fn SPECIALS.include [ast scope parent opts]
1181  (compiler.assert (= (length ast) 2) "expected one argument" ast)
1182  (let [modexpr (match (pcall resolve-module-name ast scope parent opts)
1183                  ;; if we're in a dofile and not a require, then module-name
1184                  ;; will be nil and we will not be able to successfully
1185                  ;; compile relative requires into includes, but we can still
1186                  ;; emit a runtime relative require.
1187                  (true modname) (utils.expr (string.format "%q" modname) :literal)
1188                  _ (. (compiler.compile1 (. ast 2) scope parent {:nval 1}) 1))]
1189    (if (or (not= modexpr.type :literal) (not= (: (. modexpr 1) :byte) 34))
1190        (if opts.fallback
1191            (opts.fallback modexpr)
1192            (compiler.assert false "module name must be string literal" ast))
1193        (let [mod ((load-code (.. "return " (. modexpr 1))))
1194              oldmod utils.root.options.module-name
1195              _ (set utils.root.options.module-name mod)
1196              res (or (and (utils.member? mod (or utils.root.options.skipInclude []))
1197                           (utils.expr "nil --[[SKIPPED INCLUDE]]--" :literal))
1198                      (include-circular-fallback mod modexpr opts.fallback ast)
1199                      (. utils.root.scope.includes mod) ; check cache
1200                      ;; Find path to Fennel or Lua source; prefering Fennel
1201                      (match (search-module mod)
1202                        fennel-path (include-path ast opts fennel-path mod true)
1203                        _ (let [lua-path (search-module mod package.path)]
1204                            (if lua-path (include-path ast opts lua-path mod false)
1205                                opts.fallback (opts.fallback modexpr)
1206                                (compiler.assert false (.. "module not found " mod) ast)))))]
1207          (set utils.root.options.module-name oldmod)
1208          res))))
1209
1210(doc-special :include [:module-name-literal]
1211             "Like require but load the target module during compilation and embed it in the
1212Lua output. The module must be a string literal and resolvable at compile time.")
1213
1214(fn eval-compiler* [ast scope parent]
1215  (let [env (make-compiler-env ast scope parent)
1216        opts (utils.copy utils.root.options)]
1217    (set opts.scope (compiler.make-scope compiler.scopes.compiler))
1218    (set opts.allowedGlobals (current-global-names env))
1219    ((load-code (compiler.compile ast opts) (wrap-env env)) opts.module-name
1220                                                            ast.filename)))
1221
1222(fn SPECIALS.macros [ast scope parent]
1223  (compiler.assert (= (length ast) 2) "Expected one table argument" ast)
1224  (add-macros (eval-compiler* (. ast 2) scope parent) ast scope parent))
1225
1226(doc-special :macros
1227             ["{:macro-name-1 (fn [...] ...) ... :macro-name-N macro-body-N}"]
1228             "Define all functions in the given table as macros local to the current scope.")
1229
1230(fn SPECIALS.eval-compiler [ast scope parent]
1231  (let [old-first (. ast 1)]
1232    (tset ast 1 (utils.sym :do))
1233    (let [val (eval-compiler* ast scope parent)]
1234      (tset ast 1 old-first)
1235      val)))
1236
1237(doc-special :eval-compiler ["..."]
1238             "Evaluate the body at compile-time. Use the macro system instead if possible."
1239             true)
1240
1241{:doc doc*
1242 : current-global-names
1243 : load-code
1244 : macro-loaded
1245 : macro-searchers
1246 : make-compiler-env
1247 : search-module
1248 : make-searcher
1249 : wrap-env}
1250