Source file fr_carray.ml

(*****************************************************************************)
(*                                                                           *)
(* MIT License                                                               *)
(* Copyright (c) 2022 Nomadic Labs <contact@nomadic-labs.com>                *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

module Stubs = struct
  type t = Bigstringaf.t
  type elt = Bls12_381.Fr.t

  let size = Bls12_381.Fr.size_in_bytes
  let allocate_elt () = Bls12_381.Fr.(copy one)

  external get : elt -> t -> int -> unit = "caml_polynomial_get_stubs"
    [@@noalloc]
  (** [get res p i] writes [i]-th element of a given array [p] in [res]

  - requires: [0 <= i < size p]
  - ensures: [res = p[i]] *)

  external set : t -> elt -> int -> unit = "caml_polynomial_set_stubs"
    [@@noalloc]
  (** [set p fr i] writes [fr] in the [i]-th element of [p]

  - requires: [0 <= i < size p]
  - ensures: [fr = p[i]] *)
end

module type Stubs_sig = sig
  type t = Bigstringaf.t
  type elt

  val allocate_elt : unit -> elt
  val size : int
  val get : elt -> t -> int -> unit
  val set : t -> elt -> int -> unit
end

module type Carray_sig = sig
  type t = Bigstringaf.t
  type elt

  val encoding : t Data_encoding.t

  val allocate : int -> t
  (** [allocate len] creates a C array of size [len] intialized with zeros of blst_fr *)

  val length : t -> int
  (** [length c] returns the length of a C array [c] *)

  val get : t -> int -> elt
  (** [get c i] returns the [i]-th element of a C array [c] *)

  val set : t -> elt -> int -> unit

  val copy : ?offset:int -> ?len:int -> t -> t
  (** [copy c] copies [len] elements from a C array [c] starting from position [offset] *)

  val blit : t -> src_off:int -> t -> dst_off:int -> len:int -> unit
  (** [blit src src_off dst dst_off len] copies [len] elements from [src] to
      [dst] starting at the respective offsets. *)

  val to_array : ?offset:int -> ?len:int -> t -> elt array
  (** [to_array c] converts a C array [c] to an OCaml array *)

  val of_array : elt array -> t
  (** [of_array c] converts an OCaml array [c] to a C array *)
end

module Make (S : Stubs_sig) : Carray_sig with type elt = S.elt = struct
  type t = S.t
  type elt = S.elt

  let length a = Bigstringaf.length a / S.size

  let allocate n =
    if n < 1 then raise @@ Invalid_argument "allocate: size should be >= 1";
    let size = S.size * n in
    let res = Bigstringaf.create size in
    res

  let get p i =
    if i < 0 || i >= length p then
      raise @@ Invalid_argument "get: index out of bounds";
    let res = S.allocate_elt () in
    S.get res p i;
    res

  let set p fr i =
    if i < 0 || i >= length p then
      raise @@ Invalid_argument "set: index out of bounds";
    S.set p fr i

  let to_array ?(offset = 0) ?len p =
    let len = Option.value ~default:(length p - offset) len in
    if len < 0 || offset < 0 || length p - offset < len then
      raise
      @@ Invalid_argument
           (Format.sprintf "to_array: invalid len %d or offset %d for size %d"
              len offset (length p));
    Array.init len (fun i -> get p i)

  let of_array caml_array =
    let n = Array.length caml_array in
    let res = allocate n in
    Array.iteri (fun i g -> set res g i) caml_array;
    res

  let encoding =
    Data_encoding.(
      conv
        (fun bs -> Bigstringaf.to_string bs)
        (fun s ->
          let len = String.length s in
          Bigstringaf.of_string ~off:0 ~len s)
        string)

  let blit src ~src_off dst ~dst_off ~len =
    let src_off = src_off * S.size in
    let dst_off = dst_off * S.size in
    let len = len * S.size in
    Bigstringaf.blit src ~src_off dst ~dst_off ~len

  let copy ?(offset = 0) ?len p =
    let len = Option.value ~default:(length p - offset) len in
    if len < 0 || offset < 0 || length p - offset < len then
      raise @@ Invalid_argument "copy: invalid len or offset";
    let res = allocate len in
    blit p ~src_off:offset res ~dst_off:0 ~len;
    res
end

include Make (Stubs)