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

include Lang_core
module CS = Csir.CS
module VS = Linear_algebra.Make_VectorSpace (S)
module Tables = Csir.Tables

type wire = A | B | C [@@deriving repr]
type 'a tagged = Input of 'a | Output of 'a [@@deriving repr]

type arith_desc = {
  a : int;
  b : int;
  c : int;
  ql : S.t;
  qr : S.t;
  qm : S.t;
  qc : S.t;
  qo : S.t;
  to_solve : wire;
}
[@@deriving repr]

type pow5_desc = { a : int; c : int } [@@deriving repr]
type wires_desc = { a : int; b : int; c : int } [@@deriving repr]

type add5_desc = {
  o : int;
  x1 : int;
  x2 : int;
  x3 : int;
  x4 : int;
  x5 : int;
  k1 : S.t;
  k2 : S.t;
  k3 : S.t;
  k4 : S.t;
  k5 : S.t;
}
[@@deriving repr]

type lookup_desc = {
  a : int tagged;
  b : int tagged;
  c : int tagged;
  table : string;
}
[@@deriving repr]

type ws_desc = { x1 : int; y1 : int; x2 : int; y2 : int; x3 : int; y3 : int }
[@@deriving repr]

type ed_desc = {
  a : S.t;
  d : S.t;
  x1 : int;
  y1 : int;
  x2 : int;
  y2 : int;
  x3 : int;
  y3 : int;
}
[@@deriving repr]

type bits_desc = { nb_bits : int; shift : Utils.Z.t; l : int; bits : int list }
[@@deriving repr]

type pos128full_desc = {
  x0 : int;
  y0 : int;
  x1 : int;
  y1 : int;
  x2 : int;
  y2 : int;
  k : VS.t array;
  variant : Variants.t;
}

(* we define this by hand to avoid doing all linear algebra, VS.t is actually
   S.t but for some reason Repr does not see this equality. *)
let pos128full_desc_t =
  let open Repr in
  record "pos128full_desc" (fun x0 y0 x1 y1 x2 y2 k variant ->
      { x0; y0; x1; y1; x2; y2; k; variant })
  |+ field "x0" int (fun t -> t.x0)
  |+ field "y0" int (fun t -> t.y0)
  |+ field "x1" int (fun t -> t.x1)
  |+ field "y1" int (fun t -> t.y1)
  |+ field "x2" int (fun t -> t.x2)
  |+ field "y2" int (fun t -> t.y2)
  |+ field "k" (array S.t) (fun t -> t.k)
  |+ field "variant" Variants.t (fun t -> t.variant)
  |> sealr

type pos128partial_desc = {
  a : int;
  b : int;
  c : int;
  a_5 : int;
  b_5 : int;
  c_5 : int;
  x0 : int;
  y0 : int;
  x1 : int;
  y1 : int;
  x2 : int;
  y2 : int;
  (* Can we share these? *)
  k_cols : VS.matrix array;
  variant : Variants.t;
}

let pos128partial_desc_t =
  let open Repr in
  record "pos128partial_desc"
    (fun a b c a_5 b_5 c_5 x0 y0 x1 y1 x2 y2 k_cols variant ->
      { a; b; c; a_5; b_5; c_5; x0; y0; x1; y1; x2; y2; k_cols; variant })
  |+ field "a" int (fun t -> t.a)
  |+ field "b" int (fun t -> t.b)
  |+ field "c" int (fun t -> t.c)
  |+ field "a_5" int (fun t -> t.a_5)
  |+ field "b_5" int (fun t -> t.b_5)
  |+ field "c_5" int (fun t -> t.c_5)
  |+ field "x0" int (fun t -> t.x0)
  |+ field "y0" int (fun t -> t.y0)
  |+ field "x1" int (fun t -> t.x1)
  |+ field "y1" int (fun t -> t.y1)
  |+ field "x2" int (fun t -> t.x2)
  |+ field "y2" int (fun t -> t.y2)
  |+ field "k_cols" (array (array (array S.t))) (fun t -> t.k_cols)
  |+ field "variant" Variants.t (fun t -> t.variant)
  |> sealr

type anemoi_desc = {
  x0 : int;
  y0 : int;
  w : int;
  v : int;
  x1 : int;
  y1 : int;
  kx : S.t;
  ky : S.t;
}

let anemoi_desc_t =
  let open Repr in
  record "anemoi_desc" (fun x0 y0 w v x1 y1 kx ky ->
      { x0; y0; w; v; x1; y1; kx; ky })
  |+ field "x0" int (fun t -> t.x0)
  |+ field "y0" int (fun t -> t.y0)
  |+ field "w" int (fun t -> t.w)
  |+ field "v" int (fun t -> t.v)
  |+ field "x1" int (fun t -> t.x1)
  |+ field "y1" int (fun t -> t.y1)
  |+ field "kx" S.t (fun t -> t.kx)
  |+ field "ky" S.t (fun t -> t.ky)
  |> sealr

type anemoi_double_desc = {
  x0 : int;
  y0 : int;
  w0 : int;
  w1 : int;
  y1 : int;
  x2 : int;
  y2 : int;
  kx1 : S.t;
  ky1 : S.t;
  kx2 : S.t;
  ky2 : S.t;
}

let anemoi_double_desc_t =
  let open Repr in
  record "anemoi_double_desc" (fun x0 y0 w0 w1 y1 x2 y2 kx1 ky1 kx2 ky2 ->
      { x0; y0; w0; w1; y1; x2; y2; kx1; ky1; kx2; ky2 })
  |+ field "x0" int (fun t -> t.x0)
  |+ field "y0" int (fun t -> t.y0)
  |+ field "w0" int (fun t -> t.w0)
  |+ field "w1" int (fun t -> t.w1)
  |+ field "y1" int (fun t -> t.y1)
  |+ field "x2" int (fun t -> t.x2)
  |+ field "y2" int (fun t -> t.y2)
  |+ field "kx1" S.t (fun t -> t.kx1)
  |+ field "kx2" S.t (fun t -> t.kx2)
  |+ field "ky1" S.t (fun t -> t.ky1)
  |+ field "ky2" S.t (fun t -> t.ky2)
  |> sealr

type solver_desc =
  | Arith of arith_desc
  | Pow5 of pow5_desc
  | IsZero of wires_desc
  | IsNotZero of wires_desc
  | Add5 of add5_desc
  | Lookup of lookup_desc
  | Ecc_Ws of ws_desc
  | Ecc_Ed of ed_desc
  | Skip
  | BitsOfS of bits_desc
  | Poseidon128Full of pos128full_desc
  | Poseidon128Partial of pos128partial_desc
  | AnemoiRound of anemoi_desc
  | AnemoiDoubleRound of anemoi_double_desc
  | Updater of Optimizer.trace_info
[@@deriving repr]

type solvers = solver_desc list [@@deriving repr]

type t = { solvers : solvers; initial_size : int; final_size : int }
[@@deriving repr]

let empty_solver = { solvers = []; initial_size = 0; final_size = 0 }
let append_solver sd t = { t with solvers = sd :: t.solvers }
let untag = function Input a -> a | Output a -> a
let from_tagged = function Input i -> Some i | Output _ -> None
let map3 f (a, b, c) = (f a, f b, f c)

let solve_one trace solver =
  (match solver with
  | Skip -> ()
  | Arith { a; b; c; ql; qr; qm; qc; qo; to_solve } -> (
      match to_solve with
      | C ->
          let av = trace.(a) in
          let bv = trace.(b) in
          trace.(c) <-
            S.(((ql * av) + (qr * bv) + (qm * av * bv) + qc) / negate qo)
      | A ->
          let cv = trace.(c) in
          let bv = trace.(b) in
          trace.(a) <-
            S.(negate ((qr * bv) + (qo * cv) + qc) / (ql + (bv * qm)))
      | B ->
          let cv = trace.(c) in
          let av = trace.(a) in
          trace.(b) <-
            S.(negate ((ql * av) + (qo * cv) + qc) / (qr + (av * qm))))
  | Pow5 { a; c } -> trace.(c) <- S.pow trace.(a) (Z.of_int 5)
  | Lookup { a; b; c; table } ->
      let tbl = Tables.find table Csir.table_registry in
      let wa, wb, wc = map3 untag (a, b, c) in
      let a, b, c = map3 from_tagged (a, b, c) in
      let a, b, c = map3 (Option.map (fun i -> trace.(i))) (a, b, c) in
      let entry = Option.get Csir.Table.(find { a; b; c } tbl) in
      trace.(wa) <- entry.a;
      trace.(wb) <- entry.b;
      trace.(wc) <- entry.c
  | IsZero { a; b; c } ->
      let av = trace.(a) in
      trace.(c) <- S.(if av = zero then one else zero);
      trace.(b) <- S.(if av = zero then one else S.div_exn one av)
  | IsNotZero { a; b; c } ->
      let av = trace.(a) in
      trace.(c) <- S.(if av = zero then zero else one);
      trace.(b) <- S.(if av = zero then one else S.div_exn one av)
  | Add5 { o; x1; x2; x3; x4; x5; k1; k2; k3; k4; k5 } ->
      let x1, x2, x3, x4, x5 =
        (trace.(x1), trace.(x2), trace.(x3), trace.(x4), trace.(x5))
      in
      trace.(o) <- S.((k1 * x1) + (k2 * x2) + (k3 * x3) + (k4 * x4) + (k5 * x5))
  | Ecc_Ws { x1; y1; x2; y2; x3; y3 } ->
      let x1, y1 = (trace.(x1), trace.(y1)) in
      let x2, y2 = (trace.(x2), trace.(y2)) in
      let lambda = S.(sub y2 y1 / sub x2 x1) in
      let x3_v = S.(sub (lambda * lambda) (x1 + x2)) in
      trace.(x3) <- x3_v;
      trace.(y3) <- S.(sub (lambda * sub x1 x3_v) y1)
  | Ecc_Ed { a; d; x1; y1; x2; y2; x3; y3 } ->
      let x1, y1 = (trace.(x1), trace.(y1)) in
      let x2, y2 = (trace.(x2), trace.(y2)) in
      let x1x2 = S.(mul x1 x2) in
      let y1y2 = S.(mul y1 y2) in
      let denom = S.(d * x1x2 * y1y2) in
      let x_res = S.(add (x1 * y2) (x2 * y1) / add one denom) in
      let y_res = S.(sub y1y2 (a * x1x2) / sub one denom) in
      trace.(x3) <- x_res;
      trace.(y3) <- y_res
  | BitsOfS { nb_bits; shift; l; bits } ->
      let x = trace.(l) |> S.to_z in
      let x = Z.(x + shift) in
      let binary_decomposition = Utils.bool_list_of_z ~nb_bits x in
      List.iter2
        (fun b value -> trace.(b) <- (if value then S.one else S.zero))
        bits binary_decomposition
  | Updater ti -> ignore @@ Optimizer.trace_updater ti trace
  | Poseidon128Full { x0; y0; x1; y1; x2; y2; k; variant } ->
      let matrix =
        Array.map (Array.map S.of_string)
        @@ if variant = PFull128 then Mds_full.v else Mds_128.v
      in
      let pow5 x = S.pow trace.(x) (Z.of_int 5) in
      let x_vec = [| Array.map pow5 [| x0; x1; x2 |] |] |> VS.transpose in
      let y_vec = VS.mul matrix x_vec in
      List.iteri
        (fun i yi -> trace.(yi) <- S.add k.(i) @@ y_vec.(i).(0))
        [ y0; y1; y2 ]
  | Poseidon128Partial
      { a; b; c; a_5; b_5; c_5; x0; y0; x1; y1; x2; y2; k_cols; variant } ->
      let matrix =
        Array.map (Array.map S.of_string)
        @@ if variant = PFull128 then Mds_full.v else Mds_128.v
      in
      let pow5 x = S.pow x (Z.of_int 5) in
      let ppow5 v = [| v.(0); v.(1); [| pow5 v.(2).(0) |] |] in
      let x_vec = [| [| trace.(x0) |]; [| trace.(x1) |]; [| trace.(x2) |] |] in
      let a_vec = VS.(add (mul matrix @@ ppow5 x_vec) k_cols.(0)) in
      let b_vec = VS.(add (mul matrix @@ ppow5 a_vec) k_cols.(1)) in
      let c_vec = VS.(add (mul matrix @@ ppow5 b_vec) k_cols.(2)) in
      let y_vec = VS.(add (mul matrix @@ ppow5 c_vec) k_cols.(3)) in
      trace.(a) <- a_vec.(2).(0);
      trace.(b) <- b_vec.(2).(0);
      trace.(c) <- c_vec.(2).(0);
      trace.(a_5) <- pow5 trace.(a);
      trace.(b_5) <- pow5 trace.(b);
      trace.(c_5) <- pow5 trace.(c);
      trace.(y0) <- y_vec.(0).(0);
      trace.(y1) <- y_vec.(1).(0);
      trace.(y2) <- y_vec.(2).(0)
  | AnemoiRound { x0; y0; w; v; x1; y1; kx; ky } ->
      let open Bls12_381_hash.Anemoi.Parameters in
      let g2_p_1 = S.((g * g) + one) in
      let _w_5', w', v', _u', x1', y1' =
        Anemoi_jive128_1.compute_one_round
          (module S : Ff_sig.PRIME with type t = S.t)
          beta gamma delta g alpha_inv g2_p_1 trace.(x0) trace.(y0) kx ky
      in
      trace.(w) <- w';
      trace.(v) <- v';
      trace.(x1) <- x1';
      trace.(y1) <- y1'
  | AnemoiDoubleRound { x0; y0; w0; w1; y1; x2; y2; kx1; ky1; kx2; ky2 } ->
      (* First round *)
      let open Bls12_381_hash.Anemoi.Parameters in
      let g2_p_1 = S.((g * g) + one) in
      let _w_5', w', _v', _u', x1', y1' =
        Anemoi_jive128_1.compute_one_round
          (module S : Ff_sig.PRIME with type t = S.t)
          beta gamma delta g alpha_inv g2_p_1 trace.(x0) trace.(y0) kx1 ky1
      in
      (* Computing w *)
      trace.(w0) <- w';
      trace.(y1) <- y1';
      (* Second round *)
      let open Bls12_381_hash.Anemoi.Parameters in
      let g2_p_1 = S.((g * g) + one) in
      let _w_5', w', _v', _u', x2', y2' =
        Anemoi_jive128_1.compute_one_round
          (module S : Ff_sig.PRIME with type t = S.t)
          beta gamma delta g alpha_inv g2_p_1 x1' y1' kx2 ky2
      in
      trace.(w1) <- w';
      trace.(x2) <- x2';
      trace.(y2) <- y2');
  trace

let solve : t -> S.t array -> S.t array =
 fun { solvers; initial_size; final_size } inputs ->
  if Array.length inputs <> initial_size then
    failwith
      (Printf.sprintf "input size (= %d) != initial_size (= %d)"
         (Array.length inputs) initial_size);
  let dummy =
    Array.(append inputs (init (final_size - length inputs) (fun _ -> S.zero)))
  in
  List.fold_left solve_one dummy (List.rev solvers)