From 2301b16c1d810588a36585f928891116c4111a37 Mon Sep 17 00:00:00 2001
From: Nenad <nenad.misic@nethermind.io>
Date: Mon, 24 Jun 2024 14:08:27 +0200
Subject: [PATCH] Fix funcs 'new' and 'difference'

---
 exercises/practice/custom-set/src/lib.cairo   | 37 +++++--------------
 exercises/practice/custom-set/src/tests.cairo |  1 +
 2 files changed, 10 insertions(+), 28 deletions(-)

diff --git a/exercises/practice/custom-set/src/lib.cairo b/exercises/practice/custom-set/src/lib.cairo
index b5a11713..aea1cd0f 100644
--- a/exercises/practice/custom-set/src/lib.cairo
+++ b/exercises/practice/custom-set/src/lib.cairo
@@ -4,7 +4,7 @@ use core::box::BoxTrait;
 
 #[derive(Drop, Debug)]
 pub struct CustomSet<T> {
-    collection: Array<T>,
+    pub collection: Array<T>,
 }
 
 pub impl CustomSetEq<
@@ -28,7 +28,7 @@ pub impl CustomSetImpl<
 > of CustomSetTrait<T> {
     fn new(inputs: @Array<T>) -> CustomSet<T> {
         let mut set = CustomSet::<T> { collection: array![], };
-        let mut i = inputs.len();
+        let mut i = 0;
         while let Option::Some(val) = inputs
             .get(i) {
                 let unboxed = val.unbox();
@@ -147,39 +147,20 @@ pub impl CustomSetImpl<
 
     #[must_use]
     fn difference(self: @CustomSet<T>, other: @CustomSet<T>) -> CustomSet<T> {
-        let mut set = CustomSetImpl::<T>::new(@array![]);
-
-        let mut smaller_set = self;
-        let mut larger_set = other;
-        if smaller_set.collection.len() > larger_set.collection.len() {
-            smaller_set = other;
-            larger_set = self;
-        };
+        let mut diff = CustomSetImpl::<T>::new(@array![]);
 
         let mut i = 0;
-        while i < smaller_set
-            .collection
-            .len() {
-                if !larger_set.contains(smaller_set.collection[i]) {
-                    set.add(smaller_set.collection[i].clone());
-                }
-                if !smaller_set.contains(larger_set.collection[i]) {
-                    set.add(larger_set.collection[i].clone());
-                }
-                i += 1;
-            };
-
-        // iterate the remaining items in the larger set
-        while i < larger_set
+        while let Option::Some(val) = self
             .collection
-            .len() {
-                if !smaller_set.contains(larger_set.collection[i]) {
-                    set.add(larger_set.collection[i].clone());
+            .get(i) {
+                let unboxed = val.unbox();
+                if !other.contains(unboxed) {
+                    diff.add(unboxed.clone());
                 }
                 i += 1;
             };
 
-        set
+        diff
     }
 }
 
diff --git a/exercises/practice/custom-set/src/tests.cairo b/exercises/practice/custom-set/src/tests.cairo
index 7e461ff6..2a188796 100644
--- a/exercises/practice/custom-set/src/tests.cairo
+++ b/exercises/practice/custom-set/src/tests.cairo
@@ -294,3 +294,4 @@ fn union_of_non_empty_sets_contains_all_unique_elements() {
     let expected = U32Set::new(@array![3, 2, 1]);
     assert_eq!(set_1.union(@set_2), expected);
 }
+