From dc52f07a28b5f9218c81c2c418d09675af0d4380 Mon Sep 17 00:00:00 2001
From: Pankaj Garg <pankaj@solana.com>
Date: Mon, 10 Apr 2023 10:34:53 +0530
Subject: [PATCH] Fix finding ancestor during program cache pruning

---
 program-runtime/src/loaded_programs.rs | 45 +++++++++++++++++++++++---
 1 file changed, 41 insertions(+), 4 deletions(-)

diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
index 4a3c8b5b1a5a81..24d4541b2539aa 100644
--- a/program-runtime/src/loaded_programs.rs
+++ b/program-runtime/src/loaded_programs.rs
@@ -283,7 +283,7 @@ impl LoadedPrograms {
     /// Before rerooting the blockstore this removes all programs of orphan forks
     pub fn prune<F: ForkGraph>(&mut self, fork_graph: &F, new_root: Slot) {
         self.entries.retain(|_key, second_level| {
-            let mut first_ancestor = true;
+            let mut first_ancestor_found = false;
             *second_level = second_level
                 .iter()
                 .rev()
@@ -291,9 +291,9 @@ impl LoadedPrograms {
                     let relation = fork_graph.relationship(entry.deployment_slot, new_root);
                     if entry.deployment_slot >= new_root {
                         matches!(relation, BlockRelation::Equal | BlockRelation::Descendant)
-                    } else if first_ancestor {
-                        first_ancestor = false;
-                        matches!(relation, BlockRelation::Ancestor)
+                    } else if !first_ancestor_found && matches!(relation, BlockRelation::Ancestor) {
+                        first_ancestor_found = true;
+                        first_ancestor_found
                     } else {
                         false
                     }
@@ -1134,4 +1134,41 @@ mod tests {
         // program3 was deployed on slot 25, which has been pruned
         assert!(missing.contains(&program3));
     }
+
+    #[test]
+    fn test_fork_prune_find_first_ancestor() {
+        let mut cache = LoadedPrograms::default();
+
+        // Fork graph created for the test
+        //                   0
+        //                 /   \
+        //                10    5
+        //                |
+        //                20
+
+        // Deploy program on slot 0, and slot 5.
+        // Prune the fork that has slot 5. The cache should still have the program
+        // deployed at slot 0.
+        let mut fork_graph = TestForkGraphSpecific::default();
+        fork_graph.insert_fork(&[0, 10, 20]);
+        fork_graph.insert_fork(&[0, 5]);
+
+        let program1 = Pubkey::new_unique();
+        assert!(!cache.replenish(program1, new_test_loaded_program(0, 1)).0);
+        assert!(!cache.replenish(program1, new_test_loaded_program(5, 6)).0);
+
+        cache.prune(&fork_graph, 10);
+
+        let working_slot = TestWorkingSlot::new(20, &[0, 10, 20]);
+        let (found, _missing) = cache.extract(&working_slot, vec![program1].into_iter());
+
+        // The cache should have the program deployed at slot 0
+        assert_eq!(
+            found
+                .get(&program1)
+                .expect("Did not find the program")
+                .deployment_slot,
+            0
+        );
+    }
 }