performance(sierra-to-casm): Small map optimization of vars map. (#8480)

Author: orizi

crates/cairo-lang-sierra-to-casm/src/annotations.rs

@@ -1,5 +1,3 @@
-use std::iter;
-
 use cairo_lang_casm::ap_change::{ApChangeError, ApplyApChange};
 use cairo_lang_sierra::edit_state::{put_results, take_args};
 use cairo_lang_sierra::ids::{ConcreteTypeId, FunctionId, VarId};
@@ -154,7 +152,7 @@ pub struct ProgramAnnotations {
 impl ProgramAnnotations {
     fn new(n_statements: usize, backwards_jump_indices: UnorderedHashSet<StatementIdx>) -> Self {
         ProgramAnnotations {
-            per_statement_annotations: iter::repeat_with(|| None).take(n_statements).collect(),
+            per_statement_annotations: vec![None; n_statements],
             backwards_jump_indices,
         }
     }

crates/cairo-lang-sierra-to-casm/src/references.rs

@@ -5,7 +5,7 @@ use cairo_lang_sierra::ids::{ConcreteTypeId, VarId};
 use cairo_lang_sierra::program::{Function, StatementIdx};
 use cairo_lang_sierra_type_size::TypeSizeMap;
 use cairo_lang_utils::casts::IntoOrPanic;
-use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
+use cairo_lang_utils::small_ordered_map::SmallOrderedMap;
 use cairo_lang_utils::write_comma_separated;
 use thiserror::Error;
 
@@ -23,7 +23,7 @@ pub enum ReferencesError {
     UnknownType(ConcreteTypeId),
 }
 
-pub type StatementRefs = OrderedHashMap<VarId, ReferenceValue>;
+pub type StatementRefs = SmallOrderedMap<VarId, ReferenceValue>;
 
 /// A Sierra reference to a value.
 /// Corresponds to an argument or return value of a Sierra statement.

crates/cairo-lang-sierra/src/edit_state.rs

@@ -1,4 +1,4 @@
-use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
+use cairo_lang_utils::small_ordered_map::SmallOrderedMap;
 use thiserror::Error;
 
 use crate::ids::VarId;
@@ -25,12 +25,12 @@ impl EditStateError {
 
 /// Given a map with var ids as keys, extracts out the given ids, failing if some id is missing.
 pub fn take_args<'a, V: 'a>(
-    mut state: OrderedHashMap<VarId, V>,
+    mut state: SmallOrderedMap<VarId, V>,
     ids: impl Iterator<Item = &'a VarId>,
-) -> Result<(OrderedHashMap<VarId, V>, Vec<V>), EditStateError> {
+) -> Result<(SmallOrderedMap<VarId, V>, Vec<V>), EditStateError> {
     let mut vals = vec![];
     for id in ids {
-        match state.swap_remove(id) {
+        match state.remove(id) {
             None => {
                 return Err(EditStateError::MissingReference(id.clone()));
             }
@@ -44,9 +44,9 @@ pub fn take_args<'a, V: 'a>(
 
 /// Adds the given pairs to map with var ids as keys, failing if some variable is overridden.
 pub fn put_results<'a, V>(
-    mut state: OrderedHashMap<VarId, V>,
+    mut state: SmallOrderedMap<VarId, V>,
     results: impl Iterator<Item = (&'a VarId, V)>,
-) -> Result<OrderedHashMap<VarId, V>, EditStateError> {
+) -> Result<SmallOrderedMap<VarId, V>, EditStateError> {
     for (id, v) in results {
         if state.insert(id.clone(), v).is_some() {
             return Err(EditStateError::VariableOverride(id.clone()));

crates/cairo-lang-sierra/src/edit_state_test.rs

@@ -1,10 +1,10 @@
 use cairo_lang_test_utils::test;
-use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
+use cairo_lang_utils::small_ordered_map::SmallOrderedMap;
 
 use crate::edit_state::{EditStateError, put_results, take_args};
 use crate::ids::VarId;
 
-pub type State = OrderedHashMap<VarId, i64>;
+pub type State = SmallOrderedMap<VarId, i64>;
 
 #[test]
 fn empty() {
@@ -15,19 +15,19 @@ fn empty() {
 #[test]
 fn basic_mapping() {
     assert_eq!(
-        take_args(State::from([("arg".into(), 0)]), vec![&"arg".into()].into_iter(),),
+        take_args(State::from_iter([("arg".into(), 0)]), vec![&"arg".into()].into_iter(),),
         Ok((State::default(), vec![0]))
     );
     assert_eq!(
         put_results(State::default(), vec![(&"res".into(), 1)].into_iter(),),
-        Ok(State::from([("res".into(), 1)]))
+        Ok(State::from_iter([("res".into(), 1)]))
     );
     assert_eq!(
         take_args(State::default(), vec![&"arg".into()].into_iter(),),
         Err(EditStateError::MissingReference("arg".into()))
     );
     assert_eq!(
-        put_results(State::from([("res".into(), 1)]), vec![(&"res".into(), 1)].into_iter(),),
+        put_results(State::from_iter([("res".into(), 1)]), vec![(&"res".into(), 1)].into_iter(),),
         Err(EditStateError::VariableOverride("res".into()))
     );
 }

crates/cairo-lang-sierra/src/simulation/mod.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
+use cairo_lang_utils::small_ordered_map::SmallOrderedMap;
 use itertools::izip;
 use thiserror::Error;
 
@@ -83,7 +83,7 @@ impl SimulationContext<'_> {
                 actual: inputs.len(),
             });
         }
-        let mut state = OrderedHashMap::<VarId, CoreValue>::from_iter(
+        let mut state = SmallOrderedMap::<VarId, CoreValue>::from_iter(
             izip!(func.params.iter(), inputs).map(|(param, input)| (param.id.clone(), input)),
         );
         loop {