Foreign fields 5: Optimized multiplication of sums

CHANGELOG.md

@@ -19,6 +19,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/1ad7333e9e...HEAD)
 
+# Added
+
+- `Gadgets.ForeignField.assertMul()` for efficiently constraining products of sums in non-native arithmetic https://github.com/o1-labs/o1js/pull/1262
+- `Unconstrained` for safely maintaining unconstrained values in provable code https://github.com/o1-labs/o1js/pull/1262
+
 ### Changed
 
 - Change precondition APIs to use "require" instead of "assert" as the verb, to distinguish them from provable assertions.

src/index.ts

@@ -8,6 +8,7 @@ export type {
   FlexibleProvable,
   FlexibleProvablePure,
   InferProvable,
+  Unconstrained,
 } from './lib/circuit_value.js';
 export {
   CircuitValue,

src/lib/circuit_value.ts

@@ -1,5 +1,5 @@
 import 'reflect-metadata';
-import { ProvablePure } from '../snarky.js';
+import { ProvablePure, Snarky } from '../snarky.js';
 import { Field, Bool, Scalar, Group } from './core.js';
 import {
   provable,
@@ -15,6 +15,8 @@ import type {
   IsPure,
 } from '../bindings/lib/provable-snarky.js';
 import { Provable } from './provable.js';
+import { assert } from './errors.js';
+import { inCheckedComputation } from './provable-context.js';
 
 // external API
 export {
@@ -43,6 +45,7 @@ export {
   HashInput,
   InferJson,
   InferredProvable,
+  Unconstrained,
 };
 
 type ProvableExtension<T, TJson = any> = {
@@ -454,6 +457,93 @@ function Struct<
   return Struct_ as any;
 }
 
+/**
+ * Container which holds an unconstrained value. This can be used to pass values
+ * between the out-of-circuit blocks in provable code.
+ *
+ * Invariants:
+ * - An `Unconstrained`'s value can only be accessed in auxiliary contexts.
+ * - An `Unconstrained` can be empty when compiling, but never empty when running as the prover.
+ *   (there is no way to create an empty `Unconstrained` in the prover)
+ *
+ * @example
+ * ```ts
+ * let x = Unconstrained.from(0n);
+ *
+ * class MyContract extends SmartContract {
+ *   `@method` myMethod(x: Unconstrained<bigint>) {
+ *
+ *     Provable.witness(Field, () => {
+ *       // we can access and modify `x` here
+ *       let newValue = x.get() + otherField.toBigInt();
+ *       x.set(newValue);
+ *
+ *       // ...
+ *     });
+ *
+ *     // throws an error!
+ *     x.get();
+ *   }
+ * ```
+ */
+class Unconstrained<T> {
+  private option:
+    | { isSome: true; value: T }
+    | { isSome: false; value: undefined };
+
+  private constructor(isSome: boolean, value?: T) {
+    this.option = { isSome, value: value as any };
+  }
+
+  /**
+   * Read an unconstrained value.
+   *
+   * Note: Can only be called outside provable code.
+   */
+  get(): T {
+    if (inCheckedComputation() && !Snarky.run.inProverBlock())
+      throw Error(`You cannot use Unconstrained.get() in provable code.
+
+The only place where you can read unconstrained values is in Provable.witness()
+and Provable.asProver() blocks, which execute outside the proof.
+`);
+    assert(this.option.isSome, 'Empty `Unconstrained`'); // never triggered
+    return this.option.value;
+  }
+
+  /**
+   * Modify the unconstrained value.
+   */
+  set(value: T) {
+    this.option = { isSome: true, value };
+  }
+
+  /**
+   * Create an `Unconstrained` with the given `value`.
+   */
+  static from<T>(value: T) {
+    return new Unconstrained(true, value);
+  }
+
+  /**
+   * Create an `Unconstrained` from a witness computation.
+   */
+  static witness<T>(compute: () => T) {
+    return Provable.witness(
+      Unconstrained.provable,
+      () => new Unconstrained(true, compute())
+    );
+  }
+
+  static provable: Provable<Unconstrained<any>> = {
+    sizeInFields: () => 0,
+    toFields: () => [],
+    toAuxiliary: (t?: any) => [t ?? new Unconstrained(false)],
+    fromFields: (_, [t]) => t,
+    check: () => {},
+  };
+}
+
 let primitives = new Set([Field, Bool, Scalar, Group]);
 function isPrimitive(obj: any) {
   for (let P of primitives) {

src/lib/circuit_value.unit-test.ts

@@ -1,4 +1,4 @@
-import { provable, Struct } from './circuit_value.js';
+import { provable, Struct, Unconstrained } from './circuit_value.js';
 import { UInt32 } from './int.js';
 import { PrivateKey, PublicKey } from './signature.js';
 import { expect } from 'expect';
@@ -96,6 +96,7 @@ class MyStructPure extends Struct({
 class MyTuple extends Struct([PublicKey, String]) {}
 
 let targetString = 'some particular string';
+let targetBigint = 99n;
 let gotTargetString = false;
 
 // create a smart contract and pass auxiliary data to a method
@@ -106,11 +107,22 @@ class MyContract extends SmartContract {
   // this works because MyStructPure only contains field elements
   @state(MyStructPure) x = State<MyStructPure>();
 
-  @method myMethod(value: MyStruct, tuple: MyTuple, update: AccountUpdate) {
+  @method myMethod(
+    value: MyStruct,
+    tuple: MyTuple,
+    update: AccountUpdate,
+    unconstrained: Unconstrained<bigint>
+  ) {
     // check if we can pass in string values
     if (value.other === targetString) gotTargetString = true;
     value.uint[0].assertEquals(UInt32.zero);
 
+    // cannot access unconstrained values in provable code
+    if (Provable.inCheckedComputation())
+      expect(() => unconstrained.get()).toThrow(
+        'You cannot use Unconstrained.get() in provable code.'
+      );
+
     Provable.asProver(() => {
       let err = 'wrong value in prover';
       if (tuple[1] !== targetString) throw Error(err);
@@ -119,6 +131,9 @@ class MyContract extends SmartContract {
       if (update.lazyAuthorization?.kind !== 'lazy-signature') throw Error(err);
       if (update.lazyAuthorization.privateKey?.toBase58() !== key.toBase58())
         throw Error(err);
+
+      // check if we can pass in unconstrained values
+      if (unconstrained.get() !== targetBigint) throw Error(err);
     });
   }
 }
@@ -141,7 +156,8 @@ let tx = await transaction(() => {
       uint: [UInt32.from(0), UInt32.from(10)],
     },
     [address, targetString],
-    accountUpdate
+    accountUpdate,
+    Unconstrained.from(targetBigint)
   );
 });
 

src/lib/field.ts

@@ -11,10 +11,12 @@ export { Field };
 
 // internal API
 export {
-  ConstantField,
   FieldType,
   FieldVar,
   FieldConst,
+  ConstantField,
+  VarField,
+  VarFieldVar,
   isField,
   withMessage,
   readVarMessage,
@@ -69,6 +71,7 @@ type FieldVar =
   | [FieldType.Scale, FieldConst, FieldVar];
 
 type ConstantFieldVar = [FieldType.Constant, FieldConst];
+type VarFieldVar = [FieldType.Var, number];
 
 const FieldVar = {
   constant(x: bigint | FieldConst): ConstantFieldVar {
@@ -78,6 +81,9 @@ const FieldVar = {
   isConstant(x: FieldVar): x is ConstantFieldVar {
     return x[0] === FieldType.Constant;
   },
+  isVar(x: FieldVar): x is VarFieldVar {
+    return x[0] === FieldType.Var;
+  },
   add(x: FieldVar, y: FieldVar): FieldVar {
     if (FieldVar.isConstant(x) && x[1][1] === 0n) return y;
     if (FieldVar.isConstant(y) && y[1][1] === 0n) return x;
@@ -101,6 +107,7 @@ const FieldVar = {
 };
 
 type ConstantField = Field & { value: ConstantFieldVar };
+type VarField = Field & { value: VarFieldVar };
 
 /**
  * A {@link Field} is an element of a prime order [finite field](https://en.wikipedia.org/wiki/Finite_field).
@@ -1039,7 +1046,7 @@ class Field {
   seal() {
     if (this.isConstant()) return this;
     let x = Snarky.field.seal(this.value);
-    return new Field(x);
+    return VarField(x);
   }
 
   /**
@@ -1360,3 +1367,7 @@ there is \`Provable.asProver(() => { ... })\` which allows you to use ${varName}
 Warning: whatever happens inside asProver() will not be part of the zk proof.
 `;
 }
+
+function VarField(x: VarFieldVar): VarField {
+  return new Field(x) as VarField;
+}

src/lib/gadgets/basic.ts

@@ -0,0 +1,79 @@
+/**
+ * Basic gadgets that only use generic gates
+ */
+import type { Field, VarField } from '../field.js';
+import { existsOne, toVar } from './common.js';
+import { Gates } from '../gates.js';
+import { TupleN } from '../util/types.js';
+
+export { assertOneOf };
+
+// TODO: create constant versions of these and expose on Gadgets
+
+/**
+ * Assert that a value equals one of a finite list of constants:
+ * `(x - c1)*(x - c2)*...*(x - cn) === 0`
+ *
+ * TODO: what prevents us from getting the same efficiency with snarky DSL code?
+ */
+function assertOneOf(x: Field, allowed: [bigint, bigint, ...bigint[]]) {
+  let xv = toVar(x);
+  let [c1, c2, ...c] = allowed;
+  let n = c.length;
+  if (n === 0) {
+    // (x - c1)*(x - c2) === 0
+    assertBilinear(xv, xv, [1n, -(c1 + c2), 0n, c1 * c2]);
+    return;
+  }
+  // z = (x - c1)*(x - c2)
+  let z = bilinear(xv, xv, [1n, -(c1 + c2), 0n, c1 * c2]);
+
+  for (let i = 0; i < n; i++) {
+    if (i < n - 1) {
+      // z = z*(x - c)
+      z = bilinear(z, xv, [1n, -c[i], 0n, 0n]);
+    } else {
+      // z*(x - c) === 0
+      assertBilinear(z, xv, [1n, -c[i], 0n, 0n]);
+    }
+  }
+}
+
+// low-level helpers to create generic gates
+
+/**
+ * Compute bilinear function of x and y:
+ * `z = a*x*y + b*x + c*y + d`
+ */
+function bilinear(x: VarField, y: VarField, [a, b, c, d]: TupleN<bigint, 4>) {
+  let z = existsOne(() => {
+    let x0 = x.toBigInt();
+    let y0 = y.toBigInt();
+    return a * x0 * y0 + b * x0 + c * y0 + d;
+  });
+  // b*x + c*y - z + a*x*y + d === 0
+  Gates.generic(
+    { left: b, right: c, out: -1n, mul: a, const: d },
+    { left: x, right: y, out: z }
+  );
+  return z;
+}
+
+/**
+ * Assert bilinear equation on x, y and z:
+ * `a*x*y + b*x + c*y + d === z`
+ *
+ * The default for z is 0.
+ */
+function assertBilinear(
+  x: VarField,
+  y: VarField,
+  [a, b, c, d]: TupleN<bigint, 4>,
+  z?: VarField
+) {
+  // b*x + c*y - z? + a*x*y + d === 0
+  Gates.generic(
+    { left: b, right: c, out: z === undefined ? 0n : -1n, mul: a, const: d },
+    { left: x, right: y, out: z === undefined ? x : z }
+  );
+}

src/lib/gadgets/common.ts

@@ -1,5 +1,5 @@
 import { Provable } from '../provable.js';
-import { Field, FieldConst, FieldVar, FieldType } from '../field.js';
+import { Field, FieldConst, FieldVar, VarField } from '../field.js';
 import { Tuple, TupleN } from '../util/types.js';
 import { Snarky } from '../../snarky.js';
 import { MlArray } from '../ml/base.js';
@@ -12,6 +12,7 @@ export {
   existsOne,
   toVars,
   toVar,
+  isVar,
   assert,
   bitSlice,
   witnessSlice,
@@ -21,7 +22,7 @@ export {
 
 function existsOne(compute: () => bigint) {
   let varMl = Snarky.existsVar(() => FieldConst.fromBigint(compute()));
-  return new Field(varMl);
+  return VarField(varMl);
 }
 
 function exists<N extends number, C extends () => TupleN<bigint, N>>(
@@ -31,7 +32,7 @@ function exists<N extends number, C extends () => TupleN<bigint, N>>(
   let varsMl = Snarky.exists(n, () =>
     MlArray.mapTo(compute(), FieldConst.fromBigint)
   );
-  let vars = MlArray.mapFrom(varsMl, (v) => new Field(v));
+  let vars = MlArray.mapFrom(varsMl, VarField);
   return TupleN.fromArray(n, vars);
 }
 
@@ -43,20 +44,24 @@ function exists<N extends number, C extends () => TupleN<bigint, N>>(
  *
  * Same as `Field.seal()` with the difference that `seal()` leaves constants as is.
  */
-function toVar(x: Field | bigint) {
+function toVar(x: Field | bigint): VarField {
   // don't change existing vars
-  if (x instanceof Field && x.value[1] === FieldType.Var) return x;
+  if (isVar(x)) return x;
   let xVar = existsOne(() => Field.from(x).toBigInt());
   xVar.assertEquals(x);
   return xVar;
 }
 
+function isVar(x: Field | bigint): x is VarField {
+  return x instanceof Field && FieldVar.isVar(x.value);
+}
+
 /**
  * Apply {@link toVar} to each element of a tuple.
  */
 function toVars<T extends Tuple<Field | bigint>>(
   fields: T
-): { [k in keyof T]: Field } {
+): { [k in keyof T]: VarField } {
   return Tuple.map(fields, toVar);
 }