From 893849870cc3faade29418561c53195b935d389c Mon Sep 17 00:00:00 2001 From: Logan Date: Mon, 26 Jun 2023 17:40:19 -0500 Subject: [PATCH] extended to all range types --- .gitignore | 1 + Cargo.lock | 7 ++ Cargo.toml | 8 +++ src/lib.rs | 5 ++ src/nditer.rs | 84 ++++++++++++++++++++++ src/ndrange.rs | 186 +++++++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 291 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 src/lib.rs create mode 100644 src/nditer.rs create mode 100644 src/ndrange.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..fb7f5e2 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,7 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "ndrange" +version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..09b4d28 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "ndrange" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..26378b3 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,5 @@ +#![feature(step_trait)] +pub mod nditer; +pub mod ndrange; + +pub use ndrange::*; diff --git a/src/nditer.rs b/src/nditer.rs new file mode 100644 index 0000000..51f7c47 --- /dev/null +++ b/src/nditer.rs @@ -0,0 +1,84 @@ +use std::{ + iter::{FusedIterator, Step}, + ops::RangeBounds, +}; + +use crate::{get_bound_start, ndrange::NdRange}; + +pub enum NdIter +where + R: RangeBounds, + R: ExactSizeIterator, + Self: Sized, +{ + Iterating { + ndrange: NdRange, + current: [T; N], + }, + Done, +} + +impl Iterator for NdIter +where + R: RangeBounds, + R: ExactSizeIterator, + T: Step, +{ + type Item = [T; N]; + + fn next(&mut self) -> Option { + match self { + NdIter::Iterating { ndrange, current } => { + let ret = current.clone(); + for i in 0..current.len() { + let next = T::forward(current[i].clone(), 1); + if ndrange.ranges[i].contains(&next) { + current[i] = next.clone(); + break; + } else { + current[i] = get_bound_start(&ndrange.ranges[i]); + if i == (current.len() - 1) { + *self = Self::Done; + break; + } + } + } + Some(ret) + }, + NdIter::Done => None, + } + } + + fn size_hint(&self) -> (usize, Option) { + let mut size = 1; + match self { + NdIter::Iterating { ndrange, .. } => { + for i in &ndrange.ranges { + size *= i.len(); + } + (size, Some(size)) + }, + NdIter::Done => (0, None), + } + } +} + +impl ExactSizeIterator for NdIter +where + R: RangeBounds, + R: ExactSizeIterator, + T: Step, +{ + fn len(&self) -> usize { + self.size_hint().0 + } +} + +impl FusedIterator for NdIter +where + R: RangeBounds, + R: ExactSizeIterator, + R: FusedIterator, + T: Step, +{ +} diff --git a/src/ndrange.rs b/src/ndrange.rs new file mode 100644 index 0000000..fde8350 --- /dev/null +++ b/src/ndrange.rs @@ -0,0 +1,186 @@ +use std::{iter::Step, marker::PhantomData, ops::RangeBounds}; + +use crate::nditer::NdIter; + +pub(crate) fn get_bound_start(s: &R) -> T +where + R: RangeBounds, + R: ExactSizeIterator, + T: Step, + T: Clone, +{ + match s.start_bound() { + std::ops::Bound::Included(s) => s.clone(), + std::ops::Bound::Excluded(s) => Step::forward(s.clone(), 1), + std::ops::Bound::Unbounded => unreachable!( + "Implementing exact size iterator means the range cannot be unbounded \ + at the start" + ), + } +} + +/// A N-Dimensional extension to Range and its derivatives +/// +/// # Examples +/// ``` +/// assert_eq!(ndrange!(0..5, 0..5), NdRange::new([0..5, 0..5])); +/// assert_eq!(ndrange!(0..3, 0..3).len(), 9); +/// +/// let mut iter = ndrange!(0..2, 0..2).into_iter(); +/// assert_eq!(iter.next(), Some([0, 0])); +/// assert_eq!(iter.next(), Some([1, 0])); +/// assert_eq!(iter.next(), Some([0, 1])); +/// assert_eq!(iter.next(), Some([1, 1])); +/// assert_eq!(iter.next(), None); +/// +/// assert!(ndrange!(0..10, 0..10).contains(&[2, 4])); +/// +/// assert_eq!(ndrange!().len(), 0); +/// assert_eq!(ndrange!(0..7, 0..0).len(), 0); +/// assert!(!ndrange!(0..3, 0..2, 0..0).contains(&[2, 1, 0])); +/// +/// for i in NdRange::new([0..5, 0..5]) { +/// println!("{:?}", i); +/// } +/// +/// for _ in ndrange!() { +/// panic!("Unreachable, empty NdRange returns an empty +/// iterator") } +/// +/// for i in ndrange!(0..5, 0..5) { +/// println!("{i:?}"); +/// } +/// ``` +#[derive(Clone, Eq, Hash)] +pub struct NdRange +where + R: RangeBounds, +{ + pub(crate) ranges: [R; N], + pub(crate) _phantom: PhantomData, +} + +impl std::fmt::Debug for NdRange +where + R: RangeBounds, + R: std::fmt::Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[ ")?; + for r in &self.ranges { + r.fmt(f)?; + write!(f, ", ")?; + } + write!(f, "]")?; + Ok(()) + } +} + +impl PartialEq for NdRange +where + R: RangeBounds, + R: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + for i in 0..N { + if self.ranges[i] != other.ranges[i] { + return false; + } + } + return true; + } +} + +impl NdRange +where + R: RangeBounds, +{ + pub fn new(ranges: [R; N]) -> Self { + Self { + ranges, + _phantom: PhantomData, + } + } +} + +impl NdRange +where + R: RangeBounds, + Idx: PartialOrd, +{ + pub fn contains(&self, other: &[U; N]) -> bool + where + Idx: PartialOrd, + U: PartialOrd, + { + for i in 0..N { + if !self.ranges[i].contains(&other[i]) { + return false; + } + } + return true; + } +} + +impl NdRange +where + Self: Sized, + R: RangeBounds, + R: ExactSizeIterator, + R: Clone, + Idx: Step, +{ + pub fn len(&self) -> usize { + self.clone().into_iter().len() + } +} + +impl IntoIterator for NdRange +where + R: RangeBounds, + R: ExactSizeIterator, + Idx: Step, +{ + type IntoIter = NdIter; + type Item = [Idx; N]; + + fn into_iter(self) -> Self::IntoIter { + let mut empty = N == 0; + for i in 0..N { + if self.ranges[i].len() == 0 { + empty = true; + break; + } + } + + match empty { + true => NdIter::Done, + false => NdIter::Iterating { + current: std::array::from_fn(|i| get_bound_start(&self.ranges[i])), + ndrange: self, + }, + } + } +} + +impl Default for NdRange +where + R: RangeBounds, + R: Default, +{ + fn default() -> Self { + Self { + ranges: std::array::from_fn(|_| R::default()), + _phantom: PhantomData, + } + } +} + +#[macro_export] +macro_rules! ndrange { + + ( $( $x:expr ), + $(,)? ) => { + NdRange::new([ $($x,)* ]) + }; + +}