Add rewrite with Theorem 1

src/function.rs

@@ -43,8 +43,8 @@ impl<T: Function + ?Sized> Function for Rc<T> {
 
 /// `f ∘ g`: applies `g`, then `f` to the result.
 pub struct Composition<F, G> {
-    f: F,
-    g: G,
+    pub(crate) f: F,
+    pub(crate) g: G,
 }
 
 /// Composes `f` after `g`. Composing a tensor after an index map re-indexes it,
@@ -89,7 +89,7 @@ where
 /// of its elements, repeats included, so it goes from the multisets of the
 /// domain of `f` to those of its codomain.
 pub struct MultiFunction<F> {
-    f: F,
+    pub(crate) f: F,
 }
 
 pub fn multiset_lift<F: Function>(f: F) -> MultiFunction<F> {

src/index_functions.rs

@@ -331,10 +331,30 @@ impl IndexFunction for Where {
     }
 }
 
+/// Whether a multi-index function sends each index to at most one index, which
+/// makes a sum over it a sparse tensor (Definition 4 of the paper). The answer
+/// may be a conservative `false`, never a wrong `true`.
+pub trait Sparsity {
+    fn is_sparse(&self) -> bool;
+}
+
+impl<T: Sparsity + ?Sized> Sparsity for Rc<T> {
+    fn is_sparse(&self) -> bool {
+        (**self).is_sparse()
+    }
+}
+
 pub struct EmptyMapping {
     signature: MultiIndexFunctionSignature,
 }
 
+// Every index is sent to no index at all.
+impl Sparsity for EmptyMapping {
+    fn is_sparse(&self) -> bool {
+        true
+    }
+}
+
 impl EmptyMapping {
     pub fn new(domain: HRect, codomain: HRect) -> Self {
         Self {
@@ -374,6 +394,18 @@ impl COO {
     }
 }
 
+// A partial coordinate matches every index sharing its prefix, so two stored
+// coordinates land on a common index when one is a prefix of the other, equal
+// ones included. Sorted, the coordinates starting with a given one come right
+// after it, so only neighbours need comparing.
+impl Sparsity for COO {
+    fn is_sparse(&self) -> bool {
+        let mut coords: Vec<&Vec<usize>> = self.inverse_coords.iter().collect();
+        coords.sort();
+        coords.windows(2).all(|pair| !pair[1].starts_with(pair[0]))
+    }
+}
+
 impl IndexFunction for COO {
     type Codomain = Multisets<HRect>;
 
@@ -434,6 +466,13 @@ impl IndexFunction for Diagonal {
     }
 }
 
+// An index on the diagonal is sent to one index, and any other to none.
+impl Sparsity for Diagonal {
+    fn is_sparse(&self) -> bool {
+        true
+    }
+}
+
 /// Projection dropping axis `dim`.
 /// e.g. with dim=1, f(m, n, p) = [m, p]^T
 pub struct Projection {
@@ -516,6 +555,13 @@ impl IndexFunction for InvProjection {
     }
 }
 
+// Each index is sent to the whole line along `dim`.
+impl Sparsity for InvProjection {
+    fn is_sparse(&self) -> bool {
+        self.signature.codomain().set().shape()[self.dim] <= 1
+    }
+}
+
 /// The affine map `x ↦ basis·x + offset`. `basis` has a row per dimension of the
 /// codomain and a column per dimension of the domain. Its entries and `offset`
 /// may be negative, as long as every index of the domain lands in the codomain.
@@ -642,3 +688,11 @@ impl IndexFunction for InvAffine {
         &self.signature
     }
 }
+
+// Each index is sent to its preimage under the affine map, which holds at most
+// one index when the map is one-to-one.
+impl Sparsity for InvAffine {
+    fn is_sparse(&self) -> bool {
+        self.preimage.is_injective()
+    }
+}

src/lattice.rs

@@ -38,6 +38,12 @@ impl Preimage {
         }
     }
 
+    /// Whether `x ↦ basis·x` is one-to-one, i.e. the basis has no integer
+    /// kernel. A kernel the box is too small to reach still counts.
+    pub(crate) fn is_injective(&self) -> bool {
+        self.h_pivots.len() == self.extents.len()
+    }
+
     /// Every index `x` of the box with `basis·x = c`.
     pub(crate) fn solve(&self, c: &[i64]) -> Vec<Vec<usize>> {
         let mut z = vec![0; self.h_pivots.len()];

src/lib.rs

@@ -8,6 +8,7 @@ pub mod index_functions;
 mod lattice;
 pub mod multiset;
 pub mod reductions;
+pub mod rewrite;
 pub mod sets;
 pub mod signature;
 pub mod tensor;
@@ -21,13 +22,15 @@ pub use functions::{
 };
 pub use index_functions::{
     Affine, COO, Concat, Diagonal, EmptyMapping, IndexFunction, InvAffine, InvProjection,
-    MultiIndexFunctionSignature, Permutation, Projection, Squeeze, Stride, Unsqueeze, Where,
+    MultiIndexFunctionSignature, Permutation, Projection, Sparsity, Squeeze, Stride, Unsqueeze,
+    Where,
 };
 pub use multiset::Multiset;
 pub use reductions::{
     Inf, InfReduction, Product, ProductReduction, Reduction, ReductionFunction, Sum, SumReduction,
     Sup, SupReduction,
 };
+pub use rewrite::{Rewrite, rewrite};
 pub use sets::{HRect, IndexSet, JaggedDomain, Multisets, Set, Values};
 pub use signature::{HRectSignature, Signature};
 pub use tensor::{JaggedTensor, Physical, ReductionTensor, Tensor};

src/rewrite.rs

@@ -0,0 +1,62 @@
+//! Rewriting a tensor into an equivalent one that is cheaper to evaluate, with
+//! the identities between compositions of Section 4 of the paper. Only
+//! Theorem 1 is used so far.
+
+use crate::function::{
+    Composition, Function, MultiFunction, ValueFunction, compose, multiset_lift,
+};
+use crate::index_functions::Sparsity;
+use crate::reductions::{Reduction, Sum, SumReduction};
+use crate::sets::{Multisets, Values};
+use crate::tensor::ReductionTensor;
+
+/// A tensor that an identity of the paper can rewrite into an equivalent one.
+pub trait Rewrite: Sized {
+    type Rewritten;
+
+    /// The equivalent tensor, or `Err` with the tensor unchanged when the
+    /// identity does not apply to it.
+    fn rewrite(self) -> Result<Self::Rewritten, Self>;
+}
+
+/// Rewrites `tensor` into an equivalent one that is cheaper to evaluate, or
+/// gives it back unchanged as `Err` when no identity applies to it.
+pub fn rewrite<T: Rewrite>(tensor: T) -> Result<T::Rewritten, T> {
+    tensor.rewrite()
+}
+
+/// Theorem 1: a function on values after a sparse tensor only has to be applied
+/// to its stored values, `f ∘ Σ ∘ Ū ∘ ψ = Σ ∘ (f ∘ U)‾ ∘ ψ`. This holds when
+/// `f(0) = 0` and `ψ` sends each index to at most one index, which is checked.
+impl<A, U, M> Rewrite for Composition<ValueFunction<A>, ReductionTensor<Sum<A>, U, M>>
+where
+    A: Clone + PartialEq,
+    SumReduction: Reduction<A>,
+    U: Function<Codomain = Values<A>>,
+    M: Function<Codomain = Multisets<U::Domain>> + Sparsity,
+{
+    type Rewritten = ReductionTensor<Sum<A>, Composition<ValueFunction<A>, U>, M>;
+
+    fn rewrite(self) -> Result<Self::Rewritten, Self> {
+        let zero = <SumReduction as Reduction<A>>::NEUTRAL;
+        if self.f.apply(&zero) != zero || !self.g.g.g.is_sparse() {
+            return Err(self);
+        }
+        let Composition {
+            f,
+            g:
+                Composition {
+                    g:
+                        Composition {
+                            f: MultiFunction { f: u },
+                            g: psi,
+                        },
+                    ..
+                },
+        } = self;
+        Ok(compose(
+            Sum::new(),
+            compose(multiset_lift(compose(f, u)), psi),
+        ))
+    }
+}

tests/index_functions.rs

@@ -2,8 +2,8 @@ use std::{assert_eq, collections::HashMap, rc::Rc};
 
 use cten::{
     Affine, COO, Concat, Diagonal, EmptyMapping, HRect, IndexFunction, IndexSet, InvAffine,
-    InvProjection, JaggedDomain, Multiset, Permutation, Physical, Projection, Squeeze, Stride,
-    Unsqueeze, Where,
+    InvProjection, JaggedDomain, Multiset, Permutation, Physical, Projection, Sparsity, Squeeze,
+    Stride, Unsqueeze, Where,
 };
 
 /// The indices of a multiset in sorted order, to compare multisets whose
@@ -407,3 +407,46 @@ fn a_where_sends_each_index_to_the_part_its_condition_picks() {
     assert_eq!(picks.map(&[1]), vec![1, 1]);
     assert_eq!(picks.map(&[2]), vec![0, 2]);
 }
+
+// A partial coordinate matches every index sharing its prefix, so it meets every
+// coordinate it is a prefix of.
+#[test]
+fn a_coo_is_sparse_when_no_two_coordinates_meet() {
+    let coo = |coords: Vec<Vec<usize>>| COO::new(HRect::new(vec![3, 3]), coords);
+
+    assert!(coo(vec![vec![0, 2], vec![1, 0], vec![1, 2]]).is_sparse());
+    assert!(!coo(vec![vec![1, 2], vec![0, 0], vec![1, 2]]).is_sparse());
+    assert!(!coo(vec![vec![2], vec![2, 1]]).is_sparse());
+    assert!(coo(vec![vec![2], vec![1, 1]]).is_sparse());
+}
+
+#[test]
+fn diagonal_and_empty_index_functions_are_sparse() {
+    assert!(Diagonal::new(HRect::new(vec![3, 3])).is_sparse());
+    assert!(EmptyMapping::new(HRect::new(vec![3]), HRect::new(vec![0])).is_sparse());
+}
+
+#[test]
+fn an_inv_projection_is_sparse_only_along_an_axis_of_extent_one() {
+    assert!(InvProjection::new(HRect::new(vec![2, 1]), 1).is_sparse());
+    assert!(!InvProjection::new(HRect::new(vec![2, 3]), 1).is_sparse());
+}
+
+#[test]
+fn an_inv_affine_is_sparse_when_its_affine_map_is_one_to_one() {
+    let diagonal = InvAffine::new(Affine::new(
+        HRect::new(vec![3]),
+        vec![vec![1], vec![1]],
+        vec![0, 0],
+        HRect::new(vec![3, 3]),
+    ));
+    let projection = InvAffine::new(Affine::new(
+        HRect::new(vec![2, 3, 4]),
+        vec![vec![1, 0, 0], vec![0, 0, 1]],
+        vec![0, 0],
+        HRect::new(vec![2, 4]),
+    ));
+
+    assert!(diagonal.is_sparse());
+    assert!(!projection.is_sparse());
+}