lattice_graph/lattice_abstract/
mod.rs

1//! Module for Abstract 2D Lattice Graph. It is used inside by other lattice graph in other modules like [`hex`](`crate::hex`).
2//! Use it when you want to define your own lattice graph, or to use the concreate visit iterator structs for traits in [`visit`](`petgraph::visit`).
3
4use 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)]
25/// Abstract Lattice Graph.
26/// It holds the node and edge weight data.
27/// The actural behaviour is dependent on [`Shape`](`shapes::Shape`).
28pub 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    /// Creates a graph from raw data. This api might change.
36    #[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    /// Creates a graph with uninitalized node and edge weight data.
42    /// It is extremely unsafe so should use with [`MaybeUninit`](`core::mem::MaybeUninit`) and use [`assume_init`](`Self::assume_init`).
43    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    /// Creates a graph with node and edge weight data set to [`default`](`Default::default`).
55    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    /// Creates a graph with node and edge weight data from the coordinate.
64    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, // change to E ?
68    {
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    /// Get a reference to the lattice graph's s.
93    pub fn shape(&self) -> &S {
94        &self.s
95    }
96}
97
98impl<N, E, S: Shape + Default> LatticeGraph<N, E, S> {
99    /// Creates a graph with node and edge weight data set to [`default`](`Default::default`) with [`Shape`] from default.
100    pub fn new_s() -> Self
101    where
102        N: Default,
103        E: Default,
104    {
105        Self::new(S::default())
106    }
107
108    /// Creates a graph with uninitalized node and edge weight data with [`Shape`] from default.
109    /// It is extremely unsafe so should use with [`MaybeUninit`](`core::mem::MaybeUninit`) and use [`assume_init`](`Self::assume_init`).
110    pub unsafe fn new_uninit_s() -> LatticeGraph<MaybeUninit<N>, MaybeUninit<E>, S> {
111        Self::new_uninit(S::default())
112    }
113
114    /// Creates a graph with node and edge weight data from the coordinate with [`Shape`] from default.
115    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    /**
126    Assume the underlying nodes and edges to be initialized.
127    ```
128    # use lattice_graph::hex::axial_based::*;
129    # use core::mem::MaybeUninit;
130    # use petgraph::data::*;
131    let mut hex = unsafe { HexGraphConst::<f32, (), OddR, 5, 5>::new_uninit_s() };
132    for i in 0..5{
133        for j in 0..5{
134            let offset = Offset::new(i, j);
135            let coord = hex.shape().offset_to_coordinate(offset);
136            if let Some(ref mut n) = hex.node_weight_mut(coord){
137                **n = MaybeUninit::new((i + j) as f32);
138            }
139        }
140    }
141    let hex_init = unsafe{ hex.assume_init() };
142    ```
143    */
144    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 e is drop type, drop manually to prevent dropping for invalid (uninitialized) edges.
160        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        // SAFETY : offset must be checked in `to_offset`
203        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        // SAFETY : offset must be checked in `to_offset`
241        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        // SAFETY : offset must be checked in `to_offset`
283        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        // SAFETY : offset must be checked in `to_offset`
325        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///Wrapper for [`Axis`] to be [`EdgeType`].
345#[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/// [`VisitMap`] of [`LatticeGraph`].
358#[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}