From 993238e6d7095bbe76a46c475558651a8fa4c2ed Mon Sep 17 00:00:00 2001
From: Dev Null <dev.null@curtin.edu.au>
Date: Fri, 10 May 2024 16:57:58 +0800
Subject: [PATCH] check for bad baseline_idxs in

---
 src/context.rs   |  3 ++-
 src/selection.rs | 23 +++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/context.rs b/src/context.rs
index c1f65a5..50c7f4b 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -195,7 +195,8 @@ impl MwaObsContext {
             delays: meta_ctx.delays.clone(),
         };
 
-        for (ant, mut input, number, receiver, mut slot, mut length) in izip!(
+        #[allow(unused_mut)]
+        for (ant, mut input, mut number, mut receiver, mut slot, mut length) in izip!(
             ants,
             result.ant_inputs.outer_iter_mut(),
             result.ant_numbers.iter_mut(),
diff --git a/src/selection.rs b/src/selection.rs
index 4129808..ff7dee9 100644
--- a/src/selection.rs
+++ b/src/selection.rs
@@ -337,6 +337,16 @@ impl VisSelection {
             });
         };
 
+        let max_bl_idx = corr_ctx.metafits_context.baselines.len();
+        // check all selected baseline idxs are < max_bl_idx
+        if self.baseline_idxs.iter().any(|&idx| idx >= max_bl_idx) {
+            return Err(SelectionError::BadBaselineIdx {
+                function: "VisSelection::read_mwalib".to_string(),
+                expected: format!(" < {max_bl_idx}"),
+                received: format!("{:?}", self.baseline_idxs.clone()),
+            });
+        }
+
         // since we are using read_by_baseline_into_buffer, the visibilities are read in order:
         // baseline,frequency,pol,r,i
 
@@ -650,6 +660,19 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_read_mwalib_bad_baseline_idxs() {
+        let corr_ctx = get_mwax_context();
+        let mut vis_sel = VisSelection::from_mwalib(&corr_ctx).unwrap();
+        vis_sel.baseline_idxs = vec![99999999];
+        // Create a blank array to store flags and visibilities
+        let fine_chans_per_coarse = corr_ctx.metafits_context.num_corr_fine_chans_per_coarse;
+        let mut flag_array = vis_sel.allocate_flags(fine_chans_per_coarse).unwrap();
+        let mut jones_array = vis_sel.allocate_jones(fine_chans_per_coarse).unwrap();
+        // read visibilities out of the gpubox files
+        assert!(vis_sel.read_mwalib(&corr_ctx, jones_array.view_mut(), flag_array.view_mut()).is_err());
+    }
+
     #[test]
     #[allow(clippy::unnecessary_cast)]
     fn test_read_mwalib_mwa_legacy() {