1use std::{
2 fmt,
3 iter::FusedIterator,
4 ops::{Add, ControlFlow, Div, Mul, Sub},
5};
6
7#[cfg(feature = "bigint")]
8use num_bigint::BigUint;
9
10#[derive(Clone)]
11pub struct Fraction<I> {
12 pub numer: I,
13 pub denom: I,
14}
15
16impl<I> Fraction<I> {
17 fn mediant(&self, other: &Self) -> Self
18 where
19 for<'a> &'a I: Add<&'a I, Output = I>,
20 {
21 Self {
22 numer: &self.numer + &other.numer,
23 denom: &self.denom + &other.denom,
24 }
25 }
26 fn k_mediant(&self, other: &Self, k: &I) -> Self
27 where
28 for<'a> &'a I: Add<&'a I, Output = I> + Mul<&'a I, Output = I>,
29 {
30 let tmp = Self { numer: &other.numer * k, denom: &other.denom * k };
31 self.mediant(&tmp)
32 }
33}
34
35impl<I: fmt::Display> fmt::Display for Fraction<I> {
36 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37 write!(f, "{}/{}", self.numer, self.denom)
38 }
39}
40
41impl<I: fmt::Display> fmt::Debug for Fraction<I> {
42 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43 write!(f, "{}/{}", self.numer, self.denom)
44 }
45}
46
47#[derive(Clone, Debug)]
48pub enum ApproxBound<T> {
49 Lower(T),
50 Upper(T),
51}
52
53impl<T> ApproxBound<T> {
54 pub fn as_ref(&self) -> &T {
55 match self {
56 Self::Lower(x) | Self::Upper(x) => x,
57 }
58 }
59 pub fn into_inner(self) -> T {
60 match self {
61 Self::Lower(x) | Self::Upper(x) => x,
62 }
63 }
64}
65
66pub struct SbTreeUnsigned<I, F> {
67 lower: Fraction<I>,
68 upper: Fraction<I>,
69 pred: F,
70 bound: I,
71}
72
73impl<I: SbUnsignedInt, F> SbTreeUnsigned<I, F>
74where
75 for<'a> &'a I: Add<&'a I, Output = I>
76 + Sub<&'a I, Output = I>
77 + Mul<&'a I, Output = I>
78 + Div<&'a I, Output = I>,
79 F: FnMut(&Fraction<I>) -> bool,
80{
81 fn new(pred: F, bound: I) -> Self {
82 Self {
83 lower: I::frac_0(),
84 upper: I::frac_oo(),
85 pred,
86 bound,
87 }
88 }
89 fn next_approx(
90 &mut self,
91 ) -> ControlFlow<ApproxBound<Fraction<I>>, ApproxBound<Fraction<I>>> {
92 let pred = &mut self.pred;
93 let cur = self.lower.mediant(&self.upper);
94 assert!(cur.denom <= self.bound);
95
96 let init_tf = pred(&cur);
97 let (from, to) = if init_tf {
98 (&self.lower, &self.upper)
99 } else {
100 (&self.upper, &self.lower)
101 };
102
103 let mut lo = I::const_1();
104 let mut hi = I::const_2();
105 while pred(&(from.k_mediant(&to, &hi))) == init_tf {
106 lo = &lo + &lo;
107 hi = &hi + &hi;
108 let tmp = from.k_mediant(&to, &lo);
109 if tmp.denom > self.bound {
110 let k = &(&self.bound - &from.denom) / &to.denom;
112 let res = from.k_mediant(&to, &k);
113 return if init_tf {
114 self.lower = res.clone();
115 ControlFlow::Break(ApproxBound::Lower(res))
116 } else {
117 self.upper = res.clone();
118 ControlFlow::Break(ApproxBound::Upper(res))
119 };
120 }
121 }
122
123 while &hi - &lo > I::const_1() {
124 let mid = &lo + &(&(&hi - &lo) / &I::const_2());
125 let tmp = from.k_mediant(&to, &mid);
126 let cur = tmp.denom <= self.bound && pred(&tmp) == init_tf;
127 *(if cur { &mut lo } else { &mut hi }) = mid;
128 }
129
130 let next = from.k_mediant(&to, &lo);
131 let res = if init_tf {
132 self.lower = next.clone();
133 ApproxBound::Lower(next)
134 } else {
135 self.upper = next.clone();
136 ApproxBound::Upper(next)
137 };
138 if self.lower.mediant(&self.upper).denom <= self.bound {
139 ControlFlow::Continue(res)
140 } else {
141 ControlFlow::Break(res)
142 }
143 }
144}
145
146impl<I: SbUnsignedInt, F> Iterator for SbTreeUnsigned<I, F>
147where
148 for<'a> &'a I: Add<&'a I, Output = I>
149 + Sub<&'a I, Output = I>
150 + Mul<&'a I, Output = I>
151 + Div<&'a I, Output = I>,
152 F: FnMut(&Fraction<I>) -> bool,
153{
154 type Item = ControlFlow<ApproxBound<Fraction<I>>, ApproxBound<Fraction<I>>>;
155
156 fn next(&mut self) -> Option<Self::Item> {
157 (self.lower.mediant(&self.upper).denom <= self.bound)
158 .then(|| self.next_approx())
159 }
160}
161
162impl<I: SbUnsignedInt, F> FusedIterator for SbTreeUnsigned<I, F>
163where
164 for<'a> &'a I: Add<&'a I, Output = I>
165 + Sub<&'a I, Output = I>
166 + Mul<&'a I, Output = I>
167 + Div<&'a I, Output = I>,
168 F: FnMut(&Fraction<I>) -> bool,
169{
170}
171
172pub trait FracApprox<F>: Sized {
173 fn approx(self, pred: F) -> [ApproxBound<Fraction<Self>>; 2];
174 fn approx_iter(self, pred: F) -> SbTreeUnsigned<Self, F>;
175}
176
177impl<I: SbUnsignedInt, F> FracApprox<F> for I
178where
179 for<'a> &'a I: Add<&'a I, Output = I>
180 + Sub<&'a I, Output = I>
181 + Mul<&'a I, Output = I>
182 + Div<&'a I, Output = I>,
183 F: FnMut(&Fraction<I>) -> bool,
184{
185 fn approx(self, pred: F) -> [ApproxBound<Fraction<I>>; 2] {
186 use ApproxBound::*;
187 let mut lower = Lower(I::frac_0());
188 let mut upper = Upper(I::frac_oo());
189 for x in self.approx_iter(pred) {
190 let x = match x {
191 ControlFlow::Continue(x) | ControlFlow::Break(x) => x,
192 };
193 match x {
194 Lower(x) => lower = Lower(x),
195 Upper(x) => upper = Upper(x),
196 }
197 }
198 [lower, upper]
199 }
200 fn approx_iter(self, pred: F) -> SbTreeUnsigned<Self, F> {
201 SbTreeUnsigned::new(pred, self)
202 }
203}
204
205pub trait SbUnsignedInt: Clone + Ord {
206 fn const_0() -> Self;
207 fn const_1() -> Self;
208 fn const_2() -> Self;
209 fn frac_0() -> Fraction<Self> {
210 Fraction { numer: Self::const_0(), denom: Self::const_1() }
211 }
212 fn frac_oo() -> Fraction<Self> {
213 Fraction { numer: Self::const_1(), denom: Self::const_0() }
214 }
215}
216
217macro_rules! impl_uint {
218 ( $($ty:ty)* ) => { $(
219 impl SbUnsignedInt for $ty {
220 fn const_0() -> $ty { 0 }
221 fn const_1() -> $ty { 1 }
222 fn const_2() -> $ty { 2 }
223 }
224 )* }
225}
226
227impl_uint! { u8 u16 u32 u64 u128 usize }
228
229#[cfg(feature = "bigint")]
230impl SbUnsignedInt for BigUint {
231 fn const_0() -> Self { Self::from(0_u32) }
232 fn const_1() -> Self { Self::from(1_u32) }
233 fn const_2() -> Self { Self::from(2_u32) }
234}