From 77fba5677f9d1466d9d08c2eddc57149f9010db4 Mon Sep 17 00:00:00 2001
From: kevaundray <kevtheappdev@gmail.com>
Date: Thu, 8 Jun 2023 22:06:25 +0100
Subject: [PATCH] feat: Make for-loop range be a polymorphic integer instead of
 just Field in unconstrained functions (#1583)

* make ranges be polymorphic integers

* make behavior consistent

* remove closure

* change index_type

* Update crates/noirc_frontend/src/hir/type_check/expr.rs

---------

Co-authored-by: jfecher <jake@aztecprotocol.com>
---
 .../src/ssa_refactor/ssa_gen/mod.rs           |  3 +-
 .../noirc_frontend/src/hir/type_check/expr.rs | 42 +++++++++++--------
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs
index 386648c87f0..e6edf3ec930 100644
--- a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs
+++ b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/mod.rs
@@ -236,7 +236,8 @@ impl<'a> FunctionContext<'a> {
         let loop_end = self.builder.insert_block();
 
         // this is the 'i' in `for i in start .. end { block }`
-        let loop_index = self.builder.add_block_parameter(loop_entry, Type::field());
+        let index_type = Self::convert_non_tuple_type(&for_expr.index_type);
+        let loop_index = self.builder.add_block_parameter(loop_entry, index_type);
 
         let start_index = self.codegen_non_tuple_expression(&for_expr.start_range);
         let end_index = self.codegen_non_tuple_expression(&for_expr.end_range);
diff --git a/crates/noirc_frontend/src/hir/type_check/expr.rs b/crates/noirc_frontend/src/hir/type_check/expr.rs
index 5a9077c2a7e..feb85f90a21 100644
--- a/crates/noirc_frontend/src/hir/type_check/expr.rs
+++ b/crates/noirc_frontend/src/hir/type_check/expr.rs
@@ -153,25 +153,33 @@ impl<'interner> TypeChecker<'interner> {
                 let start_span = self.interner.expr_span(&for_expr.start_range);
                 let end_span = self.interner.expr_span(&for_expr.end_range);
 
-                let mut unify_loop_range = |actual_type, span| {
-                    let expected_type = if self.is_unconstrained() {
-                        Type::FieldElement(CompTime::new(self.interner))
-                    } else {
-                        Type::comp_time(Some(span))
-                    };
+                // Check that start range and end range have the same types
+                let range_span = start_span.merge(end_span);
+                self.unify(&start_range_type, &end_range_type, range_span, || {
+                    TypeCheckError::TypeMismatch {
+                        expected_typ: start_range_type.to_string(),
+                        expr_typ: end_range_type.to_string(),
+                        expr_span: range_span,
+                    }
+                });
 
-                    self.unify(actual_type, &expected_type, span, || {
-                        TypeCheckError::TypeCannotBeUsed {
-                            typ: start_range_type.clone(),
-                            place: "for loop",
-                            span,
-                        }
-                        .add_context("The range of a loop must be known at compile-time")
-                    });
+                let expected_comptime = if self.is_unconstrained() {
+                    CompTime::new(self.interner)
+                } else {
+                    CompTime::Yes(Some(range_span))
                 };
-
-                unify_loop_range(&start_range_type, start_span);
-                unify_loop_range(&end_range_type, end_span);
+                let fresh_id = self.interner.next_type_variable_id();
+                let type_variable = Shared::new(TypeBinding::Unbound(fresh_id));
+                let expected_type = Type::PolymorphicInteger(expected_comptime, type_variable);
+
+                self.unify(&start_range_type, &expected_type, range_span, || {
+                    TypeCheckError::TypeCannotBeUsed {
+                        typ: start_range_type.clone(),
+                        place: "for loop",
+                        span: range_span,
+                    }
+                    .add_context("The range of a loop must be known at compile-time")
+                });
 
                 self.interner.push_definition_type(for_expr.identifier.id, start_range_type);