(** Encoder and decoder for Memtrace traces. *)

(** {1 Time} *)

(** Timestamps. *)
module Timestamp : sig
  type t

  val now : unit -> t
  (** [now ()] is the current timestamp. *)

  val of_int64 : int64 -> t
  (** [of_int64 us] is a timestamp from microseconds since the Unix epoch. *)

  val to_int64 : t -> int64
  (** [to_int64 t] is [t] as microseconds since the Unix epoch. *)

  val to_float : t -> float
  (** [to_float t] is [t] in the Unix module's float format. *)

  val of_float : float -> t
  (** [of_float f] is a timestamp from the Unix module's float format. *)
end

(** Time deltas measured from the start of the trace. *)
module Timedelta : sig
  type t

  val to_int64 : t -> int64
  (** [to_int64 t] is [t] as microseconds since the start of the trace. *)

  val offset : Timestamp.t -> t -> Timestamp.t
  (** [offset base delta] is [base] offset by [delta]. *)
end

(** {1 Identifiers} *)

(** Source locations in the traced program. *)
module Location : sig
  type t = {
    filename : string;
    line : int;
    start_char : int;
    end_char : int;
    defname : string;
  }

  val to_string : t -> string
  (** [to_string t] is a human-readable representation of [t]. *)

  val unknown : t
  (** [unknown] is a location representing an unknown source position. *)
end

(** Identifiers to represent allocations. *)
module Obj_id : sig
  type t = private int

  module Tbl : Hashtbl.SeededS with type key = t
  (** Hashtable keyed by object ID. *)

  module Expert : sig
    val of_int : int -> t
    (** [of_int n] is an object ID from an integer. *)
  end
end

(** Identifiers to represent domains. *)
module Domain_id : sig
  type t = private int

  val main_domain : t
  (** [main_domain] is the ID of the initial domain. *)

  module Expert : sig
    val of_int : int -> t
    (** [of_int n] is a domain ID from an integer. *)
  end
end

(** Codes for subsequences of locations in a backtrace. *)
module Location_code : sig
  type t = private int

  module Tbl : Hashtbl.SeededS with type key = t
  (** Hashtable keyed by location code. *)

  module Expert : sig
    val of_int : int -> t
    (** [of_int n] is a location code from an integer. *)
  end
end

(** Types of allocation. *)
module Allocation_source : sig
  type t = Minor | Major | External
end

(** {1 Events} *)

(** Trace events. *)
module Event : sig
  type t =
    | Alloc of {
        obj_id : Obj_id.t;
            (** An identifier for this allocation, used to refer to it in other
                events. These identifiers are generated in allocation order.

                Identifiers are not reused across domains. *)
        length : int;
            (** The ID of the domain on which this event occurred.

                All events with the same obj_id will occur on the same domain,
                unless that domain terminates before collecting this object (in
                which case promotion or collection events may occur on an
                arbitrary domain, but will always appear in order
                Alloc-Promote-Collect for a given obj_id).

                Events from different domains and different obj_ids may appear
                out-of-order in a trace: timestamps and object identifiers are
                only guaranteed to increase monotonically among events from a
                single domain. Readers that require a global order of events
                must merge the streams of events from different domains. *)
        domain : Domain_id.t;
            (** Length of the sampled allocation, in words, not including header
                word. *)
        nsamples : int;
            (** Number of samples made in this allocation. At least 1. *)
        source : Allocation_source.t;
            (** How this object was initially allocated. *)
        backtrace_buffer : Location_code.t array;
            (** Backtrace of the allocation. The backtrace elements are stored
                in order from caller to callee. The first element is the main
                entrypoint and the last is the allocation.

                NB: this is a mutable buffer, reused between events. Entries at
                indices beyond [backtrace_length - 1] are not meaningful. If you
                want to store backtraces, you must copy them using:
                [Array.sub backtrace_buffer 0 backtrace_length]. *)
        backtrace_length : int;  (** Length of the backtrace. *)
        common_prefix : int;
            (** A prefix of this length has not changed since the last
                allocation event that occurred on the same domain. *)
      }
    | Promote of Obj_id.t * Domain_id.t
    | Collect of Obj_id.t * Domain_id.t

  val domain : t -> Domain_id.t
  (** [domain ev] is the domain on which [ev] occurred. *)

  val to_string : (Location_code.t -> Location.t list) -> t -> string
  (** [to_string lookup ev] is a human-readable representation of [ev]. *)
end

(** {1 Trace metadata} *)

(** Global trace info. *)
module Info : sig
  type t = {
    sample_rate : float;
    word_size : int;
    executable_name : string;
    host_name : string;
    ocaml_runtime_params : string;
    pid : Int64.t;
    initial_domain : Domain_id.t;
    start_time : Timestamp.t;
    context : string option;
  }
end

(** {1 Writing traces} *)

(** Writing traces. *)
module Writer : sig
  type t

  exception Pid_changed

  val create : Unix.file_descr -> ?getpid:(unit -> int64) -> Info.t -> t
  (** [create fd info] is a trace writer to [fd]. *)

  val domain : t -> Domain_id.t
  (** [domain t] is the domain this writer writes events for. *)

  val for_domain : t -> domain:Domain_id.t -> t
  (** [for_domain t ~domain] is a writer for [domain] sharing [t]'s fd. *)

  (** All of the functions below may raise Unix_error if writing to the file
      descriptor fails, or Pid_changed if getpid returns a different value. *)

  val put_alloc :
    t ->
    Timestamp.t ->
    length:int ->
    nsamples:int ->
    source:Allocation_source.t ->
    callstack:Location_code.t array ->
    decode_callstack_entry:(Location_code.t -> Location.t list) ->
    Obj_id.t
  (** [put_alloc t time ~length ~nsamples ~source ~callstack
       ~decode_callstack_entry] writes an allocation event. *)

  val put_alloc_with_raw_backtrace :
    t ->
    Timestamp.t ->
    length:int ->
    nsamples:int ->
    source:Allocation_source.t ->
    callstack:Printexc.raw_backtrace ->
    Obj_id.t
  (** [put_alloc_with_raw_backtrace t time ~length ~nsamples ~source ~callstack]
      writes an allocation event using a raw backtrace. *)

  val put_alloc_backtrace_suffix :
    t ->
    Timestamp.t ->
    length:int ->
    nsamples:int ->
    source:Allocation_source.t ->
    callstack:Printexc.raw_backtrace ->
    drop_slots:int ->
    Obj_id.t
  (** [put_alloc_backtrace_suffix t time ~length ~nsamples ~source ~callstack
       ~drop_slots] writes an allocation, dropping the first [drop_slots]
      backtrace entries. *)

  val put_collect : t -> Timestamp.t -> Obj_id.t -> unit
  (** [put_collect t time obj_id] writes a collection event. *)

  val put_promote : t -> Timestamp.t -> Obj_id.t -> unit
  (** [put_promote t time obj_id] writes a promotion event. *)

  val flush : t -> unit
  (** [flush t] flushes currently buffered events. The file descriptor is not
      closed, and more events may be written after flush. *)

  val close : t -> unit
  (** [close t] flushes and closes the underlying file descriptor. *)

  module Multiplexed_domains : sig
    type t
    (** A set of per-domain writers, allowing events from different domains to
        be multiplexed into a single stream. *)

    val create : Unix.file_descr -> ?getpid:(unit -> int64) -> Info.t -> t
    (** [create fd info] is a multiplexed writer to [fd]. *)

    val put_event :
      t ->
      decode_callstack_entry:(Location_code.t -> Location.t list) ->
      Timestamp.t ->
      Event.t ->
      unit
    (** [put_event t ~decode_callstack_entry time ev] writes [ev]. *)

    val next_alloc_id : t -> domain:Domain_id.t -> Obj_id.t
    (** [next_alloc_id t ~domain] is the next allocation ID for [domain]. *)

    val flush : t -> unit
    (** [flush t] flushes all per-domain writers. *)
  end
end

(** {1 Reading traces} *)

(** Reading traces. *)
module Reader : sig
  type t

  val create : Unix.file_descr -> t
  (** [create fd] is a trace reader from [fd]. *)

  val info : t -> Info.t
  (** [info t] is the trace metadata. *)

  val lookup_location_code : t -> Location_code.t -> Location.t list
  (** [lookup_location_code t code] resolves [code] to source locations. *)

  val iter :
    t -> ?parse_backtraces:bool -> (Timedelta.t -> Event.t -> unit) -> unit
  (** [iter t f] calls [f] for each event in the trace. *)

  val open_ : filename:string -> t
  (** [open_ ~filename] opens a trace file. *)

  val size_bytes : t -> int64
  (** [size_bytes t] is the size of the trace in bytes. *)

  val close : t -> unit
  (** [close t] closes the reader. *)
end

(** {1:private Private} *)

module Private : sig
  val set_tracer_module : string -> unit
  (** [set_tracer_module name] sets the module name for backtrace filtering. *)

  val obj_ids_per_chunk : int Atomic.t
  (** For testing: number of object IDs per chunk. *)
end