nekolib/utils/
rand_gen_macro.rs

1//! 乱数生成マクロ。
2
3use super::bitop;
4
5#[macro_export]
6macro_rules! rand_gen_builder {
7    ( @gen { [ $($cur:tt)* ] } ) => {
8        $crate::rand_gen_builder!(@vec @gen {} @rest $($cur)*)
9    };
10    ( @gen { [ $($cur:tt)* ] where { $($f:ident $(= $a:expr)?),* } } ) => {
11        $crate::rand_gen_builder!(@vec @gen {} @rest $($cur)*)
12            .options() $(.$f($($a)?))*
13    };
14    ( @gen { [ $($cur:tt)* ] where { $($f:ident $(= $a:expr)?),*, } } ) => {
15        $crate::rand_gen_builder!(@vec @gen {} @rest $($cur)*)
16            .options() $(.$f($($a)?))*
17    };
18    ( @vec @gen { $($x:tt)* } @rest ; $($rest:tt)* ) => {
19        $crate::rand_gen_builder!(@vec @gen { $($x)* } @len { $($rest)* })
20    };
21    ( @vec @gen { $($cur:tt)* } @rest $tt:tt $($rest:tt)* ) => {
22        $crate::rand_gen_builder!(@vec @gen { $($cur)* $tt } @rest $($rest)*)
23    };
24    ( @vec @gen { $($x:tt)* } @len { $($len:tt)* }) => {
25        VecMarker::new($crate::rand_gen_builder!(@gen { $($x)* }), $($len)*)
26    };
27
28    ( @gen { $($x:tt)* } ) => {
29        $crate::rand_gen_builder!(@gen {} @rest $($x)*)
30    };
31    ( @gen { $($cur:tt)* } @rest where { $($where:tt)* } ) => {
32        $crate::rand_gen_builder!(@gen { $($cur)* } @where { $($where)* } @rest)
33    };
34    ( @gen { $($cur:tt)* } @rest $tt:tt $($rest:tt)* ) => {
35        $crate::rand_gen_builder!(@gen { $($cur)* $tt } @rest $($rest)*)
36    };
37    ( @gen { $($cur:tt)* } @where { $($f:ident $(= $a:expr)?),* } @rest ) => {
38        $crate::rand_gen_builder!(@gen { $($cur)* } @rest)
39            .options() $(.$f($($a)?))*
40    };
41    ( @gen { $($cur:tt)* } @where { $($f:ident $(= $a:expr)?),*, } @rest ) => {
42        $crate::rand_gen_builder!(@gen { $($cur)* } @rest)
43            .options() $(.$f($($a)?))*
44    };
45    ( @gen { $($x:tt)* } @rest ) => {
46        ($($x)*)
47    };
48}
49
50/// 乱数生成マクロ。
51///
52/// # Notes
53/// 作りかけなのでいろいろ足りない。
54///
55/// # Examples
56/// ```
57/// use rand::SeedableRng;
58/// use rand_chacha::ChaCha20Rng;
59///
60/// use nekolib::rand_gen;
61/// use nekolib::utils::ascii::*;
62/// use nekolib::utils::rand_gen_macro::*;
63///
64/// rand_gen! {
65///     rng: ChaCha20Rng;
66///
67///     n in 1_usize..=10;
68///     a in [1_i64..=100; n];
69///     s in AsciiString(16) where {
70///         distribution = &[
71///             (ASCII_LOWERCASE, 10),
72///             (ASCII_UPPERCASE, 10),
73///             (ASCII_DIGIT, 6),
74///             (charset(b"~!@#$%+?|()^*_-=[]{};:,./"), 5),
75///         ],
76///     };
77/// }
78///
79/// assert_eq!(a.len(), n);
80/// assert!(a.iter().all(|ai| (1..=100).contains(ai)));
81/// assert_eq!(s.len(), 16);
82/// // Possible value of `s`: `"3e)xIjos2^M/XI1T"`, `"X52dhjDk%i6)p1F9"`
83/// ```
84///
85/// 以下のような出力が得られるため、これを用いて再現することができる。
86///
87/// ````text
88/// To reproduce:
89///
90/// ```
91/// rand_gen! {
92///     rng = ChaCha20Rng::from_seed([250, 120, 31, 164, 15, 176, 41, 144, 61, 59, 224, 119, 135, 238, 14, 193, 149, 124, 228, 39, 107, 208, 243, 180, 7, 177, 21, 88, 19, 5, 225, 3]);
93///     // ...
94/// }
95/// ```
96/// ````
97///
98/// ```
99/// use rand::SeedableRng;
100/// use rand_chacha::ChaCha20Rng;
101///
102/// use nekolib::rand_gen;
103/// use nekolib::utils::ascii::*;
104/// use nekolib::utils::rand_gen_macro::*;
105///
106/// rand_gen! {
107///     rng = ChaCha20Rng::from_seed([250, 120, 31, 164, 15, 176, 41, 144, 61, 59, 224, 119, 135, 238, 14, 193, 149, 124, 228, 39, 107, 208, 243, 180, 7, 177, 21, 88, 19, 5, 225, 3]);
108///
109///     n in 1_usize..=10;
110///     a in [1_i64..=100; n];
111///     s in AsciiString(16) where {
112///         distribution = &[
113///             (ASCII_LOWERCASE, 10),
114///             (ASCII_UPPERCASE, 10),
115///             (ASCII_DIGIT, 6),
116///             (charset(b"~!@#$%+?|()^*_-=[]{};:,./"), 5),
117///         ],
118///     };
119///     b in [[[1_u8..100; 3]; 2]; 2];
120///     p in Permutation(5);
121///     mut x in 1_i64..10;
122/// }
123/// x += 1;
124///
125/// assert_eq!(a, [32, 86, 41, 68, 66, 46, 56, 82, 40, 1]);
126/// assert_eq!(s, "X52dhjDk%i6)p1F9");
127/// assert_eq!(b, [[[75, 20, 23], [63, 21, 58]], [[12, 6, 57], [51, 95, 70]]]);
128/// assert_eq!(p, [3, 1, 0, 4, 2]);
129/// assert_eq!(x, 5);
130/// ```
131///
132/// `oj` を用いて、hack をするのに使うとよい。
133/// 以下の例は、想定解が Rust で標的が Python のプログラムのもの。
134///
135/// ```sh
136/// % oj g/i --hack-expected target/release/x --hack-actual 'python3 x-to-be-hacked.py' target/release/x-gen
137/// ```
138#[macro_export]
139macro_rules! rand_gen {
140    ( @seed $seed:ident { mut $a:ident in $($r:tt)* } @rest ) => {
141        let mut $a = $seed.generate($crate::rand_gen_builder!(@gen { $($r)* }));
142    };
143    ( @seed $seed:ident { $a:ident in $($r:tt)* } @rest ) => {
144        let $a = $seed.generate($crate::rand_gen_builder!(@gen { $($r)* }));
145    };
146    ( @seed $seed:ident { $a:ident in $($r:tt)* } @rest ; $($rest:tt)* ) => {
147        rand_gen!(@seed $seed { $a in $($r)* } @rest);
148        rand_gen!(@seed $seed {} @rest $($rest)*);
149    };
150    ( @seed $seed:ident { mut $a:ident in $($r:tt)* } @rest ; $($rest:tt)* ) => {
151        rand_gen!(@seed $seed { mut $a in $($r)* } @rest);
152        rand_gen!(@seed $seed {} @rest $($rest)*);
153    };
154    ( @seed $seed:ident { $($cur:tt)* } @rest $tt:tt $($rest:tt)* ) => {
155        rand_gen!(@seed $seed { $($cur)* $tt } @rest $($rest)* );
156    };
157    ( @seed $seed:ident {} @rest ) => {};
158    ( $seed:ident: $ty:ty; $($rest:tt)* ) => {
159        // let mut $seed = <$ty as RandomWordGenerator>::auto_init();
160        let mut $seed = <$ty>::from_entropy();
161        eprintln!(r#"
162To reproduce:
163
164```
165rand_gen! {{
166    {} = {};
167    // ...
168}}
169```
170"#, stringify!($seed), $seed.inspect());
171        rand_gen!(@seed $seed {} @rest $($rest)*);
172    };
173    ( $seed:ident = $s:expr; $($rest:tt)* ) => {
174        let mut $seed = $s;
175        // $seed.inspect(stringify!($seed));
176        rand_gen!(@seed $seed {} @rest $($rest)*);
177    };
178}
179
180// .options() いらないかも
181pub trait GenOptions {
182    type OptionType;
183    fn options(self) -> Self::OptionType;
184}
185
186// ---
187// ```
188// pub trait RandomGenerator<Input> {
189//     type Output;
190//     fn generate(&mut self) -> Self::Output;
191// }
192// ```
193// みたいなのを作って、impl RandChaCha とかする？
194// Input = RangeInclusive<i64>, Output = i64 とかを想定している。
195//
196// 今だと impl RandGen for Input して rand_gen の引数に &mut RandChaCha とかの
197// を渡すようにしているのか、むー。でもこれが適当な impl trait なのでつらいんだね
198//
199// なんかできてそう。
200// あとは、過去のを消しつつ、options をよしなに解決すればよさそう。
201
202#[derive(Clone, Copy)]
203pub struct AsciiString(pub usize);
204
205impl GenOptions for AsciiString {
206    type OptionType = Self;
207    fn options(self) -> Self { self }
208}
209
210#[derive(Clone, Copy)]
211pub struct AsciiStringOfCharset(AsciiGen, usize);
212
213#[derive(Clone)]
214pub struct AsciiStringOfDistribution(BTreeMap<u32, AsciiGen>, u32, usize);
215
216impl AsciiString {
217    pub fn charset(self, cs: u128) -> AsciiStringOfCharset {
218        AsciiStringOfCharset(AsciiGen::new(cs), self.0)
219    }
220    pub fn distribution(self, d: &[(u128, u32)]) -> AsciiStringOfDistribution {
221        let mut map = BTreeMap::new();
222        let mut acc = 0;
223        for &(mask, p) in d.iter().filter(|&&(_, p)| p > 0) {
224            acc += p;
225            map.insert(acc, AsciiGen::new(mask));
226        }
227        AsciiStringOfDistribution(map, acc, self.0)
228    }
229}
230
231impl RandomGenerator<AsciiStringOfCharset> for ChaCha20Rng {
232    type Output = String;
233    fn generate(&mut self, subject: AsciiStringOfCharset) -> String {
234        let AsciiStringOfCharset(gen, len) = subject;
235        if len == 0 {
236            "".to_owned()
237        } else {
238            (0..len).map(|_| self.generate(gen)).collect()
239        }
240    }
241}
242
243impl RandomGenerator<AsciiStringOfDistribution> for ChaCha20Rng {
244    type Output = String;
245    fn generate(&mut self, subject: AsciiStringOfDistribution) -> String {
246        let AsciiStringOfDistribution(ref gen_map, acc, len) = subject;
247        let large = Uniform::from(0..acc);
248        (0..len)
249            .map(|_| {
250                let c = large.sample(self);
251                let small = gen_map.range(c..).next().unwrap().1;
252                self.generate(*small)
253            })
254            .collect()
255    }
256}
257
258#[derive(Clone, Copy)]
259pub struct Ascii;
260
261impl GenOptions for Ascii {
262    type OptionType = Self;
263    fn options(self) -> Ascii { self }
264}
265
266impl Ascii {
267    pub fn charset(self, cs: u128) -> AsciiGen { AsciiGen::new(cs) }
268}
269
270#[derive(Clone, Copy)]
271pub struct AsciiGen(PdepPextMaskU128, u32);
272
273impl AsciiGen {
274    pub fn new(mask: u128) -> Self {
275        Self(PdepPextMaskU128::new(mask), mask.count_ones())
276    }
277}
278
279impl RandomGenerator<AsciiGen> for ChaCha20Rng {
280    type Output = char;
281    fn generate(&mut self, subject: AsciiGen) -> char {
282        let AsciiGen(mask, pop) = subject;
283        assert_ne!(pop, 0, "empty charset");
284        let c = Uniform::from(0..pop).sample(self);
285        (1 << c).pdep(mask).trailing_zeros() as u8 as char
286    }
287}
288
289use std::collections::{BTreeMap, BTreeSet};
290use std::ops::{Range, RangeInclusive};
291
292use rand::distributions::{Distribution, Uniform};
293use rand_chacha::ChaCha20Rng;
294
295use bitop::{Pdep, PdepPextMaskU128};
296
297pub trait SeedableRngInspect {
298    fn inspect(&self) -> String;
299}
300
301impl SeedableRngInspect for ChaCha20Rng {
302    fn inspect(&self) -> String {
303        format!("ChaCha20Rng::from_seed({:?})", self.get_seed())
304    }
305}
306
307pub trait RandomGenerator<Input> {
308    type Output;
309    fn generate(&mut self, subject: Input) -> Self::Output;
310}
311
312macro_rules! impl_range {
313    ( $($t:ty)* ) => { $(
314        impl RandomGenerator<RangeInclusive<$t>> for ChaCha20Rng {
315            type Output = $t;
316            fn generate(&mut self, s: RangeInclusive<$t>) -> $t {
317                let between = Uniform::from(s);
318                between.sample(self)
319            }
320        }
321        impl RandomGenerator<Range<$t>> for ChaCha20Rng {
322            type Output = $t;
323            fn generate(&mut self, s: Range<$t>) -> $t {
324                let between = Uniform::from(s);
325                between.sample(self)
326            }
327        }
328    )* }
329}
330
331impl_range! { u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize }
332
333#[derive(Clone)]
334pub struct VecMarker<T> {
335    inner: T,
336    len: usize,
337}
338
339impl<T> VecMarker<T> {
340    pub fn new(inner: T, len: usize) -> Self { Self { inner, len } }
341}
342
343// sorted は Ord、distinct は Eq が必要になるけど、そもそも
344// 範囲で取ってランダムに生成とか言っている時点でそれらは仮定していそう。
345pub struct VecOptionsMarker<T> {
346    inner: T,
347    len: usize,
348    sorted: bool,
349    distinct: bool,
350}
351
352impl<T> GenOptions for VecMarker<T> {
353    type OptionType = VecOptionsMarker<T>;
354    fn options(self) -> VecOptionsMarker<T> {
355        let Self { inner, len } = self;
356        VecOptionsMarker { inner, len, sorted: false, distinct: false }
357    }
358}
359
360impl<T> VecOptionsMarker<T> {
361    pub fn sorted(mut self) -> Self {
362        self.sorted = true;
363        self
364    }
365    pub fn distinct(mut self) -> Self {
366        self.distinct = true;
367        self
368    }
369}
370
371// impl<T> U<Vec<T>> for S where S implements U<T> { ... }
372// みたいなのってもしかしてできない？
373// マクロで定数段だけやればとりあえずいいか。
374
375macro_rules! impl_vec_range {
376    ( $($t:ty)* ) => { $(
377        impl RandomGenerator<VecMarker<RangeInclusive<$t>>> for ChaCha20Rng {
378            type Output = Vec<$t>;
379            fn generate(
380                &mut self,
381                s: VecMarker<RangeInclusive<$t>>
382            ) -> Self::Output {
383                let VecMarker { inner, len } = s;
384                let between = Uniform::from(inner);
385                (0..len).map(|_| between.sample(self)).collect()
386            }
387        }
388        impl RandomGenerator<VecMarker<Range<$t>>> for ChaCha20Rng {
389            type Output = Vec<$t>;
390            fn generate(&mut self, s: VecMarker<Range<$t>>) -> Self::Output {
391                let VecMarker { inner, len } = s;
392                let between = Uniform::from(inner);
393                (0..len).map(|_| between.sample(self)).collect()
394            }
395        }
396    )* }
397}
398
399impl_vec_range! { u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize }
400
401macro_rules! impl_vecvec_range {
402    ( $( ($marker:ty, $vec:ty), )* ) => { $(
403        impl RandomGenerator<$marker> for ChaCha20Rng {
404            type Output = $vec;
405            fn generate(&mut self, s: $marker) -> Self::Output {
406                let VecMarker { inner, len } = s;
407                (0..len).map(|_| self.generate(inner.clone())).collect()
408            }
409        }
410    )* };
411    ( $( $t:ty )* ) => { $(
412        impl_vecvec_range! {
413            // 2
414            (VecMarker<VecMarker<RangeInclusive<$t>>>, Vec<Vec<$t>>),
415            (VecMarker<VecMarker<Range<$t>>>, Vec<Vec<$t>>),
416
417            // 3
418            (
419                VecMarker<VecMarker<VecMarker<RangeInclusive<$t>>>>,
420                Vec<Vec<Vec<$t>>>
421            ),
422            (VecMarker<VecMarker<VecMarker<Range<$t>>>>, Vec<Vec<Vec<$t>>>),
423
424            // 4
425            (
426                VecMarker<VecMarker<VecMarker<VecMarker<RangeInclusive<$t>>>>>,
427                Vec<Vec<Vec<Vec<$t>>>>
428            ),
429            (
430                VecMarker<VecMarker<VecMarker<VecMarker<Range<$t>>>>>,
431                Vec<Vec<Vec<Vec<$t>>>>
432            ),
433        }
434    )* }
435}
436
437impl_vecvec_range! { u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize }
438
439impl RandomGenerator<VecMarker<Permutation>> for ChaCha20Rng {
440    type Output = Vec<Vec<usize>>;
441    fn generate(&mut self, s: VecMarker<Permutation>) -> Self::Output {
442        let VecMarker { inner, len } = s;
443        (0..len).map(|_| self.generate(inner)).collect()
444    }
445}
446
447impl RandomGenerator<VecOptionsMarker<RangeInclusive<i64>>> for ChaCha20Rng {
448    type Output = Vec<i64>;
449    fn generate(
450        &mut self,
451        subject: VecOptionsMarker<RangeInclusive<i64>>,
452    ) -> Vec<i64> {
453        let VecOptionsMarker { inner, len, sorted, distinct } = subject;
454
455        if len == 0 {
456            return vec![];
457        }
458
459        let start = *inner.start();
460        let end = *inner.end();
461        if start > end {
462            panic!("Emptyset");
463        }
464        if distinct {
465            let u_len = end - start;
466            // 全域だったらオーバーフローするので +1 しないでおく。
467            // len == 0 は上で処理しているので、len-1 はオーバーフローしない。
468            // 符号つきだから u_len の時点でオーバーフローしうるじゃんね。
469            // 全域でなければ Range の方を使うようにする？ 全域なら別で考える？
470            // 嫌だけど。
471
472            // if end - start + 1 < len {}
473            if u_len == 0 && len == 1 {
474                return vec![start];
475            }
476            if (u_len as u128) < len as u128 - 1 {
477                panic!("by pigeonhole principle, it is infeasible");
478            }
479
480            // (end - start + 1, len)
481            // (_, 0) => []
482            // (1, 1) => [start]
483            // (x, y) if x < y => panic!()
484
485            let mut res = vec![];
486            // k >~ n/log(n) くらいなら dense だと思っていい？
487            // ゼロ割りは嫌なので k log(n) >~ n。
488            let lg_len = len.next_power_of_two().trailing_zeros() as i64;
489            let sparse = u_len * lg_len < len as i64;
490
491            if sparse {
492                let between = Uniform::from(inner);
493                let mut seen = BTreeSet::new();
494                while res.len() < len {
495                    // 失敗しまくったら下のやり方に fallback した方がいい？
496                    let cur = between.sample(self);
497                    if seen.insert(cur) {
498                        res.push(cur);
499                    }
500                }
501            } else {
502                let mut pool: Vec<_> = (start..=end).collect();
503                for _ in 0..len {
504                    let u = Uniform::from(0..pool.len());
505                    let i = u.sample(self);
506                    res.push(pool.swap_remove(i));
507                }
508            }
509
510            if sorted {
511                res.sort_unstable();
512            }
513            return res;
514        }
515
516        if !sorted {
517            return self.generate(VecMarker { inner, len });
518        }
519
520        // todo うまくやる
521        // distinct のオプションがあるとやりやすいかも
522        // 関数が同じだから再帰でやるかどうしようか
523        // [start..=end + len - 1; len] where { distinct }
524        // を作り、ソートして、a[i] -= i して返す。
525        //
526        // distinct に作るパートが難しくて、sparse か dense かで分ける？
527        // infeasible なら panic!() で。
528
529        // let mut tmp = self.generate(VecMarker { inner, len });
530        // tmp.sort_unstable();
531        // tmp
532
533        let start = *inner.start();
534        let end = inner.end() + len as i64 - 1;
535        let mut res = self.generate(VecOptionsMarker {
536            inner: start..=end,
537            len,
538            sorted: true,
539            distinct: true,
540        });
541        for i in 0..len {
542            res[i] -= i as i64;
543        }
544        res
545    }
546}
547
548#[derive(Clone, Copy)]
549pub struct Permutation(pub usize);
550
551impl RandomGenerator<Permutation> for ChaCha20Rng {
552    type Output = Vec<usize>;
553    fn generate(&mut self, subject: Permutation) -> Vec<usize> {
554        let n = subject.0;
555        let mut res: Vec<_> = (0..n).collect();
556        for i in (1..n).rev() {
557            let j = self.generate(0..=i);
558            res.swap(j, i);
559        }
560        res
561    }
562}
nekolib/utils/rand_gen_macro.rs

nekolib/utils/
rand_gen_macro.rs
