s-expression constant folding experiment

constant folding.scm

(import test matchable)

;; how to turn
`(* 2 (+ x 3 (+ 4 10)) y 10)
;;  into `(* 20 (+ 17 x ) y) ?

`(+ (+ 1 2 x) (+ 3 4 x) 5) =>
`(+ 15 x x)

(begin

  (define (constant? x)
    (or (number? x)
        (string? x)
        (boolean? x)))

  (define (cfold x env)
    ;;(print "folding " x)
    (match x
      ((? number?) x)

      ((? symbol?)
       (cond ((assoc x env) => (lambda (pair)
                                 ;;(print " sy " pair)
                                 (cdr pair)))
             (else x)))

      (('- num)
       (let ((num (cfold num env)))
         (cond ((number? num) #;(print " cp " (- num)) (- num))
               (else (error `(- ,num))))))

      (('- num rest ...)
       (cfold `(+ ,num ,@(map (lambda (x) `(- ,x)) rest)) env))

      (('/ a b)
       (let ((a (cfold a env))
             (b (cfold b env)))
         (if (and (number? a)
                  (number? b))
             (/ a b)
             `(/ ,a ,b))))

      (((and fn (or '+ '*)) rest ...)
       (let loop ((x rest)
                  (u '())
                  (k (cond ((eq? fn '+) 0)
                           ((eq? fn '*) 1))))
         ;;(print "x: " x " u:" u " k:" k " ")
         (if (pair? x)
             (let ((e (cfold (car x) env)))
               (cond
                ;; flatten nested +
                ((and (pair? e)
                      (eq? fn (car e)))
                 ;;(print " fl " e)
                 (loop (append (cdr e) (cdr x))
                       u k))
                ;; handle constant
                ((number? e)
                 (loop (cdr x)
                       u
                       ( (cond ((eq? '+ fn) +)
                               ((eq? '* fn) *)) e k)))
                (else
                 (loop (cdr x)
                       (cons e u)
                       k))))
             (if (pair? u)
                 ;; nonconstants present
                 (cons fn (cons k u))
                 ;; only constants
                 k))))

      (('> a b)
       (let ((a (cfold a env))
             (b (cfold b env)))
         (if (and (constant? a)
                  (constant? b))
             (> a b)
             `(> ,a ,b))))

      (('if test then else)
       (let ((test (cfold test env))
             (then (cfold then env))
             (else (cfold else env)))
         (cond ((eq? #t test) then)
               ((eq? #f test) else)
               (else `(if ,test ,then ,else)))))

      (('let* () body)
       (cfold body env))

      (('let* ((variable value) lets ...) body ...)
       (let ((value (cfold value env)))
         ;;(print " == " variable " : " value)
         (cfold `(let* ,lets ,@body)
                (cons (cons variable value) env))))

      (('let* () body ...)
       (error "only single-form let* body supported"))

      (else (error "TODO don't know how to fold" x))))

  (test-group
   "cfold"
   (test 10 (cfold 10 '()))
   (test 111 (cfold `(+ 1 10 100) '()))

   (test 330 (cfold `(+ 20 200 (+ 10 100)) '()))

   (test `(+ 0 x) (cfold `(+ x) '()))

   (test `(+ 15 x x)
         (cfold `(+ (+ 1 2 x) (+ 3 4 x) 5) '()))

   ;; example problem ported from
   ;; from https://en.wikipedia.org/wiki/Constant_folding
   (test "big example from wikipedia"
         4
         (cfold
          ` (let* ((a 30)
                   (b (- 9 (/ a 5))) ;; TODO + => -
                   (c (* b 4))
                   (c (if (> c 10)
                          (- c 10) ;; TODO + => -
                          c)))
              (* c (/ 60 a)))
          '()))))

;; TODO: ensure constants is first argument.
(match '(+ 0 (+ 1 x) (+ 2 x))
  (('+ A a ... ('+ B b ...))
   (list A B)))

;; can we do anything here?
`(* i (+ (* a b) (* c d))) =>