1use crate::{
22    calendar_arithmetic::{ArithmeticDate, CalendarArithmetic, PrecomputedDataSource},
23    provider::chinese_based::{ChineseBasedCacheV1, PackedChineseBasedYearInfo},
24    types::{FormattableMonth, MonthCode},
25    Calendar, CalendarError, Iso,
26};
27
28use calendrical_calculations::chinese_based::{self, ChineseBased, YearBounds};
29use calendrical_calculations::rata_die::RataDie;
30use core::marker::PhantomData;
31use core::num::NonZeroU8;
32use tinystr::tinystr;
33
34pub(crate) trait ChineseBasedWithDataLoading: Calendar {
38    type CB: ChineseBased;
39    fn get_precomputed_data(&self) -> ChineseBasedPrecomputedData<'_, Self::CB>;
42}
43
44#[derive(Debug, Eq, PartialEq, PartialOrd, Ord)]
46pub(crate) struct ChineseBasedDateInner<C: CalendarArithmetic>(pub(crate) ArithmeticDate<C>);
47
48impl<C: CalendarArithmetic> Copy for ChineseBasedDateInner<C> {}
50impl<C: CalendarArithmetic> Clone for ChineseBasedDateInner<C> {
51    fn clone(&self) -> Self {
52        *self
53    }
54}
55
56#[derive(Default)]
59pub(crate) struct ChineseBasedPrecomputedData<'a, CB: ChineseBased> {
60    data: Option<&'a ChineseBasedCacheV1<'a>>,
61    _cb: PhantomData<CB>,
62}
63
64fn compute_cache<CB: ChineseBased>(extended_year: i32) -> ChineseBasedYearInfo {
66    let mid_year = chinese_based::fixed_mid_year_from_year::<CB>(extended_year);
67    let year_bounds = YearBounds::compute::<CB>(mid_year);
68    compute_cache_with_yb::<CB>(extended_year, year_bounds)
69}
70
71fn compute_cache_with_yb<CB: ChineseBased>(
73    extended_year: i32,
74    year_bounds: YearBounds,
75) -> ChineseBasedYearInfo {
76    let YearBounds { new_year, .. } = year_bounds;
77
78    let days_in_prev_year = chinese_based::days_in_prev_year::<CB>(new_year);
79
80    let packed_data = compute_packed_with_yb::<CB>(extended_year, year_bounds);
81
82    ChineseBasedYearInfo {
83        days_in_prev_year,
84        packed_data,
85    }
86}
87
88fn compute_packed_with_yb<CB: ChineseBased>(
89    extended_year: i32,
90    year_bounds: YearBounds,
91) -> PackedChineseBasedYearInfo {
92    let YearBounds {
93        new_year,
94        next_new_year,
95        ..
96    } = year_bounds;
97    let (month_lengths, leap_month) =
98        chinese_based::month_structure_for_year::<CB>(new_year, next_new_year);
99
100    let related_iso = CB::iso_from_extended(extended_year);
101    let iso_ny = calendrical_calculations::iso::fixed_from_iso(related_iso, 1, 1);
102
103    let ny_offset = new_year - iso_ny - i64::from(PackedChineseBasedYearInfo::FIRST_NY) + 1;
105    let ny_offset = if let Ok(ny_offset) = u8::try_from(ny_offset) {
106        ny_offset
107    } else {
108        debug_assert!(
109            false,
110            "Expected small new years offset, got {ny_offset} in ISO year {related_iso}"
111        );
112        0
113    };
114    PackedChineseBasedYearInfo::new(month_lengths, leap_month, ny_offset)
115}
116
117#[cfg(feature = "datagen")]
118pub(crate) fn compute_many_packed<CB: ChineseBased>(
119    extended_years: core::ops::Range<i32>,
120) -> alloc::vec::Vec<PackedChineseBasedYearInfo> {
121    extended_years
122        .map(|extended_year| {
123            let mid_year = chinese_based::fixed_mid_year_from_year::<CB>(extended_year);
124            let year_bounds = YearBounds::compute::<CB>(mid_year);
125
126            compute_packed_with_yb::<CB>(extended_year, year_bounds)
127        })
128        .collect()
129}
130
131impl<'b, CB: ChineseBased> PrecomputedDataSource<ChineseBasedYearInfo>
132    for ChineseBasedPrecomputedData<'b, CB>
133{
134    fn load_or_compute_info(&self, extended_year: i32) -> ChineseBasedYearInfo {
135        self.data
136            .and_then(|d| d.get_for_extended_year(extended_year))
137            .unwrap_or_else(|| compute_cache::<CB>(extended_year))
138    }
139}
140
141impl<'b, CB: ChineseBased> ChineseBasedPrecomputedData<'b, CB> {
142    pub(crate) fn new(data: Option<&'b ChineseBasedCacheV1<'b>>) -> Self {
143        Self {
144            data,
145            _cb: PhantomData,
146        }
147    }
148    fn load_or_compute_info_for_iso(
151        &self,
152        fixed: RataDie,
153        iso: ArithmeticDate<Iso>,
154    ) -> (ChineseBasedYearInfo, i32) {
155        let cached = self.data.and_then(|d| d.get_for_iso::<CB>(iso));
156        if let Some(cached) = cached {
157            return cached;
158        };
159        let extended_year = CB::extended_from_iso(iso.year);
162        let mid_year = chinese_based::fixed_mid_year_from_year::<CB>(extended_year);
163        let year_bounds = YearBounds::compute::<CB>(mid_year);
164        let YearBounds { new_year, .. } = year_bounds;
165        if fixed >= new_year {
166            (
167                compute_cache_with_yb::<CB>(extended_year, year_bounds),
168                extended_year,
169            )
170        } else {
171            let extended_year = extended_year - 1;
172            (compute_cache::<CB>(extended_year), extended_year)
173        }
174    }
175}
176#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord)]
178pub(crate) struct ChineseBasedYearInfo {
180    days_in_prev_year: u16,
181    packed_data: PackedChineseBasedYearInfo,
186}
187
188impl ChineseBasedYearInfo {
189    pub(crate) fn new(days_in_prev_year: u16, packed_data: PackedChineseBasedYearInfo) -> Self {
190        Self {
191            days_in_prev_year,
192            packed_data,
193        }
194    }
195
196    pub(crate) fn new_year<CB: ChineseBased>(self, extended_year: i32) -> RataDie {
198        self.packed_data.ny_rd(CB::iso_from_extended(extended_year))
199    }
200
201    fn next_new_year<CB: ChineseBased>(self, extended_year: i32) -> RataDie {
204        self.new_year::<CB>(extended_year) + i64::from(self.packed_data.days_in_year())
205    }
206
207    pub(crate) fn leap_month(self) -> Option<NonZeroU8> {
212        self.packed_data.leap_month_idx()
213    }
214
215    fn last_day_of_previous_month(self, month: u8) -> u16 {
222        debug_assert!((1..=13).contains(&month), "Month out of bounds!");
223        if month == 1 {
226            0
227        } else {
228            self.packed_data.last_day_of_month(month - 1)
229        }
230    }
231
232    fn days_in_year(self) -> u16 {
233        self.packed_data.days_in_year()
234    }
235
236    fn days_in_prev_year(self) -> u16 {
237        self.days_in_prev_year
238    }
239
240    fn last_day_of_month(self, month: u8) -> u16 {
247        debug_assert!((1..=13).contains(&month), "Month out of bounds!");
248        self.packed_data.last_day_of_month(month)
249    }
250
251    fn days_in_month(self, month: u8) -> u8 {
252        let ret =
253            u8::try_from(self.last_day_of_month(month) - self.last_day_of_previous_month(month));
254        debug_assert!(ret.is_ok(), "Month too big!");
255        ret.unwrap_or(30)
256    }
257}
258
259impl<C: ChineseBasedWithDataLoading + CalendarArithmetic<YearInfo = ChineseBasedYearInfo>>
260    ChineseBasedDateInner<C>
261{
262    fn chinese_based_date_from_info(
264        date: RataDie,
265        year_info: ChineseBasedYearInfo,
266        extended_year: i32,
267    ) -> ChineseBasedDateInner<C> {
268        debug_assert!(
269            date < year_info.next_new_year::<C::CB>(extended_year),
270            "Stored date {date:?} out of bounds!"
271        );
272        let day_of_year = u16::try_from(date - year_info.new_year::<C::CB>(extended_year) + 1);
274        debug_assert!(day_of_year.is_ok(), "Somehow got a very large year in data");
275        let day_of_year = day_of_year.unwrap_or(1);
276        let mut month = 1;
277        for iter_month in 1..=13 {
279            month = iter_month;
280            if year_info.last_day_of_month(iter_month) >= day_of_year {
281                break;
282            }
283        }
284
285        debug_assert!((1..=13).contains(&month), "Month out of bounds!");
286
287        debug_assert!(
288            month < 13 || year_info.leap_month().is_some(),
289            "Cannot have 13 months in a non-leap year!"
290        );
291        let day_before_month_start = year_info.last_day_of_previous_month(month);
292        let day_of_month = day_of_year - day_before_month_start;
293        let day_of_month = u8::try_from(day_of_month);
294        debug_assert!(day_of_month.is_ok(), "Month too big!");
295        let day_of_month = day_of_month.unwrap_or(1);
296
297        ChineseBasedDateInner(ArithmeticDate::new_unchecked_with_info(
302            extended_year,
303            month,
304            day_of_month,
305            year_info,
306        ))
307    }
308
309    pub(crate) fn chinese_based_date_from_fixed(
312        cal: &C,
313        fixed: RataDie,
314        iso: ArithmeticDate<Iso>,
315    ) -> ChineseBasedDateInner<C> {
316        let data = cal.get_precomputed_data();
317
318        let (year_info, extended_year) = data.load_or_compute_info_for_iso(fixed, iso);
319
320        Self::chinese_based_date_from_info(fixed, year_info, extended_year)
321    }
322
323    pub(crate) fn new_year(self) -> RataDie {
324        self.0.year_info.new_year::<C::CB>(self.0.year)
325    }
326
327    pub(crate) fn fixed_from_chinese_based_date_inner(date: ChineseBasedDateInner<C>) -> RataDie {
332        let first_day_of_year = date.new_year();
333        let day_of_year = date.day_of_year(); first_day_of_year + i64::from(day_of_year) - 1
335    }
336
337    pub(crate) fn new_from_ordinals(
341        year: i32,
342        month: u8,
343        day: u8,
344        year_info: ChineseBasedYearInfo,
345    ) -> Result<ArithmeticDate<C>, CalendarError> {
346        let max_month = Self::months_in_year_with_info(year_info);
347        if !(1..=max_month).contains(&month) {
348            return Err(CalendarError::Overflow {
349                field: "month",
350                max: max_month as usize,
351            });
352        }
353
354        let max_day = year_info.days_in_month(month);
355        if day > max_day {
356            return Err(CalendarError::Overflow {
357                field: "day",
358                max: max_day as usize,
359            });
360        }
361
362        Ok(ArithmeticDate::<C>::new_unchecked_with_info(
364            year, month, day, year_info,
365        ))
366    }
367
368    pub(crate) fn months_in_year_inner(&self) -> u8 {
370        Self::months_in_year_with_info(self.0.year_info)
371    }
372
373    fn months_in_year_with_info(year_info: ChineseBasedYearInfo) -> u8 {
376        if year_info.leap_month().is_some() {
377            13
378        } else {
379            12
380        }
381    }
382
383    pub(crate) fn days_in_month_inner(&self) -> u8 {
385        self.0.year_info.days_in_month(self.0.month)
386    }
387
388    pub(crate) fn fixed_mid_year_from_year(year: i32) -> RataDie {
389        chinese_based::fixed_mid_year_from_year::<C::CB>(year)
390    }
391
392    pub(crate) fn days_in_year_inner(&self) -> u16 {
394        self.0.year_info.days_in_year()
395    }
396    pub(crate) fn days_in_prev_year(&self) -> u16 {
398        self.0.year_info.days_in_prev_year()
399    }
400
401    pub(crate) fn day_of_year(&self) -> u16 {
404        self.0.year_info.last_day_of_previous_month(self.0.month) + u16::from(self.0.day)
405    }
406
407    pub(crate) fn month(&self) -> FormattableMonth {
412        let ordinal = self.0.month;
413        let leap_month_option = self.0.year_info.leap_month();
414
415        let leap_month = if let Some(leap) = leap_month_option {
419            leap.get()
420        } else {
421            14
423        };
424        let code_inner = if leap_month == ordinal {
425            debug_assert!((2..=13).contains(&ordinal));
428            match ordinal {
429                2 => tinystr!(4, "M01L"),
430                3 => tinystr!(4, "M02L"),
431                4 => tinystr!(4, "M03L"),
432                5 => tinystr!(4, "M04L"),
433                6 => tinystr!(4, "M05L"),
434                7 => tinystr!(4, "M06L"),
435                8 => tinystr!(4, "M07L"),
436                9 => tinystr!(4, "M08L"),
437                10 => tinystr!(4, "M09L"),
438                11 => tinystr!(4, "M10L"),
439                12 => tinystr!(4, "M11L"),
440                13 => tinystr!(4, "M12L"),
441                _ => tinystr!(4, "und"),
442            }
443        } else {
444            let mut adjusted_ordinal = ordinal;
445            if ordinal > leap_month {
446                debug_assert!((2..=13).contains(&ordinal));
451                adjusted_ordinal -= 1;
452            }
453            debug_assert!((1..=12).contains(&adjusted_ordinal));
454            match adjusted_ordinal {
455                1 => tinystr!(4, "M01"),
456                2 => tinystr!(4, "M02"),
457                3 => tinystr!(4, "M03"),
458                4 => tinystr!(4, "M04"),
459                5 => tinystr!(4, "M05"),
460                6 => tinystr!(4, "M06"),
461                7 => tinystr!(4, "M07"),
462                8 => tinystr!(4, "M08"),
463                9 => tinystr!(4, "M09"),
464                10 => tinystr!(4, "M10"),
465                11 => tinystr!(4, "M11"),
466                12 => tinystr!(4, "M12"),
467                _ => tinystr!(4, "und"),
468            }
469        };
470        let code = MonthCode(code_inner);
471        FormattableMonth {
472            ordinal: ordinal as u32,
473            code,
474        }
475    }
476}
477
478impl<C: ChineseBasedWithDataLoading> CalendarArithmetic for C {
479    type YearInfo = ChineseBasedYearInfo;
480
481    fn month_days(_year: i32, month: u8, year_info: ChineseBasedYearInfo) -> u8 {
482        year_info.days_in_month(month)
483    }
484
485    fn months_for_every_year(_year: i32, year_info: ChineseBasedYearInfo) -> u8 {
487        if year_info.leap_month().is_some() {
488            13
489        } else {
490            12
491        }
492    }
493
494    fn is_leap_year(_year: i32, year_info: ChineseBasedYearInfo) -> bool {
496        year_info.leap_month().is_some()
497    }
498
499    fn last_month_day_in_year(_year: i32, year_info: ChineseBasedYearInfo) -> (u8, u8) {
504        if year_info.leap_month().is_some() {
505            (13, year_info.days_in_month(13))
506        } else {
507            (12, year_info.days_in_month(12))
508        }
509    }
510
511    fn days_in_provided_year(_year: i32, year_info: ChineseBasedYearInfo) -> u16 {
512        year_info.last_day_of_month(13)
513    }
514}
515
516pub(crate) fn chinese_based_ordinal_lunar_month_from_code(
518    code: MonthCode,
519    year_info: ChineseBasedYearInfo,
520) -> Option<u8> {
521    let leap_month = if let Some(leap) = year_info.leap_month() {
522        leap.get()
523    } else {
524        14
527    };
528
529    if code.0.len() < 3 {
530        return None;
531    }
532    let bytes = code.0.all_bytes();
533    if bytes[0] != b'M' {
534        return None;
535    }
536    if code.0.len() == 4 && bytes[3] != b'L' {
537        return None;
538    }
539    let mut unadjusted = 0;
541    if bytes[1] == b'0' {
542        if bytes[2] >= b'1' && bytes[2] <= b'9' {
543            unadjusted = bytes[2] - b'0';
544        }
545    } else if bytes[1] == b'1' && bytes[2] >= b'0' && bytes[2] <= b'2' {
546        unadjusted = 10 + bytes[2] - b'0';
547    }
548    if bytes[3] == b'L' {
549        if unadjusted + 1 != leap_month {
551            return None;
552        } else {
553            return Some(unadjusted + 1);
555        }
556    }
557    if unadjusted != 0 {
558        if unadjusted + 1 > leap_month {
561            return Some(unadjusted + 1);
562        } else {
563            return Some(unadjusted);
564        }
565    }
566    None
567}
568
569#[cfg(test)]
570mod test {
571    use super::*;
572
573    fn packed_roundtrip_single(
574        mut month_lengths: [bool; 13],
575        leap_month_idx: Option<NonZeroU8>,
576        ny_offset: u8,
577    ) {
578        if leap_month_idx.is_none() {
579            month_lengths[12] = false;
581        }
582        let packed = PackedChineseBasedYearInfo::new(month_lengths, leap_month_idx, ny_offset);
583
584        assert_eq!(
585            ny_offset,
586            packed.ny_offset(),
587            "Roundtrip with {month_lengths:?}, {leap_month_idx:?}, {ny_offset}"
588        );
589        assert_eq!(
590            leap_month_idx,
591            packed.leap_month_idx(),
592            "Roundtrip with {month_lengths:?}, {leap_month_idx:?}, {ny_offset}"
593        );
594        let mut month_lengths_roundtrip = [false; 13];
595        for (i, len) in month_lengths_roundtrip.iter_mut().enumerate() {
596            *len = packed.month_has_30_days(i as u8 + 1);
597        }
598        assert_eq!(
599            month_lengths, month_lengths_roundtrip,
600            "Roundtrip with {month_lengths:?}, {leap_month_idx:?}, {ny_offset}"
601        );
602    }
603
604    #[test]
605    fn test_roundtrip_packed() {
606        const SHORT: [bool; 13] = [false; 13];
607        const LONG: [bool; 13] = [true; 13];
608        const ALTERNATING1: [bool; 13] = [
609            false, true, false, true, false, true, false, true, false, true, false, true, false,
610        ];
611        const ALTERNATING2: [bool; 13] = [
612            true, false, true, false, true, false, true, false, true, false, true, false, true,
613        ];
614        const RANDOM1: [bool; 13] = [
615            true, true, false, false, true, true, false, true, true, true, true, false, true,
616        ];
617        const RANDOM2: [bool; 13] = [
618            false, true, true, true, true, false, true, true, true, false, false, true, false,
619        ];
620        packed_roundtrip_single(SHORT, None, 5);
621        packed_roundtrip_single(SHORT, None, 10);
622        packed_roundtrip_single(SHORT, NonZeroU8::new(11), 15);
623        packed_roundtrip_single(LONG, NonZeroU8::new(12), 15);
624        packed_roundtrip_single(ALTERNATING1, None, 2);
625        packed_roundtrip_single(ALTERNATING1, NonZeroU8::new(3), 5);
626        packed_roundtrip_single(ALTERNATING2, None, 9);
627        packed_roundtrip_single(ALTERNATING2, NonZeroU8::new(7), 26);
628        packed_roundtrip_single(RANDOM1, None, 29);
629        packed_roundtrip_single(RANDOM1, NonZeroU8::new(12), 29);
630        packed_roundtrip_single(RANDOM1, NonZeroU8::new(2), 21);
631        packed_roundtrip_single(RANDOM2, None, 25);
632        packed_roundtrip_single(RANDOM2, NonZeroU8::new(2), 19);
633        packed_roundtrip_single(RANDOM2, NonZeroU8::new(5), 2);
634        packed_roundtrip_single(RANDOM2, NonZeroU8::new(12), 5);
635    }
636}