refactor prove/setup args

prover/cmd/prover/cmd/prove.go

@@ -15,11 +15,14 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var (
-	fInput  string
-	fOutput string
-	fLarge  bool
-)
+type proverArgsT struct {
+	input      string
+	output     string
+	large      bool
+	configFile string
+}
+
+var proverArgs proverArgsT
 
 // proveCmd represents the prove command
 var proveCmd = &cobra.Command{
@@ -31,71 +34,75 @@ var proveCmd = &cobra.Command{
 func init() {
 	rootCmd.AddCommand(proveCmd)
 
-	proveCmd.Flags().StringVar(&fInput, "in", "", "input file")
-	proveCmd.Flags().StringVar(&fOutput, "out", "", "output file")
-	proveCmd.Flags().BoolVar(&fLarge, "large", false, "run the large execution circuit")
-
+	proveCmd.Flags().StringVar(&proverArgs.input, "in", "", "input file")
+	proveCmd.Flags().StringVar(&proverArgs.output, "out", "", "output file")
+	proveCmd.Flags().BoolVar(&proverArgs.large, "large", false, "run the large execution circuit")
 }
 
 func cmdProve(cmd *cobra.Command, args []string) error {
+	proverArgs.configFile = fConfigFile
+	return Prove(cmd.Name(), proverArgs)
+}
+
+func Prove(cmdName string, args proverArgsT) error {
 	// TODO @gbotrel with a specific flag, we could compile the circuit and compare with the checksum of the
 	// asset we deserialize, to make sure we are using the circuit associated with the compiled binary and the setup.
 
 	// read config
 	cfg, err := config.NewConfigFromFile(fConfigFile)
 	if err != nil {
-		return fmt.Errorf("%s failed to read config file: %w", cmd.Name(), err)
+		return fmt.Errorf("%s failed to read config file: %w", cmdName, err)
 	}
 
 	// discover the type of the job from the input file name
-	jobExecution := strings.Contains(fInput, "getZkProof")
-	jobBlobDecompression := strings.Contains(fInput, "getZkBlobCompressionProof")
-	jobAggregation := strings.Contains(fInput, "getZkAggregatedProof")
+	jobExecution := strings.Contains(args.input, "getZkProof")
+	jobBlobDecompression := strings.Contains(args.input, "getZkBlobCompressionProof")
+	jobAggregation := strings.Contains(args.input, "getZkAggregatedProof")
 
 	if jobExecution {
 		req := &execution.Request{}
-		if err := readRequest(fInput, req); err != nil {
-			return fmt.Errorf("could not read the input file (%v): %w", fInput, err)
+		if err := readRequest(args.input, req); err != nil {
+			return fmt.Errorf("could not read the input file (%v): %w", args.input, err)
 		}
 		// we use the large traces in 2 cases;
-		// 1. the user explicitly asked for it (fLarge)
+		// 1. the user explicitly asked for it (args.large)
 		// 2. the job contains the large suffix and we are a large machine (cfg.Execution.CanRunLarge)
-		large := fLarge || (strings.Contains(fInput, "large") && cfg.Execution.CanRunFullLarge)
+		large := args.large || (strings.Contains(args.input, "large") && cfg.Execution.CanRunFullLarge)
 
 		resp, err := execution.Prove(cfg, req, large)
 		if err != nil {
 			return fmt.Errorf("could not prove the execution: %w", err)
 		}
 
-		return writeResponse(fOutput, resp)
+		return writeResponse(args.output, resp)
 	}
 
 	if jobBlobDecompression {
 		req := &blobdecompression.Request{}
-		if err := readRequest(fInput, req); err != nil {
-			return fmt.Errorf("could not read the input file (%v): %w", fInput, err)
+		if err := readRequest(args.input, req); err != nil {
+			return fmt.Errorf("could not read the input file (%v): %w", args.input, err)
 		}
 
 		resp, err := blobdecompression.Prove(cfg, req)
 		if err != nil {
 			return fmt.Errorf("could not prove the blob decompression: %w", err)
 		}
 
-		return writeResponse(fOutput, resp)
+		return writeResponse(args.output, resp)
 	}
 
 	if jobAggregation {
 		req := &aggregation.Request{}
-		if err := readRequest(fInput, req); err != nil {
-			return fmt.Errorf("could not read the input file (%v): %w", fInput, err)
+		if err := readRequest(args.input, req); err != nil {
+			return fmt.Errorf("could not read the input file (%v): %w", args.input, err)
 		}
 
 		resp, err := aggregation.Prove(cfg, req)
 		if err != nil {
 			return fmt.Errorf("could not prove the aggregation: %w", err)
 		}
 
-		return writeResponse(fOutput, resp)
+		return writeResponse(args.output, resp)
 	}
 
 	return errors.New("unknown job type")

prover/cmd/prover/cmd/setup.go

@@ -32,12 +32,15 @@ import (
 	"github.com/consensys/gnark/backend/plonk"
 )
 
-var (
-	fForce     bool
-	fCircuits  string
-	fDictPath  string
-	fAssetsDir string
-)
+type setupArgsT struct {
+	force      bool
+	circuits   string
+	dictPath   string
+	assetsDir  string
+	configFile string
+}
+
+var setupArgs setupArgsT
 
 // setupCmd represents the setup command
 var setupCmd = &cobra.Command{
@@ -59,25 +62,30 @@ var allCircuits = []string{
 
 func init() {
 	rootCmd.AddCommand(setupCmd)
-	setupCmd.Flags().BoolVar(&fForce, "force", false, "overwrites existing files")
-	setupCmd.Flags().StringVar(&fCircuits, "circuits", strings.Join(allCircuits, ","), "comma separated list of circuits to setup")
-	setupCmd.Flags().StringVar(&fDictPath, "dict", "", "path to the dictionary file used in blob (de)compression")
-	setupCmd.Flags().StringVar(&fAssetsDir, "assets-dir", "", "path to the directory where the assets are stored (override conf)")
+	setupCmd.Flags().BoolVar(&setupArgs.force, "force", false, "overwrites existing files")
+	setupCmd.Flags().StringVar(&setupArgs.circuits, "circuits", strings.Join(allCircuits, ","), "comma separated list of circuits to setup")
+	setupCmd.Flags().StringVar(&setupArgs.dictPath, "dict", "", "path to the dictionary file used in blob (de)compression")
+	setupCmd.Flags().StringVar(&setupArgs.assetsDir, "assets-dir", "", "path to the directory where the assets are stored (override conf)")
 
 	viper.BindPFlag("assets_dir", setupCmd.Flags().Lookup("assets-dir"))
 }
 
 func cmdSetup(cmd *cobra.Command, args []string) error {
+	setupArgs.configFile = fConfigFile
+	return Setup(cmd.Name(), cmd.Context(), setupArgs)
+}
+
+func Setup(cmdName string, context context.Context, args setupArgsT) error {
 	// read config
-	cfg, err := config.NewConfigFromFile(fConfigFile)
+	cfg, err := config.NewConfigFromFile(args.configFile)
 	if err != nil {
-		return fmt.Errorf("%s failed to read config file: %w", cmd.Name(), err)
+		return fmt.Errorf("%s failed to read config file: %w", cmdName, err)
 	}
 
-	if fDictPath != "" {
+	if args.dictPath != "" {
 		// fail early if the dictionary file is not found but was specified.
-		if _, err := os.Stat(fDictPath); err != nil {
-			return fmt.Errorf("%s dictionary file not found: %w", cmd.Name(), err)
+		if _, err := os.Stat(args.dictPath); err != nil {
+			return fmt.Errorf("%s dictionary file not found: %w", cmdName, err)
 		}
 	}
 
@@ -86,10 +94,10 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 	for _, c := range allCircuits {
 		inCircuits[circuits.CircuitID(c)] = false
 	}
-	_inCircuits := strings.Split(fCircuits, ",")
+	_inCircuits := strings.Split(args.circuits, ",")
 	for _, c := range _inCircuits {
 		if _, ok := inCircuits[circuits.CircuitID(c)]; !ok {
-			return fmt.Errorf("%s unknown circuit: %s", cmd.Name(), c)
+			return fmt.Errorf("%s unknown circuit: %s", cmdName, c)
 		}
 		inCircuits[circuits.CircuitID(c)] = true
 	}
@@ -101,7 +109,7 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 	var srsProvider circuits.SRSProvider
 	srsProvider, err = circuits.NewSRSStore(cfg.PathForSRS())
 	if err != nil {
-		return fmt.Errorf("%s failed to create SRS provider: %w", cmd.Name(), err)
+		return fmt.Errorf("%s failed to create SRS provider: %w", cmdName, err)
 	}
 
 	// for each circuit, we start by compiling the circuit
@@ -128,9 +136,9 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 			zkEvm := zkevm.FullZkEvm(&limits)
 			builder = execution.NewBuilder(zkEvm)
 		case circuits.BlobDecompressionV0CircuitID, circuits.BlobDecompressionV1CircuitID:
-			dict, err = os.ReadFile(fDictPath)
+			dict, err = os.ReadFile(args.dictPath)
 			if err != nil {
-				return fmt.Errorf("%s failed to read dictionary file: %w", cmd.Name(), err)
+				return fmt.Errorf("%s failed to read dictionary file: %w", cmdName, err)
 			}
 
 			if c == circuits.BlobDecompressionV0CircuitID {
@@ -151,14 +159,14 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 			continue // dummy, aggregation, emulation or public input circuits are handled later
 		}
 
-		if err := updateSetup(cmd.Context(), cfg, srsProvider, c, builder, extraFlags); err != nil {
+		if err := updateSetup(context, cfg, args.force, srsProvider, c, builder, extraFlags); err != nil {
 			return err
 		}
 		if dict != nil {
 			// we save the dictionary to disk
 			dictPath := filepath.Join(cfg.PathForSetup(string(c)), config.DictionaryFileName)
 			if err := os.WriteFile(dictPath, dict, 0600); err != nil {
-				return fmt.Errorf("%s failed to write dictionary file: %w", cmd.Name(), err)
+				return fmt.Errorf("%s failed to write dictionary file: %w", cmdName, err)
 			}
 		}
 
@@ -172,7 +180,7 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 	// get verifying key for public-input circuit
 	piSetup, err := circuits.LoadSetup(cfg, circuits.PublicInputInterconnectionCircuitID)
 	if err != nil {
-		return fmt.Errorf("%s failed to load public input interconnection setup: %w", cmd.Name(), err)
+		return fmt.Errorf("%s failed to load public input interconnection setup: %w", cmdName, err)
 	}
 
 	// first, we need to collect the verifying keys
@@ -196,7 +204,7 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 				return fmt.Errorf("unknown dummy circuit: %s", allowedInput)
 			}
 
-			vk, err := getDummyCircuitVK(cmd.Context(), cfg, srsProvider, circuits.CircuitID(allowedInput), dummy.NewBuilder(mockID, curveID.ScalarField()))
+			vk, err := getDummyCircuitVK(context, cfg, srsProvider, circuits.CircuitID(allowedInput), dummy.NewBuilder(mockID, curveID.ScalarField()))
 			if err != nil {
 				return err
 			}
@@ -209,15 +217,15 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 		vkPath := filepath.Join(setupPath, config.VerifyingKeyFileName)
 		vk := plonk.NewVerifyingKey(ecc.BLS12_377)
 		if err := circuits.ReadVerifyingKey(vkPath, vk); err != nil {
-			return fmt.Errorf("%s failed to read verifying key for circuit %s: %w", cmd.Name(), allowedInput, err)
+			return fmt.Errorf("%s failed to read verifying key for circuit %s: %w", cmdName, allowedInput, err)
 		}
 
 		allowedVkForAggregation = append(allowedVkForAggregation, vk)
 	}
 
 	// we need to compute the digest of the verifying keys & store them in the manifest
 	// for the aggregation circuits to be able to check compatibility at run time with the proofs
-	allowedVkForAggregationDigests := listOfCheckum(allowedVkForAggregation)
+	allowedVkForAggregationDigests := listOfChecksums(allowedVkForAggregation)
 	extraFlagsForAggregationCircuit := map[string]any{
 		"allowedVkForAggregationDigests": allowedVkForAggregationDigests,
 	}
@@ -229,7 +237,7 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 		logrus.Infof("setting up %s (numProofs=%d)", c, numProofs)
 
 		builder := aggregation.NewBuilder(numProofs, cfg.Aggregation.AllowedInputs, piSetup, allowedVkForAggregation)
-		if err := updateSetup(cmd.Context(), cfg, srsProvider, c, builder, extraFlagsForAggregationCircuit); err != nil {
+		if err := updateSetup(context, cfg, args.force, srsProvider, c, builder, extraFlagsForAggregationCircuit); err != nil {
 			return err
 		}
 
@@ -238,7 +246,7 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 		vkPath := filepath.Join(setupPath, config.VerifyingKeyFileName)
 		vk := plonk.NewVerifyingKey(ecc.BW6_761)
 		if err := circuits.ReadVerifyingKey(vkPath, vk); err != nil {
-			return fmt.Errorf("%s failed to read verifying key for circuit %s: %w", cmd.Name(), c, err)
+			return fmt.Errorf("%s failed to read verifying key for circuit %s: %w", cmdName, c, err)
 		}
 
 		allowedVkForEmulation = append(allowedVkForEmulation, vk)
@@ -248,7 +256,7 @@ func cmdSetup(cmd *cobra.Command, args []string) error {
 	c := circuits.EmulationCircuitID
 	logrus.Infof("setting up %s", c)
 	builder := emulation.NewBuilder(allowedVkForEmulation)
-	return updateSetup(cmd.Context(), cfg, srsProvider, c, builder, nil)
+	return updateSetup(context, cfg, args.force, srsProvider, c, builder, nil)
 
 }
 
@@ -281,7 +289,7 @@ func getDummyCircuitVK(ctx context.Context, cfg *config.Config, srsProvider circ
 // and if so, if the checksums match.
 // if the files already exist and the checksums match, it skips the setup.
 // else it does the setup and writes the assets to disk.
-func updateSetup(ctx context.Context, cfg *config.Config, srsProvider circuits.SRSProvider, circuit circuits.CircuitID, builder circuits.Builder, extraFlags map[string]any) error {
+func updateSetup(ctx context.Context, cfg *config.Config, force bool, srsProvider circuits.SRSProvider, circuit circuits.CircuitID, builder circuits.Builder, extraFlags map[string]any) error {
 	if extraFlags == nil {
 		extraFlags = make(map[string]any)
 	}
@@ -297,7 +305,7 @@ func updateSetup(ctx context.Context, cfg *config.Config, srsProvider circuits.S
 	setupPath := cfg.PathForSetup(string(circuit))
 	manifestPath := filepath.Join(setupPath, config.ManifestFileName)
 
-	if !fForce {
+	if !force {
 		// we may want to skip setup if the files already exist
 		// and the checksums match
 		// read manifest if already exists
@@ -325,12 +333,13 @@ func updateSetup(ctx context.Context, cfg *config.Config, srsProvider circuits.S
 	return setup.WriteTo(setupPath)
 }
 
-// listOfCheckum Computes a list of SHA256 checksums for a list of assets, the result is given
+// listOfChecksums Computes a list of SHA256 checksums for a list of assets, the result is given
 // in hexstring.
-func listOfCheckum[T io.WriterTo](assets []T) []string {
+func listOfChecksums[T io.WriterTo](assets []T) []string {
 	res := make([]string, len(assets))
+	h := sha256.New()
 	for i := range assets {
-		h := sha256.New()
+		h.Reset()
 		_, err := assets[i].WriteTo(h)
 		if err != nil {
 			// It is unexpected that writing in a hasher could possibly fail.