Migrate fulltext_network_to_record to rust

by jelmer review by jelmer

breezy/bzr/versionedfile.py

@@ -19,7 +19,6 @@
 import functools
 import itertools
 import os
-import struct
 from copy import copy
 from io import BytesIO
 from typing import Any, Optional, Tuple
@@ -49,6 +48,7 @@
 ChunkedContentFactory = _versionedfile_rs.ChunkedContentFactory
 AbsentContentFactory = _versionedfile_rs.AbsentContentFactory
 record_to_fulltext_bytes = _versionedfile_rs.record_to_fulltext_bytes
+fulltext_network_to_record = _versionedfile_rs.fulltext_network_to_record
 
 
 adapter_registry = Registry[Tuple[str, str], Any, None]()
@@ -2068,17 +2068,6 @@ def read(self):
             yield from self._kind_factory[storage_kind](storage_kind, bytes, line_end)
 
 
-def fulltext_network_to_record(kind, bytes, line_end):
-    """Convert a network fulltext record to record."""
-    (meta_len,) = struct.unpack("!L", bytes[line_end : line_end + 4])
-    record_meta = bytes[line_end + 4 : line_end + 4 + meta_len]
-    key, parents = bencode.bdecode_as_tuple(record_meta)
-    if parents == b"nil":
-        parents = None
-    fulltext = bytes[line_end + 4 + meta_len :]
-    return [FulltextContentFactory(key, parents, None, fulltext)]
-
-
 def sort_groupcompress(parent_map):
     """Sort and group the keys in parent_map into groupcompress order.
 

crates/bazaar-py/src/versionedfile.rs

@@ -130,12 +130,28 @@ impl AbsentContentFactory {
     }
 }
 
+#[pyfunction]
+fn fulltext_network_to_record(
+    py: Python,
+    _kind: &str,
+    bytes: &[u8],
+    line_end: usize,
+) -> Vec<PyObject> {
+    let record = bazaar::versionedfile::fulltext_network_to_record(bytes, line_end);
+
+    let sub = PyClassInitializer::from(AbstractContentFactory(Box::new(record)))
+        .add_subclass(FulltextContentFactory);
+
+    vec![Py::new(py, sub).unwrap().to_object(py)]
+}
+
 pub(crate) fn _versionedfile_rs(py: Python) -> PyResult<&PyModule> {
     let m = PyModule::new(py, "versionedfile")?;
     m.add_class::<AbstractContentFactory>()?;
     m.add_class::<FulltextContentFactory>()?;
     m.add_class::<ChunkedContentFactory>()?;
     m.add_class::<AbsentContentFactory>()?;
     m.add_function(wrap_pyfunction!(record_to_fulltext_bytes, m)?)?;
+    m.add_function(wrap_pyfunction!(fulltext_network_to_record, m)?)?;
     Ok(m)
 }

crates/bazaar/src/versionedfile.rs

@@ -3,6 +3,7 @@ use byteorder::{BigEndian, WriteBytesExt};
 use pyo3::types::{PyBytes, PyTuple};
 use std::borrow::Cow;
 use std::collections::HashMap;
+use std::convert::TryInto;
 
 #[derive(Debug)]
 pub enum Error {
@@ -716,3 +717,72 @@ fn length_prefix(data: &[u8]) -> Vec<u8> {
 
     length_bytes
 }
+
+pub fn fulltext_network_to_record(bytes: &[u8], line_end: usize) -> FulltextContentFactory {
+    // Extract meta_len from the network fulltext record
+    let meta_len_bytes: [u8; 4] = bytes[line_end..line_end + 4]
+        .try_into()
+        .expect("Expected 4 bytes for meta_len");
+    let meta_len = u32::from_be_bytes(meta_len_bytes) as usize;
+
+    // Extract record_meta using meta_len
+    let record_meta = &bytes[line_end + 4..line_end + 4 + meta_len];
+
+    // Decode record_meta using Bencode
+    let mut decoder = bendy::decoding::Decoder::new(record_meta);
+
+    let mut tuple = decoder
+        .next_object()
+        .expect("Failed to decode record_meta using Bencode")
+        .expect("Failed to decode tuple using Bencode")
+        .try_into_list()
+        .unwrap();
+
+    fn decode_key(o: bendy::decoding::Object) -> Key {
+        let mut ret = vec![];
+
+        let mut l = o.try_into_list().unwrap();
+
+        while let Some(b) = l.next_object().unwrap() {
+            ret.push(b.try_into_bytes().unwrap().to_vec());
+        }
+
+        Key::Fixed(ret)
+    }
+
+    let key = decode_key(
+        tuple
+            .next_object()
+            .expect("Failed to decode record_meta using Bencode")
+            .expect("Failed to decode key using Bencode"),
+    );
+
+    let parents = tuple
+        .next_object()
+        .expect("Failed to decode record_meta using Bencode")
+        .expect("Failed to decode parents using Bencode");
+
+    // Convert parents from "nil" to None
+    let parents = match parents {
+        bendy::decoding::Object::Bytes(bytes) => {
+            if bytes == b"nil" {
+                None
+            } else {
+                panic!("Expected parents to be a list or nil");
+            }
+        }
+        bendy::decoding::Object::List(mut l) => {
+            let mut parents = vec![];
+            while let Some(parent) = l.next_object().unwrap() {
+                parents.push(decode_key(parent));
+            }
+            Some(parents)
+        }
+        _ => panic!("Expected parents to be a list or nil"),
+    };
+
+    // Extract fulltext from the remaining bytes
+    let fulltext = &bytes[line_end + 4 + meta_len..];
+
+    FulltextContentFactory::new(None, key, parents, fulltext.to_vec())
+}