tensorism 0.5.0

//! Tensorism is a library built on top of `ndarray` that provides a domain-specific language (DSL)
//! for expressing tensor computations using named indexes.
//!
//! The main goal of tensorism is to make multi-index array expressions explicit, readable and
//! compositional, while remaining compatible with the `ndarray` ecosystem.
//!
//! At the moment, tensorism focuses on expressiveness and correctness. The evaluation strategy
//! is primarily iterator-based, and no guarantee is made regarding optimal performance. The DSL
//! and its internal translation strategy are considered experimental and may evolve in
//! non-backward-compatible ways in future versions.
//!
//! # Indexed expressions and the `for` construct
//!
//! Tensorism introduces a special construct of the form `for ⟨indexes⟩ => ⟨body⟩`, where
//! `⟨indexes⟩` is a space-separated list of Rust identifiers.
//!
//! These identifiers are called **[Ricci indexes](https://en.wikipedia.org/wiki/Gregorio_Ricci-Curbastro)**
//! and may be any valid Rust identifier. Names such as `i`, `j` or `k` are used throughout the
//! documentation purely as conventional examples.
//!
//! A Ricci index represents a variable of type `usize`. Its range of values is determined by the
//! arrays it indexes inside the associated expression. When several arrays use the same index,
//! their corresponding dimensions must agree, otherwise the program panics at runtime.
//!
//! Example:
//!
//! ```ignore
//! let x: Array2<u8> = new_ndarray! { for i j => a[i, j] + b[j] };
//! ```
//!
//! Here, index `i` ranges over the first dimension of `a`, while index `j` ranges over the second
//! dimension of `a` and the only dimension of `b`, which must be compatible.
//!
//! # Generating new arrays
//!
//! When a `for ⟨indexes⟩ => ⟨body⟩` construct appears at the top level of a `new_ndarray!`
//! invocation, it generates a new `ndarray::Array` whose number of dimensions matches the number
//! of indexes.
//!
//! For instance, the previous example produces an `Array2<T>`, where `T` is the type of the
//! expression `a[i, j] + b[j]`.
//!
//! More complex expressions involving conditionals and multiple input arrays are also supported:
//!
//! ```ignore
//! let y = new_ndarray! {
//!     for i j k =>
//!         if p[i, j] - 0.3 < 0.4 * q[j, k] {
//!             r[j] * q[j, k] + 0.2
//!         } else {
//!             0.5 * s[i, j, k]
//!         }
//! };
//! ```
//!
//! # Iterators and aggregation
//!
//! A `for ⟨indexes⟩ => ⟨body⟩` construct does not have to appear at the top level.
//! When used as a sub-expression, it evaluates to a Rust iterator over the successive values
//! of ⟨body⟩ as the indexes vary.
//!
//! This makes it possible to express aggregations by passing such iterators to user-defined
//! functions or standard iterator consumers such as `sum`, `min`, or `fold`:
//!
//! ```ignore
//! let x: i64 = new_ndarray! {
//!     Iterator::sum(for i => Iterator::min(for j => a[i, j]).unwrap())
//! };
//! ```
//!
//! Any function or method that consumes an iterator of the appropriate item type can be used.
//!
//! # Reindexing
//!
//! In addition to direct indexing, tensorism provides explicit reindexing objects such as
//! `Reindexing1`, `Reindexing2`, and higher-order variants. These objects represent immutable,
//! bounded index transformations and can be safely composed inside indexed expressions.
//!
//! Reindexing allows expressing indirect access patterns while preserving runtime guarantees
//! on index validity.
//!
//! # Experimental status
//!
//! Tensorism is currently experimental. In particular, the iterator-based evaluation model and
//! the interaction with `ndarray`'s high-level APIs may change in future releases in order to
//! enable more aggressive optimizations.
extern crate rand;
extern crate tensorism_gen;

use rand::Rng;
use std::{alloc::Layout, fmt::Display, ops::Index, slice};

/**
 * `new_ndarray!` is the main macro of the Tensorism crate for constructing arrays from Rust
 * expressions using the Tensorism DSL.
 *
 * # Basic principles of `new_ndarray!` syntax
 *
 * The macro accepts any Rust expression as argument. Within the argument, the `for ⟨indexes⟩ => ⟨body⟩`
 * structure can appear multiple times, and it can be either at the **top-level** of the macro
 * or **nested** inside another expression.
 *
 * ## Top-level `for ⟨indexes⟩ => ⟨body⟩`
 *
 * When `for ⟨indexes⟩ => ⟨body⟩` appears at the very start of the macro argument, it generates
 * a multi-dimensional array.
 *
 * ```ignore
 * let x = new_ndarray! { for i j => a[i, j] + b[j] };
 * ```
 *
 * Here, the top-level `for i j =>` defines the axes of the resulting array.
 * This produces an `Array2` whose dimensions are inferred from the arrays used
 * in the expression.
 *
 * ## Nested `for ⟨indexes⟩ => ⟨body⟩` inside expressions
 *
 * A `for ⟨indexes⟩ => ⟨body⟩` can appear inside another expression, producing an iterator
 * rather than a full array:
 *
 * ```ignore
 * let nested_sum: u32 = new_ndarray! { Iterator::sum(for i => a[i] * 2) };
 * ```
 *
 * ```ignore
 * let similar = new_ndarray! {
 *     compare_iterators(
 *         for i => a[i] + i,
 *         for i j => b[i, j] - c[j]
 *     )
 * };
 * ```
 *
 * Nested usage can also appear inside other iterators:
 *
 * ```ignore
 * let nested_sum: f64 = new_ndarray! {
 *     Iterator::sum(for j => Iterator::min(for i => m[i, j]))
 * };
 * ```
 *
 * ```ignore
 * let vector = new_ndarray! {
 *     Iterator::collect::<Vec<f64>>(for i => sqrt(median(for j => a[j] * a[i]) + 1.0))
 * };
 * ```
 *
 * ## Ricci Indexes
 *
 * Tensorism expressions rely on **Ricci indexes** to represent iteration variables
 * over array dimensions. Any valid Rust identifier can serve as a Ricci index.
 * Each Ricci index always has type `usize` and takes values from `0` up to
 * the maximum allowed by the arrays it indexes. Inside the `⟨body⟩`,
 * Ricci indexes are used to access arrays using syntax `array[i1, i2, ...]`. Note
 * that no tuple is used when multiple indexes are involved;
 * commas are used directly inside the square brackets.
 *
 * Example:
 *
 * ```ignore
 * let x = new_ndarray! { for i j => a[i, j] + b[j] };
 * ```
 *
 * In this example:
 * - `i` iterates over the first dimension of `a`, from `0` to `a.dim().0 - 1`
 * - `j` iterates over the second dimension of `a`, from `0` to `a.dim().1 - 1`, which must
 *   also match `b.dim() - 1`.
 *
 * In case of dimension mismatches, the program panics at runtime.
 *
 * # Advanced features
 *
 * ## Indexers
 *
 * In addition to ordinary index variables, `tensorism` supports **indexers**, which are special
 * forms of indexation that control how an index value is produced or interpreted. Indexers
 * appear inside the brackets of an array access and refine the way an index expression is evaluated.
 * An indexer always produces a value of type `usize`, but it does so according to a specific
 * rule that is explicit in the syntax.
 *
 * Some indexers are built in:
 *
 * * `rev: i` iterates over the corresponding axis in reverse order. Conceptually, it maps
 *   the logical index `i` to `dim - 1 - i`, where `dim` is the size of the indexed axis.
 * * `plain: ⟨expr⟩` indexes an array using a `usize` expression that is independent of the
 *   surrounding iterations. The expression is evaluated once and reused, rather than being
 *   recomputed for each iteration.
 *
 * User-defined indexers are represented by values of types such as `Reindexing1`, `Reindexing2`,
 * etc. They are invoked using bracket syntax, for example `reindexing[i, j]`. Such indexers
 * can depend on one or more Ricci indices and may themselves be composed or nested. Their
 * defining property is that they are immutable and guarantee that the produced indices lie
 * within a statically known bound, which can be checked against the shape of the indexed array.
 *
 * Indexers make it possible to express non-trivial access patterns, including reordering,
 * indirection, and axis reversal, while remaining within the declarative structure of the Tensorism
 * DSL. Example:
 * ```ignore
 * let array = new_ndarray! {
 *     for i j => a[sorting_reindexing[i], rev: j] + b[j, plain: 3 * SIZE + 1]
 * };
 * ```
 * Note that indexers can be nested:
 * ```ignore
 * let array = new_ndarray! {
 *   for i j => a[reindexing1[reindexing2[rev: i, j]], i]
 * };
 * ```
 *
 * ## Aliases
 *
 * Within a `for ⟨indexes⟩ => ⟨body⟩` construct, Tensorism allows the introduction of
 * **aliases** using the `let` keyword. An alias binds a fresh index name to an index-producing
 * expression and makes it reusable inside the scope of the `⟨body⟩`. Example:
 * ```ignore
 * let array = new_ndarray! {
 *   for i j let k = reindexing2[rev: i, j] => x[k, j] + y[k]
 * };
 * ```
 * Several aliases may be introduced in sequence by inserting consecutive `let ⟨alias⟩ = ⟨indexers⟩`,
 * and each alias may depend on previously declared indices or aliases.
 *
 * ### Conditionals
 *
 * Nested `for ⟨indexes⟩ => ⟨body⟩` constructs may include a **conditional guard** introduced
 * by the `if` keyword. Example:
 * ```ignore
 * let conditional_sum = new_ndarray! {
 *   Iterator::sum(for i j if u[i] < v[i, j] => w[i, j] * u[i])
 * };
 * ```
 * A conditional restricts the set of index combinations for which the `⟨body⟩` is evaluated.
 * A conditional is only legal on **nested** level constructs. It is explicitly forbidden
 * at the top level. The reason is structural: a top-level `for ⟨indexes⟩ => ⟨body⟩` must always
 * produce a fully populated array, whereas a conditional may selectively discard index combinations and
 * therefore break rectangularity.
 *
 * Conditionals make it possible to express triangular domains, masked computations, and other index-dependent
 * restrictions, while preserving the invariant that only top-level constructions define the shape of the
 * resulting array.
 */
pub use tensorism_gen::new_ndarray;

/**
 * A type representing an immutable mapping from integer interval `0..input0_bound` to integer interval `0..output_bound`.
 *
 * Mathematically this encodes a function `0..input0_bound → 0..output_bound`.
 */
pub struct Reindexing1 {
    input0_bound: usize,
    output_bound: usize,
    ptr: *mut usize,
}

impl Reindexing1 {
    /// Creates a new `Reindexing1` from a function `f: usize -> usize` and bounds on the input and output intervals.
    /// The function `f` is evaluated at each integer from `0` to `input0_bound - 1`, and the results are reduced modulo `output_bound` to produce the reindexing.
    /// If `input0_bound` is zero, the resulting reindexing is empty and the function `f` is not evaluated at all. If `output_bound` is zero, the function panics.
    /// The resulting reindexing is immutable and guarantees that for each `i` in `0..input0_bound`, the value of `reindexing[i]` is in `0..output_bound`.
    pub fn new(
        input0_bound: usize,
        output_bound: usize,
        mut f: impl FnMut(usize) -> usize,
    ) -> Self {
        if input0_bound == 0 {
            return Self {
                input0_bound: 0,
                output_bound,
                ptr: std::ptr::null_mut(),
            };
        }
        if output_bound == 0 {
            panic!("output_bound must be greater than 0");
        }

        let layout = Layout::array::<usize>(input0_bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..input0_bound {
                ptr.add(i).write(f(i) % output_bound);
            }
        }
        Self {
            input0_bound,
            output_bound,
            ptr,
        }
    }

    /// Creates a random permutation reindexing of the interval `0..bound`.
    pub fn shuffle(bound: usize, rand: &mut impl Rng) -> Self {
        if bound == 0 {
            return Self {
                input0_bound: 0,
                output_bound: 0,
                ptr: std::ptr::null_mut(),
            };
        }
        let layout = Layout::array::<usize>(bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..bound {
                ptr.add(i).write(i);
            }
            for upper_index in (1..bound).rev() {
                let random_index = rand.random_range(0..=upper_index);
                if random_index != upper_index {
                    ptr.add(random_index).swap(ptr.add(upper_index));
                }
            }
        }
        Self {
            input0_bound: bound,
            output_bound: bound,
            ptr,
        }
    }

    /// Creates a random injective reindexing from `0..input0_bound` into `0..output_bound`.
    ///
    /// # Panics
    ///
    /// Panics if `input0_bound > output_bound`.
    pub fn shuffle_partially(
        input0_bound: usize,
        output_bound: usize,
        rand: &mut impl Rng,
    ) -> Self {
        if input0_bound > output_bound {
            panic!("input0_bound must be less than or equal to output_bound");
        }
        if input0_bound == 0 {
            return Self {
                input0_bound: 0,
                output_bound,
                ptr: std::ptr::null_mut(),
            };
        }
        let layout = Layout::array::<usize>(input0_bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..input0_bound {
                ptr.add(i).write(i);
            }
            for i in input0_bound..output_bound {
                let j = rand.random_range(0..=i);
                if j < input0_bound {
                    ptr.add(j).write(i);
                }
            }

            for i in (1..input0_bound).rev() {
                let j = rand.random_range(0..=i);
                ptr.add(j).swap(ptr.add(i));
            }
        }
        Self {
            input0_bound,
            output_bound,
            ptr,
        }
    }

    #[doc(hidden)]
    pub fn get_input0_bound(r: &Self) -> usize {
        r.input0_bound
    }
    #[doc(hidden)]
    pub fn get_output_bound(r: &Self) -> usize {
        r.output_bound
    }

    /// Returns the shape of the input domain as a 1-tuple.
    pub fn input_shape(&self) -> (usize,) {
        (self.input0_bound,)
    }
    /// Returns the bound of the output interval.
    pub fn output_bound(&self) -> usize {
        self.output_bound
    }

    #[doc(hidden)]
    pub unsafe fn get_unchecked(r: &Self, i: usize) -> usize {
        unsafe { r.ptr.add(i).read() }
    }
}

impl PartialEq for Reindexing1 {
    fn eq(&self, other: &Self) -> bool {
        if self.input0_bound != other.input0_bound || self.output_bound != other.output_bound {
            return false;
        }
        for i in 0..self.input0_bound {
            if unsafe { self.ptr.add(i).read() } != unsafe { other.ptr.add(i).read() } {
                return false;
            }
        }
        true
    }
}

impl Eq for Reindexing1 {}

impl Clone for Reindexing1 {
    fn clone(&self) -> Self {
        if self.ptr.is_null() {
            return Self {
                input0_bound: 0,
                output_bound: self.output_bound,
                ptr: std::ptr::null_mut(),
            };
        }
        let layout = Layout::array::<usize>(self.input0_bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..self.input0_bound {
                ptr.add(i).write(self.ptr.add(i).read());
            }
        }
        Self {
            input0_bound: self.input0_bound,
            output_bound: self.output_bound,
            ptr,
        }
    }
}

impl std::convert::AsRef<[usize]> for Reindexing1 {
    fn as_ref(&self) -> &[usize] {
        if self.ptr.is_null() {
            return &[];
        }
        unsafe { slice::from_raw_parts(self.ptr, self.input0_bound) }
    }
}

impl Index<usize> for Reindexing1 {
    type Output = usize;

    fn index(&self, index: usize) -> &Self::Output {
        if index >= self.input0_bound {
            panic!("Index out of bound.")
        }
        unsafe { self.ptr.add(index).as_ref().unwrap() }
    }
}

impl Drop for Reindexing1 {
    fn drop(&mut self) {
        if self.ptr.is_null() {
            return;
        }
        let layout = Layout::array::<usize>(self.input0_bound).unwrap();
        unsafe { std::alloc::dealloc(self.ptr as *mut u8, layout) }
    }
}

impl Display for Reindexing1 {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "[")?;
        for i in 0..self.input0_bound {
            if i != 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}↦{}", i, unsafe { self.ptr.add(i).read() })?;
        }
        write!(f, "]")
    }
}

unsafe impl Send for Reindexing1 {}
unsafe impl Sync for Reindexing1 {}

/**
 * A type representing an immutable mapping from a pair of  integer intervals `0..input0_bound` and `0..input1_bound` to  integer interval `0..output_bound`.
 *
 * Mathematically this encodes a function `0..input0_bound × 0..input1_bound → 0..output_bound`.
 */
pub struct Reindexing2 {
    input0_bound: usize,
    input1_bound: usize,
    output_bound: usize,
    ptr: *mut usize,
}

impl Reindexing2 {
    /// Creates a new `Reindexing2` from a function `f: (usize, usize) -> usize` and bounds on the input and output intervals.
    /// The function `f` is evaluated at each pair of integers from `0` to `input0_bound - 1` and `0` to `input1_bound - 1`, and the results are reduced modulo `output_bound` to produce the reindexing.
    /// If either `input0_bound` or `input1_bound` is zero, the resulting reindexing is empty and the function `f` is not evaluated at all. If `output_bound` is zero, the function panics.
    /// The resulting reindexing is immutable and guarantees that for each pair `(i0, i1)` in `0..input0_bound × 0..input1_bound`, the value of `reindexing[i0, i1]` is in `0..output_bound`.
    pub fn new(
        input0_bound: usize,
        input1_bound: usize,
        output_bound: usize,
        mut f: impl FnMut(usize, usize) -> usize,
    ) -> Self {
        if input0_bound == 0 || input1_bound == 0 {
            return Self {
                input0_bound: 0,
                input1_bound: 0,
                output_bound,
                ptr: std::ptr::null_mut(),
            };
        }
        if output_bound == 0 {
            panic!("output_bound must be greater than 0");
        }
        let bound = input0_bound
            .checked_mul(input1_bound)
            .expect("Saturated value");

        let layout = Layout::array::<usize>(bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            let mut i = 0;
            for i0 in 0..input0_bound {
                for i1 in 0..input1_bound {
                    ptr.add(i).write(f(i0, i1) % output_bound);
                    i += 1;
                }
            }
        }
        Self {
            input0_bound,
            input1_bound,
            output_bound,
            ptr,
        }
    }

    /// Creates a random permutation reindexing of the product domain `0..input0_bound × 0..input1_bound`.
    /// The output range is `0..(input0_bound * input1_bound)`.
    pub fn shuffle(input0_bound: usize, input1_bound: usize, rand: &mut impl Rng) -> Self {
        let bound = input0_bound
            .checked_mul(input1_bound)
            .expect("Saturated value");
        let layout = Layout::array::<usize>(bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..bound {
                ptr.add(i).write(i);
            }
            for i in 0..(bound - 2) {
                let shrinking_bound = bound - i;
                let last_valid_index = shrinking_bound - 1;
                let random_index = rand.random_range(0..shrinking_bound);
                if random_index < last_valid_index {
                    ptr.add(random_index).swap(ptr.add(last_valid_index));
                }
            }
        }

        Self {
            input0_bound,
            input1_bound,
            output_bound: bound,
            ptr,
        }
    }

    #[doc(hidden)]
    pub fn get_input0_bound(r: &Self) -> usize {
        r.input0_bound
    }
    #[doc(hidden)]
    pub fn get_input1_bound(r: &Self) -> usize {
        r.input1_bound
    }
    #[doc(hidden)]
    pub fn get_output_bound(r: &Self) -> usize {
        r.output_bound
    }
    /// Returns the shape of the input domain as a 2-tuple.
    pub fn input_shape(&self) -> (usize, usize) {
        (self.input0_bound, self.input1_bound)
    }
    /// Returns the bound of the output interval.
    pub fn output_bound(&self) -> usize {
        self.output_bound
    }

    #[doc(hidden)]
    pub unsafe fn get_unchecked(r: &Self, i0: usize, i1: usize) -> usize {
        unsafe { r.ptr.add(i0 * r.input1_bound + i1).read() }
    }
}

impl AsRef<[usize]> for Reindexing2 {
    fn as_ref(&self) -> &[usize] {
        unsafe { slice::from_raw_parts(self.ptr, self.input0_bound * self.input1_bound) }
    }
}

impl Index<(usize, usize)> for Reindexing2 {
    type Output = usize;

    fn index(&self, index: (usize, usize)) -> &Self::Output {
        if index.0 >= self.input0_bound || index.1 >= self.input1_bound {
            panic!("Index out of bound.")
        }
        unsafe {
            self.ptr
                .add(index.0 * self.input1_bound + index.1)
                .as_ref()
                .unwrap()
        }
    }
}

impl Drop for Reindexing2 {
    fn drop(&mut self) {
        if self.ptr.is_null() {
            return;
        }
        let layout = Layout::array::<usize>(self.input0_bound * self.input1_bound).unwrap();
        unsafe { std::alloc::dealloc(self.ptr as *mut u8, layout) }
    }
}

/**
 * A type representing an immutable mapping from a triple of  integer intervals `0..input0_bound`, `0..input1_bound`, and `0..input2_bound` to  integer interval `0..output_bound`.
 *
 * Mathematically this encodes a function `0..input0_bound × 0..input1_bound × 0..input2_bound → 0..output_bound`.
 */
pub struct Reindexing3 {
    input0_bound: usize,
    input1_bound: usize,
    input2_bound: usize,
    output_bound: usize,
    ptr: *mut usize,
}

impl Reindexing3 {
    /// Creates a new `Reindexing3` from a function `f: (usize, usize, usize) -> usize` and bounds on the input and output intervals.
    /// The function `f` is evaluated at each triple of integers from `0` to `input0_bound - 1`, `0` to `input1_bound - 1`, and `0` to `input2_bound - 1`, and the results are reduced modulo `output_bound` to produce the reindexing.
    /// If any of `input0_bound`, `input1_bound`, or `input2_bound` is zero, the resulting reindexing is empty and the function `f` is not evaluated at all. If `output_bound` is zero, the function panics.
    /// The resulting reindexing is immutable and guarantees that for each triple `(i0, i1, i2)` in `0..input0_bound × 0..input1_bound × 0..
    pub fn new(
        input0_bound: usize,
        input1_bound: usize,
        input2_bound: usize,
        output_bound: usize,
        mut f: impl FnMut(usize, usize, usize) -> usize,
    ) -> Self {
        if input0_bound == 0 || input1_bound == 0 || input2_bound == 0 {
            return Self {
                input0_bound: 0,
                input1_bound: 0,
                input2_bound: 0,
                output_bound,
                ptr: std::ptr::null_mut(),
            };
        }
        if output_bound == 0 {
            panic!("output_bound must be greater than 0");
        }
        let count = input0_bound
            .checked_mul(input1_bound)
            .expect("Saturated value")
            .checked_mul(input2_bound)
            .expect("Saturated value");

        let layout = Layout::array::<usize>(count).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            let mut i = 0;
            for i0 in 0..input0_bound {
                for i1 in 0..input1_bound {
                    for i2 in 0..input2_bound {
                        ptr.add(i).write(f(i0, i1, i2) % output_bound);
                        i += 1;
                    }
                }
            }
        }
        Self {
            input0_bound,
            input1_bound,
            input2_bound,
            output_bound,
            ptr,
        }
    }

    /// Creates a random permutation reindexing of the product domain `0..input0_bound × 0..input1_bound × 0..input2_bound`.
    /// The output range is `0..(input0_bound * input1_bound * input2_bound)`.
    pub fn shuffle(
        input0_bound: usize,
        input1_bound: usize,
        input2_bound: usize,
        rand: &mut impl Rng,
    ) -> Self {
        let bound = input0_bound
            .checked_mul(input1_bound)
            .expect("Saturated value")
            .checked_mul(input2_bound)
            .expect("Saturated value");
        let layout = Layout::array::<usize>(bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..bound {
                ptr.add(i).write(i);
            }
            for upper_index in (1..bound).rev() {
                let random_index = rand.random_range(0..=upper_index);
                if random_index != upper_index {
                    ptr.add(random_index).swap(ptr.add(upper_index));
                }
            }
        }

        Self {
            input0_bound,
            input1_bound,
            input2_bound,
            output_bound: bound,
            ptr,
        }
    }

    #[doc(hidden)]
    pub fn get_input0_bound(r: &Self) -> usize {
        r.input0_bound
    }
    #[doc(hidden)]
    pub fn get_input1_bound(r: &Self) -> usize {
        r.input1_bound
    }
    #[doc(hidden)]
    pub fn get_input2_bound(r: &Self) -> usize {
        r.input2_bound
    }
    #[doc(hidden)]
    pub fn get_output_bound(r: &Self) -> usize {
        r.output_bound
    }

    /// Returns the shape of the input domain as a 3-tuple.
    pub fn input_shape(&self) -> (usize, usize, usize) {
        (self.input0_bound, self.input1_bound, self.input2_bound)
    }
    /// Returns the bound of the output interval.
    pub fn output_bound(&self) -> usize {
        self.output_bound
    }

    #[doc(hidden)]
    pub unsafe fn get_unchecked(r: &Self, i0: usize, i1: usize, i2: usize) -> usize {
        unsafe {
            r.ptr
                .add(i0 * r.input1_bound * r.input2_bound + i1 * r.input2_bound + i2)
                .read()
        }
    }
}

impl AsRef<[usize]> for Reindexing3 {
    fn as_ref(&self) -> &[usize] {
        unsafe {
            slice::from_raw_parts(
                self.ptr,
                self.input0_bound * self.input1_bound * self.input2_bound,
            )
        }
    }
}

impl Index<(usize, usize, usize)> for Reindexing3 {
    type Output = usize;

    fn index(&self, index: (usize, usize, usize)) -> &Self::Output {
        if index.0 >= self.input0_bound
            || index.1 >= self.input1_bound
            || index.2 >= self.input2_bound
        {
            panic!("Index out of bound.")
        }
        let index = (index.0 * self.input1_bound + index.1) * self.input2_bound + index.2;
        unsafe { self.ptr.add(index).as_ref().unwrap() }
    }
}

impl Drop for Reindexing3 {
    fn drop(&mut self) {
        if self.ptr.is_null() {
            return;
        }
        let layout =
            Layout::array::<usize>(self.input0_bound * self.input1_bound * self.input2_bound)
                .unwrap();
        unsafe { std::alloc::dealloc(self.ptr as *mut u8, layout) }
    }
}

/**
 * A type representing an immutable mapping from a quadruple of  integer intervals `0..input0_bound`, `0..input1_bound`, `0..input2_bound`, and `0..input3_bound` to  integer interval `0..output_bound`.
 *
 * Mathematically this encodes a function `0..input0_bound × 0..input1_bound × 0..input2_bound × 0..input3_bound → 0..output_bound`.
 */
pub struct Reindexing4 {
    input0_bound: usize,
    input1_bound: usize,
    input2_bound: usize,
    input3_bound: usize,
    output_bound: usize,
    ptr: *mut usize,
}

impl Reindexing4 {
    /// Creates a new `Reindexing4` from a function `f: (usize, usize, usize, usize) -> usize` and bounds on the input and output intervals.
    /// The function `f` is evaluated at each quadruple of integers from `0` to `input0_bound - 1`, `0` to `input1_bound - 1`, `0` to `input2_bound - 1`, and `0` to `input3_bound - 1`, and the results are reduced modulo `output_bound` to produce the reindexing.
    /// If any of `input0_bound`, `input1_bound`, `input2_bound`, or `input3_bound` is zero, the resulting reindexing is empty and the function `f` is not evaluated at all. If `output_bound` is zero, the function panics.
    /// The resulting reindexing is immutable and guarantees that for each quadruple `(i0, i1, i2, i3)` in `0..input0_bound × 0..input1_bound × 0..input2_bound × 0..input3_bound`, the value of `reindexing[i0, i1, i2, i3]` is in `0..output_bound`.
    pub fn new(
        input0_bound: usize,
        input1_bound: usize,
        input2_bound: usize,
        input3_bound: usize,
        output_bound: usize,
        mut f: impl FnMut(usize, usize, usize, usize) -> usize,
    ) -> Self {
        if input0_bound == 0 || input1_bound == 0 || input2_bound == 0 || input3_bound == 0 {
            return Self {
                input0_bound: 0,
                input1_bound: 0,
                input2_bound: 0,
                input3_bound: 0,
                output_bound,
                ptr: std::ptr::null_mut(),
            };
        }
        if output_bound == 0 {
            panic!("output_bound must be greater than 0");
        }
        let count = input0_bound
            .checked_mul(input1_bound)
            .expect("Saturated value")
            .checked_mul(input2_bound)
            .expect("Saturated value")
            .checked_mul(input3_bound)
            .expect("Saturated value");

        let layout = Layout::array::<usize>(count).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            let mut i = 0;
            for i0 in 0..input0_bound {
                for i1 in 0..input1_bound {
                    for i2 in 0..input2_bound {
                        for i3 in 0..input3_bound {
                            ptr.add(i).write(f(i0, i1, i2, i3) % output_bound);
                            i += 1;
                        }
                    }
                }
            }
        }
        Self {
            input0_bound,
            input1_bound,
            input2_bound,
            input3_bound,
            output_bound,
            ptr,
        }
    }

    /// Creates a random permutation reindexing of the product domain `0..input0_bound × 0..input1_bound × 0..input2_bound × 0..input3_bound`.
    /// The output range is `0..(input0_bound * input1_bound * input2_bound * input3_bound)`.
    pub fn shuffle(
        input0_bound: usize,
        input1_bound: usize,
        input2_bound: usize,
        input3_bound: usize,
        rand: &mut impl Rng,
    ) -> Self {
        let bound = input0_bound
            .checked_mul(input1_bound)
            .expect("Saturated value")
            .checked_mul(input2_bound)
            .expect("Saturated value")
            .checked_mul(input3_bound)
            .expect("Saturated value");
        let layout = Layout::array::<usize>(bound).unwrap();

        let ptr: *mut usize;
        unsafe {
            ptr = std::alloc::alloc(layout) as *mut usize;

            if ptr.is_null() {
                std::alloc::handle_alloc_error(layout);
            }
            for i in 0..bound {
                ptr.add(i).write(i);
            }
            for upper_index in (1..bound).rev() {
                let random_index = rand.random_range(0..=upper_index);
                if random_index != upper_index {
                    ptr.add(random_index).swap(ptr.add(upper_index));
                }
            }
        }

        Self {
            input0_bound,
            input1_bound,
            input2_bound,
            input3_bound,
            output_bound: bound,
            ptr,
        }
    }

    #[doc(hidden)]
    pub fn get_input0_bound(r: &Self) -> usize {
        r.input0_bound
    }
    #[doc(hidden)]
    pub fn get_input1_bound(r: &Self) -> usize {
        r.input1_bound
    }
    #[doc(hidden)]
    pub fn get_input2_bound(r: &Self) -> usize {
        r.input2_bound
    }
    #[doc(hidden)]
    pub fn get_input3_bound(r: &Self) -> usize {
        r.input3_bound
    }
    #[doc(hidden)]
    pub fn get_output_bound(r: &Self) -> usize {
        r.output_bound
    }

    /// Returns the shape of the input domain as a 4-tuple.
    pub fn input_shape(&self) -> (usize, usize, usize, usize) {
        (
            self.input0_bound,
            self.input1_bound,
            self.input2_bound,
            self.input3_bound,
        )
    }
    /// Returns the bound of the output interval.
    pub fn output_bound(&self) -> usize {
        self.output_bound
    }

    #[doc(hidden)]
    pub unsafe fn get_unchecked(r: &Self, i0: usize, i1: usize, i2: usize, i3: usize) -> usize {
        unsafe {
            r.ptr
                .add(((i0 * r.input1_bound + i1) * r.input2_bound + i2) * r.input3_bound + i3)
                .read()
        }
    }
}

impl AsRef<[usize]> for Reindexing4 {
    fn as_ref(&self) -> &[usize] {
        unsafe {
            slice::from_raw_parts(
                self.ptr,
                self.input0_bound * self.input1_bound * self.input2_bound * self.input3_bound,
            )
        }
    }
}

impl Index<(usize, usize, usize, usize)> for Reindexing4 {
    type Output = usize;

    fn index(&self, index: (usize, usize, usize, usize)) -> &Self::Output {
        if index.0 >= self.input0_bound
            || index.1 >= self.input1_bound
            || index.2 >= self.input2_bound
            || index.3 >= self.input3_bound
        {
            panic!("Index out of bound.")
        }
        let index = ((index.0 * self.input1_bound + index.1) * self.input2_bound + index.2)
            * self.input3_bound
            + index.3;
        unsafe { self.ptr.add(index).as_ref().unwrap() }
    }
}

impl Drop for Reindexing4 {
    fn drop(&mut self) {
        if self.ptr.is_null() {
            return;
        }
        let layout = Layout::array::<usize>(
            self.input0_bound * self.input1_bound * self.input2_bound * self.input3_bound,
        )
        .unwrap();
        unsafe { std::alloc::dealloc(self.ptr as *mut u8, layout) }
    }
}

#[cfg(test)]
mod tests {
    use rand::SeedableRng;

    use super::*;

    const SEED: [u8; 32] = [
        7u8, 38u8, 19u8, 50u8, 11u8, 22u8, 33u8, 44u8, 133u8, 144u8, 155u8, 166u8, 177u8, 188u8,
        199u8, 200u8, 55u8, 66u8, 77u8, 88u8, 99u8, 100u8, 111u8, 122u8, 211u8, 222u8, 233u8,
        244u8, 255u8, 0u8, 2u8, 3u8,
    ];

    #[test]
    fn test_reindexing1() {
        let reindexing = Reindexing1::new(16, 7, |i| i * 5 + 4);

        assert_eq!(16, Reindexing1::get_input0_bound(&reindexing));
        assert_eq!(7, Reindexing1::get_output_bound(&reindexing));

        assert_eq!(
            &[4, 2, 0, 5, 3, 1, 6, 4, 2, 0, 5, 3, 1, 6, 4, 2],
            reindexing.as_ref()
        );
        assert_eq!(2, reindexing[15])
    }

    #[should_panic]
    #[test]
    fn test_panicking_reindexing1() {
        let reindexing = Reindexing1::new(16, 7, |i| i * 5 + 4);
        let _ = reindexing[16];
    }

    #[test]
    fn test_shuffle_reindexing1() {
        let mut rand = rand::rngs::StdRng::from_seed(SEED);
        let reindexing = Reindexing1::shuffle(22, &mut rand);
        assert_eq!(
            &[
                14, 2, 11, 15, 12, 4, 1, 18, 9, 16, 21, 0, 5, 10, 13, 19, 3, 20, 17, 7, 6, 8
            ],
            reindexing.as_ref()
        );
        let reindexing = Reindexing1::shuffle(22, &mut rand);
        assert_eq!(
            &[
                5, 11, 0, 3, 10, 18, 15, 4, 7, 13, 9, 16, 20, 1, 12, 8, 17, 19, 21, 6, 2, 14
            ],
            reindexing.as_ref()
        );
    }

    #[test]
    fn test_shuffle_partially_reindexing1() {
        let mut rand = rand::rngs::StdRng::from_seed(SEED);
        let reindexing = Reindexing1::shuffle_partially(7, 22, &mut rand);
        assert_eq!(&[9, 12, 17, 13, 5, 11, 1], reindexing.as_ref());

        let reindexing = Reindexing1::shuffle_partially(7, 22, &mut rand);
        assert_eq!(&[7, 6, 0, 16, 8, 19, 12], reindexing.as_ref());

        let reindexing = Reindexing1::shuffle_partially(7, 22, &mut rand);
        assert_eq!(&[2, 17, 8, 0, 12, 11, 20], reindexing.as_ref());

        let reindexing = Reindexing1::shuffle_partially(7, 22, &mut rand);
        assert_eq!(&[16, 7, 0, 19, 20, 4, 3], reindexing.as_ref());

        let reindexing = Reindexing1::shuffle_partially(7, 22, &mut rand);
        assert_eq!(&[16, 13, 18, 5, 2, 17, 0], reindexing.as_ref());
    }

    #[test]
    fn test_shuffle_partially_reindexing1_uniformity() {
        let mut rand = rand::rngs::StdRng::from_seed(SEED);
        let mut counts = [0; 22];
        const NUMBER_OF_SAMPLES: usize = 50_000;
        for _ in 0..NUMBER_OF_SAMPLES {
            let reindexing = Reindexing1::shuffle_partially(7, 22, &mut rand);
            for index in reindexing.as_ref() {
                counts[*index] += 1;
            }
        }
        let min = counts.iter().min().unwrap();
        let max = counts.iter().max().unwrap();
        assert!(
            max - min <= NUMBER_OF_SAMPLES / 100,
            "Counts are not uniform: min = {}, max = {}",
            min,
            max
        );
    }

    #[test]
    fn test_reindexing2() {
        let reindexing = Reindexing2::new(6, 4, 24, |i0, i1| 17 * (i0 * 4 + i1) + 11);

        assert_eq!(6, Reindexing2::get_input0_bound(&reindexing));
        assert_eq!(4, Reindexing2::get_input1_bound(&reindexing));
        assert_eq!(24, Reindexing2::get_output_bound(&reindexing));

        assert_eq!(
            &[
                11, 4, 21, 14, 7, 0, 17, 10, 3, 20, 13, 6, 23, 16, 9, 2, 19, 12, 5, 22, 15, 8, 1,
                18
            ],
            reindexing.as_ref(),
        );
        assert_eq!(18, reindexing[(5, 3)])
    }

    #[should_panic]
    #[test]
    fn test_panicking_reindexing2() {
        let reindexing = Reindexing2::new(6, 4, 24, |i0, i1| i0 * 4 + i1);
        let _ = reindexing[(5, 4)];
    }

    #[test]
    fn test_shuffle_reindexing2() {
        let mut rand = rand::rngs::StdRng::from_seed(SEED);
        let reindexing = Reindexing2::shuffle(3, 5, &mut rand);
        assert_eq!(
            &[8, 6, 9, 3, 0, 2, 12, 7, 11, 1, 13, 10, 14, 4, 5],
            reindexing.as_ref()
        );
    }
}