From 4438116264eeece05c9fdbdf73c8f33757a108e4 Mon Sep 17 00:00:00 2001
From: Freya Murphy <freya@freyacat.org>
Date: Mon, 26 Feb 2024 20:39:39 -0500
Subject: sort n rand

---
 matrix-stdlib/src/core.rs | 119 +++++++++++++++++++++++++++++++++++++++++++++-
 matrix-stdlib/src/iter.rs |   8 +---
 matrix-stdlib/src/lib.rs  |   7 +++
 3 files changed, 126 insertions(+), 8 deletions(-)

(limited to 'matrix-stdlib/src')

diff --git a/matrix-stdlib/src/core.rs b/matrix-stdlib/src/core.rs
index 183c142..2c76497 100644
--- a/matrix-stdlib/src/core.rs
+++ b/matrix-stdlib/src/core.rs
@@ -2,7 +2,8 @@ use std::hash::{DefaultHasher, Hash, Hasher};
 
 use matrix::{vm::Vm, value::Value, unpack_args, Result, unpack_varargs};
 use matrix_macros::native_func;
-use crate::{VmArgs, error};
+use rand::Rng;
+use crate::{VmArgs, next, error};
 
 
 fn to_radix(r: i64, mut n: i64) -> String {
@@ -162,6 +163,116 @@ fn str(_: VmArgs, args: Vec<Value>) -> Result<Value> {
     Ok(Value::String(format!("{value}").into()))
 }
 
+fn partition(list: &mut [Value], lo: usize, hi: usize) -> (usize, usize) {
+    let pivot = list[(lo + hi) / 2].clone();
+    let mut lt = lo;
+    let mut eq = lo;
+    let mut gt = hi;
+
+    while eq <= gt {
+        if list[eq] < pivot {
+            list.swap(eq, lt);
+            lt += 1;
+            eq += 1;
+        } else if list[eq] > pivot {
+            list.swap(eq, gt);
+            gt -= 1;
+        } else {
+            eq += 1;
+        }
+    }
+
+    (lt, gt)
+}
+
+fn quicksort(list: &mut [Value], lo: usize, hi: usize) {
+    if lo < hi {
+        let (lt, gt) = partition(list, lo, hi);
+        if lt > 0 {
+            quicksort(list, lo, lt - 1);
+        }
+        quicksort(list, gt + 1, hi);
+    }
+}
+
+#[native_func(1)]
+fn sort(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+    let [value] = unpack_args!(args);
+    let Value::List(mut list) = value else {
+        return error!("sort requires a list")
+    };
+    let hi = list.len() - 1;
+    quicksort(list.as_mut(), 0, hi);
+    Ok(Value::List(list))
+}
+
+#[native_func(0)]
+fn rand(_: VmArgs, _: Vec<Value>) -> Result<Value> {
+    let mut rng = rand::thread_rng();
+    Ok(Value::Float(rng.gen()))
+}
+
+#[native_func(0..)]
+fn rand_range(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+    use Value as V;
+    let ([], varargs) = unpack_varargs!(args);
+    let (min, max, step) = match varargs.as_slice() {
+        [V::Int(max)] => {
+            (0, *max, 1)
+        },
+        [V::Int(min), V::Int(max)] => {
+            (*min, *max, 1)
+        },
+        [V::Int(min), V::Int(max), V::Int(step)] => {
+            (*min, *max, *step)
+        },
+        _ => return error!("rand_range takes 1 to 3 [Int]'s")
+    };
+    if max <= min {
+        return Ok(Value::Int(min))
+    }
+    let count = ((max - 1 - min) / step) + 1;
+    let rng = (rand::thread_rng().gen::<u32>() % (count as u32)) as i64;
+    let i = min + (rng * step);
+    Ok(Value::Int(i))
+}
+
+#[native_func(2)]
+fn rand_int(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+    use Value as V;
+    let (min, max) = match args.as_slice() {
+        [V::Int(min), V::Int(max)] => {
+            (*min, *max + 1)
+        },
+        _ => return error!("rand_int takes 2 [Int]'s")
+    };
+    if max <= min {
+        return Ok(Value::Int(min))
+    }
+
+    let count = max - min;
+    let rng = (rand::thread_rng().gen::<u32>() % (count as u32)) as i64;
+    let i = min + rng;
+    Ok(Value::Int(i))
+}
+
+#[native_func(1)]
+fn rand_in((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+    let [iter] = unpack_args!(args);
+    let iter = iter.into_iter_fn()?;
+    let mut vals = Vec::new();
+    loop {
+        let val = next!(vm, frame, iter)?;
+        if val == Value::Nil { break }
+        vals.push(val);
+    }
+    if vals.is_empty() {
+        return Ok(Value::Nil)
+    }
+    let idx = rand::thread_rng().gen::<usize>() % vals.len();
+    Ok(vals[idx].clone())
+}
+
 pub fn load(vm: &mut Vm) {
     vm.load_global_fn(bin(), "bin");
     vm.load_global_fn(sex(), "sex");
@@ -178,4 +289,10 @@ pub fn load(vm: &mut Vm) {
     vm.load_global_fn(hash(), "hash");
     vm.load_global_fn(ord(), "ord");
     vm.load_global_fn(chr(), "chr");
+
+    vm.load_global_fn(sort(), "sort");
+    vm.load_global_fn(rand(), "rand");
+    vm.load_global_fn(rand_range(), "rand_range");
+    vm.load_global_fn(rand_int(), "rand_int");
+    vm.load_global_fn(rand_in(), "rand_in");
 }
diff --git a/matrix-stdlib/src/iter.rs b/matrix-stdlib/src/iter.rs
index 74056ce..630e52c 100644
--- a/matrix-stdlib/src/iter.rs
+++ b/matrix-stdlib/src/iter.rs
@@ -1,16 +1,10 @@
 use std::{cell::RefCell, rc::Rc};
 use matrix::{iter, vm::Vm, value::Value, Result, unpack_varargs, unpack_args};
 use matrix_macros::native_func;
-use crate::{error, VmArgs};
+use crate::{error, next, VmArgs};
 
 use Value as V;
 
-macro_rules! next {
-    ($vm:expr, $frame:expr, $iter:expr) => {
-        $vm.run_fn($frame, $iter.clone(), vec![])
-    };
-}
-
 #[native_func(1)]
 fn len(_: VmArgs, args: Vec<Value>) -> Result<Value> {
     let [value] = unpack_args!(args);
diff --git a/matrix-stdlib/src/lib.rs b/matrix-stdlib/src/lib.rs
index 334de90..b4ab658 100644
--- a/matrix-stdlib/src/lib.rs
+++ b/matrix-stdlib/src/lib.rs
@@ -14,7 +14,14 @@ macro_rules! error {
     };
 }
 
+macro_rules! next {
+    ($vm:expr, $frame:expr, $iter:expr) => {
+        $vm.run_fn($frame, $iter.clone(), vec![])
+    };
+}
+
 pub(crate) use error;
+pub(crate) use next;
 
 pub fn load(vm: &mut Vm) {
     core::load(vm);
-- 
cgit v1.2.3-freya