From b74e0b6dc6db49faeb40bacea524d4b542c020da Mon Sep 17 00:00:00 2001
From: schaeff <thibaut@schaeff.fr>
Date: Fri, 18 Jan 2019 18:56:56 +0100
Subject: [PATCH] add support for ifelse array access, clean up semantic
 treatment

---
 zokrates_core/src/flatten/mod.rs | 25 +++++++++++++++-
 zokrates_core/src/semantics.rs   | 49 +++++++++++++++++---------------
 2 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/zokrates_core/src/flatten/mod.rs b/zokrates_core/src/flatten/mod.rs
index 3a004ef6..525867fd 100644
--- a/zokrates_core/src/flatten/mod.rs
+++ b/zokrates_core/src/flatten/mod.rs
@@ -864,7 +864,30 @@ impl Flattener {
                         FieldElementArrayExpression::FunctionCall(..) => {
                             unimplemented!("please use intermediate variables for now")
                         }
-                        FieldElementArrayExpression::IfElse(..) => unimplemented!(""),
+                        FieldElementArrayExpression::IfElse(
+                            condition,
+                            consequence,
+                            alternative,
+                        ) => {
+                            // [if cond then [a, b] else [c, d]][1] == if cond then [a, b][1] else [c, d][1]
+                            self.flatten_field_expression(
+                                functions_flattened,
+                                arguments_flattened,
+                                statements_flattened,
+                                FieldElementExpression::IfElse(
+                                    condition,
+                                    box FieldElementExpression::Select(
+                                        consequence,
+                                        box FieldElementExpression::Number(n.clone()),
+                                    ),
+                                    box FieldElementExpression::Select(
+                                        alternative,
+                                        box FieldElementExpression::Number(n),
+                                    ),
+                                ),
+                            )
+                            .apply_recursive_substitution(&self.substitution)
+                        }
                     },
                     e => {
                         let size = array.size();
diff --git a/zokrates_core/src/semantics.rs b/zokrates_core/src/semantics.rs
index 2fa8e194..35f5131d 100644
--- a/zokrates_core/src/semantics.rs
+++ b/zokrates_core/src/semantics.rs
@@ -540,29 +540,32 @@ impl Checker {
                 let consequence_checked = self.check_expression(&consequence)?;
                 let alternative_checked = self.check_expression(&alternative)?;
 
-                match (condition_checked, consequence_checked, alternative_checked) {
-					(TypedExpression::Boolean(condition), TypedExpression::FieldElement(consequence), TypedExpression::FieldElement(alternative)) => {
-						Ok(FieldElementExpression::IfElse(box condition, box consequence, box alternative).into())
-					},
-                    (TypedExpression::Boolean(condition), TypedExpression::FieldElementArray(consequence), TypedExpression::FieldElementArray(alternative)) => {
-                        assert!(consequence.get_type() == alternative.get_type());
-                        Ok(FieldElementArrayExpression::IfElse(box condition, box consequence, box alternative).into())
-                    },
-					(condition, consequence, alternative) =>
-						Err(
-							Error {
-								message:
-									format!("if {{condition}} then {{consequence}} else {{alternative}} should have types {}, {}, {}, found {}, {}, {}",
-										Type::Boolean,
-										Type::FieldElement,
-										Type::FieldElement,
-										condition.get_type(),
-										consequence.get_type(),
-										alternative.get_type(),
-									)
-							}
-						)
-				}
+                match condition_checked {
+                    TypedExpression::Boolean(condition) => {
+                        let consequence_type = consequence_checked.get_type();
+                        let alternative_type = alternative_checked.get_type();
+                        match consequence_type == alternative_type {
+                            true => match (consequence_checked, alternative_checked) {
+                                (TypedExpression::FieldElement(consequence), TypedExpression::FieldElement(alternative)) => {
+                                    Ok(FieldElementExpression::IfElse(box condition, box consequence, box alternative).into())
+                                },
+                                (TypedExpression::FieldElementArray(consequence), TypedExpression::FieldElementArray(alternative)) => {
+                                    Ok(FieldElementArrayExpression::IfElse(box condition, box consequence, box alternative).into())
+                                },
+                                _ => unimplemented!()
+                            }
+                            false => Err(Error {
+                                message: format!("{{consequence}} and {{alternative}} in `if/else` expression should have the same type, found {}, {}", consequence_type, alternative_type)
+                            })
+                        }
+                    }
+                    c => Err(Error {
+                        message: format!(
+                            "{{condition}} after `if` should be a boolean, found {}",
+                            c.get_type()
+                        ),
+                    }),
+                }
             }
             &Expression::Number(ref n) => Ok(FieldElementExpression::Number(n.clone()).into()),
             &Expression::FunctionCall(ref fun_id, ref arguments) => {