#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 ...) (parsepost 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)]))

(: parsepost : (Sexpr (Listof PUWAE) -> PUWAE))
; Parse a post expression into a PUWAE.
(define (parsepost sexpr glob)
  (println (format "parsing '~s' with glob '~s'" sexpr glob))
  (match sexpr
    ['() (if (= 1 (length glob)) (first glob) (error 'parsepost "finished with messy glob: ~s" glob))] ; There shouldn't be anything else left in the stack.
    [(cons '+ rest) (parsepost rest (parsepost-sum glob))]
    [(cons '- rest) (parsepost rest (parsepost-sub glob))]
    [(cons '* rest) (parsepost rest (parsepost-mul glob))]
    [(cons '/ rest) (parsepost rest (parsepost-div glob))]
    [(cons (symbol: id) rest) (parsepost rest (cons (Id id) glob))] ; Capture non-post symbols (variables).
    [(cons exp rest) (parsepost rest (cons (sexpr->puwae exp) glob))])) ; Assume anything else must be a PUWAE.

(: glob-swap : Symbol (Listof PUWAE) (PUWAE PUWAE -> PUWAE) -> (Listof PUWAE))
; Validates the glob length before swapping with result of globberator.
(define (glob-swap operation glob globberator)
  (if (>= (length glob) 2)
      (cons (globberator (second glob) (first glob)) (rest (rest glob))) ; Replace the two arguments, while preserving whatever else was on the stack.
      (error 'glob-swap "too few arguments globbed for ~s: ~s" operation glob)))

(: parsepost-sum : ((Listof PUWAE) -> (Listof PUWAE)))
; Parse a post sum.
(define (parsepost-sum glob)
  (println (format "summing '~s'" glob))
  (glob-swap 'sum glob Sum))

(: parsepost-mul : ((Listof PUWAE) -> (Listof PUWAE)))
; Parse a post product.
(define (parsepost-mul glob)
  (println (format "mulling '~s'" glob))
  (glob-swap 'mul glob Mul))

(: parsepost-sub : ((Listof PUWAE) -> (Listof PUWAE)))
; Parse a post subtration.
(define (parsepost-sub glob)
  (println (format "subbing '~s'" glob))
  (glob-swap 'sub glob Sub))

(: parsepost-div : ((Listof PUWAE) -> (Listof PUWAE)))
; Parse a post quotient.
(define (parsepost-div glob)
  (println (format "divving '~s'" glob))
  (glob-swap 'div glob Div))

(: 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))]
    [(With id pre-val bound-body) (eval (subst id (Num (eval pre-val)) bound-body))]))

(: 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))]
    [(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)))]))

(: 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> "glob-swap: too few arguments globbed for sum: ((Num 3))")
(test (run "{post 3 3 3 +}") =error> "parsepost: finished with messy glob: ((Sum (Num 3) (Num 3)) (Num 3))")
(test (run "{post 1 2 + 2}") =error> "parsepost: finished with messy glob: ((Num 2) (Sum (Num 1) (Num 2)))")

(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")