75 lines
2.1 KiB
Rust
75 lines
2.1 KiB
Rust
use rand::{seq::SliceRandom, Rng};
|
|
use talc_lang::{symbol::SYM_TYPE_ERROR, throw, value::{range::RangeType, Value}, vmcalliter, Vm, exception::Result};
|
|
use talc_macros::native_func;
|
|
|
|
use crate::unpack_args;
|
|
|
|
pub fn load(vm: &mut Vm) {
|
|
vm.set_global_name("rand", rand().into());
|
|
vm.set_global_name("rand_in", rand_in().into());
|
|
vm.set_global_name("shuf", shuf().into());
|
|
}
|
|
|
|
#[native_func(0)]
|
|
pub fn rand(_: &mut Vm, _: Vec<Value>) -> Result<Value> {
|
|
Ok(Value::Float(rand::thread_rng().gen()))
|
|
}
|
|
|
|
#[native_func(1)]
|
|
pub fn rand_in(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
|
|
let [_, col] = unpack_args!(args);
|
|
match col {
|
|
Value::List(l) => {
|
|
let l = l.borrow();
|
|
let Some(v) = l.choose(&mut rand::thread_rng()) else {
|
|
throw!(*SYM_TYPE_ERROR, "rand_in: empty list")
|
|
};
|
|
Ok(v.clone())
|
|
},
|
|
Value::Table(t) => {
|
|
let t = t.borrow();
|
|
if t.is_empty() {
|
|
throw!(*SYM_TYPE_ERROR, "rand_in: empty table")
|
|
};
|
|
let i = rand::thread_rng().gen_range(0..t.len());
|
|
let key = t.keys().nth(i).unwrap();
|
|
Ok(key.clone().into_inner())
|
|
},
|
|
Value::Range(r) => {
|
|
if r.is_empty() {
|
|
throw!(*SYM_TYPE_ERROR, "rand_in: empty range")
|
|
}
|
|
match r.ty {
|
|
RangeType::Open => Ok(Value::Int(
|
|
rand::thread_rng().gen_range(r.start..r.stop))),
|
|
RangeType::Closed => Ok(Value::Int(
|
|
rand::thread_rng().gen_range(r.start..=r.stop))),
|
|
RangeType::Endless => throw!(*SYM_TYPE_ERROR, "rand_in: endless range"),
|
|
}
|
|
}
|
|
col => {
|
|
let iter = col.to_iter_function()?;
|
|
let mut values = Vec::new();
|
|
while let Some(v) = vmcalliter!(vm; iter.clone())? {
|
|
values.push(v);
|
|
}
|
|
if values.is_empty() {
|
|
throw!(*SYM_TYPE_ERROR, "rand_in: empty iterator")
|
|
}
|
|
let i = rand::thread_rng().gen_range(0..values.len());
|
|
let v = values.swap_remove(i);
|
|
Ok(v)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[native_func(1)]
|
|
pub fn shuf(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
|
|
let [_, list] = unpack_args!(args);
|
|
let Value::List(list) = list else {
|
|
throw!(*SYM_TYPE_ERROR, "shuf expected list, got {list:#}")
|
|
};
|
|
list.borrow_mut().shuffle(&mut rand::thread_rng());
|
|
Ok(Value::List(list))
|
|
}
|