1#[cfg(feature = "parallel")]
2use rayon::prelude::*;
3use serde::{Deserialize, Serialize};
4use std::ops::{Add, Div, Index, IndexMut, Mul, Neg, Range, Sub};
5
6#[derive(Debug, Clone)]
7pub enum Grid2dError {
8 DifferentDimensions((usize, usize), (usize, usize)),
9}
10
11#[derive(Debug)]
12pub struct Grid2dNeighborSample<T> {
13 cells: [T; 9],
14 cols: usize,
15 rows: usize,
16}
17
18impl<T> Grid2dNeighborSample<T>
19where
20 T: Clone,
21{
22 #[inline]
23 pub fn cols(&self) -> usize {
24 self.cols
25 }
26
27 #[inline]
28 pub fn rows(&self) -> usize {
29 self.rows
30 }
31
32 #[inline]
33 pub fn size(&self) -> (usize, usize) {
34 (self.cols, self.rows)
35 }
36
37 #[inline]
38 pub fn len(&self) -> usize {
39 self.cols * self.rows
40 }
41
42 #[inline]
43 pub fn is_empty(&self) -> bool {
44 self.cols == 0 && self.rows == 0
45 }
46
47 #[inline]
48 pub fn cells(&self) -> &[T] {
49 &self.cells
50 }
51
52 #[inline]
53 pub fn cell(&self, col: usize, row: usize) -> Option<&T> {
54 if col < self.cols && row < self.rows {
55 Some(&self.cells[row * self.cols + col])
56 } else {
57 None
58 }
59 }
60
61 #[inline]
62 pub fn cell_mut(&mut self, col: usize, row: usize) -> Option<&mut T> {
63 if col < self.cols && row < self.rows {
64 Some(&mut self.cells[row * self.cols + col])
65 } else {
66 None
67 }
68 }
69
70 #[inline]
71 pub fn get(&self, col: usize, row: usize) -> Option<T> {
72 if col < self.cols && row < self.rows {
73 Some(self.cells[row * self.cols + col].clone())
74 } else {
75 None
76 }
77 }
78
79 #[inline]
80 pub fn set(&mut self, col: usize, row: usize, value: T) {
81 if col < self.cols && row < self.rows {
82 self.cells[row * self.cols + col] = value;
83 }
84 }
85
86 pub fn map<F, R>(&self, mut f: F) -> Grid2dNeighborSample<R>
87 where
88 F: FnMut(usize, usize, &T) -> R,
89 R: Default + Copy,
90 {
91 let mut cells = [R::default(); 9];
92 let mut index = 0;
93 for (i, v) in self.cells.iter().enumerate() {
94 cells[index] = f(i % self.cols, i / self.cols, v);
95 index += 1;
96 }
97 Grid2dNeighborSample::<R> {
98 cells,
99 cols: self.cols,
100 rows: self.rows,
101 }
102 }
103
104 pub fn with<F>(&mut self, mut f: F)
105 where
106 F: FnMut(usize, usize, &T) -> T,
107 {
108 for (i, v) in self.cells.iter_mut().enumerate() {
109 *v = f(i % self.cols, i / self.cols, v);
110 }
111 }
112
113 #[inline]
114 pub fn iter(&self) -> impl Iterator<Item = &T> {
115 self.cells.iter()
116 }
117
118 #[inline]
119 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
120 self.cells.iter_mut()
121 }
122}
123
124impl<T> Index<(usize, usize)> for Grid2dNeighborSample<T>
125where
126 T: Clone + Send + Sync,
127{
128 type Output = T;
129
130 fn index(&self, (col, row): (usize, usize)) -> &T {
131 self.cell(col, row).unwrap()
132 }
133}
134
135impl<T> Index<[usize; 2]> for Grid2dNeighborSample<T>
136where
137 T: Clone + Send + Sync,
138{
139 type Output = T;
140
141 fn index(&self, [col, row]: [usize; 2]) -> &T {
142 self.cell(col, row).unwrap()
143 }
144}
145
146impl<T> IndexMut<(usize, usize)> for Grid2dNeighborSample<T>
147where
148 T: Clone + Send + Sync,
149{
150 fn index_mut(&mut self, (col, row): (usize, usize)) -> &mut T {
151 self.cell_mut(col, row).unwrap()
152 }
153}
154
155impl<T> IndexMut<[usize; 2]> for Grid2dNeighborSample<T>
156where
157 T: Clone + Send + Sync,
158{
159 fn index_mut(&mut self, [col, row]: [usize; 2]) -> &mut T {
160 self.cell_mut(col, row).unwrap()
161 }
162}
163
164impl<I, T> From<(usize, I)> for Grid2dNeighborSample<T>
165where
166 I: Iterator<Item = T>,
167 T: Default + Copy,
168{
169 fn from((cols, iter): (usize, I)) -> Self {
170 let mut cells = [T::default(); 9];
171 let mut index = 0;
172 for v in iter.take(9) {
173 cells[index] = v;
174 index += 1;
175 }
176 Self {
177 cells,
178 cols,
179 rows: index / cols,
180 }
181 }
182}
183
184#[derive(Debug, Default, Clone, Serialize, Deserialize)]
185pub struct Grid2d<T> {
186 cells: Vec<T>,
187 cols: usize,
188 rows: usize,
189}
190
191impl<T> Grid2d<T>
192where
193 T: Clone + Send + Sync,
194{
195 pub fn new(cols: usize, rows: usize, fill: T) -> Self {
196 Self {
197 cells: vec![fill; cols * rows],
198 cols,
199 rows,
200 }
201 }
202
203 pub fn with_cells(cols: usize, cells: Vec<T>) -> Self {
204 if cells.len() % cols != 0 {
205 panic!(
206 "cells does not fill grid with desired cols number. cells: {}, cols: {}",
207 cells.len(),
208 cols
209 );
210 }
211 let rows = cells.len() / cols;
212 Self { cells, cols, rows }
213 }
214
215 pub fn resize(&mut self, cols: usize, rows: usize, default: T) {
216 if cols == self.cols && rows == self.rows {
217 return;
218 }
219 self.cells.resize(cols * rows, default);
220 self.cols = cols;
221 self.rows = rows;
222 }
223
224 #[inline]
225 pub fn cols(&self) -> usize {
226 self.cols
227 }
228
229 #[inline]
230 pub fn rows(&self) -> usize {
231 self.rows
232 }
233
234 #[inline]
235 pub fn size(&self) -> (usize, usize) {
236 (self.cols, self.rows)
237 }
238
239 #[inline]
240 pub fn len(&self) -> usize {
241 self.cols * self.rows
242 }
243
244 #[inline]
245 pub fn is_empty(&self) -> bool {
246 self.cols == 0 && self.rows == 0
247 }
248
249 #[inline]
250 pub fn cells(&self) -> &[T] {
251 &self.cells
252 }
253
254 #[inline]
255 pub fn cells_mut(&mut self) -> &mut [T] {
256 &mut self.cells
257 }
258
259 #[inline]
260 pub fn cell(&self, col: usize, row: usize) -> Option<&T> {
261 if col < self.cols && row < self.rows {
262 Some(&self.cells[row * self.cols + col])
263 } else {
264 None
265 }
266 }
267
268 #[inline]
269 pub fn cell_mut(&mut self, col: usize, row: usize) -> Option<&mut T> {
270 if col < self.cols && row < self.rows {
271 Some(&mut self.cells[row * self.cols + col])
272 } else {
273 None
274 }
275 }
276
277 #[inline]
278 pub fn get(&self, col: usize, row: usize) -> Option<T> {
279 if col < self.cols && row < self.rows {
280 Some(self.cells[row * self.cols + col].clone())
281 } else {
282 None
283 }
284 }
285
286 #[inline]
287 pub fn set(&mut self, col: usize, row: usize, value: T) {
288 if col < self.cols && row < self.rows {
289 self.cells[row * self.cols + col] = value;
290 }
291 }
292
293 pub fn get_col_cells(&self, index: usize) -> Option<Vec<T>> {
294 if index < self.cols {
295 let mut result = Vec::with_capacity(self.rows);
296 for row in 0..self.rows {
297 result.push(self.cells[row * self.cols + index].clone());
298 }
299 Some(result)
300 } else {
301 None
302 }
303 }
304
305 pub fn get_row_cells(&self, index: usize) -> Option<Vec<T>> {
306 if index < self.rows {
307 let start = index * self.cols;
308 let end = start + self.cols;
309 Some(self.cells.as_slice()[start..end].to_owned())
310 } else {
311 None
312 }
313 }
314
315 pub fn copy_part(&self, mut range: Range<(usize, usize)>, result: &mut Self)
316 where
317 T: Default,
318 {
319 range.end.0 = range.end.0.min(self.cols);
320 range.end.1 = range.end.1.min(self.rows);
321 range.start.0 = range.start.0.min(range.end.0);
322 range.start.1 = range.start.1.min(range.end.1);
323 let cols = range.end.0 - range.start.0;
324 let rows = range.end.1 - range.start.1;
325 result.resize(cols, rows, T::default());
326 for row in range.start.1..range.end.1 {
327 for col in range.start.0..range.end.0 {
328 result.set(col, row, self.cells[row * self.cols + col].clone());
329 }
330 }
331 }
332
333 pub fn get_part(&self, mut range: Range<(usize, usize)>) -> Self {
334 range.end.0 = range.end.0.min(self.cols);
335 range.end.1 = range.end.1.min(self.rows);
336 range.start.0 = range.start.0.min(range.end.0);
337 range.start.1 = range.start.1.min(range.end.1);
338 let cols = range.end.0 - range.start.0;
339 let rows = range.end.1 - range.start.1;
340 let mut result = Vec::with_capacity(cols * rows);
341 for row in range.start.1..range.end.1 {
342 for col in range.start.0..range.end.0 {
343 result.push(self.cells[row * self.cols + col].clone());
344 }
345 }
346 Self::with_cells(cols, result)
347 }
348
349 pub fn get_view(&self, mut range: Range<(usize, usize)>) -> Grid2d<&T> {
350 range.end.0 = range.end.0.min(self.cols);
351 range.end.1 = range.end.1.min(self.rows);
352 range.start.0 = range.start.0.min(range.end.0);
353 range.start.1 = range.start.1.min(range.end.1);
354 let cols = range.end.0 - range.start.0;
355 let rows = range.end.1 - range.start.1;
356 let mut result = Vec::with_capacity(cols * rows);
357 for row in range.start.1..range.end.1 {
358 for col in range.start.0..range.end.0 {
359 result.push(&self.cells[row * self.cols + col]);
360 }
361 }
362 Grid2d::with_cells(cols, result)
363 }
364
365 pub fn copy_sample(&self, (col, row): (usize, usize), margin: usize, result: &mut Self)
366 where
367 T: Default,
368 {
369 let min = (col.max(margin) - margin, row.max(margin) - margin);
370 let max = (col + margin + 1, row + margin + 1);
371 self.copy_part(min..max, result);
372 }
373
374 pub fn sample(&self, (col, row): (usize, usize), margin: usize) -> Self {
375 let min = (col.max(margin) - margin, row.max(margin) - margin);
376 let max = (col + margin + 1, row + margin + 1);
377 self.get_part(min..max)
378 }
379
380 pub fn view_sample(&self, (col, row): (usize, usize), margin: usize) -> Grid2d<&T> {
381 let min = (col.max(margin) - margin, row.max(margin) - margin);
382 let max = (col + margin + 1, row + margin + 1);
383 self.get_view(min..max)
384 }
385
386 pub fn sin_map<F, R>(&self, mut f: F) -> Grid2d<R>
387 where
388 F: FnMut(usize, usize, &T) -> R,
389 R: Clone + Send + Sync,
390 {
391 Grid2d::<R>::with_cells(
392 self.cols,
393 self.cells
394 .iter()
395 .enumerate()
396 .map(|(i, v)| f(i % self.cols, i / self.cols, v))
397 .collect(),
398 )
399 }
400
401 #[cfg(feature = "parallel")]
402 pub fn par_map<F, R>(&self, f: F) -> Grid2d<R>
403 where
404 F: FnMut(usize, usize, &T) -> R,
405 F: Clone + Send + Sync,
406 R: Clone + Send + Sync,
407 {
408 let cols = self.cols;
409 let cells = self
410 .cells
411 .par_iter()
412 .enumerate()
413 .map(|(i, v)| f.clone()(i % cols, i / cols, v))
414 .collect();
415 Grid2d::<R>::with_cells(cols, cells)
416 }
417
418 pub fn map<F, R>(&self, f: F) -> Grid2d<R>
419 where
420 F: FnMut(usize, usize, &T) -> R,
421 F: Clone + Send + Sync,
422 R: Clone + Send + Sync,
423 {
424 #[cfg(not(feature = "parallel"))]
425 {
426 self.sin_map(f)
427 }
428 #[cfg(feature = "parallel")]
429 {
430 self.par_map(f)
431 }
432 }
433
434 pub fn sin_with<F>(&mut self, mut f: F)
435 where
436 F: FnMut(usize, usize, &T) -> T,
437 {
438 for (i, v) in self.cells.iter_mut().enumerate() {
439 *v = f(i % self.cols, i / self.cols, v);
440 }
441 }
442
443 #[cfg(feature = "parallel")]
444 pub fn par_with<F>(&mut self, f: F)
445 where
446 F: FnMut(usize, usize, &T) -> T,
447 F: Clone + Send + Sync,
448 {
449 let cols = self.cols;
450 self.cells.par_iter_mut().enumerate().for_each(|(i, v)| {
451 *v = f.clone()(i % cols, i / cols, v);
452 });
453 }
454
455 pub fn with<F>(&mut self, f: F)
456 where
457 F: FnMut(usize, usize, &T) -> T,
458 F: Clone + Send + Sync,
459 {
460 #[cfg(not(feature = "parallel"))]
461 self.sin_with(f);
462 #[cfg(feature = "parallel")]
463 self.par_with(f);
464 }
465
466 #[inline]
467 pub fn iter(&self) -> impl Iterator<Item = &T> {
468 self.cells.iter()
469 }
470
471 #[cfg(feature = "parallel")]
472 #[inline]
473 pub fn par_iter(&self) -> impl IndexedParallelIterator<Item = &T> {
474 self.cells.par_iter()
475 }
476
477 pub fn iter_view(
478 &self,
479 mut range: Range<(usize, usize)>,
480 ) -> impl Iterator<Item = (usize, usize, &T)> {
481 range.end.0 = range.end.0.min(self.cols);
482 range.end.1 = range.end.1.min(self.rows);
483 range.start.0 = range.start.0.min(range.end.0);
484 range.start.1 = range.start.1.min(range.end.1);
485 let cols = range.end.0 - range.start.0;
486 let rows = range.end.1 - range.start.1;
487 (0..(cols * rows)).map(move |i| {
488 let lc = i % cols;
489 let lr = i / cols;
490 let gc = range.start.0 + lc;
491 let gr = range.start.1 + lr;
492 (gc, gr, &self.cells[gr * self.cols + gc])
493 })
494 }
495
496 pub fn iter_sample<'a>(
497 &'a self,
498 mut range: Range<(usize, usize)>,
499 margin: usize,
500 ) -> impl Iterator<Item = (usize, usize, Grid2d<&T>)> + 'a {
501 range.end.0 = range.end.0.min(self.cols);
502 range.end.1 = range.end.1.min(self.rows);
503 range.start.0 = range.start.0.min(range.end.0);
504 range.start.1 = range.start.1.min(range.end.1);
505 let cols = range.end.0 - range.start.0;
506 let rows = range.end.1 - range.start.1;
507 (0..(cols * rows)).map(move |i| {
508 let lc = i % cols;
509 let lr = i / cols;
510 let gc = range.start.0 + lc;
511 let gr = range.start.1 + lr;
512 (gc, gr, self.view_sample((gc, gr), margin))
513 })
514 }
515
516 #[inline]
517 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
518 self.cells.iter_mut()
519 }
520
521 #[cfg(feature = "parallel")]
522 #[inline]
523 pub fn par_iter_mut(&mut self) -> impl IndexedParallelIterator<Item = &mut T> {
524 self.cells.par_iter_mut()
525 }
526
527 pub fn iter_view_mut(
528 &mut self,
529 mut range: Range<(usize, usize)>,
530 ) -> impl Iterator<Item = (usize, usize, &mut T)> {
531 range.end.0 = range.end.0.min(self.cols);
532 range.end.1 = range.end.1.min(self.rows);
533 range.start.0 = range.start.0.min(range.end.0);
534 range.start.1 = range.start.1.min(range.end.1);
535 let cols = self.cols;
536 self.cells.iter_mut().enumerate().filter_map(move |(i, v)| {
537 let c = i % cols;
538 let r = i / cols;
539 if c >= range.start.0 && c < range.end.0 && r >= range.start.1 && r < range.end.1 {
540 Some((c, r, v))
541 } else {
542 None
543 }
544 })
545 }
546
547 pub fn neighbor_sample(&self, (col, row): (usize, usize)) -> Grid2dNeighborSample<T>
548 where
549 T: Default + Copy,
550 {
551 let min = (col.max(1) - 1, row.max(1) - 1);
552 let max = ((col + 2).min(self.cols), (row + 2).min(self.rows));
553 let cols = max.0 - min.0;
554 let rows = max.1 - min.1;
555 let mut cells = [T::default(); 9];
556 let mut index = 0;
557 for row in min.1..max.1 {
558 for col in min.0..max.0 {
559 cells[index] = self.cells[row * self.cols + col];
560 index += 1;
561 }
562 }
563 Grid2dNeighborSample { cells, cols, rows }
564 }
565
566 pub fn into_inner(self) -> (usize, usize, Vec<T>) {
567 (self.cols, self.rows, self.cells)
568 }
569}
570
571impl<T> Index<(usize, usize)> for Grid2d<T>
572where
573 T: Clone + Send + Sync,
574{
575 type Output = T;
576
577 fn index(&self, (col, row): (usize, usize)) -> &T {
578 self.cell(col, row).unwrap()
579 }
580}
581
582impl<T> Index<[usize; 2]> for Grid2d<T>
583where
584 T: Clone + Send + Sync,
585{
586 type Output = T;
587
588 fn index(&self, [col, row]: [usize; 2]) -> &T {
589 self.cell(col, row).unwrap()
590 }
591}
592
593impl<T> IndexMut<(usize, usize)> for Grid2d<T>
594where
595 T: Clone + Send + Sync,
596{
597 fn index_mut(&mut self, (col, row): (usize, usize)) -> &mut T {
598 self.cell_mut(col, row).unwrap()
599 }
600}
601
602impl<T> IndexMut<[usize; 2]> for Grid2d<T>
603where
604 T: Clone + Send + Sync,
605{
606 fn index_mut(&mut self, [col, row]: [usize; 2]) -> &mut T {
607 self.cell_mut(col, row).unwrap()
608 }
609}
610
611impl<T> Add for &Grid2d<T>
612where
613 T: Clone + Send + Sync + Add<Output = T>,
614{
615 type Output = Result<Grid2d<T>, Grid2dError>;
616
617 fn add(self, other: &Grid2d<T>) -> Self::Output {
618 if self.cols() == other.cols() && self.rows() == other.rows() {
619 let cells = self
620 .cells
621 .iter()
622 .zip(other.cells.iter())
623 .map(|(a, b)| a.clone() + b.clone())
624 .collect::<Vec<T>>();
625 Ok(Grid2d::with_cells(self.cols(), cells))
626 } else {
627 Err(Grid2dError::DifferentDimensions(
628 (self.cols(), self.rows()),
629 (other.cols(), other.rows()),
630 ))
631 }
632 }
633}
634
635impl<T> Sub for &Grid2d<T>
636where
637 T: Clone + Send + Sync + Sub<Output = T>,
638{
639 type Output = Result<Grid2d<T>, Grid2dError>;
640
641 fn sub(self, other: &Grid2d<T>) -> Self::Output {
642 if self.cols() == other.cols() && self.rows() == other.rows() {
643 let cells = self
644 .cells
645 .iter()
646 .zip(other.cells.iter())
647 .map(|(a, b)| a.clone() - b.clone())
648 .collect::<Vec<T>>();
649 Ok(Grid2d::with_cells(self.cols(), cells))
650 } else {
651 Err(Grid2dError::DifferentDimensions(
652 (self.cols(), self.rows()),
653 (other.cols(), other.rows()),
654 ))
655 }
656 }
657}
658
659impl<T> Mul for &Grid2d<T>
660where
661 T: Clone + Send + Sync + Mul<Output = T>,
662{
663 type Output = Result<Grid2d<T>, Grid2dError>;
664
665 fn mul(self, other: &Grid2d<T>) -> Self::Output {
666 if self.cols() == other.cols() && self.rows() == other.rows() {
667 let cells = self
668 .cells
669 .iter()
670 .zip(other.cells.iter())
671 .map(|(a, b)| a.clone() * b.clone())
672 .collect::<Vec<T>>();
673 Ok(Grid2d::with_cells(self.cols(), cells))
674 } else {
675 Err(Grid2dError::DifferentDimensions(
676 (self.cols(), self.rows()),
677 (other.cols(), other.rows()),
678 ))
679 }
680 }
681}
682
683impl<T> Div for &Grid2d<T>
684where
685 T: Clone + Send + Sync + Div<Output = T>,
686{
687 type Output = Result<Grid2d<T>, Grid2dError>;
688
689 fn div(self, other: &Grid2d<T>) -> Self::Output {
690 if self.cols() == other.cols() && self.rows() == other.rows() {
691 let cells = self
692 .cells
693 .iter()
694 .zip(other.cells.iter())
695 .map(|(a, b)| a.clone() / b.clone())
696 .collect::<Vec<T>>();
697 Ok(Grid2d::with_cells(self.cols(), cells))
698 } else {
699 Err(Grid2dError::DifferentDimensions(
700 (self.cols(), self.rows()),
701 (other.cols(), other.rows()),
702 ))
703 }
704 }
705}
706
707impl<T> Neg for &Grid2d<T>
708where
709 T: Clone + Send + Sync + Neg<Output = T>,
710{
711 type Output = Grid2d<T>;
712
713 fn neg(self) -> Self::Output {
714 let cells = self.cells.iter().map(|v| -v.clone()).collect::<Vec<T>>();
715 Grid2d::with_cells(self.cols(), cells)
716 }
717}
718
719impl<T> IntoIterator for Grid2d<T>
720where
721 T: Clone + Send + Sync,
722{
723 type Item = T;
724 type IntoIter = ::std::vec::IntoIter<T>;
725
726 fn into_iter(self) -> Self::IntoIter {
727 self.cells.into_iter()
728 }
729}
730
731impl<I, T> From<(usize, I)> for Grid2d<T>
732where
733 I: Iterator<Item = T>,
734 T: Clone + Send + Sync,
735{
736 fn from((cols, iter): (usize, I)) -> Self {
737 Self::with_cells(cols, iter.collect::<Vec<T>>())
738 }
739}
740
741impl<T> Into<Vec<T>> for Grid2d<T>
742where
743 T: Clone + Send + Sync,
744{
745 fn into(self) -> Vec<T> {
746 self.cells
747 }
748}
749
750impl<T> Into<(usize, usize, Vec<T>)> for Grid2d<T>
751where
752 T: Clone + Send + Sync,
753{
754 fn into(self) -> (usize, usize, Vec<T>) {
755 self.into_inner()
756 }
757}