From f0c68f4f885775a48ab3204284731585da2f11e1 Mon Sep 17 00:00:00 2001
From: TomerStarkware <tomer@starkware.co>
Date: Sun, 22 Dec 2024 12:07:03 +0200
Subject: [PATCH] added calculation for closure type size

---
 .../test/language_features/closure_test.cairo | 26 ++++++++++++++++---
 crates/cairo-lang-lowering/src/db.rs          |  4 ++-
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/corelib/src/test/language_features/closure_test.cairo b/corelib/src/test/language_features/closure_test.cairo
index 96ef3ff94ad..5f2fa9ecef6 100644
--- a/corelib/src/test/language_features/closure_test.cairo
+++ b/corelib/src/test/language_features/closure_test.cairo
@@ -57,7 +57,9 @@ fn option_map_test() {
     assert_eq!(option_map(Option::Some(2), |x| Option::Some(x)), Option::Some(Option::Some(2)));
 }
 
-fn array_map<T, F, impl Fn: core::ops::Fn<F, (T,)>, +Drop<T>, +Drop<F>, +Drop<Fn::Output>>(
+fn fix_sized_array_map<
+    T, F, impl Fn: core::ops::Fn<F, (T,)>, +Drop<T>, +Drop<F>, +Drop<Fn::Output>,
+>(
     arr: [T; 2], f: F,
 ) -> [core::ops::Fn::<F, (T,)>::Output; 2] {
     let [a, b] = arr;
@@ -65,7 +67,25 @@ fn array_map<T, F, impl Fn: core::ops::Fn<F, (T,)>, +Drop<T>, +Drop<F>, +Drop<Fn
 }
 
 #[test]
-fn array_map_test() {
-    assert_eq!(array_map([2, 3], |x| x + 3), [5, 6]);
+fn fix_sized_array_map_test() {
+    assert_eq!(fix_sized_array_map([2, 3], |x| x + 3), [5, 6]);
 }
 
+#[generate_trait]
+impl ArrayExt of ArrayExtTrait {
+    fn map<T, +Drop<T>, F, +Drop<F>, impl func: core::ops::Fn<F, (T,)>, +Drop<func::Output>>(
+        self: Array<T>, f: F,
+    ) -> Array<func::Output> {
+        let mut output: Array<func::Output> = array![];
+        for elem in self {
+            output.append(f(elem));
+        };
+        output
+    }
+}
+#[test]
+fn array_map_test() {
+    let arr = array![1, 2, 3];
+    let result = arr.map(|x| x + 1);
+    assert_eq!(result, array![2, 3, 4]);
+}
diff --git a/crates/cairo-lang-lowering/src/db.rs b/crates/cairo-lang-lowering/src/db.rs
index 3144b33c019..d62fd71bdec 100644
--- a/crates/cairo-lang-lowering/src/db.rs
+++ b/crates/cairo-lang-lowering/src/db.rs
@@ -798,7 +798,9 @@ fn type_size(db: &dyn LoweringGroup, ty: TypeId) -> usize {
                     .to_usize()
                     .unwrap()
         }
-        TypeLongId::Closure(_) => unimplemented!(),
+        TypeLongId::Closure(closure_ty) => {
+            closure_ty.captured_types.iter().map(|ty| db.type_size(*ty)).sum()
+        }
         TypeLongId::Coupon(_) => 0,
         TypeLongId::GenericParameter(_)
         | TypeLongId::Var(_)