Algebraic simplifier

simplify.clj

;; Algebraic simplifier for fields
(declare simplify)

(defn- var->sym
  [^clojure.lang.Var v]
  (symbol (str (.name (.ns v)))
          (str (.sym v))))

(defn- resolve-sym[x]
  (let [z (resolve x)]
    (if z (var->sym z) x)))

(defn- canonical [a b]
  "Canonical ordering between two expressions a and b"
  (cond (and (number? a) (seq? b))    -1
        (and (number? a) (symbol? b)) -1
        (and (symbol? a) (number? b))  1        
        (and (symbol? a) (seq? b))    -1
        (and (seq? a)    (number? b))  1
        (and (seq? a)    (symbol? b))  1
        (and (seq? a)    (seq? b))
        (canonical (first a) (first b)) ; fixme- should be better at sorting.
        :else (compare a b)))

(defn- associative-transform [op form]
  "If form starts with op, rewrite (op (op a b) c (op d)) => (op a b c d)"
  (if (and (seq? form) (= (first form) op))
    (cons op (mapcat (fn [x]
                       (if (and (seq? x) (= (first x) op))
                         (rest x)
                         (list x)))
                     (rest form)))
    form))

(defn- collect-constants [reducefn form]
  "Collect all the constants in form reducing them with reducefn"
  (if (seq? form)
    (let [k (reduce reducefn (filter number? (rest form)))]
      (if k (list* (first form) k (filter (complement number?) (rest form)))
          form))
    form))

(defn- singleton [k form]
  (cond (and (seq? form) (= (count form) 2))
        (second form)
        (and (seq? form) (= (count form) 1))
        k
        :else form))

(defn- commutative [expr]
  (if (seq? expr)
    (cons (first expr)
          (sort canonical (map commutative (rest expr))))
    expr))

;; fixpoint of f(x)
(defn- fixpoint [f x]
  (loop [x x]
    (let [res (f x)]
      (if (= res x) x
          (recur res)))))


(defn- identity-elem [x form]
  (if (and (seq? form) (some #{x} form))
    (filter (complement #{x}) form)
    form))

(defn- zero-elem [x form]
  (if (and (seq? form) (some #{x} form))
    x
    form))


(defn- expand-minus [form]
  (cond (= (count form) 1) 0
        (= (count form) 2) `(~(var->sym #'*) -1 ~(second form))
        (> (count form) 2) (list* (var->sym #'+)
                                  (second form)
                                  (map (fn [x] `(~(var->sym #'*) -1 ~x)) (rest (rest form))))))

(defn- without-coeff [x]
  (if (seq? x)
    (cons (first x) (rest (filter #(not (number? %)) x)))
    x))

(defn- get-coeff [none x]
  (if (and (seq? x) (number? (second x)))
    (second x)
    none))

(defn- like-terms
  "Checks if x = y without constants at it's toplevel"
  [x y]
  (= (without-coeff y) (without-coeff x)))

(defn- collect-like-terms1 [op fn args]
  (if (seq args)
    (cond (and (seq? (first args))
               (symbol? (ffirst args))
               (= (name (ffirst args)) (name op)))
          (let [[t & rst] args
                like (without-coeff t)
                likes (filter #(like-terms like %) args)
                nlikes (filter #(not (like-terms like %)) args)]
            (cons `(~op ~(reduce fn (map (partial get-coeff 1) likes))
                       ~@(rest (without-coeff t)))
                  (collect-like-terms1 op fn nlikes)))
          :else
          (cons (first args) (collect-like-terms1 op fn (rest args))))
    ()))

(defn- starts-with[c x]
  (and (seq? x) (symbol? (first x))
       (= (name (first x)) (name c))))

(defn- expand-simple[x]
  (cond (or (number? x) (starts-with '* x)) x
        :else
      `(~(var->sym #'*) 1 ~x)))

(defn- collect-like-terms [op fn argmap expr]
  (if (seq? expr)
    (cons (first expr)
          (collect-like-terms1 op fn (map argmap (rest expr))))
    expr))
        
(declare simplify)

(defn- simplify-args[x]
  (cons (first x)
        (map simplify (rest x))))

(def ^:private simplification
  {"-" [[:--expand-minus expand-minus]]
   "+" [[:+-associative-transform (partial associative-transform '+)]
        [:+-collect-constants (partial collect-constants +)]
        [:+-collect-like-terms (partial collect-like-terms '* + expand-simple)]
        [:+-simplify-args simplify-args]
        [:+-identity-elem (partial identity-elem 0)]
        [:+-singleton (partial singleton 0)]
        [:+-commutative commutative]]
   "*" [[:*-associative-transform (partial associative-transform '*)]
        [:*-collect-constants (partial collect-constants *)]
        [:*-zero-elem (partial zero-elem 0)]
        [:*-identity-elem (partial identity-elem 1)]
        [:*-singleton (partial singleton 1)]
        [:*-commutative commutative]]})

(defn- maybe-resolve [x]
  (if (and (symbol? x)
           (resolve x))
    (symbol (var->sym (resolve x)))
    x))


(defn- apply-simplification [form]
  (if (seq? form)
    (let [head (name (maybe-resolve (first form)))]
      (loop [[[nm f] & rest] (get simplification head)
             expr form]
        (if (and (seq? expr) f (= head (name (maybe-resolve (first expr)))))
          (let [next (f expr)]
            (recur rest next))
          expr)))
    form))

(defn- simplify1 [a]
  (if (seq? a)
    (let [sa (cons (first a)
                   (map simplify (rest a)))]
      (fixpoint apply-simplification sa))
    a))

(defn simplify [a]
  (fixpoint simplify1 a))

;; TODO:
;; basic integer power transforms (* x x x) => (Math/pow x 3)
;; basic quotient transform
;;    (/ u v) => (* u (Math/pow v -1))


(comment
  (simplify
   '(+ (* y 2 x z)
       (* 3 x (* y z))
       (* 4 (* z y) x))))
  => (* 9 x y z))