Enforce ordering zkp inputs: constants,public,private

Honestly this feels a little unnatural to me. Might want to reverse: private, public, constant and then change all the prove/verify methods to that ordering
Sunscreen-tech · Jul 15, 2023 · 2fd8c90 · 2fd8c90
1 parent 5dc513c
commit 2fd8c90
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 16 deletions.
diff --git a/benchmarks/bfv_zkp/src/bfv.rs b/benchmarks/bfv_zkp/src/bfv.rs
@@ -203,15 +203,15 @@ type BfvPoly<F> = RnsRingPolynomial<F, POLY_DEGREE, 1>;
 
 #[zkp_program]
 fn prove_enc<F: BackendField>(
-    m: BfvPoly<F>,
-    e_1: BfvPoly<F>,
-    e_2: BfvPoly<F>,
-    u: BfvPoly<F>,
     #[constant] expected_c_0: BfvPoly<F>,
     #[constant] expected_c_1: BfvPoly<F>,
     #[constant] p_0: BfvPoly<F>,
     #[constant] p_1: BfvPoly<F>,
     #[constant] delta: NativeField<F>,
+    m: BfvPoly<F>,
+    e_1: BfvPoly<F>,
+    e_2: BfvPoly<F>,
+    u: BfvPoly<F>,
 ) {
     let q = NativeField::<F>::from(CIPHER_MODULUS).into_program_node();
 

diff --git a/examples/ordering_zkp/src/main.rs b/examples/ordering_zkp/src/main.rs
@@ -5,7 +5,7 @@ use sunscreen::{
 };
 
 #[zkp_program]
-fn greater_than<F: BackendField>(a: NativeField<F>, #[constant] b: NativeField<F>) {
+fn greater_than<F: BackendField>(#[constant] b: NativeField<F>, a: NativeField<F>) {
     a.constrain_gt_bounded(b, 32)
 }
 

diff --git a/sunscreen/tests/zkp_program_tests.rs b/sunscreen/tests/zkp_program_tests.rs
@@ -171,9 +171,9 @@ fn can_declare_array_inputs() {
 fn builder_methods_work() {
     #[zkp_program]
     fn arbitrary<F: BackendField>(
+        #[constant] zss: [[NativeField<F>; 9]; 64],
         x: NativeField<F>,
         ys: [NativeField<F>; 9],
-        #[constant] zss: [[NativeField<F>; 9]; 64],
     ) {
         for y in ys {
             x.constrain_eq(y);

diff --git a/sunscreen_compiler_macros/src/zkp_program.rs b/sunscreen_compiler_macros/src/zkp_program.rs
@@ -106,27 +106,51 @@ fn parse_inner(_attr_params: ZkpProgramAttrs, input_fn: ItemFn) -> Result<TokenS
         }
     };
 
+    let mut private_seen = false;
+    let mut public_seen = false;
     let unwrapped_inputs = match extract_fn_arguments(inputs) {
         Ok(args) => {
             args.iter().map(|a| {
                 let mut arg_kind = ArgumentKind::Private;
 
-                match a.0.len() {
-                    0 => {},
-                    1 => {
-                        let ident = a.0[0].path().get_ident();
+                match &a.0[..] {
+                    [] => {
+                        private_seen = true;
+                    },
+                    [attr] => {
+                        let ident = attr.path().get_ident();
 
                         match ident.map(|x| x.to_string()).as_deref() {
-                            Some("private") => {},
-                            Some("public") => arg_kind = ArgumentKind::Public,
-                            Some("constant") => arg_kind = ArgumentKind::Constant,
+                            Some("private") => {
+                                private_seen = true;
+                            },
+                            Some("public") => {
+                                if private_seen {
+                                    return Err(Error::compile_error(attr.path().span(), 
+                                        "#[public] arguments must be specified before #[private] arguments "
+                                    ));
+                                }
+                                arg_kind = ArgumentKind::Public;
+                                public_seen = true;
+                            },
+                            Some("constant") => {
+                                if public_seen || private_seen {
+                                    return Err(Error::compile_error(attr.path().span(),
+                                        "#[constant] arguments must be specified before #[public] and #[private] arguments"
+                                    ));
+                                }
+                                arg_kind = ArgumentKind::Constant;
+                            },
                             _ => {
-                                return Err(Error::compile_error(a.0[0].path().span(), &format!("Expected #[private], #[public] or #[constant], found {}", a.0[0].path().to_token_stream())));
+                                return Err(Error::compile_error(attr.path().span(), &format!(
+                                    "Expected #[private], #[public] or #[constant], found {}",
+                                    attr.path().to_token_stream()
+                                )));
                             }
                         }
                     },
-                    _ => {
-                        return Err(Error::compile_error(a.0[1].span(), "ZKP program arguments may only have one attribute (#[private], #[public] or #[constant])."));
+                    [_, attr, ..] => {
+                        return Err(Error::compile_error(attr.span(), "ZKP program arguments may only have one attribute (#[private], #[public] or #[constant])."));
                     }
                 };