Witness overrides (#10)

* add witnessOverrides support and passConstraints check

* Update versions & add unit test

---------

Co-authored-by: Oleh Komendant <oleg02kom@gmail.com>

Author: mllwchrry

README.md

@@ -51,6 +51,11 @@ await expect(matrix).with.witnessInputs({ a, b, c }).to.have.strict.witnessOutpu
 await expect(expect(matrix).with.witnessInputs({ a, b, c }).to.have.witnessOutputs({ e })).to.be.rejectedWith(
   `Expected output "e" to be "[[2,5,0],[17,26,0],[0,0,0]]", but got "[[1,4,0],[16,25,0],[0,0,0]]"`,
 );
+
+// witness overrides intentionally break constraints to test failure case
+await expect(expect(matrix).with.witnessInputs({ a, b, c }, { "main.a": 10n }).to.passConstraints()).to.be.rejectedWith(
+  "Expected witness to pass constraints, but it doesn't",
+);
 ```
 
 ### Proof testing
@@ -63,10 +68,11 @@ await expect(matrix).to.generateProof({ a, b, c });
 await expect(matrix).to.not.generateProof({ b, c, d });
 
 const proof = await matrix.generateProof({ a, b, c });
+const invalidProof = await matrix.generateProof({ a, b, c }, { "main.a": 10n });
 
 // proof verification assertion
 await expect(matrix).to.verifyProof(proof);
-await expect(matrix).to.not.verifyProof(otherProof);
+await expect(matrix).to.not.verifyProof(invalidProof);
 
 // use generated solidity verifier to verify the proof
 await expect(matrix).to.useSolidityVerifier(matrixVerifier).and.verifyProof(proof);

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "@solarity/chai-zkit",
-  "version": "0.3.2",
+  "version": "0.3.3",
   "license": "MIT",
   "author": "Distributed Lab",
   "readme": "README.md",
@@ -44,14 +44,14 @@
     "snarkjs": "0.7.5"
   },
   "peerDependencies": {
-    "@solarity/hardhat-zkit": ">=0.5.15 <0.6.0",
-    "@solarity/zkit": ">=0.3.6 <0.4.0",
+    "@solarity/hardhat-zkit": ">=0.5.16 <0.6.0",
+    "@solarity/zkit": ">=0.3.7 <0.4.0",
     "chai": "^4.0.0"
   },
   "devDependencies": {
     "@nomicfoundation/hardhat-ethers": "3.0.5",
-    "@solarity/hardhat-zkit": "^0.5.15",
-    "@solarity/zkit": "^0.3.6",
+    "@solarity/hardhat-zkit": "^0.5.16",
+    "@solarity/zkit": "^0.3.7",
     "@types/chai": "^4.3.16",
     "@types/chai-as-promised": "^7.1.8",
     "@types/mocha": "^10.0.6",

package.json

@@ -5,5 +5,6 @@ export const VERIFY_PROOF_METHOD: string = "verifyProof";
 export const STRICT_PROPERTY: string = "strict";
 export const WITNESS_INPUTS_METHOD: string = "witnessInputs";
 export const WITNESS_OUTPUTS_METHOD: string = "witnessOutputs";
+export const WITNESS_PASS_CONSTRAINTS_METHOD: string = "passConstraints";
 
 export const BN128_CURVE_NAME: string = "bn128";

src/constants.ts

@@ -4,34 +4,37 @@ import { GENERATE_PROOF_METHOD, USE_SOLIDITY_VERIFIER_METHOD, VERIFY_PROOF_METHO
 import { checkCircuitZKit } from "../utils";
 
 export function proof(chai: Chai.ChaiStatic, utils: Chai.ChaiUtils): void {
-  chai.Assertion.addMethod(GENERATE_PROOF_METHOD, function (this: any, inputs: Signals) {
-    const obj = utils.flag(this, "object");
+  chai.Assertion.addMethod(
+    GENERATE_PROOF_METHOD,
+    function (this: any, inputs: Signals, witnessOverrides?: Record<string, bigint>) {
+      const obj = utils.flag(this, "object");
 
-    checkCircuitZKit(obj, GENERATE_PROOF_METHOD);
+      checkCircuitZKit(obj, GENERATE_PROOF_METHOD);
 
-    const promise = (this.then === undefined ? Promise.resolve() : this).then(async () => {
-      let isGenerated = true;
+      const promise = (this.then === undefined ? Promise.resolve() : this).then(async () => {
+        let isGenerated = true;
 
-      try {
-        const proof = await obj.generateProof(inputs);
+        try {
+          const proof = await obj.generateProof(inputs, witnessOverrides);
 
-        utils.flag(this, "generatedProof", proof);
-      } catch (e) {
-        isGenerated = false;
-      }
+          utils.flag(this, "generatedProof", proof);
+        } catch (e) {
+          isGenerated = false;
+        }
 
-      this.assert(
-        isGenerated,
-        "Expected proof generation to be successful, but it isn't",
-        "Expected proof generation NOT to be successful, but it is",
-      );
-    });
+        this.assert(
+          isGenerated,
+          "Expected proof generation to be successful, but it isn't",
+          "Expected proof generation NOT to be successful, but it is",
+        );
+      });
 
-    this.then = promise.then.bind(promise);
-    this.catch = promise.catch.bind(promise);
+      this.then = promise.then.bind(promise);
+      this.catch = promise.catch.bind(promise);
 
-    return this;
-  });
+      return this;
+    },
+  );
 
   chai.Assertion.addMethod(USE_SOLIDITY_VERIFIER_METHOD, function (this: any, verifierContract: any) {
     const obj = utils.flag(this, "object");

src/core/proof.ts

@@ -1,6 +1,13 @@
-import { CircuitZKit, NumberLike, ProvingSystemType, Signals } from "@solarity/zkit";
+import * as snarkjs from "snarkjs";
 
-import { STRICT_PROPERTY, WITNESS_INPUTS_METHOD, WITNESS_OUTPUTS_METHOD } from "../constants";
+import { CircuitZKit, NumberLike, ProvingSystemType, Signals, writeWitnessFile } from "@solarity/zkit";
+
+import {
+  STRICT_PROPERTY,
+  WITNESS_INPUTS_METHOD,
+  WITNESS_OUTPUTS_METHOD,
+  WITNESS_PASS_CONSTRAINTS_METHOD,
+} from "../constants";
 import { checkCircuitZKit, loadOutputs, outputSignalsCompare } from "../utils";
 
 export function witness(chai: Chai.ChaiStatic, utils: Chai.ChaiUtils): void {
@@ -10,23 +17,26 @@ export function witness(chai: Chai.ChaiStatic, utils: Chai.ChaiUtils): void {
     return this;
   });
 
-  chai.Assertion.addMethod(WITNESS_INPUTS_METHOD, function (this: any, inputs: Signals) {
-    const obj = utils.flag(this, "object");
+  chai.Assertion.addMethod(
+    WITNESS_INPUTS_METHOD,
+    function (this: any, inputs: Signals, witnessOverrides?: Record<string, bigint>) {
+      const obj = utils.flag(this, "object");
 
-    checkCircuitZKit(obj, WITNESS_INPUTS_METHOD);
+      checkCircuitZKit(obj, WITNESS_INPUTS_METHOD);
 
-    const promise = (this.then === undefined ? Promise.resolve() : this).then(async () => {
-      const witness = await obj.calculateWitness(inputs);
+      const promise = (this.then === undefined ? Promise.resolve() : this).then(async () => {
+        const witness = await obj.calculateWitness(inputs, witnessOverrides);
 
-      utils.flag(this, "inputs", inputs);
-      utils.flag(this, "witness", witness);
-    });
+        utils.flag(this, "inputs", inputs);
+        utils.flag(this, "witness", witness);
+      });
 
-    this.then = promise.then.bind(promise);
-    this.catch = promise.catch.bind(promise);
+      this.then = promise.then.bind(promise);
+      this.catch = promise.catch.bind(promise);
 
-    return this;
-  });
+      return this;
+    },
+  );
 
   chai.Assertion.addMethod(WITNESS_OUTPUTS_METHOD, function (this: any, outputs: Signals | NumberLike[]) {
     const obj = utils.flag(this, "object");
@@ -56,4 +66,39 @@ export function witness(chai: Chai.ChaiStatic, utils: Chai.ChaiUtils): void {
 
     return this;
   });
+
+  chai.Assertion.addMethod(WITNESS_PASS_CONSTRAINTS_METHOD, function (this: any) {
+    const obj = utils.flag(this, "object");
+
+    checkCircuitZKit(obj, WITNESS_PASS_CONSTRAINTS_METHOD);
+
+    const promise = (this.then === undefined ? Promise.resolve() : this).then(async () => {
+      const witness = utils.flag(this, "witness");
+
+      if (!witness) {
+        throw new Error("`passConstraints` is expected to be called after `witnessInputs`");
+      }
+
+      const r1csFile = obj.mustGetArtifactsFilePath("r1cs");
+      const witnessFile = obj.getTemporaryWitnessPath();
+
+      await writeWitnessFile(witnessFile, witness);
+
+      const constraintsPassed = await snarkjs.wtns.check(r1csFile, witnessFile, {
+        info: () => {},
+        warn: () => {},
+      });
+
+      this.assert(
+        constraintsPassed,
+        "Expected witness to pass constraints, but it doesn't",
+        "Expected witness NOT to pass constraints, but it does",
+      );
+    });
+
+    this.then = promise.then.bind(promise);
+    this.catch = promise.catch.bind(promise);
+
+    return this;
+  });
 }