Toy bytecode compiler for integer arithmetic DSL

core.clj

(ns clojr.core
  (:require [instaparse.core :as insta])
  (:import (clojure.lang DynamicClassLoader)))

(import '[clojure.asm ClassWriter Type Opcodes]
        '[clojure.asm.commons Method GeneratorAdapter])

(def parser
  (insta/parser
    "prog = expr (<#'[\\n]+'> expr?)*
     <expr> = <spaces> add-sub <spaces>
     <add-sub> = mul-div | add | sub
     add = add-sub <spaces> <'+'> <spaces> mul-div
     sub = add-sub <spaces> <'-'> <spaces> mul-div
     <mul-div> = term | mul | div
     mul = mul-div <spaces> <'*'> <spaces> term
     div = mul-div <spaces> <'/'> <spaces> term
     <term> = number | get-var | argument | assignment | <'('> <spaces> add-sub <spaces> <')'>
     <spaces> = <#'[ ]*'>
     number = #'-?\\d+'
     varname = #'[A-Za-z]\\w*'
     argument = <'%'> #'\\d+'
     get-var = varname
     assignment = varname <spaces> <'='> expr"))

(defn nb-args
  "Returns the nb of args of the code in the ast. It is the highest value XX of the %XX parsed."
  [ast]
  (->> ast
       (insta/transform {:varname (fn [&d] nil)
                         :number (fn [&d] nil)
                         :argument #(Integer/parseInt %)})
       flatten
       (remove (some-fn nil? keyword?))
       (into [0]);; to handle no args as 0 args
       (apply max)))

(defn var-idx
  "Returns the map of local variables to idx starting at 1.
  Note that we use the fact that all vars are int so they take 1 slot so idx are consecutive."
  [ast]
  (zipmap
   (->> ast
        (insta/transform (assoc content-transf
                           :varname identity
                           :number (fn [&d] nil)
                           :argument (fn [&d] nil)) )
        flatten
        (remove (some-fn nil? keyword?)))
   (iterate inc 1)))

(defn ast-type
  "Transforms the ast to replaces numerical values, arg num and local variables
  by Integers (values or idx). Also swap lhs and rhs of assignment as we will want to first
  compute the value to assign and then store it."
  [ast]
  (let [vars-idx (var-idx ast)]
    (insta/transform {:number (fn [n] [:number (Integer/parseInt n)])
                      :argument (fn [n] [:argument (dec (Integer/parseInt n))])
                      :varname (fn [n] [:varname (vars-idx n)])
                      :assignment (fn [lhs rhs] [:assignment rhs lhs])}
                     ast)))

(defn prototype
  "creates the prototype for the function with the implementation in the given ast,
  computing the number of arguments"
  [f-name n]
  (str "int " f-name "(" (apply str (interpose ","  (repeat n "int "))) ")"))

(defn transform
  "like instaparse/transform, but also works on typed AST with non String data (Integer)"
  [trans-map [tag & data]]
  (let [children (map #(if(vector? %) (transform trans-map %) %) data)]
    (if-let [trans (trans-map tag)]
      (apply trans children)
      (into  [tag]  children))))

(defn fold
  "Performs constants folding by interpreting the AST parts that contains only numbers and binary ops."
  [typed-ast]
  (let[op-fold (fn [[tag op]]
                 [tag  (fn [& data]
                         (if (every? #(= (first %) :number) data)
                           [:number  (apply op (map second data))]
                           (into [tag ] data)))])]
    (transform  (->> [[:add +] [:sub -] [:mul *] [:div /]] (map op-fold) (reduce #(apply assoc %1 %2) {})) typed-ast)))

(defn ast-compile!
  "Generates the implementation of the code from typed-ast into the GeneratorAdapter."
  [typed-ast ga]
  (let [int-t (Type/getType Integer/TYPE)]
    (->> typed-ast (transform {:number (fn [n] (.push ga (int n)))
                               :add  (fn [& exprs] (.math ga GeneratorAdapter/ADD int-t))
                               :sub  (fn [& exprs] (.math ga GeneratorAdapter/SUB int-t))
                               :div  (fn [& exprs] (.math ga GeneratorAdapter/DIV int-t))
                               :mul  (fn [& exprs] (.math ga GeneratorAdapter/MUL int-t))
                               :assignment (fn [v [t n]] (do (.storeLocal ga n int-t)
                                                            (.loadLocal ga n int-t)));;we want the value to stay available
                               :argument (fn [n] (.loadArg ga n))
                               :get-var (fn [[t n]] (.loadLocal ga n int-t))}))))

(defn generate-cstor!
  "Generates the default constructor in the given ClassWriter."
  [cw]
  (let [init (Method/getMethod "void <init>()")]
     (doto (GeneratorAdapter. Opcodes/ACC_PUBLIC init nil nil cw)
      (.visitCode)
      (.loadThis)
      (.invokeConstructor (Type/getType Object) init)
      (.returnValue)
      (.endMethod))))

(defn generate-class!
  "Generates class implementing the given code in a static method named run of a class of the given name."
  [code class-name]
  (let [cw (doto (ClassWriter. ClassWriter/COMPUTE_FRAMES)
             (.visit Opcodes/V1_6 Opcodes/ACC_PUBLIC (.replace class-name \. \/) nil "java/lang/Object" nil)
             (generate-cstor!))
        ast (insta/parse parser code)
        n (nb-args ast)
        compile! (->> ast ast-type fold (partial ast-compile!))
        method (Method/getMethod (prototype "run" n))
        _ (doto (GeneratorAdapter. (+ Opcodes/ACC_PUBLIC Opcodes/ACC_STATIC) method nil nil cw)
            (.visitCode)
            (compile!)
            (.returnValue)
            (.endMethod))]
    (do
      (.visitEnd cw)
      (.defineClass (DynamicClassLoader.) class-name (.toByteArray cw) nil)
      (comment (doto (java.io.FileOutputStream. (str class-name ".class"))
                 (.write (.toByteArray cw))
                 (.close)))
      n)))

(comment
  (generate-class! "a=(%1 + 2) *3
b=a / 5 + %2" "cmp.Test")
  (cmp.Test/run 3 2)
  )

(defmacro compile-code [code]
  (let [class-name (str (gensym "compiled.C"))
        n (generate-class! code class-name)
        args (repeatedly n (partial gensym "arg"))]
    `(fn [~@args] (~(symbol (str class-name "/run")) ~@args))))

(comment
  (def regular-function (compile-code
                         "a=(%1 + 2) *3
                        b=a / 5 + %2"))
  (regular-function 5 2)
  )