khadhroony-solana-project/crates/ksp-logging-lib/src/writer.rs

// file: crates/ksp-logging-lib/src/writer.rs
// version: 5

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum StripAnsiState {
    Text,
    Escape,
    Csi,
    Osc,
    OscEscape,
    String,
    StringEscape,
}

/// Crate-internal `StripAnsiWriter` state shared across the owning crate.
pub(crate) struct StripAnsiWriter<W> {
    inner: W,
    state: StripAnsiState,
}

impl<W> StripAnsiWriter<W> {
    /// Creates a new `StripAnsiWriter` value.
    pub(crate) const fn new(inner: W) -> Self {
        return Self { inner, state: StripAnsiState::Text };
    }

    #[cfg(test)]
    fn into_inner(self) -> W {
        return self.inner;
    }
}

impl<W> std::io::Write for StripAnsiWriter<W>
where
    W: std::io::Write,
{
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let mut stripped = std::vec::Vec::with_capacity(buf.len());
        for byte in buf {
            self.consume_byte(*byte, &mut stripped);
        }
        let write_result = std::io::Write::write_all(&mut self.inner, stripped.as_slice());
        return match write_result {
            std::result::Result::Ok(()) => std::result::Result::Ok(buf.len()),
            std::result::Result::Err(error) => std::result::Result::Err(error),
        };
    }

    fn flush(&mut self) -> std::io::Result<()> {
        return std::io::Write::flush(&mut self.inner);
    }
}

impl<W> StripAnsiWriter<W> {
    fn consume_byte(&mut self, byte: u8, output: &mut std::vec::Vec<u8>) {
        self.state = match self.state {
            StripAnsiState::Text => {
                if byte == 0x1B {
                    StripAnsiState::Escape
                } else {
                    output.push(byte);
                    StripAnsiState::Text
                }
            },
            StripAnsiState::Escape => match byte {
                b'[' => StripAnsiState::Csi,
                b']' => StripAnsiState::Osc,
                b'P' | b'X' | b'^' | b'_' => StripAnsiState::String,
                0x1B => StripAnsiState::Escape,
                _ => StripAnsiState::Text,
            },
            StripAnsiState::Csi => {
                if (0x40..=0x7E).contains(&byte) {
                    StripAnsiState::Text
                } else {
                    StripAnsiState::Csi
                }
            },
            StripAnsiState::Osc => match byte {
                0x07 => StripAnsiState::Text,
                0x1B => StripAnsiState::OscEscape,
                _ => StripAnsiState::Osc,
            },
            StripAnsiState::OscEscape => match byte {
                b'\\' => StripAnsiState::Text,
                0x1B => StripAnsiState::OscEscape,
                _ => StripAnsiState::Osc,
            },
            StripAnsiState::String => {
                if byte == 0x1B {
                    StripAnsiState::StringEscape
                } else {
                    StripAnsiState::String
                }
            },
            StripAnsiState::StringEscape => match byte {
                b'\\' => StripAnsiState::Text,
                0x1B => StripAnsiState::StringEscape,
                _ => StripAnsiState::String,
            },
        };
    }
}

/// Crate-internal `RouteMakeWriter` state shared across the owning crate.
#[derive(Clone)]
pub(crate) struct RouteMakeWriter<W> {
    inner: W,
    filter: crate::OutputFilter,
}

impl<W> RouteMakeWriter<W> {
    /// Creates a new `RouteMakeWriter` value.
    pub(crate) fn new(inner: W, filter: crate::OutputFilter) -> Self {
        return Self { inner, filter };
    }
}

/// Crate-internal `RoutedWriter` variants used by the owning crate.
pub(crate) enum RoutedWriter<W> {
    Enabled(W),
    Disabled,
}

impl<W> std::io::Write for crate::RoutedWriter<W>
where
    W: std::io::Write,
{
    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
        return match self {
            Self::Enabled(writer) => std::io::Write::write(writer, buffer),
            Self::Disabled => std::result::Result::Ok(buffer.len()),
        };
    }

    fn flush(&mut self) -> std::io::Result<()> {
        return match self {
            Self::Enabled(writer) => std::io::Write::flush(writer),
            Self::Disabled => std::result::Result::Ok(()),
        };
    }
}

impl<'writer, W> tracing_subscriber::fmt::MakeWriter<'writer> for RouteMakeWriter<W>
where
    W: tracing_subscriber::fmt::MakeWriter<'writer>,
{
    type Writer = crate::RoutedWriter<W::Writer>;

    fn make_writer(&'writer self) -> Self::Writer {
        return crate::RoutedWriter::Enabled(tracing_subscriber::fmt::MakeWriter::make_writer(&self.inner));
    }

    fn make_writer_for(&'writer self, metadata: &tracing::Metadata<'_>) -> Self::Writer {
        if metadata_matches_filter(metadata, &self.filter) {
            return crate::RoutedWriter::Enabled(tracing_subscriber::fmt::MakeWriter::make_writer_for(&self.inner, metadata));
        }
        return crate::RoutedWriter::Disabled;
    }
}

fn metadata_matches_filter(metadata: &tracing::Metadata<'_>, filter: &crate::OutputFilter) -> bool {
    if !level_is_enabled(metadata.level(), filter.level()) {
        return false;
    }
    if !filter.targets().iter().any(|selector| -> bool {
        return selector == "*" || metadata.target().starts_with(selector.as_str());
    }) {
        return false;
    }
    return crate::current_domain_matches(filter.domains());
}

fn level_is_enabled(level: &tracing::Level, filter: crate::LogFilterLevel) -> bool {
    return match filter {
        crate::LogFilterLevel::Off => false,
        crate::LogFilterLevel::Error => level_rank(level) <= 1,
        crate::LogFilterLevel::Warn => level_rank(level) <= 2,
        crate::LogFilterLevel::Info => level_rank(level) <= 3,
        crate::LogFilterLevel::Debug => level_rank(level) <= 4,
        crate::LogFilterLevel::Trace => level_rank(level) <= 5,
    };
}

fn level_rank(level: &tracing::Level) -> u8 {
    if level == &tracing::Level::ERROR {
        return 1;
    }
    if level == &tracing::Level::WARN {
        return 2;
    }
    if level == &tracing::Level::INFO {
        return 3;
    }
    if level == &tracing::Level::DEBUG {
        return 4;
    }
    return 5;
}

#[cfg(test)]
#[path = "../unit_tests/writer.rs"]
mod tests;