Skip to content

Commit 1b32d07

Browse files
committed
perf: eliminate per-file heap allocation in sample buffer
Replace the Vec<u8> returned by the sampling path with a stack-allocated SampleBuf ([u8; 8192] + len), removing one heap allocation per file processed during binary/language detection.
1 parent 8ca6fd9 commit 1b32d07

1 file changed

Lines changed: 35 additions & 14 deletions

File tree

src/analysis/file_io.rs

Lines changed: 35 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,26 @@ use crate::langs::Language;
1818
const MMAP_THRESHOLD: u64 = 256 * 1024;
1919
/// Size of sample chunks extracted from files for binary/language detection. For large files, we sample from both the start and middle of the file.
2020
const SAMPLE_SIZE: usize = 4 * 1024;
21+
const MAX_SAMPLE: usize = SAMPLE_SIZE * 2;
22+
23+
/// Stack-allocated sample buffer; avoids a heap allocation per file during detection.
24+
pub(super) struct SampleBuf {
25+
data: [u8; MAX_SAMPLE],
26+
len: usize,
27+
}
28+
29+
impl SampleBuf {
30+
const fn new() -> Self {
31+
Self { data: [0u8; MAX_SAMPLE], len: 0 }
32+
}
33+
}
34+
35+
impl std::ops::Deref for SampleBuf {
36+
type Target = [u8];
37+
fn deref(&self) -> &[u8] {
38+
&self.data[..self.len]
39+
}
40+
}
2141

2242
/// Helper to create error context for file opening operations.
2343
fn open_file_context(path: &Path) -> String {
@@ -110,7 +130,7 @@ impl FileSource {
110130
}
111131
}
112132

113-
pub(super) fn sample(&mut self, file_size: u64) -> Result<Vec<u8>> {
133+
pub(super) fn sample(&mut self, file_size: u64) -> Result<SampleBuf> {
114134
match self {
115135
Self::Buffered(file) => sample_file(file, file_size),
116136
Self::Mapped(mmap) => Ok(sample_from_slice(mmap)),
@@ -153,32 +173,33 @@ fn sample_ranges(file_len: u64) -> (usize, Option<(u64, usize)>) {
153173
(start_len, Some((mid_offset, mid_len)))
154174
}
155175

156-
fn sample_file(file: &mut File, file_size: u64) -> Result<Vec<u8>> {
157-
let mut buffer = Vec::with_capacity(SAMPLE_SIZE * 2);
158-
let mut chunk = [0u8; SAMPLE_SIZE];
176+
fn sample_file(file: &mut File, file_size: u64) -> Result<SampleBuf> {
177+
let mut buf = SampleBuf::new();
159178
let (start_len, mid_range) = sample_ranges(file_size);
160-
let read_start = file.read(&mut chunk[..start_len])?;
161-
buffer.extend_from_slice(&chunk[..read_start]);
179+
let read_start = file.read(&mut buf.data[..start_len])?;
180+
buf.len = read_start;
162181
if let Some((mid_offset, mid_len)) = mid_range {
163182
file.seek(SeekFrom::Start(mid_offset))?;
164-
let read_mid = file.read(&mut chunk[..mid_len])?;
165-
buffer.extend_from_slice(&chunk[..read_mid]);
183+
let read_mid = file.read(&mut buf.data[buf.len..buf.len + mid_len])?;
184+
buf.len += read_mid;
166185
}
167186
// Reset for actual reading.
168187
file.rewind()?;
169-
Ok(buffer)
188+
Ok(buf)
170189
}
171190

172-
fn sample_from_slice(file_bytes: &[u8]) -> Vec<u8> {
173-
let mut samples = Vec::with_capacity(SAMPLE_SIZE * 2);
191+
fn sample_from_slice(file_bytes: &[u8]) -> SampleBuf {
192+
let mut buf = SampleBuf::new();
174193
let (start_len, mid_range) = sample_ranges(file_bytes.len() as u64);
175-
samples.extend_from_slice(&file_bytes[..start_len]);
194+
buf.data[..start_len].copy_from_slice(&file_bytes[..start_len]);
195+
buf.len = start_len;
176196
if let Some((mid_offset, mid_len)) = mid_range {
177197
let offset =
178198
usize::try_from(mid_offset).expect("mid_offset derives from file_bytes.len() which is already a usize");
179-
samples.extend_from_slice(&file_bytes[offset..offset + mid_len]);
199+
buf.data[buf.len..buf.len + mid_len].copy_from_slice(&file_bytes[offset..offset + mid_len]);
200+
buf.len += mid_len;
180201
}
181-
samples
202+
buf
182203
}
183204

184205
fn process_file_buffered(

0 commit comments

Comments
 (0)