Filter Mapper by source inference

crates/cairo-lang-semantic/src/expr/inference/canonic.rs

@@ -535,6 +535,9 @@ impl<'db> SemanticRewriter<TypeLongId<'db>, MapperError> for Mapper<'db, '_> {
         let TypeLongId::Var(var) = value else {
             return value.default_rewrite(self);
         };
+        if var.inference_id != self.mapping.source_inference_id {
+            return value.default_rewrite(self);
+        }
         let id = self
             .mapping
             .type_var_mapping
@@ -553,6 +556,9 @@ impl<'db> SemanticRewriter<ConstValue<'db>, MapperError> for Mapper<'db, '_> {
         let ConstValue::Var(var, ty) = value else {
             return value.default_rewrite(self);
         };
+        if var.inference_id != self.mapping.source_inference_id {
+            return value.default_rewrite(self);
+        }
         let id = self
             .mapping
             .const_var_mapping
@@ -582,6 +588,9 @@ impl<'db> SemanticRewriter<ImplLongId<'db>, MapperError> for Mapper<'db, '_> {
             return value.default_rewrite(self);
         };
         let var = var_id.long(self.get_db());
+        if var.inference_id != self.mapping.source_inference_id {
+            return value.default_rewrite(self);
+        }
         let id = self
             .mapping
             .impl_var_mapping
@@ -614,6 +623,9 @@ impl<'db> SemanticRewriter<NegativeImplLongId<'db>, MapperError> for Mapper<'db,
             return value.default_rewrite(self);
         };
         let var = var_id.long(self.get_db());
+        if var.inference_id != self.mapping.source_inference_id {
+            return value.default_rewrite(self);
+        }
         let id = self
             .mapping
             .negative_impl_var_mapping
@@ -627,3 +639,68 @@ impl<'db> SemanticRewriter<NegativeImplLongId<'db>, MapperError> for Mapper<'db,
         Ok(RewriteResult::Modified)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use cairo_lang_utils::Intern;
+    use salsa::Database;
+
+    use super::{Mapper, VarMapping};
+    use crate::expr::inference::{
+        ConstVar, InferenceId, LocalConstVarId, LocalTypeVarId, TypeVar,
+    };
+    use crate::items::constant::ConstValue;
+    use crate::test_utils::SemanticDatabaseForTesting;
+    use crate::types::TypeLongId;
+
+    #[test]
+    fn mapper_maps_type_var_from_source_inference_context() {
+        let db = SemanticDatabaseForTesting::default();
+        let db: &dyn Database = &db;
+        let mut mapping = VarMapping::new_to_canonic(InferenceId::NoContext);
+        mapping.type_var_mapping.insert(LocalTypeVarId(0), LocalTypeVarId(7));
+
+        let original = TypeLongId::Var(TypeVar {
+            inference_id: InferenceId::NoContext,
+            id: LocalTypeVarId(0),
+        });
+
+        let mapped = Mapper::map(db, original, &mapping).unwrap();
+        assert_eq!(
+            mapped,
+            TypeLongId::Var(TypeVar { inference_id: InferenceId::Canonical, id: LocalTypeVarId(7) })
+        );
+    }
+
+    #[test]
+    fn mapper_ignores_type_var_from_other_inference_context() {
+        let db = SemanticDatabaseForTesting::default();
+        let db: &dyn Database = &db;
+        let mut mapping = VarMapping::new_to_canonic(InferenceId::NoContext);
+        mapping.type_var_mapping.insert(LocalTypeVarId(0), LocalTypeVarId(7));
+
+        let original = TypeLongId::Var(TypeVar {
+            inference_id: InferenceId::Canonical,
+            id: LocalTypeVarId(0),
+        });
+
+        let mapped = Mapper::map(db, original.clone(), &mapping).unwrap();
+        assert_eq!(mapped, original);
+    }
+
+    #[test]
+    fn mapper_ignores_const_var_from_other_inference_context() {
+        let db = SemanticDatabaseForTesting::default();
+        let db: &dyn Database = &db;
+        let mut mapping = VarMapping::new_to_canonic(InferenceId::NoContext);
+        mapping.const_var_mapping.insert(LocalConstVarId(0), LocalConstVarId(3));
+        let ty = TypeLongId::Tuple(vec![]).intern(db);
+        let original = ConstValue::Var(
+            ConstVar { inference_id: InferenceId::Canonical, id: LocalConstVarId(0) },
+            ty,
+        );
+
+        let mapped = Mapper::map(db, original.clone(), &mapping).unwrap();
+        assert_eq!(mapped, original);
+    }
+}