#lang plai-typed

(define-type ExprC
  [numC (n : number)]
  [idC (s : symbol)]  
  [plusC (l : ExprC) (r : ExprC)]
  [multC (l : ExprC) (r : ExprC)]
  [appC (fun : symbol) (arg : ExprC)])

(define-type ExprS
  [numS (n : number)]
  [idS (s : symbol)]  
  [plusS (l : ExprS) (r : ExprS)]
  [uminusS (e : ExprS)]
  [bminusS (l : ExprS) (r : ExprS)]
  [multS (l : ExprS) (r : ExprS)]
  [appS (fun : symbol) (arg : ExprS)])

(define (desugar [as : ExprS]) : ExprC
  (type-case ExprS as
    [numS (n) (numC n)]
    [idS (s) (idC s)]
    [plusS (l r) (plusC (desugar l) (desugar r))]
    [multS (l r) (multC (desugar l) (desugar r))]
    [uminusS (e) (desugar (multS (numS -1) e))]
    [bminusS (l r) (plusC (desugar l) (multC (numC -1) (desugar r)))]
    [appS (f a) (appC f (desugar a))]))
    
(define (parseS [s : s-expression]) : ExprS
  (cond
    [(s-exp-number? s) (numS (s-exp->number s))]
    [(s-exp-symbol? s) (idS (s-exp->symbol s))]
    [(s-exp-list? s)
     (let ([sl (s-exp->list s)])
       (case (s-exp->symbol (first sl))
         [(+) (plusS (parseS (second sl)) (parseS (third sl)))]
         [(*) (multS (parseS (second sl)) (parseS (third sl)))]
         [(u-) (uminusS (parseS (second sl)))]
         [(-) (bminusS (parseS (second sl)) (parseS (third sl)))]
         [else (appS (s-exp->symbol (first sl)) (parseS (second sl)))]))]
    [else (error ‘parseS "invalid input")]))
    
(define-type FunDefC [fdC (name : symbol) (arg : symbol) (body : ExprC)])

(define (get-fundef [n : symbol] [fds : (listof FunDefC)]) : FunDefC
  (cond
    [(empty? fds) (error ‘get-fundef "reference to undefined function")]
    [(cons? fds) (cond
                   [(equal? n (fdC-name (first fds))) (first fds)]
                   [else (get-fundef n (rest fds))])]))

(define (parse-fundef [s : s-expression]) : FunDefC
  (cond
    [(s-exp-list? s)
     (let ([sl (s-exp->list s)])
       (case (s-exp->symbol (first sl))
         [(define) (fdC (s-exp->symbol (first (s-exp->list (second sl))))
                        (s-exp->symbol (second (s-exp->list (second sl))))
                        (desugar (parseS (third sl))))]
         [else (error ‘parse-fundef "invalid list")]))]
     [else (error ‘parse-fundef "invalid input")]))
 
(define (subst [what : ExprC] [for : symbol] [in : ExprC]) : ExprC
  (type-case ExprC in
    [numC (n) in]
    [idC (s) (cond
               [(symbol=? s for) what]
               [else in])]
    [appC (f a) (appC f (subst what for a))]
    [plusC (l r) (plusC (subst what for l) (subst what for r))]
    [multC (l r) (multC (subst what for l) (subst what for r))]))
(define (interp [e : ExprC] [fds : (listof FunDefC)]) : number
  (type-case ExprC e
    [numC (n) n]
    [idC (_) (error ‘interpC "shouldn‘t get here")]
    [appC (f a) (local ([define fd (get-fundef f fds)])
                  (interp (subst a (fdC-arg fd) (fdC-body fd)) fds))]
    [plusC (l r) (+ (interp l fds) (interp r fds))]
    [multC (l r) (* (interp l fds) (interp r fds))]))

(define (parse [s : s-expression]) : ExprC
  (desugar (parseS s)))

(define (main [s : s-expression] [fs : (listof s-expression)]) : number
  (interp (parse s) (map parse-fundef fs)))


(define l (list `(define (f x) (+ x x)) `(define (g x) (* x 3))))

(define s1 ‘(+ (f 2) (* 2 3)))
(main s1 l)

(define s2 ‘(+ (u- (f 2)) (- (g 5) (* 2 (+ 2 4)))))
(main s2 l)