1;; This is the core compiler module responsible for taking a parsed AST
2;; and turning it into Lua code. Main entry points are `compile` (which
3;; takes an AST), `compile-stream` and `compile-string`.
4
5(local utils (require :fennel.utils))
6(local parser (require :fennel.parser))
7(local friend (require :fennel.friend))
8
9(local unpack (or table.unpack _G.unpack))
10
11(local scopes [])
12
13(fn make-scope [?parent]
14  "Create a new Scope, optionally under a parent scope.
15Scopes are compile time constructs that are responsible for keeping track of
16local variables, name mangling, and macros.  They are accessible to user code
17via the 'eval-compiler' special form (may change). They use metatables to
18implement nesting. "
19  (let [parent (or ?parent scopes.global)]
20    {:includes (setmetatable [] {:__index (and parent parent.includes)})
21     :macros (setmetatable [] {:__index (and parent parent.macros)})
22     :manglings (setmetatable [] {:__index (and parent parent.manglings)})
23     :specials (setmetatable [] {:__index (and parent parent.specials)})
24     :symmeta (setmetatable [] {:__index (and parent parent.symmeta)})
25     :unmanglings (setmetatable [] {:__index (and parent parent.unmanglings)})
26     :gensyms (setmetatable [] {:__index (and parent parent.gensyms)})
27     :autogensyms (setmetatable [] {:__index (and parent parent.autogensyms)})
28     :vararg (and parent parent.vararg)
29     :depth (if parent (+ (or parent.depth 0) 1) 0)
30     :hashfn (and parent parent.hashfn)
31     :refedglobals {}
32     : parent}))
33
34(fn assert-msg [ast msg]
35  (let [ast-tbl (if (= :table (type ast)) ast {})
36        m (getmetatable ast)
37        filename (or (and m m.filename) ast-tbl.filename :unknown)
38        line (or (and m m.line) ast-tbl.line "?")
39        target (tostring (or (utils.sym? (. ast-tbl 1)) (. ast-tbl 1) "()"))]
40    (string.format "%s:%s: Compile error in '%s': %s" filename line target msg)))
41
42;; If you add new calls to this function, please update fennel.friend
43;; as well to add suggestions for how to fix the new error!
44(fn assert-compile [condition msg ast]
45  "Assert a condition and raise a compile error with line numbers.
46The ast arg should be unmodified so that its first element is the form called."
47  (when (not condition)
48    (let [{: source : unfriendly} (or utils.root.options {})]
49      ;; allow plugins to override assert-compile
50      (when (= nil (utils.hook :assert-compile condition msg ast
51                               utils.root.reset))
52        (utils.root.reset)
53        (if (or unfriendly (not friend) (not _G.io) (not _G.io.read))
54            ;; if we use regular `assert' we can't set level to 0
55            (error (assert-msg ast msg) 0)
56            (friend.assert-compile condition msg ast source)))))
57  condition)
58
59(set scopes.global (make-scope))
60(set scopes.global.vararg true)
61(set scopes.compiler (make-scope scopes.global))
62(set scopes.macro scopes.global)
63
64;; Allow printing a string to Lua, also keep as 1 line.
65(local serialize-subst {"\a" "\\a"
66                        "\b" "\\b"
67                        "\t" "\\t"
68                        "\n" :n
69                        "\v" "\\v"
70                        "\f" "\\f"})
71
72(fn serialize-string [str]
73  (-> (string.format "%q" str)
74      (string.gsub "." serialize-subst)
75      (string.gsub "[\x80-\xff]" #(.. "\\" ($:byte)))))
76
77(fn global-mangling [str]
78  "Mangler for global symbols. Does not protect against collisions,
79but makes them unlikely. This is the mangling that is exposed to the world."
80  (if (utils.valid-lua-identifier? str)
81      str
82      (.. :__fnl_global__ (str:gsub "[^%w]" #(string.format "_%02x" ($:byte))))))
83
84(fn global-unmangling [identifier]
85  "Reverse a global mangling.
86Takes a Lua identifier and returns the Fennel symbol string that created it."
87  (match (string.match identifier "^__fnl_global__(.*)$")
88    rest (pick-values 1
89                      (string.gsub rest "_[%da-f][%da-f]"
90                                   #(string.char (tonumber ($:sub 2) 16))))
91    _ identifier))
92
93(var allowed-globals nil)
94
95(fn global-allowed? [name]
96  "If there's a provided list of allowed globals, don't let references thru that
97aren't on the list. This list is set at the compiler entry points of compile
98and compile-stream."
99  (or (not allowed-globals) (utils.member? name allowed-globals)))
100
101(fn unique-mangling [original mangling scope append]
102  (if (and (. scope.unmanglings mangling) (not (. scope.gensyms mangling)))
103      (unique-mangling original (.. original append) scope (+ append 1))
104      mangling))
105
106(fn local-mangling [str scope ast ?temp-manglings]
107  "Creates a symbol from a string by mangling it. ensures that the generated
108symbol is unique if the input string is unique in the scope."
109  (assert-compile (not (utils.multi-sym? str))
110                  (.. "unexpected multi symbol " str) ast)
111  (let [;; Mapping mangling to a valid Lua identifier
112        raw (if (or (. utils.lua-keywords str) (str:match "^%d"))
113                (.. "_" str)
114                str)
115        mangling (-> raw
116                     (string.gsub "-" "_")
117                     (string.gsub "[^%w_]" #(string.format "_%02x" ($:byte))))
118        unique (unique-mangling mangling mangling scope 0)]
119    (tset scope.unmanglings unique str)
120    (let [manglings (or ?temp-manglings scope.manglings)]
121      (tset manglings str unique))
122    unique))
123
124(fn apply-manglings [scope new-manglings ast]
125  "Calling this function will mean that further compilation in scope will use
126these new manglings instead of the current manglings."
127  (each [raw mangled (pairs new-manglings)]
128    (assert-compile (not (. scope.refedglobals mangled))
129                    (.. "use of global " raw " is aliased by a local") ast)
130    (tset scope.manglings raw mangled)))
131
132(fn combine-parts [parts scope]
133  "Combine parts of a symbol."
134  (var ret (or (. scope.manglings (. parts 1)) (global-mangling (. parts 1))))
135  (for [i 2 (length parts)]
136    (if (utils.valid-lua-identifier? (. parts i))
137        (if (and parts.multi-sym-method-call (= i (length parts)))
138            (set ret (.. ret ":" (. parts i)))
139            (set ret (.. ret "." (. parts i))))
140        (set ret (.. ret "[" (serialize-string (. parts i)) "]"))))
141  ret)
142
143(fn next-append []
144  (set utils.root.scope.gensym-append (+ (or utils.root.scope.gensym-append 0) 1))
145  (.. "_" utils.root.scope.gensym-append "_"))
146
147(fn gensym [scope ?base ?suffix]
148  "Generates a unique symbol in the scope."
149  (var mangling (.. (or ?base "") (next-append) (or ?suffix "")))
150  (while (. scope.unmanglings mangling)
151    (set mangling (.. (or ?base "") (next-append) (or ?suffix ""))))
152  (tset scope.unmanglings mangling (or ?base true))
153  (tset scope.gensyms mangling true)
154  mangling)
155
156(fn autogensym [base scope]
157  "Generates a unique symbol in the scope based on the base name. Calling
158repeatedly with the same base and same scope will return existing symbol
159rather than generating new one."
160  (match (utils.multi-sym? base)
161    parts (do
162            (tset parts 1 (autogensym (. parts 1) scope))
163            (table.concat parts (or (and parts.multi-sym-method-call ":") ".")))
164    _ (or (. scope.autogensyms base)
165          (let [mangling (gensym scope (base:sub 1 (- 2)) :auto)]
166            (tset scope.autogensyms base mangling)
167            mangling))))
168
169(fn check-binding-valid [symbol scope ast]
170  "Check to see if a symbol will be overshadowed by a special."
171  (let [name (tostring symbol)]
172    ;; we can't block in the parser because & is still ok in symbols like &as
173    (assert-compile (not (name:find "&")) "illegal character &")
174    (assert-compile (not (or (. scope.specials name) (. scope.macros name)))
175                    (: "local %s was overshadowed by a special form or macro"
176                       :format name) ast)
177    (assert-compile (not (utils.quoted? symbol))
178                    (string.format "macro tried to bind %s without gensym" name)
179                    symbol)))
180
181(fn declare-local [symbol meta scope ast ?temp-manglings]
182  "Declare a local symbol"
183  (check-binding-valid symbol scope ast)
184  (let [name (tostring symbol)]
185    (assert-compile (not (utils.multi-sym? name))
186                    (.. "unexpected multi symbol " name) ast)
187    (tset scope.symmeta name meta)
188    (local-mangling name scope ast ?temp-manglings)))
189
190(fn hashfn-arg-name [name multi-sym-parts scope]
191  (if (not scope.hashfn) nil
192      (= name "$") :$1
193      multi-sym-parts (do
194                        (when (and multi-sym-parts
195                                   (= (. multi-sym-parts 1) "$"))
196                          (tset multi-sym-parts 1 :$1))
197                        (table.concat multi-sym-parts "."))))
198
199(fn symbol-to-expression [symbol scope ?reference?]
200  "Convert symbol to Lua code. Will only work for local symbols
201if they have already been declared via declare-local"
202  (utils.hook :symbol-to-expression symbol scope ?reference?)
203  (let [name (. symbol 1)
204        multi-sym-parts (utils.multi-sym? name)
205        name (or (hashfn-arg-name name multi-sym-parts scope) name)]
206    (let [parts (or multi-sym-parts [name])
207          etype (or (and (> (length parts) 1) :expression) :sym)
208          local? (. scope.manglings (. parts 1))]
209      (when (and local? (. scope.symmeta (. parts 1)))
210        (tset (. scope.symmeta (. parts 1)) :used true))
211      (assert-compile (not (. scope.macros (. parts 1)))
212                      (.. "tried to reference a macro at runtime") symbol)
213      ;; if it's a reference and not a symbol which introduces a new binding
214      ;; then we need to check for allowed globals
215      (assert-compile (or (not ?reference?) local? (= :_ENV (. parts 1))
216                          (global-allowed? (. parts 1)))
217                      (.. "unknown identifier in strict mode: "
218                          (tostring (. parts 1))) symbol)
219      (when (and allowed-globals (not local?) scope.parent)
220        (tset scope.parent.refedglobals (. parts 1) true))
221      (utils.expr (combine-parts parts scope) etype))))
222
223(fn emit [chunk out ?ast]
224  "Emit Lua code."
225  (if (= (type out) :table)
226      (table.insert chunk out)
227      (table.insert chunk {:ast ?ast :leaf out})))
228
229(fn peephole [chunk]
230  "Do some peephole optimization."
231  (if chunk.leaf chunk
232      (and (>= (length chunk) 3) (= (. (. chunk (- (length chunk) 2)) :leaf)
233                                    :do)
234           (not (. (. chunk (- (length chunk) 1)) :leaf))
235           (= (. (. chunk (length chunk)) :leaf) :end))
236      (let [kid (peephole (. chunk (- (length chunk) 1)))
237            new-chunk {:ast chunk.ast}]
238        (for [i 1 (- (length chunk) 3)]
239          (table.insert new-chunk (peephole (. chunk i))))
240        (for [i 1 (length kid)]
241          (table.insert new-chunk (. kid i)))
242        new-chunk) (utils.map chunk peephole)))
243
244(fn flatten-chunk-correlated [main-chunk options]
245  "Correlate line numbers in input with line numbers in output."
246  (fn flatten [chunk out last-line file]
247    (var last-line last-line)
248    (if chunk.leaf
249        (tset out last-line (.. (or (. out last-line) "") " " chunk.leaf))
250        (each [_ subchunk (ipairs chunk)]
251          (when (or subchunk.leaf (> (length subchunk) 0)) ; ignore empty chunks
252            ;; don't increase line unless it's from the same file
253            (let [source (utils.ast-source subchunk.ast)]
254              (when (= file source.filename)
255                (set last-line (math.max last-line (or source.line 0))))
256              (set last-line (flatten subchunk out last-line file))))))
257    last-line)
258
259  (let [out []
260        last (flatten main-chunk out 1 options.filename)]
261    (for [i 1 last]
262      (when (= (. out i) nil)
263        (tset out i "")))
264    (table.concat out "\n")))
265
266(fn flatten-chunk [sm chunk tab depth]
267  "Flatten a tree of indented Lua source code lines.
268Tab is what is used to indent a block."
269  (if chunk.leaf
270      (let [code chunk.leaf
271            info chunk.ast]
272        (when sm
273          (table.insert sm [(and info info.filename) (and info info.line)]))
274        code)
275      (let [tab (match tab
276                  true "  "
277                  false ""
278                  tab tab
279                  nil "")]
280        (fn parter [c]
281          (when (or c.leaf (> (length c) 0))
282            (let [sub (flatten-chunk sm c tab (+ depth 1))]
283              (if (> depth 0)
284                  (.. tab (sub:gsub "\n" (.. "\n" tab)))
285                  sub))))
286
287        (table.concat (utils.map chunk parter) "\n"))))
288
289;; Some global state for all fennel sourcemaps. For the time being, this seems
290;; the easiest way to store the source maps.  Sourcemaps are stored with source
291;; being mapped as the key, prepended with '@' if it is a filename (like
292;; debug.getinfo returns for source).  The value is an array of mappings for
293;; each line.
294(local sourcemap [])
295
296(fn make-short-src [source]
297  (let [source (source:gsub "\n" " ")]
298    (if (<= (length source) 49)
299        (.. "[fennel \"" source "\"]")
300        (.. "[fennel \"" (source:sub 1 46) "...\"]"))))
301
302(fn flatten [chunk options]
303  "Return Lua source and source map table."
304  (let [chunk (peephole chunk)]
305    (if options.correlate
306        (values (flatten-chunk-correlated chunk options) [])
307        (let [sm []
308              ret (flatten-chunk sm chunk options.indent 0)]
309          (when sm
310            (set sm.short_src
311                 (or options.filename (make-short-src (or options.source ret))))
312            (set sm.key (if options.filename (.. "@" options.filename) ret))
313            (tset sourcemap sm.key sm))
314          (values ret sm)))))
315
316(fn make-metadata []
317  "Make module-wide state table for metadata."
318  (setmetatable [] {:__index {:get (fn [self tgt key]
319                                     (when (. self tgt)
320                                       (. (. self tgt) key)))
321                              :set (fn [self tgt key value]
322                                     (tset self tgt (or (. self tgt) []))
323                                     (tset (. self tgt) key value)
324                                     tgt)
325                              :setall (fn [self tgt ...]
326                                        (let [kv-len (select "#" ...)
327                                              kvs [...]]
328                                          (when (not= (% kv-len 2) 0)
329                                            (error "metadata:setall() expected even number of k/v pairs"))
330                                          (tset self tgt (or (. self tgt) []))
331                                          (for [i 1 kv-len 2]
332                                            (tset (. self tgt) (. kvs i)
333                                                  (. kvs (+ i 1))))
334                                          tgt))}
335                    :__mode :k}))
336
337(fn exprs1 [exprs]
338  "Convert expressions to Lua string."
339  (table.concat (utils.map exprs tostring) ", "))
340
341(fn keep-side-effects [exprs chunk start ast]
342  "Compile side effects for a chunk."
343  (let [start (or start 1)]
344    (for [j start (length exprs)]
345      (let [se (. exprs j)]
346        ;; Avoid the rogue 'nil' expression (nil is usually a literal,
347        ;; but becomes an expression if a special form returns 'nil')
348        (if (and (= se.type :expression) (not= (. se 1) :nil))
349            (emit chunk (string.format "do local _ = %s end" (tostring se)) ast)
350            (= se.type :statement)
351            (let [code (tostring se)
352                  disambiguated (if (= (code:byte) 40) (.. "do end " code) code)]
353              (emit chunk disambiguated ast)))))))
354
355(fn handle-compile-opts [exprs parent opts ast]
356  "Does some common handling of returns and register targets for special
357forms. Also ensures a list expression has an acceptable number of expressions
358if opts contains the nval option."
359  (when opts.nval
360    (let [n opts.nval
361          len (length exprs)]
362      (when (not= n len)
363        (if (> len n)
364            (do
365              ; drop extra
366              (keep-side-effects exprs parent (+ n 1) ast)
367              (for [i (+ n 1) len]
368                (tset exprs i nil)))
369            (for [i (+ (length exprs) 1) n] ; pad with nils
370              (tset exprs i (utils.expr :nil :literal)))))))
371  (when opts.tail
372    (emit parent (string.format "return %s" (exprs1 exprs)) ast))
373  (when opts.target
374    (let [result (exprs1 exprs)]
375      (emit parent
376            (string.format "%s = %s" opts.target (if (= result "") :nil result))
377            ast)))
378  (if (or opts.tail opts.target)
379      ;; Prevent statements and expression from being used twice if they
380      ;; have side-effects. Since if the target or tail options are set,
381      ;; the expressions are already emitted, we should not return them. This
382      ;; is fine, as when these options are set, the caller doesn't need the
383      ;; result anyways.
384      {:returned true}
385      (doto exprs (tset :returned true))))
386
387(fn find-macro [ast scope multi-sym-parts]
388  (fn find-in-table [t i]
389    (if (<= i (length multi-sym-parts))
390        (find-in-table (and (utils.table? t) (. t (. multi-sym-parts i)))
391                       (+ i 1))
392        t))
393
394  (let [macro* (and (utils.sym? (. ast 1))
395                    (. scope.macros (tostring (. ast 1))))]
396    (if (and (not macro*) multi-sym-parts)
397        (let [nested-macro (find-in-table scope.macros 1)]
398          (assert-compile (or (not (. scope.macros (. multi-sym-parts 1)))
399                              (= (type nested-macro) :function))
400                          "macro not found in imported macro module" ast)
401          nested-macro)
402        macro*)))
403
404(fn propagate-trace-info [{: filename : line : bytestart : byteend} _index node]
405  "The stack trace info should be based on the macro caller, not the macro AST."
406  (when (and (= :table (type node)) (not= filename node.filename))
407    (let [src (utils.ast-source node)]
408      (set (src.filename src.line) (values filename line))
409      (set (src.bytestart src.byteend) (values bytestart byteend))))
410  (= :table (type node)))
411
412(fn macroexpand* [ast scope ?once]
413  "Expand macros in the ast. Only do one level if once is true."
414  (match (if (utils.list? ast)
415             (find-macro ast scope (utils.multi-sym? (. ast 1))))
416    false ast
417    macro* (let [old-scope scopes.macro
418                 _ (set scopes.macro scope)
419                 ;; TODO: we want to pass a traceback level, but it only
420                 ;; supports trimming the trace from the wrong direction.
421                 (ok transformed) (xpcall #(macro* (unpack ast 2))
422                                          debug.traceback)]
423             (utils.walk-tree transformed (partial propagate-trace-info ast))
424             (set scopes.macro old-scope)
425             (assert-compile ok transformed ast)
426             (if (or ?once (not transformed))
427                 transformed
428                 (macroexpand* transformed scope)))
429    _ ast))
430
431(fn compile-special [ast scope parent opts special]
432  (let [exprs (or (special ast scope parent opts) (utils.expr :nil :literal))
433        ;; Be very accepting of strings or expressions as well as lists
434        ;; or expressions
435        exprs (if (not= :table (type exprs))
436                  (utils.expr exprs :expression)
437                  exprs)
438        exprs (if (utils.expr? exprs) [exprs] exprs)]
439    ;; Unless the special form explicitly handles the target, tail,
440    ;; and nval properties, (indicated via the 'returned' flag),
441    ;; handle these options.
442    (if (not exprs.returned) (handle-compile-opts exprs parent opts ast)
443        (or opts.tail opts.target) {:returned true}
444        exprs)))
445
446(fn compile-function-call [ast scope parent opts compile1 len]
447  (let [fargs [] ; regular function call
448        fcallee (. (compile1 (. ast 1) scope parent {:nval 1}) 1)]
449    (assert-compile (or (= :string (type (. ast 1))) ; strings can have __call metamethod
450                        (not= fcallee.type :literal))
451                    (.. "cannot call literal value " (tostring (. ast 1))) ast)
452    (for [i 2 len]
453      (let [subexprs (compile1 (. ast i) scope parent
454                               {:nval (if (not= i len) 1)})]
455        (table.insert fargs (or (. subexprs 1) (utils.expr :nil :literal)))
456        (if (= i len)
457            ;; Add multivalues to function args
458            (for [j 2 (length subexprs)]
459              (table.insert fargs (. subexprs j)))
460            ;; Emit sub expression only for side effects
461            (keep-side-effects subexprs parent 2 (. ast i)))))
462    (let [pat (if (= :string (type (. ast 1)))
463                  "(%s)(%s)"  ; ("a")() is valid Lua call, "a"() isn't
464                  "%s(%s)")   ; regular literal call
465          call (string.format pat (tostring fcallee) (exprs1 fargs))]
466      (handle-compile-opts [(utils.expr call :statement)] parent opts ast))))
467
468(fn compile-call [ast scope parent opts compile1]
469  (utils.hook :call ast scope)
470  (let [len (length ast)
471        first (. ast 1)
472        multi-sym-parts (utils.multi-sym? first)
473        special (and (utils.sym? first) (. scope.specials (tostring first)))]
474    (assert-compile (> len 0) "expected a function, macro, or special to call"
475                    ast)
476    (if special
477        (compile-special ast scope parent opts special)
478        (and multi-sym-parts multi-sym-parts.multi-sym-method-call)
479        (let [table-with-method (table.concat [(unpack multi-sym-parts 1
480                                                       (- (length multi-sym-parts)
481                                                          1))]
482                                              ".")
483              method-to-call (. multi-sym-parts (length multi-sym-parts))
484              new-ast (utils.list (utils.sym ":" nil scope)
485                                  (utils.sym table-with-method nil scope)
486                                  method-to-call (select 2 (unpack ast)))]
487          (compile1 new-ast scope parent opts))
488        (compile-function-call ast scope parent opts compile1 len))))
489
490(fn compile-varg [ast scope parent opts]
491  (assert-compile scope.vararg "unexpected vararg" ast)
492  (handle-compile-opts [(utils.expr "..." :varg)] parent opts ast))
493
494(fn compile-sym [ast scope parent opts]
495  (let [multi-sym-parts (utils.multi-sym? ast)]
496    (assert-compile (not (and multi-sym-parts
497                              multi-sym-parts.multi-sym-method-call))
498                    "multisym method calls may only be in call position" ast)
499    ;; Handle nil as special symbol - it resolves to the nil literal
500    ;; rather than being unmangled. Alternatively, we could remove it
501    ;; from the lua keywords table.
502    (let [e (if (= (. ast 1) :nil)
503                (utils.expr :nil :literal)
504                (symbol-to-expression ast scope true))]
505      (handle-compile-opts [e] parent opts ast))))
506
507;; We do gsub transformation because some locales use , for
508;; decimal separators, which will not be accepted by Lua.
509(fn serialize-number [n]
510  (pick-values 1 (-> (tostring n)
511                     (string.gsub "," "."))))
512
513(fn compile-scalar [ast _scope parent opts]
514  (let [serialize (match (type ast)
515                    :nil tostring
516                    :boolean tostring
517                    :string serialize-string
518                    :number serialize-number)]
519    (handle-compile-opts [(utils.expr (serialize ast) :literal)] parent opts)))
520
521(fn compile-table [ast scope parent opts compile1]
522  (let [buffer []]
523    (fn write-other-values [k]
524      (when (or (not= (type k) :number) (not= (math.floor k) k) (< k 1)
525                (> k (length ast)))
526        (if (and (= (type k) :string) (utils.valid-lua-identifier? k)) [k k]
527            (let [[compiled] (compile1 k scope parent {:nval 1})
528                  kstr (.. "[" (tostring compiled) "]")]
529              [kstr k]))))
530
531    (let [keys (icollect [k v (utils.stablepairs ast)]
532                 (write-other-values k v))]
533      (utils.map keys
534                 (fn [[k1 k2]]
535                   (let [[v] (compile1 (. ast k2) scope parent {:nval 1})]
536                     (string.format "%s = %s" k1 (tostring v))))
537                 buffer))
538
539    (for [i 1 (length ast)] ; write numeric keyed values
540      (let [nval (and (not= i (length ast)) 1)]
541        (table.insert buffer
542                      (exprs1 (compile1 (. ast i) scope parent {: nval})))))
543
544    (handle-compile-opts [(utils.expr (.. "{" (table.concat buffer ", ") "}")
545                                      :expression)]
546                         parent opts ast)))
547
548(fn compile1 [ast scope parent ?opts]
549  "Compile an AST expression in the scope into parent, a tree of lines that is
550eventually compiled into Lua code. Also returns some information about the
551evaluation of the compiled expression, which can be used by the calling
552function. Macros are resolved here, as well as special forms in that order.
553
554* the `ast` param is the root AST to compile
555* the `scope` param is the scope in which we are compiling
556* the `parent` param is the table of lines that we are compiling into.
557add lines to parent by appending strings. Add indented blocks by appending
558tables of more lines.
559* the `opts` param contains info about where the form is being compiled
560
561Fields of `opts` include:
562  target: mangled name of symbol(s) being compiled to.
563     Could be one variable, 'a', or a list, like 'a, b, _0_'.
564  tail: boolean indicating tail position if set. If set, form will generate
565     a return instruction.
566  nval: The number of values to compile to if it is known to be a fixed value.
567
568In Lua, an expression can evaluate to 0 or more values via multiple returns. In
569many cases, Lua will drop extra values and convert a 0 value expression to
570nil. In other cases, Lua will use all of the values in an expression, such as
571in the last argument of a function call. Nval is an option passed to compile1
572to say that the resulting expression should have at least n values. It lets us
573generate better code, because if we know we are only going to use 1 or 2 values
574from an expression, we can create 1 or 2 locals to store intermediate results
575rather than turn the expression into a closure that is called immediately,
576which we have to do if we don't know."
577  (let [opts (or ?opts [])
578        ast (macroexpand* ast scope)]
579    (if (utils.list? ast)
580        (compile-call ast scope parent opts compile1)
581        (utils.varg? ast)
582        (compile-varg ast scope parent opts)
583        (utils.sym? ast)
584        (compile-sym ast scope parent opts)
585        (= (type ast) :table)
586        (compile-table ast scope parent opts compile1)
587        (or (= (type ast) :nil) (= (type ast) :boolean) (= (type ast) :number)
588            (= (type ast) :string))
589        (compile-scalar ast scope parent opts)
590        (assert-compile false
591                        (.. "could not compile value of type " (type ast)) ast))))
592
593;; You may be tempted to clean up and refactor this function because it's so
594;; huge and stateful but it really needs to get replaced; it is too tightly
595;; coupled to the way the compiler outputs Lua; it should be split into general
596;; data-driven parts vs Lua-emitting parts.
597(fn destructure [to from ast scope parent opts]
598  "Implements destructuring for forms like let, bindings, etc.
599  Takes a number of opts to control behavior.
600  * var: Whether or not to mark symbols as mutable
601  * declaration: begin each assignment with 'local' in output
602  * nomulti: disallow multisyms in the destructuring. for (local) and (global)
603  * noundef: Don't set undefined bindings. (set)
604  * forceglobal: Don't allow local bindings
605  * symtype: the type of syntax calling the destructuring, for lua output names"
606  (let [opts (or opts {})
607        {: isvar : declaration : forceglobal : forceset : symtype} opts
608        symtype (.. "_" (or symtype :dst))
609        setter (if declaration "local %s = %s" "%s = %s")
610        new-manglings []]
611    (fn getname [symbol up1]
612      "Get Lua source for symbol, and check for errors"
613      (let [raw (. symbol 1)]
614        (assert-compile (not (and opts.nomulti (utils.multi-sym? raw)))
615                        (.. "unexpected multi symbol " raw) up1)
616        (if declaration
617            ;; Technically this is too early to declare the local, but leaving
618            ;; out the meta table and setting it later works around the problem.
619            (declare-local symbol nil scope symbol new-manglings)
620            (let [parts (or (utils.multi-sym? raw) [raw])
621                  meta (. scope.symmeta (. parts 1))]
622              (assert-compile (not (raw:find ":")) "cannot set method sym" symbol)
623              (when (and (= (length parts) 1) (not forceset))
624                (assert-compile (not (and forceglobal meta))
625                                (string.format "global %s conflicts with local"
626                                               (tostring symbol))
627                                symbol)
628                (assert-compile (not (and meta (not meta.var)))
629                                (.. "expected var " raw) symbol)
630                (assert-compile (or meta (not opts.noundef))
631                                (.. "expected local " (. parts 1)) symbol))
632              (when forceglobal
633                (assert-compile (not (. scope.symmeta (. scope.unmanglings raw)))
634                                (.. "global " raw " conflicts with local")
635                                symbol)
636                (tset scope.manglings raw (global-mangling raw))
637                (tset scope.unmanglings (global-mangling raw) raw)
638                (when allowed-globals
639                  (table.insert allowed-globals raw)))
640              (. (symbol-to-expression symbol scope) 1)))))
641
642    (fn compile-top-target [lvalues]
643      "Compile the outer most form. We can generate better Lua in this case."
644      ;; Calculate initial rvalue
645      (let [inits (utils.map lvalues #(if (. scope.manglings $) $ :nil))
646            init (table.concat inits ", ")
647            lvalue (table.concat lvalues ", ")]
648        (var (plen plast) (values (length parent) (. parent (length parent))))
649        (local ret (compile1 from scope parent {:target lvalue}))
650        (when declaration
651          ;; A single leaf emitted at the end of the parent chunk means a
652          ;; simple assignment a = x was emitted, and we can just splice
653          ;; "local " onto the front of it. However, we can't just check
654          ;; based on plen, because some forms (such as include) insert new
655          ;; chunks at the top of the parent chunk rather than just at the
656          ;; end; this loop checks for this occurance and updates plen to be
657          ;; the index of the last thing in the parent before compiling the
658          ;; new value.
659          (for [pi plen (length parent)]
660            (when (= (. parent pi) plast)
661              (set plen pi)))
662          (if (and (= (length parent) (+ plen 1))
663                   (. (. parent (length parent)) :leaf))
664              (tset (. parent (length parent)) :leaf
665                    (.. "local " (. (. parent (length parent)) :leaf)))
666              (= init :nil)
667              (table.insert parent (+ plen 1) {: ast :leaf (.. "local " lvalue)})
668              (table.insert parent (+ plen 1)
669                            {: ast :leaf (.. "local " lvalue " = " init)})))
670        ret))
671
672    (fn destructure-sym [left rightexprs up1 top?]
673      (let [lname (getname left up1)]
674        (check-binding-valid left scope left)
675        (if top?
676            (compile-top-target [lname])
677            (emit parent (setter:format lname (exprs1 rightexprs)) left))
678        ;; We have to declare meta for the left *after* compiling the right
679        ;; see https://todo.sr.ht/~technomancy/fennel/12
680        (when declaration
681          (tset scope.symmeta (tostring left) {:var isvar}))))
682
683    (fn destructure-table [left rightexprs top? destructure1]
684      (let [s (gensym scope symtype)
685            right (match (if top?
686                             (exprs1 (compile1 from scope parent))
687                             (exprs1 rightexprs))
688                    "" :nil
689                    right right)]
690        (emit parent (string.format "local %s = %s" s right) left)
691        (each [k v (utils.stablepairs left)]
692          (when (not (and (= :number (type k))
693                          (: (tostring (. left (- k 1))) :find "^&")))
694            (if (and (utils.sym? v) (= (tostring v) "&"))
695                (let [unpack-str "(function (t, k)
696                                      local mt = getmetatable(t)
697                                      if \"table\" == type(mt) and mt.__fennelrest then
698                                         return mt.__fennelrest(t, k)
699                                      else
700                                         return {(table.unpack or unpack)(t, k)}
701                                      end
702                                   end)(%s, %s)"
703                      formatted (string.format (string.gsub unpack-str "\n%s*" " ") s k)
704                      subexpr (utils.expr formatted :expression)]
705                  (assert-compile (and (utils.sequence? left)
706                                       (= nil (. left (+ k 2))))
707                                  "expected rest argument before last parameter"
708                                  left)
709                  (destructure1 (. left (+ k 1)) [subexpr] left))
710                (and (utils.sym? k) (= (tostring k) :&as))
711                (destructure-sym v [(utils.expr (tostring s))] left)
712                (and (utils.sequence? left) (= (tostring v) :&as))
713                (let [(_ next-sym trailing) (select k (unpack left))]
714                  (assert-compile (= nil trailing)
715                                  "expected &as argument before last parameter"
716                                  left)
717                  (destructure-sym next-sym [(utils.expr (tostring s))] left))
718                (let [key (if (= (type k) :string) (serialize-string k) k)
719                      subexpr (utils.expr (string.format "%s[%s]" s key)
720                                          :expression)]
721                  (destructure1 v [subexpr] left)))))))
722
723    (fn destructure-values [left up1 top? destructure1]
724      (let [(left-names tables) (values [] [])]
725        (each [i name (ipairs left)]
726          (if (utils.sym? name) ; binding directly to a name
727              (table.insert left-names (getname name up1))
728              (let [symname (gensym scope symtype)]
729                ;; further destructuring of tables inside values
730                (table.insert left-names symname)
731                (tset tables i [name (utils.expr symname :sym)]))))
732        (assert-compile top? "can't nest multi-value destructuring" left)
733        (compile-top-target left-names)
734        (when declaration
735          (each [_ sym (ipairs left)]
736            (when (utils.sym? sym)
737              (tset scope.symmeta (tostring sym) {:var isvar}))))
738        ;; recurse if left-side tables found
739        (each [_ pair (utils.stablepairs tables)]
740          (destructure1 (. pair 1) [(. pair 2)] left))))
741
742    (fn destructure1 [left rightexprs up1 top?]
743      "Recursive auxiliary function"
744      (if (and (utils.sym? left) (not= (. left 1) :nil))
745          (destructure-sym left rightexprs up1 top?)
746          (utils.table? left)
747          (destructure-table left rightexprs top? destructure1)
748          (utils.list? left)
749          (destructure-values left up1 top? destructure1)
750          (assert-compile false
751                          (string.format "unable to bind %s %s" (type left)
752                                         (tostring left))
753                          (or (and (= (type (. up1 2)) :table) (. up1 2)) up1)))
754      (when top?
755        {:returned true}))
756
757    (let [ret (destructure1 to nil ast true)]
758      (utils.hook :destructure from to scope)
759      (apply-manglings scope new-manglings ast)
760      ret)))
761
762(fn require-include [ast scope parent opts]
763  (fn opts.fallback [e]
764    (utils.warn (: "include module not found, falling back to require: %s"
765                   :format (tostring e)))
766    (utils.expr (string.format "require(%s)" (tostring e)) :statement))
767
768  (scopes.global.specials.include ast scope parent opts))
769
770(fn compile-stream [strm options]
771  (let [opts (utils.copy options)
772        old-globals allowed-globals
773        scope (or opts.scope (make-scope scopes.global))
774        vals []
775        chunk []]
776    (utils.root:set-reset)
777    (set allowed-globals opts.allowedGlobals)
778    (when (= opts.indent nil)
779      (set opts.indent "  "))
780    (when opts.requireAsInclude
781      (set scope.specials.require require-include))
782    (set (utils.root.chunk utils.root.scope utils.root.options)
783         (values chunk scope opts))
784    (each [_ val (parser.parser strm opts.filename opts)]
785      (table.insert vals val))
786    (for [i 1 (length vals)]
787      (let [exprs (compile1 (. vals i) scope chunk
788                            {:nval (or (and (< i (length vals)) 0) nil)
789                             :tail (= i (length vals))})]
790        (keep-side-effects exprs chunk nil (. vals i))
791        (when (= i (length vals))
792          (utils.hook :chunk (. vals i) scope))))
793    (set allowed-globals old-globals)
794    (utils.root.reset)
795    (flatten chunk opts)))
796
797(fn compile-string [str opts]
798  (compile-stream (parser.string-stream str) (or opts {})))
799
800(fn compile [ast opts]
801  (let [opts (utils.copy opts)
802        old-globals allowed-globals
803        chunk []
804        scope (or opts.scope (make-scope scopes.global))]
805    (utils.root:set-reset)
806    (set allowed-globals opts.allowedGlobals)
807    (when (= opts.indent nil)
808      (set opts.indent "  "))
809    (when opts.requireAsInclude
810      (set scope.specials.require require-include))
811    (set (utils.root.chunk utils.root.scope utils.root.options)
812         (values chunk scope opts))
813    (let [exprs (compile1 ast scope chunk {:tail true})]
814      (keep-side-effects exprs chunk nil ast)
815      (utils.hook :chunk ast scope)
816      (set allowed-globals old-globals)
817      (utils.root.reset)
818      (flatten chunk opts))))
819
820(fn traceback-frame [info]
821  (if (and (= info.what :C) info.name)
822      (string.format "  [C]: in function '%s'" info.name)
823      (= info.what :C)
824      "  [C]: in ?"
825      (let [remap (. sourcemap info.source)]
826        (when (and remap (. remap info.currentline))
827          ;; And some global info
828          (set info.short_src
829               (if (. remap info.currentline 1)
830                   (. sourcemap (.. "@" (. remap info.currentline 1))
831                      :short_src)
832                   remap.short_src))
833          ;; Overwrite info with values from the mapping
834          (set info.currentline (or (. remap info.currentline 2) -1)))
835        (if (= info.what :Lua)
836            (string.format "  %s:%d: in function %s" info.short_src
837                           info.currentline
838                           (if info.name (.. "'" info.name "'") "?"))
839            (= info.short_src "(tail call)")
840            "  (tail call)"
841            (string.format "  %s:%d: in main chunk" info.short_src
842                           info.currentline)))))
843
844(fn traceback [msg start]
845  "A custom traceback function for Fennel that looks similar to debug.traceback.
846Use with xpcall to produce fennel specific stacktraces. Skips frames from the
847compiler by default; these can be re-enabled with export FENNEL_DEBUG=trace."
848  (let [msg (tostring (or msg ""))]
849    (if (and (or (msg:find "^Compile error") (msg:find "^Parse error"))
850             (not (utils.debug-on? :trace)))
851        msg ; skip the trace because it's compiler internals.
852        (let [lines []]
853          (if (or (msg:find ":%d+: Compile error") (msg:find ":%d+: Parse error"))
854              (table.insert lines msg)
855              (let [newmsg (msg:gsub "^[^:]*:%d+:%s+" "runtime error: ")]
856                (table.insert lines newmsg)))
857          (table.insert lines "stack traceback:")
858          (var (done? level) (values false (or start 2)))
859          ;; This would be cleaner factored out into its own recursive
860          ;; function, but that would interfere with the traceback itself!
861          (while (not done?)
862            (match (debug.getinfo level :Sln)
863              nil (set done? true)
864              info (table.insert lines (traceback-frame info)))
865            (set level (+ level 1)))
866          (table.concat lines "\n")))))
867
868(fn entry-transform [fk fv]
869  "Make a transformer for key / value table pairs, preserving all numeric keys"
870  (fn [k v]
871    (if (= (type k) :number)
872        (values k (fv v))
873        (values (fk k) (fv v)))))
874
875(fn mixed-concat [t joiner]
876  (let [seen []]
877    (var (ret s) (values "" ""))
878    (each [k v (ipairs t)]
879      (table.insert seen k)
880      (set ret (.. ret s v))
881      (set s joiner))
882    (each [k v (utils.stablepairs t)]
883      (when (not (. seen k))
884        (set ret (.. ret s "[" k "]" "=" v))
885        (set s joiner)))
886    ret))
887
888;; TODO: too long
889(fn do-quote [form scope parent runtime?]
890  "Expand a quoted form into a data literal, evaluating unquote"
891  (fn q [x]
892    (do-quote x scope parent runtime?))
893
894  (if (utils.varg? form)
895      (do
896        (assert-compile (not runtime?)
897                        "quoted ... may only be used at compile time" form)
898        :_VARARG)
899      (utils.sym? form) ; symbol
900      (let [filename (if form.filename (string.format "%q" form.filename) :nil)
901            symstr (tostring form)]
902        (assert-compile (not runtime?)
903                        "symbols may only be used at compile time" form)
904        ;; We should be able to use "%q" for this but Lua 5.1 throws an error
905        ;; when you try to format nil, because it's extremely bad.
906        (if (or (symstr:find "#$") (symstr:find "#[:.]")) ; autogensym
907            (string.format "sym('%s', {filename=%s, line=%s})"
908                           (autogensym symstr scope) filename
909                           (or form.line :nil))
910            ;; prevent non-gensymed symbols from being bound as an identifier
911            (string.format "sym('%s', {quoted=true, filename=%s, line=%s})"
912                           symstr filename (or form.line :nil))))
913      (and (utils.list? form) ; unquote
914           (utils.sym? (. form 1)) (= (tostring (. form 1)) :unquote))
915      (let [payload (. form 2)
916            res (unpack (compile1 payload scope parent))]
917        (. res 1))
918      (utils.list? form)
919      (let [mapped (utils.kvmap form (entry-transform #nil q))
920            filename (if form.filename (string.format "%q" form.filename) :nil)]
921        (assert-compile (not runtime?) "lists may only be used at compile time"
922                        form)
923        ;; Constructing a list and then adding file/line data to it triggers a
924        ;; bug where it changes the value of # for lists that contain nils in
925        ;; them; constructing the list all in one go with the source data and
926        ;; contents is how we construct lists in the parser and works around
927        ;; this problem; allowing # to work in a way that lets us see the nils.
928        (string.format (.. "setmetatable({filename=%s, line=%s, bytestart=%s, %s}"
929                           ", getmetatable(list()))")
930                       filename (or form.line :nil) (or form.bytestart :nil)
931                       (mixed-concat mapped ", ")))
932      (utils.sequence? form)
933      (let [mapped (utils.kvmap form (entry-transform q q))
934            source (getmetatable form)
935            filename (if source.filename (string.format "%q" source.filename)
936                         :nil)]
937        ;; need to preserve the sequence marker in the metatable here
938        (string.format "setmetatable({%s}, {filename=%s, line=%s, sequence=%s})"
939                       (mixed-concat mapped ", ") filename
940                       (if source source.line :nil)
941                       "(getmetatable(sequence()))['sequence']"))
942      (= (type form) :table) ; table
943      (let [mapped (utils.kvmap form (entry-transform q q))
944            source (getmetatable form)
945            filename (if source.filename (string.format "%q" source.filename)
946                         :nil)]
947        (string.format "setmetatable({%s}, {filename=%s, line=%s})"
948                       (mixed-concat mapped ", ") filename
949                       (if source source.line :nil)))
950      (= (type form) :string)
951      (serialize-string form)
952      (tostring form)))
953
954{: compile
955 : compile1
956 : compile-stream
957 : compile-string
958 : emit
959 : destructure
960 : require-include
961 : autogensym
962 : gensym
963 : do-quote
964 : global-mangling
965 : global-unmangling
966 : apply-manglings
967 :macroexpand macroexpand*
968 : declare-local
969 : make-scope
970 : keep-side-effects
971 : symbol-to-expression
972 ;; general
973 :assert assert-compile
974 : scopes
975 : traceback
976 :metadata (make-metadata)
977 : sourcemap}
978