round_key.rs

use std::fmt;

#[derive(Eq, PartialEq)]
enum Element {
    Index(usize),
    GCall(Vec<Element>),
}

impl fmt::Display for Element {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Element::*;
        match self {
            &Index(i) => write!(f, "W[{}]", i),
            &GCall(ref el) => {
                write!(f, "g(")?;
                for i in 0..el.len() - 1 {
                    write!(f, "{} ^ ", el[i])?;
                }
                write!(f, "{})", el[el.len() - 1])
            }
        }
    }
}

fn remove_pairs(v: &mut Vec<Element>) {
    let mut to_remove: Vec<(usize, usize)> = vec![];

    for (i, x) in v.iter().enumerate() {
        if let Some((j, _)) = v[i + 1..]
            .iter()
            .enumerate()
            .map(|(j, y)| (j + i + 1, y))
            .filter(|&(_, y)| x == y)
            .next()
        {
            if let None = to_remove
                .iter()
                .filter(|&&b| b.0 == i || b.0 == j || b.1 == i || b.1 == j)
                .next()
            {
                to_remove.push((i, j));
            }
        }
    }

    let t = v.drain(..)
        .enumerate()
        .filter(|&(i, _)| {
            if let Some(_) = to_remove.iter().filter(|&&(a, b)| a == i || b == i).next() {
                false
            } else {
                true
            }
        })
        .map(|(_, x)| x)
        .collect();
    *v = t;
}

// Finds the equation for W[index] in terms of W[40] ... W[43]
fn equation_for(index: usize) -> Vec<Element> {
    if index >= 40 {
        assert!(index <= 43); // bounds check

        return vec![Element::Index(index)];
    }

    let mut res = equation_for(index + 4);

    if index % 4 == 0 {
        // W[i] = W[i + 4] ^ g(W[i + 3])
        res.push(Element::GCall(equation_for(index + 3)));
    } else {
        // W[i] = W[i + 4] ^ W[i + 3]
        res.extend(equation_for(index + 3).into_iter());
    }

    remove_pairs(&mut res);

    res
}

fn main() {
    print!("W[0] = ");
    let res = equation_for(0);

    for i in 0..res.len() - 1 {
        print!("{} ^ ", res[i]);
    }
    println!("{}", res[res.len() - 1]);
}