# Constantine
# Copyright (c) 2018-2019    Status Research & Development GmbH
# Copyright (c) 2020-Present Mamy André-Ratsimbazafy
# Licensed and distributed under either of
#   * MIT license (license terms in the root directory or at http://opensource.org/licenses/MIT).
#   * Apache v2 license (license terms in the root directory or at http://www.apache.org/licenses/LICENSE-2.0).
# at your option. This file may not be copied, modified, or distributed except according to those terms.

import
  # Internals
  ../constantine/[
    blssig_pop_on_bls12381_g2,
    ethereum_eip2333_bls12381_key_derivation],
  ../constantine/math/arithmetic,
  # Helpers
  ../helpers/prng_unsafe,
  ./bench_blueprint

proc separator*() = separator(167)

proc report(op, curve: string, startTime, stopTime: MonoTime, startClk, stopClk: int64, iters: int) =
  let ns = inNanoseconds((stopTime-startTime) div iters)
  let throughput = 1e9 / float64(ns)
  when SupportsGetTicks:
    echo &"{op:<75} {curve:<15} {throughput:>15.3f} ops/s     {ns:>9} ns/op     {(stopClk - startClk) div iters:>9} CPU cycles (approx)"
  else:
    echo &"{op:<75} {curve:<15} {throughput:>15.3f} ops/s     {ns:>9} ns/op"

template bench(op: string, curve: string, iters: int, body: untyped): untyped =
  measure(iters, startTime, stopTime, startClk, stopClk, body)
  report(op, curve, startTime, stopTime, startClk, stopClk, iters)

proc demoKeyGen(): tuple[seckey: SecretKey, pubkey: PublicKey] =
  # Don't do this at home, this is for benchmarking purposes
  # The RNG is NOT cryptographically secure
  # The API for keygen is not ready in blssig_pop_on_bls12381_g2
  let ikm = rng.random_byte_seq(32)
  doAssert cast[ptr BigInt[255]](result.seckey.addr)[].derive_master_secretKey(ikm)
  let ok = result.pubkey.derive_public_key(result.seckey)
  doAssert ok == cttBLS_Success

proc benchDeserPubkey*(iters: int) =
  let (sk, pk) = demoKeyGen()
  var pk_comp{.noInit.}: array[48, byte]

  # Serialize compressed
  let ok = pk_comp.serialize_public_key_compressed(pk)
  doAssert ok == cttBLS_Success

  var pk2{.noInit.}: PublicKey

  bench("Pubkey deserialization (full checks)", "BLS12_381 G1", iters):
    let status = pk2.deserialize_public_key_compressed(pk_comp)

proc benchDeserPubkeyUnchecked*(iters: int) =
  let (sk, pk) = demoKeyGen()
  var pk_comp{.noInit.}: array[48, byte]

  # Serialize compressed
  let ok = pk_comp.serialize_public_key_compressed(pk)
  doAssert ok == cttBLS_Success

  var pk2{.noInit.}: PublicKey

  bench("Pubkey deserialization (skip checks)", "BLS12_381 G1", iters):
    let status = pk2.deserialize_public_key_compressed_unchecked(pk_comp)

proc benchDeserSig*(iters: int) =
  let (sk, pk) = demoKeyGen()
  const msg = "abcdef0123456789"

  var
    sig_comp{.noInit.}: array[96, byte]
    sig {.noInit.}: Signature

  let status = sig.sign(sk, msg)
  doAssert status == cttBLS_Success

  # Serialize compressed
  let ok = sig_comp.serialize_signature_compressed(sig)
  doAssert ok == cttBLS_Success

  var sig2{.noInit.}: Signature

  bench("Signature deserialization (full checks)", "BLS12_381 G2", iters):
    let status = sig2.deserialize_signature_compressed(sig_comp)

proc benchDeserSigUnchecked*(iters: int) =
  let (sk, pk) = demoKeyGen()
  const msg = "abcdef0123456789"

  var
    sig_comp{.noInit.}: array[96, byte]
    sig {.noInit.}: Signature

  let status = sig.sign(sk, msg)
  doAssert status == cttBLS_Success

  # Serialize compressed
  let ok = sig_comp.serialize_signature_compressed(sig)
  doAssert ok == cttBLS_Success

  var sig2{.noInit.}: Signature

  bench("Signature deserialization (skip checks)", "BLS12_381 G2", iters):
    let status = sig2.deserialize_signature_compressed_unchecked(sig_comp)

proc benchSign*(iters: int) =
  let (sk, pk) = demoKeyGen()
  let msg = "Mr F was here"

  var sig: Signature

  bench("BLS signature", "BLS12_381 G2", iters):
    let status = sig.sign(sk, msg)

proc benchVerify*(iters: int) =
  let (sk, pk) = demoKeyGen()
  let msg = "Mr F was here"

  var sig: Signature
  let ok = sig.sign(sk, msg)
  doAssert ok == cttBLS_Success

  bench("BLS verification", "BLS12_381", iters):
    let valid = pk.verify(msg, sig)

proc benchFastAggregateVerify*(numKeys, iters: int) =
  ## Verification of N pubkeys signing 1 message
  let msg = "Mr F was here"

  var validators = newSeq[PublicKey](numKeys)
  var sigs = newSeq[Signature](numKeys)
  var aggSig: Signature

  for i in 0 ..< numKeys:
    let (sk, pk) = demoKeyGen()
    validators[i] = pk
    let status = sigs[i].sign(sk, msg)
    doAssert status == cttBLS_Success

  aggSig.aggregate_signatures(sigs)

  bench("BLS agg verif of 1 msg by " & $numKeys & " pubkeys", "BLS12_381", iters):
    let valid = validators.fast_aggregate_verify(msg, aggSig)

proc benchVerifyMulti*(numSigs, iters: int) =
  ## Verification of N pubkeys signing for N messages

  var triplets: seq[tuple[pubkey: PublicKey, msg: array[32, byte], sig: Signature]]

  var hashedMsg: array[32, byte]
  var sig: Signature

  for i in 0 ..< numSigs:
    let (sk, pk) = demoKeyGen()
    sha256.hash(hashedMsg, "msg" & $i)
    let status = sig.sign(sk, hashedMsg)
    doAssert status == cttBLS_Success
    triplets.add (pk, hashedMsg, sig)

  bench("BLS verif of " & $numSigs & " msgs by "& $numSigs & " pubkeys", "BLS12_381", iters):
    for i in 0 ..< triplets.len:
      let ok = triplets[i].pubkey.verify(triplets[i].msg, triplets[i].sig)
      doAssert ok == cttBLS_Success

proc benchVerifyBatched*(numSigs, iters: int) =
  ## Verification of N pubkeys signing for N messages

  var
    pubkeys: seq[PublicKey]
    messages: seq[array[32, byte]]
    signatures: seq[Signature]

  var hashedMsg: array[32, byte]
  var sig: Signature

  for i in 0 ..< numSigs:
    let (sk, pk) = demoKeyGen()
    sha256.hash(hashedMsg, "msg" & $i)
    let status = sig.sign(sk, hashedMsg)
    doAssert status == cttBLS_Success

    pubkeys.add pk
    messages.add hashedMsg
    signatures.add sig

  let secureBlindingBytes = sha256.hash("Mr F was here")

  bench("BLS serial batch verify of " & $numSigs & " msgs by "& $numSigs & " pubkeys (with blinding)", "BLS12_381", iters):
    let ok = batch_verify(pubkeys, messages, signatures, secureBlindingBytes)
    doAssert ok == cttBLS_Success

const Iters = 1000

proc main() =
  separator()
  benchDeserPubkey(Iters)
  benchDeserPubkeyUnchecked(Iters)
  benchDeserSig(Iters)
  benchDeserSigUnchecked(Iters)
  separator()
  benchSign(Iters)
  benchVerify(Iters)
  separator()
  benchFastAggregateVerify(numKeys = 128, iters = 10)
  separator()

  # Simulate Block verification (at most 6 signatures per block)
  benchVerifyMulti(numSigs = 6, iters = 10)
  benchVerifyBatched(numSigs = 6, iters = 10)
  separator()

  # Simulate 10 blocks verification
  benchVerifyMulti(numSigs = 60, iters = 10)
  benchVerifyBatched(numSigs = 60, iters = 10)
  separator()

  # Simulate 30 blocks verification
  benchVerifyMulti(numSigs = 180, iters = 10)
  benchVerifyBatched(numSigs = 180, iters = 10)
  separator()

main()
notes()