From ac4d054525315ab55122101b4876c115bac99a17 Mon Sep 17 00:00:00 2001
From: Lior Goldberg <lior@starkware.co>
Date: Mon, 26 Aug 2024 17:35:18 +0300
Subject: [PATCH] Add Range type to Sierra.

commit-id:0a594c95
---
 crates/cairo-lang-sierra-type-size/src/lib.rs |  3 +-
 .../cairo-lang-sierra/src/extensions/core.rs  |  2 +
 .../src/extensions/modules/mod.rs             |  1 +
 .../src/extensions/modules/range.rs           | 64 +++++++++++++++++++
 4 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 crates/cairo-lang-sierra/src/extensions/modules/range.rs

diff --git a/crates/cairo-lang-sierra-type-size/src/lib.rs b/crates/cairo-lang-sierra-type-size/src/lib.rs
index ee596f110a9..8477253487e 100644
--- a/crates/cairo-lang-sierra-type-size/src/lib.rs
+++ b/crates/cairo-lang-sierra-type-size/src/lib.rs
@@ -53,7 +53,8 @@ pub fn get_type_size_map(
             CoreTypeConcrete::Array(_)
             | CoreTypeConcrete::Span(_)
             | CoreTypeConcrete::EcPoint(_)
-            | CoreTypeConcrete::SquashedFelt252Dict(_) => Some(2),
+            | CoreTypeConcrete::SquashedFelt252Dict(_)
+            | CoreTypeConcrete::Range(_) => Some(2),
             CoreTypeConcrete::NonZero(wrapped_ty)
             | CoreTypeConcrete::Snapshot(wrapped_ty)
             | CoreTypeConcrete::Uninitialized(wrapped_ty) => {
diff --git a/crates/cairo-lang-sierra/src/extensions/core.rs b/crates/cairo-lang-sierra/src/extensions/core.rs
index bda63468f33..4f47e81cfa7 100644
--- a/crates/cairo-lang-sierra/src/extensions/core.rs
+++ b/crates/cairo-lang-sierra/src/extensions/core.rs
@@ -41,6 +41,7 @@ use super::modules::unconditional_jump::UnconditionalJumpLibfunc;
 use super::nullable::{NullableLibfunc, NullableType};
 use super::pedersen::{PedersenLibfunc, PedersenType};
 use super::poseidon::{PoseidonLibfunc, PoseidonType};
+use super::range::RangeType;
 use super::range_check::{RangeCheck96Type, RangeCheckType};
 use super::segment_arena::SegmentArenaType;
 use super::snapshot::{SnapshotTakeLibfunc, SnapshotType};
@@ -94,6 +95,7 @@ define_type_hierarchy! {
         Snapshot(SnapshotType),
         Bytes31(Bytes31Type),
         BoundedInt(BoundedIntType),
+        Range(RangeType),
     }, CoreTypeConcrete
 }
 
diff --git a/crates/cairo-lang-sierra/src/extensions/modules/mod.rs b/crates/cairo-lang-sierra/src/extensions/modules/mod.rs
index 03b86277960..0afa8378295 100644
--- a/crates/cairo-lang-sierra/src/extensions/modules/mod.rs
+++ b/crates/cairo-lang-sierra/src/extensions/modules/mod.rs
@@ -35,6 +35,7 @@ pub mod non_zero;
 pub mod nullable;
 pub mod pedersen;
 pub mod poseidon;
+pub mod range;
 pub mod range_check;
 pub mod segment_arena;
 pub mod snapshot;
diff --git a/crates/cairo-lang-sierra/src/extensions/modules/range.rs b/crates/cairo-lang-sierra/src/extensions/modules/range.rs
new file mode 100644
index 00000000000..084a960369e
--- /dev/null
+++ b/crates/cairo-lang-sierra/src/extensions/modules/range.rs
@@ -0,0 +1,64 @@
+use super::bounded_int::BoundedIntType;
+use super::int::signed::{Sint16Type, Sint32Type, Sint64Type, Sint8Type};
+use super::int::signed128::Sint128Type;
+use super::int::unsigned::{Uint16Type, Uint32Type, Uint64Type, Uint8Type};
+use super::int::unsigned128::Uint128Type;
+use crate::extensions::type_specialization_context::TypeSpecializationContext;
+use crate::extensions::types::{
+    GenericTypeArgGenericType, GenericTypeArgGenericTypeWrapper, TypeInfo,
+};
+use crate::extensions::{NamedType, SpecializationError};
+use crate::ids::GenericTypeId;
+use crate::program::GenericArg;
+
+fn check_inner_type(ty_info: &TypeInfo) -> Result<(), SpecializationError> {
+    // Note: the implementation assumes the following types are of size 1.
+    match (&ty_info.long_id.generic_id, &ty_info.long_id.generic_args[..]) {
+        (id, []) if *id == Uint8Type::id() => (),
+        (id, []) if *id == Uint16Type::id() => (),
+        (id, []) if *id == Uint32Type::id() => (),
+        (id, []) if *id == Uint64Type::id() => (),
+        (id, []) if *id == Uint128Type::id() => (),
+        (id, []) if *id == Sint8Type::id() => (),
+        (id, []) if *id == Sint16Type::id() => (),
+        (id, []) if *id == Sint32Type::id() => (),
+        (id, []) if *id == Sint64Type::id() => (),
+        (id, []) if *id == Sint128Type::id() => (),
+        (id, [GenericArg::Value(_), GenericArg::Value(_)]) if *id == BoundedIntType::id() => (),
+        _ => return Err(SpecializationError::UnsupportedGenericArg),
+    };
+    Ok(())
+}
+
+/// Type for `Range(x, y)` where `x <= y`.
+#[derive(Default)]
+pub struct RangeTypeWrapped {}
+impl GenericTypeArgGenericType for RangeTypeWrapped {
+    const ID: GenericTypeId = GenericTypeId::new_inline("Range");
+
+    fn calc_info(
+        &self,
+        _context: &dyn TypeSpecializationContext,
+        long_id: crate::program::ConcreteTypeLongId,
+        wrapped_info: TypeInfo,
+    ) -> Result<TypeInfo, SpecializationError> {
+        check_inner_type(&wrapped_info)?;
+
+        // The following assert is a sanity check. It should follow from the fact that
+        // `check_inner_type` passed.
+        assert!(
+            wrapped_info.storable
+                && wrapped_info.duplicatable
+                && wrapped_info.droppable
+                && !wrapped_info.zero_sized
+        );
+        Ok(TypeInfo {
+            long_id,
+            duplicatable: true,
+            droppable: true,
+            storable: true,
+            zero_sized: false,
+        })
+    }
+}
+pub type RangeType = GenericTypeArgGenericTypeWrapper<RangeTypeWrapped>;