flense 0.4.0

Purpose-oriented lensing
Documentation
//! Tests for lensing `ndarray` views at arity 1 (slice), 2 (rect), and 3
//! (cube), both shared and mutable.

#![cfg(feature = "ndarray")]
#![allow(
    clippy::float_cmp,
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation
)]

use assert2::assert;
use flense::prelude::*;
use ndarray::{
    Array1,
    Array2,
    Array3,
    s,
};

use super::*;

/// Builds an `Array1<Vertex>` whose every field is a distinct function of the
/// index, so a mis-mapped stride is observable.
fn line(len: usize) -> Array1<Vertex> {
    Array1::from_shape_fn(len, |i| Vertex {
        position: [i as f32, 0.0, 0.0],
        color: [(i + 1) as f32, 0.0, 0.0],
        normal: [i as u8, 0],
        other: i as u64,
    })
}

/// Builds an `Array2<Vertex>` whose every field is a distinct function of the
/// `(row, col)` index, so a mis-mapped stride or axis is observable.
fn grid(rows: usize, cols: usize) -> Array2<Vertex> {
    Array2::from_shape_fn((rows, cols), |(i, j)| Vertex {
        position: [i as f32, j as f32, (i * cols + j) as f32],
        color: [(i + 1) as f32, (j + 1) as f32, 0.0],
        normal: [i as u8, j as u8],
        other: (i as u64) << 32 | j as u64,
    })
}

/// Builds an `Array3<Vertex>` whose every field is a distinct function of the
/// `(i, j, k)` index.
fn cube(d0: usize, d1: usize, d2: usize) -> Array3<Vertex> {
    Array3::from_shape_fn((d0, d1, d2), |(i, j, k)| Vertex {
        position: [i as f32, j as f32, k as f32],
        color: [(i + 1) as f32, (j + 1) as f32, (k + 1) as f32],
        normal: [i as u8, j as u8],
        other: (i as u64) << 32 | (j as u64) << 16 | k as u64,
    })
}

// --- arity 1: slice ---------------------------------------------------------

#[test]
fn ndarray_slice_indexes_contiguous() {
    let arr = line(5);
    let view = arr.view();
    let lens: LensSlice<'_, (Position, Other)> = view.lens_slice();

    assert!(lens.len() == 5);
    for i in 0..5 {
        let v = &arr[i];
        assert!(lens.get::<Position, _>(i).unwrap() == &v.position);
        assert!(lens.get::<Other, _>(i).unwrap() == &v.other);
    }
    assert!(lens.get::<Position, _>(5).is_none());
}

#[test]
fn ndarray_slice_handles_strided_view() {
    let arr = line(6);
    // Every other element: length 3 with a stride of 2.
    let sub = arr.slice(s![..;2]);
    let lens: LensSlice<'_, (Position, Other)> = sub.view().lens_slice();

    assert!(lens.len() == 3);
    for i in 0..3 {
        let v = &arr[i * 2];
        assert!(lens.get::<Position, _>(i).unwrap() == &v.position);
        assert!(lens.get::<Other, _>(i).unwrap() == &v.other);
    }
}

#[test]
fn ndarray_slice_mut_writes() {
    let mut arr = line(4);
    {
        let mut lens: LensSliceMut<'_, (Position, Other)> = arr.view_mut().lens_slice_mut();
        for i in 0..4 {
            *lens.get_mut::<Position, _>(i).unwrap() = [10.0 + i as f32, 0.0, 0.0];
            *lens.get_mut::<Other, _>(i).unwrap() = 0x1000 + i as u64;
        }
        assert!(lens.get_mut::<Position, _>(4).is_none());
    }
    for i in 0..4 {
        assert!(arr[i].position == [10.0 + i as f32, 0.0, 0.0]);
        assert!(arr[i].other == 0x1000 + i as u64);
    }
}

// --- arity 2: rect ----------------------------------------------------------

#[test]
fn ndarray_rect_indexes_contiguous() {
    let rows = 2;
    let cols = 3;
    let arr = grid(rows, cols);
    let view = arr.view();
    let lens: LensRect<'_, (Position, Color, Normal, Other)> = view.lens_rect();

    for i in 0..rows {
        for j in 0..cols {
            let v = &arr[[i, j]];
            assert!(lens.get::<Position, _>([i, j]).unwrap() == &v.position);
            assert!(lens.get::<Color, _>([i, j]).unwrap() == &v.color);
            assert!(lens.get::<Normal, _>([i, j]).unwrap() == &v.normal);
            assert!(lens.get::<Other, _>([i, j]).unwrap() == &v.other);
        }
    }
}

#[test]
fn ndarray_rect_out_of_bounds_is_none() {
    let arr = grid(2, 3);
    let view = arr.view();
    let lens: LensRect<'_, (Position, Other)> = view.lens_rect();

    assert!(lens.get::<Position, _>([2, 0]).is_none());
    assert!(lens.get::<Position, _>([0, 3]).is_none());
    assert!(lens.get::<Other, _>([2, 3]).is_none());
}

#[test]
fn ndarray_rect_handles_transposed_view() {
    let arr = grid(4, 5);
    // Transpose: shape (5, 4) with non-unit strides on both axes.
    let view = arr.view().reversed_axes();
    let lens: LensRect<'_, (Position, Other)> = view.lens_rect();

    for r in 0..5 {
        for c in 0..4 {
            // view[r, c] == arr[c, r]
            let v = &arr[[c, r]];
            assert!(lens.get::<Position, _>([r, c]).unwrap() == &v.position);
            assert!(lens.get::<Other, _>([r, c]).unwrap() == &v.other);
        }
    }
}

#[test]
fn ndarray_rect_handles_strided_view() {
    let arr = grid(4, 6);
    // Every other column: shape (4, 3) with a column stride of 2.
    let sub = arr.slice(s![.., ..;2]);
    let lens: LensRect<'_, (Position, Other)> = sub.view().lens_rect();

    for i in 0..4 {
        for j in 0..3 {
            let v = &arr[[i, j * 2]];
            assert!(lens.get::<Position, _>([i, j]).unwrap() == &v.position);
            assert!(lens.get::<Other, _>([i, j]).unwrap() == &v.other);
        }
    }
}

#[test]
fn ndarray_rect_mut_reads_contiguous() {
    let rows = 2;
    let cols = 3;
    let arr = grid(rows, cols);
    let mut view = arr.clone();
    let lens: LensRectMut<'_, (Position, Color, Normal, Other)> = view.view_mut().lens_rect_mut();

    for i in 0..rows {
        for j in 0..cols {
            let v = &arr[[i, j]];
            assert!(lens.get::<Position, _>([i, j]).unwrap() == &v.position);
            assert!(lens.get::<Color, _>([i, j]).unwrap() == &v.color);
            assert!(lens.get::<Normal, _>([i, j]).unwrap() == &v.normal);
            assert!(lens.get::<Other, _>([i, j]).unwrap() == &v.other);
        }
    }
}

#[test]
fn ndarray_rect_mut_writes_contiguous() {
    let rows = 2;
    let cols = 3;
    let mut arr = grid(rows, cols);
    {
        let mut lens: LensRectMut<'_, (Position, Other)> = arr.view_mut().lens_rect_mut();
        for i in 0..rows {
            for j in 0..cols {
                *lens.get_mut::<Position, _>([i, j]).unwrap() = [10.0 + i as f32, j as f32, 0.0];
                *lens.get_mut::<Other, _>([i, j]).unwrap() = 0x1000 + (i * cols + j) as u64;
                // A shared accessor observes the just-written value.
                assert!(
                    lens.get::<Position, _>([i, j]).unwrap() == &[10.0 + i as f32, j as f32, 0.0]
                );
            }
        }
    }
    for i in 0..rows {
        for j in 0..cols {
            assert!(arr[[i, j]].position == [10.0 + i as f32, j as f32, 0.0]);
            assert!(arr[[i, j]].other == 0x1000 + (i * cols + j) as u64);
        }
    }
}

#[test]
fn ndarray_rect_mut_out_of_bounds_is_none() {
    let mut arr = grid(2, 3);
    let mut lens: LensRectMut<'_, (Position, Other)> = arr.view_mut().lens_rect_mut();

    assert!(lens.get_mut::<Position, _>([2, 0]).is_none());
    assert!(lens.get_mut::<Position, _>([0, 3]).is_none());
    assert!(lens.get_mut::<Other, _>([2, 3]).is_none());
}

#[test]
fn ndarray_rect_mut_handles_transposed_view() {
    let arr = grid(4, 5);
    let mut target = arr.clone();
    {
        // Transpose: shape (5, 4) with non-unit strides on both axes.
        let mut view = target.view_mut().reversed_axes();
        let mut lens: LensRectMut<'_, (Position, Other)> = view.view_mut().lens_rect_mut();
        for r in 0..5 {
            for c in 0..4 {
                *lens.get_mut::<Other, _>([r, c]).unwrap() = 0xBEEF;
            }
        }
        // view[r, c] == target[c, r]
        for r in 0..5 {
            for c in 0..4 {
                let v = &arr[[c, r]];
                assert!(lens.get::<Position, _>([r, c]).unwrap() == &v.position);
            }
        }
    }
    for v in &target {
        assert!(v.other == 0xBEEF);
    }
}

#[test]
fn ndarray_rect_mut_handles_strided_view() {
    let arr = grid(4, 6);
    let mut target = arr.clone();
    {
        // Every other column: shape (4, 3) with a column stride of 2.
        let mut sub = target.slice_mut(s![.., ..;2]);
        let mut lens: LensRectMut<'_, (Position, Other)> = sub.view_mut().lens_rect_mut();
        for i in 0..4 {
            for j in 0..3 {
                let v = &arr[[i, j * 2]];
                assert!(lens.get::<Position, _>([i, j]).unwrap() == &v.position);
                *lens.get_mut::<Other, _>([i, j]).unwrap() = 0xCAFE;
            }
        }
    }
    for i in 0..4 {
        for j in 0..6 {
            if j % 2 == 0 {
                assert!(target[[i, j]].other == 0xCAFE);
            } else {
                assert!(target[[i, j]].other == arr[[i, j]].other);
            }
        }
    }
}

#[test]
fn ndarray_rect_mut_get_all_mut() {
    let mut arr = grid(2, 2);
    {
        let mut lens: LensRectMut<'_, (Position, Color, Normal, Other)> =
            arr.view_mut().lens_rect_mut();
        for i in 0..2 {
            for j in 0..2 {
                let mut l = lens.get_all_mut([i, j]).unwrap();
                *l.as_mut::<Position, _>() = [100.0 + i as f32, j as f32, 0.0];
                *l.as_mut::<Other, _>() = 0x2000 + (i * 2 + j) as u64;
            }
        }
        assert!(lens.get_all_mut([2, 0]).is_none());
    }
    for i in 0..2 {
        for j in 0..2 {
            assert!(arr[[i, j]].position == [100.0 + i as f32, j as f32, 0.0]);
            assert!(arr[[i, j]].other == 0x2000 + (i * 2 + j) as u64);
        }
    }
}

// --- arity 3: cube ----------------------------------------------------------

#[test]
fn ndarray_cube_indexes_contiguous() {
    let (d0, d1, d2) = (2, 3, 4);
    let arr = cube(d0, d1, d2);
    let view = arr.view();
    let lens: LensCube<'_, (Position, Color, Other)> = view.lens_cube();

    for i in 0..d0 {
        for j in 0..d1 {
            for k in 0..d2 {
                let v = &arr[[i, j, k]];
                assert!(lens.get::<Position, _>([i, j, k]).unwrap() == &v.position);
                assert!(lens.get::<Color, _>([i, j, k]).unwrap() == &v.color);
                assert!(lens.get::<Other, _>([i, j, k]).unwrap() == &v.other);
            }
        }
    }
    assert!(lens.get::<Position, _>([d0, 0, 0]).is_none());
    assert!(lens.get::<Position, _>([0, d1, 0]).is_none());
    assert!(lens.get::<Position, _>([0, 0, d2]).is_none());
}

#[test]
fn ndarray_cube_handles_permuted_view() {
    let arr = cube(2, 3, 4);
    // Permute axes: shape (4, 3, 2) with non-unit strides on every axis.
    let view = arr.view().permuted_axes([2, 1, 0]);
    let lens: LensCube<'_, (Position, Other)> = view.lens_cube();

    for i in 0..4 {
        for j in 0..3 {
            for k in 0..2 {
                // view[i, j, k] == arr[k, j, i]
                let v = &arr[[k, j, i]];
                assert!(lens.get::<Position, _>([i, j, k]).unwrap() == &v.position);
                assert!(lens.get::<Other, _>([i, j, k]).unwrap() == &v.other);
            }
        }
    }
}

#[test]
fn ndarray_cube_mut_writes() {
    let (d0, d1, d2) = (2, 2, 3);
    let mut arr = cube(d0, d1, d2);
    {
        let mut lens: LensCubeMut<'_, (Position, Other)> = arr.view_mut().lens_cube_mut();
        for i in 0..d0 {
            for j in 0..d1 {
                for k in 0..d2 {
                    *lens.get_mut::<Position, _>([i, j, k]).unwrap() =
                        [i as f32, j as f32, k as f32 + 100.0];
                    *lens.get_mut::<Other, _>([i, j, k]).unwrap() = 0x2000 + (i + j + k) as u64;
                }
            }
        }
        assert!(lens.get_mut::<Position, _>([d0, 0, 0]).is_none());
    }
    for i in 0..d0 {
        for j in 0..d1 {
            for k in 0..d2 {
                assert!(arr[[i, j, k]].position == [i as f32, j as f32, k as f32 + 100.0]);
                assert!(arr[[i, j, k]].other == 0x2000 + (i + j + k) as u64);
            }
        }
    }
}

#[test]
fn ndarray_cube_mut_handles_strided_view() {
    let arr = cube(4, 3, 4);
    let mut target = arr.clone();
    {
        // Every other slab on axis 0: shape (2, 3, 4) with a stride of 2.
        let mut sub = target.slice_mut(s![..;2, .., ..]);
        let mut lens: LensCubeMut<'_, (Position, Other)> = sub.view_mut().lens_cube_mut();
        for i in 0..2 {
            for j in 0..3 {
                for k in 0..4 {
                    let v = &arr[[i * 2, j, k]];
                    assert!(lens.get::<Position, _>([i, j, k]).unwrap() == &v.position);
                    *lens.get_mut::<Other, _>([i, j, k]).unwrap() = 0xCAFE;
                }
            }
        }
    }
    for i in 0..4 {
        for j in 0..3 {
            for k in 0..4 {
                if i % 2 == 0 {
                    assert!(target[[i, j, k]].other == 0xCAFE);
                } else {
                    assert!(target[[i, j, k]].other == arr[[i, j, k]].other);
                }
            }
        }
    }
}