Source file 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 Fr = Bls12_381.Fr

module Stubs = struct
  type fr = Fr.t
  type fr_array

  external allocate_fr_array : int -> fr_array
    = "caml_polynomial_allocate_fr_array_stubs"
  (** [allocate_fr_array n] allocates an OCaml custom block to hold a C array
  containing [n] zeros of blst_fr

  - requires: [n > 0]
  - ensures: [size res = n] and [res] is initialized with zeros of blst_fr *)

  external erase_fr_array : fr_array -> int -> unit
    = "caml_polynomial_erase_fr_array_stubs"
  (** [erase_fr_array p n] writes [n] zeros of blst_fr in a given array [p]

   - requires: [n <= size p]
   - ensures: a prefix of size [n] of an array [p] is filled with zeros *)

  external of_fr_array : fr array -> fr_array -> int -> unit
    = "caml_polynomial_of_fr_array_stubs"
    [@@noalloc]
  (** [of_fr_array res p n] converts a given C array [p] of at least size [n] to
  an OCaml array [res] of size [n]

  - requires: [n <= size p]
  - ensures: [size res = n] and [res] is a prefix of [p] *)

  external get : fr -> fr_array -> 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 copy : fr_array -> fr_array -> int -> int -> unit
    = "caml_polynomial_copy_stubs"
    [@@noalloc]
  (** [copy a b offset len] starting from position [offset] copies [len]
      elements from array [b] to array [a]

  requires:
  - [a] and [b] be disjoint
  - [len <= size a]
  - [offset + len <= size b]
      
  ensures: [a[i] = b[offset + i]] for [i = 0..(len - 1)] *)

  external of_dense : fr_array -> fr array -> int -> unit
    = "caml_polynomial_of_dense_stubs"
    [@@noalloc]
  (** [of_dense res p n] converts an OCaml array [p] of size [n] to a C array [res] of
  size [n]

  requires:
  - [size res = n]
  - [size p = n] *)

  external eq : fr_array -> fr_array -> int -> int -> bool
    = "caml_polynomial_eq_stubs"
    [@@noalloc]
  (** [eq a b size_a size_b] checks whether a polynomial [a] is equal to a polynomial [b]

  requires:
  - [size a = size_a]
  - [size b = size_b] *)

  external is_zero : fr_array -> int -> bool = "caml_polynomial_is_zero_stubs"
    [@@noalloc]
  (** [is_zero p n] checks whether a prefix of size [n] of an array [p] is filled
      with zeros of blst_fr

  - requires: [size p >= n] *)
end

module Carray : sig
  type scalar = Bls12_381.Fr.t
  type t = Stubs.fr_array * int

  val encoding : t Data_encoding.t

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

  val erase : t -> unit
  (** [erase c] overwrites a C array [c] with zeros of blst_fr *)

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

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

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

  val to_string : t -> string
  (** [to_string c] returns the string representation of a C array [c] *)

  val to_array : t -> scalar array
  (** [to_array c] converts a C array [c] to an OCaml array *)

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

  val equal : t -> t -> bool
  (** [equal a b] checks whether a C array [a] is equal to a C array [b]

      Note: [equal] is defined as polynomial equality, i.e., [a] and [b] can have
      different sizes and padded with zeros of blst_fr *)
end = struct
  type scalar = Bls12_381.Fr.t

  (* Invariant: size of a P is always > 0 *)
  type t = Stubs.fr_array * int

  let erase (p, n) = Stubs.erase_fr_array p n
  let length (_, n) = n

  let get (p, n) i =
    if i < 0 || i >= n then raise @@ Invalid_argument "get: index out of bounds";
    let res = Fr.(copy zero) in
    Stubs.get res p i;
    res

  let allocate n =
    if n < 1 then raise @@ Invalid_argument "allocate: should be greater than 1"
    else Stubs.allocate_fr_array n

  let to_array (p, n) =
    let d = Array.init n (fun _ -> Fr.(copy zero)) in
    Stubs.of_fr_array d p n;
    d

  let of_array coeff =
    let n = Array.length coeff in
    let res = allocate n in
    Stubs.of_dense res coeff n;
    (res, n)

  let encoding : t Data_encoding.t =
    let open Data_encoding in
    let fr_encoding =
      conv Fr.to_bytes Fr.of_bytes_exn (Fixed.bytes Fr.size_in_bytes)
    in
    conv to_array of_array (array fr_encoding)

  let to_string p =
    String.concat " ; " (List.map Fr.to_string (Array.to_list @@ to_array p))

  let copy ?(offset = 0) ?len (coefficients, n) =
    let len = Option.value ~default:n len in
    assert (offset + len <= n);
    let res = allocate len in
    Stubs.copy res coefficients offset len;
    (res, len)

  let equal (poly_1, size_1) (poly_2, size_2) =
    Stubs.eq poly_1 poly_2 size_1 size_2
end

include Carray