Skip to content

Commit e43b9e1

Browse files
feat(vm): add 'key=' and 'reverse=' to 'sorted()' and 'list.sort()'.
1 parent 6587b37 commit e43b9e1

5 files changed

Lines changed: 71 additions & 18 deletions

File tree

compiler/src/modules/vm/builtins/sequence.rs

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,21 +74,53 @@ impl<'a> VM<'a> {
7474
self.push(Val::int(n)); Ok(())
7575
}
7676

77-
pub fn call_sorted(&mut self) -> Result<(), VmErr> {
77+
pub fn call_sorted(&mut self, reverse: bool) -> Result<(), VmErr> {
7878
let o = self.pop()?;
7979
let mut items = self.extract_iter(o, false)?;
8080
self.sort_by_lt(&mut items)?;
81+
if reverse { items.reverse(); }
8182
self.alloc_and_push_list(items)
8283
}
8384

84-
/* sorted(iterable, key=fn), decorate-sort-undecorate via a parallel keys vec. key=None delegates to the no-key path. */
85-
pub fn call_sorted_with_key(&mut self, key: Option<Val>, chunk: &crate::modules::parser::SSAChunk, slots: &mut [Val]) -> Result<(), VmErr> {
85+
/* sorted(iterable, key=fn, reverse=False) — delegates to call_sorted when key is absent. */
86+
pub fn call_sorted_with_key(&mut self, key: Option<Val>, reverse: bool, chunk: &crate::modules::parser::SSAChunk, slots: &mut [Val]) -> Result<(), VmErr> {
8687
let key = match key {
8788
Some(k) if !k.is_none() => k,
88-
_ => return self.call_sorted(),
89+
_ => return self.call_sorted(reverse),
8990
};
9091
let o = self.pop()?;
9192
let items = self.extract_iter(o, false)?;
93+
let mut sorted = self.sort_by_key(items, key, chunk, slots)?;
94+
if reverse { sorted.reverse(); }
95+
self.alloc_and_push_list(sorted)
96+
}
97+
98+
/* list.sort(key=fn, reverse=False) in-place. Snapshots list before key calls so heap borrow ends before exec_call. */
99+
pub fn call_list_sort_keyed(&mut self, recv: Val, key: Option<Val>, reverse: bool, chunk: &crate::modules::parser::SSAChunk, slots: &mut [Val]) -> Result<(), VmErr> {
100+
let items = match self.heap.get(recv) {
101+
HeapObj::List(rc) => rc.borrow().clone(),
102+
_ => return Err(cold_type("sort: receiver is not a list")),
103+
};
104+
let mut result = if let Some(k) = key.filter(|k| !k.is_none()) {
105+
self.sort_by_key(items, k, chunk, slots)?
106+
} else {
107+
let mut s = items;
108+
self.sort_by_lt(&mut s)?;
109+
s
110+
};
111+
if reverse { result.reverse(); }
112+
let rc = match self.heap.get(recv) {
113+
HeapObj::List(rc) => rc.clone(),
114+
_ => return Err(cold_type("sort: receiver is not a list")),
115+
};
116+
*rc.borrow_mut() = result;
117+
self.mark_impure();
118+
self.push(Val::none());
119+
Ok(())
120+
}
121+
122+
/* Decorate-sort-undecorate: applies key fn to each item, sorts by resulting keys, returns reordered items. */
123+
fn sort_by_key(&mut self, items: Vec<Val>, key: Val, chunk: &crate::modules::parser::SSAChunk, slots: &mut [Val]) -> Result<Vec<Val>, VmErr> {
92124
let mut keys: Vec<Val> = Vec::with_capacity(items.len());
93125
for &item in &items {
94126
self.push(key);
@@ -111,8 +143,7 @@ impl<'a> VM<'a> {
111143
}
112144
});
113145
if let Some(e) = sort_err { return Err(e); }
114-
let sorted: Vec<Val> = indices.into_iter().map(|i| items[i]).collect();
115-
self.alloc_and_push_list(sorted)
146+
Ok(indices.into_iter().map(|i| items[i]).collect())
116147
}
117148

118149
/* In-place sort via `lt_vals`. Stashes the first error and surfaces it after the sort, `sort_by` closures can't return Result. */

compiler/src/modules/vm/handlers/function.rs

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ impl<'a> VM<'a> {
3030
OpCode::CallType => self.call_type(),
3131
OpCode::CallChr => self.call_chr(),
3232
OpCode::CallOrd => self.call_ord(),
33-
OpCode::CallSorted => self.call_sorted(),
33+
OpCode::CallSorted => self.call_sorted(false),
3434
OpCode::CallList => self.call_list(chunk, slots),
3535
OpCode::CallTuple => self.call_tuple(chunk, slots),
3636
OpCode::CallEnumerate => self.call_enumerate(),
@@ -241,6 +241,23 @@ impl<'a> VM<'a> {
241241
if let HeapObj::BoundMethod(recv, id) = self.heap.get(callee) {
242242
let recv = *recv;
243243
let id = *id;
244+
if id.name() == "sort" && !kw_flat.is_empty() {
245+
if !positional.is_empty() {
246+
return Err(cold_type("list.sort() takes no positional arguments"));
247+
}
248+
let mut sort_key: Option<Val> = None;
249+
let mut sort_reverse = false;
250+
for pair in kw_flat.chunks(2) {
251+
let (name_v, val_v) = (pair[0], pair[1]);
252+
let is_key = name_v.is_heap() && matches!(self.heap.get(name_v), HeapObj::Str(s) if s == "key");
253+
let is_reverse = name_v.is_heap() && matches!(self.heap.get(name_v), HeapObj::Str(s) if s == "reverse");
254+
if is_key { sort_key = Some(val_v); }
255+
else if is_reverse { sort_reverse = self.truthy(val_v); }
256+
else { return Err(cold_type("list.sort() got unexpected keyword argument")); }
257+
}
258+
self.call_list_sort_keyed(recv, sort_key, sort_reverse, chunk, slots)?;
259+
return Ok(true);
260+
}
244261
self.exec_bound_method(recv, id, positional, kw_flat)?;
245262
return Ok(true);
246263
}
@@ -594,20 +611,19 @@ impl<'a> VM<'a> {
594611
pub(crate) fn dispatch_native(&mut self, id: super::super::types::NativeFnId, positional: &[Val], kw: &[Val], chunk: &SSAChunk, slots: &mut [Val]) -> Result<(), VmErr> {
595612
use super::super::types::NativeFnId::*;
596613

597-
// `sorted()` is the only builtin taking kwargs (`key=fn`); extract it before the no-kw guard.
614+
// `sorted()` is the only builtin taking kwargs (`key=`, `reverse=`); extract them before the no-kw guard.
598615
let mut sort_key: Option<Val> = None;
616+
let mut sort_reverse = false;
599617
let leftover_storage: Vec<Val>;
600618
let kw_remaining: &[Val] = if id == Sorted {
601619
let mut leftover: Vec<Val> = Vec::new();
602620
for chunk_pair in kw.chunks(2) {
603621
let (name_v, val_v) = (chunk_pair[0], chunk_pair[1]);
604622
let is_key = name_v.is_heap() && matches!(self.heap.get(name_v), HeapObj::Str(s) if s == "key");
605-
if is_key {
606-
sort_key = Some(val_v);
607-
} else {
608-
leftover.push(name_v);
609-
leftover.push(val_v);
610-
}
623+
let is_reverse = name_v.is_heap() && matches!(self.heap.get(name_v), HeapObj::Str(s) if s == "reverse");
624+
if is_key { sort_key = Some(val_v); }
625+
else if is_reverse { sort_reverse = self.truthy(val_v); }
626+
else { leftover.push(name_v); leftover.push(val_v); }
611627
}
612628
leftover_storage = leftover;
613629
&leftover_storage
@@ -679,7 +695,7 @@ impl<'a> VM<'a> {
679695
Type => self.call_type(),
680696
Chr => self.call_chr(),
681697
Ord => self.call_ord(),
682-
Sorted => self.call_sorted_with_key(sort_key, chunk, slots),
698+
Sorted => self.call_sorted_with_key(sort_key, sort_reverse, chunk, slots),
683699
Enumerate => self.call_enumerate(),
684700
List => self.call_list(chunk, slots),
685701
Tuple => self.call_tuple(chunk, slots),

compiler/tests/cases/vm.json

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1776,5 +1776,9 @@
17761776
{"src": "print((1, 2, 3) < (1, 2))", "output": ["False"]},
17771777
{"src": "print(b'abc' < b'abd')", "output": ["True"]},
17781778
{"src": "print(sorted([[2], [1], [1, 0]]))", "output": ["[[1], [1, 0], [2]]"]},
1779-
{"src": "print(sorted([10, 5, 8, 1, 67], key=lambda x: -x))", "output": ["[67, 10, 8, 5, 1]"]}
1779+
{"src": "print(sorted([10, 5, 8, 1, 67], key=lambda x: -x))", "output": ["[67, 10, 8, 5, 1]"]},
1780+
{"src": "print(sorted([3, 1, 2], reverse=True))", "output": ["[3, 2, 1]"]},
1781+
{"src": "print(sorted([3, 1, 2], key=lambda x: -x, reverse=False))", "output": ["[3, 2, 1]"]},
1782+
{"src": "lst = [3, 1, 2]; lst.sort(reverse=True); print(lst)", "output": ["[3, 2, 1]"]},
1783+
{"src": "lst = [3, 1, 2]; lst.sort(key=lambda x: -x); print(lst)", "output": ["[3, 2, 1]"]}
17801784
]

docs/pages/reference/builtins.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -297,11 +297,13 @@ print(list(range(10, 0, -2)))
297297

298298
### sorted
299299

300-
New sorted list. No `key=` / `reverse=`, sort by derived value via a precomputed list of `(key, value)` tuples. Numbers, strings, bytes, and tuples/lists order lexicographically; mixed un-orderable types raise `TypeError`.
300+
New sorted list. Accepts `key=fn` and `reverse=True/False`. Numbers, strings, bytes, and tuples/lists order lexicographically; mixed un-orderable types raise `TypeError`.
301301

302302
```python
303303
print(sorted([3, 1, 4, 1, 5]))
304304
print(sorted("hello"))
305+
print(sorted([3, 1, 4, 1, 5], reverse=True))
306+
print(sorted(["banana", "apple", "kiwi"], key=len))
305307
```
306308

307309
```text Output

docs/pages/reference/methods.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ print(ys)
215215

216216
### Mutating
217217

218-
Return `None`, mutate in place. `extend` accepts any iterable. `sort` has no `key=` / `reverse=`, sort by derived key via precomputed tuples.
218+
Return `None`, mutate in place. `extend` accepts any iterable. `sort` accepts `key=fn` and `reverse=True/False`.
219219

220220
```python
221221
xs = [1, 2, 3]

0 commit comments

Comments
 (0)