pe.ml

type expr =
| Var of string
| Int of int
| Add of (expr * expr)
| Mult of (expr * expr)
| Let of (string * expr * expr)
| IfZero of (expr * expr * expr)
| Pair of (expr * expr)
| Zro of expr
| Fst of expr
| Left of expr
| Right of expr
| Match of (expr * (string * expr) * (string * expr))
| Abs of (string * expr)
| App of (expr * expr)
| Count
| Unit
| MkRef of expr
| GetRef of expr
| SetRef of (expr * expr)

module ENV = Map.Make(String)
module STORE = Map.Make(Int)

type value =
| VInt of int
| VPair of (value * value)
| VLeft of value
| VRight of value
| VFunc of (value -> value)
| VRef of value ref
| VUnit

let rec eval (e : value ENV.t) (c : int ref) : expr -> value =
  let recurse = fun expr -> eval e c expr in
  function
  | Var v -> ENV.find v e
  | Int i -> VInt i
  | Add (x, y) -> (match (recurse x, recurse y) with (VInt x, VInt y) -> VInt (x + y) | _ -> assert false)
  | Mult (x, y) -> (match (recurse x, recurse y) with (VInt x, VInt y) -> VInt (x * y) | _ -> assert false)
  | Let (var, v, body) ->
      eval (ENV.add var (recurse v) e) c body
  | IfZero (c, x, y) ->
      (match recurse c with VInt ci -> if ci = 0 then recurse x else recurse y | _ -> assert false)
  | Pair (x, y) -> VPair (recurse x, recurse y)
  | Zro x -> (match recurse x with VPair (l, r) -> l | _ -> assert false)
  | Fst x -> (match recurse x with VPair (l, r) -> r | _ -> assert false)
  | Left x -> VLeft (recurse x)
  | Right x -> VRight (recurse x)
  | Match (x, (lv, lb), (rv, rb)) ->
      (match recurse x with
      | VLeft l -> eval (ENV.add lv l e) c lb
      | VRight r -> eval (ENV.add lv r e) c rb
      | _ -> assert false)
  | Abs (v, b) ->
      VFunc (fun p -> eval (ENV.add v p e) c b)
  | App (f, x) ->
      (match recurse f with VFunc f_ -> f_ (recurse x) | _ -> assert false)
  | Count -> let v = !c in c := (v + 1); VInt v
  | Unit -> VUnit
  | MkRef x -> VRef (ref (recurse x))
  | GetRef r -> (match recurse r with VRef r_ -> !r_ | _ -> assert false)
  | SetRef (r, x) -> (match recurse r with VRef r_ -> r_ := recurse x; VUnit | _ -> assert false)

let genCounter () =
  let cnt = ref 0 in
  let gen () =
    let ret = !cnt in
    cnt := ret + 1 ;
    ret
  in
  gen

let freshVar = genCounter ()

let freshCounter = genCounter ()

let freshName () = "x" ^ string_of_int (freshVar())

type letlist = (string * expr) list ref

let new_letlist () : letlist = ref []

let push_letlist (ll : letlist) (x : expr) : string =
  let name = freshName () in
  ll := (name, x) :: !ll;
  name

let let_letlist (ll : letlist) (x : expr) : expr = List.fold_left (fun x (n, v) -> Let (n, v, x)) x !ll

let with_letlist (f : letlist -> expr) : expr =
  let ll = new_letlist () in
  let_letlist ll (f ll)

let with_letlist' (f : letlist -> (expr * 'a)) : (expr * 'a) =
  let ll = new_letlist () in
  let (e, a) = f ll in
  (let_letlist ll e, a)

type sValue =
| SInt of int
| SPair of (pValue * pValue)
| SLeft of pValue
| SRight of pValue
| SFunc of (letlist -> pValue -> world -> (pValue * world))
| SUnit
| SRef of int
and pValue = {static: sValue option; dyn: string}
and world = {counter: int option; store: pValue STORE.t}

let mergeCounter a b =
  match (a, b) with
  | Some a_, Some b_ -> if a_ == b_ then Some a_ else None
  | _ -> None

let mergeStore a b =
  STORE.merge (fun _ av bv ->
    match (av, bv) with
    | (Some apv, Some bpv) -> if apv.dyn == bpv.dyn then av else None
    | _ -> None) a b

let mergeWorld a b = {counter = mergeCounter a.counter b.counter; store = mergeStore a.store b.store}

let emptyWorld = {counter = None; store = STORE.empty}

let static ll s d = {static = Some s; dyn = push_letlist ll d}

let dynamic ll d = {static = None; dyn = push_letlist ll d}

let pInt ll i = static ll (SInt i) (Int i)

let pPair ll x y = static ll (SPair (x, y)) (Pair (Var x.dyn, Var y.dyn))

let pLeft ll x = static ll (SLeft x) (Left (Var x.dyn))

let pRight ll x = static ll (SRight x) (Right (Var x.dyn))

let pFunc ll f =
  let name = freshName () in
  static ll (SFunc f) (Abs (name, with_letlist(fun ll -> Var (fst (f ll (dynamic ll (Var name)) emptyWorld)).dyn)))

let pUnit ll = static ll SUnit Unit

let rec pEval (e : pValue ENV.t) (ll : letlist) : expr -> world -> (pValue * world) =
  let recurse = fun expr -> pEval e ll expr in
  let return x = fun s -> (x, s) in
  let ( >>= ) x f = fun s -> let (x', s') = x s in f x' s' in
  let ( >> ) a b = a >>= (fun _ -> b) in
  let get = fun s -> (s, s) in
  let put s = fun _ -> ((), s) in
  let dynS = fun (a, s) -> (Var a.dyn, s) in
  function
  | Var v -> return (ENV.find v e)
  | Int i -> return (pInt ll i)
  | Add (x, y) ->
      recurse x >>= (fun px ->
      recurse y >>= (fun py ->
      return
        (match (px.static, py.static) with
        | Some (SInt xi), Some (SInt yi) -> pInt ll (xi + yi)
        | Some (SInt 0), y -> py
        | x, Some (SInt 0) -> px
        | _ -> dynamic ll (Add (Var px.dyn, Var py.dyn)))))
  | Mult (x, y) ->
      recurse x >>= (fun px ->
      recurse y >>= (fun py ->
      return
        (match (px.static, py.static) with
        | Some (SInt xi), Some (SInt yi) -> pInt ll (xi * yi)
        | Some (SInt 0), _ -> pInt ll 0
        | _, Some (SInt 0) -> pInt ll 0
        | Some (SInt 1), y -> py
        | x, Some (SInt 1) -> px
        | _ -> dynamic ll (Mult (Var px.dyn, Var py.dyn)))))
  | Let (var, v, body) ->
      recurse v >>= (fun pv ->
        pEval (ENV.add var pv e) ll body)
  | IfZero (c, l, r) ->
      recurse c >>= (fun pc ->
      (match pc.static with
      | Some (SInt cv) -> if cv = 0 then recurse l else recurse r
      | None ->
          get >>= (fun c ->
          let (le, lw) = with_letlist' (fun ll -> dynS (pEval e ll l c)) in
          let (re, rw) = with_letlist' (fun ll -> dynS (pEval e ll r c)) in
          put (mergeWorld lw rw) >>
          return (dynamic ll (IfZero (Var pc.dyn, le, re))))
      | _ -> assert false))
  | Pair (x, y) ->
      recurse x >>= (fun px ->
      recurse y >>= (fun py ->
      return (pPair ll px py)))
  | Zro x ->
      recurse x >>= (fun px ->
      return
        (match px.static with
        | Some (SPair (l, r)) -> l
        | None -> dynamic ll (Zro (Var px.dyn))
        | _ -> assert false))
  | Fst x ->
      recurse x >>= (fun px ->
      return
        (match px.static with
        | Some (SPair (l, r)) -> r
        | None -> dynamic ll (Fst (Var px.dyn))
        | _ -> assert false))
  | Left x -> recurse x >>= (fun px -> return (pLeft ll px))
  | Right x -> recurse x >>= (fun px -> return (pRight ll px))
  | Match (x, (lv, lb), (rv, rb)) ->
      recurse x >>= (fun px ->
      (match px.static with
      | Some (SLeft l) -> pEval (ENV.add lv l e) ll lb
      | Some (SRight r) -> pEval (ENV.add lv r e) ll rb
      | None ->
          get >>= (fun c ->
          let (le, lw) = with_letlist' (fun ll -> dynS (pEval (ENV.add lv (dynamic ll (Var lv)) e) ll lb c)) in
          let (re, rw) = with_letlist' (fun ll -> dynS (pEval (ENV.add rv (dynamic ll (Var rv)) e) ll rb c)) in
          put (mergeWorld lw rw) >>
          return (dynamic ll (Match (Var px.dyn, (lv, le), (rv, re)))))
      | _ -> assert false))
  | Abs (v, b) ->
      return (pFunc ll (fun ll p -> pEval (ENV.add v p e) ll b))
  | App (f, x) ->
      recurse f >>= (fun pf ->
      recurse x >>= (fun px ->
      (match pf.static with
      | Some (SFunc f_) -> f_ ll px
      | None ->
          put emptyWorld >>
          return (dynamic ll (App (Var pf.dyn, Var px.dyn)))
      | _ -> assert false)))
  | Count ->
      get >>= (fun s ->
      (match s.counter with
      | Some i -> put {counter = Some (i + 1); store = s.store} >> return (pInt ll i)
      | None -> return (dynamic ll Count)))
  | Unit -> return (pUnit ll)
  | MkRef x ->
      recurse x >>= (fun px ->
      let storeId = freshVar () in
      get >>= (fun s ->
      put {counter = s.counter; store = STORE.add storeId px s.store} >>
      return (static ll (SRef 1) (MkRef (Var px.dyn)))))
  | GetRef r ->
      recurse r >>= (fun pr ->
      (match pr.static with
      | Some (SRef storeId) ->
          get >>= (fun s ->
          return
            (match STORE.find_opt storeId s.store with | Some x -> x | None -> dynamic ll (GetRef (Var pr.dyn))))
      | None -> return (dynamic ll (GetRef (Var pr.dyn)))
      | _ -> assert false))
  | SetRef (r, x) ->
      recurse r >>= (fun pr ->
      recurse x >>= (fun px ->
      (match pr.static with
      | Some (SRef storeId) ->
          get >>= (fun s ->
          put {counter = s.counter; store = STORE.add storeId px s.store} >>
          return (pUnit ll))
      | None ->
          get >>= (fun s ->
          put {counter = s.counter; store = STORE.empty} >>
          return (pUnit ll))
      | _ -> assert false)))