steiner-tree 0.0.1

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

//! Set operations on iterators with sorted elements.
//!
//! Set operations like union, difference, XOR and intersection can be computed efficiently
//! on two iterators which have their elements sorted.

// /// Specify the set operation.
// pub enum SetOperationE {
//     /// S = A ⋃ B
//     Union,
//     /// S = A \ B
//     Difference,
//     /// (Disjunctive union or XOR)
//     /// S = (A \ B) ⋃ (B \ A)
//     SymmetricDifference,
//     /// S = A ⋂ B
//     Intersection,
// }

use std::marker::PhantomData;
use std::iter::Peekable;

pub trait SetOperation {
    /// Tell if some element should be in the output set.
    fn eval(is_in_a: bool, is_in_b: bool) -> bool;
}

/// S = A ⋃ B
pub struct Union;

impl SetOperation for Union {
    fn eval(is_in_a: bool, is_in_b: bool) -> bool {
        is_in_a || is_in_b
    }
}

/// S = A \ B
pub struct Difference;

impl SetOperation for Difference {
    fn eval(is_in_a: bool, is_in_b: bool) -> bool {
        is_in_a && !is_in_b
    }
}

/// (Disjunctive union or XOR)
/// S = (A \ B) ⋃ (B \ A)
pub struct SymmetricDifference;

impl SetOperation for SymmetricDifference {
    fn eval(is_in_a: bool, is_in_b: bool) -> bool {
        is_in_a ^ is_in_b
    }
}

/// S = A ⋂ B
pub struct Intersection;

impl SetOperation for Intersection {
    fn eval(is_in_a: bool, is_in_b: bool) -> bool {
        is_in_a && is_in_b
    }
}

/// Yield the results of a set operation of A and B.
/// A and B must be sorted otherwise the result is not correct.
pub struct SetOpIter<A, B, Op>
    where A: Iterator,
          B: Iterator {
    a: Peekable<A>,
    b: Peekable<B>,
    _op: PhantomData<Op>,
}

fn net_set_op<A, B, E, Op>(a: A, b: B) -> SetOpIter<A::IntoIter, B::IntoIter, Op>
    where A: IntoIterator<Item=E>,
          B: IntoIterator<Item=E>,
          E: Eq + Ord,
          Op: SetOperation {
    SetOpIter {
        a: a.into_iter().peekable(),
        b: b.into_iter().peekable(),
        _op: Default::default(),
    }
}

/// Compute the set intersection of two sets A and B. Both sets must be sorted iterators.
pub fn intersection<A, B, E>(a: A, b: B) -> SetOpIter<A::IntoIter, B::IntoIter, Intersection>
    where A: IntoIterator<Item=E>,
          B: IntoIterator<Item=E>,
          E: Eq + Ord {
    net_set_op(a, b)
}

/// Compute the set union of two sets A and B. Both sets must be sorted iterators.
pub fn union<A, B, E>(a: A, b: B) -> SetOpIter<A::IntoIter, B::IntoIter, Union>
    where A: IntoIterator<Item=E>,
          B: IntoIterator<Item=E>,
          E: Eq + Ord {
    net_set_op(a, b)
}

/// Compute the set difference of two sets A and B. Both sets must be sorted iterators.
pub fn difference<A, B, E>(a: A, b: B) -> SetOpIter<A::IntoIter, B::IntoIter, Difference>
    where A: IntoIterator<Item=E>,
          B: IntoIterator<Item=E>,
          E: Eq + Ord {
    net_set_op(a, b)
}

/// Compute the set symmetric-difference of two sets A and B. Both sets must be sorted iterators.
pub fn symmetric_difference<A, B, E>(a: A, b: B) -> SetOpIter<A::IntoIter, B::IntoIter, SymmetricDifference>
    where A: IntoIterator<Item=E>,
          B: IntoIterator<Item=E>,
          E: Eq + Ord {
    net_set_op(a, b)
}

impl<A, B, E, Op> Iterator for SetOpIter<A, B, Op>
    where A: Iterator<Item=E>,
          B: Iterator<Item=E>,
          E: Eq + Ord,
          Op: SetOperation {
    type Item = E;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Take the smaller front value of the both iterators or take both front elements
            // if they are equal.
            let (a, b) = match (self.a.peek(), self.b.peek()) {
                (Some(_), None) => (self.a.next(), None),
                (None, Some(_)) => (None, self.b.next()),
                (Some(a), Some(b)) => {
                    if a < b {
                        (self.a.next(), None)
                    } else if b < a {
                        (None, self.b.next())
                    } else {
                        debug_assert!(a == b);
                        (self.a.next(), self.b.next())
                    }
                }
                (None, None) => break None,
            };

            // Check that A is sorted.
            debug_assert!(
                if let (Some(a), Some(a_next)) = (&a, self.a.peek()) {
                    a <= a_next
                } else {
                    true
                },
                "iterator A is not sorted"
            );
            // Check that B is sorted.
            debug_assert!(
                if let (Some(b), Some(b_next)) = (&b, self.b.peek()) {
                    b <= b_next
                } else {
                    true
                },
                "iterator B is not sorted"
            );

            let output_element = match (a, b) {
                (Some(a), None) if Op::eval(true, false) => Some(a),
                (None, Some(b)) if Op::eval(false, true) => Some(b),
                (Some(a), Some(b)) if Op::eval(true, true) => {
                    debug_assert!(a == b);
                    Some(a)
                }
                (None, None) => unreachable!(),
                _ => None
            };

            if output_element.is_some() {
                break output_element;
            }
        }
    }
}

#[test]
fn test_union() {
    let a = vec![1, 2, 3];
    let b = vec![2, 4];
    let union: Vec<_> = union(a, b).collect();
    assert_eq!(union, vec![1, 2, 3, 4])
}

#[test]
fn test_difference() {
    let a = vec![1, 2, 3];
    let b = vec![2, 4];
    let diff: Vec<_> = difference(a, b).collect();
    assert_eq!(diff, vec![1, 3])
}

#[test]
fn test_symmetric_difference() {
    let a = vec![1, 2, 3];
    let b = vec![2, 4];
    let xor: Vec<_> = symmetric_difference(a, b).collect();
    assert_eq!(xor, vec![1, 3, 4])
}

#[test]
fn test_intersection() {
    let a = vec![1, 2, 3];
    let b = vec![2, 4];
    let intersection: Vec<_> = intersection(a, b).collect();
    assert_eq!(intersection, vec![2])
}