use alloc::{boxed::Box, vec};
use core::mem;
use ax_driver::{AxBlockDevice, PartitionBlockDevice, PartitionRegion, prelude::*};
fn take<'a>(buf: &mut &'a [u8], cnt: usize) -> &'a [u8] {
let (first, rem) = buf.split_at(cnt);
*buf = rem;
first
}
fn take_mut<'a>(buf: &mut &'a mut [u8], cnt: usize) -> &'a mut [u8] {
let (first, rem) = mem::take(buf).split_at_mut(cnt);
*buf = rem;
first
}
pub struct SeekableDisk {
dev: PartitionBlockDevice<AxBlockDevice>,
block_id: u64,
offset: usize,
block_size_log2: u8,
read_buffer: Box<[u8]>,
write_buffer: Box<[u8]>,
write_buffer_dirty: bool,
}
impl SeekableDisk {
pub fn new(dev: AxBlockDevice, region: PartitionRegion) -> Self {
assert!(dev.block_size().is_power_of_two());
let block_size = dev.block_size();
let block_size_log2 = block_size.trailing_zeros() as u8;
let read_buffer = vec![0u8; block_size].into_boxed_slice();
let write_buffer = vec![0u8; block_size].into_boxed_slice();
Self {
dev: PartitionBlockDevice::new(dev, region),
block_id: 0,
offset: 0,
block_size_log2,
read_buffer,
write_buffer,
write_buffer_dirty: false,
}
}
pub fn size(&self) -> u64 {
self.dev.num_blocks() << self.block_size_log2
}
pub fn block_size(&self) -> usize {
1 << self.block_size_log2
}
pub fn position(&self) -> u64 {
(self.block_id << self.block_size_log2) + self.offset as u64
}
pub fn set_position(&mut self, pos: u64) -> DevResult<()> {
self.flush()?;
self.block_id = pos >> self.block_size_log2;
self.offset = pos as usize & (self.block_size() - 1);
Ok(())
}
pub fn flush(&mut self) -> DevResult<()> {
if self.write_buffer_dirty {
self.dev.write_block(self.block_id, &self.write_buffer)?;
self.write_buffer_dirty = false;
}
Ok(())
}
fn read_partial(&mut self, buf: &mut &mut [u8]) -> DevResult<usize> {
self.flush()?;
self.dev.read_block(self.block_id, &mut self.read_buffer)?;
let data = &self.read_buffer[self.offset..];
let length = buf.len().min(data.len());
take_mut(buf, length).copy_from_slice(&data[..length]);
self.offset += length;
if self.offset == self.block_size() {
self.block_id += 1;
self.offset = 0;
}
Ok(length)
}
pub fn read(&mut self, mut buf: &mut [u8]) -> DevResult<usize> {
let mut read = 0;
if self.offset != 0 {
read += self.read_partial(&mut buf)?;
}
if buf.len() >= self.block_size() {
let blocks = buf.len() >> self.block_size_log2;
let length = blocks << self.block_size_log2;
self.dev
.read_block(self.block_id, take_mut(&mut buf, length))?;
read += length;
self.block_id = self
.block_id
.checked_add(blocks as u64)
.ok_or(DevError::BadState)?;
}
if !buf.is_empty() {
read += self.read_partial(&mut buf)?;
}
Ok(read)
}
fn write_partial(&mut self, buf: &mut &[u8]) -> DevResult<usize> {
if !self.write_buffer_dirty {
self.dev.read_block(self.block_id, &mut self.write_buffer)?;
self.write_buffer_dirty = true;
}
let data = &mut self.write_buffer[self.offset..];
let length = buf.len().min(data.len());
data[..length].copy_from_slice(take(buf, length));
self.offset += length;
if self.offset == self.block_size() {
self.flush()?;
self.block_id += 1;
self.offset = 0;
}
Ok(length)
}
pub fn write(&mut self, mut buf: &[u8]) -> DevResult<usize> {
let mut written = 0;
if self.offset != 0 {
written += self.write_partial(&mut buf)?;
}
if buf.len() >= self.block_size() {
let blocks = buf.len() >> self.block_size_log2;
let length = blocks << self.block_size_log2;
self.dev
.write_block(self.block_id, take(&mut buf, length))?;
written += length;
self.block_id = self
.block_id
.checked_add(blocks as u64)
.ok_or(DevError::BadState)?;
}
if !buf.is_empty() {
written += self.write_partial(&mut buf)?;
}
Ok(written)
}
}