;; ** The Picky interpreter, no explicit types

#lang pl

#|
The grammar:
  <PICKY> ::= <num>
            | <id>
            | { + <PICKY> <PICKY> }
            | { - <PICKY> <PICKY> }
            | { = <PICKY> <PICKY> }
            | { < <PICKY> <PICKY> }
            | { fun { <id> } <PICKY> }
            | { call <PICKY> <PICKY> }
            | { with { <id> <PICKY> } <PICKY> }
            | { if <PICKY> <PICKY> <PICKY> }

The types are no longer part of the input syntax.

Evaluation rules:
  eval(N,env)                = N
  eval(x,env)                = lookup(x,env)
  eval({+ E1 E2},env)        = eval(E1,env) + eval(E2,env)
  eval({- E1 E2},env)        = eval(E1,env) - eval(E2,env)
  eval({= E1 E2},env)        = eval(E1,env) = eval(E2,env)
  eval({< E1 E2},env)        = eval(E1,env) < eval(E2,env)
  eval({fun {x} E},env)      = <{fun {x} E}, env>
  eval({call E1 E2},env1)    = eval(B,extend(x,eval(E2,env1),env2))
                              if eval(E1,env1) = <{fun {x} B}, env2>
                             = error!  otherwise <-- never happens
  eval({with {x E1} E2},env) = eval(E2,extend(x,eval(E1,env),env))
  eval({if E1 E2 E3},env)    = eval(E2,env)  if eval(E1,env) is true
                             = eval(E3,env)  otherwise

Type checking rules (note the extra complexity in the `fun' rule):

  Γ ⊢ n : Number

  Γ ⊢ x : Γ(x)

  Γ ⊢ A : Number   Γ ⊢ B : Number
  ———————————————————————————————
       Γ ⊢ {+ A B} : Number

  Γ ⊢ A : Number   Γ ⊢ B : Number
  ———————————————————————————————
       Γ ⊢ {< A B} : Boolean

       Γ[x:=τ₁] ⊢ E : τ₂
  ————————————————————————————
  Γ ⊢ {fun {x} E} : (τ₁ -> τ₂)

  Γ ⊢ F : (τ₁ -> τ₂)  Γ ⊢ V : τ₁
  ——————————————————————————————
       Γ ⊢ {call F V} : τ₂

  Γ ⊢ C : Boolean   Γ ⊢ T : τ   Γ ⊢ E : τ
  ———————————————————————————————————————
            Γ ⊢ {if C T E} : τ

  Γ ⊢ V : τ₁   Γ[x:=τ₁] ⊢ E : τ₂
  ——————————————————————————————
     Γ ⊢ {with {x V} E} : τ₂

|#

(define-type PICKY
  [Num   Number]
  [Id    Symbol]
  [Add   PICKY PICKY]
  [Sub   PICKY PICKY]
  [Equal PICKY PICKY]
  [Less  PICKY PICKY]
  [Fun   Symbol PICKY] ; no types even here
  [Call  PICKY PICKY]
  [With  Symbol PICKY PICKY]
  [If    PICKY PICKY PICKY])

(: parse-sexpr : Sexpr -> PICKY)
;; parses s-expressions into PICKYs
(define (parse-sexpr sexpr)
  (match sexpr
    [(number: n)    (Num n)]
    [(symbol: name) (Id name)]
    [(list '+ lhs rhs) (Add   (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list '- lhs rhs) (Sub   (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list '= lhs rhs) (Equal (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list '< lhs rhs) (Less  (parse-sexpr lhs) (parse-sexpr rhs))]
    [(list 'call fun arg)
                       (Call  (parse-sexpr fun) (parse-sexpr arg))]
    [(list 'if c t e)
     (If (parse-sexpr c) (parse-sexpr t) (parse-sexpr e))]
    [(cons 'fun more)
     (match sexpr
       [(list 'fun (list (symbol: name)) body)
        (Fun name (parse-sexpr body))]
       [else (error 'parse-sexpr "bad `fun' syntax in ~s" sexpr)])]
    [(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)])]
    [else (error 'parse-sexpr "bad expression syntax: ~s" sexpr)]))

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

;; Typechecker and related types and helpers

;; this is not a part of the AST now, and it also has a new variant
;; for type variables (see `same-type' for how it's used)
(define-type TYPE
  [NumT]
  [BoolT]
  [FunT TYPE TYPE]
  [?T (Boxof (U TYPE #f))])

;; this is similar to ENV, but it holds type information for the
;; identifiers during typechecking; it is essentially "Γ"
(define-type TYPEENV
  [EmptyTypeEnv]
  [ExtendTypeEnv Symbol TYPE TYPEENV])

(: type-lookup : Symbol TYPEENV -> TYPE)
;; similar to `lookup' for type environments; note that the
;; error is phrased as a typecheck error, since this indicates
;; a failure at the type checking stage
(define (type-lookup name typeenv)
  (cases typeenv
    [(EmptyTypeEnv) (error 'typecheck "no binding for ~s" name)]
    [(ExtendTypeEnv id type rest-env)
     (if (eq? id name) type (type-lookup name rest-env))]))

(: typecheck : PICKY TYPE TYPEENV -> Void)
;; Checks that the given expression has the specified type. Used
;; only for side-effects, so return a void value. There are two
;; side-effects that it can do: throw an error if the input
;; expression doesn't typecheck, and type variables can be mutated
;; once their values are known -- this is done by the `types='
;; utility function that follows.
(define (typecheck expr type type-env)
  ;; convenient helpers
  (: type= : TYPE -> Void)
  (define (type= type2) (types= type type2 expr))
  (: two-nums : PICKY PICKY -> Void)
  (define (two-nums e1 e2)
    (typecheck e1 (NumT) type-env)
    (typecheck e2 (NumT) type-env))
  (cases expr
    [(Num n)     (type= (NumT))]
    [(Id name)   (type= (type-lookup name type-env))]
    [(Add   l r) (type= (NumT))  (two-nums l r)] ; note that the
    [(Sub   l r) (type= (NumT))  (two-nums l r)] ; order in these
    [(Equal l r) (type= (BoolT)) (two-nums l r)] ; things can be
    [(Less  l r) (type= (BoolT)) (two-nums l r)] ; swapped...
    [(Fun bound-id bound-body)
     (let (;; the identity of these type variables is important!
           [itype (?T (box #f))]
           [otype (?T (box #f))])
       (type= (FunT itype otype))
       (typecheck bound-body otype
                  (ExtendTypeEnv bound-id itype type-env)))]
    [(Call fun arg)
     (let ([type2 (?T (box #f))]) ; same here
       (typecheck arg type2 type-env)
       (typecheck fun (FunT type2 type) type-env))]
    [(With bound-id named-expr bound-body)
     (let ([type2 (?T (box #f))]) ; and here
       (typecheck named-expr type2 type-env)
       (typecheck bound-body type
                  (ExtendTypeEnv bound-id type2 type-env)))]
    [(If cond-expr then-expr else-expr)
     (typecheck cond-expr (BoolT) type-env)
     (typecheck then-expr type type-env)
     (typecheck else-expr type type-env)]))

(: types= : TYPE TYPE PICKY -> Void)
;; Compares the two input types, and throw an error if they don't
;; match. This function is the core of `typecheck', and it is used
;; only for its side-effect. Another side effect in addition to
;; throwing an error is when type variables are present -- they will
;; be mutated in an attempt to make the typecheck succeed. Note that
;; the two type arguments are not symmetric: the first type is the
;; expected one, and the second is the one that the code implies
;; -- but this matters only for the error messages. Also, the
;; expression input is used only for these errors. As the code
;; clearly shows, the main work is done by `same-type' below.
(define (types= type1 type2 expr)
  (unless (same-type type1 type2)
    (error 'typecheck "type error for ~s: expecting ~a, got ~a"
           expr (type->string type1) (type->string type2))))

(: type->string : TYPE -> String)
;; Convert a TYPE to a human readable string,
;; used for error messages
(define (type->string type)
  (format "~s" type)
  ;; The code below would be useful, but unfortunately it doesn't
  ;; work in some cases. To see the problem, try to run the example
  ;; below that applies identity on itself. It's left here so you
  ;; can try it out when you're not running into this problem.
  #|
  (cases type
    [(NumT)  "Num"]
    [(BoolT) "Bool"]
    [(FunT i o)
     (string-append (type->string i) " -> " (type->string o))]
    [(?T box)
     (let ([t (unbox box)])
       (if t (type->string t) "?"))])
  |#)

;; Convenience type to make it possible to have a single `cases'
;; dispatch on two types instead of nesting `cases' in each branch
(define-type 2TYPES [PairT TYPE TYPE])

(: same-type : TYPE TYPE -> Boolean)
;; Compares the two input types, return true or false whether
;; they're the same. The process might involve mutating ?T type
;; variables.
(define (same-type type1 type2)
  ;; the `PairT' type is only used to conveniently match on both
  ;; types in a single `cases', it's not used in any other way
  (cases (PairT type1 type2)
    ;; flatten the first type, or set it to the second if it's unset
    [(PairT (?T box) type2)
     (let ([t1 (unbox box)])
       (if t1
         (same-type t1 type2)
         (begin (set-box! box type2) #t)))]
    ;; do the same for the second (reuse the above case)
    [(PairT type1 (?T box)) (same-type type2 type1)]
    ;; the rest are obvious
    [(PairT (NumT) (NumT)) #t]
    [(PairT (BoolT) (BoolT)) #t]
    [(PairT (FunT i1 o1) (FunT i2 o2))
     (and (same-type i1 i2) (same-type o1 o2))]
    [else #f]))

;; Evaluator and related types and helpers

(define-type ENV
  [EmptyEnv]
  [Extend Symbol VAL ENV])

(define-type VAL
  [NumV  Number]
  [BoolV Boolean]
  [FunV  Symbol PICKY ENV])

(: lookup : Symbol ENV -> VAL)
;; lookup a symbol in an environment, 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)]
    [(Extend id val rest-env)
     (if (eq? id name) val (lookup name rest-env))]))

(: strip-numv : Symbol VAL -> Number)
;; converts a VAL to a Racket number if possible, throws an error if
;; not using the given name for the error message
(define (strip-numv name val)
  (cases val
    [(NumV n) n]
    ;; this error will never be reached, see below for more
    [else (error name "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 (strip-numv 'arith-op val1)
            (strip-numv 'arith-op val2))))

(: bool-op : (Number Number -> Boolean) VAL VAL -> VAL)
;; gets a Racket numeric binary predicate, and uses it
;; within a BoolV wrapper
(define (bool-op op val1 val2)
  (BoolV (op (strip-numv 'bool-op val1)
             (strip-numv 'bool-op val2))))

(: eval : PICKY ENV -> VAL)
;; evaluates PICKY expressions by reducing them to values
(define (eval expr env)
  (cases expr
    [(Num n) (NumV n)]
    [(Id name) (lookup name env)]
    [(Add   l r) (arith-op + (eval l env) (eval r env))]
    [(Sub   l r) (arith-op - (eval l env) (eval r env))]
    [(Equal l r) (bool-op  = (eval l env) (eval r env))]
    [(Less  l r) (bool-op  < (eval l env) (eval r env))]
    [(Fun bound-id bound-body) (FunV bound-id bound-body env)]
    [(Call fun-expr arg-expr)
     (define fval (eval fun-expr env))
     (cases fval
       [(FunV bound-id bound-body f-env)
        (eval bound-body
              (Extend bound-id (eval arg-expr env) f-env))]
       ;; `cases' requires complete coverage of all variants, but
       ;; this `else' is never used since we typecheck programs
       [else (error 'eval "`call' expects a function, got: ~s"
                    fval)])]
    [(With bound-id named-expr bound-body)
     (eval bound-body (Extend bound-id (eval named-expr env) env))]
    [(If cond-expr then-expr else-expr)
     (let ([bval (eval cond-expr env)])
       (if (cases bval
             [(BoolV b) b]
             ;; same as above: this case is never reached
             [else (error 'eval "`if' expects a boolean, got: ~s"
                                bval)])
         (eval then-expr env)
         (eval else-expr env)))]))

(: run : String -> Number)
;; evaluate a PICKY program contained in a string
(define (run str)
  (let ([prog (parse str)])
    (typecheck prog (NumT) (EmptyTypeEnv))
    (let ([result (eval prog (EmptyEnv))])
      (cases result
        [(NumV n) n]
        ;; this error is never reached, since we make sure
        ;; that the program always evaluates to a number above
        [else (error 'run "evaluation returned a non-number: ~s"
                          result)]))))

;; tests -- including translations of the FLANG tests
(test (run "5") => 5)
(test (run "{fun {x} {+ x 1}}") =error> "type error")
(test (run "{call {fun {x} {+ x 1}} 4}") => 5)
(test (run "{with {x 3} {+ x 1}}") => 4)
(test (run "{with {identity {fun {x} x}} {call identity 1}}") => 1)
(test (run "{with {add3 {fun {x} {+ x 3}}} {call add3 1}}") => 4)
(test (run "{with {add3 {fun {x} {+ x 3}}}
              {with {add1 {fun {x} {+ x 1}}}
                {with {x 3}
                  {call add1 {call add3 x}}}}}")
      => 7)
(test (run "{with {identity {fun {x} x}}
              {with {foo {fun {x} {+ x 1}}}
                {call {call identity foo} 123}}}")
      => 124)
(test (run "{with {x 3}
              {with {f {fun {y} {+ x y}}}
                {with {x 5} {call f 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}
                {call f 4}}}")
      => 7)
(test (run "{call {call {fun {x} {call x 1}}
                        {fun {x} {fun {y} {+ x y}}}}
                  123}")
      => 124)
(test (run "{call {fun {x} {if {< x 2} {+ x 5} {+ x 6}}} 1}") => 6)
(test (run "{call {fun {x} {if {< x 2} {+ x 5} {+ x 6}}} 2}") => 8)

;; Note that we still have a language with the same type system,
;; even though it looks like it could be more flexible -- for
;; example, the following two examples work:
(test (run "{with {identity {fun {x} x}}
              {call identity 1}}")
      => 1)
(test (run "{with {identity {fun {x} x}}
              {if {call identity {< 1 2}} 1 2}}")
      => 1)
;; but this doesn't, since identity can not be used with different
;; types:
(test (run "{with {identity {fun {x} x}}
              {if {call identity {< 1 2}}
                {call identity 1}
                2}}")
      =error> "type error")
;; this doesn't work either -- with an interesting error message:
(test (run "{with {identity {fun {x} x}}
              {call {call identity identity} 1}}")
      =error> "type error")
;; ... but these two work fine:
(test (run "{with {identity1 {fun {x} x}}
              {with {identity2 {fun {x} x}}
                {+ {call identity1 1}
                   {if {call identity2 {< 1 2}} 1 2}}}}")
      => 2)
(test (run "{with {identity1 {fun {x} x}}
              {with {identity2 {fun {x} x}}
                {call {call identity1 identity2} 1}}}")
      => 1)