From ef667bb4041aecfda4f2cd239450d91116c78c44 Mon Sep 17 00:00:00 2001
From: Paul Miller <paul@paulmillr.com>
Date: Mon, 21 Aug 2023 15:48:34 +0000
Subject: [PATCH] poseidon: refactor validateOpts, fix tests

---
 src/abstract/poseidon.ts | 59 ++++++++++++++++++++--------------------
 test/poseidon.test.js    |  4 +++
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/src/abstract/poseidon.ts b/src/abstract/poseidon.ts
index bf326c0..2c986dc 100644
--- a/src/abstract/poseidon.ts
+++ b/src/abstract/poseidon.ts
@@ -15,33 +15,36 @@ export type PoseidonOpts = {
 };
 
 export function validateOpts(opts: PoseidonOpts) {
-  const { Fp } = opts;
+  const { Fp, mds, reversePartialPowIdx: rev, roundConstants: rc } = opts;
+  const { roundsFull, roundsPartial, sboxPower, t } = opts;
+
   validateField(Fp);
   for (const i of ['t', 'roundsFull', 'roundsPartial'] as const) {
     if (typeof opts[i] !== 'number' || !Number.isSafeInteger(opts[i]))
       throw new Error(`Poseidon: invalid param ${i}=${opts[i]} (${typeof opts[i]})`);
   }
-  const rev = opts.reversePartialPowIdx;
+
+  // MDS is TxT matrix
+  if (!Array.isArray(mds) || mds.length !== t) throw new Error('Poseidon: wrong MDS matrix');
+  const _mds = mds.map((mdsRow) => {
+    if (!Array.isArray(mdsRow) || mdsRow.length !== t)
+      throw new Error(`Poseidon MDS matrix row: ${mdsRow}`);
+    return mdsRow.map((i) => {
+      if (typeof i !== 'bigint') throw new Error(`Poseidon MDS matrix value=${i}`);
+      return Fp.create(i);
+    });
+  });
+
   if (rev !== undefined && typeof rev !== 'boolean')
     throw new Error(`Poseidon: invalid param reversePartialPowIdx=${rev}`);
-  let { sboxPower } = opts;
-  if (!sboxPower || ![3, 5, 7].includes(sboxPower))
-    throw new Error(`Poseidon wrong sboxPower=${sboxPower}`);
 
-  const _sboxPower = BigInt(sboxPower);
-  let sboxFn = (n: bigint) => FpPow(Fp, n, _sboxPower);
-  // Unwrapped sbox power for common cases (195->142μs)
-  if (sboxPower === 3) sboxFn = (n: bigint) => Fp.mul(Fp.sqrN(n), n);
-  else if (sboxPower === 5) sboxFn = (n: bigint) => Fp.mul(Fp.sqrN(Fp.sqrN(n)), n);
+  if (roundsFull % 2 !== 0) throw new Error(`Poseidon roundsFull is not even: ${roundsFull}`);
+  const rounds = roundsFull + roundsPartial;
 
-  if (opts.roundsFull % 2 !== 0)
-    throw new Error(`Poseidon roundsFull is not even: ${opts.roundsFull}`);
-  const rounds = opts.roundsFull + opts.roundsPartial;
-
-  if (!Array.isArray(opts.roundConstants) || opts.roundConstants.length !== rounds)
+  if (!Array.isArray(rc) || rc.length !== rounds)
     throw new Error('Poseidon: wrong round constants');
-  const roundConstants = opts.roundConstants.map((rc) => {
-    if (!Array.isArray(rc) || rc.length !== opts.t)
+  const roundConstants = rc.map((rc) => {
+    if (!Array.isArray(rc) || rc.length !== t)
       throw new Error(`Poseidon wrong round constants: ${rc}`);
     return rc.map((i) => {
       if (typeof i !== 'bigint' || !Fp.isValid(i))
@@ -49,18 +52,16 @@ export function validateOpts(opts: PoseidonOpts) {
       return Fp.create(i);
     });
   });
-  // MDS is TxT matrix
-  if (!Array.isArray(opts.mds) || opts.mds.length !== opts.t)
-    throw new Error('Poseidon: wrong MDS matrix');
-  const mds = opts.mds.map((mdsRow) => {
-    if (!Array.isArray(mdsRow) || mdsRow.length !== opts.t)
-      throw new Error(`Poseidon MDS matrix row: ${mdsRow}`);
-    return mdsRow.map((i) => {
-      if (typeof i !== 'bigint') throw new Error(`Poseidon MDS matrix value=${i}`);
-      return Fp.create(i);
-    });
-  });
-  return Object.freeze({ ...opts, rounds, sboxFn, roundConstants, mds });
+
+  if (!sboxPower || ![3, 5, 7].includes(sboxPower))
+    throw new Error(`Poseidon wrong sboxPower=${sboxPower}`);
+  const _sboxPower = BigInt(sboxPower);
+  let sboxFn = (n: bigint) => FpPow(Fp, n, _sboxPower);
+  // Unwrapped sbox power for common cases (195->142μs)
+  if (sboxPower === 3) sboxFn = (n: bigint) => Fp.mul(Fp.sqrN(n), n);
+  else if (sboxPower === 5) sboxFn = (n: bigint) => Fp.mul(Fp.sqrN(Fp.sqrN(n)), n);
+
+  return Object.freeze({ ...opts, rounds, sboxFn, roundConstants, mds: _mds });
 }
 
 export function splitConstants(rc: bigint[], t: number) {
diff --git a/test/poseidon.test.js b/test/poseidon.test.js
index 98bcc52..38a2249 100644
--- a/test/poseidon.test.js
+++ b/test/poseidon.test.js
@@ -163,6 +163,7 @@ should('poseidonperm_x5_255_3', () => {
     t,
     roundsFull: 8,
     roundsPartial: 57,
+    sboxPower: 5,
     mds,
     roundConstants,
   });
@@ -229,6 +230,7 @@ should('poseidonperm_x5_255_5', () => {
     t,
     roundsFull: 8,
     roundsPartial: 60,
+    sboxPower: 5,
     mds,
     roundConstants,
   });
@@ -280,6 +282,7 @@ should('poseidonperm_x5_254_3', () => {
     t,
     roundsFull: 8,
     roundsPartial: 57,
+    sboxPower: 5,
     mds,
     roundConstants,
   });
@@ -347,6 +350,7 @@ should('poseidonperm_x5_254_5', () => {
     t,
     roundsFull: 8,
     roundsPartial: 60,
+    sboxPower: 5,
     mds,
     roundConstants,
   });