diff --git a/src/analyzer/expression_analyzer.rs b/src/analyzer/expression_analyzer.rs
index ef92c0f1..7f4fdd77 100644
--- a/src/analyzer/expression_analyzer.rs
+++ b/src/analyzer/expression_analyzer.rs
@@ -926,4 +926,6 @@ pub(crate) fn add_decision_dataflow(
         (expr_pos.start_offset() as u32, expr_pos.end_offset() as u32),
         Rc::new(cond_type),
     );
+
+    analysis_data.data_flow_graph.add_node(decision_node);
 }
diff --git a/src/analyzer/stmt/if_conditional_analyzer.rs b/src/analyzer/stmt/if_conditional_analyzer.rs
index cbcb99a8..32ee51ed 100644
--- a/src/analyzer/stmt/if_conditional_analyzer.rs
+++ b/src/analyzer/stmt/if_conditional_analyzer.rs
@@ -353,9 +353,7 @@ pub(crate) fn add_branch_dataflow(
                 );
             }
 
-            if analysis_data.data_flow_graph.kind == GraphKind::FunctionBody {
-                analysis_data.data_flow_graph.add_node(branch_node);
-            }
+            analysis_data.data_flow_graph.add_node(branch_node);
         }
     }
 }
diff --git a/src/code_info/data_flow/graph.rs b/src/code_info/data_flow/graph.rs
index 4177148a..489a7227 100644
--- a/src/code_info/data_flow/graph.rs
+++ b/src/code_info/data_flow/graph.rs
@@ -157,7 +157,13 @@ impl DataFlowGraph {
 
         let mut origin_nodes = vec![];
 
-        let mut child_node_ids = vec![assignment_node_id.clone()];
+        let mut child_node_ids = vec![];
+
+        if self.vertices.contains_key(assignment_node_id)
+            || self.sources.contains_key(assignment_node_id)
+        {
+            child_node_ids.push(assignment_node_id.clone());
+        }
 
         for _ in 0..50 {
             let mut all_parent_nodes = vec![];