use crate::core::{Graph, IgraphError, IgraphResult, VertexId};
pub fn square_lattice(
dim: &[u32],
nei: u32,
directed: bool,
mutual: bool,
periodic: Option<&[bool]>,
) -> IgraphResult<Graph> {
if nei >= 2 {
return Err(IgraphError::InvalidArgument(format!(
"square_lattice: nei={nei} not yet supported (only nei in {{0, 1}}); \
higher radii require igraph_connect_neighborhood (future AWU)"
)));
}
if let Some(p) = periodic {
if p.len() != dim.len() {
return Err(IgraphError::InvalidArgument(format!(
"square_lattice: periodic vector length {} must match dim length {}",
p.len(),
dim.len()
)));
}
}
let dims = dim.len();
let mut vcount: u32 = 1;
for &d in dim {
vcount = vcount.checked_mul(d).ok_or_else(|| {
IgraphError::InvalidArgument(format!(
"square_lattice: vertex count Π dim overflows u32 for dim={dim:?}"
))
})?;
}
if vcount == 0 {
return Graph::new(0, directed);
}
let mut weights: Vec<u32> = Vec::with_capacity(dims);
let mut acc: u32 = 1;
for &d in dim {
weights.push(acc);
acc = acc.checked_mul(d).ok_or_else(|| {
IgraphError::InvalidArgument(format!("square_lattice: weight overflow for dim={dim:?}"))
})?;
}
let mut coords: Vec<u32> = vec![0; dims];
let mut edges: Vec<(VertexId, VertexId)> = Vec::new();
let is_periodic = |idx: usize| -> bool { periodic.is_some_and(|p| p[idx]) };
for i in 0..vcount {
for j in 0..dims {
let dj = dim[j];
let pj = is_periodic(j);
let cj = coords[j];
if pj || cj != dj - 1 {
let new_nei: u32 = if cj == dj - 1 {
i - (dj - 1) * weights[j]
} else {
i + weights[j]
};
if new_nei != i && (dj != 2 || cj != 1 || directed) {
edges.push((i, new_nei));
}
}
if mutual && directed && (pj || cj != 0) {
let new_nei: u32 = if cj != 0 {
i - weights[j]
} else {
i + (dj - 1) * weights[j]
};
if new_nei != i && (dj != 2 || !pj) {
edges.push((i, new_nei));
}
}
}
for pos in 0..dims {
if coords[pos] != dim[pos] - 1 {
coords[pos] += 1;
break;
}
coords[pos] = 0;
}
}
let mut g = Graph::new(vcount, directed)?;
g.add_edges(edges)?;
Ok(g)
}
#[cfg(test)]
mod tests {
use super::*;
fn dump_edges(g: &Graph) -> Vec<(u32, u32)> {
let m = u32::try_from(g.ecount()).expect("ecount fits u32 in tests");
(0..m)
.map(|e| g.edge(e).expect("edge id in bounds"))
.collect()
}
fn degree_of(g: &Graph, v: VertexId) -> usize {
g.degree(v).expect("vertex in range")
}
#[test]
fn zero_dim_is_singleton() {
let g = square_lattice(&[], 1, false, false, None).expect("0-dim");
assert_eq!(g.vcount(), 1);
assert_eq!(g.ecount(), 0);
}
#[test]
fn dim_one_is_singleton() {
let g = square_lattice(&[1], 1, false, false, None).expect("dim=[1]");
assert_eq!(g.vcount(), 1);
assert_eq!(g.ecount(), 0);
}
#[test]
fn zero_in_dim_collapses_to_null() {
let g = square_lattice(&[0, 3], 1, false, false, None).expect("null");
assert_eq!(g.vcount(), 0);
assert_eq!(g.ecount(), 0);
}
#[test]
fn one_d_path() {
let g = square_lattice(&[3], 1, false, false, None).expect("P_3");
assert_eq!(g.vcount(), 3);
assert_eq!(g.ecount(), 2);
assert_eq!(dump_edges(&g), vec![(0, 1), (1, 2)]);
}
#[test]
fn one_d_periodic_is_cycle() {
let g = square_lattice(&[3], 1, false, false, Some(&[true])).expect("C_3");
assert_eq!(g.vcount(), 3);
assert_eq!(g.ecount(), 3);
assert_eq!(dump_edges(&g), vec![(0, 1), (1, 2), (0, 2)]);
}
#[test]
fn two_d_2x2_is_four_cycle() {
let g = square_lattice(&[2, 2], 1, false, false, None).expect("2x2");
assert_eq!(g.vcount(), 4);
assert_eq!(g.ecount(), 4);
assert_eq!(dump_edges(&g), vec![(0, 1), (0, 2), (1, 3), (2, 3)]);
}
#[test]
fn two_d_3x3_grid() {
let g = square_lattice(&[3, 3], 1, false, false, None).expect("3x3");
assert_eq!(g.vcount(), 9);
assert_eq!(g.ecount(), 12);
assert_eq!(degree_of(&g, 0), 2);
assert_eq!(degree_of(&g, 1), 3);
assert_eq!(degree_of(&g, 4), 4);
assert_eq!(degree_of(&g, 8), 2);
}
#[test]
fn two_d_3x3_torus() {
let g = square_lattice(&[3, 3], 1, false, false, Some(&[true, true])).expect("torus");
assert_eq!(g.vcount(), 9);
assert_eq!(g.ecount(), 18);
for v in 0..g.vcount() {
assert_eq!(degree_of(&g, v), 4, "torus vertex {v} should be 4-regular");
}
}
#[test]
fn two_d_2x2_torus_does_not_double_edge() {
let g = square_lattice(&[2, 2], 1, false, false, Some(&[true, true])).expect("torus 2x2");
assert_eq!(g.vcount(), 4);
assert_eq!(g.ecount(), 4);
assert_eq!(dump_edges(&g), vec![(0, 1), (0, 2), (1, 3), (2, 3)]);
}
#[test]
fn three_d_2x2x2_is_cube() {
let g = square_lattice(&[2, 2, 2], 1, false, false, None).expect("Q_3");
assert_eq!(g.vcount(), 8);
assert_eq!(g.ecount(), 12);
for v in 0..g.vcount() {
assert_eq!(degree_of(&g, v), 3, "Q_3 vertex {v} should be 3-regular");
}
}
#[test]
fn directed_low_to_high_only() {
let g = square_lattice(&[3], 1, true, false, None).expect("dir P_3");
assert!(g.is_directed());
assert_eq!(g.ecount(), 2);
assert_eq!(dump_edges(&g), vec![(0, 1), (1, 2)]);
}
#[test]
fn directed_mutual_emits_both_arcs() {
let g = square_lattice(&[3], 1, true, true, None).expect("dir mut P_3");
assert!(g.is_directed());
assert_eq!(g.ecount(), 4);
let edges = dump_edges(&g);
assert!(edges.contains(&(0, 1)));
assert!(edges.contains(&(1, 0)));
assert!(edges.contains(&(1, 2)));
assert!(edges.contains(&(2, 1)));
}
#[test]
fn directed_mutual_periodic_3cycle() {
let g = square_lattice(&[3], 1, true, true, Some(&[true])).expect("dir mut C_3");
assert!(g.is_directed());
assert_eq!(g.ecount(), 6);
}
#[test]
fn mixed_periodicity() {
let g =
square_lattice(&[3, 3], 1, false, false, Some(&[true, false])).expect("mixed periodic");
assert_eq!(g.vcount(), 9);
assert_eq!(g.ecount(), 15);
}
#[test]
fn nei_two_or_more_rejected() {
let err = square_lattice(&[3, 3], 2, false, false, None).expect_err("nei=2");
match err {
IgraphError::InvalidArgument(msg) => assert!(msg.contains("nei")),
other => panic!("unexpected error: {other:?}"),
}
}
#[test]
fn periodic_length_mismatch_rejected() {
let err = square_lattice(&[3, 3], 1, false, false, Some(&[true])).expect_err("mismatch");
match err {
IgraphError::InvalidArgument(msg) => assert!(msg.contains("periodic")),
other => panic!("unexpected error: {other:?}"),
}
}
#[test]
fn vcount_overflow_rejected() {
let err = square_lattice(&[65536, 65536], 1, false, false, None).expect_err("overflow");
match err {
IgraphError::InvalidArgument(msg) => assert!(msg.contains("overflow")),
other => panic!("unexpected error: {other:?}"),
}
}
}
#[cfg(all(test, feature = "proptest-harness"))]
mod proptests {
use super::*;
use proptest::prelude::*;
proptest! {
#[test]
fn two_d_grid_edge_count(a in 1u32..=6, b in 1u32..=6) {
let g = square_lattice(&[a, b], 1, false, false, None).unwrap();
let expected = a * (b - 1) + b * (a - 1);
prop_assert_eq!(g.vcount(), a * b);
prop_assert_eq!(g.ecount() as u32, expected);
}
#[test]
fn two_d_torus_is_four_regular(a in 3u32..=6, b in 3u32..=6) {
let g = square_lattice(
&[a, b], 1, false, false, Some(&[true, true]),
).unwrap();
for v in 0..g.vcount() {
prop_assert_eq!(g.degree(v).unwrap(), 4);
}
}
#[test]
fn power_of_two_lattice_is_hypercube(k in 0u32..=5) {
use crate::hypercube;
let dims = vec![2u32; k as usize];
let g = square_lattice(&dims, 1, false, false, None).unwrap();
let q = hypercube(k, false).unwrap();
prop_assert_eq!(g.vcount(), q.vcount());
prop_assert_eq!(g.ecount(), q.ecount());
}
#[test]
fn directed_mutual_doubles_undirected(a in 2u32..=5, b in 2u32..=5) {
let u = square_lattice(&[a, b], 1, false, false, None).unwrap();
let dm = square_lattice(&[a, b], 1, true, true, None).unwrap();
prop_assert_eq!(dm.ecount(), u.ecount() * 2);
}
}
}