|
1 use hg::errors::HgError; |
|
2 use hg::revlog::inner_revlog::RevisionBuffer; |
1 use pyo3::buffer::{Element, PyBuffer}; |
3 use pyo3::buffer::{Element, PyBuffer}; |
2 use pyo3::exceptions::PyValueError; |
4 use pyo3::exceptions::PyValueError; |
3 use pyo3::prelude::*; |
5 use pyo3::prelude::*; |
4 use pyo3::types::PyDict; |
6 use pyo3::types::{PyBytes, PyDict}; |
|
7 |
5 /// Create the module, with `__package__` given from parent |
8 /// Create the module, with `__package__` given from parent |
6 /// |
9 /// |
7 /// According to PyO3 documentation, which links to |
10 /// According to PyO3 documentation, which links to |
8 /// <https://github.com/PyO3/pyo3/issues/1517>, the same convoluted |
11 /// <https://github.com/PyO3/pyo3/issues/1517>, the same convoluted |
9 /// write to sys.modules has to be made as with the `cpython` crate. |
12 /// write to sys.modules has to be made as with the `cpython` crate. |
75 )); |
78 )); |
76 }; |
79 }; |
77 |
80 |
78 Ok((buf, Box::new(bytes))) |
81 Ok((buf, Box::new(bytes))) |
79 } |
82 } |
|
83 |
|
84 /// Takes an initialization function `init` which writes bytes to a |
|
85 /// Python-backed buffer, to save on a (potentially large) memory allocation |
|
86 /// and copy. If `init` fails to write the full expected length `len`, an error |
|
87 /// is raised. |
|
88 pub fn with_pybytes_buffer<F>( |
|
89 py: Python, |
|
90 len: usize, |
|
91 init: F, |
|
92 ) -> Result<Py<PyBytes>, hg::revlog::RevlogError> |
|
93 where |
|
94 F: FnOnce( |
|
95 &mut dyn RevisionBuffer<Target = Py<PyBytes>>, |
|
96 ) -> Result<(), hg::revlog::RevlogError>, |
|
97 { |
|
98 // Largely inspired by code in PyO3 |
|
99 // https://pyo3.rs/main/doc/pyo3/types/struct.pybytes#method.new_bound_with |
|
100 unsafe { |
|
101 let pyptr = pyo3::ffi::PyBytes_FromStringAndSize( |
|
102 std::ptr::null(), |
|
103 len as pyo3::ffi::Py_ssize_t, |
|
104 ); |
|
105 let pybytes = Bound::from_owned_ptr_or_err(py, pyptr) |
|
106 .map_err(|e| HgError::abort_simple(e.to_string()))? |
|
107 .downcast_into_unchecked(); |
|
108 let buffer: *mut u8 = pyo3::ffi::PyBytes_AsString(pyptr).cast(); |
|
109 debug_assert!(!buffer.is_null()); |
|
110 let mut rev_buf = PyRevisionBuffer::new(pybytes.unbind(), buffer, len); |
|
111 // Initialise the bytestring in init |
|
112 // If init returns an Err, the buffer is deallocated by `pybytes` |
|
113 init(&mut rev_buf).map(|_| rev_buf.finish()) |
|
114 } |
|
115 } |
|
116 |
|
117 /// Wrapper around a Python-provided buffer into which the revision contents |
|
118 /// will be written. Done for speed in order to save a large allocation + copy. |
|
119 struct PyRevisionBuffer { |
|
120 py_bytes: Py<PyBytes>, |
|
121 _buf: *mut u8, |
|
122 len: usize, |
|
123 current_buf: *mut u8, |
|
124 current_len: usize, |
|
125 } |
|
126 |
|
127 impl PyRevisionBuffer { |
|
128 /// # Safety |
|
129 /// |
|
130 /// `buf` should be the start of the allocated bytes of `bytes`, and `len` |
|
131 /// exactly the length of said allocated bytes. |
|
132 #[inline] |
|
133 unsafe fn new(bytes: Py<PyBytes>, buf: *mut u8, len: usize) -> Self { |
|
134 Self { |
|
135 py_bytes: bytes, |
|
136 _buf: buf, |
|
137 len, |
|
138 current_len: 0, |
|
139 current_buf: buf, |
|
140 } |
|
141 } |
|
142 |
|
143 /// Number of bytes that have been copied to. Will be different to the |
|
144 /// total allocated length of the buffer unless the revision is done being |
|
145 /// written. |
|
146 #[inline] |
|
147 fn current_len(&self) -> usize { |
|
148 self.current_len |
|
149 } |
|
150 } |
|
151 |
|
152 impl RevisionBuffer for PyRevisionBuffer { |
|
153 type Target = Py<PyBytes>; |
|
154 |
|
155 #[inline] |
|
156 fn extend_from_slice(&mut self, slice: &[u8]) { |
|
157 assert!(self.current_len + slice.len() <= self.len); |
|
158 unsafe { |
|
159 // We cannot use `copy_from_nonoverlapping` since it's *possible* |
|
160 // to create a slice from the same Python memory region using |
|
161 // [`PyBytesDeref`]. Probable that LLVM has an optimization anyway? |
|
162 self.current_buf.copy_from(slice.as_ptr(), slice.len()); |
|
163 self.current_buf = self.current_buf.add(slice.len()); |
|
164 } |
|
165 self.current_len += slice.len() |
|
166 } |
|
167 |
|
168 #[inline] |
|
169 fn finish(self) -> Self::Target { |
|
170 // catch unzeroed bytes before it becomes undefined behavior |
|
171 assert_eq!( |
|
172 self.current_len(), |
|
173 self.len, |
|
174 "not enough bytes read for revision" |
|
175 ); |
|
176 self.py_bytes |
|
177 } |
|
178 } |