type instruction =
  | Store of int * int64 | Load of int * int | Fence
  | Move of int * int64 | Branch of int * int64 * int
  | Enter | Leave | Exchange of int * int * int64

type goal = Overlap | Outcome of {
  registers : (int * int * int64) list;
  memory : (int * int64) list;
  drained : bool;
}

type model = Sc | Tso
type location = int * int
type t = {
  name : string;
  names : string array;
  bound : int;
  threads : instruction array array;
  candidates : location list;
  goal : goal;
}

let model = function "sc" -> Sc | "tso" -> Tso | _ -> invalid_arg "unknown memory model"
let model_name = function Sc -> "sc" | Tso -> "tso"
let location value = match Sexp.list value with
  | [t; pc] -> Sexp.integer t, Sexp.integer pc
  | _ -> invalid_arg "invalid fence location"

let instruction value =
  let open Sexp in
  match value with
  | Store (a, v) -> List [Atom "store"; i a; w v]
  | Load (a, r) -> List [Atom "load"; i a; i r]
  | Fence -> List [Atom "fence"]
  | Move (r, v) -> List [Atom "move"; i r; w v]
  | Branch (r, v, target) -> List [Atom "branch"; i r; w v; i target]
  | Enter -> List [Atom "enter"]
  | Leave -> List [Atom "leave"]
  | Exchange (a, r, v) -> List [Atom "exchange"; i a; i r; w v]

let decode_instruction value =
  let open Sexp in
  match list value with
  | [Atom "store"; a; v] -> Store (integer a, word v)
  | [Atom "load"; a; r] -> Load (integer a, integer r)
  | [Atom "fence"] -> Fence
  | [Atom "move"; r; v] -> Move (integer r, word v)
  | [Atom "branch"; r; v; target] -> Branch (integer r, word v, integer target)
  | [Atom "enter"] -> Enter
  | [Atom "leave"] -> Leave
  | [Atom "exchange"; a; r; v] -> Exchange (integer a, integer r, word v)
  | _ -> invalid_arg "invalid instruction"

let validate case =
  let check condition message = if not condition then invalid_arg message in
  check (Array.length case.threads = 2) "exactly two threads required";
  check (Array.length case.names > 0 && Array.length case.names <= 16) "address count outside 1..16";
  check (case.bound > 0 && case.bound <= 64) "step bound outside 1..64";
  let address a = check (a >= 0 && a < Array.length case.names) "address outside model" in
  let register r = check (r >= 0 && r < 4) "register outside model" in
  Array.iter (fun program ->
    check (Array.length program > 0 && Array.length program <= 64) "thread length outside 1..64";
    Array.iter (function
      | Store (a, _) -> address a
      | Load (a, r) | Exchange (a, r, _) -> address a; register r
      | Move (r, _) -> register r
      | Branch (r, _, target) -> register r;
          check (target >= 0 && target <= Array.length program) "branch target outside thread"
      | Fence | Enter | Leave -> ()) program) case.threads;
  let thread t = check (t >= 0 && t < 2) "thread outside model" in
  List.iter (fun (t, pc) -> thread t;
    check (pc >= 0 && pc < Array.length case.threads.(t)) "fence outside thread") case.candidates;
  check (List.length case.candidates = List.length (List.sort_uniq compare case.candidates)) "duplicate fence";
  match case.goal with
  | Overlap -> ()
  | Outcome goal ->
      List.iter (fun (t, r, _) -> thread t; register r) goal.registers;
      List.iter (fun (a, _) -> address a) goal.memory

let singleton = function [value] -> value | _ -> invalid_arg "expected one field value"

let fields allowed values =
  List.iter (function
    | Sexp.List (Sexp.Atom name :: _) when List.mem name allowed -> ()
    | Sexp.List (Sexp.Atom name :: _) -> invalid_arg ("unknown field " ^ name)
    | _ -> invalid_arg "malformed field") values

let decode value =
  let open Sexp in
  let entries = match value with
    | List (Atom "case" :: fields) -> fields
    | _ -> invalid_arg "expected case"
  in
  fields ["name"; "names"; "bound"; "threads"; "candidates"; "goal"] entries;
  let get key = singleton (field key entries) in
  let goal = match field "goal" entries with
    | [Atom "overlap"] -> Overlap
    | [List (Atom "outcome" :: entries)] ->
        fields ["registers"; "memory"; "drained"] entries;
        let registers = List.map (fun value -> match list value with
          | [t; r; v] -> integer t, integer r, word v
          | _ -> invalid_arg "invalid register goal") (field "registers" entries) in
        let memory = List.map (fun value -> match list value with
          | [a; v] -> integer a, word v
          | _ -> invalid_arg "invalid memory goal") (field "memory" entries) in
        let drained = match field "drained" entries with
          | [Atom "true"] -> true | [Atom "false"] -> false
          | _ -> invalid_arg "invalid drain flag" in
        Outcome {registers; memory; drained}
    | _ -> invalid_arg "invalid goal"
  in
  let case = {
    name = atom (get "name");
    names = Array.of_list (List.map atom (field "names" entries));
    bound = integer (get "bound");
    threads = Array.of_list (List.map (fun thread ->
      Array.of_list (List.map decode_instruction (list thread))) (field "threads" entries));
    candidates = List.map location (field "candidates" entries);
    goal;
  } in
  validate case;
  case

let encode case =
  let open Sexp in
  let goal = match case.goal with
    | Overlap -> Atom "overlap"
    | Outcome goal -> List [Atom "outcome";
        List (Atom "registers" :: List.map (fun (t, r, v) -> List [i t; i r; w v]) goal.registers);
        List (Atom "memory" :: List.map (fun (a, v) -> List [i a; w v]) goal.memory);
        List [Atom "drained"; Atom (string_of_bool goal.drained)]] in
  List [Atom "case"; List [Atom "name"; Atom case.name];
    List (Atom "names" :: Array.to_list (Array.map a case.names));
    List [Atom "bound"; i case.bound];
    List (Atom "threads" :: Array.to_list (Array.map (fun thread ->
      List (Array.to_list (Array.map instruction thread))) case.threads));
    List (Atom "candidates" :: List.map (fun (t, pc) -> List [i t; i pc]) case.candidates);
    List [Atom "goal"; goal]]

let load path = decode (Sexp.one (Store.read path))