From 48b66e17f990ee01c08efd788c9d6301694ab439 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Wed, 24 Nov 2021 19:49:59 +0100
Subject: [PATCH] Squeeze circuit implemented

---
 circuits/keccak256.circom             | 15 +++++++++
 go-keccak256-bits-impl/keccak.go      | 26 +++++++++++----
 go-keccak256-bits-impl/keccak_test.go | 47 +++++++++++++++++++++++++++
 test/circuits/squeeze_test.circom     |  5 +++
 test/keccak256.js                     | 42 ++++++++++++++++++++++++
 5 files changed, 129 insertions(+), 6 deletions(-)
 create mode 100644 test/circuits/squeeze_test.circom

diff --git a/circuits/keccak256.circom b/circuits/keccak256.circom
index be5716b..d19a1f7 100644
--- a/circuits/keccak256.circom
+++ b/circuits/keccak256.circom
@@ -121,6 +121,21 @@ template Final(nBits) {
     }
 }
 
+template Squeeze(nBits) {
+    signal input s[25*64];
+    signal output out[nBits];
+    var i;
+    var j;
+
+    for (i=0; i<25; i++) {
+        for (j=0; j<64; j++) {
+            if (i*64+j<nBits) {
+                out[i*64+j] <== s[i*64+j];
+            }
+        }
+    }
+}
+
 template Keccakf() {
     signal input in[25*64];
     signal output out[25*64];
diff --git a/go-keccak256-bits-impl/keccak.go b/go-keccak256-bits-impl/keccak.go
index f54b6c0..098822b 100644
--- a/go-keccak256-bits-impl/keccak.go
+++ b/go-keccak256-bits-impl/keccak.go
@@ -60,12 +60,26 @@ func absorb(s [25 * 64]bool, block []bool) [25 * 64]bool {
 }
 
 func squeeze(s [25 * 64]bool) []bool {
-	b := make([]bool, 8*8*len(s))
-	n := size * 8
-	for i := 0; i < len(s)/64; i++ {
-		copy(b[i*8*8:i*8*8+64], s[i*64:i*64+64])
-	}
-	return b[:n]
+	// option1
+	// b := make([]bool, 8*8*len(s))
+	// for i := 0; i < 25; i++ {
+	//         copy(b[i*64:i*64+64], s[i*64:i*64+64])
+	// }
+	// return b[:size*8]
+
+	// option2
+	// out := make([]bool, size*8)
+	// for i := 0; i < 25; i++ {
+	//         for j := 0; j < 64; j++ {
+	//                 if i*64+j < size*8 {
+	//                         out[i*64+j] = s[i*64+j]
+	//                 }
+	//         }
+	// }
+	// return out
+
+	// option3
+	return s[:size*8]
 }
 
 func keccakfRound(s [25 * 64]bool, r int) [25 * 64]bool {
diff --git a/go-keccak256-bits-impl/keccak_test.go b/go-keccak256-bits-impl/keccak_test.go
index 9b5ee32..304dc3c 100644
--- a/go-keccak256-bits-impl/keccak_test.go
+++ b/go-keccak256-bits-impl/keccak_test.go
@@ -210,3 +210,50 @@ func TestFinal(t *testing.T) {
 			1486295134719870048, 9161068934282437999, 8245604251371175619,
 			8421994351908003183})
 }
+
+func TestSqueeze(t *testing.T) {
+	in := []uint64{16852464862333879129, 9588646233186836430, 693207875935078627,
+		6545910230963382296, 3599194178366828471, 13130606490077331384,
+		10374798023615518933, 7285576075118720444, 4097382401500492461,
+		3968685317688314807, 3350659309646210303, 640023485234837464,
+		2550030127986774041, 8948768022010378840, 10678227883444996205,
+		1395278318096830339, 2744077813166753978, 13362598477502046010,
+		14601579319881128511, 4070707967569603186, 16833768365875755098,
+		1486295134719870048, 9161068934282437999, 8245604251371175619,
+		8421994351908003183}
+	inBits := u64ArrayToBits(in)
+	var inBits1600 [25 * 64]bool
+	copy(inBits1600[:], inBits[:])
+
+	outBits := squeeze(inBits1600)
+
+	// printU64Array("in", in)
+	// printBytes("out", bitsToBytes(outBits[:]))
+
+	qt.Assert(t, bitsToBytes(outBits[:]), qt.DeepEquals,
+		[]byte{89, 195, 41, 13, 129, 251, 223, 233, 206, 31, 253, 61,
+			242, 182, 17, 133, 227, 8, 157, 240, 227, 196, 158, 9, 24, 232,
+			42, 96, 172, 190, 215, 90})
+
+	// 2nd test
+	in = []uint64{16953415415620100490, 7495738965189503699,
+		12723370805759944158, 3295955328722933810, 12121371508560456016,
+		174876831679863147, 15944933357501475584, 7502339663607726274,
+		12048918224562833898, 16715284461100269102, 15582559130083209842,
+		1743886467337678829, 2424196198791253761, 1116417308245482383,
+		10367365997906434042, 1849801549382613906, 13294939539683415102,
+		4478091053375708790, 2969967870313332958, 14618962068930014237,
+		2721742233407503451, 12003265593030191290, 8109318293656735684,
+		6346795302983965746, 12210038122000333046}
+	inBits = u64ArrayToBits(in)
+	copy(inBits1600[:], inBits[:])
+
+	outBits = squeeze(inBits1600)
+
+	// printU64Array("in", in)
+	// printBytes("out", bitsToBytes(outBits[:]))
+	qt.Assert(t, bitsToBytes(outBits[:]), qt.DeepEquals,
+		[]byte{138, 225, 170, 89, 127, 161, 70, 235, 211, 170, 44, 237,
+			223, 54, 6, 104, 222, 165, 229, 38, 86, 126, 146, 176, 50, 24,
+			22, 164, 232, 149, 189, 45})
+}
diff --git a/test/circuits/squeeze_test.circom b/test/circuits/squeeze_test.circom
new file mode 100644
index 0000000..2c0f35c
--- /dev/null
+++ b/test/circuits/squeeze_test.circom
@@ -0,0 +1,5 @@
+pragma circom 2.0.0;
+
+include "../../circuits/keccak256.circom";
+
+component main = Squeeze(32*8);
diff --git a/test/keccak256.js b/test/keccak256.js
index 4fe3938..7ce642a 100644
--- a/test/keccak256.js
+++ b/test/keccak256.js
@@ -479,3 +479,45 @@ describe("Keccak-Final test", function () {
 	assert.deepEqual(stateOutU64, expectedOut);
     });
 });
+
+describe("Keccak-Squeeze test", function () {
+    this.timeout(100000);
+
+    let cir;
+    before(async () => {
+	cir = await c_tester(path.join(__dirname, "circuits", "squeeze_test.circom"));
+	await cir.loadConstraints();
+	console.log("n_constraints", cir.constraints.length);
+    });
+
+    it ("Squeeze 1 (testvector generated from go)", async () => {
+	const input = strsToBigInts(["16852464862333879129", "9588646233186836430", "693207875935078627", "6545910230963382296", "3599194178366828471", "13130606490077331384", "10374798023615518933", "7285576075118720444", "4097382401500492461", "3968685317688314807", "3350659309646210303", "640023485234837464", "2550030127986774041", "8948768022010378840", "10678227883444996205", "1395278318096830339", "2744077813166753978", "13362598477502046010", "14601579319881128511", "4070707967569603186", "16833768365875755098", "1486295134719870048", "9161068934282437999", "8245604251371175619", "8421994351908003183"]);
+
+	const expectedOut = [89, 195, 41, 13, 129, 251, 223, 233, 206, 31, 253, 61, 242, 182, 17, 133, 227, 8, 157, 240, 227, 196, 158, 9, 24, 232, 42, 96, 172, 190, 215, 90];
+
+	const inIn = u64ArrayToBits(input);
+	const expectedOutBits = bytesToBits(expectedOut);
+
+	const witness = await cir.calculateWitness({ "s": inIn }, true);
+
+	const stateOut = witness.slice(1, 1+(32*8));
+	const stateOutBytes = bitsToBytes(stateOut);
+	// console.log(stateOutBytes, expectedOut);
+	assert.deepEqual(stateOutBytes, expectedOut);
+    });
+    it ("Squeeze 2 (testvector generated from go)", async () => {
+	const input = strsToBigInts(["16953415415620100490", "7495738965189503699", "12723370805759944158", "3295955328722933810", "12121371508560456016", "174876831679863147", "15944933357501475584", "7502339663607726274", "12048918224562833898", "16715284461100269102", "15582559130083209842", "1743886467337678829", "2424196198791253761", "1116417308245482383", "10367365997906434042", "1849801549382613906", "13294939539683415102", "4478091053375708790", "2969967870313332958", "14618962068930014237", "2721742233407503451", "12003265593030191290", "8109318293656735684", "6346795302983965746", "12210038122000333046"]);
+    
+	const expectedOut = [138, 225, 170, 89, 127, 161, 70, 235, 211, 170, 44, 237, 223, 54, 6, 104, 222, 165, 229, 38, 86, 126, 146, 176, 50, 24, 22, 164, 232, 149, 189, 45];
+    
+	const inIn = u64ArrayToBits(input);
+	const expectedOutBits = bytesToBits(expectedOut);
+    
+	const witness = await cir.calculateWitness({ "s": inIn }, true);
+    
+	const stateOut = witness.slice(1, 1+(32*8));
+	const stateOutBytes = bitsToBytes(stateOut);
+	// console.log(stateOutBytes, expectedOut);
+	assert.deepEqual(stateOutBytes, expectedOut);
+    });
+});