Source file gadget_anemoi.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.                                                 *)
(*                                                                           *)
(*****************************************************************************)

open Lang_core
open Lang_stdlib

module Make (L : LIB) = struct
  open L

  let nb_rounds, _state_size, matrix, rc =
    Bls12_381_hash.Anemoi.Parameters.state_size_1

  let rc =
    Array.init
      (Array.length rc + 2)
      (fun i ->
        (* We pad with two zeroes *)
        if i / 2 = nb_rounds then Bls12_381.Fr.zero
        else if i mod 2 = 0 then rc.(i / 2)
        else rc.(nb_rounds + (i / 2)))

  let rec repeat : n:int -> ('a -> 'a t) -> 'a -> 'a t =
   fun ~n f e ->
    if n <= 0 then ret e
    else
      let* x = f e in
      repeat ~n:(n - 1) f x

  let round :
      scalar repr * scalar repr * int -> (scalar repr * scalar repr * int) t =
   fun (xi, yi, i) ->
    let kx, ky = (rc.(i), rc.(i + 1)) in
    with_label ~label:"Anemoi.round"
    @@ let* res = Anemoi.anemoi_round ~kx ~ky (xi, yi) in
       let xj, yj = of_pair res in
       ret @@ (xj, yj, i + 2)

  let init_state_for_rounds x0 y0 =
    let* x00 =
      Num.add
        ~ql:matrix.(0).(0)
        ~qr:matrix.(0).(1)
        ~qc:S.((matrix.(0).(0) * rc.(0)) + (matrix.(0).(1) * rc.(1)))
        x0 y0
    in
    let* y00 =
      Num.add
        ~ql:matrix.(1).(0)
        ~qr:matrix.(1).(1)
        ~qc:S.((matrix.(1).(0) * rc.(0)) + (matrix.(1).(1) * rc.(1)))
        x0 y0
    in
    ret (x00, y00)

  let compress : scalar repr -> scalar repr -> scalar repr t =
   fun x0 y0 ->
    with_label ~label:"Anemoi.compress"
    @@ let* x00, y00 = init_state_for_rounds x0 y0 in
       let* xn, yn, _i = repeat ~n:nb_rounds round (x00, y00, 2) in
       Num.add5 ~k5:S.zero xn yn x0 y0 y0

  let double_round :
      scalar repr * scalar repr * int -> (scalar repr * scalar repr * int) t =
   fun (xi, yi, i) ->
    let kx1, ky1 = (rc.(i), rc.(i + 1)) in
    let kx2, ky2 = (rc.(i + 2), rc.(i + 3)) in
    with_label ~label:"Anemoi.double_round"
    @@ let* res = Anemoi.anemoi_double_round ~kx1 ~ky1 ~kx2 ~ky2 (xi, yi) in
       let xj, yj = of_pair res in
       ret @@ (xj, yj, i + 4)

  let compress_two : scalar repr -> scalar repr -> scalar repr t =
   fun x0 y0 ->
    with_label ~label:"Anemoi.compress_two"
    @@ let* x00, y00 = init_state_for_rounds x0 y0 in
       let* xn, yn, n = repeat ~n:(nb_rounds / 2) double_round (x00, y00, 2) in
       let* xnn, ynn, _ =
         if nb_rounds mod 2 = 0 then ret (xn, yn, 0) else round (xn, yn, n)
       in
       Num.add5 ~k5:S.zero xnn ynn x0 y0 y0

  let digest : ?input_length:int -> scalar list repr -> scalar repr t =
   fun ?input_length:_ inputs ->
    match of_list inputs with
    | [] ->
        constant_scalar (Bls12_381_hash.Anemoi.jive128_1_compress S.zero S.zero)
    | [ x ] ->
        let* zero = constant_scalar S.zero in
        compress_two zero x
    | x :: rest -> foldM compress_two x rest
end

module Anemoi128 = struct
  module P : Hash_sig.P_HASH = struct
    type scalar = S.t

    let direct ?input_length:_ inputs =
      match Array.to_list inputs with
      | [] -> Bls12_381_hash.Anemoi.jive128_1_compress S.zero S.zero
      | [ x ] -> Bls12_381_hash.Anemoi.jive128_1_compress S.zero x
      | x :: rest ->
          List.fold_left Bls12_381_hash.Anemoi.jive128_1_compress x rest

    type ctxt = S.t

    let init ?input_length:_ () = S.zero

    let digest ctxt inputs =
      ignore ctxt;
      direct inputs

    let get ctxt = ctxt
  end

  module V : Hash_sig.HASH = Make
end