Merge pull request #1751 from o1-labs/feature/infer-nested

Accept .provable and some other type improvements

CHANGELOG.md

@@ -23,7 +23,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Changed
 
-- Reduced maximum bit length for `xor`, `not`, and `and`, operations from 254 to 240 bits to improve performance and simplify implementation. https://github.com/o1-labs/o1js/pull/1745
+- Reduced maximum bit length for `xor`, `not`, and `and`, operations from 254 to 240 bits to prevent overflow vulnerabilities. https://github.com/o1-labs/o1js/pull/1745
+- Allow using `Type` instead of `Type.provable` in APIs that expect a provable type https://github.com/o1-labs/o1js/pull/1751
+  - Example: `Provable.witness(Bytes32, () => bytes)`
+- Automatically wrap and unwrap `Unconstrained` in `fromValue` and `toValue`, so that we don't need to deal with "unconstrained" values outside provable code https://github.com/o1-labs/o1js/pull/1751
 
 ## [1.5.0](https://github.com/o1-labs/o1js/compare/ed198f305...1c736add) - 2024-07-09
 

benchmark/benchmarks/ecdsa.ts

@@ -27,9 +27,9 @@ const EcdsaBenchmarks = benchmark(
 
     tic('witness generation');
     await Provable.runAndCheck(async () => {
-      let message_ = Provable.witness(Bytes32.provable, () => message);
-      let signature_ = Provable.witness(Ecdsa.provable, () => signature);
-      let publicKey_ = Provable.witness(Secp256k1.provable, () => publicKey);
+      let message_ = Provable.witness(Bytes32, () => message);
+      let signature_ = Provable.witness(Ecdsa, () => signature);
+      let publicKey_ = Provable.witness(Secp256k1, () => publicKey);
       await keccakAndEcdsa.rawMethods.verifyEcdsa(
         message_,
         signature_,

src/examples/benchmarks/foreign-field.ts

@@ -5,14 +5,8 @@ class ForeignScalar extends createForeignField(
 ) {}
 
 function main() {
-  let s = Provable.witness(
-    ForeignScalar.Canonical.provable,
-    ForeignScalar.random
-  );
-  let t = Provable.witness(
-    ForeignScalar.Canonical.provable,
-    ForeignScalar.random
-  );
+  let s = Provable.witness(ForeignScalar.Canonical, ForeignScalar.random);
+  let t = Provable.witness(ForeignScalar.Canonical, ForeignScalar.random);
   s.mul(t);
 }
 

src/examples/benchmarks/keccak-witness.ts

@@ -4,7 +4,7 @@ let Bytes32 = Bytes(32);
 
 console.time('keccak witness');
 await Provable.runAndCheck(() => {
-  let bytes = Provable.witness(Bytes32.provable, () => Bytes32.random());
+  let bytes = Provable.witness(Bytes32, () => Bytes32.random());
   Hash.Keccak256.hash(bytes);
 });
 console.timeEnd('keccak witness');

src/examples/crypto/ecdsa/ecdsa.ts

@@ -16,12 +16,12 @@ class Bytes32 extends Bytes(32) {}
 
 const keccakAndEcdsa = ZkProgram({
   name: 'ecdsa',
-  publicInput: Bytes32.provable,
+  publicInput: Bytes32,
   publicOutput: Bool,
 
   methods: {
     verifyEcdsa: {
-      privateInputs: [Ecdsa.provable, Secp256k1.provable],
+      privateInputs: [Ecdsa, Secp256k1],
       async method(message: Bytes32, signature: Ecdsa, publicKey: Secp256k1) {
         return signature.verifyV2(message, publicKey);
       },
@@ -31,12 +31,12 @@ const keccakAndEcdsa = ZkProgram({
 
 const ecdsa = ZkProgram({
   name: 'ecdsa-only',
-  publicInput: Scalar.provable,
+  publicInput: Scalar,
   publicOutput: Bool,
 
   methods: {
     verifySignedHash: {
-      privateInputs: [Ecdsa.provable, Secp256k1.provable],
+      privateInputs: [Ecdsa, Secp256k1],
       async method(message: Scalar, signature: Ecdsa, publicKey: Secp256k1) {
         return signature.verifySignedHashV2(message, publicKey);
       },

src/examples/crypto/foreign-field.ts

@@ -88,7 +88,7 @@ assert(uCanonical instanceof SmallField.Canonical);
 class AlmostSmallField extends SmallField.AlmostReduced {}
 
 class MyContract extends SmartContract {
-  @state(AlmostSmallField.provable) x = State<AlmostSmallField>();
+  @state(AlmostSmallField) x = State<AlmostSmallField>();
 
   @method async myMethod(y: AlmostSmallField) {
     let x = y.mul(2);

src/examples/crypto/rsa/rsa.ts

@@ -1,14 +1,7 @@
 /**
  * RSA signature verification with o1js
  */
-import {
-  Field,
-  Gadgets,
-  Provable,
-  Struct,
-  Unconstrained,
-  provable,
-} from 'o1js';
+import { Field, Gadgets, Provable, Struct, Unconstrained } from 'o1js';
 
 export { Bigint2048, rsaVerify65537 };
 
@@ -21,7 +14,7 @@ const Field18 = Provable.Array(Field, 18);
 
 class Bigint2048 extends Struct({
   fields: Field18,
-  value: Unconstrained.provable as Provable<Unconstrained<bigint>>,
+  value: Unconstrained.withEmpty(0n),
 }) {
   modMul(x: Bigint2048, y: Bigint2048) {
     return multiply(x, y, this);
@@ -66,7 +59,8 @@ function multiply(
   // witness q, r so that x*y = q*p + r
   // this also adds the range checks in `check()`
   let { q, r } = Provable.witness(
-    provable({ q: Bigint2048, r: Bigint2048 }),
+    // TODO Struct() should be unnecessary
+    Struct({ q: Bigint2048, r: Bigint2048 }),
     () => {
       let xy = x.toBigint() * y.toBigint();
       let p0 = p.toBigint();

src/examples/crypto/sha256/sha256.ts

@@ -6,10 +6,10 @@ class Bytes12 extends Bytes(12) {}
 
 let SHA256Program = ZkProgram({
   name: 'sha256',
-  publicOutput: Bytes(32).provable,
+  publicOutput: Bytes(32),
   methods: {
     sha256: {
-      privateInputs: [Bytes12.provable],
+      privateInputs: [Bytes12],
       async method(xs: Bytes12) {
         return Gadgets.SHA256.hash(xs);
       },

src/lib/mina/account-update.ts

@@ -3,7 +3,11 @@ import {
   FlexibleProvable,
   StructNoJson,
 } from '../provable/types/struct.js';
-import { provable, provablePure } from '../provable/types/provable-derivers.js';
+import {
+  provable,
+  provableExtends,
+  provablePure,
+} from '../provable/types/provable-derivers.js';
 import {
   memoizationContext,
   memoizeWitness,
@@ -1347,7 +1351,7 @@ type AccountUpdateForestBase = MerkleListBase<AccountUpdateTreeBase>;
 
 const AccountUpdateTreeBase = StructNoJson({
   id: RandomId,
-  accountUpdate: HashedAccountUpdate.provable,
+  accountUpdate: HashedAccountUpdate,
   children: MerkleListBase<AccountUpdateTreeBase>(),
 });
 
@@ -1368,10 +1372,29 @@ class AccountUpdateForest extends MerkleList.create(
   AccountUpdateTreeBase,
   merkleListHash
 ) {
+  static provable = provableExtends(AccountUpdateForest, super.provable);
+
+  push(update: AccountUpdate | AccountUpdateTreeBase) {
+    return super.push(
+      update instanceof AccountUpdate ? AccountUpdateTree.from(update) : update
+    );
+  }
+  pushIf(condition: Bool, update: AccountUpdate | AccountUpdateTreeBase) {
+    return super.pushIf(
+      condition,
+      update instanceof AccountUpdate ? AccountUpdateTree.from(update) : update
+    );
+  }
+
   static fromFlatArray(updates: AccountUpdate[]): AccountUpdateForest {
     let simpleForest = accountUpdatesToCallForest(updates);
     return this.fromSimpleForest(simpleForest);
   }
+
+  toFlatArray(mutate = true, depth = 0) {
+    return AccountUpdateForest.toFlatArray(this, mutate, depth);
+  }
+
   static toFlatArray(
     forest: AccountUpdateForestBase,
     mutate = true,
@@ -1410,6 +1433,17 @@ class AccountUpdateForest extends MerkleList.create(
       });
     });
   }
+
+  // fix static methods
+  static empty() {
+    return AccountUpdateForest.provable.empty();
+  }
+  static from(array: AccountUpdateTreeBase[]) {
+    return new AccountUpdateForest(super.from(array));
+  }
+  static fromReverse(array: AccountUpdateTreeBase[]) {
+    return new AccountUpdateForest(super.fromReverse(array));
+  }
 }
 
 /**
@@ -1427,8 +1461,8 @@ class AccountUpdateForest extends MerkleList.create(
  */
 class AccountUpdateTree extends StructNoJson({
   id: RandomId,
-  accountUpdate: HashedAccountUpdate.provable,
-  children: AccountUpdateForest.provable,
+  accountUpdate: HashedAccountUpdate,
+  children: AccountUpdateForest,
 }) {
   /**
    * Create a tree of account updates which only consists of a root.
@@ -1570,9 +1604,7 @@ class UnfinishedForest {
   }
 
   witnessHash(): UnfinishedForestFinal {
-    let final = Provable.witness(AccountUpdateForest.provable, () =>
-      this.finalize()
-    );
+    let final = Provable.witness(AccountUpdateForest, () => this.finalize());
     return this.setFinal(final);
   }
 
@@ -1615,8 +1647,7 @@ class UnfinishedForest {
   }
 
   toFlatArray(mutate = true, depth = 0): AccountUpdate[] {
-    if (this.isFinal())
-      return AccountUpdateForest.toFlatArray(this.final, mutate, depth);
+    if (this.isFinal()) return this.final.toFlatArray(mutate, depth);
     assert(this.isMutable(), 'final or mutable');
     let flatUpdates: AccountUpdate[] = [];
     for (let node of this.mutable) {

src/lib/mina/actions/action-types.ts

@@ -6,6 +6,7 @@ import { Actions } from '../account-update.js';
 import { Hashed } from '../../provable/packed.js';
 import { hashWithPrefix } from '../../provable/crypto/poseidon.js';
 import { prefixes } from '../../../bindings/crypto/constants.js';
+import { ProvableType } from '../../provable/types/provable-intf.js';
 
 export { MerkleActions, MerkleActionHashes, HashedAction, FlatActions };
 export { emptyActionState, emptyActionsHash };
@@ -23,7 +24,7 @@ function MerkleActions<A extends Actionable<any>>(
   fromActionState?: Field
 ) {
   return MerkleList.create(
-    MerkleActionList(actionType).provable,
+    MerkleActionList(actionType),
     (hash, actions) =>
       hashWithPrefix(prefixes.sequenceEvents, [hash, actions.hash]),
     fromActionState ?? emptyActionState
@@ -35,7 +36,7 @@ type MerkleActionList<T> = MerkleList<Hashed<T>>;
 
 function MerkleActionList<A extends Actionable<any>>(actionType: A) {
   return MerkleList.create(
-    HashedAction(actionType).provable,
+    HashedAction(actionType),
     (hash, action) =>
       hashWithPrefix(prefixes.sequenceEvents, [hash, action.hash]),
     emptyActionsHash
@@ -45,8 +46,9 @@ function MerkleActionList<A extends Actionable<any>>(actionType: A) {
 type HashedAction<T> = Hashed<T>;
 
 function HashedAction<A extends Actionable<any>>(actionType: A) {
-  return Hashed.create(actionType as Actionable<InferProvable<A>>, (action) =>
-    hashWithPrefix(prefixes.event, actionType.toFields(action))
+  let type = ProvableType.get(actionType as Actionable<InferProvable<A>>);
+  return Hashed.create(type, (action) =>
+    hashWithPrefix(prefixes.event, type.toFields(action))
   );
 }
 
@@ -55,14 +57,15 @@ function actionFieldsToMerkleList<T>(
   fields: bigint[][][],
   fromActionState?: bigint
 ) {
-  const HashedActionT = HashedAction(actionType);
-  const MerkleActionListT = MerkleActionList(actionType);
+  let type = ProvableType.get(actionType);
+  const HashedActionT = HashedAction(type);
+  const MerkleActionListT = MerkleActionList(type);
   const MerkleActionsT = MerkleActions(
-    actionType,
+    type,
     fromActionState ? Field(fromActionState) : undefined
   );
   let actions = fields.map((event) =>
-    event.map((action) => actionType.fromFields(action.map(Field)))
+    event.map((action) => type.fromFields(action.map(Field)))
   );
   let hashes = actions.map((as) => as.map((a) => HashedActionT.hash(a)));
   return MerkleActionsT.from(hashes.map((h) => MerkleActionListT.from(h)));
@@ -92,9 +95,5 @@ function MerkleActionHashes(fromActionState?: Field) {
 type FlatActions<T> = MerkleList<Hashed<T>>;
 
 function FlatActions<A extends Actionable<any>>(actionType: A) {
-  const HashedAction = Hashed.create(
-    actionType as Actionable<InferProvable<A>>,
-    (action) => hashWithPrefix(prefixes.event, actionType.toFields(action))
-  );
-  return MerkleList.create(HashedAction.provable);
+  return MerkleList.create(HashedAction(actionType));
 }

src/lib/mina/actions/batch-reducer.ts

@@ -24,6 +24,11 @@ import {
   MerkleActions,
   emptyActionState,
 } from './action-types.js';
+import {
+  ProvableHashable,
+  ProvablePure,
+  ProvableType,
+} from '../../provable/types/provable-intf.js';
 
 // external API
 export { BatchReducer, ActionBatch };
@@ -57,7 +62,7 @@ class BatchReducer<
   Action = InferProvable<ActionType>
 > {
   batchSize: BatchSize;
-  actionType: Actionable<Action>;
+  actionType: ProvableHashable<Action> & ProvablePure<Action>;
   Batch: ReturnType<typeof ActionBatch>;
 
   program: ActionStackProgram;
@@ -133,7 +138,8 @@ class BatchReducer<
     maxActionsPerUpdate?: number;
   }) {
     this.batchSize = batchSize;
-    this.actionType = actionType as Actionable<Action>;
+    this.actionType = ProvableType.get(actionType) as ProvableHashable<Action> &
+      ProvablePure<Action>;
     this.Batch = ActionBatch(this.actionType);
 
     this.maxUpdatesFinalProof = maxUpdatesFinalProof;
@@ -385,7 +391,7 @@ class BatchReducer<
       // we make it easier to write the reducer code by making sure dummy actions have dummy values
       hashedAction = Provable.if(
         isDummy,
-        HashedActionT.provable,
+        HashedActionT,
         emptyHashedAction,
         hashedAction
       );
@@ -549,9 +555,9 @@ function ActionBatch<A extends Actionable<any>>(actionType: A) {
     processedActionState: Field,
     onchainActionState: Field,
     onchainStack: Field,
-    stack: MerkleActions(actionType).provable,
+    stack: MerkleActions(actionType),
     isRecursive: Bool,
-    witnesses: Unconstrained.provableWithEmpty<ActionWitnesses>([]),
+    witnesses: Unconstrained.withEmpty<ActionWitnesses>([]),
   });
 }
 
@@ -769,7 +775,7 @@ function actionStackProgram(maxUpdatesPerProof: number) {
         privateInputs: [
           SelfProof,
           Bool,
-          Unconstrained.provableWithEmpty<ActionWitnesses>([]),
+          Unconstrained.withEmpty<ActionWitnesses>([]),
         ],
 
         async method(

src/lib/mina/actions/batch-reducer.unit-test.ts

@@ -86,12 +86,10 @@ class UnsafeAirdrop extends SmartContract {
    *
    * Note: This two-step process is necessary so that multiple users can claim concurrently.
    */
-  @method.returns(MerkleMap.provable)
+  @method.returns(MerkleMap)
   async settleClaims(batch: Batch, proof: BatchProof) {
     // witness merkle map and require that it matches the onchain root
-    let eligibleMap = Provable.witness(MerkleMap.provable, () =>
-      eligible.clone()
-    );
+    let eligibleMap = Provable.witness(MerkleMap, () => eligible.clone());
     this.eligibleRoot.requireEquals(eligibleMap.root);
     this.eligibleLength.requireEquals(eligibleMap.length);