Fixed correctness issue in ZkProgram typing

o1-labs · Dec 20, 2024 · 4de57a7 · 4de57a7
1 parent 17adf4b
commit 4de57a7
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 45 deletions.
diff --git a/src/lib/proof-system/zkprogram.ts b/src/lib/proof-system/zkprogram.ts
@@ -184,41 +184,40 @@ let SideloadedTag = {
   },
 };
 
-function ZkProgram<
-  Config extends {
-    publicInput?: ProvableType;
-    publicOutput?: ProvableType;
-    methods: {
-      [I in string]: {
-        privateInputs: Tuple<PrivateInput>;
-        auxiliaryOutput?: ProvableType;
-      };
+type ConfigBaseType = {
+  publicInput?: ProvableType;
+  publicOutput?: ProvableType;
+  methods: {
+    [I in string]: {
+      privateInputs: Tuple<PrivateInput>;
+      auxiliaryOutput?: ProvableType;
     };
-  },
-  Methods extends {
-    [I in keyof Config['methods']]: Method<
-      InferProvableOrUndefined<Get<Config, 'publicInput'>>,
-      InferProvableOrVoid<Get<Config, 'publicOutput'>>,
-      Config['methods'][I]
-    >;
-  },
-  // derived types for convenience
-  MethodSignatures extends Config['methods'] = Config['methods'],
-  PrivateInputs extends {
-    [I in keyof Config['methods']]: Config['methods'][I]['privateInputs'];
-  } = {
-    [I in keyof Config['methods']]: Config['methods'][I]['privateInputs'];
-  },
-  AuxiliaryOutputs extends {
-    [I in keyof MethodSignatures]: Get<MethodSignatures[I], 'auxiliaryOutput'>;
-  } = {
-    [I in keyof MethodSignatures]: Get<MethodSignatures[I], 'auxiliaryOutput'>;
-  }
->(
+  };
+};
+
+type InferMethodSignatures<Config extends ConfigBaseType> = Config['methods'];
+type InferPrivateInput<Config extends ConfigBaseType> = {
+  [I in keyof Config['methods']]: Config['methods'][I]['privateInputs'];
+};
+type InferAuxilaryOutputs<Config extends ConfigBaseType> = {
+  [I in keyof InferMethodSignatures<Config>]: Get<
+    InferMethodSignatures<Config>[I],
+    'auxiliaryOutput'
+  >;
+};
+type InferMethodType<Config extends ConfigBaseType> = {
+  [I in keyof Config['methods']]: Method<
+    InferProvableOrUndefined<Get<Config, 'publicInput'>>,
+    InferProvableOrVoid<Get<Config, 'publicOutput'>>,
+    Config['methods'][I]
+  >;
+};
+
+function ZkProgram<Config extends ConfigBaseType>(
   config: Config & {
     name: string;
     methods: {
-      [I in keyof Config['methods']]: Methods[I];
+      [I in keyof Config['methods']]: InferMethodType<Config>[I];
     };
     overrideWrapDomain?: 0 | 1 | 2;
   }
@@ -248,10 +247,10 @@ function ZkProgram<
 
   publicInputType: ProvableOrUndefined<Get<Config, 'publicInput'>>;
   publicOutputType: ProvableOrVoid<Get<Config, 'publicOutput'>>;
-  privateInputTypes: PrivateInputs;
-  auxiliaryOutputTypes: AuxiliaryOutputs;
+  privateInputTypes: InferPrivateInput<Config>;
+  auxiliaryOutputTypes: InferAuxilaryOutputs<Config>;
   rawMethods: {
-    [I in keyof Config['methods']]: Methods[I]['method'];
+    [I in keyof Config['methods']]: InferMethodType<Config>[I]['method'];
   };
 
   Proof: typeof Proof<
@@ -265,10 +264,26 @@ function ZkProgram<
   [I in keyof Config['methods']]: Prover<
     InferProvableOrUndefined<Get<Config, 'publicInput'>>,
     InferProvableOrVoid<Get<Config, 'publicOutput'>>,
-    PrivateInputs[I],
-    InferProvableOrUndefined<AuxiliaryOutputs[I]>
+    InferPrivateInput<Config>[I],
+    InferProvableOrUndefined<InferAuxilaryOutputs<Config>[I]>
   >;
 } {
+  type Methods = {
+    [I in keyof Config['methods']]: Method<
+      InferProvableOrUndefined<Get<Config, 'publicInput'>>,
+      InferProvableOrVoid<Get<Config, 'publicOutput'>>,
+      Config['methods'][I]
+    >;
+  };
+  // derived types for convenience
+  type MethodSignatures = Config['methods'];
+  type PrivateInputs = {
+    [I in keyof Config['methods']]: Config['methods'][I]['privateInputs'];
+  };
+  type AuxiliaryOutputs = {
+    [I in keyof MethodSignatures]: Get<MethodSignatures[I], 'auxiliaryOutput'>;
+  };
+
   let doProving = true;
 
   let methods = config.methods;
@@ -582,15 +597,8 @@ type ZkProgram<
         auxiliaryOutput?: ProvableType;
       };
     };
-  },
-  Methods extends {
-    [I in keyof Config['methods']]: Method<
-      InferProvableOrUndefined<Get<Config, 'publicInput'>>,
-      InferProvableOrVoid<Get<Config, 'publicOutput'>>,
-      Config['methods'][I]
-    >;
   }
-> = ReturnType<typeof ZkProgram<Config, Methods>>;
+> = ReturnType<typeof ZkProgram<Config>>;
 
 class SelfProof<PublicInput, PublicOutput> extends Proof<
   PublicInput,

diff --git a/src/lib/testing/constraint-system.ts b/src/lib/testing/constraint-system.ts
@@ -114,7 +114,7 @@ constraintSystem.fromZkProgram = function fromZkProgram<
   methodName: K,
   test: ConstraintSystemTest
 ) {
-  let program_: ZkProgram<any, any> = program as any;
+  let program_: ZkProgram<any> = program as any;
   let from: any = [...program_.privateInputTypes[methodName]];
   if (program_.publicInputType !== Undefined) {
     from.unshift(program_.publicInputType);