1package poseidon
 2
 3import (
 4	"errors"
 5	"fmt"
 6	"math/big"
 7
 8	"github.com/iden3/go-iden3-crypto/ff"
 9	"github.com/iden3/go-iden3-crypto/utils"
10)
11
12const NROUNDSF = 8 //nolint:golint
13
14var NROUNDSP = []int{56, 57, 56, 60, 60, 63, 64, 63} //nolint:golint
15
16func zero() *ff.Element {
17	return ff.NewElement()
18}
19
20// ark computes Add-Round Key, from the paper https://eprint.iacr.org/2019/458.pdf
21func ark(state []*ff.Element, c []*ff.Element, it int) {
22	for i := 0; i < len(state); i++ {
23		state[i].Add(state[i], c[it+i])
24	}
25}
26
27// exp5 performs x^5 mod p
28// https://eprint.iacr.org/2019/458.pdf page 8
29func exp5(a *ff.Element) {
30	a.Exp(*a, 5)
31}
32
33// sbox https://eprint.iacr.org/2019/458.pdf page 6
34func sbox(nRoundsF, nRoundsP int, state []*ff.Element, i int) {
35	if (i < nRoundsF/2) || (i >= nRoundsF/2+nRoundsP) {
36		for j := 0; j < len(state); j++ {
37			exp5(state[j])
38		}
39	} else {
40		exp5(state[0])
41	}
42}
43
44// mix returns [[matrix]] * [vector]
45func mix(state []*ff.Element, newState []*ff.Element, m [][]*ff.Element) {
46	mul := zero()
47	for i := 0; i < len(state); i++ {
48		newState[i].SetUint64(0)
49		for j := 0; j < len(state); j++ {
50			mul.Mul(m[j][i], state[j])
51			newState[i].Add(newState[i], mul)
52		}
53	}
54}
55
56// Hash computes the Poseidon hash for the given inputs
57func Hash(inpBI []*big.Int) (*big.Int, error) {
58	t := len(inpBI) + 1
59	if len(inpBI) == 0 || len(inpBI) >= len(NROUNDSP)-1 {
60		return nil, fmt.Errorf("invalid inputs length %d, max %d", len(inpBI), len(NROUNDSP)-1)
61	}
62	if !utils.CheckBigIntArrayInField(inpBI[:]) {
63		return nil, errors.New("inputs values not inside Finite Field")
64	}
65	inp := utils.BigIntArrayToElementArray(inpBI[:])
66	state := make([]*ff.Element, t)
67	copy(state[:], inp[:])
68	state[len(state)-1] = zero()
69
70	nRoundsF := NROUNDSF
71	nRoundsP := NROUNDSP[t-2]
72
73	newState := make([]*ff.Element, t)
74	for i := 0; i < t; i++ {
75		newState[i] = zero()
76	}
77
78	// ARK --> SBox --> M, https://eprint.iacr.org/2019/458.pdf pag.5
79	for i := 0; i < nRoundsF+nRoundsP; i++ {
80		ark(state, c.c[t-2], i*t)
81		sbox(nRoundsF, nRoundsP, state, i)
82		if i < nRoundsF+nRoundsP-1 {
83			mix(state, newState, c.m[t-2])
84			state, newState = newState, state
85		}
86	}
87	rE := state[0]
88	r := big.NewInt(0)
89	rE.ToBigIntRegular(r)
90	return r, nil
91}