(* Data definitions *)

type cae = 
    Num of int
  | Add of cae * cae
  | Sub of cae * cae
  | IfZ of cae * cae * cae
  | Arg
  | This
  | New of string * cae list
  | Get of cae * int
  | DSend of cae * int * cae
  | SSend of cae * cae * cae

and cdecl = Class of string * int * cae list

and caeValue =
    NumV of int
  | ObjV of cdecl * caeValue list

exception Failed of string

(* Number operations *)

let mkNumOp = fun op
  -> function
      (NumV(a), NumV(b)) -> NumV(op a b)
    | _ -> raise (Failed "not numbers")
          
let numPlus = mkNumOp (fun a b -> a + b)
let numMinus = mkNumOp (fun a b -> a - b)

(* Lookup *)

exception NoSuch of string * string

let rec find = fun what nameOf name vals ->
  match vals with
    [] -> raise (NoSuch (what, name))
  | a::rest -> 
      if (name = nameOf(a))
      then a
      else find what nameOf name rest
          
let rec locate = fun what nameOf name vals ->
  match vals with
    [] -> raise (NoSuch (what, name))
  | a::rest -> 
      if (name = nameOf(a))
      then 0
      else 1 + (locate what nameOf name rest)
          
let findClass = (find "class"
                   (fun (Class(name, _, _)) -> name))

(* Interpreter *)

let rec interp : (cae * cdecl list * caeValue * caeValue 
                  -> caeValue ) 
    = function (expr, cdecls, this, arg) -> 
      let recur = fun e -> interp(e, cdecls, this, arg)
      in match expr with
      Num(n) -> NumV(n)
    | Add(l, r) -> numPlus(recur l, recur r)
    | Sub(l, r) -> numMinus(recur l, recur r)
    | IfZ(tst, thn, els) ->
        if (recur tst = NumV(0))
        then recur thn
        else recur els
    | This -> this
    | Arg -> arg
    | New(name, exprs) -> 
        let decl = findClass name cdecls
        in let vals = (List.map recur exprs)
        in ObjV(decl, vals)
    | Get(expr, n) ->
        (match recur expr with
          ObjV(_, vals) -> 
            List.nth vals n
        | _ -> raise (Failed "not an object for get"))
    | DSend(expr, n, argExpr) ->
        (match recur expr with
          (ObjV(Class(_, _, methods), _) as this) -> 
            let body = List.nth methods n
            in interp(body, cdecls, this, recur argExpr)
        | _ -> raise (Failed "not an object for send"))
    | SSend(expr, body, argExpr) ->
        let this = recur expr
        in interp(body, cdecls, this, recur argExpr)

(* Examples *)

let posnMDist = Add(Get(This, 0), Get(This, 1))

let posnAddDist = Add(DSend(This, 0, Num(0)),
                      DSend(Arg, 0, Num(0)))

let posnClass = Class("posn",
                      2,
                      [posnMDist;
                       posnAddDist;
                       Add(Get(This, 0), Arg);
                       Sub(Arg, Get(This, 1));
                       New("posn", [Num(0); Num(1)])])
    
let posn3DClass = Class("posn3D",
                        3,
                        [Add(Get(This, 2), 
                             SSend(This, posnMDist, Arg));
                         SSend(This, posnAddDist, Arg)])
    
let mkPosn27 = New("posn", [Num(2); Num(7)])
let mkPosn531 = New("posn3D", [Num(5); Num(3); Num(1)])

let mdist = fun o -> DSend(o, 0, Num(0))
let addDist = fun o p -> DSend(o, 1, p)
let addX = fun o y -> DSend(o, 2, y)
let subY = fun o y -> DSend(o, 3, y)

let interpPosn = fun x -> interp(x, [posnClass;posn3DClass], NumV(0), NumV(0))

;;

(* Testing *)

exception TestFailed of caeValue * caeValue

let test = function 
    (a, b) ->
      if (a = b)
      then true
      else raise (TestFailed (a, b))
        
;;
  
(* Test cases *)

test(interpPosn(mdist mkPosn27),
     NumV(9));;
test(interpPosn(addX mkPosn27 (Num 10)),
     NumV(12));;
test(interpPosn(subY (DSend(mkPosn27, 4, Num(0))) (Num 15)),
     NumV(14));;
test(interpPosn(addDist mkPosn531  mkPosn27),
     NumV(18));;

(* ********************************************* *)

(* Source Language with Inheritance *)

type icae = 
    INum of int
  | IAdd of icae * icae
  | ISub of icae * icae
  | IIfZ of icae * icae * icae
  | IArg
  | IThis
  | INew of string * icae list
  | IGet of icae * string * string
  | ISend of icae * string * string * icae
  | ISuper of string * icae

and idecl = 
    IClass of string * string * ifield list * imeth list
      
and ifield = IField of string

and imeth = IMethod of string * icae


let findIClass = (find "class"
                    (fun (IClass(name, _, _, _)) -> name))

let findIMethod = (find "method"
                    (fun (IMethod(name, _)) -> name))

let locateIField = (locate "field"
                      (fun (IField(name)) -> name))

let locateIMethod = (locate "method"
                       (fun (IMethod(name, _)) -> name))

let rec classFieldCount = function 
    (cname, idecls) ->
      if (cname = "object")
      then 0
      else let IClass(_, sname, fields, _) = findIClass cname idecls
      in (List.length fields) + classFieldCount(sname, idecls)

let rec compileExpr = function
    (expr, thisClass, idecls) -> 
      let recur = fun expr -> compileExpr(expr, thisClass, idecls)
      in match expr with
        INum(n) -> Num(n)
      | IAdd(l, r) -> Add(recur l, recur r)
      | ISub(l, r) -> Sub(recur l, recur r)
      | IIfZ(tst, thn, els) -> IfZ(recur tst, 
                                   recur thn,
                                   recur els)
      | IArg -> Arg
      | IThis -> This
      | INew(s, exprs) -> New(s, List.map recur exprs)
      | IGet(expr, cname, fname) -> 
          let IClass(_, sname, fields, _) = findIClass cname idecls
          in Get(recur expr, ((locateIField fname fields)
                                + classFieldCount(sname, idecls)))
      | ISend(expr, cname, mname, argExpr) -> 
          let IClass(_, _, _, methods) = findIClass cname idecls
          in DSend(recur expr, locateIMethod mname methods, recur argExpr)
      | ISuper(mname, expr) ->
          let IClass(_, sname, _, _) = thisClass
          in let super = findIClass sname idecls
          in let IClass(_, _, _, methods) = super
          in let IMethod(_, body) = findIMethod mname methods
          in SSend(This, compileExpr(body, super, idecls), recur expr)

let rec compileMethods = function
    (sdecl, idecls) -> 
      let IClass(name, superName, fields, methods) = sdecl 
      in Class(name,
               List.length fields,
               List.map 
                 (fun (IMethod(name, expr)) -> 
                   compileExpr(expr,
                               sdecl,
                               idecls))
                 methods)

let addFields = function
    (superFields, fields) -> List.append superFields fields
        
let rec addReplaceMethods : (cae list * string list 
			       * cae list * string list
			     -> cae list * string list) = function
    (methods, names, [], []) -> (methods, names)
  | (methods, names, meth::mrest, name::nrest) ->
      let (methods, names) = addReplaceMethod(methods, names, meth, name)
      in addReplaceMethods(methods,
                           names,
                           mrest,
                           nrest)
  | _ -> raise (Failed "shouldn't happen")
      
and addReplaceMethod : (cae list * string list * cae * string
			-> cae list * string list) = function
    ([], [], bmeth, bname) -> ([bmeth], [bname])
  | (ameth::arest, aname::arestnames, bmeth, bname)
    -> if (aname = bname)
    then (bmeth::arest, bname::arestnames)
    else let (meths, names) 
        = addReplaceMethod (arest, arestnames, bmeth, bname)
    in (ameth::meths, aname::names)
  | _ -> raise (Failed "shouldn't happen")

let rec flattenClassNames : (cdecl * idecl list * cdecl list 
			     -> cdecl * string list) = function
    (Class(name, fields, methods), idecls, cdecls) -> 
      let IClass(_, superName, _, imethods) = findIClass name idecls
      in let (Class(_, superFields, superMethods), superMNames)
          = if (superName = "object") 
          then (Class("object", 0, []), [])
          else flattenClassNames(findClass superName cdecls,
                                 idecls, cdecls)
      in let (methods, names)
          = addReplaceMethods(superMethods, 
                              superMNames,
                              methods, 
                              (List.map 
                                 (fun (IMethod(name, _)) -> name)
                                 imethods))
      in (Class(name, 
                superFields + fields,
                methods),
          names)

let flattenClass = function x -> 
  let (c, names) = flattenClassNames x
  in c

let iinterp = function
    (idecls, expr)  -> 
      let expr 
          = compileExpr(expr, IClass("bad", "bad", [], []), idecls)
      in let cdeclsNotFlat
          = (List.map (fun sdecl -> compileMethods(sdecl, idecls))
               idecls)
      in let cdecls 
          = (List.map (fun cdecl -> flattenClass(cdecl, 
                                                 idecls,
                                                 cdeclsNotFlat))
               cdeclsNotFlat)
      in interp(expr, cdecls, NumV(0), NumV(0))
        
(* Examples *)

let sposnClass = IClass("posn", "object",
                        [IField("x"); IField("y")],
                        [IMethod("mdist", IAdd(IGet(IThis, "posn", "x"), IGet(IThis, "posn", "y")));
                         IMethod("addDist", IAdd(ISend(IThis, "posn", "mdist", INum(0)),
                                                 ISend(IArg, "posn", "mdist", INum(0))))])
    
let sposn3DClass = IClass("posn3D", "posn",
                          [IField("z")],
                          [IMethod("mdist", IAdd(IGet(IThis, "posn3D", "z"), 
                                                 ISuper("mdist", IArg)))])
    
let smkPosn27 = INew("posn", [INum(2); INum(7)])
let smkPosn531 = INew("posn3D", [INum(5); INum(3); INum(1)])

let smdist = fun o -> ISend(o, "posn", "mdist", INum(0))
let saddDist = fun o p -> ISend(o, "posn", "addDist", p)

let sinterpPosn = fun x -> iinterp([sposnClass;sposn3DClass], x)

;;

(* Test cases *)

test(sinterpPosn(smdist smkPosn27),
     NumV(9));;
test(sinterpPosn(smdist smkPosn531),
     NumV(9));;
test(sinterpPosn(saddDist smkPosn531 smkPosn27),
     NumV(18));;
test(sinterpPosn(saddDist smkPosn27 smkPosn531),
     NumV(18));;

(* ********************************************* *)

type tcdecl = TClass of string * string * tfield list * tmeth list
      
and tfield = TField of string * te
    
and tmeth = TMethod of string * te * te * icae
    
and te = 
    NumTE
  | ObjTE of string

and ty = 
    NumT
  | ObjT of string

(* Type checking *)

let findTClass = (find "class"
                    (fun (TClass(name, _, _, _)) -> name))
let findTField = (find "field"
                    (fun (TField(name, _)) -> name))
let findTMethod = (find "method"
                     (fun (TMethod(name, _, _, _)) -> name))

let parseType = function
    NumTE -> NumT
  | ObjTE(s) -> ObjT(s)

let rec getAllFieldTypes = function
    (cname, tdecls) -> 
      if (cname = "object")
      then []
      else let TClass(_, sname, fields, _) = findTClass cname tdecls
      in (List.append 
            (List.map (fun (TField(_, te)) -> parseType te)
               fields)
            (getAllFieldTypes (sname, tdecls)))

let rec findInTree = fun
  findInList extract (name, tdecl, tdecls) -> 
    let items = extract tdecl
    and TClass(_, sname, _, _) = tdecl
    in try (findInList name items)
    with exn -> if (sname == "object")
    then raise exn
    else (findInTree findInList extract
            (name, findTClass sname tdecls, tdecls))
          
let findFieldInTree = (findInTree findTField
                         (fun (TClass(_, _, fields, _)) -> fields))
        
let findMethodInTree = (findInTree findTMethod
                          (fun (TClass(_, _, _, methods)) -> methods))

let rec isSubClass = function
    (aname, bname, tdecls) -> 
      if (aname = bname)
      then true
      else if (aname = "object")
      then false
      else let TClass(_, sname, _, _) = findTClass aname tdecls
      in isSubClass(sname, bname, tdecls)

let isSubType = function
    (ObjT(aname), ObjT(bname), tdecls) ->
      isSubClass(aname, bname, tdecls)
  | (NumT, NumT, _) -> true
  | _ -> false

exception NoType of icae * string

let rec typecheckExpr = function
    (expr, tdecls, argTy, thisClass) -> 
      let recur = fun expr -> typecheckExpr(expr, tdecls, argTy, thisClass)
      in match expr with
        INum(n) -> NumT
      | IAdd(l, r) -> 
          (match (recur l, recur r) with
            (NumT, NumT) -> NumT
          | _ -> raise (NoType(expr, "not numbers")))
      | ISub(l, r) -> 
          (match (recur l, recur r) with
            (NumT, NumT) -> NumT
          | _ -> raise (NoType(expr, "not numbers")))
      | IIfZ(tst, thn, els) -> 
          (match (recur tst) with
            NumT -> 
              let thnty = recur thn
              and elsty = recur els
              in if (isSubType(thnty, elsty, tdecls))
              then elsty
              else if (isSubType(elsty, thnty, tdecls))
              then thnty
              else raise (NoType (expr, "branch mismatch"))
          | _ -> raise (NoType(expr, "test is non-number")))
      | IArg -> argTy
      | IThis -> 
          let TClass(name, _, _, _) = thisClass
          in ObjT(name)
      | INew(cname, exprs) -> 
          let argTys = List.map recur exprs
          and fieldTys = getAllFieldTypes(cname, tdecls)
          in if (argTys = fieldTys)
          then ObjT(cname)
          else raise (NoType(expr, "field type mismatch"))
      | IGet(expr, getCName, fname) -> 
          (match (recur expr) with 
            ObjT(cname) ->
              if not (isSubClass(cname, getCName, tdecls))
              then raise (NoType(expr, "field class mismatch"))
              else let TField(_, fieldTE) 
                  = findFieldInTree(fname,
                                    findTClass cname tdecls,
                                    tdecls)
              in parseType(fieldTE)
          | _ -> raise (NoType(expr, "not an object for get")))
      | ISend(expr, sendCName, mname, argExpr) -> 
          (match (recur expr) with 
            ObjT(cname) ->
              if not (isSubClass(cname, sendCName, tdecls))
              then raise (NoType(expr, "method class mismatch"))
              else typecheckSend(cname, mname, argExpr, tdecls, recur)
          | _ -> raise (NoType(expr, "not an object for send")))
      | ISuper(mname, argExpr) ->
          let TClass(_, sname, _, _) = thisClass
          in typecheckSend(sname, mname, argExpr, tdecls, recur)

and typecheckSend = function 
    (cname, mname, argExpr, tdecls, recur) -> 
      let TMethod(_, argTE, resultTE, _) 
          = findMethodInTree(mname,
                             findTClass cname tdecls,
                             tdecls)
      in if (isSubType((recur argExpr), parseType(argTE), tdecls))
      then parseType(resultTE)
      else raise (NoType(argExpr, "arg type mismatch"))

let typecheckMethod = function
    (TMethod(_, argTE, resultTE, body), tdecl, tdecls) ->
      let bodyT = typecheckExpr(body, tdecls, parseType argTE, tdecl)
      in if (isSubType(bodyT, (parseType resultTE), tdecls))
      then ()
      else raise(NoType(body, "result type mismatch"))

exception BadOverride of string

let checkOverride = function
    (TMethod(mname, argTE, resultTE, _), tdecl, tdecls) ->
      let TClass(_, sname, _, _) = tdecl
      in try let TMethod(_, sArgTE, sResultTE, _) 
          = findMethodInTree(mname,
                             findTClass sname tdecls,
                             tdecls)
      in if ((sArgTE = argTE) & (sResultTE = resultTE))
      then ()
      else raise(BadOverride(mname))
      with (NoSuch _) -> ()

let typecheck = function
    (tdecls, expr)  -> 
      let _ = (List.map 
                 (fun tdecl ->
                   let TClass(_, _, _, methods) = tdecl
                   in (List.map (fun m -> (typecheckMethod(m, tdecl, tdecls);
                                           checkOverride(m , tdecl, tdecls)))
                         methods))
                 tdecls)
      in typecheckExpr(expr, tdecls, NumT, TClass("bad", "bad", [], []))

let stripTypes = function
    TClass(name, sname, fields, methods) ->
      IClass(name, sname,
             (List.map (fun (TField(n,te)) -> IField(n)) 
                fields),
             (List.map (fun (TMethod(n,ate,rte,body)) -> IMethod(n, body)) 
                methods))

let tinterp = function
    (tdecls, expr) ->
      iinterp(List.map stripTypes tdecls, expr)

(* Examples *)

let tposnClass = TClass("posn", "object",
                        [TField("x", NumTE); TField("y", NumTE)],
                        [TMethod("mdist", NumTE, NumTE, IAdd(IGet(IThis, "posn", "x"), IGet(IThis, "posn", "y")));
                         TMethod("addDist", ObjTE("posn"), NumTE,
                                 IAdd(ISend(IThis, "posn", "mdist", INum(0)),
                                      ISend(IArg, "posn", "mdist", INum(0))))])
    
let tposn3DClass = TClass("posn3D", "posn",
                          [TField("z", NumTE)],
                          [TMethod("mdist", NumTE, NumTE,
                                   IAdd(IGet(IThis, "posn3D", "z"), 
                                        ISuper("mdist", IArg)))])
    
let tinterpPosn = fun x -> tinterp([tposnClass;tposn3DClass], 
                                   x)
let typecheckPosn = fun x -> typecheck([tposnClass;tposn3DClass], x)

(* More Testing *)

exception TyTestFailed of ty * ty

let tytest = function 
    (a, b) ->
      if (a = b)
      then true
      else raise (TyTestFailed (a, b))

let notytest = function 
    thunk -> try (thunk (); raise (Failed "no exception!?")) with 
      BadOverride _ -> true
    | NoType(_, _) -> true
    | NoSuch(_, _) -> true

;;

(* Test cases *)

test(tinterpPosn(smdist smkPosn27),
     NumV(9));;
test(tinterpPosn(smdist smkPosn531),
     NumV(9));;
test(tinterpPosn(saddDist smkPosn531 smkPosn27),
     NumV(18));;
test(tinterpPosn(saddDist smkPosn27 smkPosn531),
     NumV(18));;

tytest(typecheckPosn(smdist smkPosn27),
       NumT);;
tytest(typecheckPosn(smdist smkPosn531),
       NumT);;
tytest(typecheckPosn(saddDist smkPosn531 smkPosn27),
       NumT);;
tytest(typecheckPosn(saddDist smkPosn27 smkPosn531),
       NumT);;

notytest(fun () -> typecheckPosn(smdist (INum(10))));;
notytest(fun () -> typecheck([tposnClass;
                              TClass("other", "posn",
                                     [],
                                     [TMethod("mdist", ObjTE("object"), NumTE, INum(10))])],
                             INum(10)));;