1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use arrow2::{array::Array, bitmap::Bitmap};
use sample_std::{Random, Sample, SampleAll, Shrunk};

pub mod array;
pub mod chunk;
pub mod datatypes;
pub mod fixed_size_list;
pub mod list;
pub mod primitive;
pub mod struct_;

pub type ArrowSampler = Box<dyn Sample<Output = Box<dyn Array>> + Send + Sync>;

pub(crate) fn generate_validity<V>(
    null: &mut Option<V>,
    g: &mut Random,
    len: usize,
) -> Option<Bitmap>
where
    V: Sample<Output = bool>,
{
    null.as_mut().map(|null| {
        Bitmap::from_trusted_len_iter(std::iter::repeat(()).take(len).map(|_| !null.generate(g)))
    })
}

pub trait SetLen {
    fn set_len(&mut self, len: usize);
}

#[derive(Debug, Clone)]
pub struct AlwaysValid;

impl Sample for AlwaysValid {
    type Output = Option<Bitmap>;

    fn generate(&mut self, _: &mut sample_std::Random) -> Self::Output {
        None
    }

    fn shrink(&self, _: Self::Output) -> Shrunk<Self::Output> {
        Box::new(std::iter::empty())
    }
}

impl SetLen for AlwaysValid {
    fn set_len(&mut self, _: usize) {}
}

impl<C, I, O> SetLen for sample_std::TryConvert<C, I, O>
where
    C: SetLen,
{
    fn set_len(&mut self, len: usize) {
        self.inner.set_len(len)
    }
}

impl<S, F, I> SampleLen for sample_std::TryConvert<S, F, I>
where
    S: Sample + SetLen,
    F: Fn(S::Output) -> Box<dyn Array>,
    I: Fn(Box<dyn Array>) -> Option<S::Output>,
{
}

impl<C> SetLen for sample_std::SamplerChoice<C>
where
    C: SetLen,
{
    fn set_len(&mut self, len: usize) {
        for choice in &mut self.choices {
            choice.set_len(len);
        }
    }
}

impl SampleLen for sample_std::SamplerChoice<ArrowLenSampler> {}

impl<S: SetLen> SetLen for SampleAll<S> {
    fn set_len(&mut self, len: usize) {
        for sampler in &mut self.samplers {
            sampler.set_len(len);
        }
    }
}

impl Sample for Box<dyn SampleLen> {
    type Output = Box<dyn Array>;

    fn generate(&mut self, g: &mut Random) -> Self::Output {
        self.as_mut().generate(g)
    }

    fn shrink(&self, v: Self::Output) -> Shrunk<'_, Self::Output> {
        self.as_ref().shrink(v)
    }
}

pub trait SampleLen: Sample<Output = Box<dyn Array>> + SetLen {}

pub type ArrowLenSampler = Box<dyn SampleLen>;

impl SetLen for ArrowLenSampler {
    fn set_len(&mut self, len: usize) {
        self.as_mut().set_len(len)
    }
}

impl SampleLen for ArrowLenSampler {}

pub struct FixedLenSampler<A> {
    pub len: usize,
    pub array: A,
}

impl<A> Sample for FixedLenSampler<A>
where
    A: Sample + SetLen,
    A::Output: Clone,
{
    type Output = A::Output;

    fn generate(&mut self, g: &mut Random) -> Self::Output {
        self.array.set_len(self.len);
        self.array.generate(g)
    }

    fn shrink(&self, array: Self::Output) -> Shrunk<Self::Output> {
        self.array.shrink(array)
    }
}