: selection: routing: handle evals of selections over 0-dim slices (#507)

Summary:

handle 0d slices by canonically embedding them as 1d slices of extent 1, enabling uniform evaluation and routing logic. adds test coverage for  `eval` and `next_steps`.

Differential Revision: D78168758

Author: shayne-fletcher

ndslice/src/selection.rs

@@ -565,7 +565,20 @@ impl Selection {
         opts: &EvalOpts,
         slice: &'a Slice,
     ) -> Result<Box<dyn Iterator<Item = usize> + 'a>, ShapeError> {
-        Ok(Self::validate(self, opts, slice)?.eval_rec(slice, vec![0; slice.num_dim()], 0))
+        // Canonically embed 0D as 1D (extent 1).
+        if slice.num_dim() == 0 {
+            let slice = Slice::new(slice.offset(), vec![1], vec![1]).unwrap();
+            return Ok(Box::new(
+                self.validate(opts, &slice)?
+                    .eval_rec(&slice, vec![0; 1], 0)
+                    .collect::<Vec<_>>()
+                    .into_iter(),
+            ));
+        }
+
+        Ok(self
+            .validate(opts, slice)?
+            .eval_rec(slice, vec![0; slice.num_dim()], 0))
     }
 
     fn eval_rec<'a>(
@@ -1831,6 +1844,17 @@ mod tests {
         assert_matches!(res.as_slice(), [i, j] if *i < *j && *i < 8 && *j < 8);
     }
 
+    #[test]
+    fn test_eval_zero_dim_slice() {
+        let slice_0d = Slice::new(1, vec![], vec![]).unwrap();
+        assert_eq!(eval(true_(), &slice_0d), vec![1]);
+        assert_eq!(eval(false_(), &slice_0d), vec![]);
+        assert_eq!(eval(all(true_()), &slice_0d), vec![1]);
+        assert_eq!(eval(all(false_()), &slice_0d), vec![]);
+        assert_eq!(eval(union(true_(), true_()), &slice_0d), vec![1]);
+        assert_eq!(eval(intersection(true_(), false_()), &slice_0d), vec![]);
+    }
+
     #[test]
     fn test_selection_10() {
         let slice = &test_slice();

ndslice/src/selection/routing.rs

@@ -378,6 +378,14 @@ impl RoutingFrame {
         _chooser: &mut dyn FnMut(&Choice) -> usize,
         f: &mut dyn FnMut(RoutingStep) -> ControlFlow<()>,
     ) -> ControlFlow<()> {
+        if self.slice.num_dim() == 0 {
+            // Canonically embed 0D as 1D (extent 1).
+            let embedded = Slice::new(self.slice.offset(), vec![1], vec![1]).unwrap();
+            let mut this = self.clone();
+            this.slice = Arc::new(embedded);
+            this.here = vec![0];
+            return this.next_steps(_chooser, f);
+        }
         let selection = self
             .selection
             .clone()
@@ -1590,4 +1598,48 @@ mod tests {
             "Expected panic due to overdelivery, but no panic occurred"
         );
     }
+
+    #[test]
+    fn test_next_steps_zero_dim_slice() {
+        use std::ops::ControlFlow;
+
+        use crate::selection::dsl::*;
+
+        let slice = Slice::new(42, vec![], vec![]).unwrap();
+        let selection = true_();
+        let frame = RoutingFrame::root(selection, slice.clone());
+
+        let mut steps = vec![];
+        let _ = frame.next_steps(
+            &mut |_| panic!("Unexpected Choice in 0D test"),
+            &mut |step| {
+                steps.push(step);
+                ControlFlow::Continue(())
+            },
+        );
+
+        assert_eq!(steps.len(), 1);
+        let step = steps[0].as_forward().unwrap();
+        assert_eq!(step.here, vec![0]);
+        assert!(step.deliver_here());
+        assert_eq!(step.slice.location(&step.here).unwrap(), 42);
+
+        let selection = all(false_());
+        let frame = RoutingFrame::root(selection, slice);
+
+        let mut steps = vec![];
+        let _ = frame.next_steps(
+            &mut |_| panic!("Unexpected Choice in 0D test"),
+            &mut |step| {
+                steps.push(step);
+                ControlFlow::Continue(())
+            },
+        );
+
+        assert_eq!(steps.len(), 1);
+        let step = steps[0].as_forward().unwrap();
+        assert_eq!(step.here, vec![0]);
+        assert!(!step.deliver_here());
+        assert_eq!(step.slice.location(&step.here).unwrap(), 42);
+    }
 }