INCLUDE "debug.ml"
INCLUDE "utils.ml"
INCLUDE "trace.ml"

open Format

type instr =
  | SELF of unit
  | LEFT of State.t
  | RIGHT of State.t

type opcode =
  | OP_NOP of unit
  | OP_LEFT1 of State.t
  | OP_LEFT2 of State.t * State.t
  | OP_RIGHT1 of State.t
  | OP_RIGHT2 of State.t * State.t
  | OP_LEFT1_RIGHT1 of State.t * State.t
  | OP_LEFT2_RIGHT1 of State.t * State.t * State.t
  | OP_LEFT1_RIGHT2 of State.t * State.t * State.t
  | OP_LEFT2_RIGHT2 of State.t * State.t * State.t * State.t
  | OP_SELF of unit
  | OP_SELF_LEFT1 of State.t
  | OP_SELF_LEFT2 of State.t * State.t
  | OP_SELF_RIGHT1 of State.t
  | OP_SELF_RIGHT2 of State.t * State.t
  | OP_SELF_LEFT1_RIGHT1 of State.t * State.t
  | OP_SELF_LEFT2_RIGHT1 of State.t * State.t * State.t
  | OP_SELF_LEFT1_RIGHT2 of State.t * State.t * State.t
  | OP_SELF_LEFT2_RIGHT2 of State.t * State.t * State.t * State.t
  | OP_OTHER of instr array

type code = Nil | Cons of State.t * opcode * code

let rec length l =
  match l with
      Nil -> 0
    | Cons(_, _, t) -> 1 + length t
let debug fmt l =
  fprintf fmt "length of code is %i\n%!" (length l)


let print_instr fmt i =
  match i with
    | SELF _ -> fprintf fmt "SELF"
    | LEFT q -> fprintf fmt "LEFT{%a}" State.print q
    | RIGHT q -> fprintf fmt "RIGHT{%a}" State.print q

let print_opcode fmt code =
  match code with
    | OP_NOP _ -> fprintf fmt "OP_NOP"

    | OP_LEFT1 src ->
	fprintf fmt "OP_LEFT1{%a}" State.print src

    | OP_LEFT2 (src1, src2) ->
	fprintf fmt "OP_LEFT2{%a, %a}" State.print src1 State.print src2

    | OP_RIGHT1 src ->
	fprintf fmt "OP_RIGHT1{%a}" State.print src

    | OP_RIGHT2 (src1, src2) ->
	fprintf fmt "OP_RIGHT2{%a, %a}" State.print src1 State.print src2

    | OP_LEFT1_RIGHT1 (src1, src2) ->
	fprintf fmt "OP_LEFT1_RIGHT1{%a}{%a}" State.print src1 State.print src2

    | OP_LEFT2_RIGHT1 (src1, src2, src3) ->
	fprintf fmt "OP_LEFT2_RIGHT1{%a, %a}{%a}"
	  State.print src1 State.print src2 State.print src3

    | OP_LEFT1_RIGHT2 (src1, src2, src3) ->
	fprintf fmt "OP_LEFT1_RIGHT2{%a}{%a, %a}"
	  State.print src1 State.print src2 State.print src3

    | OP_LEFT2_RIGHT2 (src1, src2, src3, src4) ->
	fprintf fmt "OP_LEFT2_RIGHT2{%a, %a}{%a, %a}"
	  State.print src1 State.print src2 State.print src3 State.print src4

    | OP_SELF _ ->
	fprintf fmt "OP_SELF"

    | OP_SELF_LEFT1 src ->
	fprintf fmt "OP_SELF_LEFT1{%a}" State.print src

    | OP_SELF_LEFT2 (src1, src2) ->
	fprintf fmt "OP_SELF_LEFT2{%a, %a}" State.print src1 State.print src2

    | OP_SELF_RIGHT1 src ->
	fprintf fmt "OP_SELF_RIGHT1{%a}" State.print src

    | OP_SELF_RIGHT2 (src1, src2) ->
	fprintf fmt "OP_SELF_RIGHT2{%a, %a}" State.print src1 State.print src2

    | OP_SELF_LEFT1_RIGHT1 (src1, src2) ->
	fprintf fmt "OP_SELF_LEFT1_RIGHT1{%a}{%a}" State.print src1 State.print src2

    | OP_SELF_LEFT2_RIGHT1 (src1, src2, src3) ->
	fprintf fmt "OP_SELF_LEFT2_RIGHT1{%a, %a}{%a}"
	  State.print src1 State.print src2 State.print src3

    | OP_SELF_LEFT1_RIGHT2 (src1, src2, src3) ->
	fprintf fmt "OP_SELF_LEFT1_RIGHT2{%a}{%a, %a}"
	  State.print src1 State.print src2 State.print src3

    | OP_SELF_LEFT2_RIGHT2 (src1, src2, src3, src4) ->
	fprintf fmt "OP_SELF_LEFT2_RIGHT2{%a, %a}{%a, %a}"
	  State.print src1 State.print src2 State.print src3 State.print src4
    | OP_OTHER line ->
	fprintf fmt "OP_OTHER: ";
	Array.iter (fun i -> print_instr fmt i; fprintf fmt " ") line

let merge_rev equal choose l =
  match l with
    | [] -> l
    | x :: ll ->
	List.fold_left
	  (fun acc i ->
	     let j = List.hd acc in
	       if equal i j then (choose i j)::(List.tl acc)
	       else i::acc) [x] ll

let compile_instr_list l =
  let linstr = merge_rev (=) (fun i _ -> i) (List.sort (fun x y -> compare y x) l) in
    match linstr with
	[] -> OP_NOP()
      | [ LEFT q ] -> OP_LEFT1 q
      | [ LEFT q1; LEFT q2 ] -> OP_LEFT2(q2, q1)
      | [ RIGHT q ] -> OP_RIGHT1 q
      | [ RIGHT q1; RIGHT q2 ] -> OP_RIGHT2(q2, q1)
      | [ LEFT q1; RIGHT q2 ] -> OP_LEFT1_RIGHT1(q1, q2)
      | [ LEFT q1; LEFT q2; RIGHT q3 ] -> OP_LEFT2_RIGHT1 (q2, q1, q3)
      | [ LEFT q1; RIGHT q2; RIGHT q3 ] -> OP_LEFT1_RIGHT2 (q1, q3, q2)
      | [ LEFT q1; LEFT q2; RIGHT q3; RIGHT q4 ] -> OP_LEFT2_RIGHT2 (q2, q1, q4, q3)
      | [ SELF () ] -> OP_SELF()

      | [ SELF _; LEFT q ] -> OP_SELF_LEFT1 q
      | [ SELF _; LEFT q1; LEFT q2 ] -> OP_SELF_LEFT2(q2, q1)
      | [ SELF _; RIGHT q ] -> OP_SELF_RIGHT1 q
      | [ SELF _; RIGHT q1; RIGHT q2 ] -> OP_SELF_RIGHT2(q2, q1)
      | [ SELF _; LEFT q1; RIGHT q2 ] -> OP_SELF_LEFT1_RIGHT1(q1, q2)
      | [ SELF _; LEFT q1; LEFT q2; RIGHT q3 ] -> OP_SELF_LEFT2_RIGHT1 (q2, q1, q3)
      | [ SELF _; LEFT q1; RIGHT q2; RIGHT q3 ] -> OP_SELF_LEFT1_RIGHT2 (q1, q3, q2)
      | [ SELF _; LEFT q1; LEFT q2; RIGHT q3; RIGHT q4 ] ->
	  OP_SELF_LEFT2_RIGHT2 (q2, q1, q4, q3)
      | i -> OP_OTHER (Array.of_list i)


let to_list l =
  let rec loop l acc =
    match l with
	[] -> acc
      | (a, b)::ll -> loop ll (Cons(a,b, acc))
  in loop l Nil


let rec filter_uniq statel stater l =
  match l with
      [] -> []
    | (s, il)::ll ->
	let nil, nsl, nsr =
	  List.fold_left
	    (fun ((a_il, al, ar)as acc) i ->
	       match i with
		 | LEFT q ->
		     if List.mem q al then acc
		     else (i :: a_il, q::al, ar)
		 | RIGHT q ->
		     if List.mem q ar then acc
		     else (i :: a_il, al, q :: ar)
		 | _ -> (i :: a_il, al, ar)) ([], statel, stater) il
	in
	  (s, nil) :: (filter_uniq nsl nsr ll)

let compile l =
  let l = List.sort (fun (s1, _) (s2, _) -> compare s1 s2) l in
  let l = filter_uniq [] [] l in
  let l = merge_rev
    (fun (s1, _) (s2, _) -> s1 = s2)
    (fun (s1, i1) (_, i2) -> (s1, i1@i2)) l
  in
 let marking =
    List.exists
      (fun (_, l) -> List.exists (function SELF _ -> true | _ -> false) l)
      l
  in
  let l = List.map (fun (s, il) -> (s, compile_instr_list il)) l in
  let l = List.filter (fun (_, instr) -> instr <> OP_NOP ()) l in
    to_list l, not marking


type 'a update = 'a -> 'a -> 'a -> Tree.t -> Tree.node -> StateSet.t * 'a
type 'a cache = 'a update Cache.Lvl3.t

let dummy_update = fun _ _ _ _ _ -> failwith "Uninitialized L3JIT"
let show_stats (a : 'a cache) =
  let count = ref 0 in
  Cache.Lvl3.iteri (fun _ _ _ _ b -> if not b then incr count) a;
  eprintf "%!L3JIT: %i used entries\n%!" !count

let create () =
  let v = Cache.Lvl3.create 1024 dummy_update in
  if !Options.verbose then at_exit (fun () -> show_stats v);
  v

let find (t : 'a cache) tlist s1 s2 =
  Cache.Lvl3.find t
    (Uid.to_int s2.StateSet.Node.id)
    (Uid.to_int s1.StateSet.Node.id)
    (Uid.to_int tlist.Translist.Node.id)

let add (t : 'a cache) tlist s1 s2 v =
  Cache.Lvl3.add t
    (Uid.to_int s2.StateSet.Node.id)
    (Uid.to_int s1.StateSet.Node.id)
    (Uid.to_int tlist.Translist.Node.id)
    v

let eval_form auto s1 s2 f =
  let rec loop f =
    match Formula.expr f with
    | Formula.False | Formula.True | Formula.Pred _ -> f, []
    | Formula.Atom(`Left, b, q) ->
      Formula.of_bool (b == (StateSet.mem q s1)),
      if b && StateSet.mem q auto.Ata.topdown_marking_states then [LEFT q] else []
    | Formula.Atom (`Right, b, q) ->
      Formula.of_bool(b == (StateSet.mem q s2)),
      if b && StateSet.mem q auto.Ata.topdown_marking_states then [RIGHT q] else []
    | Formula.Atom (`Epsilon, _, _) -> assert false

    | Formula.Or(f1, f2) ->
      let b1, i1 = loop f1 in
      let b2, i2 = loop f2 in
      Formula.or_pred b1 b2, i1 @ i2
    | Formula.And(f1, f2) ->
      let b1, i1 = loop f1 in
      let b2, i2 = loop f2 in
      Formula.and_pred b1 b2, i1 @ i2
  in
  loop f

let eval_trans auto s1 s2 trans =
  Translist.fold
    (fun t ((a_st, a_op, a_todo) as acc)->
      let q, _, m, f = Transition.node t in
      let form, ops = eval_form auto s1 s2 f in
      match Formula.expr form with
      | Formula.True ->
	StateSet.add q a_st,
	(q, (if m then (SELF() :: ops) else ops)):: a_op,
	a_todo
      | Formula.False -> acc
      | Formula.Pred p -> a_st, a_op,
	(p.Tree.Predicate.node, q, [(q,(if m then (SELF() :: ops) else ops))]) :: a_todo
      | _ -> assert false
    ) trans (StateSet.empty, [], [])

let compile_update auto trl s1 s2 =
  let orig_s1, orig_s2 =
    Translist.fold (fun t (a1, a2) ->
      let _, _, _, f = Transition.node t in
      let  fs1, fs2 = Formula.st f in
      (StateSet.union a1 fs1, StateSet.union a2 fs2)
    ) trl (StateSet.empty, StateSet.empty)
  in
  let ns1 = StateSet.inter s1 orig_s1
  and ns2 = StateSet.inter s2 orig_s2 in
  let res, ops, todo = eval_trans auto ns1 ns2 trl in
  let code, not_marking = compile ops in
  let todo_code, todo_notmarking =
    List.fold_left (fun (l, b) (p, q, o) -> let c, b' = compile o in
					    (p, q, c)::l, b && b')
	      ([], not_marking) todo
  in
  let opcode = res, code, todo_notmarking, todo_code in
  opcode

let gen_code exec auto tlist s1 s2 =
  let res, code, not_marking, todo_code = compile_update auto tlist s1 s2 in
  let f =
    if todo_code == [] then
      if not_marking then begin fun empty_slot sl1 sl2 _ node ->
	let slot1_empty = sl1 == empty_slot
	and slot2_empty = sl2 == empty_slot in
	if slot1_empty && slot2_empty then res,sl2
	else
	  let sl =
	    if slot2_empty then
	      if slot1_empty then
		Array.copy empty_slot
	      else sl1
	    else sl2
	  in
	  exec sl sl1 sl2 node code;
	  res, sl
      end
      else (* marking *) begin fun empty_slot sl1 sl2 _ node ->
	let sl =
	  if sl2 == empty_slot  then
	    if sl1 == empty_slot then
	      Array.copy empty_slot
	    else sl1
	  else sl2
	in
	exec sl sl1 sl2 node code;
	res, sl
      end
    else (* todo != [] *)
      begin fun empty_slot sl1 sl2 tree node ->
	let sl =
	  if sl2 == empty_slot  then
	    if sl1 == empty_slot then
	      Array.copy empty_slot
	    else sl1
	  else sl2
	in
	exec sl sl1 sl2 node code;
	List.fold_left
	  (fun ares (p, q, code) ->
	    if !p tree node then begin
	      if code != Nil then exec sl sl1 sl2 node code;
	      StateSet.add q ares
	    end
	    else ares) res todo_code, sl

      end
  in
  f





DEFINE SET(a, b) = (a) <- (b)

DEFINE PRINT_TEMPLATE(ns) =
      let pr fmt (state, count) =
	fprintf fmt "%a: %i" State.print state (ns.length count)
      in
      Pretty.print_array ~sep:", " pr fmt (Array.mapi (fun x y -> (x,y)) s)

DEFINE EXEC_INSTR_TEMPLATE(ns) = fun slot1 slot2 t inst acc ->
   match inst with
    | SELF _ ->  ns.snoc acc t
    | LEFT src -> ns.concat acc slot1.(src)
    | RIGHT src -> ns.concat acc slot2.(src)


DEFINE EXEC_CODE_BODY_TEMPLATE(ns) =
  (match code with
    | OP_NOP _ -> ()

    | OP_LEFT1 src ->
      SET(slot.(dst), slot1.(src))

    | OP_LEFT2 (src1, src2) ->
      SET(slot.(dst) , ns.concat slot1.(src1) slot1.(src2))

    | OP_RIGHT1 src -> SET(slot.(dst) , slot2.(src))

    | OP_RIGHT2 (src1, src2) ->
      SET (slot.(dst) , ns.concat slot2.(src1) slot2.(src2) )

    | OP_LEFT1_RIGHT1 (src1, src2) ->
      SET (slot.(dst) , ns.concat slot1.(src1) slot2.(src2))

    | OP_LEFT2_RIGHT1 (src1, src2, src3) ->
      SET (slot.(dst) , ns.concat3 slot1.(src1) slot1.(src2) slot2.(src3))

    | OP_LEFT1_RIGHT2 (src1, src2, src3) ->
      SET (slot.(dst) , ns.concat3 slot1.(src1) slot2.(src2) slot2.(src3));

    | OP_LEFT2_RIGHT2 (src1, src2, src3, src4) ->
        SET (slot.(dst) , ns.concat4 slot1.(src1) slot1.(src2) slot2.(src3) slot2.(src4))

    | OP_SELF _ ->
	slot.(dst) <- ns.singleton t

    | OP_SELF_LEFT1 src -> slot.(dst) <- ns.cons t slot1.(src)

    | OP_SELF_LEFT2 (src1, src2) ->
	slot.(dst) <- ns.conscat t slot1.(src1) slot1.(src2)

    | OP_SELF_RIGHT1 src -> slot.(dst) <- ns.cons t slot2.(src)

    | OP_SELF_RIGHT2 (src1, src2) ->
	slot.(dst) <- ns.conscat t slot2.(src1) slot2.(src2)

    | OP_SELF_LEFT1_RIGHT1 (src1, src2) ->
	slot.(dst) <- ns.conscat t slot1.(src1) slot2.(src2)

    | OP_SELF_LEFT2_RIGHT1 (src1, src2, src3) ->
	slot.(dst) <- ns.conscat3 t slot1.(src1) slot1.(src2) slot2.(src3)

    | OP_SELF_LEFT1_RIGHT2 (src1, src2, src3) ->
	slot.(dst) <- ns.conscat3 t slot1.(src1) slot2.(src2) slot2.(src3)

    | OP_SELF_LEFT2_RIGHT2 (src1, src2, src3, src4) ->
	slot.(dst) <-
	  ns.conscat4 t slot1.(src1) slot1.(src2) slot2.(src3) slot2.(src4)
    | OP_OTHER line -> assert false (*
      let acc = ref ns.empty in
      let len = Array.length line - 1 in
      for j = 0 to len do
	acc := exec_instr slot1 slot2 t line.(j) !acc
      done;
      slot.(dst) <- !acc *) )

DEFINE EXEC_CODE_TEMPLATE(ns) = fun slot slot1 slot2 t dst code ->
  EXEC_CODE_BODY_TEMPLATE(ns)


DEFINE EXEC_REC_TEMPLATE(exec_code) =
	  (match code with
	  | Nil -> ()
	  | Cons(dst, opcode, code1) ->
	    TRACE("res-jit", 3, __ "  %a := %a\n%!"
	      State.print dst print_opcode opcode;
	    );
	    exec_code slot slot1 slot2 t dst opcode;
	    begin
	      match code1 with
	    | Nil -> ()
	    | Cons(dst, opcode, code1) ->
	      TRACE("res-jit", 3, __ "  %a := %a\n%!"
		State.print dst print_opcode opcode;
	      );
	      exec_code slot slot1 slot2 t dst opcode;
	      exec slot slot1 slot2 t code1

	    end)
let count_exec_code slot slot1 slot2 t dst code =
  EXEC_CODE_BODY_TEMPLATE(NodeSet.Count)


let count_exec slot slot1 slot2 t code =
  let rec exec slot slot1 slot2 t code =
    EXEC_REC_TEMPLATE(count_exec_code)
  in
  exec slot slot1 slot2 t code

let mat_exec_code slot slot1 slot2 t dst code =
  EXEC_CODE_BODY_TEMPLATE(NodeSet.Mat)


let mat_exec slot slot1 slot2 t code =
  let rec exec slot slot1 slot2 t code =
    EXEC_REC_TEMPLATE(mat_exec_code)
  in
  exec slot slot1 slot2 t code


DEFINE EXEC_TEMPLATE =
	  (TRACE("res-jit", 3, __ "Node %i:\n" (Node.to_int t));
	   TRACE("res-jit", 3, __ " LEFT  : %a\n" pr_slot slot1);
	   TRACE("res-jit", 3, __ " RIGHT : %a\n" pr_slot slot2);
	   exec slot slot1 slot2 t code;
	   TRACE("res-jit", 3, __ " RES   : %a\n\n%!" pr_slot slot))


DEFINE UPDATE_TEMPLATE =
	  let f = find cache tlist s1 s2 in
	  if f == dummy_update then
	    let f = gen_code exec auto tlist s1 s2 in
	    add cache tlist s1 s2 f;
	    f empty_res sl1 sl2 tree node
	  else
	    f empty_res sl1 sl2 tree node


let update exec cache auto tlist s1 s2 empty_res sl1 sl2 tree node =
  let f = find cache tlist s1 s2 in
  if f == dummy_update then
    let f = gen_code exec auto tlist s1 s2 in
    add cache tlist s1 s2 f;
    f empty_res sl1 sl2 tree node
  else
    f empty_res sl1 sl2 tree node


module type S =
  sig
    module NS : NodeSet.S
    type t = NS.t array
    val exec : t -> t -> t -> Tree.node -> code -> unit
    val update : t cache -> Ata.t -> Translist.t -> StateSet.t -> StateSet.t ->
      t -> t -> t -> Tree.t -> Tree.node -> StateSet.t * t
    val print : Format.formatter -> t -> unit
    val var : int -> t -> t
    val close : ((int*State.t, NS.t) Hashtbl.t) -> t -> t
    val is_open : t -> bool
  end

module Count =
  struct
    module NS = NodeSet.Count
    type t = NodeSet.Count.t array
    let print fmt s = PRINT_TEMPLATE(NS)
    let exec_instr = EXEC_INSTR_TEMPLATE(NodeSet.Count)
    let exec_code = EXEC_CODE_TEMPLATE(NodeSet.Count)
    let rec exec slot slot1 slot2 t code = EXEC_REC_TEMPLATE(exec_code)
    let exec slot slot1 slot2 t code = EXEC_TEMPLATE
    let update cache auto tlist s1 s2 empty_res sl1 sl2 tree node = UPDATE_TEMPLATE
    let var _ x = x
    let close _ x = x
    let is_open _ = false
  end

module Mat =
  struct
    module NS = NodeSet.Mat
    type t = NodeSet.Mat.t array
    let print fmt s = PRINT_TEMPLATE(NS)
    let exec_instr = EXEC_INSTR_TEMPLATE(NodeSet.Mat)
    let exec_code = EXEC_CODE_TEMPLATE(NodeSet.Mat)
    let rec exec slot slot1 slot2 t code = EXEC_REC_TEMPLATE(exec_code)
    let exec slot slot1 slot2 t code = EXEC_TEMPLATE
    let update cache auto tlist s1 s2 empty_res sl1 sl2 tree node = UPDATE_TEMPLATE
    let var _ x = x
    let close _ x = x
    let is_open _ = false
  end



module Make(U : NodeSet.S) =
  struct
    module NS = U
    type t = U.t array
    let print fmt s = PRINT_TEMPLATE(NS)
    let exec_instr = EXEC_INSTR_TEMPLATE(U)
    let exec_code = EXEC_CODE_TEMPLATE(U)
    let rec exec slot slot1 slot2 t code = EXEC_REC_TEMPLATE(exec_code)
    let exec slot slot1 slot2 t code = EXEC_TEMPLATE
    let update cache auto tlist s1 s2 empty_res sl1 sl2 tree node = UPDATE_TEMPLATE
    let var i t =
      Array.mapi (fun j _ -> NS.var (i,j)) t
    let close h t =
      Array.map (NS.close h) t

    let is_open t =
      List.exists NS.is_open (Array.to_list t)
  end