plenum-bot/src/date.rs

use chrono::{Datelike, Duration, NaiveDate, Weekday};
use nom::branch::alt;
use nom::bytes::complete::{tag, tag_no_case};
use nom::character::complete::{i32, multispace0 as ws};
use nom::combinator::{eof, value, verify};
use nom::multi::separated_list1;
use nom::sequence::{delimited, separated_pair, tuple};
use nom::IResult;
use std::collections::BTreeSet;
use std::error::Error;

/// Defines a day within the month. Negative numbers count from the end.
pub enum DaySpec {
    /// nth day, no matter the weekday
    DayOfMonth(i32),
    /// nth occurrence of a specific weekday
    NthWeekdayOfMonth(i32, Weekday),
}

/* ***** parsing w/ nom ***** */

fn parse_weekday(input: &str) -> IResult<&str, Weekday> {
    alt((
        value(Weekday::Mon, tag_no_case("Mon")),
        value(Weekday::Tue, tag_no_case("Tue")),
        value(Weekday::Wed, tag_no_case("Wed")),
        value(Weekday::Thu, tag_no_case("Thu")),
        value(Weekday::Fri, tag_no_case("Fri")),
        value(Weekday::Sat, tag_no_case("Sat")),
        value(Weekday::Sun, tag_no_case("Sun")),
    ))(input)
}

fn parse_nth_weekday(input: &str) -> IResult<&str, DaySpec> {
    let (remaining, (i, w)) = delimited(
        tuple((tag("("), ws)),
        separated_pair(
            verify(i32, |&i| i >= -5 && i != 0 && i <= 5),
            tuple((tag(","), ws)),
            parse_weekday,
        ),
        tuple((ws, tag(")"))),
    )(input)?;
    Ok((remaining, DaySpec::NthWeekdayOfMonth(i, w)))
}

fn parse_day_of_month(input: &str) -> IResult<&str, DaySpec> {
    let (remaining, i) = verify(i32, |&i| i >= -31 && i != 0 && i <= 31)(input)?;
    Ok((remaining, DaySpec::DayOfMonth(i)))
}

fn parse_day_spec(input: &str) -> IResult<&str, DaySpec> {
    alt((parse_nth_weekday, parse_day_of_month))(input)
}

fn parse_day_spec_list(input: &str) -> IResult<&str, Vec<DaySpec>> {
    delimited(ws, separated_list1(tuple((tag(","), ws)), parse_day_spec), tuple((ws, eof)))(input)
}

pub fn parse_spec(input: &str) -> Result<Vec<DaySpec>, Box<dyn Error>> {
    match parse_day_spec_list(input) {
        Ok((_remainder, days)) => Ok(days),
        Err(e) => Err(format!("{}", e).into()),
    }
}

/* ***** figuring out which days are within a given month ***** */

fn day_of_month(year: i32, month: u32, day: i32) -> Option<NaiveDate> {
    if day > 0 {
        NaiveDate::from_ymd_opt(year, month, day as u32)
    } else {
        NaiveDate::from_ymd_opt(
            if month == 12 { year + 1 } else { year },
            if month == 12 { 1 } else { month },
            1,
        )
        .map(|d| d + Duration::days(day as i64))
    }
}

fn nth_weekday_of_month(year: i32, month: u32, week_day: Weekday, nth: i32) -> Option<NaiveDate> {
    // collect all
    let mut dates = Vec::new();
    let mut date = NaiveDate::from_ymd_opt(year, month, 1)?;
    while date.weekday() != week_day {
        date = date.succ_opt()?;
    }
    while date.month() == month {
        dates.push(date);
        date = date + Duration::weeks(1);
    }
    // get the right one
    let index = if nth < 0 { dates.len() as i32 + nth } else { nth - 1 } as usize;
    dates.get(index).cloned()
}

pub fn get_matching_dates_around(
    reference_date: NaiveDate, day_specs: Vec<DaySpec>,
) -> Vec<NaiveDate> {
    let mut dates = BTreeSet::new();
    // compute current / previous / next month once
    let (cyear, cmonth, _day) =
        (reference_date.year_ce().1 as i32, reference_date.month(), reference_date.day());
    let (pyear, pmonth) = if cmonth == 1 { (cyear - 1, 12) } else { (cyear, cmonth - 1) };
    let (nyear, nmonth) = if cmonth == 12 { (cyear + 1, 1) } else { (cyear, cmonth + 1) };
    // collect all relevant days, deduplicating as we go
    for spec in &day_specs {
        match spec {
            DaySpec::DayOfMonth(day) => {
                dates.insert_opt(day_of_month(pyear, pmonth, *day));
                dates.insert_opt(day_of_month(cyear, cmonth, *day));
                dates.insert_opt(day_of_month(nyear, nmonth, *day));
            },
            DaySpec::NthWeekdayOfMonth(nth, week_day) => {
                dates.insert_opt(nth_weekday_of_month(pyear, pmonth, *week_day, *nth));
                dates.insert_opt(nth_weekday_of_month(cyear, cmonth, *week_day, *nth));
                dates.insert_opt(nth_weekday_of_month(nyear, nmonth, *week_day, *nth));
            },
        }
    }
    // make sorted list
    let mut result: Vec<NaiveDate> = dates.into_iter().collect();
    result.sort_unstable();
    // find closest index & return either nearest 2 or nearest 3 (if reference date is in the list)
    match result.binary_search(&reference_date) {
        Ok(index) => {
            let start = index.saturating_sub(1);
            let end = (index + 2).min(result.len());
            result[start..end].to_vec()
        },
        Err(index) => {
            let start = index.saturating_sub(1);
            let end = (index + 1).min(result.len());
            result[start..end].to_vec()
        },
    }
}

/* ***** more util ***** */

// I don't want to deal with manual unwrapping / matching…
trait InsertOpt<T> {
    fn insert_opt(&mut self, item: Option<T>) -> bool;
}

impl<T: Ord> InsertOpt<T> for BTreeSet<T> {
    fn insert_opt(&mut self, item: Option<T>) -> bool {
        if let Some(value) = item {
            self.insert(value)
        } else {
            false
        }
    }
}