Trx.init_times ()

type term = Var of string
          |  Int of int
          | Op of term * string * term

let lift x = .<x>.
let rec find v ls = match ls with
                 [] -> None
               | ((n,t)::bs) -> if v = n then Some t 
                                         else find v bs
let rec subst1 sigma pat =
      match pat with
       Var v -> (match (find v sigma) with
                  Some w -> w)
      |Int i -> Int i
      |Op(t1,s,t2) -> Op(subst1 sigma t1, s, subst1 sigma t2)
let rec matchK1 pat k msigma t =
       match (msigma) with
         None -> k None
       | Some (sigma) ->
       (match pat with
           Var u -> (match (find u sigma) with
                          None ->k( Some((u,t)::sigma))
                        | Some w -> if w = t then k(Some sigma)
                                             else k None)
               | Int n -> (match t with
                            Int u -> if u=n then k msigma
                                            else k None
                          | _ -> k None)
               | Op(t11,s1,t12) ->
                     (match t with
                       Op(t21,s2,t22) -> (if s2=s1 then
                                             (matchK1 t11 
                                             (fun msig -> matchK1 t12 k msig t22)
                                              msigma t21)
                                          else k None)
                      | _ ->k None))   
let rec substK2 sigma pat =
    match pat with
     Var v -> (match (find v sigma) with
                          Some w -> w)
   | Int i -> .< Int i>.
   | Op(t1,s,t2) -> .<Op(.~(substK2 sigma t1), s, .~(substK2 sigma t2))>.

let rewriteK1 (lhs,rhs) =
 function t -> (matchK1 lhs (function None -> t 
                               | Some s ->  (subst1 s rhs))
                   (Some []) t)

let rec matchK2 pat k msigma t =
       match (msigma) with
         None -> k None
       | Some (sigma) ->
       (match pat with
           Var u -> (match (find u sigma) with
                          None -> k( Some((u,t)::sigma))
                        | Some w -> .<if .~w = .~t then .~(k(Some sigma))
                                             else .~(k None)>.)
           | Int n -> .<(match .~t with
                            Int u -> if u= .~(lift n) then .~(k msigma)
                                            else .~(k None)
                          | _ -> .~(k None))>.
           | Op(t11,s1,t12) ->
                     .<(match .~t with
                       Op(t21,s2,t22) -> (if s2= .~(lift s1) then
                                          .~(matchK2 t11 
                                             (fun msig -> matchK2 t12 k msig 
                                                           .<t22>.)
                                              msigma .<t21>.)
                                          else .~(k None))
                       | _ -> .~(k None))>.) 

let rewriteK2 (lhs,rhs) =
      .<fun t -> .~(matchK2 lhs 
                     (function None -> .<t>.
                              | Some s ->(substK2 s rhs))
                      (Some []) .<t>.)>.

let t1 = Op(Op(Var"a","+",Var"b"),"+",Op(Var"d","+",Var"e"))
let r1 = (Op(Var "x","+",Int 0), Var "x")
let r2 = (Op(Int 0,"+",Var "x"), Var "x")
let r3 = (Op(Var "x","+",Op(Var "y","+",Var "z")),
          Op(Op(Var "x","+",Var "y"),"+",Var"z"))
let rx = (Op(Op(Op(Op(Var "x","*",Var "p"),"+",Op(Var "x","*",Var "q")),"+",
             Op(Var "x","*",Var "r")),"+",Op(Var "x","*",Var "s")),
          Op(Var "x","*",Op(Op(Var "p","+",Var "q"),"+",
                            Op(Var"r","+",Var "s"))))
let tn = Op(Op(Op(Op(Var "z","*",Var "a"),"+",Op(Var "z","*",Var "b")),"+",
             Op(Var "z","*",Var "d")),"+",Op(Var "z","*",Var "c"))
let t2 = (Op(Op(Op(Var "a","+",Int 0),"+",Int 0),"+",Int 0))

let multi f r t =f r (f r (f r t)) 
let multi' f r = .< fun t -> .~(f r) (.~(f r) (.~(f r) t))>. 
let unstage = Trx.timenew "unstaged running"
              (fun () -> rewriteK1 r3 t1)

let st1Run = Trx.timenew "stage 1 running"
              (fun () -> (rewriteK2 r3))  
let comp =
    Trx.timenew "compiling"
         (fun () -> .! st1Run)
let st2Run = Trx.timenew "stage 2 running"
         (fun () -> (comp t1))
let _ = Trx.print_times ()
