playfair.rs

//! Implementation of the [Playfair cipher].
//!
//! [Playfair cipher]: https://en.wikipedia.org/wiki/Playfair_cipher

#![cfg_attr(feature = "no_std", no_std)]
#![cfg_attr(all(docs, not(doctest)), feature(doc_auto_cfg))]

use core::fmt::{self, Display};

// Normally we'd use a `std` feature instead of a `no_std` feature (this works
// better with the additive nature of cargo features) but... we want this file
// to work in the playground (where we can't pass in extra `cfg`s) without too
// many modifications.
macro_rules! on_std {
    ($($i:item)*) => {$(
        #[cfg(not(feature = "no_std"))]
        $i
    )*};
}

/// Compact representation of the coordinates of a cell in the [`Playfair`]
/// cipher table.
///
/// Because we know that the `row` and `column` of each coordinate is in `0..5`
/// we can save some space over using `(usize, usize)` to represent these
/// coordinates.
///
/// Within this type, bits `0..3` represent the row (i.e. `y`) and bits `3..6`
/// represent the column (i.e. `x`). This ordering means the derived `Ord` impl
/// behaves as we'd expect: `x` is compared first with `y` as a tiebreaker.
//
// Really we only need 5 bits (`ceil(log2(25))`) to store these instead of 8
// and we can also probably do something clever involving overlaying the
// char-to-pos and pos-to-char (i.e. `table`; also only uses <5 bits per elem)
// tables but we'll leave those as potential optimizations for the future, if
// the need arises.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
struct TableCoord(u8);

impl TableCoord {
    const fn new(row: usize, col: usize) -> Self {
        debug_assert!(row < 5);
        debug_assert!(col < 5);

        Self(((col << 3) | row) as u8)
    }

    const fn row(self) -> usize {
        (self.0 & 0b111) as usize
    }

    const fn col(self) -> usize {
        ((self.0 >> 3) & 0b111) as usize
    }
}

impl From<TableCoord> for (usize, usize) {
    fn from(c: TableCoord) -> Self {
        (c.col(), c.row())
    }
}

impl TableCoord {
    /// Yields an iterator over the coordinates of all the cells in the cipher
    /// table, row major.
    fn iter_all() -> impl Iterator<Item = Self> {
        (0..5).flat_map(|r| (0..5).map(move |c| Self::new(r, c)))
    }

    fn idx<T>(self, table: &[[T; 5]; 5]) -> &T {
        &table[self.row()][self.col()]
    }

    fn idx_mut<T>(self, table: &mut [[T; 5]; 5]) -> &mut T {
        &mut table[self.row()][self.col()]
    }
}

impl TableCoord {
    const fn right(self) -> Self {
        Self::new(self.row(), (self.col() + 1) % 5)
    }
    const fn left(self) -> Self {
        Self::new(self.row(), (self.col() + 5 - 1) % 5)
    }

    const fn above(self) -> Self {
        Self::new((self.row() + 5 - 1) % 5, self.col())
    }
    const fn below(self) -> Self {
        Self::new((self.row() + 1) % 5, self.col())
    }
}

#[derive(thiserror::Error, Debug, PartialEq, Eq, Clone, Copy, Hash)]
#[non_exhaustive]
pub enum Error {
    #[error("Expected an ASCII alphabetical character but got: {got}")]
    ExpectedAsciiAlpha { got: char },
}

/// Implementation of the Playfair cipher as described [here][Playfair cipher].
///
/// ## Usage
///
/// String based: `encrypt_str`/`decrypt_str`.
/// Reader/Write based (`u8`, not unicode friendly):
/// Iterator + Callback based (`no_std` friendly):
///
/// ## Notes & Caveats
///
///   1) This implementation is ASCII-centric. By default this implementation
///      passes through all non-ascii alphabetical characters, unencrypted:
///      ```
///      # use playfair::Playfair;
///      # #[cfg(not(feature = "no_std"))]
///      assert_eq!(Playfair::default().encrypt_str("hi 🤗"), Ok("bm 🤗".to_string()));
///      ```
///      Though this can be customized with [`Playfair::strict`].
///
///   2) Case is preserved:
///      ```
///      # use playfair::Playfair;
///      # #[cfg(not(feature = "no_std"))] {
///      let pf = Playfair::default();
///
///      let orig = "Hello World";
///      let enc = pf.encrypt_str(orig).unwrap();
///      assert_eq!(enc, "Dmynqs Vqcrgo");
///
///      let dec = pf.decrypt_str(enc).unwrap();
///      assert_eq!(dec, "Helqoq Worldq");
///      # }
///      ```
///
///   3) By default, `Q` is used as a padding character:
///      ```
///      # use playfair::Playfair;
///      # #[cfg(not(feature = "no_std"))] {
///      let pf = Playfair::default();
///
///      // Trailing single characters are padded:
///      assert_eq!(pf.decrypt_str(pf.encrypt_str("E").unwrap()), Ok("EQ".to_string()));
///      // The case of the added character matches the that of the trailing
///      // character:
///      assert_eq!(pf.decrypt_str(pf.encrypt_str("e").unwrap()), Ok("eq".to_string()));
///
///      // Duplicate characters are also padded:
///      assert_eq!(pf.decrypt_str(pf.encrypt_str("aa").unwrap()), Ok("aq".to_string()));
///      // Here the case of the second character is used for the added
///      // character:
///      assert_eq!(pf.decrypt_str(pf.encrypt_str("aA").unwrap()), Ok("aQ".to_string()));
///      # }
///      ```
///
///      You can customize the padding character with
///      [`Playfair::set_padding_char`].
///
///   4) To squeeze the 26 ASCII alphabetical characters into the 25 element
///      table used by this cipher, `I` and `J` are treated as the same
///      character when encrypting and decrypting.
/// [Playfair cipher]: https://en.wikipedia.org/wiki/Playfair_cipher
#[derive(Debug, PartialEq, Eq, Clone, Hash)]
pub struct Playfair {
    /// Uppercase ASCII characters constituting the cipher's table.
    //
    // This lets us go from position to *character* quickly.
    table: [[u8; 5]; 5],
    /// Alphabetical character to table position map.
    //
    // This lets us go from character to *position* quickly.
    char_to_pos: [TableCoord; 25],
    /// "Uncommon" character used to pad single characters (i.e. trailing chars
    /// and digrams of duplicates like `"ee"` that only have one _unique_ char).
    ///
    /// This is also stored in uppercase.
    ///
    /// Note that trailing chars/repeated chars containing this char will be
    /// encoded as this char + 1.
    padding_char: u8,
    /// If `true`, this instance will error when given non-ASCII alphabetical
    /// characters to encrypt or decrypt.
    ///
    /// If `false` (the default), this instance will *pass through* non-ASCII
    /// alphabetical characters.
    pub strict: bool,
}

impl Display for Playfair {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for row in self.table {
            for c in row {
                write!(f, "{}", c as char)?;
            }
            writeln!(f)?;
        }

        Ok(())
    }
}

/// Constructs a [`Playfair`] instance with "playfair example" as the key and `Q` as the
/// padding character. Strict mode is disabled.
impl Default for Playfair {
    fn default() -> Self {
        Self::new("playfair example")
    }
}

mod helpers {
    // Maps `i` and `j` to the same index.
    pub(crate) const fn char_to_idx(c: char) -> usize {
        if !c.is_ascii_alphabetic() {
            panic!()
        }

        let ret = match c.to_ascii_uppercase() as u8 {
            c @ b'A'..=b'I' => c - b'A',
            c @ b'J'..=b'Z' => c - b'A' - 1,
            _ => unreachable!(),
        };

        ret as usize
    }

    // Never returns `j`.
    pub(crate) const fn idx_to_char(i: usize) -> u8 {
        match i as u8 + b'A' {
            c @ b'A'..=b'I' => c,
            c @ b'J'..=b'Y' => c + 1,
            _ => unreachable!(),
        }
    }

    pub(crate) const fn copy_case(new: u8, orig: u8) -> u8 {
        debug_assert!(new.is_ascii_alphabetic() && orig.is_ascii_alphabetic());
        if orig.is_ascii_lowercase() {
            new.to_ascii_lowercase()
        } else {
            new.to_ascii_uppercase()
        }
    }
}

impl Playfair {
    pub fn new(key: impl AsRef<str>) -> Self {
        Self::new_from_iter(key.as_ref().chars())
    }

    // Whitespace and other non-alphabetic chars are allowed in the key but they
    // are ignored.
    pub fn new_from_iter(key: impl Iterator<Item = char>) -> Self {
        let mut coords = TableCoord::iter_all();
        let mut table = [[0; 5]; 5];
        let mut char_to_pos = [None::<TableCoord>; 25];

        // Ordinarily we'd just use a `HashSet` but we want to support `no_std`,
        // so:
        let mut seen = [false; 25];
        for (char, coord) in key
            .filter(|c| c.is_ascii_alphabetic())
            .map(|c| c.to_ascii_uppercase())
            .filter(|&c| {
                let idx = helpers::char_to_idx(c);
                if seen[idx] {
                    false
                } else {
                    seen[idx] = true;
                    true
                }
            })
            // Pair each *valid* new alphabetic char with a table coordinate:
            .zip(&mut coords)
        {
            *coord.idx_mut(&mut table) = char as u8;
            char_to_pos[helpers::char_to_idx(char)] = Some(coord);
        }

        // Fill in any remaining blank spaces with the characters *not* in the
        // key:
        for (idx, coord) in seen
            .iter()
            .enumerate()
            .filter(|(_, &s)| !s)
            .map(|(i, _)| i)
            .zip(&mut coords)
        {
            *coord.idx_mut(&mut table) = helpers::idx_to_char(idx);
            char_to_pos[idx] = Some(coord);
        }

        let char_to_pos = char_to_pos.map(Option::unwrap);

        Self {
            table,
            char_to_pos,
            padding_char: b'Q',
            strict: false,
        }
    }

    pub fn set_padding_char(&mut self, char: char) -> Result<(), Error> {
        // We'll accept lowercase chars:
        let char_upper = char.to_ascii_uppercase();
        if char_upper.is_ascii_uppercase() {
            self.padding_char = char_upper as u8;
            Ok(())
        } else {
            Err(Error::ExpectedAsciiAlpha { got: char })
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
pub enum TransformationMode {
    Encrypt,
    Decrypt,
}

impl Playfair {
    fn transform_pair(&self, _digram @ (a, b): (u8, u8), mode: TransformationMode) -> (u8, u8) {
        use TransformationMode::*;
        debug_assert!(a.is_ascii_alphabetic() && b.is_ascii_alphabetic());

        // Need to swap out b if it's the same as a:
        let b = if a.to_ascii_uppercase() == b.to_ascii_uppercase() {
            if b.to_ascii_uppercase() as u8 == self.padding_char {
                // uh-oh; we got a pair of our padding char...
                // use the next char up:
                helpers::copy_case(((self.padding_char - b'A' + 1) % 26) + b'A', b)
            } else {
                helpers::copy_case(self.padding_char, b)
            }
        } else {
            b
        };

        let a_coord = self.char_to_pos[helpers::char_to_idx(a as char)];
        let b_coord = self.char_to_pos[helpers::char_to_idx(b as char)];

        let (a_coord, b_coord) = match (a_coord, b_coord, mode) {
            (ac, bc, mode) if ac.row() == bc.row() => match mode {
                Encrypt => (ac.right(), bc.right()),
                Decrypt => (ac.left(), bc.left()),
            },
            (ac, bc, mode) if ac.col() == bc.col() => match mode {
                Encrypt => (ac.below(), bc.below()),
                Decrypt => (ac.above(), bc.above()),
            },
            (ac, bc, _) => {
                let (a_col, a_row) = ac.into();
                let (b_col, b_row) = bc.into();

                (TableCoord::new(a_row, b_col), TableCoord::new(b_row, a_col))
            }
        };

        let a_ = *a_coord.idx(&self.table);
        let b_ = *b_coord.idx(&self.table);
        (helpers::copy_case(a_, a), helpers::copy_case(b_, b))
    }

    fn transform_single(&self, single: u8, mode: TransformationMode) -> (u8, u8) {
        self.transform_pair(
            (single, helpers::copy_case(self.padding_char, single)),
            mode,
        )
    }

    pub fn transform(
        &self,
        inp: impl Iterator<Item = char>,
        mut output_func: impl FnMut(char),
        mode: TransformationMode,
    ) -> Result<(), Error> {
        // Essentially a bespoke inline SmallVec; should probably turn this into
        // it's own type..
        let mut buf_len = 0;
        let mut buf = [0; 2];

        for next in inp {
            // Process a new character:
            if next.is_ascii_alphabetic() {
                // Put chars we can transform in the buffer:
                debug_assert!(buf_len < 2);
                buf[buf_len] = next as u8;
                buf_len += 1;
            } else {
                // For chars we can't transform:
                if self.strict {
                    return Err(Error::ExpectedAsciiAlpha { got: next });
                }

                // If we have elements in our buffer we need to drain them
                // before we can emit this un-transform-able character:
                if buf_len != 0 {
                    assert_eq!(buf_len, 1);
                    let (a, b) = self.transform_single(buf[0], mode);
                    output_func(a as char);
                    output_func(b as char);

                    buf_len = 0;
                }

                output_func(next);
            }

            // If we've filled up our two element buffer we've got a digraph
            // and can do a transformation:
            if buf_len == 2 {
                let (a, b) = self.transform_pair((buf[0], buf[1]), mode);
                output_func(a as char);
                output_func(b as char);

                buf_len = 0;
            }
        }

        // Clear the remaining elements in the buffer:
        match buf_len {
            0 => {},
            1 => {
                let (a, b) = self.transform_single(buf[0], mode);
                output_func(a as char);
                output_func(b as char);
            },
            _ => unreachable!(),
        }

        Ok(())
    }
}

on_std! {
    impl Playfair {
        fn transform_str(
            &self,
            inp: impl AsRef<str>,
            mode: TransformationMode,
        ) -> Result<String, Error> {
            let mut output = String::with_capacity(inp.as_ref().len());
            let func = |c| { output.push(c); };
            self.transform(inp.as_ref().chars(), func, mode)?;

            Ok(output)
        }

        pub fn encrypt_str(&self, inp: impl AsRef<str>) -> Result<String, Error> {
            self.transform_str(inp, TransformationMode::Encrypt)
        }

        pub fn decrypt_str(&self, inp: impl AsRef<str>) -> Result<String, Error> {
            self.transform_str(inp, TransformationMode::Decrypt)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use TransformationMode::*;

    #[test]
    fn pairs() {
        fn roundtrip(inp: (u8, u8), exp: (u8, u8), out: Option<(u8, u8)>) {
            let pf = Playfair::default();
            println!("{pf}");
            let enc = pf.transform_pair(inp, Encrypt);
            assert_eq!(
                enc,
                exp,
                "E: ({}, {}) -> ({}, {}); got: ({}, {})",
                inp.0 as char,
                inp.1 as char,
                exp.0 as char,
                exp.1 as char,
                enc.0 as char,
                enc.1 as char
            );

            let got = pf.transform_pair(enc, Decrypt);
            let exp = out.unwrap_or(inp);
            assert_eq!(
                got,
                exp,
                "D: ({}, {}) -> ({}, {}); got: ({}, {})",
                enc.0 as char,
                enc.1 as char,
                inp.0 as char,
                inp.1 as char,
                exp.0 as char,
                exp.1 as char
            );
        }

        roundtrip((b'H', b'I'), (b'B', b'M'), None);
        roundtrip((b'D', b'E'), (b'O', b'D'), None);
        roundtrip((b'T', b'H'), (b'Z', b'B'), None);
        roundtrip((b'E', b'G'), (b'X', b'D'), None);
        roundtrip((b'E', b'E'), (b'X', b'O'), Some((b'E', b'Q')));
        roundtrip((b'E', b'e'), (b'X', b'o'), Some((b'E', b'q')));
    }

    #[test]
    #[cfg(not(feature = "no_std"))]
    fn string() {
        let inp = "hide the gold in the tree stump";
        let enc = "bmod zbxo dqac rk zbxo uixo kzzryk";
        let pf = Playfair::default();

        let res = pf.encrypt_str(inp).unwrap();
        // eprintln!("{res}");
        assert_eq!(res, enc);

        let dec_res = "hide theq gold in theq treq stumpq";
        let res = pf.decrypt_str(enc).unwrap();
        assert_eq!(res, dec_res);
    }
}

// TODO: fuzz `encrypt_str`/`decrypt_str`.

// TODO: proptest (with inputs without repeated chars or `i`/`j` and only even
// length strings) that asserts that we always roundtrip.


fn main() {
    let pf = Playfair::new("playfair example");
    eprintln!("{pf}");

    let inp = "Hide the gold in the tree stump! 💀";
    eprintln!("{inp}");

    let res = pf.encrypt_str(inp).unwrap();
    eprintln!("{res}");

    let res = pf.decrypt_str(res).unwrap();
    eprintln!("{res}");

}