mal-to-llvm.mal

(load-file "utils.mal")
(load-file "macros-eval.mal")
(def! *macros-eval-env* (new-env repl-env))

(def! *malc-path* (atom "."))
(def! *compile-mode* (atom "release"))

(def! *current-reg* (atom 0))
(def! new-reg
  (fn* []
    (str "%r" (swap! *current-reg* inc))))

(def! *current-stack-reg* (atom 0))
(def! new-stack-reg
  (fn* []
    (str "%sr" (swap! *current-stack-reg* inc))))

(def! *current-label* (atom 0))
(def! new-label
  (fn* [role]
    (str "label" (swap! *current-label* inc) role)))

(def! *current-func-name* (atom 0))
(def! new-func-name
  (fn* []
    (str "@malfunc" (swap! *current-func-name* inc))))

(def! *functions-definitions-list* (atom []))
(def! add-function-definition
  (fn* [code-str]
    (swap! *functions-definitions-list* conj code-str)))

(def! escape-llvm-string
  (fn* [s]
    (replace-chars s {"\n" "\\0a" "\"" "\\22" "\\" "\\5c"})))

(def! linkage-type
  (fn* []
    (if (= @*compile-mode* "debug") "external" "private")))

(def! *current-string-global* (atom 0))
(def! *strings-list* (atom []))
(def! *strings-map* (atom {}))
(def! add-string
  (fn* [val]
    (let* [str-val (str val)]
      (if (contains? @*strings-map* str-val)
        (get @*strings-map* str-val)
        (let* [str-num (swap! *current-string-global* inc)
               str-global (str "@str." str-num)
               len (inc (string/length str-val))
               llvm-escaped (escape-llvm-string str-val)
               def-line (str str-global "= private constant [" len " x i8] c\"" llvm-escaped "\\00\"")]
          (do
            (swap! *strings-list* conj def-line)
            (swap! *strings-map* assoc str-val str-global)
            str-global))))))

(def! *native-funcs* (atom {}))
(def! add-native-func
  (fn* [name llvm-name llvm-args]
    (swap! *native-funcs* assoc (str name) ['llvm-func llvm-name llvm-args])))

(def! *exported-native-funcs* (atom {}))
(def! add-exported-native-func
  (fn* [name]
    (swap! *exported-native-funcs* assoc (str name) true)))

(def! debug-line
  (fn* [message & args]
    (line "; DEBUG" message ":" (pr-str args))))

(def! append-lines
  (fn* [& args]
    args))

(def! line
  (fn* [& args]
    (string/join " " args)))

(def! emit
  (fn* [& args]
    (let* [emit1 (fn* [arg]
                      (if (sequential? arg)
                        (map emit1 arg)
                        (println arg)))]
    (map emit1 args))))

(def! line-label
  (fn* [name]
    (line (str name ":"))))

(def! line-assign
  (fn* [result & args]
    (apply line result "=" args)))

(def! line-assign-to-reg
  (fn* [result arg]
    (line-assign result "call %mal_obj @identity(%mal_obj " arg ")")))

(def! compile-const-integer
  (fn* [result num]
    (line-assign result "call %mal_obj @make_integer(i64" num ")")))

(def! compile-nil
  (fn* [result]
    (line-assign result "call %mal_obj @make_nil()")))

(def! compile-true
  (fn* [result]
    (line-assign result "call %mal_obj @make_true()")))

(def! compile-false
  (fn* [result]
    (line-assign result "call %mal_obj @make_false()")))

(def! compile-const-bytearray
  (fn* [result type-num body]
    (let* [str-len (string/length body)
           str-global (add-string body)]
      (line result "= call %mal_obj @mal_make_bytearray_obj(i32 " type-num ", i32 " str-len ", i8* getelementptr([" (inc str-len) " x i8], [" (inc str-len) " x i8]* " str-global ", i32 0, i32 0))"))))

(def! native-func?
  (fn* [funcname]
    (contains? @*native-funcs* (str funcname))))

(def! get-native-func
  (fn* [funcname]
    (get @*native-funcs* (str funcname))))

(def! exported-native-func?
  (fn* [funcname]
    (contains? @*exported-native-funcs* (str funcname))))

(def! compile-apply-native-func
  (fn* [result ast env-type]
    (let* [funcname (first ast)
           funcobj (get-native-func funcname)
           llvm-func-name (nth funcobj 1)
           argregs (n-entries new-reg (count (rest ast)))
           argregs-with-exprs (zip argregs (rest ast))
           argregs-with-types (map (fn* [a] (str "%mal_obj " a)) argregs)]
      (append-lines
        (debug-line "compile-apply-native-func" result ast)
        (map (fn* [re] (compile (first re) (nth re 1) env-type)) argregs-with-exprs)
        (line result "= call %mal_obj" (str "@" llvm-func-name) "(" (string/join ", " argregs-with-types) ")")))))

(def! compile-apply-mal-func
  (fn* [result ast env-type]
    (let* [funcobj-reg (new-reg)
           funcenv-reg (new-reg)
           binds-reg (new-reg)
           exprs-reg (new-reg)
           funcptr-reg (new-reg)
           casted-funcptr-reg (new-reg)
           callenv-reg (new-reg)]
      (append-lines
        (debug-line "compile-apply-mal-func" result ast)
        (compile funcobj-reg (first ast) env-type)
        (line-assign funcenv-reg "call %mal_obj @fn_env(%mal_obj " funcobj-reg ")")
        (line-assign binds-reg "call %mal_obj @fn_args_names(%mal_obj " funcobj-reg ")")
        (compile-vector exprs-reg (rest ast) env-type)
        (line-assign callenv-reg "call %mal_obj @new_env(%mal_obj" funcenv-reg ", %mal_obj" binds-reg ", %mal_obj" exprs-reg ")")
        (line-assign funcptr-reg "call %mal_obj @fn_func_ptr(%mal_obj " funcobj-reg ")")
        (line-assign casted-funcptr-reg "inttoptr %mal_obj" funcptr-reg "to %mal_obj(%mal_obj)*")
        (line-assign result "call %mal_obj" casted-funcptr-reg "(%mal_obj " callenv-reg ")")))))

(def! compile-apply
  (fn* [result ast env-type]
    (let* [funcname (first ast)]
      (append-lines
        (debug-line "compile-apply" result ast)
        (cond
          (= env-type 'native-func)
            (cond
              (native-func? funcname)
                (compile-apply-native-func result ast env-type)
              :else
                (throw (str "Unknown function '" funcname "'")))

          (= env-type 'mal-env)
            (if (exported-native-func? funcname)
              (compile-apply-native-func result ast env-type) ; Optimized function application for built-in functions
              (compile-apply-mal-func result ast env-type))

          :else
            (throw "Unknown env-type"))))))

(def! compile-native-func
  (fn* [name llvm-name params & body]
    (let* [params-with-types (map (fn* [p] (str "%mal_obj %"  p)) params)]
      (do
        (add-native-func name llvm-name params)
        (if (empty? body)
          (append-lines (line ""))
          (append-lines
            (line "")
            (line "")
            (line "define" (linkage-type) "%mal_obj" (str "@" llvm-name) "(" (string/join ", " params-with-types) ") {")
            (compile "%funcresult" (first body) 'native-func)
            (line "ret %mal_obj %funcresult")
            (line "} ; end of native func " llvm-name)
            (line "")))))))

(def! llvm-func-type
  (fn* [args-num]
    (cond
      (= args-num 0) "%mal_obj()*"
      (= args-num 1) "%mal_obj(%mal_obj)*"
      (= args-num 2) "%mal_obj(%mal_obj,%mal_obj)*"
      (= args-num 3) "%mal_obj(%mal_obj,%mal_obj,%mal_obj)*"
      :else          (throw "Too many arguments to native func"))))

(def! compile-nativefn-value
  (fn* [result llvm-name args-names]
    (let* [args-num (count args-names)
           args-names-reg (new-reg)
           llvm-name-reg (new-reg)
           func-ptr-reg (new-reg)
           n0 (new-reg)
           n1 (new-reg)
           n2 (new-reg)]
      (append-lines
        (compile-literal-list args-names-reg args-names 'native-func)
        (compile-literal llvm-name-reg llvm-name 'native-func)
        (line-assign result "call %mal_obj @make_native_func()")
        (compile-const-integer n0 0)
        (compile-const-integer n1 1)
        (compile-const-integer n2 2)
        (line "call %mal_obj @mal_set_elementarray_item("
              "%mal_obj" result ", %mal_obj" n0 ", %mal_obj" args-names-reg ")")
        (line "call %mal_obj @mal_set_elementarray_item("
              "%mal_obj" result ", %mal_obj" n1 ", %mal_obj" llvm-name-reg ")")
        (line-assign func-ptr-reg "ptrtoint" (llvm-func-type args-num) (str "@" llvm-name) "to %mal_obj")
        (line "call %mal_obj @mal_set_elementarray_item("
              "%mal_obj" result ", %mal_obj" n2 ", %mal_obj" func-ptr-reg ")")))))

(def! compile-if
  (fn* [result ast env-type]
    (let* [cond-exp (nth ast 1)
           true-exp (nth ast 2)
           false-exp (if (> (count ast) 3) (nth ast 3) nil)
           result-stack-reg (new-stack-reg)
           cond-reg (new-reg)
           tmp-true-result (new-reg)
           tmp-false-result (new-reg)
           true-label (new-label ".if_true")
           false-label (new-label ".if_false")
           after-if-label (new-label ".if_end")]
      (append-lines
        (line-assign result-stack-reg "alloca %mal_obj")
        (compile cond-reg cond-exp env-type)
        (line "switch %mal_obj" cond-reg ", label" (str "%" true-label)
              "[ %mal_obj 2, label" (str "%" false-label) "\n"
              "%mal_obj 4, label" (str "%" false-label) "]")
        (line-label true-label)
        (compile tmp-true-result true-exp env-type)
        (line "store %mal_obj" tmp-true-result ", %mal_obj*" result-stack-reg)
        (line "br label" (str "%" after-if-label))
        (line-label false-label)
        (compile tmp-false-result false-exp env-type)
        (line "store %mal_obj" tmp-false-result ", %mal_obj*" result-stack-reg)
        (line "br label" (str "%" after-if-label))
        (line-label after-if-label)
        (line-assign result "load %mal_obj, %mal_obj*" result-stack-reg)))))

(def! compile-try-catch
  (fn* [result ast env-type]
    (let* [try-exp (nth ast 1)
           catch-clause (nth ast 2)
           catch-exn-var-name (nth catch-clause 1)
           catch-exp (nth catch-clause 2)
           try-func-name (new-func-name)
           catch-func-name (new-func-name)
           result-stack-reg (new-stack-reg)
           try-result (new-reg)
           landingpad-exn-reg (new-reg)
           catch-exn-reg (new-reg)
           catch-exn-var-reg (new-reg)
           catch-env-reg (new-reg)
           catch-result (new-reg)
           cont-label (new-label ".try_cont")
           catch-label (new-label ".try_catch")
           after-try-label (new-label ".try_end")]
      (do
        (add-function-definition (define-mal-func try-func-name try-exp))
        (add-function-definition (define-mal-func catch-func-name catch-exp))
        (append-lines
          (line-assign result-stack-reg "alloca %mal_obj")
          (line-assign try-result "invoke %mal_obj" try-func-name "(%mal_obj %env)"
                       "to label" (str "%" cont-label)
                       "unwind label" (str "%" catch-label))
          (line-label cont-label)
          (line "store %mal_obj" try-result ", %mal_obj*" result-stack-reg)
          (line "br label" (str "%" after-try-label))
          (line-label catch-label)
          (line-assign landingpad-exn-reg "landingpad {i8*,i32} catch i8* null")
          (line-assign catch-exn-reg "call %mal_obj @mal_get_exception_from_landing_pad({i8*,i32}" landingpad-exn-reg ")")
          (compile-literal catch-exn-var-reg catch-exn-var-name env-type)
          (line-assign catch-env-reg "call %mal_obj @new_singlevar_env(%mal_obj %env, %mal_obj" catch-exn-var-reg ", %mal_obj" catch-exn-reg ")")
          (line-assign catch-result "call %mal_obj" catch-func-name "(%mal_obj " catch-env-reg ")")
          (line "store %mal_obj" catch-result ", %mal_obj*" result-stack-reg)
          (line "br label" (str "%" after-try-label))
          (line-label after-try-label)
          (line-assign result "load %mal_obj, %mal_obj*" result-stack-reg))))))

; AST manipulation: rewrite let* as an anonymous fn* with def! inside
; Converts (let* [x 8 y (+ x 2)] (body ...))
; to ((fn* [] (do (def! x 8) (def! y (+ x 2)) (body ...))))
(def! convert-let*-to-fn*
  (fn* [ast]
    (let* [vars-to-defs (fn* [vars defs]
                          (if (empty? vars)
                            defs
                            (let* [var1-name (nth vars 0)
                                   var1-value (nth vars 1)
                                   rest-vars (rest (rest vars))]
                              (vars-to-defs rest-vars (concat defs [(list 'def! var1-name var1-value)])))))
           defs (vars-to-defs (nth ast 1) [])
           body (nth ast 2)]
      (list (list 'fn* [] (concat (list 'do) defs [body]))))))

(def! compile-load-file
  (fn* [result filename env-type]
    (append-lines
      (debug-line "load-file" filename)
      (compile result (read-string (str "(do " (slurp filename) ")")) env-type))))

(def! compile-list
  (fn* [result ast env-type]
    (let* [a0 (first ast)]
      (cond
        (empty? ast)
          (compile-literal-list result ast env-type)

        (= 'defmacro! a0)
          (do
            (EVAL ast *macros-eval-env*)
            (if result
              (compile-nil result)
              (line "")))

        (= 'macroexpand a0)
          (compile-literal result (expand-macros (nth ast 1)) env-type)

        (= 'load-file a0)
          (compile-load-file result (nth ast 1) env-type)

        (= 'def! a0)
          (let* [var-reg (new-reg)
                 val-reg (new-reg)]
            (append-lines
              (compile-literal var-reg (nth ast 1) env-type)
              (compile val-reg (nth ast 2) env-type)
              (line-assign result "call %mal_obj @env_set(%mal_obj %env, %mal_obj " var-reg ", %mal_obj " val-reg ")")))

        (= 'let* a0)
          (compile result (convert-let*-to-fn* ast) env-type)

        (= 'quote a0)
          (compile-literal result (nth ast 1) env-type)

        (= 'do a0)
          (let* [exprs (rest ast)
                 regs-and-lines (map (fn* [e] (compile-temp e env-type)) exprs)
                 lines (map (fn* [p] (nth p 1)) regs-and-lines)]
            (append-lines
              lines
              (if result
                (line-assign-to-reg result (first (last regs-and-lines)))
                (line ""))))

        (= 'if a0)
          (compile-if result ast env-type)

        (= 'try* a0)
          (if (> (count ast) 2)
            (compile-try-catch result ast env-type)
            (compile result (nth ast 1) env-type)) ; catchless try*

        (= 'fn* a0)
          (compile-fn result ast env-type)

        (= 'defnativefn a0)
          (apply compile-native-func (rest ast))

        (= 'exportnativefn a0)
          (do
            (map add-exported-native-func (rest ast))
            (line ""))

        :else
          (compile-apply result ast env-type)))))

(def! compile-temp
  (fn* [ast env-type]
    (let* [r (new-reg)]
      [r (compile r ast env-type)])))

(def! compile-variable
  (fn* [result var-name env-type]
    (cond
      (= env-type 'native-func)
        (line-assign-to-reg result (str "%" var-name))

      (= env-type 'mal-env)
        (if (exported-native-func? var-name)
          (let* [funcobj (get-native-func var-name)
                 llvm-func-name (nth funcobj 1)
                 llvm-func-args (nth funcobj 2)]
            (compile-nativefn-value result llvm-func-name llvm-func-args))
          (let* [var-name-reg (new-reg)]
            (append-lines
              (compile-literal var-name-reg var-name env-type)
              (line-assign result "call %mal_obj @env_get(%mal_obj %env, %mal_obj" var-name-reg ")"))))

      :else
        (throw "Unknown env-type"))))

(def! set-elementarray-item
  (fn* [result index item-reg]
    (let* [index-reg (new-reg)]
      (append-lines
        (compile-const-integer index-reg index)
        (line "call %mal_obj @mal_set_elementarray_item("
              "%mal_obj" result
              ", %mal_obj" index-reg
              ", %mal_obj" item-reg ")")))))

(def! compile-elementarray
  (fn* [result type-num itemsregs env-type]
    (let* [type-reg (new-reg)
           len-reg (new-reg)]
      (append-lines
        (compile type-reg type-num env-type)
        (compile len-reg (count itemsregs) env-type)
        (line result "= call %mal_obj @mal_make_elementarray_obj(%mal_obj " type-reg ", %mal_obj " len-reg ")")
        (vector-map-with-index (fn* [r index] (set-elementarray-item result index r)) itemsregs)))))

(def! compile-literal-elementarray
  (fn* [type-num result ast env-type]
    (let* [regs (n-entries new-reg (count ast))
           regs-and-exprs (zip regs ast)
           lines (map (fn* [re] (compile-literal (first re) (nth re 1) env-type)) regs-and-exprs)]
      (append-lines
        lines
        (compile-elementarray result type-num regs env-type)))))

(def! compile-literal-list
  (fn* [result ast env-type]
    (compile-literal-elementarray 33 result ast env-type)))

(def! compile-literal-vector
  (fn* [result ast env-type]
    (compile-literal-elementarray 34 result ast env-type)))

(def! compile-vector
  (fn* [result ast env-type]
    (let* [regs (n-entries new-reg (count ast))
           regs-and-exprs (zip regs ast)
           lines (map (fn* [re] (compile (first re) (nth re 1) env-type)) regs-and-exprs)]
      (append-lines
        lines
        (compile-elementarray result 34 regs env-type)))))

(def! map-to-list-helper
  (fn* [hm result ks]
    (if (empty? ks)
      result
      (map-to-list-helper hm (concat result [(first ks) (get hm (first ks))]) (rest ks)))))

(def! map-to-list
  (fn* [hm]
    (map-to-list-helper hm [] (keys hm))))

(def! compile-hash-map
  (fn* [result ast env-type]
    (let* [lst (map-to-list ast)
           regs (n-entries new-reg (count lst))
           regs-and-exprs (zip regs lst)
           lines (map (fn* [re] (compile (first re) (nth re 1) env-type)) regs-and-exprs)]
      (append-lines
        lines
        (compile-elementarray result 35 regs env-type)))))

(def! compile-fn
  (fn* [result ast env-type]
    (let* [func-name (new-func-name)
           args-names (nth ast 1)
           args-names-reg (new-reg)
           body (nth ast 2)
           func-ptr-reg (new-reg)
           n0 (new-reg)
           n1 (new-reg)
           n2 (new-reg)]
      (do
        (add-function-definition (define-mal-func func-name body))
        (append-lines
          (compile-literal-list args-names-reg args-names env-type)
          (line-assign result "call %mal_obj @make_func()")
          (compile-const-integer n0 0)
          (compile-const-integer n1 1)
          (compile-const-integer n2 2)
          (line "call %mal_obj @mal_set_elementarray_item("
                "%mal_obj" result ", %mal_obj" n0 ", %mal_obj" args-names-reg ")")
          (line "call %mal_obj @mal_set_elementarray_item("
                "%mal_obj" result ", %mal_obj" n1 ", %mal_obj %env)")
          (line-assign func-ptr-reg "ptrtoint %mal_obj(%mal_obj)*" func-name "to %mal_obj")
          (line "call %mal_obj @mal_set_elementarray_item("
                "%mal_obj" result ", %mal_obj" n2 ", %mal_obj" func-ptr-reg ")"))))))

(def! compile-literal
  (fn* [result ast env-type]
    (append-lines
      (debug-line "compile-literal" result ast)
      (cond
        (nil? ast)     (compile-nil result)
        (true? ast)    (compile-true result)
        (false? ast)   (compile-false result)
        (number? ast)  (compile-const-integer result ast)
        (symbol? ast)  (compile-const-bytearray result 17 ast)
        (string? ast)  (compile-const-bytearray result 18 ast)
        (keyword? ast) (compile-const-bytearray result 19 ast)
        (list? ast)    (compile-literal-list result ast env-type)
        (vector? ast)  (compile-literal-vector result ast env-type)))))

(def! literal?
  (fn* [ast]
    (cond
      (nil? ast)     true
      (true? ast)    true
      (false? ast)   true
      (number? ast)  true
      (symbol? ast)  true
      (string? ast)  true
      (keyword? ast) true
      :else          false)))

(def! compile-without-macro-expand
  (fn* [result ast env-type]
    (append-lines
      (cond
        (symbol? ast)  (compile-variable result ast env-type)
        (literal? ast) (compile-literal result ast env-type)
        (list? ast)    (compile-list result ast env-type)
        (vector? ast)  (compile-vector result ast env-type)
        (map? ast)     (compile-hash-map result ast env-type)))))

(def! expand-macros
  (fn* [ast]
    (if (list? ast)
      (MACROEXPAND ast *macros-eval-env*)
      ast)))

(def! compile
  (fn* [result ast env-type]
    (let* [exp-ast (expand-macros ast)]
    (append-lines
      (debug-line "compile     ast=" ast)
      (debug-line "compile exp-ast=" exp-ast)
      (compile-without-macro-expand result exp-ast env-type)))))

(def! define-mal-func
  (fn* [name body]
    (append-lines
      (line "")
      (line "")
      (line "define" (linkage-type) "%mal_obj" name "(%mal_obj %env) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {")
      (compile "%funcresult" body 'mal-env)
      (line "ret %mal_obj %funcresult")
      (line "} ; end of function" name)
      (line ""))))

(def! emit-strings-list
  (fn* []
    (map emit @*strings-list*)))

(def! emit-functions-definitions
  (fn* []
    (map emit @*functions-definitions-list*)))

(def! malc-file
  (fn* [filename]
    (str @*malc-path* "/" filename)))

(def! compile-program-string
  (fn* [s]
    (let* [runtime-native-funcs-ast (read-string (str "(do " (slurp (malc-file "runtime-native-funcs.mal")) ")"))
           runtime-core-funcs-ast (read-string (str "(do " (slurp (malc-file "runtime-core-funcs.mal")) ")"))
           tokenizer-ast (read-string (str "(do " (slurp (malc-file "tokenizer.mal")) ")"))
           reader-ast (read-string (str "(do " (slurp (malc-file "reader.mal")) ")"))
           user-prog-ast (read-string (str "(do " s ")"))
           whole-ast (list 'do runtime-core-funcs-ast tokenizer-ast reader-ast user-prog-ast)]
      (do
        (emit (slurp (malc-file "runtime-header.ll")))
        (emit (compile nil runtime-native-funcs-ast 'native-func))
        (emit (define-mal-func "@mal_prog_main" whole-ast))
        (emit-functions-definitions)
        (emit-strings-list)
        (emit (slurp (malc-file "runtime-footer.ll")))))))

(def! compile-program-file
  (fn* [filename]
    (compile-program-string (slurp filename))))

(def! -main
  (fn* [compile-mode malc-path source-filename]
    (do
      (reset! *compile-mode* compile-mode)
      (reset! *malc-path* malc-path)
      (compile-program-file source-filename))))

(if (< (count *ARGV*) 3)
  (println "Usage: malc <debug|release> path-to-malc sourcefile.mal")
  (apply -main *ARGV*))