performance(sierra-calcs): Made topological order not require vec alloc.

SIERRA_UPDATE_PATCH_CHANGE_TAG=Just performance gain - no interface effect.

Author: orizi

crates/cairo-lang-sierra-ap-change/src/compute.rs

@@ -191,25 +191,18 @@ impl<'a, TokenUsages: Fn(StatementIdx, CostTokenType) -> usize>
             (0..self.program.statements.len()).map(StatementIdx),
             self.program.statements.len(),
             |idx| {
-                Ok(self.branches[idx.0]
-                    .iter()
-                    .flat_map(|(ap_change, target)| match ap_change {
-                        ApChange::Unknown => None,
-                        ApChange::FunctionCall(id) => {
-                            if self.function_ap_change.contains_key(id) {
-                                Some(*target)
-                            } else {
-                                None
-                            }
-                        }
-                        ApChange::Known(_)
-                        | ApChange::DisableApTracking
-                        | ApChange::FromMetadata
-                        | ApChange::AtLocalsFinalization(_)
-                        | ApChange::FinalizeLocals
-                        | ApChange::EnableApTracking => Some(*target),
-                    })
-                    .collect())
+                Ok(self.branches[idx.0].iter().flat_map(|(ap_change, target)| match ap_change {
+                    ApChange::Unknown => None,
+                    ApChange::FunctionCall(id) => {
+                        self.function_ap_change.contains_key(id).then_some(*target)
+                    }
+                    ApChange::Known(_)
+                    | ApChange::DisableApTracking
+                    | ApChange::FromMetadata
+                    | ApChange::AtLocalsFinalization(_)
+                    | ApChange::FinalizeLocals
+                    | ApChange::EnableApTracking => Some(*target),
+                }))
             },
             |_| unreachable!("Cycle isn't an error."),
         )

crates/cairo-lang-sierra-gas/src/compute_costs.rs

@@ -7,7 +7,6 @@ use cairo_lang_sierra::program::{BranchInfo, Invocation, Program, Statement, Sta
 use cairo_lang_utils::casts::IntoOrPanic;
 use cairo_lang_utils::iterators::zip_eq3;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
-use cairo_lang_utils::ordered_hash_set::OrderedHashSet;
 use cairo_lang_utils::unordered_hash_map::{Entry, UnorderedHashMap};
 use cairo_lang_utils::unordered_hash_set::UnorderedHashSet;
 use itertools::zip_eq;
@@ -180,12 +179,19 @@ fn get_branch_requirements_dependencies(
     idx: &StatementIdx,
     invocation: &Invocation,
     libfunc_cost: &[BranchCost],
-) -> OrderedHashSet<StatementIdx> {
-    let mut res: OrderedHashSet<StatementIdx> = Default::default();
+) -> Vec<StatementIdx> {
+    let mut res = vec![];
+    // Adds to the result if not already in it.
+    // Since rather small - more efficient using a Vec than a Map.
+    let mut add_to_res = |idx| {
+        if !res.contains(&idx) {
+            res.push(idx);
+        }
+    };
     for (branch_info, branch_cost) in zip_eq(&invocation.branches, libfunc_cost) {
         match branch_cost {
             BranchCost::FunctionCost { const_cost: _, function, sign: _ } => {
-                res.insert(function.entry_point);
+                add_to_res(function.entry_point);
             }
             BranchCost::WithdrawGas(WithdrawGasBranchInfo {
                 success: true,
@@ -197,9 +203,8 @@ fn get_branch_requirements_dependencies(
             }
             _ => {}
         }
-        res.insert(idx.next(&branch_info.target));
+        add_to_res(idx.next(&branch_info.target));
     }
-
     res
 }
 
@@ -454,13 +459,11 @@ impl<CostType: CostTypeTrait> CostContext<'_, CostType> {
                         // Return has no dependencies.
                         vec![]
                     }
-                    Statement::Invocation(invocation) => {
-                        let libfunc_cost = &self.branch_costs[current_idx.0];
-
-                        get_branch_requirements_dependencies(current_idx, invocation, libfunc_cost)
-                            .into_iter()
-                            .collect()
-                    }
+                    Statement::Invocation(invocation) => get_branch_requirements_dependencies(
+                        current_idx,
+                        invocation,
+                        &self.branch_costs[current_idx.0],
+                    ),
                 }
             },
         )?;
@@ -674,10 +677,13 @@ impl<CostType: CostTypeTrait> CostContext<'_, CostType> {
 /// Generates a topological ordering of the statements according to the given dependencies_callback.
 ///
 /// Each statement appears in the ordering after its dependencies.
-fn compute_reverse_topological_order(
+fn compute_reverse_topological_order<
+    Dependencies: IntoIterator<Item = StatementIdx>,
+    DependenciesCallback: Fn(&StatementIdx) -> Dependencies,
+>(
     n_statements: usize,
     detect_cycles: bool,
-    dependencies_callback: impl Fn(&StatementIdx) -> Vec<StatementIdx>,
+    dependencies_callback: DependenciesCallback,
 ) -> Result<Vec<StatementIdx>, CostError> {
     reverse_topological_ordering(
         detect_cycles,

crates/cairo-lang-sierra-gas/src/generate_equations.rs

@@ -110,14 +110,12 @@ fn get_reverse_topological_ordering(program: &Program) -> Result<Vec<StatementId
         program.statements.len(),
         |idx| {
             Ok(match program.get_statement(&idx).unwrap() {
-                Statement::Invocation(invocation) => invocation
-                    .branches
-                    .iter()
-                    .rev()
-                    .map(|branch| idx.next(&branch.target))
-                    .collect(),
-                Statement::Return(_) => vec![],
-            })
+                Statement::Invocation(invocation) => invocation.branches.as_slice(),
+                Statement::Return(_) => &[],
+            }
+            .iter()
+            .rev()
+            .map(move |branch| idx.next(&branch.target)))
         },
         |_| unreachable!("Cycles are not detected."),
     )

crates/cairo-lang-sierra/src/algorithm/topological_order.rs

@@ -21,11 +21,15 @@ enum TopologicalOrderStatus {
 /// `out_of_bounds_err` - a function that returns an error for a node is out of bounds.
 /// `cycle_err` - a function that returns an error for a node that is part of a cycle.
 /// Note: Will only work properly if the nodes are in the range [0, node_count).
-pub fn reverse_topological_ordering<E>(
+pub fn reverse_topological_ordering<
+    E,
+    Children: IntoIterator<Item = StatementIdx>,
+    GetChildren: Fn(StatementIdx) -> Result<Children, E>,
+>(
     detect_cycles: bool,
     roots: impl Iterator<Item = StatementIdx>,
     node_count: usize,
-    get_children: impl Fn(StatementIdx) -> Result<Vec<StatementIdx>, E>,
+    get_children: GetChildren,
     cycle_err: impl Fn(StatementIdx) -> E,
 ) -> Result<Vec<StatementIdx>, E> {
     let mut ordering = vec![];
@@ -45,12 +49,16 @@ pub fn reverse_topological_ordering<E>(
 
 /// Calculates the reverse topological ordering starting from `root`. For more info see
 /// `reverse_topological_ordering`.
-fn calculate_reverse_topological_ordering<E>(
+fn calculate_reverse_topological_ordering<
+    E,
+    Children: IntoIterator<Item = StatementIdx>,
+    GetChildren: Fn(StatementIdx) -> Result<Children, E>,
+>(
     detect_cycles: bool,
     ordering: &mut Vec<StatementIdx>,
     status: &mut [TopologicalOrderStatus],
     root: StatementIdx,
-    get_children: &impl Fn(StatementIdx) -> Result<Vec<StatementIdx>, E>,
+    get_children: &GetChildren,
     cycle_err: &impl Fn(StatementIdx) -> E,
 ) -> Result<(), E> {
     // A stack of statements to visit.