Add list.sort, sorted()

adrian17 · adrian17 · commit 45da6b8af19c · 2019-02-26T22:03:45.000+01:00
diff --git a/tests/snippets/list.py b/tests/snippets/list.py
@@ -117,3 +117,42 @@ def __eq__(self, x):
 assert a == b
 
 assert [foo] == [foo]
+
+for size in [1, 2, 3, 4, 5, 8, 10, 100, 1000]:
+   lst = list(range(size))
+   orig = lst[:]
+   lst.sort()
+   assert lst == orig
+   assert sorted(lst) == orig
+   assert_raises(ZeroDivisionError, lambda: sorted(lst, key=lambda x: 1/x))
+   lst.reverse()
+   assert sorted(lst) == orig
+   assert sorted(lst, reverse=True) == lst
+   assert sorted(lst, key=lambda x: -x) == lst
+   assert sorted(lst, key=lambda x: -x, reverse=True) == orig
+
+assert sorted([(1, 2, 3), (0, 3, 6)]) == [(0, 3, 6), (1, 2, 3)]
+assert sorted([(1, 2, 3), (0, 3, 6)], key=lambda x: x[0]) == [(0, 3, 6), (1, 2, 3)]
+assert sorted([(1, 2, 3), (0, 3, 6)], key=lambda x: x[1]) == [(1, 2, 3), (0, 3, 6)]
+assert sorted([(1, 2), (), (5,)], key=len) == [(), (5,), (1, 2)]
+
+lst = [3, 1, 5, 2, 4]
+class C:
+  def __init__(self, x): self.x = x
+  def __lt__(self, other): return self.x < other.x
+lst.sort(key=C)
+assert lst == [1, 2, 3, 4, 5]
+
+lst = [3, 1, 5, 2, 4]
+class C:
+  def __init__(self, x): self.x = x
+  def __gt__(self, other): return self.x > other.x
+lst.sort(key=C)
+assert lst == [1, 2, 3, 4, 5]
+
+lst = [5, 1, 2, 3, 4]
+def f(x):
+    lst.append(1)
+    return x
+assert_raises(ValueError, lambda: lst.sort(key=f)) # "list modified during sort"
+assert lst == [1, 2, 3, 4, 5]
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -689,7 +689,16 @@ fn builtin_setattr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 // builtin_slice
-// builtin_sorted
+
+fn builtin_sorted(vm: &mut VirtualMachine, mut args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(iterable, None)]);
+    let items = vm.extract_elements(iterable)?;
+    let lst = vm.ctx.new_list(items);
+
+    args.shift();
+    vm.call_method_pyargs(&lst, "sort", args)?;
+    Ok(lst)
+}
 
 fn builtin_sum(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(iterable, None)]);
@@ -763,6 +772,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
         "round" => ctx.new_rustfunc(builtin_round),
         "set" => ctx.set_type(),
         "setattr" => ctx.new_rustfunc(builtin_setattr),
+        "sorted" => ctx.new_rustfunc(builtin_sorted),
         "slice" => ctx.slice_type(),
         "staticmethod" => ctx.staticmethod_type(),
         "str" => ctx.str_type(),
diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs
@@ -3,8 +3,8 @@ use std::cell::{Cell, RefCell};
 use super::objbool;
 use super::objint;
 use super::objsequence::{
-    get_elements, get_item, get_mut_elements, seq_equal, seq_ge, seq_gt, seq_le, seq_lt, seq_mul,
-    PySliceableSequence,
+    get_elements, get_elements_cell, get_item, get_mut_elements, seq_equal, seq_ge, seq_gt, seq_le,
+    seq_lt, seq_mul, PySliceableSequence,
 };
 use super::objstr;
 use super::objtype;
@@ -334,12 +334,99 @@ fn list_reverse(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
 }
 
+fn quicksort(
+    vm: &mut VirtualMachine,
+    keys: &mut [PyObjectRef],
+    values: &mut [PyObjectRef],
+) -> PyResult<()> {
+    let len = values.len();
+    if len >= 2 {
+        let pivot = partition(vm, keys, values)?;
+        quicksort(vm, &mut keys[0..pivot], &mut values[0..pivot])?;
+        quicksort(vm, &mut keys[pivot + 1..len], &mut values[pivot + 1..len])?;
+    }
+    Ok(())
+}
+
+fn partition(
+    vm: &mut VirtualMachine,
+    keys: &mut [PyObjectRef],
+    values: &mut [PyObjectRef],
+) -> PyResult<usize> {
+    let len = values.len();
+    let pivot = len / 2;
+
+    values.swap(pivot, len - 1);
+    keys.swap(pivot, len - 1);
+
+    let mut store_idx = 0;
+    for i in 0..len - 1 {
+        let result = vm._lt(keys[i].clone(), keys[len - 1].clone())?;
+        let boolval = objbool::boolval(vm, result)?;
+        if boolval {
+            values.swap(i, store_idx);
+            keys.swap(i, store_idx);
+            store_idx += 1;
+        }
+    }
+
+    values.swap(store_idx, len - 1);
+    keys.swap(store_idx, len - 1);
+    Ok(store_idx)
+}
+
+fn do_sort(
+    vm: &mut VirtualMachine,
+    values: &mut Vec<PyObjectRef>,
+    key_func: Option<PyObjectRef>,
+    reverse: bool,
+) -> PyResult<()> {
+    // build a list of keys. If no keyfunc is provided, it's a copy of the list.
+    let mut keys: Vec<PyObjectRef> = vec![];
+    for x in values.iter() {
+        keys.push(match &key_func {
+            None => x.clone(),
+            Some(ref func) => vm.invoke(
+                (*func).clone(),
+                PyFuncArgs {
+                    args: vec![x.clone()],
+                    kwargs: vec![],
+                },
+            )?,
+        });
+    }
+
+    quicksort(vm, &mut keys, values)?;
+
+    if reverse {
+        values.reverse();
+    }
+
+    Ok(())
+}
+
 fn list_sort(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(list, Some(vm.ctx.list_type()))]);
-    let mut _elements = get_mut_elements(list);
-    unimplemented!("TODO: figure out how to invoke `sort_by` on a Vec");
-    // elements.sort_by();
-    // Ok(vm.get_none())
+    let key_func = args.get_optional_kwarg("key");
+    let reverse = args.get_optional_kwarg("reverse");
+    let reverse = match reverse {
+        None => false,
+        Some(val) => objbool::boolval(vm, val)?,
+    };
+
+    let elements_cell = get_elements_cell(list);
+    // replace list contents with [] for duration of sort.
+    // this prevents keyfunc from messing with the list and makes it easy to
+    // check if it tries to append elements to it.
+    let mut elements = elements_cell.replace(vec![]);
+    do_sort(vm, &mut elements, key_func, reverse)?;
+    let temp_elements = elements_cell.replace(elements);
+
+    if !temp_elements.is_empty() {
+        return Err(vm.new_value_error("list modified during sort".to_string()));
+    }
+
+    Ok(vm.get_none())
 }
 
 fn list_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
diff --git a/vm/src/obj/objsequence.rs b/vm/src/obj/objsequence.rs
@@ -302,6 +302,14 @@ pub fn seq_mul(elements: &[PyObjectRef], product: &PyObjectRef) -> Vec<PyObjectR
     new_elements
 }
 
+pub fn get_elements_cell<'a>(obj: &'a PyObjectRef) -> &'a RefCell<Vec<PyObjectRef>> {
+    if let PyObjectPayload::Sequence { ref elements } = obj.payload {
+        elements
+    } else {
+        panic!("Cannot extract elements from non-sequence");
+    }
+}
+
 pub fn get_elements<'a>(obj: &'a PyObjectRef) -> impl Deref<Target = Vec<PyObjectRef>> + 'a {
     if let PyObjectPayload::Sequence { ref elements } = obj.payload {
         elements.borrow()
