Simplify sorted tree items a bit

crates/objects/src/lib.rs

@@ -19,7 +19,6 @@
  */
 
 use memchr::memchr;
-use std::borrow::Cow;
 
 use pyo3::exceptions::PyTypeError;
 use pyo3::import_exception;
@@ -88,14 +87,20 @@ fn parse_tree(
     Ok(entries)
 }
 
-fn name_with_suffix(mode: u32, name: &[u8]) -> Cow<[u8]> {
-    if mode & S_IFDIR != 0 {
-        let mut v = name.to_vec();
-        v.push(b'/');
-        v.into()
-    } else {
-        name.into()
+fn cmp_with_suffix(a: (u32, &[u8]), b: (u32, &[u8])) -> std::cmp::Ordering {
+    let len = std::cmp::min(a.1.len(), b.1.len());
+    let cmp = a.1[..len].cmp(&b.1[..len]);
+    if cmp != std::cmp::Ordering::Equal {
+        return cmp;
     }
+
+    let c1 =
+        a.1.get(len)
+            .map_or_else(|| if a.0 & S_IFDIR != 0 { b'/' } else { 0 }, |&c| c);
+    let c2 =
+        b.1.get(len)
+            .map_or_else(|| if b.0 & S_IFDIR != 0 { b'/' } else { 0 }, |&c| c);
+    c1.cmp(&c2)
 }
 
 /// Iterate over a tree entries dictionary.
@@ -114,19 +119,19 @@ fn sorted_tree_items(
     entries: &Bound<PyDict>,
     name_order: bool,
 ) -> PyResult<Vec<PyObject>> {
-    let mut qsort_entries = Vec::new();
-    for (name, e) in entries.iter() {
-        let (mode, sha): (u32, Vec<u8>) = e
-            .extract()
-            .map_err(|e| PyTypeError::new_err((format!("invalid type: {}", e),)))?;
-        qsort_entries.push((name.extract::<Vec<u8>>().unwrap(), mode, sha));
-    }
+    let mut qsort_entries = entries
+        .iter()
+        .map(|(name, value)| -> PyResult<(Vec<u8>, u32, Vec<u8>)> {
+            let value = value
+                .extract::<(u32, Vec<u8>)>()
+                .map_err(|e| PyTypeError::new_err((format!("invalid type: {}", e),)))?;
+            Ok((name.extract::<Vec<u8>>().unwrap(), value.0, value.1))
+        })
+        .collect::<PyResult<Vec<(Vec<u8>, u32, Vec<u8>)>>>()?;
     if name_order {
         qsort_entries.sort_by(|a, b| a.0.cmp(&b.0));
     } else {
-        qsort_entries.sort_by(|a, b| {
-            name_with_suffix(a.1, a.0.as_slice()).cmp(&name_with_suffix(b.1, b.0.as_slice()))
-        });
+        qsort_entries.sort_by(|a, b| cmp_with_suffix((a.1, a.0.as_slice()), (b.1, b.0.as_slice())));
     }
     let objectsm = py.import_bound("dulwich.objects")?;
     let tree_entry_cls = objectsm.getattr("TreeEntry")?;

dulwich/objects.py

@@ -1667,7 +1667,12 @@ _parse_tree_py = parse_tree
 _sorted_tree_items_py = sorted_tree_items
 try:
     # Try to import Rust versions
-    from dulwich._objects import parse_tree as _parse_tree_rs, sorted_tree_items as _sorted_tree_items_rs
+    from dulwich._objects import (
+        parse_tree as _parse_tree_rs,
+    )
+    from dulwich._objects import (
+        sorted_tree_items as _sorted_tree_items_rs,
+    )
 except ImportError:
     pass
 else:

tests/test_objects.py

@@ -834,14 +834,20 @@ _SORTED_TREE_ITEMS = [
 
 
 _TREE_ITEMS_BUG_1325 = {
-    b'dir': (stat.S_IFDIR | 0o644, b'5944b31ff85b415573d1a43eb942e2dea30ab8be'),
-    b'dira': (0o100644, b'cf7a729ca69bfabd0995fc9b083e86a18215bd91'),
+    b"dir": (stat.S_IFDIR | 0o644, b"5944b31ff85b415573d1a43eb942e2dea30ab8be"),
+    b"dira": (0o100644, b"cf7a729ca69bfabd0995fc9b083e86a18215bd91"),
 }
 
 
 _SORTED_TREE_ITEMS_BUG_1325 = [
-    TreeEntry(path=b'dir', mode=stat.S_IFDIR | 0o644, sha=b'5944b31ff85b415573d1a43eb942e2dea30ab8be'),
-    TreeEntry(path=b'dira', mode=0o100644, sha=b'cf7a729ca69bfabd0995fc9b083e86a18215bd91'),
+    TreeEntry(
+        path=b"dir",
+        mode=stat.S_IFDIR | 0o644,
+        sha=b"5944b31ff85b415573d1a43eb942e2dea30ab8be",
+    ),
+    TreeEntry(
+        path=b"dira", mode=0o100644, sha=b"cf7a729ca69bfabd0995fc9b083e86a18215bd91"
+    ),
 ]
 
 
@@ -898,7 +904,9 @@ class TreeTests(ShaFileCheckTests):
         )
 
     test_parse_tree = functest_builder(_do_test_parse_tree, _parse_tree_py)
-    test_parse_tree_extension = ext_functest_builder(_do_test_parse_tree, _parse_tree_rs)
+    test_parse_tree_extension = ext_functest_builder(
+        _do_test_parse_tree, _parse_tree_rs
+    )
 
     def _do_test_sorted_tree_items(self, sorted_tree_items):
         def do_sort(entries, name_order):
@@ -909,7 +917,10 @@ class TreeTests(ShaFileCheckTests):
         self.assertIsInstance(actual[0], TreeEntry)
 
         actual = do_sort(_TREE_ITEMS_BUG_1325, False)
-        self.assertEqual(key_entry((b"a", (0o40644, b"cf7a729ca69bfabd0995fc9b083e86a18215bd91"))), b"a/")
+        self.assertEqual(
+            key_entry((b"a", (0o40644, b"cf7a729ca69bfabd0995fc9b083e86a18215bd91"))),
+            b"a/",
+        )
         self.assertEqual(_SORTED_TREE_ITEMS_BUG_1325, actual)
         self.assertIsInstance(actual[0], TreeEntry)
 
@@ -920,7 +931,7 @@ class TreeTests(ShaFileCheckTests):
         # AttributeError.
         errors = (TypeError, ValueError, AttributeError)
         self.assertRaises(errors, do_sort, b"foo", False)
-        self.assertRaises(errors, do_sort, {b"foo": (1, 2, 3)},  False)
+        self.assertRaises(errors, do_sort, {b"foo": (1, 2, 3)}, False)
 
         myhexsha = b"d80c186a03f423a81b39df39dc87fd269736ca86"
         self.assertRaises(errors, do_sort, {b"foo": (b"xxx", myhexsha)}, False)
@@ -963,8 +974,6 @@ class TreeTests(ShaFileCheckTests):
                     0o100755,
                     b"d80c186a03f423a81b39df39dc87fd269736ca86",
                 ),
-
-
             ],
             list(sorted_tree_items(_TREE_ITEMS, True)),
         )