Edits to color_wheel.rs WIP

color_wheel.rs

/*
 *   Copyright (c) 2023 R3BL LLC
 *   All rights reserved.
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

use std::{borrow::Cow, mem};

use crate::*;

// FIXME: add state in struct to hold the index for current color of the gradient (which is ring buffer)
pub struct ColorWheel;

const START_COLOR: TuiColor = TuiColor::Rgb { r: 255, g: 0, b: 0 };
const END_COLOR: TuiColor = TuiColor::Rgb { r: 0, g: 255, b: 0 };

impl ColorWheel {
    /// Colorizes text in-place if style's lolcat field is true, or leaves text alone if colorize
    /// fails.
    pub fn lolcat_from_style(maybe_style: &Option<Style>, mut_text: &mut Cow<str>) {
        match maybe_style {
            Some(style) if style.lolcat => {
                let colorized_text =
                    ColorWheel::colorize_into_ansi_string(mut_text, &START_COLOR, &END_COLOR);
                if let Ok(mut new_text) = colorized_text {
                    mem::swap(&mut new_text, mut_text.to_mut());
                }
            }
            _ => (),
        }
    }

    /// Returns colorized version of input string, or a copy of the unicode's string if colorize
    /// failed.
    pub fn lolcat_unicode(text: &UnicodeString) -> String {
        match ColorWheel::colorize_into_ansi_string(&text.string, &START_COLOR, &END_COLOR) {
            Ok(new_text) => new_text,
            Err(_) => text.string.clone(),
        }
    }

    /// Returns colorized version of input string, or a copy of the string if colorize failed.
    pub fn lolcat_str(text: &str) -> Cow<'_, str> {
        match ColorWheel::colorize_into_ansi_string(text, &START_COLOR, &END_COLOR) {
            Ok(new_text) => Cow::Owned(new_text),
            Err(_) => Cow::Borrowed(text),
        }
    }

    /// Returns an iterations-sized gradient between start and end colors.
    ///
    /// When possible (iterations > 1), gradient is inclusive of start and end colors. A "single
    /// iteration" gradient is just the start color.
    ///
    /// # Arguments
    /// - `start_color` - RGB start color.
    /// - `end_color` - RGB end color.
    /// - `iterations` - count of elements in result.
    pub fn generate_gradient(
        start_color: &TuiColor,
        end_color: &TuiColor,
        iterations: usize,
    ) -> CommonResult<Vec<TuiColor>> {
        let mut result = Vec::with_capacity(iterations);

        if iterations > 0 {
            let (start_r, start_g, start_b) = ColorWheel::get_rgb_from_color(start_color)?;
            let (end_r, end_g, end_b) = ColorWheel::get_rgb_from_color(end_color)?;

            for step in 1..=iterations {
                result.push(TuiColor::Rgb {
                    r: ColorWheel::tween(start_r, end_r, iterations, step - 1),
                    g: ColorWheel::tween(start_g, end_g, iterations, step - 1),
                    b: ColorWheel::tween(start_b, end_b, iterations, step - 1),
                });
            }
        }

        Ok(result)
    }

    /// Returns a gradient-colored string.
    ///
    /// # Arguments
    /// - `text` - string to color.
    /// - `start_color` - RGB start color.
    /// - `end_color` - RGB end color.
    // FIXME: this should colorize into styled_text
    // FIXME: this should only accept UnicodeString as input
    pub fn colorize_into_ansi_string(
        text: &str,
        start_color: &TuiColor,
        end_color: &TuiColor,
    ) -> CommonResult<String> {
        let mut components = Vec::with_capacity(text.chars().count());

        let gradient = ColorWheel::generate_gradient(start_color, end_color, text.chars().count())?;
        let mut gradient_iter = gradient.iter();

        // FIXME: zip iterators for the gradient & UnicodeString.vec_segments & replace this code below
        text.chars().try_for_each(|x| -> CommonResult<()> {
            // Would like to use "ok_or" to transform Option<TuiColor> into Result<TuiColor>, but this
            // CommonError crap is in my way. let c =
            // i.next().ok_or(CommonError::new(CommonErrorType::StackUnderflow, "Gradient
            // underflow"))?;
            let c = match gradient_iter.next() {
                Some(c) => Ok(c),
                None => CommonError::new(CommonErrorType::StackUnderflow, "Gradient underflow!"),
            }?;

            let (r, g, b) = match c {
                TuiColor::Rgb { r, g, b } => Ok((r, g, b)),
                _ => CommonError::new(
                    CommonErrorType::InvalidValue,
                    "Unable to extract RGB from {c:?}",
                ),
            }?;

            // FIXME: emit styled_text rather than ANSI
            components.push(format!("\x1b[38;2;{};{};{}m{x}", r, g, b));
            Ok(())
        })?;

        if !components.is_empty() {
            // Suffix a reset.
            components.push("\x1b[0m".to_string());
        }

        Ok(components.join(""))
    }

    /// Returns value for given "step" between "start" and "end" arguments.
    ///
    /// There is no constraint on the ordering of "start" and "end" arguments.
    ///
    /// # Arguments
    /// - `start` - inclusive starting value.
    /// - `end` - inclusive ending value.
    /// - `steps` - total number of distinct steps between a and b, should be > 0.
    /// - `step` - requested value calculation.  Should be 0 <= step < steps.
    fn tween(start: u8, end: u8, steps: usize, step: usize) -> u8 {
        (start as f32 + ((end as f32 - start as f32) / (steps - 1) as f32) * step as f32).round()
            as u8
    }

    // FIXME: add support for "named" colors
    /// Returns RGB components from a [TuiColor].
    ///
    /// Currently pukes on non-RGB (eg, "Named") Colors.
    ///
    /// # Arguments
    /// - `color` - an RGB Color.
    fn get_rgb_from_color(color: &TuiColor) -> CommonResult<(u8, u8, u8)> {
        match color {
            TuiColor::Rgb {
                r: ar,
                g: ag,
                b: ab,
            } => Ok((*ar, *ag, *ab)),
            _ => {
                CommonError::new(CommonErrorType::InvalidValue, "Pass me a real color")
                // Ok((255,255,255))
            }
        }
    }
}

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

    const BLACK: TuiColor = TuiColor::Rgb { r: 0, g: 0, b: 0 };
    const GRAY: TuiColor = TuiColor::Rgb {
        r: 128,
        g: 128,
        b: 128,
    };
    const WHITE: TuiColor = TuiColor::Rgb {
        r: 255,
        g: 255,
        b: 255,
    };

    #[test]
    fn generates_empty_gradient() -> CommonResult<()> {
        // Providing "0" steps results in empty result.
        assert!(ColorWheel::generate_gradient(&BLACK, &WHITE, 0)?.is_empty());
        Ok(())
    }

    #[test]
    fn generates_single_element_gradient_of_start_color() -> CommonResult<()> {
        // Providing "1" step results in a single element result of the "start" color.
        assert_eq!(
            ColorWheel::generate_gradient(&BLACK, &WHITE, 1)?,
            vec!(BLACK)
        );
        Ok(())
    }

    #[test]
    fn generates_two_element_inclusive_gradient() -> CommonResult<()> {
        // Providing "2" steps generates an inclusive gradient of the "start" and "end" colors.
        assert_eq!(
            ColorWheel::generate_gradient(&BLACK, &WHITE, 2)?,
            vec!(BLACK, WHITE)
        );
        Ok(())
    }

    #[test]
    fn generates_three_element_inclusive_gradient() -> CommonResult<()> {
        // Demonstrates consistent gradient between start and end colors.
        assert_eq!(
            ColorWheel::generate_gradient(&BLACK, &WHITE, 3)?,
            vec!(BLACK, GRAY, WHITE)
        );
        Ok(())
    }

    #[test]
    fn empty_colorize_result_from_empty_input() -> CommonResult<()> {
        // Colorizing an empty string results in an empty string.
        assert_eq!(
            ColorWheel::colorize_into_ansi_string("", &BLACK, &WHITE)?,
            ""
        );
        Ok(())
    }

    #[test]
    fn single_glyph_colorize_result_colored_start() -> CommonResult<()> {
        // "0;0;0" is BLACK, and the bit after "A" ("\x1b[0m") is "reset".
        assert_eq!(
            ColorWheel::colorize_into_ansi_string("A", &BLACK, &WHITE)?,
            "\x1b[38;2;0;0;0mA\x1b[0m"
        );
        Ok(())
    }

    #[test]
    fn two_glyph_colorize_result_colored_inclusively() -> CommonResult<()> {
        // "0;0;0" is BLACK, "255;255;255" is WHITE
        assert_eq!(
            ColorWheel::colorize_into_ansi_string("AB", &BLACK, &WHITE)?,
            "\x1b[38;2;0;0;0mA\x1b[38;2;255;255;255mB\x1b[0m"
        );
        Ok(())
    }

    #[test]
    fn three_glyph_colorize_result_colored_continuously() -> CommonResult<()> {
        // "0;0;0" is BLACK, "128;128;128" is in the middle, "255;255;255" is WHITE
        assert_eq!(
            ColorWheel::colorize_into_ansi_string("ABC", &BLACK, &WHITE)?,
            "\x1b[38;2;0;0;0mA\x1b[38;2;128;128;128mB\x1b[38;2;255;255;255mC\x1b[0m"
        );
        Ok(())
    }

    #[test]
    fn foreign_glyphs_colorized_correctly() -> CommonResult<()> {
        // Don't know what 226 represents here (its all Chinese to me), but hopefully it is
        // consistent.
        assert_eq!(
            ColorWheel::colorize_into_ansi_string("我想要一个很酷的中文纹身", &BLACK, &WHITE)?
                .chars()
                .count(),
            226
        );
        Ok(())
    }

    #[test]
    fn colorize_pukes_on_non_rgb_color_arg() {
        // FIXME: honor non-RGB colors.
        assert_eq!(
            ColorWheel::colorize_into_ansi_string("Nope", &BLACK, &TuiColor::Green)
                .err()
                .unwrap()
                .downcast::<CommonError>()
                .unwrap()
                .err_msg
                .unwrap(),
            "Pass me a real color",
        );
    }
}