From e1338393d5f78c3c672360411119bd9c0224f79a Mon Sep 17 00:00:00 2001
From: TomerStarkware <144585788+TomerStarkware@users.noreply.github.com>
Date: Sun, 29 Dec 2024 17:22:33 +0200
Subject: [PATCH] fix rewrites of ordered_hash_map so the key is also rewritten
 (#6945)

---
 .../cairo-lang-semantic/src/expr/inference.rs |  7 ++---
 .../src/expr/inference/canonic.rs             |  5 +++-
 .../cairo-lang-semantic/src/substitution.rs   | 27 ++++++++++++++++---
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/crates/cairo-lang-semantic/src/expr/inference.rs b/crates/cairo-lang-semantic/src/expr/inference.rs
index af1d0b1bf9b..17d0df63c06 100644
--- a/crates/cairo-lang-semantic/src/expr/inference.rs
+++ b/crates/cairo-lang-semantic/src/expr/inference.rs
@@ -41,7 +41,10 @@ use crate::items::imp::{
     GeneratedImplId, GeneratedImplItems, GeneratedImplLongId, ImplId, ImplImplId, ImplLongId,
     ImplLookupContext, UninferredGeneratedImplId, UninferredGeneratedImplLongId, UninferredImpl,
 };
-use crate::items::trt::{ConcreteTraitGenericFunctionId, ConcreteTraitGenericFunctionLongId};
+use crate::items::trt::{
+    ConcreteTraitGenericFunctionId, ConcreteTraitGenericFunctionLongId, ConcreteTraitTypeId,
+    ConcreteTraitTypeLongId,
+};
 use crate::substitution::{HasDb, RewriteResult, SemanticRewriter, SubstitutionRewriter};
 use crate::types::{
     ClosureTypeLongId, ConcreteEnumLongId, ConcreteExternTypeLongId, ConcreteStructLongId,
@@ -933,14 +936,12 @@ impl<'db> Inference<'db> {
             }
             _ => {}
         };
-
         let (canonical_trait, canonicalizer) = CanonicalTrait::canonicalize(
             self.db,
             self.inference_id,
             concrete_trait_id,
             impl_var_trait_item_mappings,
         );
-
         // impl_type_bounds order is deterimend by the generic params of the function and therefore
         // is consistent.
         let solution_set = match self.db.canonic_trait_solutions(
diff --git a/crates/cairo-lang-semantic/src/expr/inference/canonic.rs b/crates/cairo-lang-semantic/src/expr/inference/canonic.rs
index aa8346d813b..34f374c66e9 100644
--- a/crates/cairo-lang-semantic/src/expr/inference/canonic.rs
+++ b/crates/cairo-lang-semantic/src/expr/inference/canonic.rs
@@ -23,7 +23,10 @@ use crate::items::imp::{
     GeneratedImplId, GeneratedImplItems, GeneratedImplLongId, ImplId, ImplImplId, ImplLongId,
     UninferredGeneratedImplId, UninferredGeneratedImplLongId, UninferredImpl,
 };
-use crate::items::trt::{ConcreteTraitGenericFunctionId, ConcreteTraitGenericFunctionLongId};
+use crate::items::trt::{
+    ConcreteTraitGenericFunctionId, ConcreteTraitGenericFunctionLongId, ConcreteTraitTypeId,
+    ConcreteTraitTypeLongId,
+};
 use crate::substitution::{HasDb, RewriteResult, SemanticObject, SemanticRewriter};
 use crate::types::{
     ClosureTypeLongId, ConcreteEnumLongId, ConcreteExternTypeLongId, ConcreteStructLongId,
diff --git a/crates/cairo-lang-semantic/src/substitution.rs b/crates/cairo-lang-semantic/src/substitution.rs
index 0929c303953..e42b1d488aa 100644
--- a/crates/cairo-lang-semantic/src/substitution.rs
+++ b/crates/cairo-lang-semantic/src/substitution.rs
@@ -1,4 +1,5 @@
 use std::collections::VecDeque;
+use std::hash::Hash;
 use std::ops::{Deref, DerefMut};
 
 use cairo_lang_defs::ids::{
@@ -28,7 +29,10 @@ use crate::items::imp::{
     GeneratedImplId, GeneratedImplItems, GeneratedImplLongId, ImplId, ImplImplId, ImplLongId,
     UninferredGeneratedImplId, UninferredGeneratedImplLongId, UninferredImpl,
 };
-use crate::items::trt::{ConcreteTraitGenericFunctionId, ConcreteTraitGenericFunctionLongId};
+use crate::items::trt::{
+    ConcreteTraitGenericFunctionId, ConcreteTraitGenericFunctionLongId, ConcreteTraitTypeId,
+    ConcreteTraitTypeLongId,
+};
 use crate::types::{
     ClosureTypeLongId, ConcreteEnumLongId, ConcreteExternTypeLongId, ConcreteStructLongId,
     ImplTypeId,
@@ -198,19 +202,32 @@ impl<T, E, TRewriter: SemanticRewriter<T, E>> SemanticRewriter<Box<T>, E> for TR
     }
 }
 
-impl<T: Clone, V: Clone, E, TRewriter: SemanticRewriter<V, E>>
+impl<T: Clone + Hash + Eq, V: Clone, E, TRewriter: SemanticRewriter<V, E> + SemanticRewriter<T, E>>
     SemanticRewriter<OrderedHashMap<T, V>, E> for TRewriter
 {
     fn internal_rewrite(&mut self, value: &mut OrderedHashMap<T, V>) -> Result<RewriteResult, E> {
         let mut result = RewriteResult::NoChange;
-        for el in value.iter_mut() {
-            match self.internal_rewrite(el.1)? {
+        let mut changed_key = Vec::new();
+        for (k, v) in value.iter_mut() {
+            let mut temp_key = k.clone();
+            match self.internal_rewrite(&mut temp_key)? {
+                RewriteResult::Modified => {
+                    changed_key.push((k.clone(), temp_key));
+                    result = RewriteResult::Modified;
+                }
+                RewriteResult::NoChange => {}
+            }
+            match self.internal_rewrite(v)? {
                 RewriteResult::Modified => {
                     result = RewriteResult::Modified;
                 }
                 RewriteResult::NoChange => {}
             }
         }
+        for (old_key, new_key) in changed_key {
+            let v = value.swap_remove(&old_key).unwrap();
+            value.insert(new_key, v);
+        }
 
         Ok(result)
     }
@@ -344,6 +361,8 @@ macro_rules! add_basic_rewrites {
         $crate::prune_single!(__regular_helper, ConcreteExternTypeLongId, $($exclude)*);
         $crate::prune_single!(__regular_helper, ConcreteTraitId, $($exclude)*);
         $crate::prune_single!(__regular_helper, ConcreteTraitLongId, $($exclude)*);
+        $crate::prune_single!(__regular_helper, ConcreteTraitTypeId, $($exclude)*);
+        $crate::prune_single!(__regular_helper, ConcreteTraitTypeLongId, $($exclude)*);
         $crate::prune_single!(__regular_helper, ConcreteImplId, $($exclude)*);
         $crate::prune_single!(__regular_helper, ConcreteImplLongId, $($exclude)*);
         $crate::prune_single!(__regular_helper, ConcreteTraitGenericFunctionLongId, $($exclude)*);