open Core.Std open Qptypes (** New job : Request to create a new multi-tasked job *) module State : sig type t val of_string : string -> t val to_string : t -> string end = struct type t = string let of_string x = x let to_string x = x end module Newjob_msg : sig type t = { state: State.t; address_tcp: Address.Tcp.t ; address_inproc: Address.Inproc.t; } val create : address_tcp:string -> address_inproc:string -> state:string -> t val to_string : t -> string end = struct type t = { state: State.t; address_tcp: Address.Tcp.t ; address_inproc: Address.Inproc.t; } let create ~address_tcp ~address_inproc ~state = { state = State.of_string state; address_tcp = Address.Tcp.of_string address_tcp ; address_inproc = Address.Inproc.of_string address_inproc ; } let to_string t = Printf.sprintf "new_job %s %s %s" ( State.to_string t.state ) ( Address.Tcp.to_string t.address_tcp ) ( Address.Inproc.to_string t.address_inproc ) end module Endjob_msg : sig type t = { state: State.t; } val create : state:string -> t val to_string : t -> string end = struct type t = { state: State.t; } let create ~state = { state = State.of_string state; } let to_string t = Printf.sprintf "end_job %s" ( State.to_string t.state ) end (** Connect : connect a new client to the task server *) module Connect_msg : sig type t = Tcp | Inproc | Ipc val create : typ:string -> t val to_string : t -> string end = struct type t = Tcp | Inproc | Ipc let create ~typ = match typ with | "tcp" -> Tcp | "inproc" -> Inproc | "ipc" -> Ipc | _ -> assert false let to_string = function | Tcp -> "connect tcp" | Inproc -> "connect inproc" | Ipc -> "connect ipc" end (** ConnectReply : Reply to the connect messsage *) module ConnectReply_msg : sig type t = { client_id: Id.Client.t ; state: State.t ; push_address: Address.t; } val create : state:State.t -> client_id:Id.Client.t -> push_address:Address.t -> t val to_string : t -> string end = struct type t = { client_id: Id.Client.t ; state: State.t ; push_address: Address.t; } let create ~state ~client_id ~push_address = { client_id ; state ; push_address } let to_string x = Printf.sprintf "connect_reply %s %d %s" (State.to_string x.state) (Id.Client.to_int x.client_id) (Address.to_string x.push_address) end (** Disconnect : disconnect a client from the task server *) module Disconnect_msg : sig type t = { client_id: Id.Client.t ; state: State.t ; } val create : state:string -> client_id:string -> t val to_string : t -> string end = struct type t = { client_id: Id.Client.t ; state: State.t ; } let create ~state ~client_id = { client_id = Id.Client.of_string client_id ; state = State.of_string state } let to_string x = Printf.sprintf "disconnect %s %d" (State.to_string x.state) (Id.Client.to_int x.client_id) end module DisconnectReply_msg : sig type t = { state: State.t ; } val create : state:State.t -> t val to_string : t -> string end = struct type t = { state: State.t ; } let create ~state = { state } let to_string x = Printf.sprintf "disconnect_reply %s" (State.to_string x.state) end (** AddTask : Add a new task to the queue *) module AddTask_msg : sig type t = { state: State.t; task: string; } val create : state:string -> task:string -> t val to_string : t -> string end = struct type t = { state: State.t; task: string; } let create ~state ~task = { state = State.of_string state ; task } let to_string x = Printf.sprintf "add_task %s %s" (State.to_string x.state) x.task end (** AddTaskReply : Reply to the AddTask message *) module AddTaskReply_msg : sig type t val create : task_id:Id.Task.t -> t val to_string : t -> string end = struct type t = Id.Task.t let create ~task_id = task_id let to_string x = Printf.sprintf "add_task_reply %d" (Id.Task.to_int x) end (** DelTask : Remove a task from the queue *) module DelTask_msg : sig type t = { state: State.t; task_id: Id.Task.t } val create : state:string -> task_id:string -> t val to_string : t -> string end = struct type t = { state: State.t; task_id: Id.Task.t } let create ~state ~task_id = { state = State.of_string state ; task_id = Id.Task.of_string task_id } let to_string x = Printf.sprintf "del_task %s %d" (State.to_string x.state) (Id.Task.to_int x.task_id) end (** DelTaskReply : Reply to the DelTask message *) module DelTaskReply_msg : sig type t val create : task_id:Id.Task.t -> more:bool -> t val to_string : t -> string end = struct type t = { task_id : Id.Task.t ; more : bool; } let create ~task_id ~more = { task_id ; more } let to_string x = let more = if x.more then "more" else "done" in Printf.sprintf "del_task_reply %s %d" more (Id.Task.to_int x.task_id) end (** GetTask : get a new task to do *) module GetTask_msg : sig type t = { client_id: Id.Client.t ; state: State.t ; } val create : state:string -> client_id:string -> t val to_string : t -> string end = struct type t = { client_id: Id.Client.t ; state: State.t ; } let create ~state ~client_id = { client_id = Id.Client.of_string client_id ; state = State.of_string state } let to_string x = Printf.sprintf "get_task %s %d" (State.to_string x.state) (Id.Client.to_int x.client_id) end (** GetTaskReply : Reply to the GetTask message *) module GetTaskReply_msg : sig type t val create : task_id:Id.Task.t option -> task:string option -> t val to_string : t -> string end = struct type t = { task_id: Id.Task.t option ; task : string option ; } let create ~task_id ~task = { task_id ; task } let to_string x = match x.task_id, x.task with | Some task_id, Some task -> Printf.sprintf "get_task_reply %d %s" (Id.Task.to_int task_id) task | _ -> Printf.sprintf "get_task_reply 0" end (** GetPsi : get the current variational wave function *) module GetPsi_msg : sig type t = { client_id: Id.Client.t ; } val create : client_id:string -> t val to_string : t -> string end = struct type t = { client_id: Id.Client.t ; } let create ~client_id = { client_id = Id.Client.of_string client_id } let to_string x = Printf.sprintf "get_psi %d" (Id.Client.to_int x.client_id) end module Psi : sig type t = { n_state : Strictly_positive_int.t ; n_det : Strictly_positive_int.t ; psi_det_size : Strictly_positive_int.t ; n_det_generators : Strictly_positive_int.t option; n_det_selectors : Strictly_positive_int.t option; psi_det : string ; psi_coef : string ; energy : string; } val create : n_state:Strictly_positive_int.t -> n_det:Strictly_positive_int.t -> psi_det_size:Strictly_positive_int.t -> n_det_generators:Strictly_positive_int.t option -> n_det_selectors:Strictly_positive_int.t option -> psi_det:string -> psi_coef:string -> energy:string -> t end = struct type t = { n_state : Strictly_positive_int.t ; n_det : Strictly_positive_int.t ; psi_det_size : Strictly_positive_int.t ; n_det_generators : Strictly_positive_int.t option; n_det_selectors : Strictly_positive_int.t option; psi_det : string ; psi_coef : string ; energy : string ; } let create ~n_state ~n_det ~psi_det_size ~n_det_generators ~n_det_selectors ~psi_det ~psi_coef ~energy = assert (Strictly_positive_int.to_int n_det <= Strictly_positive_int.to_int psi_det_size ); { n_state; n_det ; psi_det_size ; n_det_generators ; n_det_selectors ; psi_det ; psi_coef ; energy } end (** GetPsiReply_msg : Reply to the GetPsi message *) module GetPsiReply_msg : sig type t = { client_id : Id.Client.t ; psi : Psi.t } val create : client_id:Id.Client.t -> psi:Psi.t -> t val to_string_list : t -> string list val to_string : t -> string end = struct type t = { client_id : Id.Client.t ; psi : Psi.t } let create ~client_id ~psi = { client_id ; psi } let to_string x = let g, s = match x.psi.Psi.n_det_generators, x.psi.Psi.n_det_selectors with | Some g, Some s -> Strictly_positive_int.to_int g, Strictly_positive_int.to_int s | _ -> -1, -1 in Printf.sprintf "get_psi_reply %d %d %d %d %d %d" (Id.Client.to_int x.client_id) (Strictly_positive_int.to_int x.psi.Psi.n_state) (Strictly_positive_int.to_int x.psi.Psi.n_det) (Strictly_positive_int.to_int x.psi.Psi.psi_det_size) g s let to_string_list x = [ to_string x ; x.psi.Psi.psi_det ; x.psi.Psi.psi_coef ; x.psi.Psi.energy ] end (** PutPsi : put the current variational wave function *) module PutPsi_msg : sig type t = { client_id : Id.Client.t ; n_state : Strictly_positive_int.t ; n_det : Strictly_positive_int.t ; psi_det_size : Strictly_positive_int.t ; n_det_generators : Strictly_positive_int.t option; n_det_selectors : Strictly_positive_int.t option; psi : Psi.t option } val create : client_id:string -> n_state:string -> n_det:string -> psi_det_size:string -> psi_det:string option -> psi_coef:string option -> n_det_generators: string option -> n_det_selectors:string option -> energy:string option -> t val to_string_list : t -> string list val to_string : t -> string end = struct type t = { client_id : Id.Client.t ; n_state : Strictly_positive_int.t ; n_det : Strictly_positive_int.t ; psi_det_size : Strictly_positive_int.t ; n_det_generators : Strictly_positive_int.t option; n_det_selectors : Strictly_positive_int.t option; psi : Psi.t option } let create ~client_id ~n_state ~n_det ~psi_det_size ~psi_det ~psi_coef ~n_det_generators ~n_det_selectors ~energy = let n_state, n_det, psi_det_size = Int.of_string n_state |> Strictly_positive_int.of_int , Int.of_string n_det |> Strictly_positive_int.of_int , Int.of_string psi_det_size |> Strictly_positive_int.of_int in assert (Strictly_positive_int.to_int psi_det_size >= Strictly_positive_int.to_int n_det); let n_det_generators, n_det_selectors = match n_det_generators, n_det_selectors with | Some x, Some y -> Some (Strictly_positive_int.of_int @@ Int.of_string x), Some (Strictly_positive_int.of_int @@ Int.of_string y) | _ -> None, None in let psi = match (psi_det, psi_coef, energy) with | (Some psi_det, Some psi_coef, Some energy) -> Some (Psi.create ~n_state ~n_det ~psi_det_size ~psi_det ~psi_coef ~n_det_generators ~n_det_selectors ~energy) | _ -> None in { client_id = Id.Client.of_string client_id ; n_state ; n_det ; psi_det_size ; n_det_generators ; n_det_selectors ; psi } let to_string x = match x.n_det_generators, x.n_det_selectors with | Some g, Some s -> Printf.sprintf "put_psi %d %d %d %d %d %d" (Id.Client.to_int x.client_id) (Strictly_positive_int.to_int x.n_state) (Strictly_positive_int.to_int x.n_det) (Strictly_positive_int.to_int x.psi_det_size) (Strictly_positive_int.to_int g) (Strictly_positive_int.to_int s) | _, _ -> Printf.sprintf "put_psi %d %d %d %d %d %d" (Id.Client.to_int x.client_id) (Strictly_positive_int.to_int x.n_state) (Strictly_positive_int.to_int x.n_det) (Strictly_positive_int.to_int x.psi_det_size) (-1) (-1) let to_string_list x = match x.psi with | Some psi -> [ to_string x ; psi.Psi.psi_det ; psi.Psi.psi_coef ; psi.Psi.energy ] | None -> [ to_string x ; "None" ; "None" ; "None" ] end (** PutPsiReply_msg : Reply to the PutPsi message *) module PutPsiReply_msg : sig type t val create : client_id:Id.Client.t -> t val to_string : t -> string end = struct type t = { client_id : Id.Client.t ; } let create ~client_id = { client_id; } let to_string x = Printf.sprintf "put_psi_reply %d" (Id.Client.to_int x.client_id) end (** TaskDone : Inform the server that a task is finished *) module TaskDone_msg : sig type t = { client_id: Id.Client.t ; state: State.t ; task_id: Id.Task.t ; } val create : state:string -> client_id:string -> task_id:string -> t val to_string : t -> string end = struct type t = { client_id: Id.Client.t ; state: State.t ; task_id: Id.Task.t; } let create ~state ~client_id ~task_id = { client_id = Id.Client.of_string client_id ; state = State.of_string state ; task_id = Id.Task.of_string task_id; } let to_string x = Printf.sprintf "task_done %s %d %d" (State.to_string x.state) (Id.Client.to_int x.client_id) (Id.Task.to_int x.task_id) end (** Terminate *) module Terminate_msg : sig type t val create : unit -> t val to_string : t -> string end = struct type t = Terminate let create () = Terminate let to_string x = "terminate" end (** OK *) module Ok_msg : sig type t val create : unit -> t val to_string : t -> string end = struct type t = Ok let create () = Ok let to_string x = "ok" end (** Error *) module Error_msg : sig type t val create : string -> t val to_string : t -> string end = struct type t = string let create x = x let to_string x = String.concat ~sep:" " [ "error" ; x ] end (** Message *) type t = | GetPsi of GetPsi_msg.t | PutPsi of PutPsi_msg.t | GetPsiReply of GetPsiReply_msg.t | PutPsiReply of PutPsiReply_msg.t | Newjob of Newjob_msg.t | Endjob of Endjob_msg.t | Connect of Connect_msg.t | ConnectReply of ConnectReply_msg.t | Disconnect of Disconnect_msg.t | DisconnectReply of DisconnectReply_msg.t | GetTask of GetTask_msg.t | GetTaskReply of GetTaskReply_msg.t | DelTask of DelTask_msg.t | DelTaskReply of DelTaskReply_msg.t | AddTask of AddTask_msg.t | AddTaskReply of AddTaskReply_msg.t | TaskDone of TaskDone_msg.t | Terminate of Terminate_msg.t | Ok of Ok_msg.t | Error of Error_msg.t | SetStopped | SetWaiting | SetRunning let of_string s = let l = String.split ~on:' ' s |> List.filter ~f:(fun x -> (String.strip x) <> "") |> List.map ~f:String.lowercase in match l with | "add_task" :: state :: task -> AddTask (AddTask_msg.create ~state ~task:(String.concat ~sep:" " task) ) | "del_task" :: state :: task_id :: [] -> DelTask (DelTask_msg.create ~state ~task_id) | "get_task" :: state :: client_id :: [] -> GetTask (GetTask_msg.create ~state ~client_id) | "task_done" :: state :: client_id :: task_id :: [] -> TaskDone (TaskDone_msg.create ~state ~client_id ~task_id) | "disconnect" :: state :: client_id :: [] -> Disconnect (Disconnect_msg.create ~state ~client_id) | "connect" :: t :: [] -> Connect (Connect_msg.create t) | "new_job" :: state :: push_address_tcp :: push_address_inproc :: [] -> Newjob (Newjob_msg.create push_address_tcp push_address_inproc state) | "end_job" :: state :: [] -> Endjob (Endjob_msg.create state) | "terminate" :: [] -> Terminate (Terminate_msg.create () ) | "get_psi" :: client_id :: [] -> GetPsi (GetPsi_msg.create ~client_id) | "put_psi" :: client_id :: n_state :: n_det :: psi_det_size :: n_det_generators :: n_det_selectors :: [] -> PutPsi (PutPsi_msg.create ~client_id ~n_state ~n_det ~psi_det_size ~n_det_generators:(Some n_det_generators) ~n_det_selectors:(Some n_det_selectors) ~psi_det:None ~psi_coef:None ~energy:None ) | "put_psi" :: client_id :: n_state :: n_det :: psi_det_size :: [] -> PutPsi (PutPsi_msg.create ~client_id ~n_state ~n_det ~psi_det_size ~n_det_generators:None ~n_det_selectors:None ~psi_det:None ~psi_coef:None ~energy:None) | "ok" :: [] -> Ok (Ok_msg.create ()) | "error" :: rest -> Error (Error_msg.create (String.concat ~sep:" " rest)) | "set_stopped" :: [] -> SetStopped | "set_running" :: [] -> SetRunning | "set_waiting" :: [] -> SetWaiting | _ -> failwith "Message not understood" let to_string = function | GetPsi x -> GetPsi_msg.to_string x | PutPsiReply x -> PutPsiReply_msg.to_string x | Newjob x -> Newjob_msg.to_string x | Endjob x -> Endjob_msg.to_string x | Connect x -> Connect_msg.to_string x | ConnectReply x -> ConnectReply_msg.to_string x | Disconnect x -> Disconnect_msg.to_string x | DisconnectReply x -> DisconnectReply_msg.to_string x | GetTask x -> GetTask_msg.to_string x | GetTaskReply x -> GetTaskReply_msg.to_string x | DelTask x -> DelTask_msg.to_string x | DelTaskReply x -> DelTaskReply_msg.to_string x | AddTask x -> AddTask_msg.to_string x | AddTaskReply x -> AddTaskReply_msg.to_string x | TaskDone x -> TaskDone_msg.to_string x | Terminate x -> Terminate_msg.to_string x | Ok x -> Ok_msg.to_string x | Error x -> Error_msg.to_string x | PutPsi x -> PutPsi_msg.to_string x | GetPsiReply x -> GetPsiReply_msg.to_string x | SetStopped -> "set_stopped" | SetRunning -> "set_running" | SetWaiting -> "set_waiting" let to_string_list = function | PutPsi x -> PutPsi_msg.to_string_list x | GetPsiReply x -> GetPsiReply_msg.to_string_list x | _ -> assert false