Merge pull request #247 from bufferhe4d/generic-asserteq

Implement generic assert_eq

Author: mimoo

examples/fixture/asm/kimchi/generic_assert_eq.asm

@@ -0,0 +1,33 @@
+@ noname.0.7.0
+@ public inputs: 3
+
+DoubleGeneric<1>
+DoubleGeneric<1>
+DoubleGeneric<1>
+DoubleGeneric<1,0,0,0,-3>
+DoubleGeneric<1,0,0,0,-3>
+DoubleGeneric<1,0,0,0,-1>
+DoubleGeneric<1,0,0,0,-2>
+DoubleGeneric<1,0,-1,0,2>
+DoubleGeneric<1,0,0,0,-5>
+DoubleGeneric<1,0,0,0,-1>
+DoubleGeneric<1,0,0,0,-2>
+DoubleGeneric<1,0,-1,0,2>
+DoubleGeneric<0,0,-1,1>
+DoubleGeneric<1,1,-1>
+DoubleGeneric<1,1,-1>
+DoubleGeneric<1,0,0,0,-5>
+DoubleGeneric<1,0,0,0,-1>
+DoubleGeneric<1,0,0,0,-1>
+DoubleGeneric<1,0,0,0,-2>
+DoubleGeneric<1,0,0,0,-3>
+DoubleGeneric<1,0,0,0,-4>
+DoubleGeneric<1,0,0,0,-5>
+(0,0) -> (5,0) -> (9,0) -> (12,1) -> (13,0) -> (16,0) -> (17,0)
+(1,0) -> (6,0) -> (10,0) -> (14,0) -> (18,0)
+(2,0) -> (3,0) -> (4,0) -> (7,0) -> (11,0) -> (12,0) -> (14,1) -> (19,0)
+(7,2) -> (8,0)
+(11,2) -> (15,0)
+(12,2) -> (13,1)
+(13,2) -> (20,0)
+(14,2) -> (21,0)

examples/fixture/asm/r1cs/generic_assert_eq.asm

@@ -0,0 +1,18 @@
+@ noname.0.7.0
+@ public inputs: 3
+
+3 == (v_3) * (1)
+3 == (v_3) * (1)
+1 == (v_1) * (1)
+2 == (v_2) * (1)
+5 == (v_3 + 2) * (1)
+1 == (v_1) * (1)
+2 == (v_2) * (1)
+v_4 == (v_3) * (v_1)
+5 == (v_3 + 2) * (1)
+1 == (v_1) * (1)
+1 == (v_1) * (1)
+2 == (v_2) * (1)
+3 == (v_3) * (1)
+4 == (v_1 + v_4) * (1)
+5 == (v_2 + v_3) * (1)

examples/generic_assert_eq.no

@@ -0,0 +1,64 @@
+const size = 2;
+struct Thing {
+    xx: Field,
+    yy: [Field; 2],
+}
+
+struct Nestedthing {
+    xx: Field,
+    another: [Another; 2],
+}
+
+struct Another {
+    aa: Field,
+    bb: [Field; 2],
+}
+
+fn init_arr(element: Field, const LEN: Field) -> [Field; LEN] {
+    let arr = [element; LEN];
+    return arr;
+}
+
+fn main(pub public_arr: [Field; 2], pub public_input: Field) {
+    let generic_arr = init_arr(public_input, size);
+    let arr = [3, 3];
+    
+    assert_eq(generic_arr, arr);
+    let mut concrete_arr = [1, 2];
+
+    // instead of the following:
+    // assert_eq(public_arr[0], concrete_arr[0]);
+    // assert_eq(public_arr[1], concrete_arr[1]);
+    // we can write:
+    assert_eq(public_arr, concrete_arr);
+
+    let thing = Thing { xx: 5, yy: [1, 2] };
+    let other_thing = Thing { xx: generic_arr[0] + 2, yy: public_arr };
+
+    // instead of the following:
+    // assert_eq(thing.xx, other_thing.xx);
+    // assert_eq(thing.yy[0], other_thing.yy[0]);
+    // assert_eq(thing.yy[1], other_thing.yy[1]);
+    // we can write:
+    assert_eq(thing, other_thing);
+
+    let nested_thing = Nestedthing { xx: 5, another: [
+        Another { aa: public_arr[0], bb: [1, 2] }, 
+        Another { aa: generic_arr[1], bb: [4, 5] }
+    ] };
+    let other_nested_thing = Nestedthing { xx: generic_arr[0] + 2, another: [
+        Another { aa: 1, bb: public_arr }, 
+        Another { aa: 3, bb: [public_arr[0] + (public_input * public_arr[0]), public_arr[1] + public_input] }
+    ] };
+    
+    // instead of the following:
+    // assert_eq(nested_thing.xx, other_nested_thing.xx);
+    // assert_eq(nested_thing.another[0].aa, other_nested_thing.another[0].aa);
+    // assert_eq(nested_thing.another[0].bb[0], other_nested_thing.another[0].bb[0]);
+    // assert_eq(nested_thing.another[0].bb[1], other_nested_thing.another[0].bb[1]);
+    // assert_eq(nested_thing.another[1].aa, other_nested_thing.another[1].aa);
+    // assert_eq(nested_thing.another[1].bb[0], other_nested_thing.another[1].bb[0]);
+    // assert_eq(nested_thing.another[1].bb[1], other_nested_thing.another[1].bb[1]);
+    // we can write:
+    assert_eq(nested_thing, other_nested_thing);
+}

src/error.rs

@@ -45,6 +45,8 @@ pub enum ErrorKind {
     AssignmentToImmutableVariable,
     #[error("the {0} of assert_eq must be of type Field or BigInt. It was of type {1}")]
     AssertTypeMismatch(&'static str, TyKind),
+    #[error("the types in assert_eq don't match: expected {0} but got {1}")]
+    AssertEqTypeMismatch(TyKind, TyKind),
     #[error(
         "the dependency `{0}` does not appear to be listed in your manifest file `Noname.toml`"
     )]

src/stdlib/builtins.rs

@@ -13,7 +13,8 @@ use crate::{
     constants::Span,
     error::{Error, ErrorKind, Result},
     helpers::PrettyField,
-    parser::types::{GenericParameters, TyKind},
+    parser::types::{GenericParameters, ModulePath, TyKind},
+    type_checker::FullyQualified,
     var::{ConstOrCell, Value, Var},
 };
 
@@ -33,13 +34,118 @@ impl Module for BuiltinsLib {
     fn get_fns<B: Backend>() -> Vec<(&'static str, FnInfoType<B>, bool)> {
         vec![
             (ASSERT_FN, assert_fn, false),
-            (ASSERT_EQ_FN, assert_eq_fn, false),
+            (ASSERT_EQ_FN, assert_eq_fn, true),
             // true -> skip argument type checking for log
             (LOG_FN, log_fn, true),
         ]
     }
 }
 
+/// Represents a comparison that needs to be made
+enum Comparison<B: Backend> {
+    /// Compare two variables
+    Vars(B::Var, B::Var),
+    /// Compare a variable with a constant
+    VarConst(B::Var, B::Field),
+    /// Compare two constants
+    Constants(B::Field, B::Field),
+}
+
+/// Helper function to generate all comparisons
+fn assert_eq_values<B: Backend>(
+    compiler: &CircuitWriter<B>,
+    lhs_info: &VarInfo<B::Field, B::Var>,
+    rhs_info: &VarInfo<B::Field, B::Var>,
+    typ: &TyKind,
+    span: Span,
+) -> Vec<Comparison<B>> {
+    let mut comparisons = Vec::new();
+
+    match typ {
+        // Field and Bool has the same logic
+        TyKind::Field { .. } | TyKind::Bool => {
+            let lhs_var = &lhs_info.var[0];
+            let rhs_var = &rhs_info.var[0];
+            match (lhs_var, rhs_var) {
+                (ConstOrCell::Const(a), ConstOrCell::Const(b)) => {
+                    comparisons.push(Comparison::Constants(a.clone(), b.clone()));
+                }
+                (ConstOrCell::Const(cst), ConstOrCell::Cell(cvar))
+                | (ConstOrCell::Cell(cvar), ConstOrCell::Const(cst)) => {
+                    comparisons.push(Comparison::VarConst(cvar.clone(), cst.clone()));
+                }
+                (ConstOrCell::Cell(lhs), ConstOrCell::Cell(rhs)) => {
+                    comparisons.push(Comparison::Vars(lhs.clone(), rhs.clone()));
+                }
+            }
+        }
+
+        // Arrays (fixed size)
+        TyKind::Array(element_type, size) => {
+            let size = *size as usize;
+            let element_size = compiler.size_of(element_type);
+
+            // compare each element recursively
+            for i in 0..size {
+                let start = i * element_size;
+                let mut element_comparisons = assert_eq_values(
+                    compiler,
+                    &VarInfo::new(
+                        Var::new(lhs_info.var.range(start, element_size).to_vec(), span),
+                        false,
+                        Some(*element_type.clone()),
+                    ),
+                    &VarInfo::new(
+                        Var::new(rhs_info.var.range(start, element_size).to_vec(), span),
+                        false,
+                        Some(*element_type.clone()),
+                    ),
+                    element_type,
+                    span,
+                );
+                comparisons.append(&mut element_comparisons);
+            }
+        }
+
+        // Custom types (structs)
+        TyKind::Custom { module, name } => {
+            let qualified = FullyQualified::new(module, name);
+            let struct_info = compiler.struct_info(&qualified).expect("struct not found");
+
+            // compare each field recursively
+            let mut offset = 0;
+            for (_, field_type) in &struct_info.fields {
+                let field_size = compiler.size_of(field_type);
+                let mut field_comparisons = assert_eq_values(
+                    compiler,
+                    &VarInfo::new(
+                        Var::new(lhs_info.var.range(offset, field_size).to_vec(), span),
+                        false,
+                        Some(field_type.clone()),
+                    ),
+                    &VarInfo::new(
+                        Var::new(rhs_info.var.range(offset, field_size).to_vec(), span),
+                        false,
+                        Some(field_type.clone()),
+                    ),
+                    field_type,
+                    span,
+                );
+                comparisons.append(&mut field_comparisons);
+                offset += field_size;
+            }
+        }
+
+        // GenericSizedArray should be monomorphized to Array before reaching here
+        // no need to handle it seperately
+        TyKind::GenericSizedArray(_, _) => {
+            unreachable!("GenericSizedArray should be monomorphized")
+        }
+    }
+
+    comparisons
+}
+
 /// Asserts that two vars are equal.
 fn assert_eq_fn<B: Backend>(
     compiler: &mut CircuitWriter<B>,
@@ -52,67 +158,53 @@ fn assert_eq_fn<B: Backend>(
     let lhs_info = &vars[0];
     let rhs_info = &vars[1];
 
-    // they are both of type field
-    if !matches!(lhs_info.typ, Some(TyKind::Field { .. })) {
-        let lhs = lhs_info.typ.clone().ok_or_else(|| {
-            Error::new(
-                "constraint-generation",
-                ErrorKind::UnexpectedError("No type info for lhs of assertion"),
-                span,
-            )
-        })?;
-
-        Err(Error::new(
+    // get types of both arguments
+    let lhs_type = lhs_info.typ.as_ref().ok_or_else(|| {
+        Error::new(
             "constraint-generation",
-            ErrorKind::AssertTypeMismatch("rhs", lhs),
+            ErrorKind::UnexpectedError("No type info for lhs of assertion"),
             span,
-        ))?
-    }
+        )
+    })?;
 
-    if !matches!(rhs_info.typ, Some(TyKind::Field { .. })) {
-        let rhs = rhs_info.typ.clone().ok_or_else(|| {
-            Error::new(
-                "constraint-generation",
-                ErrorKind::UnexpectedError("No type info for rhs of assertion"),
-                span,
-            )
-        })?;
+    let rhs_type = rhs_info.typ.as_ref().ok_or_else(|| {
+        Error::new(
+            "constraint-generation",
+            ErrorKind::UnexpectedError("No type info for rhs of assertion"),
+            span,
+        )
+    })?;
 
-        Err(Error::new(
+    // they have the same type
+    if !lhs_type.match_expected(rhs_type, false) {
+        return Err(Error::new(
             "constraint-generation",
-            ErrorKind::AssertTypeMismatch("rhs", rhs),
+            ErrorKind::AssertEqTypeMismatch(lhs_type.clone(), rhs_type.clone()),
             span,
-        ))?
+        ));
     }
 
-    // retrieve the values
-    let lhs_var = &lhs_info.var;
-    assert_eq!(lhs_var.len(), 1);
-    let lhs_cvar = &lhs_var[0];
-
-    let rhs_var = &rhs_info.var;
-    assert_eq!(rhs_var.len(), 1);
-    let rhs_cvar = &rhs_var[0];
-
-    match (lhs_cvar, rhs_cvar) {
-        // two constants
-        (ConstOrCell::Const(a), ConstOrCell::Const(b)) => {
-            if a != b {
-                Err(Error::new(
-                    "constraint-generation",
-                    ErrorKind::AssertionFailed,
-                    span,
-                ))?
-            }
-        }
+    // first collect all comparisons needed
+    let comparisons = assert_eq_values(compiler, lhs_info, rhs_info, lhs_type, span);
 
-        // a const and a var
-        (ConstOrCell::Const(cst), ConstOrCell::Cell(cvar))
-        | (ConstOrCell::Cell(cvar), ConstOrCell::Const(cst)) => {
-            compiler.backend.assert_eq_const(cvar, *cst, span)
-        }
-        (ConstOrCell::Cell(lhs), ConstOrCell::Cell(rhs)) => {
-            compiler.backend.assert_eq_var(lhs, rhs, span)
+    // then add all the constraints
+    for comparison in comparisons {
+        match comparison {
+            Comparison::Vars(lhs, rhs) => {
+                compiler.backend.assert_eq_var(&lhs, &rhs, span);
+            }
+            Comparison::VarConst(var, constant) => {
+                compiler.backend.assert_eq_const(&var, constant, span);
+            }
+            Comparison::Constants(a, b) => {
+                if a != b {
+                    return Err(Error::new(
+                        "constraint-generation",
+                        ErrorKind::AssertionFailed,
+                        span,
+                    ));
+                }
+            }
         }
     }
 

src/tests/examples.rs

@@ -839,6 +839,25 @@ fn test_generic_array_nested(#[case] backend: BackendKind) -> miette::Result<()>
     Ok(())
 }
 
+#[rstest]
+#[case::kimchi_vesta(BackendKind::KimchiVesta(KimchiVesta::new(false)))]
+#[case::r1cs(BackendKind::R1csBls12_381(R1CS::new()))]
+fn test_generic_assert_eq(#[case] backend: BackendKind) -> miette::Result<()> {
+    let public_inputs = r#"{"public_arr": ["1", "2"], "public_input": "3"}"#;
+    let private_inputs = r#"{}"#;
+
+    test_file(
+        "generic_assert_eq",
+        public_inputs,
+        private_inputs,
+        vec![],
+        backend,
+        DEFAULT_OPTIONS,
+    )?;
+
+    Ok(())
+}
+
 #[rstest]
 #[case::kimchi_vesta(BackendKind::KimchiVesta(KimchiVesta::new(false)))]
 #[case::r1cs(BackendKind::R1csBls12_381(R1CS::new()))]