ZJIT: Deduplicate successor and predecessor sets (#15263)

Fixes https://github.com/Shopify/ruby/issues/877

I didn't consider the ability to have the successor or predecessor sets having duplicates when originally crafting the Iongraph support PR, but have added this to prevent that happening in the future.

I don't think it interferes with the underlying Iongraph implementation, but it doesn't really make sense.

I think this kind of behaviour happens when there are multiple jump instructions that go to the same basic block within a given block.

2 files changed, 31 insertions, 3 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index bbe5dd3435..face61f1f6 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -9,7 +9,7 @@ use crate::{
     cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json
 };
 use std::{
-    cell::RefCell, collections::{HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, c_int, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
+    cell::RefCell, collections::{BTreeSet, HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, c_int, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
 };
 use crate::hir_type::{Type, types};
 use crate::bitset::BitSet;
@@ -5849,7 +5849,13 @@ impl<'a> ControlFlowInfo<'a> {
             // Since ZJIT uses extended basic blocks, one must check all instructions
             // for their ability to jump to another basic block, rather than just
             // the instructions at the end of a given basic block.
-            let successors: Vec<BlockId> = block
+            //
+            // Use BTreeSet to avoid duplicates and maintain an ordering. Also
+            // `BTreeSet<BlockId>` provides conversion trivially back to an `Vec<BlockId>`.
+            // Ordering is important so that the expect tests that serialize the predecessors
+            // and successors don't fail intermittently.
+            // todo(aidenfoxivey): Use `BlockSet` in lieu of `BTreeSet<BlockId>`
+            let successors: BTreeSet<BlockId> = block
                 .insns
                 .iter()
                 .map(|&insn_id| uf.find_const(insn_id))
@@ -5867,7 +5873,8 @@ impl<'a> ControlFlowInfo<'a> {
             }
 
             // Store successors for this block.
-            successor_map.insert(block_id, successors);
+            // Convert successors from a `BTreeSet<BlockId>` to a `Vec<BlockId>`.
+            successor_map.insert(block_id, successors.iter().copied().collect());
         }
 
         Self {
diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs
index 5e6ec11892..b487352748 100644
--- a/zjit/src/hir/tests.rs
+++ b/zjit/src/hir/tests.rs
@@ -3484,6 +3484,27 @@ pub mod hir_build_tests {
         assert!(cfi.is_succeeded_by(bb1, bb0));
         assert!(cfi.is_succeeded_by(bb3, bb1));
      }
+
+     #[test]
+     fn test_cfi_deduplicated_successors_and_predecessors() {
+         let mut function = Function::new(std::ptr::null());
+
+         let bb0 = function.entry_block;
+         let bb1 = function.new_block(0);
+
+         // Construct two separate jump instructions.
+         let v1 = function.push_insn(bb0, Insn::Const { val: Const::Value(Qfalse) });
+         let _ = function.push_insn(bb0, Insn::IfTrue { val: v1, target: edge(bb1)});
+         function.push_insn(bb0, Insn::Jump(edge(bb1)));
+
+         let retval = function.push_insn(bb1, Insn::Const { val: Const::CBool(true) });
+         function.push_insn(bb1, Insn::Return { val: retval });
+
+         let cfi = ControlFlowInfo::new(&function);
+
+         assert_eq!(cfi.predecessors(bb1).collect::<Vec<_>>().len(), 1);
+         assert_eq!(cfi.successors(bb0).collect::<Vec<_>>().len(), 1);
+     }
  }
 
  /// Test dominator set computations.