type origin = Memory | Buffer
type event = Idle | Drain of int * int * int64
  | Issue of int * int * Syntax.instruction * origin option
type t = {
  pc : int array;
  regs : int64 array array;
  memory : int64 array;
  buffers : (int * int64) list array;
  critical : bool array;
}

let initial (case : Syntax.t) = {
  pc = [|0; 0|]; regs = Array.init 2 (fun _ -> Array.make 4 0L);
  memory = Array.make (Array.length case.names) 0L;
  buffers = [|[]; []|]; critical = [|false; false|];
}

let copy state = {
  pc = Array.copy state.pc; regs = Array.map Array.copy state.regs;
  memory = Array.copy state.memory; buffers = Array.copy state.buffers;
  critical = Array.copy state.critical;
}

let step (case : Syntax.t) memory_model fences state actor =
  let open Syntax in
  if actor < 0 || actor > 4 then invalid_arg "invalid scheduler choice";
  if actor = 4 then Some (state, Idle) else
  let t = actor mod 2 in
  let out = copy state in
  if actor >= 2 then
    match state.buffers.(t) with
    | [] -> None
    | (a, v) :: tail ->
        out.buffers.(t) <- tail;
        out.memory.(a) <- v;
        Some (out, Drain (t, a, v))
  else
    let pc = state.pc.(t) in
    if pc >= Array.length case.threads.(t) ||
       (List.mem (t, pc) fences && state.buffers.(t) <> []) then None
    else begin
      let ins = case.threads.(t).(pc) in
      out.pc.(t) <- pc + 1;
      let done_ origin = Some (out, Issue (t, pc, ins, origin)) in
      match ins with
      | Store (a, v) ->
          if List.length state.buffers.(t) = 8 then None
          else begin
            if memory_model = Sc then out.memory.(a) <- v
            else out.buffers.(t) <- state.buffers.(t) @ [a, v];
            done_ None
          end
      | Load (a, r) ->
          let value, origin = List.fold_left (fun (value, origin) (address, stored) ->
            if address = a then stored, Buffer else value, origin)
            (state.memory.(a), Memory) state.buffers.(t) in
          out.regs.(t).(r) <- value;
          done_ (Some origin)
      | Fence -> if state.buffers.(t) <> [] then None else done_ None
      | Move (r, v) -> out.regs.(t).(r) <- v; done_ None
      | Branch (r, v, target) ->
          if state.regs.(t).(r) = v then out.pc.(t) <- target;
          done_ None
      | Enter -> out.critical.(t) <- true; done_ None
      | Leave -> out.critical.(t) <- false; done_ None
      | Exchange (a, r, v) ->
          if state.buffers.(t) <> [] then None else begin
            out.regs.(t).(r) <- state.memory.(a);
            out.memory.(a) <- v;
            done_ None
          end
    end

let replay case model fences actors =
  let states = Array.make (Array.length actors + 1) (initial case) in
  let events = Array.make (Array.length actors) Idle in
  Array.iteri (fun i actor ->
    match step case model fences states.(i) actor with
    | None -> invalid_arg (Printf.sprintf "disabled transition at step %d" i)
    | Some (state, event) -> states.(i + 1) <- state; events.(i) <- event) actors;
  states, events

let observations state =
  List.map Int64.of_int (Array.to_list state.pc) @
  Array.to_list state.regs.(0) @ Array.to_list state.regs.(1) @
  Array.to_list state.memory @
  Array.to_list (Array.map (fun queue -> Int64.of_int (List.length queue)) state.buffers) @
  Array.to_list (Array.map (fun active -> if active then 1L else 0L) state.critical)

let bad (case : Syntax.t) state =
  match case.goal with
  | Syntax.Overlap -> state.critical.(0) && state.critical.(1)
  | Syntax.Outcome goal ->
      Array.for_all (fun t -> state.pc.(t) = Array.length case.threads.(t)) [|0; 1|] &&
      List.for_all (fun (t, r, v) -> state.regs.(t).(r) = v) goal.registers &&
      List.for_all (fun (a, v) -> state.memory.(a) = v) goal.memory &&
      (not goal.drained || Array.for_all ((=) []) state.buffers)

let reachable case model fences bound =
  let start = initial case in
  let todo = Queue.create () in
  let seen = Hashtbl.create 512 in
  Queue.add (start, 0) todo;
  Hashtbl.add seen start ();
  let rec search () =
    if Queue.is_empty todo then false, Hashtbl.length seen else
    let state, depth = Queue.take todo in
    if bad case state then true, Hashtbl.length seen else begin
      if depth < bound then
        for actor = 0 to 3 do
          match step case model fences state actor with
          | Some (next, _) when not (Hashtbl.mem seen next) ->
              Hashtbl.add seen next ();
              Queue.add (next, depth + 1) todo
          | _ -> ()
        done;
      search ()
    end
  in
  search ()

let reschedule case fences requested =
  let state = ref (initial case) in
  let actors = ref [] in
  let execute actor = match step case Syntax.Tso fences !state actor with
    | None -> false
    | Some (next, _) -> state := next; actors := actor :: !actors; true
  in
  Array.iter (fun actor ->
    if not (execute actor) then begin
      if actor < 2 then begin
        while (!state).buffers.(actor) <> [] do
          if not (execute (actor + 2)) then failwith "cannot drain pending store"
        done
      end;
      if not (execute actor) then ignore (execute 4)
    end) requested;
  Array.of_list (List.rev !actors)

type graph = { states : t array; parents : (int * int) array; edges : int; violations : int list }

let explore ?(limit=250000) ?(visit=fun _ _ -> ()) case model fences =
  let start = initial case in
  let todo = Queue.create () in
  let seen = Hashtbl.create 2048 in
  let states = ref [] in
  let violations = ref [] in
  let edges = ref 0 in
  let insert state parent actor =
    let id = Hashtbl.length seen in
    if id >= limit then failwith "state graph exceeds the exploration limit";
    Hashtbl.add seen state id;
    states := (id, state, parent, actor) :: !states;
    Queue.add (id, state) todo;
    visit state (id + 1)
  in
  insert start (-1) 4;
  while not (Queue.is_empty todo) do
    let id, state = Queue.take todo in
    if bad case state then violations := id :: !violations;
    for actor = 0 to 3 do
      match step case model fences state actor with
      | None -> ()
      | Some (next, _) ->
          incr edges;
          if not (Hashtbl.mem seen next) then insert next id actor
    done
  done;
  let values = Array.make (Hashtbl.length seen) start in
  let parents = Array.make (Array.length values) (-1, 4) in
  List.iter (fun (id, state, parent, actor) ->
    values.(id) <- state; parents.(id) <- parent, actor) !states;
  {states=values; parents; edges= !edges; violations=List.rev !violations}

let counterexample graph =
  match graph.violations with
  | [] -> None
  | id :: _ ->
      let rec path id actors =
        let parent, actor = graph.parents.(id) in
        if parent < 0 then Array.of_list actors else path parent (actor :: actors)
      in
      Some (path id [])