#lang pl

#|
The grammar:
  <TIDBIT> ::= <num>
            | { + <TIDBIT> <TIDBIT> }
            | { - <TIDBIT> <TIDBIT> }
            | { * <TIDBIT> <TIDBIT> }
            | { / <TIDBIT> <TIDBIT> }
            | { with { <id> <TIDBIT> } <TIDBIT> }
            | <id>
            | { fun <id> <TIDBIT> }
            | { call <TIDBIT> <TIDBIT> }
            | { <TIDBIT> of <TIDBIT> }
            | { fun { <id> ... } <TIDBIT> }
            | { call <TIDBIT> {<TIDBIT> ...} }
            | { <TIDBIT> of {<TIDBIT> ...} }
            | { bind {{<id> <TIDBIT>} ...} <TIDBIT> }
            | { bind* {{<id> <TIDBIT>} ...} <TIDBIT> }
|#

(define-type TIDBIT
             [Num Number]
             [Add TIDBIT TIDBIT]
             [Sub TIDBIT TIDBIT]
             [Mul TIDBIT TIDBIT]
             [Div TIDBIT TIDBIT]
             [Id Symbol]
             [With Symbol TIDBIT TIDBIT]
             [Fun (Listof Symbol) TIDBIT]
             [Call TIDBIT (Listof TIDBIT)]
             [Bind (Listof Symbol) (Listof TIDBIT) TIDBIT]
             [Bind* (Listof Symbol) (Listof TIDBIT) TIDBIT]
             [If0 TIDBIT TIDBIT TIDBIT])

(define-type Idx = Integer)

(define-type CORE
             [CNum Number]
             [CAdd CORE CORE]
             [CSub CORE CORE]
             [CMul CORE CORE]
             [CDiv CORE CORE]
             [CIdx Idx]
             [CFun CORE]
             [CCall CORE CORE]
             [CIf0 CORE CORE CORE]) ; No way to implement equality check with only arithmetic and functions.

(define-type BINDING-DEPTH = (Symbol -> Idx))

(: empty-depth : BINDING-DEPTH)
(define (empty-depth s) (error 'empty-depth "No binding for ~s." s))

(: binding-encountered : BINDING-DEPTH Symbol -> BINDING-DEPTH)
(define (binding-encountered bd id)
  (lambda ([s : Symbol]) (if (symbol=? s id) 1 (add1 (bd s)))))

(: currycall : CORE (Listof TIDBIT) BINDING-DEPTH -> CORE)
(define (currycall body args bd)
  (match args
      ['() body]
      [(cons f r) (currycall (CCall body (preprocess f bd)) r bd)]))

(: curryfn : (Listof Symbol) TIDBIT BINDING-DEPTH -> CORE)
(define (curryfn params body bd)
  (match params
    ['() (preprocess body bd)]
    [else (CFun (curryfn (rest params) body (binding-encountered bd (first params))))]))

(: binds : (Listof Symbol) (Listof TIDBIT) TIDBIT -> TIDBIT)
(define (binds names vals body)
  (Call (Fun names body) vals))

; This could also be done with functions like above, just nested instead of n-ary.
(: binds* : (Listof Symbol) (Listof TIDBIT) TIDBIT -> TIDBIT)
(define (binds* names vals body)
  (match names
    ['() body]
    [else (With (first names) (first vals) (binds* (rest names) (rest vals) body))]))

(: preprocess : TIDBIT BINDING-DEPTH -> CORE)
(define (preprocess tb bd)
  (cases tb
         [(Num n) (CNum n)]
         [(Add l r) (CAdd (preprocess l bd) (preprocess r bd))]
         [(Sub l r) (CSub (preprocess l bd) (preprocess r bd))]
         [(Mul l r) (CMul (preprocess l bd) (preprocess r bd))]
         [(Div l r) (CDiv (preprocess l bd) (preprocess r bd))]
         [(Id s) (CIdx (bd s))]
         [(With name value body) (CCall (CFun (preprocess body (binding-encountered bd name)))
                                        (preprocess value bd))]
         [(Fun params body) (curryfn params body bd)]
         [(Call body args) (currycall (preprocess body bd) args bd)]

         [(Bind names vals body) (preprocess (binds names vals body) bd)]
         [(Bind* names vals body) (preprocess (binds* names vals body) bd)]
         [(If0 n body alt) (CIf0 (preprocess n bd) (preprocess body bd) (preprocess alt bd))]))

(: parse-bind : (Listof Sexpr) (Listof Sexpr) TIDBIT -> TIDBIT)
(define (parse-bind names vals body)
  (Bind (map (lambda ([id : Sexpr])
               (match id
                 [(list (symbol: name)) name]
                 [else (error 'parse-bind "~s not a good id." id)])) names)
        (map parse-sexpr vals)
        body))

(: parse-sexpr : Sexpr -> TIDBIT)
; Parses Sexprs into TIDBITs.
(define (parse-sexpr sexpr)
  (match sexpr
    [(number: n) (Num n)]
    [(symbol: name) (Id name)]

    ; With.
    [(cons 'with more)
     (match sexpr
       [(list 'with (list (symbol: name) named) body)
        (With name (parse-sexpr named) (parse-sexpr body))]
       [else (error 'parse-sexpr "bad `with' syntax in ~s" sexpr)])]

    ; Function declaration.
    [(cons 'fun more)
     (match sexpr
       [(list 'fun (symbol: param) body) (Fun (list param) (parse-sexpr body))]
       [(list 'fun (list (symbol: params) ...) body) (Fun params (parse-sexpr body))]
       [(list 'fun body) (Fun '() (parse-sexpr body))]
       [else (error 'parse-sexpr "bad `fun' syntax in ~s" sexpr)])]

    ; Math.
    [(list '+ lhs rhs) (Add (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list '- lhs rhs) (Sub (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list '* lhs rhs) (Mul (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list '/ lhs rhs) (Div (parse-sexpr lhs) (parse-sexpr rhs))]

    ; Function calls.
    [(list 'call fun) (Call (parse-sexpr fun) '())]
    [(list 'call fun arg) (Call (parse-sexpr fun) (list (parse-sexpr arg)))]
    [(list fun 'of arg) (Call (parse-sexpr fun) (list (parse-sexpr arg)))]
    [(list 'call fun args ...) (Call (parse-sexpr fun) (map parse-sexpr args))]
    [(list fun 'of args ...) (Call (parse-sexpr fun) (map parse-sexpr args))]

    ; Binds.
    [(list 'bind more ...)
     (match sexpr
       [(list 'bind (list (list (symbol: names) (sexpr: vals)) ...) body) (Bind names (map parse-sexpr vals) (parse-sexpr body))]
       [else (error 'parse-sexpr "bad bind syntax in ~s" sexpr)])]

    [(list 'bind* more ...)
     (match sexpr
       [(list 'bind* (list (list (symbol: names) (sexpr: vals)) ...) body) (Bind* names (map parse-sexpr vals) (parse-sexpr body))]
       [else (error 'parse-sexpr "bad bind* syntax in ~s" sexpr)])]

    [(list 'if0 n body alt) (If0 (parse-sexpr n) (parse-sexpr body) (parse-sexpr alt))]

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

(: parse : String -> TIDBIT)
;; parses a string containing a TIDBIT expression to a TIDBIT AST
(define (parse str)
  (parse-sexpr (string->sexpr str)))

;; Types for environments, values, and a lookup function
(define-type VAL
  [NumV Number]
  [FunV CORE ENV])

(define-type ENV = (Listof VAL))

(: NumV->number : VAL -> Number)
;; convert a TIDBIT runtime numeric value to a Racket one
(define (NumV->number val)
  (cases val
    [(NumV n) n]
    [else (error 'arith-op "expected a number, got: ~s" val)]))

(: arith-op : (Number Number -> Number) VAL VAL -> VAL)
;; gets a Racket numeric binary operator, and uses it within a NumV
;; wrapper
(define (arith-op op val1 val2)
  (NumV (op (NumV->number val1) (NumV->number val2))))

(: eval : CORE ENV -> VAL)
;; evaluates CORE expressions by reducing them to values
(define (eval expr env)
  (cases expr
         [(CNum n) (NumV n)]
         [(CAdd l r) (arith-op + (eval l env) (eval r env))]
         [(CSub l r) (arith-op - (eval l env) (eval r env))]
         [(CMul l r) (arith-op * (eval l env) (eval r env))]
         [(CDiv l r) (arith-op / (eval l env) (eval r env))]
         [(CIdx n) (list-ref env (sub1 n))]
         [(CFun body) (FunV body env)]
         [(CCall fun-expr arg-expr)
          (let ([fval (eval fun-expr env)])
            (cases fval
                   [(FunV body f-env) (eval body (cons (eval arg-expr env) f-env))]
                   [else (error 'eval "`call' expects a function, got: ~s" fval)]))]
         [(CIf0 n body alt) (if (zero? (NumV->number (eval n env))) (eval body env) (eval alt env))]))

(: run : String -> Number)
;; evaluate a TIDBIT program contained in a string
(define (run str)
  (let ([result (eval (preprocess (parse str) empty-depth) '())])
    (cases result
      [(NumV n) n]
      [else (error 'run "evaluation returned a non-number: ~s"
                   result)])))

(test (run "{if0 0 1 2}") => 1)
(test (run "{if0 1 1 2}") => 2)


; Exercise 2
(test (run "{call {fun {x x} x} 1 2}") => 2) ; Does not check parameters are unique.
(test (run "{call {call {fun {x y} {+ x y}} 1} 2}") => 3) ; Not fulfilling the arity returns a function. Kind of a neat feature.
(test (run "{call {fun {x y} {fun {z} {+ x {+ y z}}}} 1 2 3}") => 6) ; Again arity mismatch.

(test (run "{bind {{x y} {y 1}} {+ x y}}") =error> "empty-depth: No binding for y.")


; Is this a valid way of implementing nullary functions?
(test (run "{call {fun 4}}") => 4)
(test (run "{call {fun 4} 3}") =error> "eval: `call' expects a function, got: (NumV 4)")
(test (run "{bind* {{x 4} {f {fun x}}} {call f}}") => 4)
(test (run "{bind* {{f {fun {x} {+ x 1}}} {g {fun 10}} {h {fun 6}}} {+ {call h} {call f {call g}}}}") => 17)
(test (run "{fun 10}") => 10)

(test (run "{bind* {{x 2} {y x} {z y}} {+ x {* y z}}}") => 6)
(test (run "{bind {{x 2} {y x} {z y}} {+ x {* y z}}}") =error> "empty-depth: No binding for x.")

(test (run "{bind* {{x 2} {y 3} {z 4}} {+ x {* y z}}}") => 14)
(test (run "{bind {{x 2} {y 3} {z 4}} {+ x {* y z}}}") => 14)
(test (run "{{fun {x y z} {+ x {+ y z}}} of 1 2 3}") => 6)

(test (run "{{fun {x} {+ x 1}} of 4}")
      => 5)
(test (run "{with {add3 {fun {x} {+ x 3}}}
              {add3 of 1}}")
      => 4)
(test (run "{with {x 1} {with {y 2} {+ x y}}}") => 3)
(test (run "{call {call {fun {x} {fun {y} {+ x y}}} 1} 2}") => 3)

(test (run "{with fh dhd dhdh lja}") =error> "parse-sexpr: bad `with' syntax in (with fh dhd dhdh lja)")
(test (run "{fun fh dhd dhdh lja}") =error> "parse-sexpr: bad `fun' syntax in (fun fh dhd dhdh lja)")
(test (run "{}") =error> "parse-sexpr: bad syntax in ()")
(test (run "{asdf of 2}") =error> "empty-depth: No binding for asdf.")
(test (run "{+ 1 {fun {x} x}}") =error> "arith-op: expected a number, got: (FunV (CIdx 1) ())")
(test (run "{1 of 1}") =error> "eval: `call' expects a function, got: (NumV 1)")
(test (run "{fun {x} x}") =error> "run: evaluation returned a non-number: (FunV (CIdx 1) ())")

(test (run "{with {identity {fun {x} x}}
              {with {foo {fun {x} {+ x 1}}}
                {{identity of foo} of 123}}}")
      => 124)
(test (run "{with {add3 {fun {x} {+ x 3}}}
              {with {add1 {fun {x} {+ x 1}}}
                {with {x 3}
                  {add1 of {add3 of x}}}}}")
      => 7)
(test (run "{with {x 3}
              {with {f {fun {y} {+ x y}}}
                {with {x 5}
                  {f of 4}}}}")
      => 7)
(test (run "{call {with {x 3}
                    {fun {y} {+ x y}}}
                  4}")
      => 7)
(test (run "{with {f {with {x 3} {fun {y} {+ x y}}}}
              {with {x 100}
                {f of 4}}}")
      => 7)
(test (run "{call {call {fun {x} {x of 1}}
                        {fun {x} {fun {y} {+ x y}}}}
                  123}")
      => 124)