Source file evaluations_c.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 = Polynomial_c.Stubs.fr
  type fr_array = Polynomial_c.Stubs.fr_array

  external add : fr_array -> fr_array -> fr_array -> int -> int -> unit
    = "caml_polynomial_evaluations_add_stubs"
    [@@noalloc]
  (** [add res a b size_a size_b] writes the result of polynomial addition
      of [a] and [b] using the evaluation representation in [res], where
      - [a] is evaluated on [domain_a] of size [size_a]
      - [b] is evaluated on [domain_b] of size [size_b]

   requires:
   - [size a = size_a]
   - [size b = size_b]
   - [size res = min (size_a, size_b)]
   - [res], [a] and [b] are either pairwise disjoint or equal
   - [size_b mod size_a = 0] *)

  external mul_arrays :
    fr_array ->
    (fr_array * int * int * Bytes.t * int) array ->
    int ->
    int ->
    unit = "caml_polynomial_evaluations_mul_arrays_stubs"
    [@@noalloc]
  (** [mul_arrays res eval_evallen_comp_power_powlen size_res nb_evals] writes
      the result of computing [p₁(gᶜ₁·x)ᵐ₁·p₂(gᶜ₂·x)ᵐ₂·…·pₖ(gᶜₖ·x)ᵐₖ] using
      the evaluation representation in [res], where
      - [eval_evallen_comp_power_powlen.[i] = (pᵢ, size_p_i, cᵢ, mᵢ, size_bits_m_i)]
      - a polynomial [pᵢ] is evaluated on [domain_p_i] of size [size_p_i]
      - [cᵢ] is a composition_gx; it computes [pᵢ(gᶜᵢ·x)] instead of [pᵢ(x)],
      where [g] is a primitive [size_p_i]-th root of unity
      - [mᵢ] is a power in [pᵢ(x)ᵐᵢ]
      - [size_bits_m_i] is the *exact* number of bits in [mᵢ]
      - [nb_evals] is a number of evaluations, i.e., [i = 1..nb_evals]

   requires:
   - [size res = size_res]
   - [size eval_evallen_comp_power_powlen = nb_evals]
   - [size p_i = size_p_i]
   - [size_bits m_i = size_bits_m_i]
   - [size_p_i mod size_res = 0]
   - [res] and [p_i] are either pairwise disjoint or equal when [cᵢ = 0] *)

  external linear_arrays :
    fr_array -> (fr_array * int * fr * int) array -> fr -> int -> int -> unit
    = "caml_polynomial_evaluations_linear_arrays_stubs"
    [@@noalloc]
  (** [linear_arrays res eval_evallen_coeff_comp add_constant size_res nb_evals]
      writes  the result of computing
      [λ₁·p₁(gᶜ₁·x) + λ₂·p₂(gᶜ₂·x) + … + λₖ·pₖ(gᶜₖ·x) + add_constant] using
      the evaluation representation in [res], where
      - [eval_evallen_coeff_comp.[i] = (pᵢ, size_p_i, λᵢ, cᵢ)]
      - a polynomial [pᵢ] is evaluated on [domain_p_i] of size [size_p_i]
      - [cᵢ] is a composition_gx; it computes [pᵢ(gᶜᵢ·x)] instead of [pᵢ(x)],
      where [g] is a primitive [size_p_i]-th root of unity
      - [λᵢ] is a coefficient in [λᵢ·pᵢ(x)]
      - [nb_evals] is a number of evaluations, i.e., [i = 1..nb_evals]

   requires:
   - [size res = size_res]
   - [size eval_evallen_coeff_comp = nb_evals]
   - [size p_i = size_p_i]
   - [size_p_i mod size_res = 0]
   - [res] and [p_i] are disjoint *)

  external fft_inplace :
    fr_array ->
    domain:fr_array ->
    log:int ->
    degree:int ->
    log_degree:int ->
    unit = "caml_polynomial_fft_inplace_on_stubs"
    [@@noalloc]
  (** [fft_inplace p domain log] computes Fast Fourier Transform.
  It converts the coefficient representation of a polynomial [p] to
  the evaluation representation

  requires:
  - [size p = size domain]
  - [size domain = 2^log]
  - [domain = [one; g; ..; g^{n-1}]] where [g] is a primitive
  [n]-th root of unity and [n = 2^log] (as done by {!compute_domain}) *)

  external ifft_inplace : fr_array -> domain:fr_array -> log:int -> unit
    = "caml_polynomial_ifft_inplace_on_stubs"
    [@@noalloc]
  (** [ifft_inplace p domain log] computes Inverse Fast Fourier Transform.
  It converts the evaluation representation of a polynomial [p] to
  the coefficient representation

  requires:
  - [size p = size domain]
  - [size domain = 2^log]
  - [domain = [one; g; ..; g^{n-1}]] where [g] is a primitive
  [n]-th root of unity and [n = 2^log] (as done by {!compute_domain}) *)
end

module type Evaluations_c_sig = sig
  type scalar
  type polynomial
  type t

  val make_evaluation : int * scalar array -> t
  (** [make_evaluation (d, e)] creates a value of type [t] from the evaluation
   representation of a polynomial [e] of degree [d], i.e, it converts an OCaml
   array to a C array *)

  val string_of_eval : t -> string
  (** [string_of_eval e] returns the string representation of evaluation [e] *)

  type domain

  val of_domain : domain -> t
  (** [of_domain d] converts [d] from type {!domain} to type {!t}.

  Note: [of_domain d] doesn't create a copy of [d] *)

  val to_domain : t -> domain
  (** [to_domain d] converts [d] from type {!t} to type {!domain}.

  Note: [to_domain d] doesn't create a copy of [d] *)

  val zero : t
  (** [zero] returns the evaluation representation of the zero polynomial *)

  val degree : t -> int
  (** [degree] returns the degree of a polynomial. Returns [-1] for the zero polynomial *)

  val length : t -> int
  (** [length e] returns the size of domain where a polynomial is evaluated, or equally,
  the size of a C array where evaluation [e] is stored *)

  val create : int -> t
  (** [create len] returns the evaluation representation of a zero polynomial of size [len] *)

  val copy : ?res:t -> t -> t
  (** [copy ?res a] returns a copy of evaluation [a]. The function writes the result in [res]
   if [res] has the correct size and allocates a new array for the result otherwise *)

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

  val mul_by_scalar : scalar -> t -> t
  (** [mul_by_scalar] computes muliplication of a polynomial by a blst_fr element *)

  val mul_c :
    ?res:t ->
    evaluations:t list ->
    ?composition_gx:int list * int ->
    ?powers:int list ->
    unit ->
    t
  (** [mul_c] computes [p₁(gᶜ₁·x)ᵐ₁·p₂(gᶜ₂·x)ᵐ₂·…·pₖ(gᶜₖ·x)ᵐₖ], where
   - [pᵢ = List.nth evaluations i]
   - [mᵢ = List.nth powers i]
   - [cᵢ = List.nth (fst composition_gx) i]
   - [n = snd composition_gx] is the order of generator, i.e., [gⁿ = 1]

  The function writes the result in [res] if [res] has the correct size (= min (size pᵢ))
  and allocates a new array for the result otherwise

  Note: [res] can be equal to [pᵢ] when [cᵢ = 0] *)

  val linear_c :
    ?res:t ->
    evaluations:t list ->
    ?linear_coeffs:scalar list ->
    ?composition_gx:int list * int ->
    ?add_constant:scalar ->
    unit ->
    t
  (** [linear_c] computes [λ₁·p₁(gᶜ₁·x) + λ₂·p₂(gᶜ₂·x) + … + λₖ·pₖ(gᶜₖ·x) + add_constant], where
   - [pᵢ = List.nth evaluations i]
   - [λᵢ = List.nth linear_coeffs i]
   - [cᵢ = List.nth (fst composition_gx) i]
   - [n = snd composition_gx] is the order of generator, i.e., [gⁿ = 1]

  The function writes the result in [res] if [res] has the correct size (= min (size pᵢ))
  and allocates a new array for the result otherwise

  Note: [res] and [pᵢ] are disjoint *)

  val add : ?res:t -> t -> t -> t
  (** [add ?res a b] computes polynomial addition of [a] and [b]. The function writes
   the result in [res] if [res] has the correct size (= min (size (a, b))) and allocates
   a new array for the result otherwise

  Note: [res] can be equal to either [a] or [b] *)

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

  Note: [equal] is defined as restrictive equality, i.e., the same polynomial
  evaluated on different domains are said to be different *)

  val evaluation_fft : domain -> polynomial -> t
  (** [evaluation_fft domain p] converts the coefficient representation of
  a polynomial [p] to the evaluation representation. [domain] can be obtained
  using {!Domain.build}

  Note:
  - size of domain must be a power of two
  - size of a polynomial must be less than or equal to size of domain *)

  val interpolation_fft : domain -> t -> polynomial
  (** [interpolation_fft domain p] converts the evaluation representation of
  a polynomial [p] to the coefficient representation. [domain] can be obtained
  using {!Domain.build}

  Note:
  - size of domain must be a power of two
  - size of a polynomial must be equal to size of domain *)

  (* TODO DELETE *)
  val evaluation_fft2 : domain -> polynomial -> scalar array
  val interpolation_fft2 : domain -> scalar array -> polynomial
end

module Evaluations_c_impl = struct
  module Domain_c = Domain.Stubs
  module Domain = Domain.Domain_unsafe
  module Polynomial = Polynomial_c.Polynomial_unsafe

  type scalar = Fr.t
  type polynomial = Polynomial.t

  (* degree & evaluations & length *)
  type t = int * Stubs.fr_array * int
  type domain = Domain.t

  let make_evaluation (d, p) =
    if d < 0 then
      raise @@ Invalid_argument "make_evaluation: degree must be non negative";
    if Array.length p <= d then
      raise
      @@ Invalid_argument "make_evaluation: array must be longer than degree";
    let res, len = Carray.of_array p in
    (d, res, len)

  let string_of_eval (d, e, l) =
    Printf.sprintf "%d : [%s]" d (Carray.to_string (e, l))

  let of_domain domain =
    let d, l = Domain.to_carray domain in
    (1, d, l)

  let to_domain (_, eval, l) = Domain.of_carray (eval, l)
  let zero = (-1, Carray.allocate 1, 1)
  let degree (d, _, _) = d
  let length (_, _, l) = l
  let create length = (-1, Carray.allocate length, length)

  let allocate_for_res res length_result =
    match res with
    | Some (_, e, l) when l = length_result -> e
    | _ -> Carray.allocate length_result

  let copy ?res (d, e, l) =
    let res = allocate_for_res res l in
    Carray.Stubs.copy res e 0 l;
    (d, res, l)

  let get (_, eval, l) i = Carray.get (eval, l) i

  let mul_by_scalar lambda (deg, eval, len) =
    let res = Carray.allocate len in
    Polynomial_c.Stubs.mul_by_scalar res lambda eval len;
    (deg, res, len)

  (* multiplies evaluations of all polynomials with name in poly_names,
     the resulting eval has the size of the smallest evaluation *)
  let mul_c ?res ~evaluations
      ?(composition_gx = (List.init (List.length evaluations) (fun _ -> 0), 1))
      ?(powers = List.init (List.length evaluations) (fun _ -> 1)) () =
    let domain_len = snd composition_gx in
    assert (domain_len != 0);
    let deg_result, list_array, is_zero, min_length_eval =
      (* List.fold_left for 3 lists of same size *)
      let rec fold_left3 f acc l1 l2 l3 =
        match (l1, l2, l3) with
        | [], [], [] -> acc
        | a1 :: l1, a2 :: l2, a3 :: l3 -> fold_left3 f (f acc a1 a2 a3) l1 l2 l3
        | _ ->
            raise
              (Invalid_argument "fold_left3 : lists don’t have the same size.")
      in
      fold_left3
        (fun (acc_degree, list, is_zero, min_length_eval) x composition pow ->
          let d, eval, l = x in
          let is_zero = if d = -1 then true else is_zero in
          let new_min_length_eval =
            if l < min_length_eval then l else min_length_eval
          in
          let deg_to_add = if d = -1 then 0 else d * pow in
          let pow = Z.of_int pow in
          let pow_bytes = Z.to_bits pow |> Bytes.unsafe_of_string in
          let pow_len = Z.numbits pow in
          let rescale_composition = composition * l / domain_len in
          ( acc_degree + deg_to_add,
            (* ! here we reverse but it's ok because mul is commutative*)
            (eval, l, rescale_composition, pow_bytes, pow_len) :: list,
            is_zero,
            new_min_length_eval ))
        (0, [], false, Int.max_int)
        evaluations (fst composition_gx) powers
    in
    if is_zero then zero
    else
      let log = Z.(log2up (of_int deg_result)) in
      let length_degree = Z.(to_int (one lsl log)) in
      if length_degree > min_length_eval then
        raise
          (Invalid_argument
             (Printf.sprintf
                "Utils.mul_evaluations : evaluation is too short (length=%d) \
                 for expected result size %d"
                min_length_eval length_degree))
      else
        let nb_evals = List.length evaluations in
        let res = allocate_for_res res min_length_eval in
        Stubs.mul_arrays res (Array.of_list list_array) min_length_eval nb_evals;
        (deg_result, res, min_length_eval)

  (* Adds evaluation of a1 × p1 + a2 × p2 in evaluations
     /!\ the degree may not be always accurate,
         the resulting degree may not be the max of the 2 polynomials degrees *)
  let linear_c ?res ~evaluations
      ?(linear_coeffs = List.init (List.length evaluations) (fun _ -> Fr.one))
      ?(composition_gx = (List.init (List.length evaluations) (fun _ -> 0), 1))
      ?(add_constant = Fr.zero) () =
    let domain_len = snd composition_gx in
    assert (domain_len != 0);
    let list_eval_coeff_composition =
      List.map2
        (fun (eval, coeff) composition ->
          let rescale_composition = composition * length eval / domain_len in
          (eval, coeff, rescale_composition))
        (List.combine evaluations linear_coeffs)
        (fst composition_gx)
      |> List.filter (fun ((d, _eval, _length), _, _) -> d != -1)
    in
    let l =
      List.sort
        (fun ((_, _, length_1), _, _) ((_, _, length_2), _, _) ->
          Int.compare length_1 length_2)
        list_eval_coeff_composition
    in
    let length_result =
      if List.compare_length_with l 0 = 0 then 0
      else
        let (_, _, res), _, _ = List.hd l in
        res
    in
    if length_result = 0 then zero
    else
      let l =
        List.sort
          (fun ((d1, _, _), _, _) ((d2, _, _), _, _) -> Int.compare d2 d1)
          list_eval_coeff_composition
      in
      let degree_result =
        if List.compare_length_with l 0 = 0 then -1
        else
          let (res, _, _), _, _ = List.hd l in
          res
      in
      let nb_evals = List.length list_eval_coeff_composition in
      let res = allocate_for_res res length_result in
      let list_eval_coeff_composition =
        List.map
          (fun ((_, eval, eval_len), linear_coeff, composition) ->
            (eval, eval_len, linear_coeff, composition))
          list_eval_coeff_composition
      in
      Stubs.linear_arrays res
        (Array.of_list list_eval_coeff_composition)
        add_constant length_result nb_evals;
      (degree_result, res, length_result)

  (* Adds 2 evaluations *)
  let add ?res (d1, eval1, l1) (d2, eval2, l2) =
    if d1 = -1 then
      let res = allocate_for_res res l2 in
      copy ~res:(d2, res, l2) (d2, eval2, l2)
    else if d2 = -1 then
      let res = allocate_for_res res l1 in
      copy ~res:(d1, res, l1) (d1, eval1, l1)
    else
      let deg_result = max d1 d2 in
      let length_result = min l1 l2 in
      let res = allocate_for_res res length_result in
      Stubs.add res eval1 eval2 l1 l2;
      (deg_result, res, length_result)

  (*restrictive equality, the same polynomial evaluated
    on different domains are said to be different*)
  let equal (deg1, eval1, l1) (deg2, eval2, l2) =
    if deg1 <> deg2 || l1 <> l2 then false
    else Carray.Stubs.eq eval1 eval2 l1 l2

  let evaluation_fft_internal : Domain.t -> polynomial -> Carray.t =
   fun domain poly ->
    let degree = Polynomial.degree poly in
    let log_degree = Z.log2up (Z.of_int (degree + 1)) in
    let domain, n_domain = Domain.to_carray domain in
    let poly, size_poly = Polynomial.to_carray poly in
    let log = Z.(log2up @@ of_int n_domain) in
    if not Z.(log = log2 @@ of_int n_domain) then
      raise @@ Invalid_argument "Size of domain should be a power of 2.";
    (* See MR from Marc for padding *)
    if not (size_poly <= n_domain) then
      raise @@ Invalid_argument "Size of poly should be smaller than domain.";
    let dense_evaluation = Carray.allocate n_domain in
    Carray.Stubs.copy dense_evaluation poly 0 size_poly;
    Stubs.fft_inplace dense_evaluation ~domain ~log ~degree ~log_degree;
    (dense_evaluation, n_domain)

  let evaluation_fft : domain -> polynomial -> t =
   fun domain poly ->
    let d = Polynomial.degree poly in
    let domain_length = Domain.length domain in
    if d = -1 then (d, Carray.allocate domain_length, domain_length)
    else
      let res, len = evaluation_fft_internal domain poly in
      let d = Polynomial.degree poly in
      (d, res, len)

  let evaluation_fft2 : Domain.t -> polynomial -> scalar array =
   fun domain p ->
    let d = Polynomial.degree p in
    let domain_length = Domain.length domain in
    if d = -1 then Array.init domain_length (fun _ -> Fr.(copy zero))
    else Carray.to_array (evaluation_fft_internal domain p)

  let interpolation_fft_internal : Domain.t -> Carray.t -> polynomial =
   fun domain (coefficients, n_coefficients) ->
    let domain, n_domain = Domain.to_carray domain in
    let log = Z.(log2up @@ of_int n_domain) in
    if not Z.(log = log2 @@ of_int n_domain) then
      raise @@ Invalid_argument "Size of domain should be a power of 2.";
    if not (n_coefficients = n_domain) then
      raise @@ Invalid_argument "Size of coefficients should be same as domain.";
    Stubs.ifft_inplace coefficients ~domain ~log;
    Polynomial.of_carray (coefficients, n_coefficients)

  let interpolation_fft : domain -> t -> polynomial =
   fun domain (d, evaluation, l) ->
    if d = -1 then Polynomial.zero
    else
      let length_res = Domain.length domain in
      let rescaled_eval = Carray.allocate length_res in
      Domain_c.rescale rescaled_eval evaluation length_res l;
      interpolation_fft_internal domain (rescaled_eval, length_res)

  let interpolation_fft2 : Domain.t -> scalar array -> polynomial =
   fun domain coefficients ->
    interpolation_fft_internal domain (Carray.of_array coefficients)
end

include (
  Evaluations_c_impl :
    Evaluations_c_sig
      with type scalar = Bls12_381.Fr.t
       and type domain = Domain.t
       and type polynomial = Polynomial_c.t)