add and fix tests, make size array arbitrary, optimize example

zokrates_cli/examples/dex/ringtrade_example.code

@@ -12,10 +12,10 @@ def findPairInTokenPairList(field[9] pairs, field sourceToken, field targetToken
 
 // add `volume` to element `token` in `volumeToken`
 def addVolumesForOrder(field[3] volumeToken, field token, field volume) -> (field[3]):
-  field txVolumeToken1 = volumeToken[0] + if token == 0 then volume else 0 fi
+  for field i in 0..3 do
-  field txVolumeToken2 = volumeToken[1] + if token == 1 then volume else 0 fi
+    volumeToken[i] = volumeToken[i] + if token == i then volume else 0 fi
-  field txVolumeToken3 = volumeToken[2] + if token == 2 then volume else 0 fi
+  endfor
-  return [txVolumeToken1, txVolumeToken2, txVolumeToken3]
+  return volumeToken
 
 def updateHighestTouchedOrder(field[9] highestTouchedOrder, field sourceToken, field targetToken, field limit, field volume) -> (field[9]):
 

zokrates_core/src/flatten/mod.rs

@@ -894,11 +894,13 @@ impl Flattener {
                                         }
                                     },
                                     box TypedExpression::FieldElement(ref e) => {
-                                        //let size = array.size();
-                                        let size = 9; // TODO change
                                         // we have array[e] with e an arbitrary expression
                                         // first we check that e is in 0..array.len(), so we check that sum(if e == i then 1 else 0) == 1
                                         // here depending on the size, we could use a proper range check based on bits
+                                        let size = match array.get_type() {
+                                            Type::FieldElementArray(n) => n,
+                                            _ => panic!("checker should generate array element based on non array")
+                                        };
                                         let range_check = (0..size)
                                             .map(|i| 
                                                 FieldElementExpression::IfElse(
@@ -1425,7 +1427,7 @@ mod tests {
         let arguments_flattened = vec![];
         let mut statements_flattened = vec![];
         let statement = TypedStatement::Definition(
-            Variable::field_element("a".to_string()),
+            TypedAssignee::Identifier(Variable::field_element("a")),
             TypedExpression::FieldElement(FieldElementExpression::FunctionCall("foo".to_string(), vec![]))
         );
 
@@ -1469,7 +1471,7 @@ mod tests {
             arguments: vec![Parameter { id: Variable::field_element("a"), private: true }],
             statements: vec![
                 TypedStatement::Definition(
-                    Variable::field_element("a".to_string()), 
+                    TypedAssignee::Identifier(Variable::field_element("a")),
                     FieldElementExpression::Add(
                         box FieldElementExpression::Identifier("a".to_string()), 
                         box FieldElementExpression::Number(FieldPrime::from(1))
@@ -1508,7 +1510,9 @@ mod tests {
             id: String::from("foo"),
             arguments: vec![],
             statements: vec![
-                TypedStatement::Definition(Variable::field_element("a"), FieldElementExpression::Number(FieldPrime::from(3)).into()),
+                TypedStatement::Definition(
+                    TypedAssignee::Identifier(Variable::field_element("a")),
+                    FieldElementExpression::Number(FieldPrime::from(3)).into()),
                 TypedStatement::Return(vec![
                     FieldElementExpression::Identifier(String::from("a")).into()
                     ]
@@ -1566,6 +1570,12 @@ mod tests {
 
     #[test]
     fn powers() {
+
+        // def main():
+        //     field a = 7
+        //     field b = a**4
+        //     return b
+
         // def main():
         //     _0 = 7
         //     _1 = (_0 * _0)
@@ -1577,8 +1587,11 @@ mod tests {
             id: String::from("main"),
             arguments: vec![],
             statements: vec![
-                TypedStatement::Definition(Variable::field_element("a"), FieldElementExpression::Number(FieldPrime::from(7)).into()),
+                TypedStatement::Definition(
-                TypedStatement::Definition(Variable::field_element("b"), FieldElementExpression::Pow(box FieldElementExpression::Identifier(String::from("a")), box FieldElementExpression::Number(FieldPrime::from(4))).into()),
+                    TypedAssignee::Identifier(Variable::field_element("a")),
+                    FieldElementExpression::Number(FieldPrime::from(7)).into()
+                ),
+                TypedStatement::Definition(TypedAssignee::Identifier(Variable::field_element("b")), FieldElementExpression::Pow(box FieldElementExpression::Identifier(String::from("a")), box FieldElementExpression::Number(FieldPrime::from(4))).into()),
                 TypedStatement::Return(vec![
                     FieldElementExpression::Identifier(String::from("b")).into()
                     ]
@@ -1665,7 +1678,7 @@ mod tests {
                 id: "main".to_string(),
                 arguments: vec![],
                 statements: vec![
-                    TypedStatement::Definition(Variable::field_element("a".to_string()), TypedExpression::FieldElement(FieldElementExpression::FunctionCall("foo".to_string(), vec![]))),
+                    TypedStatement::Definition(TypedAssignee::Identifier(Variable::field_element("a")), TypedExpression::FieldElement(FieldElementExpression::FunctionCall("foo".to_string(), vec![]))),
                     TypedStatement::MultipleDefinition(vec![Variable::field_element("b".to_string()), Variable::field_element("c".to_string())], TypedExpressionList::FunctionCall("foo".to_string(), vec![], vec![Type::FieldElement, Type::FieldElement])),
                     TypedStatement::Return(
                         vec![TypedExpression::FieldElement(FieldElementExpression::Number(FieldPrime::from(1)))]
@@ -1725,7 +1738,7 @@ mod tests {
         let mut functions_flattened = vec![];
         let arguments_flattened = vec![];
         let mut statements_flattened = vec![];
-        let statement = TypedStatement::Definition(Variable::field_array("foo", 3), FieldElementArrayExpression::Value(3, vec![
+        let statement = TypedStatement::Definition(TypedAssignee::Identifier(Variable::field_array("foo", 3)), FieldElementArrayExpression::Value(3, vec![
                 FieldElementExpression::Number(FieldPrime::from(1)),
                 FieldElementExpression::Number(FieldPrime::from(2)),
                 FieldElementExpression::Number(FieldPrime::from(3)),
@@ -1768,7 +1781,7 @@ mod tests {
         let arguments_flattened = vec![];
         let mut statements_flattened = vec![];
         let statement = TypedStatement::Definition(
-            Variable::field_array("foo", 3),
+            TypedAssignee::Identifier(Variable::field_array("foo", 3)),
             FieldElementArrayExpression::Value(3, vec![
                 FieldElementExpression::Number(FieldPrime::from(1)),
                 FieldElementExpression::Number(FieldPrime::from(2)),
@@ -1805,7 +1818,7 @@ mod tests {
         let arguments_flattened = vec![];
         let mut statements_flattened = vec![];
         let statement = TypedStatement::Definition(
-            Variable::field_array("foo", 3),
+            TypedAssignee::Identifier(Variable::field_array("foo", 3)),
             FieldElementArrayExpression::Value(3, vec![
                 FieldElementExpression::Number(FieldPrime::from(1)),
                 FieldElementExpression::Number(FieldPrime::from(2)),
@@ -1848,7 +1861,7 @@ mod tests {
         let arguments_flattened = vec![];
         let mut statements_flattened = vec![];
         let def = TypedStatement::Definition(
-            Variable::field_array("foo", 3),
+            TypedAssignee::Identifier(Variable::field_array("foo", 3)),
             FieldElementArrayExpression::Value(3, vec![
                 FieldElementExpression::Number(FieldPrime::from(1)),
                 FieldElementExpression::Number(FieldPrime::from(2)),
@@ -1857,7 +1870,7 @@ mod tests {
         );
 
         let sum = TypedStatement::Definition(
-            Variable::field_element("bar"),
+            TypedAssignee::Identifier(Variable::field_element("bar")),
             FieldElementExpression::Add(
                 box FieldElementExpression::Add(
                     box FieldElementExpression::Select(box FieldElementArrayExpression::Identifier(3, String::from("foo")), box FieldElementExpression::Number(FieldPrime::from(0))),

zokrates_core/src/semantics.rs

@@ -244,10 +244,16 @@ impl Checker {
 				let checked_expr = self.check_expression(&expr)?;
 				let expression_type = checked_expr.get_type();
 
-				// check that the variable is declared, has the right type and is well formed
+				// check that the assignee is declared and is well formed
-				let var = self.check_assignee(&assignee, &expression_type)?;
+				let var = self.check_assignee(&assignee)?;
 
-				Ok(TypedStatement::Definition(var, checked_expr))
+				let var_type = var.get_type();
+
+				// make sure the assignee has the same type as the rhs
+				match var_type == expression_type {
+					true => Ok(TypedStatement::Definition(var, checked_expr)),
+					false => Err( Error { message: format!("Expression of type {} cannot be assigned to {} of type {}", expression_type, var, var_type) }),
+				}
 			}
 			Statement::Condition(lhs, rhs) => {
 				let checked_lhs = self.check_expression(&lhs)?;
@@ -335,16 +341,12 @@ impl Checker {
 		}
 	}
 
-	fn check_assignee<T: Field>(&mut self, assignee: &Assignee<T>, t: &Type) -> Result<TypedAssignee<T>, Error> {
+	fn check_assignee<T: Field>(&mut self, assignee: &Assignee<T>) -> Result<TypedAssignee<T>, Error> {
 		// check that the assignee is declared
 		match assignee {
 			Assignee::Identifier(variable_name) => {
 				match self.get_scope(&variable_name) {
 					Some(var) => {
-						if *t != var.id.get_type() {
-							return Err( Error { message: format!("Expression of type {} cannot be assigned to {} of type {}", *t, var.id.id, var.id.get_type()) });
-						}
-
 						Ok(TypedAssignee::Identifier(var.id.clone()))
 					},
 					None => {
@@ -353,11 +355,8 @@ impl Checker {
 				}
 			},
 			Assignee::ArrayElement(box assignee, box index) => {
-				let expected_type = match t {
+
-					Type::FieldElement => Type::FieldElementArray(9), // TODO CHANGE
+				let checked_assignee = self.check_assignee(&assignee)?;
-					_ => panic!("cannot have brackets if its not a field element")
-				};
-				let checked_assignee = self.check_assignee(&assignee, &expected_type)?;
 				let checked_index = self.check_expression(&index)?;
 
 				Ok(TypedAssignee::ArrayElement(box checked_assignee, box checked_index))
@@ -932,7 +931,7 @@ mod tests {
 				Variable::field_element("a")
 			),
 			Statement::MultipleDefinition(
-				vec![String::from("a")],
+				vec![Assignee::Identifier(String::from("a"))],
 				Expression::FunctionCall("foo".to_string(), vec![]))
 		];
 
@@ -1008,7 +1007,7 @@ mod tests {
 				Variable::field_element("a")
 			),
 			Statement::MultipleDefinition(
-				vec![String::from("a")],
+				vec![Assignee::Identifier(String::from("a"))],
 				Expression::FunctionCall("foo".to_string(), vec![]))
 		];
 
@@ -1062,7 +1061,7 @@ mod tests {
 				Variable::field_element("b")
 			),
 			Statement::MultipleDefinition(
-				vec![String::from("a"), String::from("b")],
+				vec![Assignee::Identifier(String::from("a")), Assignee::Identifier(String::from("b"))],
 				Expression::FunctionCall("foo".to_string(), vec![
 					Expression::Identifier("x".to_string())
 				])
@@ -1161,7 +1160,7 @@ mod tests {
 				Variable::field_element("b")
 			),
 			Statement::MultipleDefinition(
-				vec![String::from("a"), String::from("b")],
+				vec![Assignee::Identifier(String::from("a")), Assignee::Identifier(String::from("b"))],
 				Expression::FunctionCall("foo".to_string(), vec![])
 			),
 			Statement::Return(
@@ -1369,4 +1368,46 @@ mod tests {
 			Err(Error { message: "Duplicate declaration for variable named a".to_string() })
 		);
 	}
+
+	mod assignee {
+		use super::*;
+
+		#[test]
+		fn identifier() {
+			// a = 42
+			let a = Assignee::Identifier::<FieldPrime>(String::from("a"));
+
+			let mut checker: Checker = Checker::new();
+			checker.check_statement::<FieldPrime>(
+				&Statement::Declaration(Variable::field_element("a")),
+				&vec![],
+			).unwrap();
+
+			assert_eq!(
+				checker.check_assignee(&a),
+				Ok(TypedAssignee::Identifier(Variable::field_element("a")))
+			);
+		}
+
+		#[test]
+		fn array_element() {
+			// field[33] a
+			// a[2] = 42
+			let a = Assignee::ArrayElement(box Assignee::Identifier(String::from("a")), box Expression::Number(FieldPrime::from(2)));
+
+			let mut checker: Checker = Checker::new();
+			checker.check_statement::<FieldPrime>(
+				&Statement::Declaration(Variable::field_array("a", 33)),
+				&vec![],
+			).unwrap();
+
+			assert_eq!(
+				checker.check_assignee(&a),
+				Ok(TypedAssignee::ArrayElement(
+					box TypedAssignee::Identifier(Variable::field_array("a",33)),
+					box FieldElementExpression::Number(FieldPrime::from(2)).into())
+				)
+			);
+		}
+	}
 }

zokrates_core/src/typed_absy/mod.rs

@@ -124,6 +124,21 @@ pub enum TypedAssignee<T: Field> {
     ArrayElement(Box<TypedAssignee<T>>, Box<TypedExpression<T>>)
 }
 
+impl<T: Field> Typed for TypedAssignee<T> {
+    fn get_type(&self) -> Type {
+        match *self {
+            TypedAssignee::Identifier(ref v) => v.get_type(),
+            TypedAssignee::ArrayElement(ref a, _) => {
+                let a_type = a.get_type();
+                match a_type {
+                    Type::FieldElementArray(_) => Type::FieldElement,
+                    _ => panic!("array element has to take array")
+                }
+            }
+        }
+    }
+}
+
 impl<T: Field> fmt::Debug for TypedAssignee<T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {