Worklist order improvements (#360)

src/cwe_checker_lib/src/analysis/backward_interprocedural_fixpoint/mod.rs

@@ -11,6 +11,7 @@
 //! The `Computation` object provides the necessary methods for the actual fixpoint computation.
 
 use super::fixpoint::Context as GeneralFPContext;
+use super::forward_interprocedural_fixpoint;
 use super::graph::*;
 use super::interprocedural_fixpoint_generic::*;
 use crate::intermediate_representation::*;
@@ -256,6 +257,42 @@ pub fn create_computation<'a, T: Context<'a>>(
     super::fixpoint::Computation::new(generalized_problem, default_value.map(NodeValue::Value))
 }
 
+/// Generate a new computation from the corresponding context and an optional default value for nodes.
+/// Uses a bottom up worklist order when computing the fixpoint.
+///
+/// The worklist order prefers callee nodes before caller nodes.
+pub fn create_computation_with_bottom_up_worklist_order<'a, T: Context<'a>>(
+    problem: T,
+    default_value: Option<T::Value>,
+) -> super::fixpoint::Computation<GeneralizedContext<'a, T>> {
+    let priority_sorted_nodes =
+        forward_interprocedural_fixpoint::create_bottom_up_worklist(problem.get_graph());
+    let generalized_problem = GeneralizedContext::new(problem);
+    super::fixpoint::Computation::from_node_priority_list(
+        generalized_problem,
+        default_value.map(NodeValue::Value),
+        priority_sorted_nodes,
+    )
+}
+
+/// Generate a new computation from the corresponding context and an optional default value for nodes.
+/// Uses a top down worklist order when computing the fixpoint.
+///
+/// The worklist order prefers caller nodes before callee nodes.
+pub fn create_computation_with_top_down_worklist_order<'a, T: Context<'a>>(
+    problem: T,
+    default_value: Option<T::Value>,
+) -> super::fixpoint::Computation<GeneralizedContext<'a, T>> {
+    let priority_sorted_nodes =
+        forward_interprocedural_fixpoint::create_top_down_worklist(problem.get_graph());
+    let generalized_problem = GeneralizedContext::new(problem);
+    super::fixpoint::Computation::from_node_priority_list(
+        generalized_problem,
+        default_value.map(NodeValue::Value),
+        priority_sorted_nodes,
+    )
+}
+
 #[cfg(test)]
 pub mod tests;
 

src/cwe_checker_lib/src/analysis/fixpoint.rs

@@ -113,7 +113,7 @@ impl<T: Context> Computation<T> {
     /// and the list of nodes of the graph ordered by the priority for the worklist algorithm.
     /// The worklist algorithm will try to stabilize the nodes with a higher index
     /// in the `priority_sorted_nodes` array before those with a lower index.
-    fn from_node_priority_list(
+    pub fn from_node_priority_list(
         fp_context: T,
         default_value: Option<T::NodeValue>,
         priority_sorted_nodes: Vec<NodeIndex>,
@@ -141,32 +141,6 @@ impl<T: Context> Computation<T> {
         }
     }
 
-    /// Create a new fixpoint computation from a fixpoint problem and an optional default value for all nodes.
-    ///
-    /// Computations created by this function use an alternate priority order for the fixpoint stabilization algorithm:
-    /// Nodes with 10 or more incoming edges will be stabilized last by the algorithm.
-    pub fn new_with_alternate_worklist_order(
-        fp_context: T,
-        default_value: Option<T::NodeValue>,
-    ) -> Self {
-        let graph = fp_context.get_graph();
-        let mut high_priority_nodes = Vec::new();
-        let mut priority_sorted_nodes = Vec::new();
-        for node in petgraph::algo::kosaraju_scc(&graph).into_iter().flatten() {
-            if graph
-                .neighbors_directed(node, petgraph::EdgeDirection::Incoming)
-                .count()
-                >= 10
-            {
-                priority_sorted_nodes.push(node);
-            } else {
-                high_priority_nodes.push(node)
-            }
-        }
-        priority_sorted_nodes.append(&mut high_priority_nodes);
-        Self::from_node_priority_list(fp_context, default_value, priority_sorted_nodes)
-    }
-
     /// Get the value of a node.
     pub fn get_node_value(&self, node: NodeIndex) -> Option<&T::NodeValue> {
         self.node_values.get(&node)
@@ -391,27 +365,5 @@ mod tests {
             computation.take_next_node_from_worklist(),
             Some(NodeIndex::new(20))
         );
-
-        let mut computation =
-            Computation::new_with_alternate_worklist_order(FPContext { graph }, Some(1));
-        assert!(computation.node_priority_list[19] < computation.node_priority_list[0]);
-        assert!(computation.node_priority_list[1] > computation.node_priority_list[20]);
-        // assert that the nodes have the correct priority ordering
-        assert_eq!(
-            computation.take_next_node_from_worklist(),
-            Some(NodeIndex::new(0))
-        );
-        for _i in 1..19 {
-            assert!(computation.take_next_node_from_worklist().unwrap().index() < 19);
-        }
-        assert_eq!(
-            computation.take_next_node_from_worklist(),
-            Some(NodeIndex::new(20))
-        );
-        // nodes with a lot of incoming edges get stabilized last in the alternate worklist order
-        assert_eq!(
-            computation.take_next_node_from_worklist(),
-            Some(NodeIndex::new(19))
-        );
     }
 }

src/cwe_checker_lib/src/analysis/forward_interprocedural_fixpoint.rs

@@ -14,6 +14,7 @@ use super::graph::*;
 use super::interprocedural_fixpoint_generic::*;
 use crate::intermediate_representation::*;
 use petgraph::graph::EdgeIndex;
+use petgraph::graph::NodeIndex;
 use std::marker::PhantomData;
 
 /// The context for an interprocedural fixpoint computation.
@@ -267,19 +268,174 @@ pub fn create_computation<'a, T: Context<'a>>(
     super::fixpoint::Computation::new(generalized_problem, default_value.map(NodeValue::Value))
 }
 
+/// Returns a node ordering with callee nodes behind caller nodes.
+pub fn create_bottom_up_worklist(graph: &Graph) -> Vec<NodeIndex> {
+    let mut graph = graph.clone();
+    graph.retain_edges(|frozen, edge| !matches!(frozen[edge], Edge::Call(..)));
+    petgraph::algo::kosaraju_scc(&graph)
+        .into_iter()
+        .flatten()
+        .collect()
+}
+
+/// Returns a node ordering with caller nodes behind callee nodes.
+pub fn create_top_down_worklist(graph: &Graph) -> Vec<NodeIndex> {
+    let mut graph = graph.clone();
+    graph.retain_edges(|frozen, edge| !matches!(frozen[edge], Edge::CrReturnStub));
+    petgraph::algo::kosaraju_scc(&graph)
+        .into_iter()
+        .flatten()
+        .collect()
+}
+
 /// Generate a new computation from the corresponding context and an optional default value for nodes.
-/// Uses the alternate worklist order when computing the fixpoint.
+/// Uses a bottom up worklist order when computing the fixpoint.
 ///
-/// The alternate worklist order moves nodes with 10 or more incoming edges to the end of the priority queue.
-/// This can improve the convergence speed for these nodes in the fixpoint algorithm.
-/// Use if you encounter convergence problems for nodes with a lot of incoming edges.
-pub fn create_computation_with_alternate_worklist_order<'a, T: Context<'a>>(
+/// The worklist order prefers callee nodes before caller nodes.
+pub fn create_computation_with_bottom_up_worklist_order<'a, T: Context<'a>>(
     problem: T,
     default_value: Option<T::Value>,
 ) -> super::fixpoint::Computation<GeneralizedContext<'a, T>> {
+    let priority_sorted_nodes: Vec<NodeIndex> = create_bottom_up_worklist(problem.get_graph());
     let generalized_problem = GeneralizedContext::new(problem);
-    super::fixpoint::Computation::new_with_alternate_worklist_order(
+    super::fixpoint::Computation::from_node_priority_list(
         generalized_problem,
         default_value.map(NodeValue::Value),
+        priority_sorted_nodes,
     )
 }
+
+/// Generate a new computation from the corresponding context and an optional default value for nodes.
+/// Uses a top down worklist order when computing the fixpoint.
+///
+/// The worklist order prefers caller nodes before callee nodes.
+pub fn create_computation_with_top_down_worklist_order<'a, T: Context<'a>>(
+    problem: T,
+    default_value: Option<T::Value>,
+) -> super::fixpoint::Computation<GeneralizedContext<'a, T>> {
+    let priority_sorted_nodes: Vec<NodeIndex> = create_top_down_worklist(problem.get_graph());
+    let generalized_problem = GeneralizedContext::new(problem);
+    super::fixpoint::Computation::from_node_priority_list(
+        generalized_problem,
+        default_value.map(NodeValue::Value),
+        priority_sorted_nodes,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+
+    use crate::{
+        analysis::{
+            expression_propagation::Context,
+            forward_interprocedural_fixpoint::{
+                create_computation_with_bottom_up_worklist_order,
+                create_computation_with_top_down_worklist_order,
+            },
+        },
+        intermediate_representation::*,
+    };
+    use std::collections::{BTreeMap, HashMap, HashSet};
+
+    fn new_block(name: &str) -> Term<Blk> {
+        Term {
+            tid: Tid::new(name),
+            term: Blk {
+                defs: vec![],
+                jmps: vec![],
+                indirect_jmp_targets: Vec::new(),
+            },
+        }
+    }
+
+    /// Creates a project with one caller function of two blocks and one callee function of one block.
+    fn mock_project() -> Project {
+        let mut callee_block = new_block("callee block");
+        callee_block.term.jmps.push(Term {
+            tid: Tid::new("ret"),
+            term: Jmp::Return(Expression::const_from_i32(42)),
+        });
+
+        let called_function = Term {
+            tid: Tid::new("called_function"),
+            term: Sub {
+                name: "called_function".to_string(),
+                blocks: vec![callee_block],
+                calling_convention: Some("_stdcall".to_string()),
+            },
+        };
+
+        let mut caller_block_2 = new_block("caller_block_2");
+
+        let mut caller_block_1 = new_block("caller_block_1");
+        caller_block_1.term.jmps.push(Term {
+            tid: Tid::new("call"),
+            term: Jmp::Call {
+                target: called_function.tid.clone(),
+                return_: Some(caller_block_2.tid.clone()),
+            },
+        });
+
+        caller_block_2.term.jmps.push(Term {
+            tid: Tid::new("jmp"),
+            term: Jmp::Branch(caller_block_1.tid.clone()),
+        });
+
+        let caller_function = Term {
+            tid: Tid::new("caller_function"),
+            term: Sub {
+                name: "caller_function".to_string(),
+                blocks: vec![caller_block_1, caller_block_2],
+                calling_convention: Some("_stdcall".to_string()),
+            },
+        };
+        let mut project = Project::mock_x64();
+        project.program.term.subs = BTreeMap::from([
+            (caller_function.tid.clone(), caller_function.clone()),
+            (called_function.tid.clone(), called_function.clone()),
+        ]);
+
+        project
+    }
+
+    #[test]
+    /// Checks if the nodes corresponding to the callee function are first in the worklist.
+    fn check_bottom_up_worklist() {
+        let project = mock_project();
+        let extern_subs = HashSet::new();
+
+        let graph = crate::analysis::graph::get_program_cfg(&project.program, extern_subs);
+        let context = Context::new(&graph);
+        let comp = create_computation_with_bottom_up_worklist_order(context, Some(HashMap::new()));
+        // The last two nodes should belong to the callee
+        for node in comp.get_worklist()[6..].iter() {
+            match graph[*node] {
+                crate::analysis::graph::Node::BlkStart(_, sub)
+                | crate::analysis::graph::Node::BlkEnd(_, sub) => {
+                    assert_eq!(sub.tid, Tid::new("called_function"))
+                }
+                _ => panic!(),
+            }
+        }
+    }
+
+    #[test]
+    fn check_top_down_worklist() {
+        let project = mock_project();
+        let extern_subs = HashSet::new();
+
+        let graph = crate::analysis::graph::get_program_cfg(&project.program, extern_subs);
+        let context = Context::new(&graph);
+        let comp = create_computation_with_top_down_worklist_order(context, Some(HashMap::new()));
+        // The first two nodes should belong to the callee
+        for node in comp.get_worklist()[..2].iter() {
+            match graph[*node] {
+                crate::analysis::graph::Node::BlkStart(_, sub)
+                | crate::analysis::graph::Node::BlkEnd(_, sub) => {
+                    assert_eq!(sub.tid, Tid::new("called_function"))
+                }
+                _ => panic!(),
+            }
+        }
+    }
+}

src/cwe_checker_lib/src/analysis/pointer_inference/mod.rs

@@ -108,7 +108,7 @@ impl<'a> PointerInference<'a> {
         let sub_to_entry_node_map = crate::analysis::graph::get_entry_nodes_of_subs(context.graph);
 
         let mut fixpoint_computation =
-            super::forward_interprocedural_fixpoint::create_computation_with_alternate_worklist_order(context, None);
+            super::forward_interprocedural_fixpoint::create_computation_with_bottom_up_worklist_order(context, None);
         if print_stats {
             let _ = log_sender.send(LogThreadMsg::Log(
                 LogMessage::new_info(format!(