1mod coo_error;
4mod coo_utils;
5
6#[cfg(feature = "alloc")]
7use alloc::vec::Vec;
8use cl_aux::{ArrayWrapper, SingleTypeStorage};
9pub use coo_error::*;
10use coo_utils::{does_not_have_duplicates_sorted, value, value_mut};
11
12pub type CooArray<DATA, const D: usize, const DN: usize> = Coo<[([usize; D], DATA); DN], D>;
14pub type CooMut<'data, DATA, const D: usize> = Coo<&'data mut [([usize; D], DATA)], D>;
16pub type CooRef<'data, DATA, const D: usize> = Coo<&'data [([usize; D], DATA)], D>;
18#[cfg(feature = "alloc")]
19pub type CooVec<DATA, const D: usize> = Coo<Vec<([usize; D], DATA)>, D>;
21
22#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
29#[derive(Clone, Debug, Default, Eq, Ord, PartialEq, PartialOrd)]
30pub struct Coo<DS, const D: usize> {
31 pub(crate) data: DS,
32 pub(crate) dims: ArrayWrapper<usize, D>,
33}
34
35impl<DS, const D: usize> Coo<DS, D> {
36 #[inline]
45 pub fn dims(&self) -> &[usize; D] {
46 &self.dims
47 }
48}
49
50impl<DATA, DS, const D: usize> Coo<DS, D>
51where
52 DS: AsRef<[<DS as SingleTypeStorage>::Item]> + SingleTypeStorage<Item = ([usize; D], DATA)>,
53{
54 #[cfg_attr(feature = "alloc", doc = "```rust")]
63 #[cfg_attr(not(feature = "alloc"), doc = "```ignore")]
64 #[inline]
72 pub fn new(dims: [usize; D], data: DS) -> crate::Result<Self> {
73 if !crate::utils::are_in_ascending_order(data.as_ref(), |a, b| [&a.0, &b.0]) {
74 return Err(CooError::InvalidIndcsOrder.into());
75 }
76 let has_invalid_indcs = !data.as_ref().iter().all(|&(indcs, _)| {
77 indcs.iter().zip(dims.iter()).all(
78 |(data_idx, dim)| {
79 if dim == &0 {
80 true
81 } else {
82 data_idx < dim
83 }
84 },
85 )
86 });
87 if has_invalid_indcs {
88 return Err(CooError::InvalidIndcs.into());
89 }
90 if !does_not_have_duplicates_sorted(data.as_ref(), |a, b| a.0[..] != b.0[..]) {
91 return Err(CooError::DuplicatedIndices.into());
92 }
93 Ok(Self { data, dims: dims.into() })
94 }
95
96 #[inline]
105 pub fn data(&self) -> &[([usize; D], DATA)] {
106 self.data.as_ref()
107 }
108
109 #[inline]
124 pub fn value(&self, indcs: [usize; D]) -> Option<&DATA> {
125 value(indcs, self.data.as_ref())
126 }
127}
128
129impl<DATA, DS, const D: usize> Coo<DS, D>
130where
131 DS: AsMut<[<DS as SingleTypeStorage>::Item]> + SingleTypeStorage<Item = ([usize; D], DATA)>,
132{
133 #[inline]
135 pub fn value_mut(&mut self, indcs: [usize; D]) -> Option<&mut DATA> {
136 value_mut(indcs, self.data.as_mut())
137 }
138}
139
140#[cfg(feature = "rand")]
141impl<DATA, DS, const D: usize> Coo<DS, D>
142where
143 DS: AsMut<[<DS as SingleTypeStorage>::Item]>
144 + AsRef<[<DS as SingleTypeStorage>::Item]>
145 + Default
146 + SingleTypeStorage<Item = ([usize; D], DATA)>
147 + cl_aux::CapacityUpperBound
148 + cl_aux::Push<<DS as SingleTypeStorage>::Item>,
149{
150 #[cfg_attr(feature = "alloc", doc = "```rust")]
161 #[cfg_attr(not(feature = "alloc"), doc = "```ignore")]
162 #[inline]
170 pub fn new_controlled_random_rand<R>(
171 dims: [usize; D],
172 nnz: usize,
173 rng: &mut R,
174 mut cb: impl FnMut(&mut R, &[usize; D]) -> DATA,
175 ) -> crate::Result<Self>
176 where
177 R: rand::Rng,
178 {
179 use rand::distr::Distribution as _;
180 if nnz > crate::utils::max_nnz(&dims) {
181 return Err(CooError::NnzGreaterThanMaximumNnz.into());
182 }
183 let mut data: DS = Default::default();
184 if nnz > data.as_ref().len() {
185 return Err(crate::Error::InsufficientCapacity);
186 }
187 for _ in 0..nnz {
188 let indcs: [usize; D] = ArrayWrapper::from_fn(|idx| {
189 let dim = *dims.get(idx).unwrap_or(&0);
190 if dim == 0 {
191 0
192 } else {
193 rand::distr::Uniform::new(0, dim).ok().map(|el| el.sample(rng)).unwrap_or_default()
194 }
195 })
196 .0;
197 if data.as_ref().iter().all(|value| value.0 != indcs) {
198 data
199 .push({
200 let element = cb(rng, &indcs);
201 (indcs, element)
202 })
203 .map_err(|_err| crate::Error::InsufficientCapacity)?;
204 }
205 }
206 data.as_mut().sort_unstable_by(|a, b| a.0.cmp(&b.0));
207 Coo::new(dims, data)
208 }
209
210 #[inline]
217 pub fn new_random_rand<R>(rng: &mut R, upper_bound: usize) -> crate::Result<Self>
218 where
219 R: rand::Rng,
220 rand::distr::StandardUniform: rand::distr::Distribution<DATA>,
221 {
222 let dims = crate::utils::valid_random_dims(rng, upper_bound);
223 let max_nnz = crate::utils::max_nnz(&dims);
224 let nnz = if max_nnz == 0 { 0 } else { rng.random_range(0..max_nnz) };
225 Self::new_controlled_random_rand(dims, nnz, rng, |r, _| r.random())
226 }
227}