From 3b492f9801573eebcfe1b6c584afe8a3beecf2b4 Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Tue, 27 Oct 2020 21:23:32 +0300
Subject: [PATCH] poseidon contract now accepts both selectors: for uint256[n]
 and bytes32[n]

---
 src/poseidon_gencontract.js | 76 +++++++++++++++++++++++++------------
 test/poseidoncontract.js    |  4 +-
 2 files changed, 53 insertions(+), 27 deletions(-)

diff --git a/src/poseidon_gencontract.js b/src/poseidon_gencontract.js
index 02ae9b2..3f5c02f 100644
--- a/src/poseidon_gencontract.js
+++ b/src/poseidon_gencontract.js
@@ -103,8 +103,13 @@ function createCode(nInputs) {
     C.push(0);
     C.calldataload();
     C.div();
-    C.push(Web3Utils.keccak256(`poseidon(uint256[${nInputs}])`).slice(0, 10)); // poseidon(uint256[])
+    C.dup(0);
+    C.push(Web3Utils.keccak256(`poseidon(uint256[${nInputs}])`).slice(0, 10)); // poseidon(uint256[n])
     C.eq();
+    C.swap(1);
+    C.push(Web3Utils.keccak256(`poseidon(bytes32[${nInputs}])`).slice(0, 10)); // poseidon(bytes32[n])
+    C.eq();
+    C.or();
     C.jmpi("start");
     C.invalid();
 
@@ -159,30 +164,51 @@ function createCode(nInputs) {
 }
 
 function generateABI(nInputs) {
-  return [
-    {
-        "constant": true,
-        "inputs": [
-            {
-                "internalType": `uint256[${nInputs}]`,
-                "name": "input",
-                "type": `uint256[${nInputs}]`
-            }
-        ],
-        "name": "poseidon",
-        "outputs": [
-            {
-                "internalType": "uint256",
-                "name": "",
-                "type": "uint256"
-            }
-        ],
-        "payable": false,
-        "stateMutability": "pure",
-        "type": "function"
-    }
-  ]
-};
+    return [
+        {
+            "constant": true,
+            "inputs": [
+                {
+                    "internalType": `bytes32[${nInputs}]`,
+                    "name": "input",
+                    "type": `bytes32[${nInputs}]`
+                }
+            ],
+            "name": "poseidon",
+            "outputs": [
+                {
+                    "internalType": "bytes32",
+                    "name": "",
+                    "type": "bytes32"
+                }
+            ],
+            "payable": false,
+            "stateMutability": "pure",
+            "type": "function"
+        },
+        {
+            "constant": true,
+            "inputs": [
+                {
+                    "internalType": `uint256[${nInputs}]`,
+                    "name": "input",
+                    "type": `uint256[${nInputs}]`
+                }
+            ],
+            "name": "poseidon",
+            "outputs": [
+                {
+                    "internalType": "uint256",
+                    "name": "",
+                    "type": "uint256"
+                }
+            ],
+            "payable": false,
+            "stateMutability": "pure",
+            "type": "function"
+        }
+    ];
+}
 
 module.exports.generateABI = generateABI;
 module.exports.createCode = createCode;
diff --git a/test/poseidoncontract.js b/test/poseidoncontract.js
index 20605f0..c54d265 100644
--- a/test/poseidoncontract.js
+++ b/test/poseidoncontract.js
@@ -14,8 +14,8 @@ describe("Poseidon Smart contract test", function () {
     let poseidon4;
     let accounts;
     this.timeout(100000);
-    let C2
-    let C4
+    let C2;
+    let C4;
 
     before(async () => {
         web3 = new Web3(ganache.provider(), null, { transactionConfirmationBlocks: 1 });