factrs 0.3.0

// Similar to gtsam: https://github.com/borglab/gtsam/blob/develop/gtsam/inference/Symbol.cpp
use std::{
    fmt::{self, Write},
    mem::size_of,
};

use crate::variables::VariableDtype;

// ------------------------- Symbol Basics ------------------------- //

/// Newtype wrap around u64
///
/// In it's final form, a Key is what is used for indexing inside of
/// Values and Factors. Generally it is created from a [Symbol]
#[derive(Clone, Copy, Eq, Hash, PartialEq, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Key(pub u64);

impl Symbol for Key {}

/// Human-readable symbol that will be turned into a [Key]
///
/// This is a requirement to be inserted into Values. Examples are
/// X(100), A(101), etc.
pub trait Symbol: fmt::Debug + Into<Key> {}

/// Adds type information to a [Symbol]
///
/// Will almost always be generated by [assign_symbols](factrs::assign_symbols)
pub trait TypedSymbol<V: VariableDtype>: Symbol {}

/// Custom formatting for keys in [Values](factrs::containers::Values) or
/// [Graph](factrs::containers::Graph)
///
/// The average user will likely never have to implement this as
/// [DefaultSymbolHandler] will almost always be sufficient, unless a custom key
/// is used.
pub trait KeyFormatter {
    fn fmt(f: &mut dyn Write, key: Key) -> fmt::Result;
}

// --------------------- Basic single char symbol --------------------- //

// Char is stored in the high CHR_BITS
// Idx is stored in the low IDX_BITS
const TOTAL_SIZE: usize = u64::BITS as usize;
const CHR_SIZE: usize = size_of::<char>() * 8;
const IDX_SIZE: usize = TOTAL_SIZE - CHR_SIZE;
const CHR_MASK: u64 = (char::MAX as u64) << IDX_SIZE;
const IDX_MASK: u64 = !CHR_MASK;

/// Default symbol handler for [Symbols](Symbol)
///
/// Specifically, this converts a char and an index into a [Key] and back
#[derive(Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct DefaultSymbolHandler;

impl DefaultSymbolHandler {
    pub fn sym_to_key(chr: char, idx: u32) -> Key {
        Key(((chr as u64) << IDX_SIZE) | (idx as u64 & IDX_MASK))
    }

    pub fn key_to_sym(k: Key) -> (char, u32) {
        let chr = ((k.0 & CHR_MASK) >> IDX_SIZE) as u8 as char;
        let idx = (k.0 & IDX_MASK) as u32;
        (chr, idx)
    }

    pub fn format(f: &mut dyn Write, chr: char, idx: u32) -> fmt::Result {
        write!(f, "{chr}{idx}")
    }
}

impl KeyFormatter for DefaultSymbolHandler {
    fn fmt(f: &mut dyn Write, key: Key) -> fmt::Result {
        let (chr, idx) = Self::key_to_sym(key);
        Self::format(f, chr, idx)
    }
}

/// Creates and assigns [Symbols](Symbol) to
/// [Variables](factrs::variables::Variable)
///
/// To reduce runtime errors, factrs symbols are tagged
/// with the type they will be used with. This macro will create a new symbol
/// and implement all the necessary traits for it to be used as a symbol.
/// ```
/// use factrs::{assign_symbols, variables::{SO2, SE2}};
/// assign_symbols!(X: SO2; Y: SE2);
/// ```
#[macro_export]
macro_rules! assign_symbols {
    ($($name:ident : $($var:ident),+);* $(;)?) => {$(
        assign_symbols!($name);

        $(
            impl $crate::containers::TypedSymbol<$var> for $name {}
        )*
    )*};

    ($($name:ident),*) => {
        $(
            #[derive(Clone, Copy)]
            pub struct $name(pub u32);

            impl From<$name> for $crate::containers::Key {
                fn from(key: $name) -> $crate::containers::Key {
                    // TODO: Could we compute this char -> int info at compile time?
                    let chr = stringify!($name).chars().next().unwrap();
                    let idx = key.0;
                    $crate::containers::DefaultSymbolHandler::sym_to_key(chr, idx)
                }
            }

            impl std::fmt::Debug for $name {
                fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                    let chr = stringify!($name).chars().next().unwrap();
                    let idx = self.0;
                    $crate::containers::DefaultSymbolHandler::format(f, chr, idx)
                }
            }

            impl $crate::containers::Symbol for $name {}
        )*
    };
}