Trx.init_times ()

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

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 match1 pat msigma t =
        match msigma with
         |None -> None
         |Some(sigma) ->
          (match pat with
		 Var u ->(match (find u sigma) with
                          None -> Some((u,t)::sigma)
                        | Some w -> if w = t then Some sigma
                                             else None)
               | Int n -> (match t with
                            Int u -> if u=n then msigma
                                            else None
                          | _ -> None)
               | Op(t11,s1,t12) ->
                     (match t with
                       Op(t21,s2,t22) -> (if s2=s1 then
                                             (match1 t11 
                                              (match1 t12 msigma t22) t21)
                                          else None)
                      | _ -> None))   
let rewrite1(lhs,rhs) t =
     match match1 lhs (Some []) t with
      None -> t
    | Some (sigma) -> subst1 sigma rhs

(*staged version? *)

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

let rewrite2 (lhs,rhs) t=
    .< match .~(match2 lhs .<Some []>. t) with
      None -> .~ t
    | Some (sigma) -> subst1 sigma rhs >.

let t1 = Op(Op(Var"a","+",Var"b"),"+",Op(Var"d","+",Var"e"))
let r3 = (Op(Var "x","+",Op(Var "y","+",Var "z")),
          Op(Op(Var "x","+",Var "y"),"+",Var"z"))

let unstage = Trx.timenew "unstaged running"
              (fun () -> rewrite1 r3 t1)

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