1use crate::{BlkError, DeviceInfo, QueueLimits};
2
3#[derive(Debug, Clone, Copy)]
4pub struct TransferRuntimeCaps {
5 pub max_transfer_bytes: usize,
6 pub max_segments: usize,
7}
8
9impl TransferRuntimeCaps {
10 pub const fn new(max_transfer_bytes: usize, max_segments: usize) -> Self {
11 Self {
12 max_transfer_bytes,
13 max_segments,
14 }
15 }
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct TransferSegment {
20 pub byte_offset: usize,
22 pub byte_len: usize,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct TransferChunk {
27 pub lba: u64,
28 pub block_count: u32,
29 pub byte_offset: usize,
30 pub byte_len: usize,
31 max_segment_size: usize,
32}
33
34impl TransferChunk {
35 pub fn segments(self) -> TransferSegments {
36 TransferSegments {
37 remaining_len: self.byte_len,
38 byte_offset: 0,
39 max_segment_size: self.max_segment_size,
40 }
41 }
42}
43
44pub struct TransferSegments {
45 remaining_len: usize,
46 byte_offset: usize,
47 max_segment_size: usize,
48}
49
50impl Iterator for TransferSegments {
51 type Item = TransferSegment;
52
53 fn next(&mut self) -> Option<Self::Item> {
54 if self.remaining_len == 0 {
55 return None;
56 }
57
58 let byte_len = self.remaining_len.min(self.max_segment_size);
59 let segment = TransferSegment {
60 byte_offset: self.byte_offset,
61 byte_len,
62 };
63 self.byte_offset += byte_len;
64 self.remaining_len -= byte_len;
65 Some(segment)
66 }
67}
68
69impl ExactSizeIterator for TransferSegments {
70 fn len(&self) -> usize {
71 self.remaining_len.div_ceil(self.max_segment_size)
72 }
73}
74
75#[derive(Debug, Clone, Copy)]
76pub struct TransferPlanner {
77 device: DeviceInfo,
78 limits: QueueLimits,
79 max_chunk_size: usize,
80}
81
82impl TransferPlanner {
83 pub fn new(
84 device: DeviceInfo,
85 limits: QueueLimits,
86 caps: TransferRuntimeCaps,
87 ) -> Result<Self, BlkError> {
88 let max_chunk_size = planned_transfer_size(device, limits, caps)?;
89
90 Ok(Self {
91 device,
92 limits,
93 max_chunk_size,
94 })
95 }
96
97 pub const fn chunk_size(&self) -> usize {
98 self.max_chunk_size
99 }
100
101 pub fn plan(&self, lba: u64, byte_len: usize) -> Result<TransferPlan, BlkError> {
102 TransferPlan::new(self.device, self.limits, self.max_chunk_size, lba, byte_len)
103 }
104}
105
106#[derive(Debug, Clone, Copy)]
107pub struct TransferPlan {
108 next_lba: u64,
109 byte_offset: usize,
110 remaining_bytes: usize,
111 block_size: usize,
112 max_chunk_size: usize,
113 max_segment_size: usize,
114}
115
116impl TransferPlan {
117 fn new(
118 device: DeviceInfo,
119 limits: QueueLimits,
120 max_chunk_size: usize,
121 lba: u64,
122 byte_len: usize,
123 ) -> Result<Self, BlkError> {
124 let block_size = device.logical_block_size;
125 if block_size == 0 || byte_len == 0 || !byte_len.is_multiple_of(block_size) {
126 return Err(BlkError::InvalidRequest);
127 }
128
129 let block_count = byte_len / block_size;
130 let block_count_u64 = u64::try_from(block_count).map_err(|_| BlkError::InvalidRequest)?;
131 if lba >= device.num_blocks
132 || lba
133 .checked_add(block_count_u64)
134 .is_none_or(|end| end > device.num_blocks)
135 {
136 return Err(BlkError::InvalidBlockIndex(lba));
137 }
138
139 Ok(Self {
140 next_lba: lba,
141 byte_offset: 0,
142 remaining_bytes: byte_len,
143 block_size,
144 max_chunk_size,
145 max_segment_size: limits.max_segment_size,
146 })
147 }
148}
149
150fn planned_transfer_size(
151 device: DeviceInfo,
152 limits: QueueLimits,
153 caps: TransferRuntimeCaps,
154) -> Result<usize, BlkError> {
155 let block_size = device.logical_block_size;
156 let max_segments = limits.max_segments.min(caps.max_segments);
157 if block_size == 0
158 || limits.max_blocks_per_request == 0
159 || max_segments == 0
160 || limits.max_segment_size == 0
161 || caps.max_transfer_bytes == 0
162 {
163 return Err(BlkError::InvalidRequest);
164 }
165
166 let max_by_blocks = block_size.saturating_mul(limits.max_blocks_per_request as usize);
167 let max_by_segments = limits.max_segment_size.saturating_mul(max_segments);
168 let max_chunk_size = [max_by_blocks, max_by_segments, caps.max_transfer_bytes]
169 .into_iter()
170 .min()
171 .ok_or(BlkError::InvalidRequest)?;
172 let max_chunk_size = align_down(max_chunk_size, block_size);
173 if max_chunk_size < block_size {
174 return Err(BlkError::InvalidRequest);
175 }
176 Ok(max_chunk_size)
177}
178
179impl Iterator for TransferPlan {
180 type Item = TransferChunk;
181
182 fn next(&mut self) -> Option<Self::Item> {
183 if self.remaining_bytes == 0 {
184 return None;
185 }
186
187 let byte_len = self.remaining_bytes.min(self.max_chunk_size);
188 let block_count = byte_len / self.block_size;
189 let block_count_u32 = block_count as u32;
190 let chunk = TransferChunk {
191 lba: self.next_lba,
192 block_count: block_count_u32,
193 byte_offset: self.byte_offset,
194 byte_len,
195 max_segment_size: self.max_segment_size,
196 };
197
198 self.next_lba += block_count as u64;
199 self.byte_offset += byte_len;
200 self.remaining_bytes -= byte_len;
201 Some(chunk)
202 }
203}
204
205fn align_down(value: usize, align: usize) -> usize {
206 value / align * align
207}
208
209#[cfg(test)]
210mod tests {
211 use alloc::vec::Vec;
212
213 use super::*;
214 use crate::{QueueInfo, RequestFlags};
215
216 fn queue_info_with(limits: QueueLimits) -> QueueInfo {
217 QueueInfo {
218 id: 0,
219 device: DeviceInfo::new(64, 512),
220 limits,
221 }
222 }
223
224 fn queue_limits(
225 max_blocks_per_request: u32,
226 max_segments: usize,
227 max_segment_size: usize,
228 ) -> QueueLimits {
229 QueueLimits {
230 dma_mask: u64::MAX,
231 dma_alignment: 512,
232 max_blocks_per_request,
233 max_segments,
234 max_segment_size,
235 supported_flags: RequestFlags::NONE,
236 supports_flush: false,
237 supports_discard: false,
238 supports_write_zeroes: false,
239 }
240 }
241
242 fn test_runtime_caps() -> TransferRuntimeCaps {
243 TransferRuntimeCaps {
244 max_transfer_bytes: 16 * 1024,
245 max_segments: 16,
246 }
247 }
248
249 fn chunk_summary(chunks: &[TransferChunk]) -> Vec<(u64, u32, usize, usize, usize)> {
250 chunks
251 .iter()
252 .map(|chunk| {
253 let segments = chunk.segments();
254 (
255 chunk.lba,
256 chunk.block_count,
257 chunk.byte_offset,
258 chunk.byte_len,
259 segments.len(),
260 )
261 })
262 .collect()
263 }
264
265 #[test]
266 fn simple_limits_allow_single_block_transfers() {
267 let info = queue_info_with(QueueLimits::simple(512, u64::MAX));
268 let planner = TransferPlanner::new(info.device, info.limits, test_runtime_caps()).unwrap();
269 let plan = planner.plan(0, 2048).unwrap();
270 let chunks: Vec<_> = plan.collect();
271
272 assert_eq!(planner.chunk_size(), 512);
273 assert_eq!(
274 chunk_summary(&chunks),
275 [
276 (0, 1, 0, 512, 1),
277 (1, 1, 512, 512, 1),
278 (2, 1, 1024, 512, 1),
279 (3, 1, 1536, 512, 1),
280 ]
281 );
282 }
283
284 #[test]
285 fn transfer_plan_chunks_by_runtime_cap() {
286 let info = queue_info_with(queue_limits(16, 4, 4096));
287 let planner = TransferPlanner::new(
288 info.device,
289 info.limits,
290 TransferRuntimeCaps {
291 max_transfer_bytes: 2048,
292 max_segments: 16,
293 },
294 )
295 .unwrap();
296 let plan = planner.plan(4, 5120).unwrap();
297 let chunks: Vec<_> = plan.collect();
298
299 assert_eq!(planner.chunk_size(), 2048);
300 assert_eq!(
301 chunk_summary(&chunks),
302 [
303 (4, 4, 0, 2048, 1),
304 (8, 4, 2048, 2048, 1),
305 (12, 2, 4096, 1024, 1),
306 ]
307 );
308 }
309
310 #[test]
311 fn transfer_chunk_segments_split_by_hard_segment_size() {
312 let info = queue_info_with(queue_limits(16, 4, 1024));
313 let planner = TransferPlanner::new(info.device, info.limits, test_runtime_caps()).unwrap();
314 let mut plan = planner.plan(0, 4096).unwrap();
315 let chunk = plan.next().unwrap();
316 let segment_iter = chunk.segments();
317 assert_eq!(segment_iter.len(), 4);
318 let segments: Vec<_> = segment_iter.collect();
319
320 assert_eq!(
321 segments,
322 [
323 TransferSegment {
324 byte_offset: 0,
325 byte_len: 1024,
326 },
327 TransferSegment {
328 byte_offset: 1024,
329 byte_len: 1024,
330 },
331 TransferSegment {
332 byte_offset: 2048,
333 byte_len: 1024,
334 },
335 TransferSegment {
336 byte_offset: 3072,
337 byte_len: 1024,
338 },
339 ]
340 );
341 assert!(plan.next().is_none());
342 }
343
344 #[test]
345 fn transfer_plan_clamps_to_hard_block_count() {
346 let info = queue_info_with(queue_limits(4, 8, 4096));
347 let planner = TransferPlanner::new(info.device, info.limits, test_runtime_caps()).unwrap();
348 let plan = planner.plan(0, 5120).unwrap();
349 let chunks: Vec<_> = plan.collect();
350
351 assert_eq!(planner.chunk_size(), 2048);
352 assert_eq!(
353 chunks
354 .iter()
355 .map(|chunk| chunk.byte_len)
356 .collect::<Vec<_>>(),
357 [2048, 2048, 1024]
358 );
359 }
360
361 #[test]
362 fn transfer_plan_clamps_to_runtime_limits() {
363 let info = queue_info_with(queue_limits(16, 8, 2048));
364 let planner = TransferPlanner::new(
365 info.device,
366 info.limits,
367 TransferRuntimeCaps {
368 max_transfer_bytes: 4096,
369 max_segments: 1,
370 },
371 )
372 .unwrap();
373 let plan = planner.plan(0, 4096).unwrap();
374 let chunks: Vec<_> = plan.collect();
375
376 assert_eq!(planner.chunk_size(), 2048);
377 assert_eq!(
378 chunks
379 .iter()
380 .map(|chunk| chunk.byte_len)
381 .collect::<Vec<_>>(),
382 [2048, 2048]
383 );
384 }
385
386 #[test]
387 fn transfer_planner_rejects_too_small_runtime_cap() {
388 let info = queue_info_with(queue_limits(16, 8, 2048));
389
390 assert_eq!(
391 TransferPlanner::new(
392 info.device,
393 info.limits,
394 TransferRuntimeCaps {
395 max_transfer_bytes: 511,
396 max_segments: 1,
397 },
398 )
399 .unwrap_err(),
400 BlkError::InvalidRequest
401 );
402 }
403
404 #[test]
405 fn transfer_planner_does_not_depend_on_queue_identity() {
406 let mut info = queue_info_with(queue_limits(16, 8, 2048));
407 let first = TransferPlanner::new(info.device, info.limits, test_runtime_caps()).unwrap();
408 info.id = 7;
409 let second = TransferPlanner::new(info.device, info.limits, test_runtime_caps()).unwrap();
410
411 assert_eq!(first.chunk_size(), second.chunk_size());
412 }
413
414 #[test]
415 fn transfer_planner_checks_range_when_creating_plan() {
416 let info = queue_info_with(QueueLimits::simple(512, u64::MAX));
417 let planner = TransferPlanner::new(info.device, info.limits, test_runtime_caps()).unwrap();
418
419 assert_eq!(
420 planner.plan(63, 1024).unwrap_err(),
421 BlkError::InvalidBlockIndex(63)
422 );
423 }
424}