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
type 'a tagged = Input of 'a | Output of 'a

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;
}

type pow5_desc = { a : int; c : int }
type wires_desc = { a : int; b : int; c : int }
type add5_desc = { o : int; x1 : int; x2 : int; x3 : int; x4 : int; x5 : int }

type lookup_desc = {
  a : int tagged;
  b : int tagged;
  c : int tagged;
  table : string;
}

type ws_desc = { x1 : int; y1 : int; x2 : int; y2 : int; x3 : int; y3 : int }

type ed_desc = {
  a : S.t;
  d : S.t;
  x1 : int;
  y1 : int;
  x2 : int;
  y2 : int;
  x3 : int;
  y3 : int;
}

type bits_desc = { nb_bits : int; shift : Z.t; l : int; bits : int list }

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

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;
}

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
  | Updater of Optimizer.trace_info

type solvers = solver_desc list
type t = { solvers : solvers; initial_size : int; final_size : int }

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 t = 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 } t) 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 } ->
      let x1, x2, x3, x4, x5 =
        (trace.(x1), trace.(x2), trace.(x3), trace.(x4), trace.(x5))
      in
      trace.(o) <- S.(x1 + x2 + x3 + x4 + 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));
  trace

let solve : t -> S.t array -> S.t array =
 fun { solvers; initial_size; final_size } inputs ->
  assert (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)

(* Encodings *)

let wire_encoding : wire Data_encoding.encoding =
  Data_encoding.(
    union ~tag_size:`Uint8
      [
        case (Tag 0) ~title:"A" unit
          (function A -> Some () | _ -> None)
          (fun () -> A);
        case (Tag 1) ~title:"B" unit
          (function B -> Some () | _ -> None)
          (fun () -> B);
        case (Tag 2) ~title:"C" unit
          (function C -> Some () | _ -> None)
          (fun () -> C);
      ])

let scalar_enc = Csir.Scalar.encoding

let arith_desc_encoding : arith_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun { a; b; c; ql; qr; qm; qc; qo; to_solve } ->
        (a, b, c, ql, qr, qm, qc, qo, to_solve))
      (fun (a, b, c, ql, qr, qm, qc, qo, to_solve) ->
        { a; b; c; ql; qr; qm; qc; qo; to_solve })
      (obj9 (req "a" int31) (req "b" int31) (req "c" int31)
         (req "ql" scalar_enc) (req "qr" scalar_enc) (req "qm" scalar_enc)
         (req "qc" scalar_enc) (req "qo" scalar_enc)
         (req "to_sove" wire_encoding)))

let pow5_desc_encoding : pow5_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ a; c } : pow5_desc) -> (a, c))
      (fun (a, c) -> { a; c })
      (obj2 (req "a" int31) (req "c" int31)))

let wires_desc_encoding : wires_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ a; b; c } : wires_desc) -> (a, b, c))
      (fun (a, b, c) -> { a; b; c })
      (obj3 (req "a" int31) (req "b" int31) (req "c" int31)))

let add5_desc_encoding : add5_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ o; x1; x2; x3; x4; x5 } : add5_desc) -> (o, x1, x2, x3, x4, x5))
      (fun (o, x1, x2, x3, x4, x5) -> { o; x1; x2; x3; x4; x5 })
      (obj6 (req "o" int31) (req "x1" int31) (req "x2" int31) (req "x3" int31)
         (req "x4" int31) (req "x5" int31)))

let tagged_int_encoding : int tagged Data_encoding.encoding =
  Data_encoding.(
    union ~tag_size:`Uint8
      [
        case (Tag 0) ~title:"Input" int31
          (function Input i -> Some i | _ -> None)
          (fun i -> Input i);
        case (Tag 1) ~title:"Output" int31
          (function Output i -> Some i | _ -> None)
          (fun i -> Output i);
      ])

let lookup_desc_encoding : lookup_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ a; b; c; table } : lookup_desc) -> (a, b, c, table))
      (fun (a, b, c, table) -> { a; b; c; table })
      (obj4
         (req "a" tagged_int_encoding)
         (req "b" tagged_int_encoding)
         (req "c" tagged_int_encoding)
         (req "table" string)))

let ws_desc_encoding : ws_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ x1; y1; x2; y2; x3; y3 } : ws_desc) -> (x1, y1, x2, y2, x3, y3))
      (fun (x1, y1, x2, y2, x3, y3) -> { x1; y1; x2; y2; x3; y3 })
      (obj6 (req "x1" int31) (req "y1" int31) (req "x2" int31) (req "y2" int31)
         (req "x3" int31) (req "y3" int31)))

let ed_desc_encoding : ed_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ a; d; x1; y1; x2; y2; x3; y3 } : ed_desc) ->
        (a, d, x1, y1, x2, y2, x3, y3))
      (fun (a, d, x1, y1, x2, y2, x3, y3) -> { a; d; x1; y1; x2; y2; x3; y3 })
      (obj8 (req "a" scalar_enc) (req "d" scalar_enc) (req "x1" int31)
         (req "y1" int31) (req "x2" int31) (req "y2" int31) (req "x3" int31)
         (req "y3" int31)))

let bits_desc_encoding : bits_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun { nb_bits; shift; l; bits } -> (nb_bits, shift, l, bits))
      (fun (nb_bits, shift, l, bits) -> { nb_bits; shift; l; bits })
      (obj4 (req "nb_bits" int31) (req "shift" z) (req "l" int31)
         (req "bits" (list int31))))

let pos128full_desc_encoding : pos128full_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ x0; y0; x1; y1; x2; y2; k; variant } : pos128full_desc) ->
        (x0, y0, x1, y1, x2, y2, k, variant))
      (fun (x0, y0, x1, y1, x2, y2, k, variant) ->
        { x0; y0; x1; y1; x2; y2; k; variant })
      (obj8 (req "x0" int31) (req "y0" int31) (req "x1" int31) (req "y1" int31)
         (req "x2" int31) (req "y2" int31)
         (req "k" (array scalar_enc))
         (req "variant" Variants.encoding)))

let pos128partial_desc_encoding : pos128partial_desc Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun ({ a; b; c; a_5; b_5; c_5; x0; y0; x1; y1; x2; y2; k_cols; variant } :
             pos128partial_desc) ->
        ((a, b, c, variant), (a_5, b_5, c_5, x0, y0, x1, y1, x2, y2, k_cols)))
      (fun ((a, b, c, variant), (a_5, b_5, c_5, x0, y0, x1, y1, x2, y2, k_cols)) ->
        { a; b; c; a_5; b_5; c_5; x0; y0; x1; y1; x2; y2; k_cols; variant })
      (merge_objs
         (obj4 (req "a" int31) (req "b" int31) (req "c" int31)
            (req "variant" Variants.encoding))
         (obj10 (req "a_5" int31) (req "b_5" int31) (req "c_5" int31)
            (req "x0" int31) (req "y0" int31) (req "x1" int31) (req "y1" int31)
            (req "x2" int31) (req "y2" int31)
            (req "k_cols" (array (array (array scalar_enc)))))))

let solver_desc_encoding : solver_desc Data_encoding.encoding =
  Data_encoding.(
    union ~tag_size:`Uint8
      [
        case (Tag 0) ~title:"arith" arith_desc_encoding
          (function Arith desc -> Some desc | _ -> None)
          (fun desc -> Arith desc);
        case (Tag 1) ~title:"pow5" pow5_desc_encoding
          (function Pow5 desc -> Some desc | _ -> None)
          (fun desc -> Pow5 desc);
        case (Tag 2) ~title:"iszero" wires_desc_encoding
          (function IsZero desc -> Some desc | _ -> None)
          (fun desc -> IsZero desc);
        case (Tag 3) ~title:"isnotzero" wires_desc_encoding
          (function IsNotZero desc -> Some desc | _ -> None)
          (fun desc -> IsNotZero desc);
        case (Tag 4) ~title:"add5" add5_desc_encoding
          (function Add5 desc -> Some desc | _ -> None)
          (fun desc -> Add5 desc);
        case (Tag 5) ~title:"lookup" lookup_desc_encoding
          (function Lookup desc -> Some desc | _ -> None)
          (fun desc -> Lookup desc);
        case (Tag 6) ~title:"ecc_ws" ws_desc_encoding
          (function Ecc_Ws desc -> Some desc | _ -> None)
          (fun desc -> Ecc_Ws desc);
        case (Tag 7) ~title:"ecc_ed" ed_desc_encoding
          (function Ecc_Ed desc -> Some desc | _ -> None)
          (fun desc -> Ecc_Ed desc);
        case (Tag 8) ~title:"skip" unit
          (function Skip -> Some () | _ -> None)
          (fun () -> Skip);
        case (Tag 9) ~title:"bits" bits_desc_encoding
          (function BitsOfS desc -> Some desc | _ -> None)
          (fun desc -> BitsOfS desc);
        case (Tag 10) ~title:"pos128full" pos128full_desc_encoding
          (function Poseidon128Full desc -> Some desc | _ -> None)
          (fun desc -> Poseidon128Full desc);
        case (Tag 11) ~title:"pos128partial" pos128partial_desc_encoding
          (function Poseidon128Partial desc -> Some desc | _ -> None)
          (fun desc -> Poseidon128Partial desc);
        case (Tag 12) ~title:"updater" Optimizer.trace_info_encoding
          (function Updater desc -> Some desc | _ -> None)
          (fun desc -> Updater desc);
      ])

let solver_encoding : t Data_encoding.encoding =
  Data_encoding.(
    conv
      (fun { solvers; initial_size; final_size } ->
        (solvers, initial_size, final_size))
      (fun (solvers, initial_size, final_size) ->
        { solvers; initial_size; final_size })
      (obj3
         (req "solvers" @@ list solver_desc_encoding)
         (req "initial_size" int31) (req "final_size" int31)))