Source file class_group.ml

open Utils

module Stubs = struct
  external caml_check_bytes_stubs :
    Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> bool = "caml_check_bytes_stubs"

  external zero : Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> unit
    = "caml_zero_stubs"

  external one : Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> unit
    = "caml_one_stubs"

  external eq : Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> Bytes.t -> int
    = "caml_eq_stubs"

  external neg : Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> Bytes.t -> int
    = "caml_neg_stubs"

  external add :
    Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> Bytes.t -> Bytes.t -> int
    = "caml_add_stubs"

  external double : Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> Bytes.t -> int
    = "caml_double_stubs"

  external mul :
    Bytes.t -> Unsigned.Size_t.t -> Bytes.t -> Bytes.t -> int -> Bytes.t -> int
    = "caml_mul_bytecode_stubs" "caml_mul_stubs"
end

module type Class_group_sig = sig
  (* BE CAREFUL : GET THE CANONICAL REPRESENTATION !!! *)
  include S.GROUP

  val discriminant : Integer.t

  val discriminant_size : Unsigned.Size_t.t
  (* val of_coefficients : Integer.t * Integer.t * Integer.t -> t

     val get_coefficients : t -> Integer.t * Integer.t * Integer.t

     val get_discriminant : t -> Integer.t *)
end

module Make (D : sig
  val discriminant : Integer.t
end) : Class_group_sig = struct
  type t = Bytes.t

  let invalid_discriminant_size =
    Invalid_argument
      "Invalid discriminant: current code only support discriminant of size 128"

  let discriminant_size_int = Integer.size_in_bytes D.discriminant

  let discriminant =
    (* Constraint from Chia's cpp code *)
    if discriminant_size_int = 128 then D.discriminant
    else raise invalid_discriminant_size

  let discriminant_size = Unsigned.Size_t.of_int discriminant_size_int

  let discriminant_bytes = Integer.to_bytes D.discriminant

  (* Constraint from Chia's cpp code *)
  let size_in_bytes = 100

  let check_bytes input =
    Stubs.caml_check_bytes_stubs discriminant_bytes discriminant_size input

  let of_bytes_exn input =
    if check_bytes input then input else failwith "Form not in group"

  let of_bytes_opt input = if check_bytes input then Some input else None

  let to_bytes f = f

  (* let of_coefficients _a _b _c = let id = Bytes.create size_in_bytes in
     Stubs.identity disc discriminant_size id ; id

     (* TODO *) let get_coefficients _f = let zero = Integer.of_bytes
     (Bytes.create 0) in (zero, zero, zero)

     (* TODO *) let get_discriminant _f = Integer.of_bytes (Bytes.create 0) *)

  let size_in_memory = 100

  let zero =
    let buffer = Bytes.create size_in_bytes in
    Stubs.zero discriminant_bytes discriminant_size buffer ;
    buffer

  let one =
    let buffer = Bytes.create size_in_bytes in
    Stubs.one discriminant_bytes discriminant_size buffer ;
    buffer

  let eq f1 f2 =
    let status = Stubs.eq discriminant_bytes discriminant_size f1 f2 in
    match status with
    | x when x = 0 -> false
    | x when x = 1 -> true
    | x when x = 2 -> raise Utils.invalid_group_element
    | _ -> raise Utils.unknown_error

  let is_zero f = eq f zero

  let is_one f = eq f one

  let neg f =
    let buffer = Bytes.create size_in_bytes in
    let status = Stubs.neg discriminant_bytes discriminant_size f buffer in
    match status with
    | x when x = 0 -> buffer
    | x when x = 2 -> raise Utils.invalid_group_element
    | _ -> raise Utils.unknown_error

  let add f1 f2 =
    let buffer = Bytes.create size_in_bytes in
    let status = Stubs.add discriminant_bytes discriminant_size f1 f2 buffer in
    match status with
    | x when x = 0 -> buffer
    | x when x = 2 -> raise Utils.invalid_group_element
    | _ -> raise Utils.unknown_error

  let double f =
    let buffer = Bytes.create size_in_bytes in
    let status = Stubs.double discriminant_bytes discriminant_size f buffer in
    match status with
    | x when x = 0 -> buffer
    | x when x = 2 -> raise Utils.invalid_group_element
    | _ -> raise Utils.unknown_error

  let mul form exponent =
    let rec aux f n =
      if Z.(equal n zero) then zero
      else if Z.(equal n one) then f
      else if Z.(lt n zero) then aux (neg f) Z.(neg n)
      else
        let buffer = Bytes.create size_in_bytes in
        let exp_bits = Z.to_bits n |> Bytes.unsafe_of_string in
        let exp_len = Z.numbits n in
        let status =
          Stubs.mul
            discriminant_bytes
            discriminant_size
            f
            exp_bits
            exp_len
            buffer
        in
        match status with
        | x when x = 0 -> buffer
        | x when x = 2 -> raise Utils.invalid_group_element
        | _ -> raise Utils.unknown_error
    in
    aux form exponent

  let random ?state () =
    let () =
      match state with
      | None -> Random.self_init ()
      | Some s -> Random.set_state s
    in
    let mul_bytes form exp_bits exp_len =
      let buffer = Bytes.create size_in_bytes in
      let status =
        Stubs.mul
          discriminant_bytes
          discriminant_size
          form
          exp_bits
          exp_len
          buffer
      in
      match status with
      | x when x = 0 -> buffer
      | x when x = 2 -> raise Utils.invalid_group_element
      | _ -> raise Utils.unknown_error
    in
    let len_exp = discriminant_size_int / 2 in
    let random_exp =
      Bytes.init len_exp (fun _ -> Random.int 256 |> char_of_int)
    in
    mul_bytes one random_exp len_exp
end