#lang plait
(define-type Exp
  [num (n : Number)]
  [plus (left : Exp) (right : Exp)]
  [times (left : Exp) (right : Exp)]
  [cnd (test : Exp) (then : Exp) (else : Exp)])

(define (calc e)
  (type-case Exp e
    [(num n) n]
    [(plus l r) (+ (calc l) (calc r))]
    [(times l r) (* (calc l) (calc r))]
    [(cnd c t e) (if (zero? (calc c))
                     (calc t)
                     (calc e))]))

(define (parse s)
  (local
      [(define (sx n) (list-ref (s-exp->list s) n))
       (define (px n) (parse (sx n)))
       (define (? pat) (s-exp-match? pat s))] 
    (cond
      [(? `NUMBER) (num (s-exp->number s))]
      [(? `(+ ANY ANY)) (plus (px 1) (px 2))]
      [(? `(* ANY ANY)) (times (px 1) (px 2))]
      [(? `(if0 ANY ANY ANY))                 ;; NEW
       (cnd (px 1) (px 2) (px 3))]
      [else (error 'parse (to-string s))])))

(run : (S-Exp -> Number))
(define (run s)
  (calc (parse s)))

(test (run `1) 1)
(test (run `2.3) 2.3)
(test (run `{+ 1 2}) 3)
(test (run `{+ {+ 1 2} 3})
      6)
(test (run `{+ 1 {+ {+ 2 3} 4}})
      10)

(test (run `{+ 1 {+ 2 3}}) 6)
(test (run `{* {+ 2 3} {+ 5 6}}) 55)
(test (run `{if0 0 (* 1 2) 1}) 2)