(* If you would like, write how many hours you spent on this homework:

  XXX

*)

module  StringMap = Map.Make(String)

type value = int
type op = | Add | Mul | Sub | Negate | Div | And | Or | LT | EQ | Not

let eval_op (o : op) (vs : value list) : value = 
  match o, vs with
  | Add, [n1; n2] -> (n1 + n2)
  | Mul, [n1; n2] -> (n1 * n2)
  | Sub, [n1; n2] -> (n1 - n2)
  | And, [n1; n2] -> (if n1 = 1 && n2 = 1 then 1 else 0)
  | Negate, [n]   -> (-n)
  | Or, [n1; n2] -> (if n1 = 1 || n2 = 1 then 1 else 0)
  | Div, [n1; n2] -> n1 / n2
  | LT,  [n1; n2] -> (if n1 < n2 then 1 else 0)
  | EQ,  [n1; n2] -> (if n1 = n2 then 1 else 0)
  | Not, [n]      -> (if n = 0 then 1 else 0)
  | _ -> failwith "eval_op: incorrect arguments"

module type Theory = sig
  type ty

  val type_of_const : value -> ty option

  val type_of_op : op -> ty list -> ty option

  (* Used for if/then/else *)
  val merge_ty : ty -> ty -> ty option

  val interpret_ty : ty -> value -> bool
end

module IntTheory : Theory = struct
  type ty = | TInt 

  let type_of_op (o : op) (args : ty list) : ty option =
    match o, args with
    | Add, [TInt; TInt] -> Some TInt
    | Mul, [TInt; TInt] -> Some TInt
    | Sub, [TInt; TInt] -> Some TInt
    (* Div is disallowed. *)
    (* | Div, [TInt; TInt] -> Some TInt *)
    | Negate,  [TInt] -> Some TInt
    | And, [TInt; TInt] -> Some TInt
    | Or,  [TInt; TInt] -> Some TInt
    | Not,  [TInt] -> Some TInt
    | EQ,  [TInt; TInt] -> Some TInt
    | _ -> None

  let type_of_const (_ : value) : ty option =
    (* All values are well-typed *)
    Some TInt


  let merge_ty (t1 : ty) (t2 : ty) = 
      if t1 = t2 then Some t1 else None
  let interpret_ty (_ : ty) (_ : value) = 
    true
end


module BasicLang (S : Theory) = struct
  type expr = | Value of value | Var of string | Let of string * expr * expr | App of op * expr list | Ite of expr * expr * expr

  let rec eval (m : value StringMap.t) (e : expr) : value = 
    match e with
    | Value v -> v
    | Var x -> StringMap.find x m (* Will throw exception if not in the map. *)
    | Let (x, e1, e2) ->
      let v = eval m e1 in 
      eval (StringMap.add x v m) e2
    | App (op, args) ->
      let vs = List.map (eval m) args in
      eval_op op vs
    | Ite (g, t, f) ->
      if eval m g = 1 then 
          eval m t 
      else 
          eval m f

  let rec typecheck (m : S.ty StringMap.t) (e : expr) : S.ty option = 
    match e with
    | Value v -> S.type_of_const v
    | Var x -> StringMap.find_opt x m
    | Let (x, e1, e2) ->
      (match typecheck m e1 with
       | None -> None
       | Some t1 ->
         let m' = StringMap.add x t1 m in
         typecheck m' e2)
    | App (op, args) ->
      let arg_tys_opt = List.map (typecheck m) args in
      if List.exists Option.is_none arg_tys_opt then None
      else
        S.type_of_op op (List.map Option.get arg_tys_opt)
    | Ite (g, e1, e2) ->
      (match typecheck m g with
       | Some _ ->
         (match (typecheck m e1, typecheck m e2) with
          | (Some t1, Some t2) -> S.merge_ty t1 t2
          | _ -> None)
       | _ -> None)
end

(* Q1: Why must division be disallowed? *)
module Q1 = struct 
  open BasicLang(IntTheory)
  let bad_div () : expr  = failwith "TODO"
end

(* Q2: A theory for Positive / Negative / Zero / Any *)
module SignTheory : Theory = struct
  type ty = | Pos | Zero | Neg | Any
  let type_of_const (v: value) : ty option = failwith "TODO"
  let type_of_op (o : op) (ts : ty list) = failwith "TODO"
  let merge_ty (t1 : ty) (t2 : ty) : ty option = failwith "TODO"
  let interpret_ty (t : ty) (v : value) : bool = failwith "TODO"

end

(* Q3: Extend with interval analysis *)
module IntervalTheory : Theory = struct
  type ty = (int * int) (* First component must be <= second component! *)
  let type_of_const (v: value) : ty option = failwith "TODO"
  let type_of_op (o : op) (ts : ty list) = failwith "TODO"
  let merge_ty (t1 : ty) (t2 : ty) : ty option = failwith "TODO"
  let interpret_ty (t : ty) (v : value) : bool = failwith "TODO"
end

(* Q4: Modulus Analysis *)
module Mod5Theory : Theory = struct
  type ty = | Mod5 of int (* Must be in range [0..4] *)| Any
  let type_of_const (v: value) : ty option = failwith "TODO"
  let type_of_op (o : op) (ts : ty list) = failwith "TODO"
  let merge_ty (t1 : ty) (t2 : ty) : ty option = failwith "TODO"
  let interpret_ty (t : ty) (v : value) : bool = failwith "TODO"
end

(* Q5: MiniPython *)
module MiniPython (S : Theory) = struct

  type expr = | Value of value | Var of string | App of op * expr list

  type cmd = | Assign of string * expr | Ite of expr * cmd * cmd | Seq of cmd * cmd | While of expr * cmd | Skip
  
  let rec eval_expr (m : value StringMap.t) (e: expr) : value = 
    match e with 
    | Value v -> v
    | Var x -> StringMap.find x m
    | App (o, es) -> eval_op o (List.map (eval_expr m) es)

  let rec typecheck_expr (m : S.ty StringMap.t) (e: expr) : S.ty option = failwith "TODO" 

  let rec eval_cmd (m : value StringMap.t) (c : cmd) : value StringMap.t = 
    match c with 
    | Assign (x, e) -> StringMap.add x (eval_expr m e) m
    | Ite (e, c1, c2) -> if eval_expr m e = 1 then eval_cmd m c1 else eval_cmd m c2
    | Seq (c1, c2) -> let m' = eval_cmd m c1 in eval_cmd m' c2
    | While (e, c) -> 
        if eval_expr m e = 1 then 
            let m' = eval_cmd m c in 
            eval_cmd m' (While (e, c))
        else 
          m
    | Skip -> m

  let rec typecheck_cmd (m : S.ty StringMap.t) (c : cmd) : bool = failwith "TODO"

end