1use crate::unreachable_debug_checked;
5use fixedbitset::FixedBitSet;
6use ndarray::Array2;
7use petgraph::{
8 data::{DataMap, DataMapMut},
9 visit::{Data, GraphBase, GraphProp, IntoNodeIdentifiers, NodeCount, VisitMap, Visitable},
10 EdgeType,
11};
12use std::{marker::PhantomData, mem::MaybeUninit, ptr::drop_in_place};
13
14mod edges;
15pub use edges::{EdgeReference, EdgeReferences, Edges, EdgesDirected};
16mod neighbors;
17pub use neighbors::*;
18mod nodes;
19pub use nodes::*;
20pub mod shapes;
21pub(crate) use shapes::*;
22pub mod square;
23
24#[derive(Debug, Clone, PartialEq, Eq, Hash)]
25pub struct LatticeGraph<N, E, S: Shape> {
29 nodes: Array2<N>,
30 edges: Vec<Array2<E>>,
31 s: S,
32}
33
34impl<N, E, S: Shape> LatticeGraph<N, E, S> {
35 #[doc(hidden)]
37 pub unsafe fn new_raw(nodes: Array2<N>, edges: Vec<Array2<E>>, s: S) -> Self {
38 Self { nodes, edges, s }
39 }
40
41 pub unsafe fn new_uninit(s: S) -> LatticeGraph<MaybeUninit<N>, MaybeUninit<E>, S> {
44 let nodes = Array2::uninit((s.horizontal(), s.vertical()));
45 let ac = S::Axis::COUNT;
46 let mut edges = Vec::with_capacity(ac);
47 for _i in 0..ac {
48 edges.push(Array2::uninit((s.horizontal(), s.vertical())))
49 }
50 debug_assert_eq!(edges.len(), S::Axis::COUNT);
51 LatticeGraph { nodes, edges, s }
52 }
53
54 pub fn new(s: S) -> Self
56 where
57 N: Default,
58 E: Default,
59 {
60 Self::new_with(s, |_| N::default(), |_, _| E::default())
61 }
62
63 pub fn new_with<FN, FE>(s: S, mut n: FN, mut e: FE) -> Self
65 where
66 FN: FnMut(S::Coordinate) -> N,
67 FE: FnMut(S::Coordinate, S::Axis) -> E, {
69 let mut uninit = unsafe { Self::new_uninit(s) };
70 let s = &uninit.s;
71 let nodes = uninit.nodes.as_slice_mut().unwrap();
72 let edges = &mut uninit.edges;
73 for i in 0..s.node_count() {
74 let offset = s.index_to_offset(i);
75 let c = s.offset_to_coordinate(offset);
76 unsafe { std::ptr::write(nodes.get_unchecked_mut(i), MaybeUninit::new(n(c))) }
77 for (j, edge) in edges.iter_mut().enumerate() {
78 let a = unsafe { <S::Axis as Axis>::from_index_unchecked(j) };
79 if s.move_coord(c, a.foward()).is_err() {
80 continue;
81 }
82 let ex = e(c, a);
83 let t = edge.get_mut((offset.horizontal, offset.vertical));
84 if let Some(x) = t {
85 unsafe { std::ptr::write(x, MaybeUninit::new(ex)) };
86 }
87 }
88 }
89 unsafe { uninit.assume_init() }
90 }
91
92 pub fn shape(&self) -> &S {
94 &self.s
95 }
96}
97
98impl<N, E, S: Shape + Default> LatticeGraph<N, E, S> {
99 pub fn new_s() -> Self
101 where
102 N: Default,
103 E: Default,
104 {
105 Self::new(S::default())
106 }
107
108 pub unsafe fn new_uninit_s() -> LatticeGraph<MaybeUninit<N>, MaybeUninit<E>, S> {
111 Self::new_uninit(S::default())
112 }
113
114 pub fn new_with_s<FN, FE>(n: FN, e: FE) -> Self
116 where
117 FN: FnMut(S::Coordinate) -> N,
118 FE: FnMut(S::Coordinate, S::Axis) -> E,
119 {
120 Self::new_with(S::default(), n, e)
121 }
122}
123
124impl<N, E, S: Shape> LatticeGraph<MaybeUninit<N>, MaybeUninit<E>, S> {
125 pub unsafe fn assume_init(self) -> LatticeGraph<N, E, S> {
145 let md = std::mem::ManuallyDrop::new(self);
146 LatticeGraph {
147 nodes: core::ptr::read(&md.nodes).assume_init(),
148 edges: core::ptr::read(&md.edges)
149 .into_iter()
150 .map(|e| e.assume_init())
151 .collect(),
152 s: core::ptr::read(&md.s),
153 }
154 }
155}
156
157impl<N, E, S: Shape> Drop for LatticeGraph<N, E, S> {
158 fn drop(&mut self) {
159 if std::mem::needs_drop::<E>() {
161 let ni = self.node_identifiers();
162 let s = &self.s;
163 let e = &mut self.edges;
164 unsafe {
165 for (di, edges) in e.drain(..).enumerate() {
166 let dir = S::Axis::from_index_unchecked(di).foward();
167 for (coord, mut e) in ni.clone().zip(edges.into_iter()) {
168 if s.move_coord(coord, dir.clone()).is_ok() {
169 drop_in_place(&mut e);
170 }
171 }
172 }
173 }
174 }
175 }
176}
177
178impl<N, E, S> Default for LatticeGraph<N, E, S>
179where
180 N: Default,
181 E: Default,
182 S: Shape + Default + Clone,
183{
184 fn default() -> Self {
185 Self::new(S::default())
186 }
187}
188
189impl<N, E, S: Shape> GraphBase for LatticeGraph<N, E, S> {
190 type NodeId = S::Coordinate;
191 type EdgeId = (S::Coordinate, S::Axis);
192}
193
194impl<N, E, S: Shape> Data for LatticeGraph<N, E, S> {
195 type NodeWeight = N;
196 type EdgeWeight = E;
197}
198
199impl<N, E, S: Shape> DataMap for LatticeGraph<N, E, S> {
200 fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
201 let offset = self.s.to_offset(id);
202 offset
204 .map(move |offset| unsafe {
205 if cfg!(debug_assertions) {
206 self.nodes
207 .get((offset.horizontal, offset.vertical))
208 .unwrap()
209 } else {
210 self.nodes
211 .get((offset.horizontal, offset.vertical))
212 .unwrap_unchecked()
213 }
214 })
215 .ok()
216 }
217
218 fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
219 let offset = self.s.to_offset(id.0);
220 let ax = id.1.to_index();
221 if let Ok(offset) = offset {
222 if self.s.move_coord(id.0, id.1.foward()).is_err() {
223 return None;
224 }
225 unsafe {
226 self.edges
227 .get_unchecked(ax)
228 .get((offset.horizontal, offset.vertical))
229 }
230 } else {
231 None
232 }
233 }
234}
235
236impl<N, E, S: Shape> DataMapMut for LatticeGraph<N, E, S> {
237 fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
238 let offset = self.s.to_offset(id);
239
240 offset
242 .map(move |offset| unsafe {
243 if cfg!(debug_assertions) {
244 self.nodes
245 .get_mut((offset.horizontal, offset.vertical))
246 .unwrap()
247 } else {
248 self.nodes
249 .get_mut((offset.horizontal, offset.vertical))
250 .unwrap_unchecked()
251 }
252 })
253 .ok()
254 }
255
256 fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
257 let offset = self.s.to_offset(id.0);
258 let ax = id.1.to_index();
259 if let Ok(offset) = offset {
260 if self.s.move_coord(id.0, id.1.foward()).is_err() {
261 return None;
262 }
263 unsafe {
264 self.edges
265 .get_unchecked_mut(ax)
266 .get_mut((offset.horizontal, offset.vertical))
267 }
268 } else {
269 None
270 }
271 }
272}
273
274impl<N, E, S: Shape> LatticeGraph<N, E, S> {
275 #[doc(hidden)]
276 #[inline]
277 pub unsafe fn node_weight_unchecked(
278 &self,
279 id: <LatticeGraph<N, E, S> as GraphBase>::NodeId,
280 ) -> &<LatticeGraph<N, E, S> as Data>::NodeWeight {
281 let offset = self.s.to_offset_unchecked(id);
282 self.node_weight_unchecked_raw(offset)
284 }
285
286 #[doc(hidden)]
287 pub unsafe fn node_weight_unchecked_raw(
288 &self,
289 offset: Offset,
290 ) -> &<LatticeGraph<N, E, S> as Data>::NodeWeight {
291 self.nodes
292 .get((offset.horizontal, offset.vertical))
293 .unwrap_unchecked()
294 }
295
296 #[doc(hidden)]
297 #[inline]
298 pub unsafe fn edge_weight_unchecked(
299 &self,
300 id: <LatticeGraph<N, E, S> as GraphBase>::EdgeId,
301 ) -> &<LatticeGraph<N, E, S> as Data>::EdgeWeight {
302 let offset = self.s.to_offset_unchecked(id.0);
303 let ax = id.1.to_index();
304 self.edge_weight_unchecked_raw((offset, ax))
305 }
306
307 #[doc(hidden)]
308 pub unsafe fn edge_weight_unchecked_raw(
309 &self,
310 (offset, ax): (Offset, usize),
311 ) -> &<LatticeGraph<N, E, S> as Data>::EdgeWeight {
312 self.edges
313 .get_unchecked(ax)
314 .get((offset.horizontal, offset.vertical))
315 .unwrap_unchecked()
316 }
317
318 #[doc(hidden)]
319 pub unsafe fn node_weight_mut_unchecked(
320 &mut self,
321 id: <LatticeGraph<N, E, S> as GraphBase>::NodeId,
322 ) -> &mut <LatticeGraph<N, E, S> as Data>::NodeWeight {
323 let offset = self.s.to_offset_unchecked(id);
324 self.nodes
326 .get_mut((offset.horizontal, offset.vertical))
327 .unwrap_unchecked()
328 }
329
330 #[doc(hidden)]
331 pub unsafe fn edge_weight_mut_unchecked(
332 &mut self,
333 id: <LatticeGraph<N, E, S> as GraphBase>::EdgeId,
334 ) -> &mut <LatticeGraph<N, E, S> as Data>::EdgeWeight {
335 let offset = self.s.to_offset_unchecked(id.0);
336 let ax = id.1.to_index();
337 self.edges
338 .get_unchecked_mut(ax)
339 .get_mut((offset.horizontal, offset.vertical))
340 .unwrap_unchecked()
341 }
342}
343
344#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
346pub struct EdgeTypeWrap<A>(PhantomData<A>);
347impl<A: Axis> EdgeType for EdgeTypeWrap<A> {
348 fn is_directed() -> bool {
349 A::DIRECTED
350 }
351}
352
353impl<N, E, S: Shape> GraphProp for LatticeGraph<N, E, S> {
354 type EdgeType = EdgeTypeWrap<S::Axis>;
355}
356
357#[derive(Debug, Clone, Hash, PartialEq, Eq)]
359pub struct VisMap<S> {
360 v: Vec<FixedBitSet>,
361 s: S,
362}
363
364impl<S: Shape> VisMap<S> {
365 pub(crate) fn new(s: S) -> Self {
366 let h = s.horizontal();
367 let v = s.vertical();
368 let mut vec = Vec::with_capacity(h);
369 for _ in 0..h {
370 vec.push(FixedBitSet::with_capacity(v));
371 }
372 Self { v: vec, s }
373 }
374}
375
376impl<S: Shape> VisitMap<S::Coordinate> for VisMap<S> {
377 fn visit(&mut self, a: S::Coordinate) -> bool {
378 let offset = self.s.to_offset(a);
379 if let Ok(a) = offset {
380 !self.v[a.horizontal].put(a.vertical)
381 } else {
382 false
383 }
384 }
385
386 fn is_visited(&self, a: &S::Coordinate) -> bool {
387 let offset = self.s.to_offset(*a);
388 if let Ok(a) = offset {
389 self.v[a.horizontal].contains(a.vertical)
390 } else {
391 false
392 }
393 }
394
395 fn unvisit(&mut self, a: S::Coordinate) -> bool {
396 let offset = self.s.to_offset(a);
397 if let Ok(offset) = offset {
398 self.v[offset.horizontal].set(offset.vertical, false);
399 true
400 } else {
401 false
402 }
403 }
404}
405
406impl<N, E, S: Shape + Clone> Visitable for LatticeGraph<N, E, S> {
407 type Map = VisMap<S>;
408
409 fn visit_map(&self) -> Self::Map {
410 VisMap::new(self.s.clone())
411 }
412
413 fn reset_map(&self, map: &mut Self::Map) {
414 map.v.iter_mut().for_each(|x| x.clear())
415 }
416}