Skip to main content

nekolib/math/
frac_approx.rs

1use std::ops::{Add, AddAssign, Mul, Neg};
2
3pub struct FracApproxIter<I, F> {
4    lower: Fraction<I>,
5    upper: Fraction<I>,
6    bound: Option<I>,
7    pred: F,
8}
9
10impl<I: SbInt, F: Fn(I, I) -> bool> FracApproxIter<I, F> {
11    fn new(pred: F, bound: Option<I>) -> Self {
12        let (lower, upper) = if !I::SIGNED {
13            (I::lowest_frac(), I::highest_frac())
14        } else if pred(I::ZERO, I::ONE) {
15            (I::zero_frac(), I::highest_frac())
16        } else {
17            (I::lowest_frac(), I::zero_frac())
18        };
19        let (lower, upper) = (lower.into(), upper.into());
20        Self { lower, upper, bound, pred }
21    }
22
23    fn iter_return(&mut self, q: I, tf: bool) -> (ApproxFrac<I>, I) {
24        (ApproxFrac::from(((self.lower, self.upper), tf)), q)
25    }
26}
27
28impl<I: SbInt, F: Fn(I, I) -> bool> Iterator for FracApproxIter<I, F> {
29    type Item = (ApproxFrac<I>, I);
30    fn next(&mut self) -> Option<Self::Item> {
31        let Self { lower, upper, bound, pred } = self;
32        let pred = |fr: Fraction<I>| {
33            let (lower, upper) = fr.into();
34            pred(lower, upper)
35        };
36        let median = *lower + *upper;
37
38        if median.is_deeper(*bound) {
39            return None;
40        }
41
42        let tf = pred(median);
43        let (from, to) = if tf { (*lower, *upper) } else { (*upper, *lower) };
44
45        let mut lo = I::ONE;
46        let mut hi = lo + I::ONE;
47        while pred(from + to * hi) == tf {
48            lo += lo;
49            hi += hi;
50            if (from + to * lo).is_deeper(*bound) {
51                let bound =
52                    if let Some(bound) = *bound { bound } else { continue };
53                let steps = bound.steps(from.into_inner(), to.into_inner());
54                let front = from + to * steps;
55                *(if tf { lower } else { upper }) = front;
56                return Some(self.iter_return(steps, tf));
57            }
58        }
59
60        while lo.lt1(hi) {
61            let mid = lo.avg(hi);
62            let cur = pred(from + to * mid) == tf;
63            *(if cur { &mut lo } else { &mut hi }) = mid;
64        }
65
66        *(if tf { lower } else { upper }) = from + to * lo;
67        Some(self.iter_return(lo, tf))
68    }
69}
70
71pub enum ApproxFrac<I> {
72    Lower((I, I)),
73    Upper((I, I)),
74}
75use ApproxFrac::{Lower, Upper};
76
77impl<I: std::fmt::Display> std::fmt::Debug for ApproxFrac<I> {
78    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        let (name, (num, den)) = match self {
80            Lower(x) => ("Lower", x),
81            Upper(x) => ("Upper", x),
82        };
83        fmt.debug_tuple(name).field(&format!("{}/{}", num, den)).finish()
84    }
85}
86
87impl<I> ApproxFrac<I> {
88    pub fn into_inner(self) -> (I, I) {
89        match self {
90            Lower(x) | Upper(x) => x,
91        }
92    }
93}
94
95impl<I: SbInt> From<((Fraction<I>, Fraction<I>), bool)> for ApproxFrac<I> {
96    fn from(((lower, upper), tf): ((Fraction<I>, Fraction<I>), bool)) -> Self {
97        if tf { Lower(lower.into()) } else { Upper(upper.into()) }
98    }
99}
100
101pub trait FracApprox<I, F> {
102    fn frac_approx_iter(self) -> FracApproxIter<I, F>;
103    fn frac_approx_iter_bound(self, bound: I) -> FracApproxIter<I, F>;
104}
105
106impl<I: SbInt, F: Fn(I, I) -> bool> FracApprox<I, F> for F {
107    fn frac_approx_iter(self) -> FracApproxIter<I, F> {
108        FracApproxIter::new(self, None)
109    }
110    fn frac_approx_iter_bound(self, bound: I) -> FracApproxIter<I, F> {
111        FracApproxIter::new(self, Some(bound))
112    }
113}
114
115pub trait SbInt:
116    Copy
117    + Eq
118    + PartialOrd<Self>
119    + AddAssign<Self>
120    + Add<Self, Output = Self>
121    + Mul<Self, Output = Self>
122    + std::fmt::Display
123    + std::fmt::Debug
124{
125    const ZERO: Self;
126    const ONE: Self;
127    const SIGNED: bool;
128    fn lt1(self, other: Self) -> bool;
129    fn avg(self, other: Self) -> Self;
130    fn abs(self) -> Self;
131    fn neg(self) -> Self;
132    fn steps(self, from: (Self, Self), to: (Self, Self)) -> Self;
133    fn lowest_frac() -> (Self, Self);
134    fn highest_frac() -> (Self, Self);
135    fn zero_frac() -> (Self, Self);
136}
137
138#[derive(Clone, Copy, Eq, PartialEq, std::fmt::Debug)]
139struct Fraction<I>(I, I);
140
141impl<I: SbInt> From<(I, I)> for Fraction<I> {
142    fn from((num, den): (I, I)) -> Self { Self(num, den) }
143}
144
145impl<I: SbInt> From<Fraction<I>> for (I, I) {
146    fn from(frac: Fraction<I>) -> Self { (frac.0, frac.1) }
147}
148
149impl<I: SbInt> Neg for Fraction<I> {
150    type Output = Self;
151    fn neg(self) -> Self { self.neg() }
152}
153
154impl<I: SbInt> Mul<I> for Fraction<I> {
155    type Output = Self;
156    fn mul(self, a: I) -> Self { Self(self.0 * a, self.1 * a) }
157}
158
159impl<I: SbInt> Add<Fraction<I>> for Fraction<I> {
160    type Output = Self;
161    fn add(self, other: Self) -> Self {
162        // mediant
163        Self(self.0 + other.0, self.1 + other.1)
164    }
165}
166
167impl<I: SbInt> Fraction<I> {
168    fn is_deeper(self, bound: Option<I>) -> bool {
169        bound.into_iter().any(|b| self.1 > b)
170    }
171    fn neg(self) -> Self { Self(self.0.neg(), self.1) }
172    fn into_inner(self) -> (I, I) { (self.0, self.1) }
173}
174
175macro_rules! impl_uint {
176    ( $($ty:ty)* ) => { $(
177        impl SbInt for $ty {
178            const ZERO: $ty = 0;
179            const ONE: $ty = 1;
180            const SIGNED: bool = false;
181            fn lt1(self, other: Self) -> bool { self + 1 < other }
182            fn avg(self, other: Self) -> Self {
183                self + (other - self) / 2
184            }
185            fn abs(self) -> Self { self }
186            fn neg(self) -> Self { self.wrapping_neg() } // not to be called
187            fn steps(self, from: (Self, Self), to: (Self, Self)) -> Self {
188                if to.0 == 0 {
189                    (self - from.1) / to.1
190                } else if to.1 == 0 {
191                    (self - from.0) / to.0
192                } else {
193                    ((self - from.0) / to.0).min((self - from.1) / to.1)
194                }
195            }
196            fn lowest_frac() -> (Self, Self) { (0, 1) }
197            fn highest_frac() -> (Self, Self) { (1, 0) }
198            fn zero_frac() -> (Self, Self) { (0, 1) }
199        }
200    )* }
201}
202
203impl_uint! { u8 u16 u32 u64 u128 usize }
204
205macro_rules! impl_int {
206    ( $($ty:ty)* ) => { $(
207        impl SbInt for $ty {
208            const ZERO: $ty = 0;
209            const ONE: $ty = 1;
210            const SIGNED: bool = true;
211            fn lt1(self, other: Self) -> bool { self + 1 < other }
212            fn avg(self, other: Self) -> Self {
213                self + (other - self) / 2
214            }
215            fn abs(self) -> Self { self.abs() }
216            fn neg(self) -> Self { -self}
217            fn steps(self, from: (Self, Self), to: (Self, Self)) -> Self {
218                if to.0 == 0 {
219                    (self - from.1) / to.1
220                } else if to.1 == 0 {
221                    (self - from.0) / to.0
222                } else {
223                    ((self - from.0) / to.0).min((self - from.1) / to.1)
224                }
225            }
226            fn lowest_frac() -> (Self, Self) { (-1, 0) }
227            fn highest_frac() -> (Self, Self) { (1, 0) }
228            fn zero_frac() -> (Self, Self) { (0, 1) }
229        }
230    )* }
231}
232
233impl_int! { i8 i16 i32 i64 i128 isize }
234
235#[test]
236fn sanity_check() {
237    let sqrt2 = |p: i128, q| p * p <= 2 * q * q;
238    for ap in sqrt2.frac_approx_iter().take(30) {
239        eprintln!("{ap:?}");
240    }
241
242    let frac = |p: i128, q| 113 * p <= 355 * q;
243    for ap in frac.frac_approx_iter_bound(113) {
244        eprintln!("{ap:?}");
245    }
246}