#lang plait #:untyped

;; On object is a tuple containing
;;   - a hash table for fields
;;   - a hash table for methods, where each method
;;     is a function that takes `this` and one argument
(define-type-alias Object
  ((Listof (Symbol * ....))
   * (Listof (Symbol * (.... .... -> ....)))))

(define-syntax-rule (object ([field-id field-expr] ...)
                            [method-id body-expr] ...)
  (values (list (values 'field-id field-expr)
                ...)
          (list (values 'method-id (method body-expr))
                ...)))

;; `(method body)` expands to `(lambda (this arg) body)`, but
;; in a way that allows `body` to refer to `this` and `arg'
(define-syntax (method stx)
  (syntax-case stx ()
    [(method body)
     (with-syntax ([this-id (datum->syntax #'body 'this)]
                   [arg-id (datum->syntax #'body 'arg)])
       #'(lambda (this-id arg-id)
           body))]))

;; Extract a field form the first hash table.
(define-syntax-rule (get o-expr f-id)
  (find (fst o-expr) 'f-id))

;; Extract a method form the second hash table.
(define-syntax-rule (send o-expr m-id arg-expr)
  (let ([o o-expr])
    ((find (snd o) 'm-id) o arg-expr)))

(define (find [l : (Listof (Symbol * 'a))] [name : Symbol]) : 'a
  (type-case (Listof (Symbol * 'a)) l
    [empty
     (error 'find (string-append "not found: " (symbol->string name)))]
    [(cons p rst-l)
     (if (symbol=? (fst p) name)
         (snd p)
         (find rst-l name))]))

;; ----------------------------------------

(define-type-alias Value Object)

(define-syntax-rule (numV init-n)
  (object ([n init-n])
          [apply (error 'interp "not a function")]
          [number (get this n)]))

;; Equivalently:
#;
(define-syntax-rule (numV init-n)
  (values
   (list (values 'n init-n))
   (list (values 'apply
                 (lambda (this arg)
                   (error 'interp "not a function")))
         (values 'number
                 (lambda (this arg)
                   (get this n))))))

(define-syntax-rule (closV init-n init-body init-c-env)
  (object ([n init-n]
           [body init-body]
           [c-env init-c-env])
          [apply (send (get this body) interp
                       (extend-env 
                        (bind (get this n) arg)
                        (get this c-env)))]
          [number (error 'interp "not a number")]))

;; ----------------------------------------

(define-type-alias ExprE Object)

(define-syntax-rule (numE init-n)
  (object ([n init-n])
          [interp (numV (get this n))]))

(define-syntax-rule (idE init-s)
  (object ([s init-s])
          [interp (send arg lookup (get this s))]))

(define-syntax-rule (plusE init-l init-r)
  (object ([l init-l]
           [r init-r])
          [interp (num+ (send (get this l) interp arg)
                        (send (get this r) interp arg))]))

(define-syntax-rule (multE init-l init-r)
  (object ([l init-l]
           [r init-r])
          [interp (num* (send (get this l) interp arg)
                        (send (get this r) interp arg))]))

(define-syntax-rule (letE init-n init-rhs init-body)
  (object ([n init-n]
           [rhs init-rhs]
           [body init-body])
          [interp (send (get this body) interp
                        (extend-env 
                         (bind (get this n)
                               (send (get this rhs) interp arg))
                         arg))]))

(define-syntax-rule (lamE init-n init-body)
  (object ([n init-n]
           [body init-body])
          [interp (closV (get this n)
                         (get this body)
                         arg)]))

(define-syntax-rule (appE init-fun init-arg)
  (object ([fun init-fun]
           [arg init-arg])
          [interp (send (send (get this fun) interp arg)
                        apply
                        (send (get this arg) interp arg))]))

;; ----------------------------------------

(define-type-alias Env Object)

(define-type Binding
  (bind [name : Symbol]
        [val : Value]))

(define mt-env
  (object ()
          [lookup (error 'interp "free variable")]))

(define-syntax-rule (extend-env init-b init-env)
  (object ([b init-b]
           [env init-env])
          [lookup (if (symbol=? arg (bind-name (get this b)))
                      (bind-val (get this b))
                      (send (get this env) lookup arg))]))

(module+ test
  (print-only-errors #t))

;; parse ----------------------------------------
(define (parse s)
  (cond
    [(s-exp-match? `NUMBER s) (numE (s-exp->number s))]
    [(s-exp-match? `SYMBOL s) (idE (s-exp->symbol s))]
    [(s-exp-match? `{+ ANY ANY} s)
     (plusE (parse (second (s-exp->list s)))
            (parse (third (s-exp->list s))))]
    [(s-exp-match? `{* ANY ANY} s)
     (multE (parse (second (s-exp->list s)))
            (parse (third (s-exp->list s))))]
    [(s-exp-match? `{let {[SYMBOL ANY]} ANY} s)
     (let ([bs (s-exp->list (first
                             (s-exp->list (second
                                           (s-exp->list s)))))])
       (letE (s-exp->symbol (first bs))
             (parse (second bs))
             (parse (third (s-exp->list s)))))]
    [(s-exp-match? `{lambda {SYMBOL} ANY} s)
     (lamE (s-exp->symbol (first (s-exp->list 
                                  (second (s-exp->list s)))))
           (parse (third (s-exp->list s))))]
    [(s-exp-match? `{ANY ANY} s)
     (appE (parse (first (s-exp->list s)))
           (parse (second (s-exp->list s))))]
    [else (error 'parse "invalid input")]))

;; Testing is a problem, because we can't easily inspect
;; function results...
#;
(module+ test
  (test (parse '2)
        (numE 2))
  (test (parse `x) ; note: backquote instead of normal quote
        (idE 'x))
  (test (parse '{+ 2 1})
        (plusE (numE 2) (numE 1)))
  (test (parse '{* 3 4})
        (multE (numE 3) (numE 4)))
  (test (parse '{+ {* 3 4} 8})
        (plusE (multE (numE 3) (numE 4))
               (numE 8)))
  (test (parse `{let {[x {+ 1 2}]}
                  y})
        (appE (lamE 'x (idE 'y))
              (plusE (numE 1) (numE 2))))
  (test (parse '{lambda {x} 9})
        (lamE 'x (numE 9)))
  (test (parse '{double 9})
        (appE (idE 'double) (numE 9)))
  (test/exn (parse '{{+ 1 2}})
            "invalid input"))

;; interp ----------------------------------------
(define (interp a env)
  (send a interp env))

(module+ test
  (test (send (interp (parse `2) mt-env)
              number 0)
        2)
  (test/exn (interp (parse `x) mt-env)
            "free variable")
  (test (send (interp (parse `x) 
                      (extend-env (bind 'x (numV 9)) mt-env))
              number 0)
        9)
  (test (send (interp (parse `{+ 2 1}) mt-env)
              number 0)
        3)
  (test (send (interp (parse `{* 2 1}) mt-env)
              number 0)
        2)
  (test (send (interp (parse `{+ {* 2 3} {+ 5 8}})
                      mt-env)
              number 0)
        19)
  ;; The following test is a problem, too, since the
  ;; expected result has a function representing
  ;; the body expression in the closure.
  #;
  (test (interp (parse '{lambda {x} {+ x x}})
                mt-env)
        (closV 'x (plusE (idE 'x) (idE 'x)) mt-env))
  (test (send (interp (parse `{let {[x 5]}
                                {+ x x}})
                      mt-env)
              number 0)
        10)
  (test (send (interp (parse `{let {[x 5]}
                                {let {[x {+ 1 x}]}
                                  {+ x x}}})
                      mt-env)
              number 0)
        12)
  (test (send (interp (parse `{let {[x 5]}
                                {let {[y 6]}
                                  x}})
                      mt-env)
              number 0)
        5)
  (test (send (interp (parse `{{lambda {x} {+ x x}} 8})
                      mt-env)
              number 0)
        16)

  (test/exn (interp (parse `{1 2}) mt-env)
            "not a function")
  (test/exn (interp (parse `{+ 1 {lambda {x} x}}) mt-env)
            "not a number")
  (test/exn (interp (parse `{let {[bad {lambda {x} {+ x y}}]}
                              {let {[y 5]}
                                {bad 2}}})
                    mt-env)
            "free variable"))

;; num+ and num* ----------------------------------------
(define (num-op op l r)
  (numV (op (send l number 0) (send r number 0))))

(define (num+ l r)
  (num-op + l r))
(define (num* l r)
  (num-op * l r))

(module+ test
  (test (send (num+ (numV 1) (numV 2)) number 0)
        3)
  (test (send (num* (numV 2) (numV 3)) number 0)
        6))

;; lookup ----------------------------------------
(define (lookup n env)
  (send env lookup n))

(module+ test
  (test/exn (lookup 'x mt-env)
            "free variable")
  (test (send (lookup 'x (extend-env (bind 'x (numV 8)) mt-env))
              number 0)
        8)
  (test (send (lookup 'x (extend-env
                          (bind 'x (numV 9))
                          (extend-env (bind 'x (numV 8)) mt-env)))
              number 0)
        9)
  (test (send (lookup 'y (extend-env
                          (bind 'x (numV 9))
                          (extend-env (bind 'y (numV 8)) mt-env)))
              number 0)
        8))