#lang pl

; Since postfix expressions are just a different way of writing normal expressions, my instinct was to have them be a kind of syntactic sugar rather than their own AST type.
; I realize now this wasn't the intended way of implementing post, though it's equivalent for the user.

#|
This is the grammar of the language

PUWAE :== <num>
  | <id>
  | {+ PUWAE PUWAE}
  | {- PUWAE PUWAE}
  | {* PUWAE PUWAE}
  | {/ PUWAE PUWAE}
  | {post POSTT...}
  | {with {<id> PUWAE} PUWAE}

POSTT :== PUWAE
  | '+'
  | '-'
  | '*'
  | '/'

|#

(define-type PUWAE
  [Num Number]
  [Id Symbol]
  [Sum PUWAE PUWAE]
  [Sub PUWAE PUWAE]
  [Mul PUWAE PUWAE]
  [Div PUWAE PUWAE]
  [Post (Listof POSTT)]
  [With Symbol PUWAE PUWAE])

(define-type POSTT = (U '+ '- '* '/ PUWAE))

(: parse : (String -> PUWAE))
;; parse the string into our language
(define (parse str)
  (sexpr->puwae (string->sexpr str)))

(: sexpr->puwae : (Sexpr -> PUWAE))
;; turn the sexpr into an puwae
(define (sexpr->puwae sexpr)
  (match sexpr
    [(number: n) (Num n)]
    [(symbol: id) (Id id)]
    [(list '+ l r) (Sum (sexpr->puwae l) (sexpr->puwae r))]
    [(list '- l r) (Sub (sexpr->puwae l) (sexpr->puwae r))]
    [(list '* l r) (Mul (sexpr->puwae l) (sexpr->puwae r))]
    [(list '/ l r) (Div (sexpr->puwae l) (sexpr->puwae r))]
    [(list 'post posts ...) (Post (sexpr-list->postts posts))]
    [(list 'with (list (symbol: id) val) bound-body)
     (With id (sexpr->puwae val) (sexpr->puwae bound-body))]
    [_ (error 'sexpr->puwae "malformed expression in ~s" sexpr)]))

(: sexpr->postt : (Sexpr -> POSTT))
;; Turn a post into POSTT.
(define (sexpr->postt sexpr)
  (match sexpr
    ['+ '+]
    ['- '-]
    ['* '*]
    ['/ '/]
    [_ (sexpr->puwae sexpr)]))

(: sexpr-list->postts : ((Listof Sexpr) -> (Listof POSTT)))
(define (sexpr-list->postts sexprs)
  (match sexprs
    ['() '()]
    [(cons token rest)
     (cons (sexpr->postt token) (sexpr-list->postts rest))]))

(: eval : (PUWAE -> Number))
;; evaluate to a final number
(define (eval puwae)
  (cases puwae
    [(Num num) num]
    [(Id id) (error 'eval "unbound identifier ~s" id)]
    [(Sum l r) (+ (eval l) (eval r))]
    [(Sub l r) (- (eval l) (eval r))]
    [(Mul l r) (* (eval l) (eval r))]
    [(Div l r) (/ (eval l) (eval r))]
    [(Post postts) (eval-post postts '())]
    [(With id pre-val bound-body) (eval (subst id (Num (eval pre-val)) bound-body))]))

(: post-stack-apply : (Symbol (Number Number -> Number) (Listof Number) -> (Listof Number)))
(define (post-stack-apply operator operation stack)
  (if (>= (length stack) 2)
      (cons (operation (second stack) (first stack)) (rest (rest stack)))
      (error 'eval-post "Too few arguments for ~s: ~s" operator stack)))

(: eval-post : ((Listof POSTT) (Listof Number) -> Number))
(define (eval-post postts stack)
  (match postts
    ['()
     (if (= 1 (length stack))
         (first stack)
         (error 'eval-post "Finished with messy stack: ~s" stack))]
    [(cons token rest)
     (if (symbol? token)
         (match token
           ['+ (eval-post rest (post-stack-apply '+ + stack))]
           ['- (eval-post rest (post-stack-apply '- - stack))]
           ['* (eval-post rest (post-stack-apply '* * stack))]
           ['/ (eval-post rest (post-stack-apply '/ / stack))]
           [_ (error 'eval-post "Bad post operator ~s" token)])
         (eval-post rest (cons (eval token) stack)))]))

(: subst : (Symbol PUWAE PUWAE -> PUWAE))
;; subst all instances of id with val in body
(define (subst id val body)
  (cases body
    [(Num num) body]
    [(Id other-id) (if (symbol=? id other-id) val body)]
    [(Sum l r) (Sum (subst id val l) (subst id val r))]
    [(Sub l r) (Sub (subst id val l) (subst id val r))]
    [(Mul l r) (Mul (subst id val l) (subst id val r))]
    [(Div l r) (Div (subst id val l) (subst id val r))]
    [(Post postts) (Post (subst-postts id val postts))]
    [(With other-id pre-val bound-body)
     (With other-id
           (subst id val pre-val)
           (if (symbol=? id other-id)
               bound-body
               (subst id val bound-body)))]))

(: subst-postts : (Symbol PUWAE (Listof POSTT) -> (Listof POSTT)))
(define (subst-postts id val postts)
  (match postts
    ['() '()]
    [(cons token rest)
     (if (symbol? token)
         (cons token (subst-postts id val rest))
         (cons (subst id val token) (subst-postts id val rest)))]))

(: run : String -> Number)
;; run the program
(define (run str)
  (eval (parse str)))

(test (run "{post 2}") => 2)
(test (run "{post 1 0 +}") => 1)
(test (run "{post 1 1 +}") => 2)
(test (run "{* {post 1 3 +} {post 5 6 +}}") => 44)
(test (run "{post 1 2 + 4 *}") => 12)
(test (run "{post 1 3 + {+ 5 6} *}") => 44)
(test (run "{post {post 1 3 +} {post 5 6 +} *}") => 44)
(test (run "{* {+ {post 1} {post 3}} {+ {post 5} {post 6}}}") => 44)
(test (run "{with {x {post 1 3 +}} {post 5 6 + x *}}") => 44)
(test (run "{post 1 3 + 5 6 + *}") => 44) ; Not 48 :P (5 + 6 ≠ 12).

(test (run "{post 3 +}") =error> "eval-post: Too few arguments for +: (3)")
(test (run "{post 3 3 3 +}") =error> "eval-post: Finished with messy stack: (6 3)")
(test (run "{post 1 2 + 2}") =error> "eval-post: Finished with messy stack: (2 3)")

(test (run "4") => 4)
(test (run "{+ 4 {* 2 3}}") => 10)
(test (run "{+ {* 7 7} {* 7 7}}") => 98)
(test (run "5") => 5)
(test (run "{+ 5 5}") => 10)

(test (parse "{with {x {* 7 7}} {+ x x}}")
      =>
      (With 'x (Mul (Num 7) (Num 7)) (Sum (Id 'x) (Id 'x))))

(test (run "x") =error> "unbound identifier x")

(test (run "{with {x 4} y}") =error> "unbound identifier y")
(test (run "{with {x x} 3}") =error> "unbound identifier x")
(test (run "{with {x 4} 3}") => 3)
(test (run "{with {x 4} x}") => 4)
(test (run "{with {x 4} {+ x 2}}") => 6)
(test (run "{with {x 4} {- x 2}}") => 2)
(test (run "{with {x 4} {with {y 3} {+ x y}}}") => 7)
(test (run "{with {x {* 7 7}} {+ x x}}") => 98)
(test (run "{with {x 4} {with {y {* x x}} {+ x y}}}") => 20)
(test (run "{with {x 4} {with {x 5} x}}") => 5) ;; tests shadowing
(test (run "{with {x 4} {with {x {* x x}} x}}") => 16)
;; we can use a variable in its own shadowing^
(test (run "{with {x {with {x 4} x}} {with {x {* x x}} x}}") => 16)
(test (run "{with {x {* 7 7}} {+ x x}}") => 98)
(test (run "{with {x {- 7 3}} {+ x x}}") => 8)
(test (run "{with {x 5} {+ x x}}") => 10)
(test (run "{with {x 5} {/ x x}}") => 1)
(test (run "{with {x {+ 5 5}} {+ x x}}") => 20)
(test (run "{with {x 5} {with {y {+ x 3}} {+ y y}}}") => 16)
(test (run "{with {x {+ 5 5}} {with {y {+ x 3}} {+ y y}}}") => 26)
(test (run "{with {x 5} {+ x {with {x 3} 10}}}") => 15)
(test (run "{with {x 5} {+ x {with {x 3} x}}}") => 8)
(test (run "{with {x 5} {+ x {with {y 3} x}}}") => 10)
(test (run "{with {x 5} {+ x {+ {with {x 3} x} x}}}") => 13)
(test (run "{with {x 5} {with {y x} y}}") => 5)
(test (run "{with {x {/ 5 2}} {+ x 1}}") => 7/2)
(test (run "{with {x 5} {with {x x} x}}") => 5)
(test (run "{with {x 1} y}") =error> "eval: unbound identifier y")