feat: support patching properly

.gitignore

@@ -1,4 +1,4 @@
-benchmarks/**/patched-*
+benchmarks/**/*.patched.circom
 benchmarks/**/*.json
 benchmarks/**/*.sym
 benchmarks/**/*.r1cs

picus.rkt

@@ -9,14 +9,16 @@
          (prefix-in pre: "./picus/precondition.rkt")
          "picus/tmpdir.rkt"
          "picus/ansi.rkt"
-         "picus/verbose.rkt")
+         "picus/verbose.rkt"
+         "picus/global-inputs.rkt")
 
 ; =====================================
 ; ======== commandline parsing ========
 ; =====================================
 ; parse command line arguments
 (define arg-r1cs #f)
 (define arg-circom #f)
+(define arg-patch? #f)
 (define arg-opt-level #f)
 (define arg-clean? #t)
 (define arg-timeout 5000)
@@ -41,6 +43,8 @@
  #:once-each
  [("--noclean") "do not clean up temporary files (default: false)"
                 (set! arg-clean? #f)]
+ [("--patch-circom") "patch circom file to add public inputs (only applicable for --circom, default: false)"
+                     (set! arg-patch? #t)]
  [("--opt-level") p-opt-level "optimization level for circom compilation (only applicable for --circom, default: 0)"
                   (set! arg-opt-level
                         (match p-opt-level
@@ -87,25 +91,70 @@
   (tokamak:exit "specify either --r1cs or --circom"))
 
 (unless (implies arg-opt-level arg-circom)
-  (tokamak:exit "optimization level only applicable for --circom"))
+  (tokamak:exit "--opt-level only applicable for --circom"))
 
-(when arg-circom
+(unless (implies arg-patch? arg-circom)
+  (tokamak:exit "--patch-circom only applicable for --circom"))
+
+;; compile-circom :: path? -> path?
+;; compile circom file to r1cs file
+(define (compile-circom circom-path)
   (unless (system* (find-executable-path "circom")
                    "-o"
                    (get-tmpdir)
                    "--r1cs"
-                   arg-circom
+                   circom-path
                    "--sym"
                    (match arg-opt-level
                      [(or #f "0") "--O0"]
                      ["1" "--O1"]
                      ["2" "--O2"]))
     (tokamak:exit "circom compilation failed"))
-  (set! arg-r1cs (~a (build-path (get-tmpdir)
-                                 (file-name-from-path
-                                  (path-replace-extension arg-circom ".r1cs"))))))
+  (~a (build-path
+       (get-tmpdir)
+       (file-name-from-path (path-replace-extension circom-path ".r1cs")))))
+
+;; compile+patch-circom :: path? boolean? -> (values path? r1cs?)
+(define (compile+patch-circom circom-path patch?)
+  (define r1cs-path (compile-circom circom-path))
+  (cond
+    [patch?
+     (define r0 (r1cs:read-r1cs r1cs-path))
+     (cond
+       ;; no public inputs
+       [(zero? (r1cs:get-npubin r0))
+        (define patched-circom-path
+          (path-replace-extension circom-path ".patched.circom"))
+
+        (with-output-to-file patched-circom-path
+          #:exists 'replace
+          (λ ()
+            ;; HACK: this assumes that the circom file has a line with "component main ="
+            ;; which is the case for most circom files with no public inputs
+            (displayln
+             (string-replace
+              (file->string circom-path)
+              "component main ="
+              (format "component main {public [~a]} ="
+                      (string-join
+                       (get-global-inputs
+                        r1cs-path
+                        (path-replace-extension r1cs-path ".sym"))
+                       ", "))))))
+
+        (define patched-r1cs-path (compile-circom patched-circom-path))
+        (values patched-r1cs-path (r1cs:read-r1cs patched-r1cs-path))]
+       ;; already has public inputs
+       [else (values r1cs-path r0)])]
+    ;; not patching
+    [else (values r1cs-path (r1cs:read-r1cs r1cs-path))]))
+
+(define-values (r1cs-path r0)
+  (cond
+    [arg-r1cs (values arg-r1cs (r1cs:read-r1cs arg-r1cs))]
+    [else (compile+patch-circom arg-circom arg-patch?)]))
 
-(printf "# r1cs file: ~a\n" arg-r1cs)
+(printf "# r1cs file: ~a\n" r1cs-path)
 (printf "# timeout: ~a\n" arg-timeout)
 (printf "# solver: ~a\n" arg-solver)
 (printf "# selector: ~a\n" arg-selector)
@@ -130,7 +179,6 @@
 ; ==================================
 ; ======== main preparation ========
 ; ==================================
-(define r0 (r1cs:read-r1cs arg-r1cs))
 (define nwires (r1cs:get-nwires r0))
 (define mconstraints (r1cs:get-mconstraints r0))
 (printf "# number of wires: ~a\n" nwires)
@@ -193,7 +241,7 @@
 ; p1cnsts
 ;    | (downstream queries)
 ;   ...
-(define path-sym (string-replace arg-r1cs ".r1cs" ".sym"))
+(define path-sym (string-replace r1cs-path ".r1cs" ".sym"))
 (define-values (res res-ks res-us res-info)
   (dpvl:apply-algorithm
    r0 nwires mconstraints
@@ -226,7 +274,7 @@
   (sort (hash->list info) string<? #:key car))
 
 (when (equal? 'unsafe res)
-  (printf "# ~a is underconstrained. Below is a counterexample:\n" arg-r1cs)
+  (printf "# ~a is underconstrained. Below is a counterexample:\n" r1cs-path)
   (match-define (list input-info output1-info output2-info other-info) res-info)
   (define output1-ordered (order output1-info))
   (define output2-ordered (order output2-info))

picus/public-inputs.rkt → picus/global-inputs.rkt

@@ -1,13 +1,13 @@
 #lang racket/base
 
-(provide get-public-inputs)
+(provide get-global-inputs)
 
 (require racket/set
          racket/match
          csv-reading
          (prefix-in r1cs: "./r1cs/r1cs-grammar.rkt"))
 
-(define (get-public-inputs r1cs-path sym-path)
+(define (get-global-inputs r1cs-path sym-path)
   (define r0 (r1cs:read-r1cs r1cs-path))
   (define input-signals (r1cs:r1cs-inputs r0))
   (define ins (list->set input-signals))
@@ -24,9 +24,9 @@
 
 (module+ test
   (require rackunit)
-  (check-equal? (get-public-inputs "../benchmarks/circomlib-cff5ab6/[email protected]"
+  (check-equal? (get-global-inputs "../benchmarks/circomlib-cff5ab6/[email protected]"
                                    "../benchmarks/circomlib-cff5ab6/[email protected]")
                 (set "in"))
-  (check-equal? (get-public-inputs "../benchmarks/circomlib-cff5ab6/[email protected]"
+  (check-equal? (get-global-inputs "../benchmarks/circomlib-cff5ab6/[email protected]"
                                    "../benchmarks/circomlib-cff5ab6/[email protected]")
                 (set "x1" "y1" "x2" "y2")))

tests/circomlib-test.rkt

@@ -10,7 +10,7 @@
            racket/format
            racket/match
            rackunit
-           "../picus/public-inputs.rkt")
+           "../picus/global-inputs.rkt")
 
   (define filenames
     (command-line
@@ -25,67 +25,15 @@
   (define (make-run-config filename #:timeout [timeout 60])
     (run-config filename timeout))
 
-  (define (get-full-path filename)
-    (~a (build-path benchmark-dir filename)))
-
   (define (flush)
     (flush-output)
     (flush-output (current-error-port)))
 
-  (define (compile-circom circom-path)
-    (flush)
-    (system* (find-executable-path "circom")
-             "-o"
-             benchmark-dir
-             "--r1cs"
-             circom-path
-             "--sym"
-             "--O0"
-             #;"--jsons")
-    (flush))
-
-  (define (has-public-input? orig-content)
-    ;; match a line starting with "component main =".
-    ;; This is very hacky, but it is the case that in our benchmarks,
-    ;; every file that contains "component main =" has no public input.
-    (not (regexp-match? #px"(?m:^component main =)" orig-content)))
-
   (define-check (check:core run-conf expected)
     (match-define (run-config filename timeout) run-conf)
     (printf "=== checking ~a ===\n" filename)
     (printf "##[group]~a\n" filename)
-    (printf "Compiling circom file\n")
-    (define circom-path (get-full-path (format "~a.circom" filename)))
-    (define r1cs-path (get-full-path (format "~a.r1cs" filename)))
-    (define sym-path (get-full-path (format "~a.sym" filename)))
-
-    ;; initially compile Circom to get information about public inputs
-    (compile-circom circom-path)
-
-    (define orig-content (file->string circom-path))
-
-    (define real-r1cs-path
-      (cond
-        [(has-public-input? orig-content) r1cs-path]
-        [else
-         (define public-inputs (get-public-inputs r1cs-path sym-path))
-         (define patched-circom-path (get-full-path (format "patched-~a.circom" filename)))
-         (define patched-r1cs-path (get-full-path (format "patched-~a.r1cs" filename)))
-
-         (with-output-to-file patched-circom-path
-           #:exists 'replace
-           (λ ()
-             (displayln
-              (string-replace
-               orig-content
-               "component main ="
-               (format "component main {public [~a]} ="
-                       (string-join public-inputs ", "))))))
-
-         (compile-circom patched-circom-path)
-         patched-r1cs-path]))
-
-    (printf "Starting Picus\n")
+    (flush)
     (define-values (in out) (make-pipe))
     (define orig-out (current-output-port))
     (define string-port (open-output-string))
@@ -99,7 +47,10 @@
                             "--timeout" "5000"
                             "--weak"
                             "--verbose" "1"
-                            "--r1cs" real-r1cs-path)]
+                            "--circom" (~a (build-path
+                                            benchmark-dir
+                                            (format "~a.circom" filename)))
+                            "--patch-circom")]
                    [current-output-port out])
       (define main-thd
         (thread (λ ()