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

//! A random number module for the Koto language

use {
    koto_runtime::{
        external_error, get_external_instance, make_external_value, num2, num4, ExternalValue,
        Value, ValueMap,
    },
    rand::{Rng, SeedableRng},
    rand_chacha::ChaCha20Rng,
    std::fmt,
};

pub fn make_module() -> ValueMap {
    use Value::*;

    // The random module contains a default generator
    let mut result = ChaChaRng::make_value_map(ChaCha20Rng::from_entropy());

    // random.generator is available to create custom generators
    result.add_fn("generator", |vm, args| match vm.get_args(args) {
        [] => Ok(Map(ChaChaRng::make_value_map(ChaCha20Rng::from_entropy()))),
        [Number(n)] => Ok(Map(ChaChaRng::make_value_map(ChaCha20Rng::seed_from_u64(
            n.to_bits(),
        )))),
        _ => external_error!("random.generator - expected no arguments, or seed number"),
    });

    result
}

#[derive(Debug)]
struct ChaChaRng(ChaCha20Rng);

impl ChaChaRng {
    fn make_value_map(rng: ChaCha20Rng) -> ValueMap {
        use Value::*;

        let mut result = ValueMap::new();

        result.add_instance_fn("bool", |vm, args| {
            let args = vm.get_args(args);
            get_external_instance!(args, "random", "bool", Self, rng, {
                Ok(Bool(rng.0.gen::<bool>()))
            })
        });

        result.add_instance_fn("number", |vm, args| {
            let args = vm.get_args(args);
            get_external_instance!(args, "random", "number", Self, rng, {
                Ok(Number(rng.0.gen::<f64>().into()))
            })
        });

        result.add_instance_fn("number2", |vm, args| {
            let args = vm.get_args(args);
            get_external_instance!(args, "random", "number2", Self, rng, {
                let result = num2::Num2(rng.0.gen::<f64>(), rng.0.gen::<f64>());
                Ok(Num2(result))
            })
        });

        result.add_instance_fn("number4", |vm, args| {
            let args = vm.get_args(args);
            get_external_instance!(args, "random", "number4", Self, rng, {
                let result = num4::Num4(
                    rng.0.gen::<f32>(),
                    rng.0.gen::<f32>(),
                    rng.0.gen::<f32>(),
                    rng.0.gen::<f32>(),
                );
                Ok(Num4(result))
            })
        });

        result.add_instance_fn("pick", |vm, args| {
            let args = vm.get_args(args);
            get_external_instance!(args, "random", "number", Self, rng, {
                match &args[1..] {
                    [List(l)] => {
                        let index = rng.0.gen_range(0, l.len());
                        Ok(l.data()[index].clone())
                    }
                    [Range(r)] => {
                        let (start, end) = if r.end > r.start {
                            (r.start, r.end)
                        } else {
                            (r.end, r.start)
                        };
                        let size = end - start;
                        let index = rng.0.gen_range(0, size);
                        Ok(Number((start + index).into()))
                    }
                    _ => external_error!("random.pick - expected list or range as argument"),
                }
            })
        });

        result.add_instance_fn("seed", |vm, args| {
            let args = vm.get_args(args);
            get_external_instance!(args, "random", "seed", Self, rng, {
                match &args[1..] {
                    [Number(n)] => {
                        *rng = ChaChaRng(ChaCha20Rng::seed_from_u64(n.to_bits()));
                        Ok(Empty)
                    }
                    _ => external_error!("random.seed - expected number as argument"),
                }
            })
        });

        result.insert(Value::ExternalDataId, make_external_value(Self(rng)));
        result
    }
}

impl ExternalValue for ChaChaRng {
    fn value_type(&self) -> String {
        "Rng".to_string()
    }
}

impl fmt::Display for ChaChaRng {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Rng")
    }
}