;;; ---<<<TOY>>>----------------------------------------------------
#lang pl
;;; ----------------------------------------------------------------
;;; Syntax

#| The AST:
<TOY> ::= <num>
       | <id>
       | { bind {{ <id> <TOY> } ... } <TOY> ... }
       | { bindrec {{ <id> <TOY> } ... } <TOY> ... }
       | { fun { <id> ... } <TOY> ... }
       | { rfun { <id> ... } <TOY> ... }
       | { <TOY> <TOY> ... }
       | { if <TOY> <TOY> <TOY> }
       | { set! <id> <TOY> }
|#

;; A matching abstract syntax tree datatype:
(define-type TOY
             [Num Number]
             [Id Symbol]
             [Bind (Listof Symbol) (Listof TOY) (Listof TOY)]
             [Bindrec (Listof Symbol) (Listof TOY) (Listof TOY)]
             [Fun (Listof Symbol) (Listof TOY)]
             [Rfun (Listof Symbol) (Listof TOY)]
             [Call TOY (Listof TOY)]
             [If TOY TOY TOY]
             [Set Symbol TOY])

; Things you can't name things.
(define BAD_WORDS (list 'bind 'bindrec 'fun 'rfun 'if 'set! '+ '- '* '/ '= '< '> 'true 'false 'void))

(: unique-list? : (Listof Any) -> Boolean)
;; Tests whether a list is unique, guards Bind and Fun values.
(define (unique-list? xs)
  (or (null? xs)
      (and (not (member (first xs) (rest xs)))
           (unique-list? (rest xs)))))

(: no-bad-words? : (Listof Symbol) -> Boolean)
;; Tests whether a list avoids all bad words.
(define (no-bad-words? names)
  (or (null? names)
      (and (not (member (first names) BAD_WORDS))
           (no-bad-words? (rest names)))))

(: parse-sexpr : Sexpr -> TOY)
;; parses s-expressions into TOYs
(define (parse-sexpr sexpr)
  (match sexpr
    ; Primitives.
    [(number: n) (Num n)]
    [(symbol: name) (Id name)]

    ; Bindrec.
    [(cons (and binder (or 'bind 'bindrec)) _) ; `binder` bound to 'bind or 'bindrec.
     ; `_` feels cooler than `more`.
     (match sexpr
       [(list b (list (list (symbol: names) (sexpr: nameds)) ...) body ...)
        (if (not (null? body))
            (if (unique-list? names)
                (if (no-bad-words? names)
                    ((if (symbol=? b 'bind) Bind Bindrec)
                     names
                     (map parse-sexpr nameds)
                     (map parse-sexpr body))
                    (error 'parse-sexpr "bad word found in `~s' : ~s" binder names))
                (error 'parse-sexpr "duplicate `~s' names: ~s" binder names))
            (error 'parse-sexpr "bind expression missing body: ~s" sexpr))]
       [else (error 'parse-sexpr "bad `~s' syntax in ~s" binder sexpr)])]

    ;
    ;
    ;
    ; RFun definition.
    [(cons (and funer (or 'fun 'rfun)) _)
     (match sexpr
       [(list f (list (symbol: parameters) ...) body ...)
        (if (not (null? body))
            (if (unique-list? parameters)
                (if (no-bad-words? parameters)
                    ((if (symbol=? f 'fun) Fun Rfun) parameters (map parse-sexpr body))
                    (error 'parse-sexpr "bad words found in function parameter names in ~s" sexpr))
                (error 'parse-sexpr "bad parameters defined for function in ~s" sexpr))
            (error 'parse-sexpr "disemobodied function in ~s" sexpr))]
       [else (error 'parse-sexpr "bad function syntax in ~s" sexpr)])]

    ; Conditional.
    [(cons 'if _)
     (match sexpr
       [(list 'if cond then else)
        (If (parse-sexpr cond)
            (parse-sexpr then)
            (parse-sexpr else))]
       [else (error 'parse-sexpr "bad `if' syntax in ~s" sexpr)])]

    ; Set!.
    [(cons 'set! _)
     (match sexpr
       [(list 'set! (symbol: thing) as)
        (Set thing (parse-sexpr as))])]

    ; Function application
    [(list fun args ...) ; other lists are applications
     (Call (parse-sexpr fun)
           (map parse-sexpr args))]

    ; Bad.
    [else (error 'parse-sexpr "bad syntax in ~s" sexpr)]))

(: parse : String -> TOY)
;; Parses a string containing an TOY expression to a TOY AST.
(define (parse str)
  (parse-sexpr (string->sexpr str)))

(define-type ENV
             [EmptyEnv]
             [FrameEnv FRAME ENV])

;; a frame is an association list of names and values.
(define-type FRAME = (Listof (List Symbol (Boxof VAL))))

(define-type VAL
             [RktV Any]
             ; Parameters, compiled body, whether it's an Rfun, environment.
             [FunV (Listof Symbol) (ENV -> VAL) Boolean ENV]
             [PrimV ((Listof VAL) -> VAL)]
             [VoidV])

(: extend : (Listof Symbol) (Listof VAL) ENV -> ENV)
(define (extend names vals env)
  (raw-extend names (map (inst box VAL) vals) env))

(: raw-extend : (Listof Symbol) (Listof (Boxof VAL)) ENV -> ENV)
;; extends an environment with a new frame.
(define (raw-extend names bvals env)
  (if (= (length names) (length bvals))
      (FrameEnv (map (lambda ([name : Symbol] [bval : (Boxof VAL)])
                       (list name bval))
                     names bvals)
                env)
      (error 'extend "arity mismatch for names: ~s" names)))

(: extend-rec : (Listof Symbol) (Listof (ENV -> VAL)) ENV -> ENV)
(define (extend-rec names compiled-vals env)
  (define newenv (extend names (map (lambda (_) (VoidV)) compiled-vals) env))
  (for-each (lambda ([name : Symbol] [compiled-val : (ENV -> VAL)])
              (set-box! (lookup name newenv) (compiled-val newenv)))
            names compiled-vals)
  newenv)

(: lookup : Symbol ENV -> (Boxof VAL))
;; lookup a symbol in an environment, frame by frame,
;; return its value or throw an error if it isn't bound
(define (lookup name env)
  (cases env
         [(EmptyEnv) (error 'lookup "no binding for ~s" name)]
         [(FrameEnv frame rest)
          (let ([cell (assq name frame)])
            (if cell
                (second cell)
                (lookup name rest)))]))

(: unwrap-rktv : VAL -> Any)
;; helper for `racket-func->prim-val': unwrap a RktV wrapper in
;; preparation to be sent to the primitive function
(define (unwrap-rktv x)
  (cases x
         [(RktV v) v]
         [else (error 'racket-func "bad input: ~s" x)]))

(: racket-func->prim-val : Function -> (Boxof VAL))
;; converts a racket function to a primitive evaluator function
;; which is a PrimV holding a ((Listof VAL) -> VAL) function.
;; (the resulting function will use the list function as is,
;; and it is the list function's responsibility to throw an error
;; if it's given a bad number of arguments or bad input types.)
(define (racket-func->prim-val racket-func)
  (define list-func (make-untyped-list-function racket-func))
  (box (PrimV (lambda (args)
                (RktV (list-func (map unwrap-rktv args)))))))

;; The global environment has a few primitives:
(: global-environment : ENV)
(define global-environment
  (FrameEnv (list (list '+ (racket-func->prim-val +))
                  (list '- (racket-func->prim-val -))
                  (list '* (racket-func->prim-val *))
                  (list '/ (racket-func->prim-val /))
                  (list '< (racket-func->prim-val <))
                  (list '> (racket-func->prim-val >))
                  (list '= (racket-func->prim-val =))
                  ;; values
                  (list 'true (box (RktV #t)))
                  (list 'false (box (RktV #f)))
                  (list 'void (box (VoidV))))
            (EmptyEnv)))

(: compiler-enabled? : (Boxof Boolean))
(define compiler-enabled? (box #f))

(: ensure-compiler-enabled! : Symbol -> Void)
(define (ensure-compiler-enabled! who)
  (if (unbox compiler-enabled?)
      (void)
      (error who "compiler is not enabled")))

(: run-compiled-body : (Listof (ENV -> VAL)) ENV -> VAL)
; Compile a body of expressions.
(define (run-compiled-body compiled env)
  (if (null? compiled)
      (error 'compile "empty body")
      (if (null? (rest compiled))
          ((first compiled) env)
          (let ([_ ((first compiled) env)]) ; We've side effects in this language.
            (run-compiled-body (rest compiled) env)))))

(: compile-body : (Listof TOY) -> (ENV -> VAL))
;; Compile a sequence of body expressions into runtime sequencing code.
(define (compile-body body)
  (ensure-compiler-enabled! 'compile-body)
  (define compiled-body (map compile body))
  (lambda ([env : ENV])
    (run-compiled-body compiled-body env)))

; Convenient helper for running compiled code.
(: caller : ENV -> (ENV -> VAL) -> VAL)
(define (caller env)
  (lambda ([compiled : (ENV -> VAL)]) (compiled env)))

(: compile : TOY -> (ENV -> VAL))
;; compileuates TOY expressions.
(define (compile expr)
  (ensure-compiler-enabled! 'compile)
  (cases expr
         [(Num n) (λ ([_ : ENV]) (RktV n))]
         [(Id name) (λ ([env : ENV]) (unbox (lookup name env)))]
         [(Bind names exprs body)
          (define compiled-exprs (map compile exprs))
          (define compiled-body (compile-body body))
          (λ ([env : ENV])
            (compiled-body
              (extend names
                      (map (caller env) compiled-exprs)
                      env)))]
         [(Bindrec names exprs body)
          (define compiled-exprs (map compile exprs))
          (define compiled-body (compile-body body))
          (λ ([env : ENV])
            (compiled-body (extend-rec names compiled-exprs env)))]
         [(Fun names body)
          (define compiled-body (compile-body body))
          (λ ([env : ENV]) (FunV names compiled-body #f env))]
         [(Rfun names body)
          (define compiled-body (compile-body body))
          (λ ([env : ENV]) (FunV names compiled-body #t env))]
         [(Call fun-expr arg-exprs)
          (define compiled-fun (compile fun-expr))
          (define compiled-args (map compile arg-exprs))
          (define compiled-get-boxes (compile-get-boxes arg-exprs))
          (λ ([env : ENV])
            (let ([fval (compiled-fun env)]
                  [arg-vals-thunk
                   (lambda ()
                     (map (caller env) compiled-args))]) ; THunk.
              (cases fval
                     [(PrimV proc) (proc (arg-vals-thunk))]
                     [(FunV names body ref? fun-env)
                      (if ref?
                          (compile-rfun names compiled-get-boxes body env fun-env)
                          (body (extend names (arg-vals-thunk) fun-env)))]
                     [else (error 'compile "function call with a non-function: ~s"
                                  fval)])))]
         [(If cond-expr then-expr else-expr)
          (define compiled-cond (compile cond-expr))
          (define compiled-then (compile then-expr))
          (define compiled-else (compile else-expr))
          (λ ([env : ENV])
            (if (cases (compiled-cond env)
                       [(RktV v) v] ; Racket value => use as boolean
                       [else #t]) ; other values are always true
                (compiled-then env)
                (compiled-else env)))]
         [(Set thing as)
          (define compiled-as (compile as))
          (λ ([env : ENV])
            (set-box! (lookup thing env) (compiled-as env))
            (VoidV))]))

(: compile-rfun : (Listof Symbol) (ENV -> (Listof (Boxof VAL))) (ENV -> VAL) ENV ENV -> VAL)
; Compile an RFun.
(define (compile-rfun params get-boxes body callerenv calleeenv)
  (body (raw-extend params (get-boxes callerenv) calleeenv)))

(: compile-get-boxes : (Listof TOY) -> (ENV -> (Listof (Boxof VAL))))
; Utility for applying rfun.
(define (compile-get-boxes exprs)
  (ensure-compiler-enabled! 'compile-get-boxes)
  (: compile-getter : TOY -> (ENV -> (Boxof VAL)))
  (define (compile-getter expr)
    (cases expr
           [(Id name) (lambda ([env : ENV]) (lookup name env))]
           [else
            (lambda ([_ : ENV]) (error 'compile-get-boxes "non-identifier"))]))
  (let ([getters (map compile-getter exprs)])
    (lambda ([env : ENV])
      (map (lambda ([getter : (ENV -> (Boxof VAL))])
             (getter env))
           getters))))

(: run : String -> Any)
;; evaluate a TOY program contained in a string
(define (run str)
  (define parsed (parse str))
  (set-box! compiler-enabled? #t)
  (define compiled (compile parsed))
  (set-box! compiler-enabled? #f)
  (let ([result (compiled global-environment)])
    (cases result
           [(RktV v) v]
           [else (error 'run "evaluation returned a bad value: ~s"
                        result)])))

(test (run "{{rfun {x} x} {/ 4 0}}") =error> "non-identifier")
(test (run "{5 {/ 6 0}}") =error> "non-function")
(test (run "
{bind
    {
        {swap!
            {rfun {x y}
                {bind {{tmp x}}
                    {set! x y}
                    {set! y tmp}
                }
            }
        }

        {a 1}

        {b 2}
    }

    {swap! a b}

    {+ a {* 10 b}}
}") => 12)

(test (run "{bind {{make-counter
{fun {}
{bind {{c 0}}
{fun {}
{set! c {+ 1 c}}
c}}}}}
{bind {{c1 {make-counter}}
{c2 {make-counter}}}
{* {c1} {c1} {c2} {c1}}}}")

      => 6)
(test (run "{bindrec {{foo {fun {}

{set! foo {fun {} 2}}
1}}}
{+ {foo} {* 10 {foo}}}}")
      => 21)

(test (run "{bindrec {{fact {fun {n}
{if {= 0 n}
1
{* n {fact {- n 1}}}}}}}

{fact 5}}")
      => 120)

(test (run "
{bind {{x 10}}
  {bind {{y {set! x 1}}}
    x
  }
}
") => 1)
(test (run "{bind {{x 2}} {set! x 1}}") =error> "run: evaluation returned a bad value: (VoidV)")

(test (run "{if {= 1 2} 0 1}") => 1)
(test (run "{if true 1 2}") => 1)
(test (run "{{fun {x} {+ x 1}} 4}")
      => 5)
(test (run "{bind {{add3 {fun {x} {+ x 3}}}} {add3 1}}")
      => 4)
(test (run "{bind {{add3 {fun {x} {+ x 3}}}
                       {add1 {fun {x} {+ x 1}}}}
                  {bind {{x 3}} {add1 {add3 x}}}}")
      => 7)
(test (run "{bind {{identity {fun {x} x}}
                       {foo {fun {x} {+ x 1}}}}
                  {{identity foo} 123}}")
      => 124)
(test (run "{bind {{x 3}}
                  {bind {{f {fun {y} {+ x y}}}}
                    {bind {{x 5}}
                      {f 4}}}}")
      => 7)
(test (run "{{{fun {x} {x 1}}
                  {fun {x} {fun {y} {+ x y}}}}
                 123}")
      => 124)

;; More tests for complete coverage
(test (run "{bind x 5 x}") =error> "bad `bind' syntax")
(test (run "{fun x x}") =error> "parse-sexpr: bad function syntax in (fun x x)")
(test (run "{if x}") =error> "bad `if' syntax")
(test (run "{}") =error> "bad syntax")
(test (run "{bind {{x 5} {x 5}} x}") =error> "parse-sexpr: duplicate `bind' names: (x x)")
(test (run "{fun {x x} x}") =error> "parse-sexpr: bad parameters defined for function in (fun (x x) x)")
(test (run "{+ x 1}") =error> "no binding for")
(test (run "{+ 1 {fun {x} x}}") =error> "bad input")
(test (run "{+ 1 {fun {x} x}}") =error> "bad input")
(test (run "{1 2}") =error> "with a non-function")
(test (run "{{fun {x} x}}") =error> "arity mismatch")
(test (run "{if {< 4 5} 6 7}") => 6)
(test (run "{if {< 5 4} 6 7}") => 7)
(test (run "{if + 6 7}") => 6)
(test (run "{fun {x} x}") =error> "returned a bad value")

(test (compile (Num 1)) =error> "compiler is not enabled")
(test (compile-body (list (Num 1))) =error> "compiler is not enabled")
(test (compile-get-boxes (list (Id 'x))) =error> "compiler is not enabled")

(test (run "{bind {{if 5}} if}") =error> "parse-sexpr: bad word found in `bind' : (if)")
(test (run "{fun {if} if}") =error> "parse-sexpr: bad words found in function parameter names in (fun (if) if)")