Super Simple Nanopass CAS

#lang racket
(require parser-tools/lex
         (prefix-in : parser-tools/lex-sre)
         parser-tools/yacc
         nanopass/base)

(define-tokens math [CNST VAR])
(define-empty-tokens math/e [PLS MIN TIM LEFT RIGHT EQ END])

(define lex-math
  (lexer
   [(:+ whitespace) (lex-math input-port)]
   [lower-case (token-VAR (string->symbol lexeme))]
   [(:or "(" "[") 'LEFT]
   [(:or ")" "]") 'RIGHT]
   ["+" 'PLS] ["-" 'MIN] ["*" 'TIM] ["=" 'EQ]
   [(:: (:+ numeric) (:? "." (:+ numeric)))
    (token-CNST (string->number lexeme))]
   [#\newline 'END]
   [(eof) 'END]))

(define-syntax define-parse
  (syntax-rules ()
    [(_ name clause ...)
     (define name
       (let ([parse (parser clause ...)])
         (λ (port)
           (parse (λ () (lex-math port))))))]))

(define-parse parse-math
  (tokens math math/e)
  (start eqn)
  (end END)
  (error (λ _ (error "invalid equation")))
  (grammar
   (eqn
    [(expr EQ expr) `(solve ,$1 ,$3)])
   (expr
    [(expr PLS term) `(op+ ,$1 ,$3)]
    [(expr MIN term) `(op- ,$1 ,$3)]
    [(term) $1])
   (term
    [(term TIM fact) `(op* ,$1 ,$3)]
    [(term nn-fact) `(op* ,$1 ,$2)]
    [(fact) $1])
   (fact
    [(MIN nn-fact) `(op~ ,$2)]
    [(nn-fact) $1])
   (nn-fact
    [(CNST) $1]
    [(VAR) $1]
    [(LEFT expr RIGHT) $2])))


;-------
(define-language Aorig
  (terminals
   (number [c])
   (nonnegative-integer [n m])
   (symbol [x y]))
  (entry Solve)
  (Arith (a)
    c
    x
    (op+ a1 a2)
    (op- a1 a2)
    (op~ a)
    (op* a1 a2))
  (Solve (so)
    (solve a1 a2)))

(define-parser parse-sexp Aorig)


;-------
(define-language Aeqzero
  (extends Aorig)
  (Solve (so)
    (- (solve a1 a2))
    (+ (solve a))))

(define-pass rhs-zero : Aorig (ir) -> Aeqzero ()
  (definitions)
  (Solve : Solve (so) -> Solve ()
    [(solve ,[a1] ,[a2]) `(solve (op- ,a1 ,a2))]))

;-------
(define-language Anosub
  (extends Aeqzero)
  (Arith (a)
    (- (op- a1 a2))))

(define-pass remove-subtract : Aeqzero (ir) -> Anosub ()
  (definitions)
  (Arith : Arith (a) -> Arith ()
    [(op- ,[a1*] ,[a2*]) `(op+ ,a1* (op~ ,a2*))]))

;-------
(define-language Anoneg
  (extends Anosub)
  (Arith (a)
    (- (op~ a))))

(define-pass remove-negate : Anosub (ir) -> Anoneg ()
  (definitions)
  (Arith : Arith (a) -> Arith ()
    [(op~ ,[a*]) `(op* -1 ,a*)]))

;-------
(define-language Atiered
  (extends Anoneg)
  (Factor (f)
    (+ c x))
  (Term (t)
    (+ (op* t1 t2) f))
  (Arith (a)
    (- c x)
    (- (op* a1 a2))
    (+ t)))

(define-pass distribute : Anoneg (ir) -> Atiered ()
  (definitions
    (define (mult lh rh)
      (nanopass-case (Atiered Arith) lh
        [(op+ ,a1 ,a2)
         (with-output-language (Atiered Arith)
           `(op+ ,(mult a1 rh) ,(mult a2 rh)))]
        [,t
         (nanopass-case (Atiered Arith) rh
           [(op+ ,a1 ,a2) (mult rh lh)]
           [,t2
            (with-output-language (Atiered Term)
              `(op* ,t ,t2))])])))
  (Arith : Arith (a) -> Arith ()
    [,c c]
    [,x x]
    [(op* ,[a1*] ,[a2*]) (mult a1* a2*)]))

;--------
(define-language Atiered*
  (extends Atiered)
  (Term (t)
    (- (op* t1 t2) f)
    (+ (prod f ...)))
  (Arith (a)
    (- (op+ a1 a2) t)
    (+ (sum t ...))))

(define-pass flat-trees : Atiered (ir) -> Atiered* ()
  (definitions
    (define (factors t)
      (nanopass-case (Atiered* Term) t
        [(prod ,f^ ...) f^]))
    (define (terms a)
      (nanopass-case (Atiered* Arith) a
        [(sum ,t^ ...) t^])))
  (Factor : Factor (f) -> Factor ())
  (Term : Term (t) -> Term ()
    [,f
     `(prod ,(Factor f))]
    [(op* ,[t1*] ,[t2*])
     `(prod ,(append (factors t1*) (factors t2*)) ...)])
  (Arith : Arith (a) -> Arith ()
    [,t
     `(sum ,(Term t))]
    [(op+ ,[a1*] ,[a2*])
     `(sum ,(append (terms a1*) (terms a2*)) ...)]))

;--------
(define-language Aexpts
  (extends Atiered*)
  (Term (t)
    (- (prod f ...))
    (+ (c [x n] ...))))

(define-pass simplify-terms : Atiered* (ir) -> Aexpts ()
  (definitions)
  (Term : Term (t) -> Term ()
    [(prod ,f^ ...)
     (let ([vars (make-hash)]
           [constant 1])
       (for ([f (in-list f^)])
         (nanopass-case (Atiered* Factor) f
           [,x (hash-update! vars x add1 0)]
           [,c (set! constant (* c constant))]))
       (let ([vars+ (sort (hash->list vars)
                          symbol<?
                          #:key car)])
         `(,constant [,(map car vars+) ,(map cdr vars+)] ...)))]))

(define-pass merge-terms : Aexpts (ir) -> Aexpts ()
  (definitions)
  (EqvC : Term (t) -> * ()
    [(,c [,x ,n] ...) (list x n)])
  (Arith : Arith (a) -> Arith ()
    [(sum ,t^ ...)
     `(sum ,(for/list ([g (in-list (group-by EqvC t^))])
              (define-values [c x n]
                (for/fold ([c0 0] [x #f] [n #f])
                          ([t (in-list g)])
                  (nanopass-case (Aexpts Term) t
                    [(,c [,x ,n] ...) (values (+ c c0) x n)])))
              (with-output-language (Aexpts Term)
                `(,c [,x ,n] ...)))
           ...)]))

(define-pass remove-zeros : Aexpts (ir) -> Aexpts ()
  (definitions)
  (Keep? : Term (t) -> * ()
    [(,c [,x ,n] ...) (not (zero? c))])
  (Arith : Arith (a) -> Arith ()
    [(sum ,t^ ...) `(sum ,(filter Keep? t^) ...)]))

;--------
(define-language Aone-var
  (extends Aexpts)
  (Term (t)
    (- (c [x n] ...))
    (+ (c n)))
  (Solve (so)
    (- (solve a))
    (+ (solve x a))))

(define-pass one-variable : Aexpts (ir) -> Aone-var ()
  (definitions)
  (Term : Term (a var) -> Term ()
    [(,c [,x^ ,n^] ...)
     (define n
       (for/fold ([n 0])
                 ([x (in-list x^)]
                  [n (in-list n^)])
         (cond
           [(not (unbox var)) (set-box! var x) n]
           [(equal? (unbox var) x) n]
           [else (error "multi-variable equations unsupported")])))
     `(,c ,n)])
  (Arith : Arith (a var) -> Arith ()
    [(sum ,[t* var -> t*] ...) `(sum ,t* ...)])
  (Solve : Solve (so [var (box #f)]) -> Solve ()
    [(solve ,[a* var -> a*])
     (cond
       [(not (unbox var)) (error "no variables found to solve")]
       [else `(solve ,(unbox var) ,a*)])]))

;--------
(define-language Apolynomial
  (extends Aone-var)
  (Term (t) (- (c n)))
  (Arith (a) (- (sum t ...)))
  (Solve (so)
    (- (solve x a))
    (+ (solve-poly x [c ...]))))

(define-pass polynomial : Aone-var (ir) -> Apolynomial ()
  (definitions)
  (Coef : Term (so) -> * ()
    [(,c ,n) (values n c)])
  (Solve : Solve (so) -> Solve ()
    [(solve ,x (sum ,t^ ...))
     (define coefs
       (for/hash ([t (in-list t^)])
         (Coef t)))
     (define degree
       (apply max (hash-keys coefs)))
     (define (coef i)
       (hash-ref coefs i 0))
     `(solve-poly ,x [,(map coef (range (add1 degree))) ...])]))

;--------
(define-language Asolution
  (extends Apolynomial)
  (Solve (so)
    (- (solve-poly x [c ...]))
    (+ (solutions [x c] ...))))

(define-pass solve : Apolynomial (ir) -> Asolution ()
  (Solve : Solve (so) -> Solve ()
    [(solve-poly ,x [,c0]) `(solutions)]
    [(solve-poly ,x [,c0 ,c1])
     ; a + bx = 0
     ; x = -a / b
     `(solutions [,x ,(/ (- c0) c1)])]
    [(solve-poly ,x [,c0 ,c1 ,c2])
     ; c + bx + ax^2 = 0
     ; x = (-b +- sqrt(q)) / 2a
     ;   q = b^2 - 4ac
     (let* ([q (- (* c1 c1) (* 4 c2 c0))]
            [solu (λ (sign)
                    (/ (- (* sign (sqrt q)) c1)
                       (+ c2 c2)))])
       (cond
         [(zero? q) `(solutions [,x ,(solu +1)])]
         [else      `(solutions [,x ,(solu +1)]
                                [,x ,(solu -1)])]))]
    [(solve-poly ,x [,c ...])
     (error "cannot solve polynomial of degree > 2")]))


;==============

(define cas
  (compose unparse-Asolution
           solve
           polynomial
           one-variable
           remove-zeros
           merge-terms
           simplify-terms
           flat-trees
           distribute
           remove-negate
           remove-subtract
           rhs-zero
           parse-sexp
           parse-math))


(module+ main

  (define (read+cas port)
    (define ln (read-line port))
    (cdr (cas (open-input-string ln))))

  (define (read/prompt prompt reader)
    (λ (port)
      (display prompt)
      (flush-output)
      (reader port)))

  (define (read/catch is-exn deflt reader)
    (λ (port)
      (with-handlers ([is-exn (λ (e)
                                (printf "; error: ~a\n" (exn-message e))
                                deflt)])
        (reader port))))

  (define reader
    (read/prompt "cas> "
                 (read/catch exn:fail? '()
                             read+cas)))

  (for* ([solus (in-port reader)]
         [solu (in-list solus)])
    (printf "~a = ~a\n"
            (first solu)
            (second solu)))
  )