;; --- MiniJava interpreter ---

(defun run-java (fname &optional (stdout t) (stderr t) (stdin nil))
  "Run the named MiniJava program"
  (mj-run (parse-java fname) stdout stderr))

(defun mj-run (ast &optional (stdout t) (stderr t) (stdin nil))
  "Interpret a MiniJava program from the AST"
  (mj-statement (fourth (second ast)) (setup-mj-env ast stdout stderr stdin)))


;; --- MiniJava type-checker ---

(defun tc-java (fname)
  "Type check the named MiniJava program"
  (mj-tc (parse-java fname)))

(defun mj-tc (ast &optional (stderr t))
  "Type check a MiniJava program from the AST"
  (let ((env (setup-mj-env ast t stderr)))
    (mj-tc-main (second ast) env)
    (mj-tc-methods env)))


;;; ------------------------------------------------------------------
;;; Preprocessing and environment setup
;;; ------------------------------------------------------------------


;; --- MiniJava environment ---

;; The mj-env structure tracks the current execution state.  That
;; means that it knows enough to do variable and method lookups, both
;; using the current stack frame and in any object instances.  Because
;; the AST is not organized well for Lisp lookups, I'm going to chop
;; it up and distribute the bits into a bunch of hash tables.

(defstruct mj-env
  stdout    ; standard output stream
  stderr    ; error output stream
  stdin     ; standard input stream
  parents   ; class-name -> parent class-name
  classes   ; class-name -> class layout
  methods   ; (class-name . method-name) -> method info
  vtables   ; class-name -> vtable (associates method -> slot, implementor)
  stack)    ; Run-time stack

(defun setup-mj-env (ast &optional (stdout t) (stderr t) (stdin nil))
  "Set up a MiniJava run-time environment"
  (let ((class-decls (cdr (third ast)))
        (env (make-mj-env)))
    (setf (mj-env-stdout  env) stdout)
    (setf (mj-env-stderr  env) stderr)
    (setf (mj-env-stdin   env) stdin)
    (setf (mj-env-parents env) (setup-mj-parents class-decls))
    (setf (mj-env-classes env) (setup-mj-classes class-decls))
    (setf (mj-env-methods env) (setup-mj-methods class-decls))
    env))

(defun id-name (id)
  "Convert various forms of 'id' to an unadorned symbol."
  (cond 
    ((atom id)                    id)
    ((eq (car id) 'id)            (second id))
    ((eq (car id) 'IdentifierExp) (second id))
    (t (error "Invalid identifier node"))))


;; --- Variable info lists ---

;; A var-info list maps names to variable information, specifically
;;   - The variable type, and
;;   - The index of the variable in a stack frame or object layout
;; Indices are assigned sequentially starting from 1 (the zero slot is
;; reserved, eg for a vtable pointer or a "this" object).  Variables
;; may be shadowed, so for each name, we keep a stack of declarations.

(defun make-var-info ()
  "Create a new binding list"
  (let ((vl (make-hash-table)))
    (setf (gethash '*count* vl) 0)
    vl))

(defun get-var-info (id vl)
  "Get var info from the list."
  (let ((id (id-name id)))
    (car (gethash id vl))))

(defun get-var-slot (id vl)
  "Get index info about a binding from the list"
  (car (get-var-info id vl)))

(defun get-var-type (id vl)
  "Get type info about a variable"
  (cdr (get-var-info id vl)))

(defun get-var-count (vl)
  "Find out how many vars are already present"
  (gethash '*count* vl))

(defun add-var-info (id type vl)
  "Allocate a new binding"
  (let ((id (id-name id)))
    (if (get-var-info id vl)
        (format t "Warning: shadowed declarations of ~A~%" id))
    (push `(,(incf (gethash '*count* vl)) . ,type)
          (gethash id vl))))

(defun map-visible-vars (func vl)
  "Call (func key slot type) on all visible vars in a layout"
  (maphash #'(lambda (key vars)
               (unless (eq key '*count*)
                 (funcall func key (caar vars) (cdar vars))))
           vl))

(defun map-vars (func vl)
  "Call (func key slot type) on all vars in a layout"
  (maphash #'(lambda (key vars)
               (unless (eq key '*count*)
                 (map 'list #'(lambda (var)
                                (funcall func key (car var) (cdr var)))
                      vars)))
           vl))


;; --- Parents ---

;; The inheritance relations among classes are stored as assoc lists
;; of (class-name . parent-class-name).  When we set up this assoc list,
;; let's also run basic sanity checks: no undefined classes and no cyclic
;; inheritance relations.  Usually inheritance isn't that deep, so there's
;; no point in using a fancy algorithm for the cycle detection.

(defun setup-mj-parents (class-decls)
  "Set up an assoc list of (class-name . parent-class)"
  (let ((parent-list nil))
    (dolist (class-decl class-decls)
      (let* ((name    (id-name (second class-decl)))
             (extends (third class-decl))
             (parent  (id-name (second extends))))
        (push (cons name parent) parent-list)))
    ; (check-mj-parents parent-list)
    parent-list))

; (defun check-mj-parents (parent-list)
;   "Check to make sure that the inheritance relations are legal")


;; --- Class object layouts ---

;; A class instance is represented as an array whose zeroth offset
;; contains a vtable entry.  The other slots are allocated in such a way
;; that subclasses and superclasses always behave compatibly, i.e.
;;
;;    [vtable, superclass data, subclass data]
;;
;; This is the way C++ usually lays out object data structures (in the
;; case of single inheritance -- multiple inheritance is a lot messier).
;; I'm just going to represent the vtable by the class name.

(defun setup-mj-classes (class-decls)
  "Set up class-name -> layout assoc mapping"
  (let ((class-info (compress-mj-class-info class-decls))
        (map nil))
    (dolist (class class-info)
      (push `(,(first class) . ,(setup-mj-class class class-info)) map))
    map))

(defun compress-mj-class-info (class-decls)
  "Assemble list of (class-name parent-name var-list) from class-decls"
  (let ((class-info nil))
    (dolist (class-decl class-decls)
      (let ((class-name  (id-name (second class-decl)))
            (parent-name (id-name (second (third class-decl))))
            (vars-list   (cdr (fourth class-decl))))
        (push `(,class-name ,parent-name ,vars-list) class-info)))
    (nreverse class-info)))

(defun setup-mj-class (class class-info)
  "Compute the layout for a class"
  (let ((layout (make-var-info)))
    (add-mj-layout class class-info layout)))

(defun add-mj-layout (class class-info layout)
  "Add this class's info to an object layout"
  (when class
    (let ((parent-name (second class))
          (vars-list   (third class)))
      (add-mj-layout (assoc parent-name class-info) class-info layout)
      (setup-mj-vars vars-list layout)))
  layout)

(defun setup-mj-vars (var-decls var-info)
  "Add a vars list into a layout"
  (dolist (var-decl var-decls)
    (let ((var-name (id-name (third var-decl)))
          (var-type (second var-decl)))
      (add-var-info var-name var-type var-info))))


;; --- Methods info ---

;; We identify methods by (class-name . method-name) cons pairs.  For
;; each method, we keep around the formals list, the code, and a table
;; that maps local names -> stack frame positions.

(defstruct mj-method
  class       ; Name of this class (for setting 'this' type)
  name        ; Method name
  formals     ; Formal parameter list (as an (id . type) assoc list)
  layout      ; Stack frame layout
  code        ; Statement list
  type        ; Return value type
  value)      ; Return value expression

(defun setup-mj-methods (class-decls)
  "Make a table of (class-name . method-name) -> method info"
  (let ((methods (make-hash-table :test #'equal)))
    (dolist (class-decl class-decls)
      (dolist (method-decl (cdr (fifth class-decl)))
        (let ((class-name  (id-name (second class-decl)))
              (method-name (id-name (third method-decl))))
          (setf (gethash `(,class-name . ,method-name) methods)
                (setup-mj-method class-name method-decl)))))
    methods))

(defun setup-mj-method (class-name method-decl)
  "Set up the record for a MiniJava method"
  (let ((method (make-mj-method)))
    (setf (mj-method-class   method) class-name)
    (setf (mj-method-name    method) (id-name (third method-decl)))
    (setf (mj-method-layout  method) (make-var-info))
    (setf (mj-method-code    method) (cdr (sixth method-decl)))
    (setf (mj-method-type    method) (second method-decl))
    (setf (mj-method-value   method) (seventh method-decl))
    (setup-mj-formals method-decl method)
    (setup-mj-vars (cdr (fifth method-decl))
                   (mj-method-layout method))
    method))

(defun setup-mj-formals (method-decl method)
  "Process the formals list for a method declaration"
  (let ((formals nil))
    (dolist (formal (cdr (fourth method-decl)))
      (let ((formal-name (id-name (third formal)))
            (formal-type (second formal)))
        (push (cons formal-name formal-type) formals)
        (add-var-info formal-name formal-type (mj-method-layout method))))
    (setf (mj-method-formals method) (nreverse formals))))


;;; ------------------------------------------------------------------
;;; Type checking
;;; ------------------------------------------------------------------

;; EXERCISE: Write the type-checking system for MiniJava.  We've
;; included functions and comments for our type-checker; you may
;; decide to use different function names and data structures.

;; --- Type comparison operations ---

;; We need not only functions to check if types are the same (which is easy),
;; but also functions to check if one type is a subtype of another.

;; --- Method checks ---

;; We have to ensure compatibility of virtual methods; that is,
;; methods parent.foo and child.foo should have the same argument
;; and return types.  I believe this is stricter than what Java
;; requires (Java would allow different methods with the name "foo",
;; so long as the method calls could be disambiguated from the types
;; of the arguments).
;; 
;; Otherwise, an individual method is okay if all its statements are
;; okay, and if the expression for the return value is compatible with
;; the return type.

;(defun mj-tc-methods (env)
;  "Type-check all methods.")


;; --- Main check ---

;(defun mj-tc-main (ast env)
;  "Type check a MiniJava main statement from the MainClass AST")


;; --- Variable lookups for type-checking ---

;; We can look up variables in the stack, or we can look them up in
;; "this."  For the type checker phase, we keep around a mock stack
;; frame of the form (stack-layout this-layout class marks) to do 
;; this lookup.  As we process a method, we will mark seen variables
;; in the "marks" array, so that at the end of the day we can check
;; for unused variables.


;; --- Type-check statements ---

;; Compare the type-checking code below to the analogous interpreter code.

;(defun mj-tc-statement (ast env)
;  "Type check a MiniJava statement")


;; --- Type-check values ---

;; There's one subtle point to notice here, which is a difference between
;; the behavior of mj-tc-call and mj-call.  The type checking for the method
;; occurs based on the lexical type of the object; but the dispatch occurs
;; based on the dynamic type of the object.  So if we want any sort of sanity,
;; it's critical that the two type signatures match (which is enforced
;; by earlier checks).

;(defun mj-tc-call (call-ast env)
;  "Typecheck a method call.")

;(defun mj-tc-exp (ast env)
;  "Get the type of a Mini-Java expression subtree")


;;; ------------------------------------------------------------------
;;; Interpretation
;;; ------------------------------------------------------------------


;; --- Object creation ---

;; MiniJava objects are represented as (type . array) pairs, where
;; the car is the lexical type.  The dynamic type (the type used to
;; originally create the object) is stored in the zeroth entry of the
;; array; as mentioned earlier, this is similar to how C++ handles
;; object layouts.  The dynamic type is used for method dispatch, but
;; the lexical type is used for variable access.  
;;
;; There's a subtle point here: if a superclass and a subclass both
;; have variables with the same name, then (assuming this isn't
;; forbidden as an error) the "right" behavior is for references to
;; the name to depend on the lexical type.  That is, methods from the
;; parent class should use the parent's version of "foo," while
;; methods from the child class should use the child's version of
;; "foo."  This is what I do now -- but I also give a warning, because
;; you probably really didn't mean to overshadow the previous declaration.

(defun mj-new-object (class-name env)
  "Allocate a new MiniJava object"
  (let ((class-layout (cdr (assoc class-name (mj-env-classes env)))))
    (when (not class-layout)
      (error "Could not find class for new object."))
    (let ((obj-data (make-array `(,(1+ (get-var-count class-layout))))))
      (setf (elt obj-data 0) class-name)             ; Dynamic type
      (map-vars #'(lambda (name slot type)           ; Default values
                    (setf (elt obj-data slot) (mj-type-default type)))
                class-layout)
      `(,class-name . ,obj-data))))


;; --- Type defaults and checks ---

;; When we check type compatibility (e.g. for passing an object), we will
;; automatically cast the requested type (if possible).

(defun mj-type-default (type)
  "Return the default value for the given MiniJava type"
  (cond
    ((eq type 'IntType)      0)
    ((eq type 'IntArrayType) nil)
    ((eq type 'BooleanType)  'false)
    ((listp type)            nil)
    (t (error "Unrecognized type"))))

(defun has-parent (type1 type2 env)
  "Check if type1 has type2 as a parent"
  (when type1
    (or (eq type1 type2)
        (let ((type1-parent (cdr (assoc type1 (mj-env-parents env)))))
          (has-parent type1-parent type2 env)))))

;; Don't have to fill this in...
(defun mj-check-type (type value env)
  "Check that the value matches the given type")

(defun mj-typecast (type value env)
  "Typecast the given value to the given type"
  ; (mj-check-type type value env)
  (if (listp type)
      `(,(id-name (second type)) . ,(cdr value))
      value))


;; --- Method calls ---

;; To call a method, we set up the stack frame, initializing all
;; the formals (including 'this') to the corresponding actuals, and
;; initializing all the local variables to their default values.
;; Then we execute all the statements in turn, evaluate the return
;; value, pop the frame, and pass back the return value.
;;
;; Note that an interpreter stack frame consists of a pair: the layout
;; table + the stack data.  In a compiled code, we would already know
;; the layout.

(defun mj-call (call-ast env)
  "Call a method."
  (labels ((e (expr) (mj-exp-eval expr env)))
    (let ((this    (e (second call-ast)))
          (name    (id-name (third call-ast)))
          (actuals (mapcar #'e (cdr (fourth call-ast)))))
      (mj-dispatch this name actuals env))))

(defun mj-dispatch (this method-name actuals env)
  "Execute a method call."
  (let ((method (obj-method-lookup this method-name env)))
    (push (mj-new-frame method this actuals env) (mj-env-stack env))
    (dolist (statement (mj-method-code method))
      (mj-statement statement env))
    (let ((return-type  (mj-method-type method))
          (return-value (mj-exp-eval (mj-method-value method) env)))
      (pop (mj-env-stack env))
      (mj-typecast return-type return-value env))))

(defun mj-new-frame (method this actuals env)
  "Set up a new stack frame"
  (let* ((class   (mj-method-class method))
         (formals (mj-method-formals method))
         (layout  (mj-method-layout method))
         (frame   (make-array `(,(1+ (get-var-count layout))))))
    (map-vars #'(lambda (name slot type)             ; Default values
                  (setf (elt frame slot) (mj-type-default type)))
              layout)
    (setf (elt frame 0) `(,class . ,(cdr this)))     ; Set this pointer
    (unless (eq (length formals) (length actuals))
      (format t "Incorrect number of arguments"))
    (dotimes (i (length formals))
      (let ((formal (pop formals))
            (actual (pop actuals)))
        ;; Old code (with type-checking in interpreter):
        ;; (setf (elt frame (1+ i)) (mj-check-type (cdr formal) actual env))))
        (setf (elt frame (1+ i)) actual)))
    `(,layout . ,frame)))

(defun obj-method-lookup (this method-name env)
  "Look up an object's method."
  ;; Note: We check the vtable to get the method used, but we check
  ;;   the declared type in order to decide whether we should have
  ;;   access to that method.
  (when (or (not this)
            (not (listp this)))
    (error "Invalid object in method invocation"))
  (when (not (method-lookup (car this) method-name env))
    (error "Invalid method name"))
  (method-lookup (elt (cdr this) 0) method-name env))

(defun method-lookup (class-name method-name env)
  "Look up the nearest method"
  (when class-name
    (or (gethash `(,class-name . ,method-name) (mj-env-methods env))
        (method-lookup (cdr (assoc class-name (mj-env-parents env)))
                       method-name
                       env))))


;; --- Variable lookup and assignment ---

;; When we look up a variable, we first try the local variable list,
;; then go to the members of 'this'.  You can't get a variable out of
;; any object other than 'this' in MiniJava.

(defun mj-get-var-info (var-name env)
  "Look up the info on a variable"
  (when (not (mj-env-stack env))
    (error "Cannot access variables in main routine"))
  (let* ((stack-frame  (car (mj-env-stack env)))
         (stack-layout (car stack-frame))
         (stack-data   (cdr stack-frame))
         (this         (elt stack-data 0))
         (this-layout  (cdr (assoc (car this) (mj-env-classes env))))
         (this-data    (cdr this))
         (stack-info   (get-var-info var-name stack-layout))
         (this-info    (get-var-info var-name this-layout)))
    (if stack-info
        (values stack-info stack-data)
        (if this-info
            (values this-info this-data)
            (error "Unknown variable")))))

(defun mj-get-value (var-name env)
  "Read a value"
  (multiple-value-bind (info array)
      (mj-get-var-info var-name env)
    (elt array (car info))))

(defun mj-set-value (var-name value env)
  "Write a value"
  (multiple-value-bind (info array)
      (mj-get-var-info var-name env)
    (setf (elt array (car info))
          (mj-check-type (cdr info) value env))))

(defun mj-this (env)
  "Look up 'this'"
  (when (not (mj-env-stack env))
    (error "Cannot access 'this' in main routine"))
  (elt (cdar (mj-env-stack env)) 0))


;; --- MiniJava statements and expressions ---

(defun mj-statement (ast env)
  "Execute a MiniJava statement"
  (labels 
      ((c (v)  (eq (car ast) v))
       (e (i)  (mj-exp-eval  (nth i ast) env))
       (s (i)  (mj-statement (nth i ast) env)))
    (cond 
      ((c 'If)          (if (eq (e 1) 'true) (s 2) (s 3)))
      ((c 'While)       (while (eq (e 1) 'true) (s 2)))
      ((c 'Print)       (format (mj-env-stdout env) "~A~%" (e 1)))
      ((c 'Assign)      (mj-set-value (id-name (second ast))
                                      (e 2) env))
      ((c 'ArrayAssign) (setf (elt (mj-get-value (id-name (second ast)) env)
                                   (e 2))
                              (e 3)))
      ((c 'Block)       (dolist (s (cdadr ast)) (mj-statement s env)))
      (t
       (pprint ast)
       (error "Unexpected statement")))))
     
(defun mj-exp-eval (ast env)
  "Evaluate a Mini-Java expression subtree"
  (labels 
      ((c  (v) (eq (car ast) v))
       (e1 ()  (mj-exp-eval (second ast) env))
       (e2 ()  (mj-exp-eval (third  ast) env)))
    (cond 
      ((eq ast 'this)      (mj-this env))
      ((atom ast)          (error "Unexpected atomic expression"))
      ((c 'Not)            (mj-not (e1)))
      ((c 'And)            (mj-and (second ast) (third ast) env))
      ((c 'LessThan)       (mj-<   (e1) (e2)))
      ((c 'Plus)           (mj-+   (e1) (e2)))
      ((c 'Minus)          (mj--   (e1) (e2)))
      ((c 'Times)          (mj-*   (e1) (e2)))
      ((c 'IntegerLiteral) (second ast))
      ((c 'BooleanLiteral) (second ast))
      ((c 'ArrayLookup)    (elt (e1) (e2)))
      ((c 'ArrayLength)    (length (e1)))
      ((c 'NewArray)       (make-array `(,(e1)) :initial-element 0))
      ((c 'NewObject)      (mj-new-object (id-name (second (second ast))) env))
      ((c 'Call)           (mj-call ast env))
      ((c 'Read)           (mj-read env))
      ((c 'IdentifierExp)  (mj-get-value (id-name (second ast)) env)))))

(defun mj-not (e1)
  (unless (or (eq e1 'true) (eq e1 'false))
    (error "Cannot take not of non-boolean value"))
  (if (eq e1 'true)
      'false
      'true))

(defun mj-bool-check (val)
  (cond ((eq val 'true) t)
        ((eq val 'false) nil)
        (t (error (format nil "Expected boolean, saw ~A" val)))))

(defun mj-and (clause1 clause2 env)
  (if (and (mj-bool-check (mj-exp-eval clause1 env))
           (mj-bool-check (mj-exp-eval clause2 env)))
      'true
      'false))

(defun mj-+ (e1 e2)
  (unless (and (numberp e1) (numberp e2))
    (error "Cannot take the sum of non-numeric values"))
  (+ e1 e2))

(defun mj-- (e1 e2)
  (unless (and (numberp e1) (numberp e2))
    (error "Cannot take the difference of non-numeric values"))
  (- e1 e2))

(defun mj-* (e1 e2)
  (unless (and (numberp e1) (numberp e2))
    (error "Cannot take the product of non-numeric values"))
  (* e1 e2))

(defun mj-< (e1 e2)
  (unless (and (numberp e1) (numberp e2))
    (error "Cannot compare non-numeric values"))
  (if (< e1 e2) 'true 'false))

(defun mj-read (env)
  "Read a MiniJava integer"
  (let ((result (if (mj-env-stdin env)
                    (read (mj-env-stdin env))
                    (progn
                      (format t "~%Enter number: ")
                      (read)))))
    (if (integerp result)
        result
        0)))