Common subexpression elimination

cse.clj

; Common sub-expression elimination

(defn- tally
  "Count occurences of every subexpression in expr into the map t1."
  [tl expr]
  (if (seq? expr)
    (merge-with + tl { expr 1 })
    tl))

(declare tally-expr)

(defn- tally-seq
  "Tally a sequence"
  [tl expr type]
  (tally (reduce tally-expr tl expr) (cons type expr)))

(defn- tally-map
  "Tally a map"
  [tl expr]
  (tally 
   (reduce (fn [t [k v]] (-> t (tally-expr k) (tally-expr v))) tl expr)
   expr))

(defn- tally-expr [tl expr]
  (cond (vector? expr) (tally-seq tl expr ::vector)
        (map? expr)    (tally-map tl expr ::map)
        (seq? expr)    (tally-seq tl expr ::seq)
        (set? expr)    (tally-seq tl expr ::set)
        :else          (tally tl expr)))

(defn- ->key
  "convert collections into a tagged seq"
  [expr]
  (cond (vector? expr) (list* ::vector expr)
        (map? expr)    (list* ::map expr)
        (seq? expr)    (list* ::seq expr)
        (set? expr)    (list* ::set expr)
        :else expr))

(defn- key->
  "convert a tagged seq back into a collection"
  [k]
  (if (seqable? k)
    (case (first k)
      ::vector (vec (rest k))
      ::map (into {} (rest k))
      ::seq (rest k)
      ::set (into #{} (rest k))
      k)
    k))

(defn- build-env
  "Takes a tally map and returns an environment of expr -> variable when the count > 1"
  [tl]
  (into {} (map (fn [[k v]] [k (gensym)]) (filter (fn [[k v]] (not= v 1)) tl))))

(defn- lookup [e r]
  (let [res (key-> (r (->key e)))]
    (cond res res
          (seq? e)    (map #(lookup % r) e)
          (vector? e) (into [] (map #(lookup % r) e))
          :else e)))

(defn- cse-lookup[x r]
  (if (seqable? x)
    (map #(lookup % r) x)
    x))        

(defn- build-defs
  "Takes an environment of (expr -> var) and builds a let definition block"
  [env]
  (loop [r env
         out []]
    (if (seq r)
      (let [in (key-> (ffirst r))
            res (lookup in env)]
        (if (not= in res)
          (recur (rest r) (concat out [res (cse-lookup in env)]))
          (recur (rest r) out)))
      (vec out))))
    
(defn cse
  "Takes a form (pure function calls only!) and produces a let form extracting the common subexpressions"
  [expr]
  (let [tl (tally-expr (array-map) expr)
        r (build-env tl)
        defs (build-defs r)]
    `(let ~defs
       ~(cse-lookup expr r))))


(comment
  (cse '(+ (* (+ x y)
              (+ x y))
           (* (+ x y)
              (+ x y))))
  =>
  (clojure.core/let
      [G__31311 (+ x y)
       G__31312 (* G__31311 G__31311)]
    (+ G__31312 G__31312)))