Merge pull request #1 from vmx/neptune-like-hashing

fix: make it produce the same results as Neptune

param.go

@@ -4,6 +4,7 @@ import (
 	"github.com/pkg/errors"
 	ff "github.com/triplewz/poseidon/bls12_381"
 	"math"
+	"math/big"
 )
 
 // security level (in bits)
@@ -75,44 +76,53 @@ func isRoundNumberSecure(t, rf, rp int) bool {
 // appendBits converts a number to the bit slice.
 // For simplicity, we use uint8 1 or 0 to represent a bit.
 func appendBits(bits []byte, n, size int) []byte {
-	for i := 0; i < size; i++ {
-		b := (n & 1)
+	for i := size - 1; i >= 0; i-- {
+		bitmask := 1 << i
+		b := (n & bitmask) >> i
 		bits = append(bits, byte(b))
-		n >>= 1
 	}
 
 	return bits
 }
 
-// genNewBits generates new 80-bits slice.
-func genNewBits(bits []byte) []byte {
+// genNewBits generates new 80-bits slice and returns the newly generated bit.
+func genNewBits(bits []byte) byte {
 	newBit := byte(bits[0] ^ bits[13] ^ bits[23] ^ bits[38] ^ bits[51] ^ bits[62])
 	newBits := append(bits, newBit)
 	copy(bits, newBits[1:])
-	return bits
+	return newBit
 }
 
 // nextByte converts bits to byte.
-func nextByte(bits []byte) byte {
+func nextByte(bits []byte, bitCount int) byte {
 	var b byte
-	for i := 0; i < 8; i++ {
+	for i := 0; i < bitCount; i++ {
+		newBit := genNewBits(bits)
+		for newBit == 0 {
+			genNewBits(bits)
+			newBit = genNewBits(bits)
+		}
+		newBit = genNewBits(bits)
+
 		b <<= 1
-		if bits[i] == 1 {
+		if newBit == 1 {
 			b += 1
 		}
 	}
 	return b
 }
 
 // getBytes generates a random byte slice.
-func getBytes(bits []byte) []byte {
+func getBytes(bits []byte, fieldsize int) []byte {
+	// Only prime fields are supported, and they always have reminder bits.
+	remainderBits := fieldsize % 8
+
 	buf := make([]byte, ff.Bytes)
-	for i := 0; i < ff.Bytes; i++ {
-		buf[i] = nextByte(bits)
-		//regen 8 bits.
-		for i := 0; i < 8; i++ {
-			genNewBits(bits)
-		}
+	buf[0] = nextByte(bits, remainderBits)
+
+	// The first byte is already set.
+	for i := 1; i < ff.Bytes; i++ {
+		buf[i] = nextByte(bits, 8)
 	}
 
 	return buf
@@ -151,13 +161,20 @@ func genRoundConstants(field, sbox int, fieldsize, t, rf, rp int) []*ff.Element
 	bits = appendBits(bits, (1<<30)-1, 30)
 
 	for i := 0; i < 160; i++ {
-		bits = genNewBits(bits)
+		genNewBits(bits)
 	}
 
 	roundConsts := make([]*ff.Element, numCons)
 	for i := 0; i < numCons; i++ {
-		buf := getBytes(bits)
-		roundConsts[i] = new(ff.Element).SetBytes(buf)
+		for {
+			buf := getBytes(bits, fieldsize)
+			bufBigint := new(big.Int).SetBytes(buf)
+			// Skip all buffers that would result in invalid field elements.
+			if ff.IsValid(bufBigint) {
+				roundConsts[i] = new(ff.Element).SetBytes(buf)
+				break
+			}
+		}
 	}
 
 	return roundConsts

poseidon.go

@@ -39,9 +39,9 @@ var PoseidonExp = new(big.Int).SetUint64(5)
 func Hash(input []*big.Int, pdsContants *PoseidonConst, hash HashMode) (*big.Int, error) {
 	state := bigToElement(input)
 
-	//state[0] = 0,state[1:width] = input
-	zero := new(ff.Element).SetZero()
-	state = append([]*ff.Element{zero}, state...)
+	// Neptune (a Rust implementation of Poseidon) is using domain tag 0x3 by default.
+	domain_tag := new(ff.Element).SetString("3")
+	state = append([]*ff.Element{domain_tag}, state...)
 
 	//pdsContants, err := genPoseidonConstants(t)
 	//if err != nil {

poseidon_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"github.com/stretchr/testify/assert"
 	ff "github.com/triplewz/poseidon/bls12_381"
+	"math/big"
 	"os"
 	"testing"
 )
@@ -115,6 +116,14 @@ func TestPoseidonHash(t *testing.T) {
 	}
 }
 
+func TestPoseidonHashFixed(t *testing.T) {
+	cons, _ := GenPoseidonConstants(3)
+	input := []*big.Int{big.NewInt(0), big.NewInt(0)}
+	hash, _ := Hash(input, cons, OptimizedStatic)
+	expected, _ := new(big.Int).SetString("48fe0b1331196f6cdb33a7c6e5af61b76fd388e1ef1d3d418be5147f0e4613d4", 16)
+	assert.Equal(t, hash, expected)
+}
+
 func benchmarkStatic(b *testing.B, str []string) {
 	cons, _ := GenPoseidonConstants(len(str) + 1)
 	input := hexToBig(str)