From 6e07fa91db62a0ef4fb8fe792127e5ed5361cbc7 Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Tue, 30 Jul 2024 19:51:14 +0300
Subject: [PATCH] dev/block: remove extra copy on aligned ops

---
 kernel/driver/block/ahci/src/command.rs     |  4 +-
 kernel/driver/block/ahci/src/port.rs        |  8 +--
 kernel/driver/block/nvme/src/drive.rs       | 22 +++---
 kernel/driver/block/nvme/src/lib.rs         |  5 +-
 kernel/driver/fs/ext2/src/dir.rs            | 26 ++-----
 kernel/driver/fs/ext2/src/file.rs           | 19 +++---
 kernel/driver/fs/ext2/src/lib.rs            | 19 ++----
 kernel/libk/libk-mm/src/lib.rs              | 46 +++++++++++++
 kernel/libk/libk-util/src/lib.rs            | 75 ++++++++++++++++++++-
 kernel/libk/libk-util/src/lru_hash_table.rs |  5 +-
 kernel/libk/src/task/sync.rs                |  1 -
 kernel/libk/src/vfs/block/cache.rs          | 34 +++++++---
 kernel/libk/src/vfs/block/device.rs         | 65 +++++++++++++++---
 kernel/libk/src/vfs/block/mod.rs            | 16 ++++-
 kernel/libk/src/vfs/block/partition.rs      | 42 ++++++++++--
 xtask/src/build/mod.rs                      |  7 +-
 16 files changed, 303 insertions(+), 91 deletions(-)

diff --git a/kernel/driver/block/ahci/src/command.rs b/kernel/driver/block/ahci/src/command.rs
index d2c57561..b98efcf7 100644
--- a/kernel/driver/block/ahci/src/command.rs
+++ b/kernel/driver/block/ahci/src/command.rs
@@ -2,7 +2,7 @@ use core::mem::{size_of, MaybeUninit};
 
 use libk_mm::{
     address::{AsPhysicalAddress, PhysicalAddress},
-    PageBox,
+    PageBox, PageSlice,
 };
 use tock_registers::register_structs;
 
@@ -85,7 +85,7 @@ impl AtaIdentify {
 }
 
 impl AtaReadDmaEx {
-    pub fn new(lba: u64, sector_count: usize, buffer: &PageBox<[MaybeUninit<u8>]>) -> Self {
+    pub fn new(lba: u64, sector_count: usize, buffer: &PageSlice<MaybeUninit<u8>>) -> Self {
         assert_eq!(buffer.len() % SECTOR_SIZE, 0);
         assert_ne!(buffer.len(), 0);
 
diff --git a/kernel/driver/block/ahci/src/port.rs b/kernel/driver/block/ahci/src/port.rs
index 9e2b78c6..86eb8016 100644
--- a/kernel/driver/block/ahci/src/port.rs
+++ b/kernel/driver/block/ahci/src/port.rs
@@ -10,7 +10,7 @@ use async_trait::async_trait;
 use bytemuck::Zeroable;
 use futures_util::task::AtomicWaker;
 use libk::vfs::block::NgBlockDevice;
-use libk_mm::{address::AsPhysicalAddress, device::DeviceMemoryIo, PageBox};
+use libk_mm::{address::AsPhysicalAddress, device::DeviceMemoryIo, PageBox, PageSlice};
 use libk_util::{sync::IrqSafeSpinlock, waker::QueueWaker, OneTimeInit};
 use tock_registers::interfaces::{Readable, Writeable};
 
@@ -297,13 +297,13 @@ impl NgBlockDevice for AhciPort {
     async fn read(
         &self,
         lba: u64,
-        buffer: &mut PageBox<[MaybeUninit<u8>]>,
+        buffer: &mut PageSlice<MaybeUninit<u8>>,
     ) -> Result<(), AhciError> {
-        let command = AtaReadDmaEx::new(lba, buffer.len() / SECTOR_SIZE, &buffer);
+        let command = AtaReadDmaEx::new(lba, buffer.len() / SECTOR_SIZE, buffer);
         self.submit(&command).await?.wait_for_completion().await
     }
 
-    async fn write(&self, _lba: u64, _buffer: PageBox<[u8]>) -> Result<(), AhciError> {
+    async fn write(&self, _lba: u64, _buffer: &PageSlice<u8>) -> Result<(), AhciError> {
         // TODO AtaDmaWriteEx
         Err(AhciError::FeatureNotImplemented)
     }
diff --git a/kernel/driver/block/nvme/src/drive.rs b/kernel/driver/block/nvme/src/drive.rs
index adf783db..30b2f8fc 100644
--- a/kernel/driver/block/nvme/src/drive.rs
+++ b/kernel/driver/block/nvme/src/drive.rs
@@ -4,7 +4,7 @@ use alloc::{boxed::Box, format};
 use async_trait::async_trait;
 use kernel_fs::devfs;
 use libk::vfs::block::{probe_partitions, NgBlockDevice, NgBlockDeviceWrapper};
-use libk_mm::{address::AsPhysicalAddress, PageBox};
+use libk_mm::{address::AsPhysicalAddress, PageSlice};
 
 use crate::{command::IdentifyNamespaceRequest, IoDirection};
 
@@ -64,30 +64,34 @@ impl NgBlockDevice for NvmeDrive {
     async fn read(
         &self,
         lba: u64,
-        buffer: &mut PageBox<[MaybeUninit<u8>]>,
+        buffer: &mut PageSlice<MaybeUninit<u8>>,
     ) -> Result<(), NvmeError> {
+        debug_assert_eq!(buffer.len() % self.block_size(), 0);
+        let lba_count = buffer.len() / self.block_size();
+
         self.controller
             .perform_io(
                 self.nsid,
                 lba,
+                lba_count,
                 unsafe { buffer.as_physical_address() },
                 IoDirection::Read,
             )
-            .await?;
-
-        Ok(())
+            .await
     }
-    async fn write(&self, lba: u64, buffer: PageBox<[u8]>) -> Result<(), NvmeError> {
+    async fn write(&self, lba: u64, buffer: &PageSlice<u8>) -> Result<(), NvmeError> {
+        debug_assert_eq!(buffer.len() % self.block_size(), 0);
+        let lba_count = buffer.len() / self.block_size();
+
         self.controller
             .perform_io(
                 self.nsid,
                 lba,
+                lba_count,
                 unsafe { buffer.as_physical_address() },
                 IoDirection::Write,
             )
-            .await?;
-
-        Ok(())
+            .await
     }
 
     fn block_size(&self) -> usize {
diff --git a/kernel/driver/block/nvme/src/lib.rs b/kernel/driver/block/nvme/src/lib.rs
index 808a879f..360d5c2b 100644
--- a/kernel/driver/block/nvme/src/lib.rs
+++ b/kernel/driver/block/nvme/src/lib.rs
@@ -260,6 +260,7 @@ impl NvmeController {
         &'static self,
         nsid: u32,
         lba: u64,
+        lba_count: usize,
         buffer_address: PhysicalAddress,
         direction: IoDirection,
     ) -> Result<(), NvmeError> {
@@ -280,7 +281,7 @@ impl NvmeController {
                 IoRead {
                     nsid,
                     lba,
-                    count: 1,
+                    count: lba_count as _,
                 },
                 &[buffer_address],
                 true,
@@ -289,7 +290,7 @@ impl NvmeController {
                 IoWrite {
                     nsid,
                     lba,
-                    count: 1,
+                    count: lba_count as _,
                 },
                 &[buffer_address],
                 true,
diff --git a/kernel/driver/fs/ext2/src/dir.rs b/kernel/driver/fs/ext2/src/dir.rs
index 219c3e66..f1f9c10c 100644
--- a/kernel/driver/fs/ext2/src/dir.rs
+++ b/kernel/driver/fs/ext2/src/dir.rs
@@ -1,19 +1,13 @@
-use core::{any::Any, mem::MaybeUninit, ops::Range, str::FromStr};
+use core::{any::Any, mem::MaybeUninit, str::FromStr};
 
-use alloc::{sync::Arc, vec};
+use alloc::sync::Arc;
 use libk::{
     block,
     error::Error,
-    task::sync::{MappedAsyncMutexGuard, Mutex},
-    vfs::{
-        block::cache::{CachedBlock, CachedBlockRef},
-        CommonImpl, DirectoryImpl, DirectoryOpenPosition, InstanceData, Metadata, Node, NodeFlags,
-        NodeRef, RegularImpl,
-    },
+    vfs::{CommonImpl, DirectoryImpl, DirectoryOpenPosition, Metadata, Node, NodeFlags, NodeRef},
 };
-use libk_util::lru_hash_table::LruCache;
 use yggdrasil_abi::{
-    io::{DirectoryEntry, FileMode, FileType, GroupId, OpenOptions, UserId},
+    io::{DirectoryEntry, FileMode, FileType, GroupId, UserId},
     util::FixedString,
 };
 
@@ -22,16 +16,10 @@ use crate::{Dirent, Ext2Fs, Inode};
 pub struct DirectoryNode {
     fs: Arc<Ext2Fs>,
     inode: Inode,
+    #[allow(unused)]
     ino: u32,
 }
 
-struct DirentIterExt<'a> {
-    fs: &'a Ext2Fs,
-    inode: &'a Inode,
-    offset: usize,
-    current: Option<CachedBlockRef>,
-}
-
 struct DirentIter<'a> {
     fs: &'a Ext2Fs,
     block: &'a [u8],
@@ -161,7 +149,7 @@ impl DirectoryNode {
 }
 
 impl CommonImpl for DirectoryNode {
-    fn size(&self, node: &NodeRef) -> Result<u64, Error> {
+    fn size(&self, _node: &NodeRef) -> Result<u64, Error> {
         Ok(self.inode.size_lower as _)
     }
 
@@ -169,7 +157,7 @@ impl CommonImpl for DirectoryNode {
         self
     }
 
-    fn metadata(&self, node: &NodeRef) -> Result<Metadata, Error> {
+    fn metadata(&self, _node: &NodeRef) -> Result<Metadata, Error> {
         Ok(Metadata {
             uid: unsafe { UserId::from_raw(self.inode.uid as _) },
             gid: unsafe { GroupId::from_raw(self.inode.gid as _) },
diff --git a/kernel/driver/fs/ext2/src/file.rs b/kernel/driver/fs/ext2/src/file.rs
index c7bee810..daa9916b 100644
--- a/kernel/driver/fs/ext2/src/file.rs
+++ b/kernel/driver/fs/ext2/src/file.rs
@@ -1,10 +1,10 @@
 use core::any::Any;
 
-use alloc::{sync::Arc, vec};
+use alloc::sync::Arc;
 use libk::{
     block,
     error::Error,
-    vfs::{block, CommonImpl, InstanceData, Metadata, Node, NodeFlags, NodeRef, RegularImpl},
+    vfs::{CommonImpl, InstanceData, Metadata, Node, NodeFlags, NodeRef, RegularImpl},
 };
 use yggdrasil_abi::io::{FileMode, GroupId, OpenOptions, UserId};
 
@@ -13,6 +13,7 @@ use crate::{Ext2Fs, Inode};
 pub struct RegularNode {
     fs: Arc<Ext2Fs>,
     inode: Inode,
+    #[allow(unused)]
     ino: u32,
 }
 
@@ -53,7 +54,7 @@ impl RegularNode {
 }
 
 impl CommonImpl for RegularNode {
-    fn metadata(&self, node: &NodeRef) -> Result<Metadata, Error> {
+    fn metadata(&self, _node: &NodeRef) -> Result<Metadata, Error> {
         Ok(Metadata {
             uid: unsafe { UserId::from_raw(self.inode.uid as _) },
             gid: unsafe { GroupId::from_raw(self.inode.gid as _) },
@@ -65,7 +66,7 @@ impl CommonImpl for RegularNode {
         self
     }
 
-    fn size(&self, node: &NodeRef) -> Result<u64, Error> {
+    fn size(&self, _node: &NodeRef) -> Result<u64, Error> {
         Ok(self.size())
     }
 }
@@ -87,16 +88,16 @@ impl RegularImpl for RegularNode {
         Ok(())
     }
 
-    fn truncate(&self, node: &NodeRef, new_size: u64) -> Result<(), Error> {
+    fn truncate(&self, _node: &NodeRef, _new_size: u64) -> Result<(), Error> {
         Err(Error::ReadOnly)
     }
 
     fn write(
         &self,
-        node: &NodeRef,
-        instance: Option<&InstanceData>,
-        pos: u64,
-        buf: &[u8],
+        _node: &NodeRef,
+        _instance: Option<&InstanceData>,
+        _pos: u64,
+        _buf: &[u8],
     ) -> Result<usize, Error> {
         Err(Error::ReadOnly)
     }
diff --git a/kernel/driver/fs/ext2/src/lib.rs b/kernel/driver/fs/ext2/src/lib.rs
index 75b72c9e..720e6255 100644
--- a/kernel/driver/fs/ext2/src/lib.rs
+++ b/kernel/driver/fs/ext2/src/lib.rs
@@ -2,16 +2,9 @@
 
 extern crate alloc;
 
-use core::{
-    any::Any,
-    mem::MaybeUninit,
-    ops::{Deref, DerefMut},
-    pin::Pin,
-    sync::atomic::AtomicBool,
-    task::{Context, Poll},
-};
+use core::ops::{Deref, DerefMut};
 
-use alloc::{boxed::Box, sync::Arc, vec, vec::Vec};
+use alloc::{boxed::Box, sync::Arc, vec};
 use bytemuck::{Pod, Zeroable};
 use dir::DirectoryNode;
 use file::RegularNode;
@@ -19,13 +12,12 @@ use libk::{
     error::Error,
     vfs::{
         block::{
-            cache::{BlockCache, CachedBlock, CachedBlockRef},
+            cache::{BlockCache, CachedBlockRef},
             BlockDevice,
         },
-        CommonImpl, Metadata, NodeRef,
+        NodeRef,
     },
 };
-use libk_mm::PageBox;
 use libk_util::OneTimeInit;
 use static_assertions::const_assert_eq;
 
@@ -237,6 +229,7 @@ impl Ext2Fs {
             return Err(Error::InvalidArgument);
         }
 
+        let bgdt_offset = 1;
         let block_size = 1024usize << superblock.block_size_log2;
 
         let bgdt_entry_count = ((superblock.total_blocks + superblock.block_group_block_count - 1)
@@ -261,7 +254,7 @@ impl Ext2Fs {
             bgdt_entry_count,
         );
         for i in 0..bgdt_block_count {
-            let disk_offset = (i as u64 + 1) * block_size as u64;
+            let disk_offset = (i as u64 + bgdt_offset) * block_size as u64;
             device
                 .read_exact(
                     disk_offset,
diff --git a/kernel/libk/libk-mm/src/lib.rs b/kernel/libk/libk-mm/src/lib.rs
index 4668fcf6..5de4e74f 100644
--- a/kernel/libk/libk-mm/src/lib.rs
+++ b/kernel/libk/libk-mm/src/lib.rs
@@ -15,6 +15,7 @@ use core::{
     fmt,
     mem::{size_of, MaybeUninit},
     ops::{Deref, DerefMut},
+    slice::SliceIndex,
 };
 
 use address::Virtualize;
@@ -67,6 +68,10 @@ pub struct PageBox<T: ?Sized> {
     page_count: usize,
 }
 
+pub struct PageSlice<T> {
+    data: [T],
+}
+
 impl<T> PageBox<T> {
     #[inline]
     fn alloc_slice(count: usize, zeroed: bool) -> Result<(PhysicalAddress, usize), Error> {
@@ -197,6 +202,14 @@ impl<T> PageBox<[T]> {
         let slice = unsafe { slice.assume_init_slice() };
         Ok(slice)
     }
+
+    pub fn as_slice(&self) -> &PageSlice<T> {
+        todo!()
+    }
+
+    pub fn as_slice_mut(&mut self) -> &mut PageSlice<T> {
+        unsafe { core::mem::transmute(&mut self[..]) }
+    }
 }
 
 impl<T> PageBox<MaybeUninit<T>> {
@@ -331,3 +344,36 @@ impl<T: ?Sized + fmt::Display> fmt::Display for PageBox<T> {
 
 unsafe impl<T: ?Sized + Send> Send for PageBox<T> {}
 unsafe impl<T: ?Sized + Sync> Sync for PageBox<T> {}
+
+impl<T> PageSlice<T> {
+    pub fn subslice_mut<R: SliceIndex<[T], Output = [T]>>(
+        &mut self,
+        index: R,
+    ) -> &mut PageSlice<T> {
+        unsafe { core::mem::transmute(&mut self.data[index]) }
+    }
+
+    pub fn subslice<R: SliceIndex<[T], Output = [T]>>(&self, index: R) -> &PageSlice<T> {
+        unsafe { core::mem::transmute(&self.data[index]) }
+    }
+}
+
+impl<T> AsPhysicalAddress for PageSlice<T> {
+    unsafe fn as_physical_address(&self) -> PhysicalAddress {
+        PhysicalAddress::from_virtualized(self.data.as_ptr().addr())
+    }
+}
+
+impl<T> Deref for PageSlice<T> {
+    type Target = [T];
+
+    fn deref(&self) -> &Self::Target {
+        &self.data
+    }
+}
+
+impl<T> DerefMut for PageSlice<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.data
+    }
+}
diff --git a/kernel/libk/libk-util/src/lib.rs b/kernel/libk/libk-util/src/lib.rs
index c7fe4137..d534f53c 100644
--- a/kernel/libk/libk-util/src/lib.rs
+++ b/kernel/libk/libk-util/src/lib.rs
@@ -16,7 +16,7 @@ extern crate alloc;
 
 use core::{
     mem::MaybeUninit,
-    ops::{Deref, DerefMut},
+    ops::{Add, Deref, DerefMut, Mul, Range, Sub},
     panic,
 };
 
@@ -37,6 +37,57 @@ pub trait IsTrue {}
 
 impl IsTrue for ConstAssert<true> {}
 
+pub trait RangeExt<N> {
+    fn chunked_range(self, step: N) -> ChunkedRangeIter<N>;
+    fn mul(self, mul: N) -> Self;
+    fn add(self, add: N) -> Self;
+}
+
+pub struct ChunkedRangeIter<N> {
+    pos: N,
+    end: N,
+    step: N,
+}
+
+impl<N> ChunkedRangeIter<N> {
+    pub fn from_range(range: Range<N>, step: N) -> Self {
+        Self {
+            pos: range.start,
+            end: range.end,
+            step,
+        }
+    }
+}
+
+impl<N: Copy + Ord + Add<Output = N> + Sub<Output = N>> Iterator for ChunkedRangeIter<N> {
+    type Item = Range<N>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.pos >= self.end {
+            return None;
+        }
+        let step = core::cmp::min(self.end - self.pos, self.step);
+        let range = self.pos..self.pos + step;
+        self.pos = self.pos + step;
+
+        Some(range)
+    }
+}
+
+impl<N: Mul<Output = N> + Add<Output = N> + Copy> RangeExt<N> for Range<N> {
+    fn chunked_range(self, step: N) -> ChunkedRangeIter<N> {
+        ChunkedRangeIter::from_range(self, step)
+    }
+
+    fn mul(self, mul: N) -> Self {
+        self.start * mul..self.end * mul
+    }
+
+    fn add(self, add: N) -> Self {
+        self.start + add..self.end + add
+    }
+}
+
 /// Statically-allocated "dynamic" vector
 pub struct StaticVector<T, const N: usize> {
     data: [MaybeUninit<T>; N],
@@ -105,3 +156,25 @@ impl<T, const N: usize> DerefMut for StaticVector<T, N> {
         unsafe { MaybeUninit::slice_assume_init_mut(&mut self.data[..self.len]) }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::RangeExt;
+
+    #[test]
+    fn chunked_range() {
+        let r = 0..30;
+        let mut i = 0;
+        for range in r.chunked_range(10) {
+            assert_eq!(range, i..i + 10);
+            i += 10;
+        }
+
+        let r = 0..28;
+        let mut it = r.chunked_range(10);
+        assert_eq!(it.next(), Some(0..10));
+        assert_eq!(it.next(), Some(10..20));
+        assert_eq!(it.next(), Some(20..28));
+        assert_eq!(it.next(), None);
+    }
+}
diff --git a/kernel/libk/libk-util/src/lru_hash_table.rs b/kernel/libk/libk-util/src/lru_hash_table.rs
index 8d9e622a..1432ca85 100644
--- a/kernel/libk/libk-util/src/lru_hash_table.rs
+++ b/kernel/libk/libk-util/src/lru_hash_table.rs
@@ -3,10 +3,7 @@ use core::{
     ops::AsyncFnOnce,
 };
 
-use alloc::{
-    collections::{linked_list, LinkedList},
-    vec::Vec,
-};
+use alloc::{collections::LinkedList, vec::Vec};
 
 use crate::hash_table::DefaultHashBuilder;
 
diff --git a/kernel/libk/src/task/sync.rs b/kernel/libk/src/task/sync.rs
index 69811381..c0d1188a 100644
--- a/kernel/libk/src/task/sync.rs
+++ b/kernel/libk/src/task/sync.rs
@@ -1,7 +1,6 @@
 use core::{
     cell::UnsafeCell,
     future::poll_fn,
-    marker::PhantomData,
     ops::{AsyncFnOnce, Deref, DerefMut},
     sync::atomic::{AtomicBool, AtomicU32, Ordering},
     task::{Context, Poll},
diff --git a/kernel/libk/src/vfs/block/cache.rs b/kernel/libk/src/vfs/block/cache.rs
index cf8323bd..5e33a3d9 100644
--- a/kernel/libk/src/vfs/block/cache.rs
+++ b/kernel/libk/src/vfs/block/cache.rs
@@ -18,7 +18,7 @@ pub struct CachedBlock {
 }
 
 pub struct CachedBlockRef {
-    entry: Arc<IrqSafeRwLock<CachedBlock>>,
+    _entry: Arc<IrqSafeRwLock<CachedBlock>>,
     lock: IrqSafeRwLockReadGuard<'static, CachedBlock>,
 }
 
@@ -48,12 +48,28 @@ impl BlockCache {
         let read = block.read();
         if read.dirty {
             log::info!("Evict block {}", address);
-            if let Err(err) = self.device.write_exact(address, &read.data).await {
+            if let Err(err) = self
+                .device
+                .write_aligned(address, read.data.as_slice())
+                .await
+            {
                 log::error!("Disk error: flushing block {}: {:?}", address, err);
             }
         }
     }
 
+    async fn fetch_block(&self, address: u64) -> Result<Arc<IrqSafeRwLock<CachedBlock>>, Error> {
+        let mut data = PageBox::new_uninit_slice(self.block_size)?;
+        self.device
+            .read_aligned(address, data.as_slice_mut())
+            .await?;
+        let data = unsafe { data.assume_init_slice() };
+        Ok(Arc::new(IrqSafeRwLock::new(CachedBlock {
+            data,
+            dirty: false,
+        })))
+    }
+
     pub async fn get<'a>(&'a self, address: u64) -> Result<CachedBlockRef, Error> {
         debug_assert_eq!(address % self.block_size as u64, 0);
 
@@ -63,14 +79,7 @@ impl BlockCache {
             .await
             .try_map_guard_async::<_, Error, _>(|cache: &'a mut LruCache<_, _>| async move {
                 let (value, evicted) = cache
-                    .try_get_or_insert_with_async(address, || async move {
-                        let mut data = PageBox::new_slice(0, self.block_size)?;
-                        self.device.read_exact(address, &mut data).await?;
-                        Ok(Arc::new(IrqSafeRwLock::new(CachedBlock {
-                            data,
-                            dirty: false,
-                        })))
-                    })
+                    .try_get_or_insert_with_async(address, || self.fetch_block(address))
                     .await?;
 
                 if evicted.is_some() {
@@ -102,7 +111,10 @@ impl CachedBlockRef {
         let entry = entry.clone();
         // Safety: ok, Arc instance is still held
         let lock = unsafe { core::mem::transmute(entry.read()) };
-        Self { lock, entry }
+        Self {
+            lock,
+            _entry: entry,
+        }
     }
 }
 
diff --git a/kernel/libk/src/vfs/block/device.rs b/kernel/libk/src/vfs/block/device.rs
index 8dabd5ea..bf5bdb35 100644
--- a/kernel/libk/src/vfs/block/device.rs
+++ b/kernel/libk/src/vfs/block/device.rs
@@ -11,13 +11,16 @@ use core::{
 use alloc::boxed::Box;
 use async_trait::async_trait;
 use futures_util::{task::AtomicWaker, Future};
-use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageBox, PageProvider};
-use libk_util::waker::QueueWaker;
+use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageBox, PageProvider, PageSlice};
+use libk_util::{waker::QueueWaker, RangeExt};
 use yggdrasil_abi::{error::Error, io::DeviceRequest};
 
-use crate::vfs::block::{
-    request::{IoOperation, IoRequest, IoSubmissionId},
-    BlockDevice,
+use crate::{
+    task::debug,
+    vfs::block::{
+        request::{IoOperation, IoRequest, IoSubmissionId},
+        BlockDevice,
+    },
 };
 
 #[async_trait]
@@ -27,9 +30,9 @@ pub trait NgBlockDevice: Sync {
     async fn read(
         &self,
         lba: u64,
-        buffer: &mut PageBox<[MaybeUninit<u8>]>,
+        buffer: &mut PageSlice<MaybeUninit<u8>>,
     ) -> Result<(), Self::Error>;
-    async fn write(&self, lba: u64, buffer: PageBox<[u8]>) -> Result<(), Self::Error>;
+    async fn write(&self, lba: u64, buffer: &PageSlice<u8>) -> Result<(), Self::Error>;
 
     fn block_size(&self) -> usize;
     fn block_count(&self) -> usize;
@@ -75,6 +78,52 @@ impl<'a, D: NgBlockDevice + 'a> NgBlockDeviceWrapper<'a, D> {
 
 #[async_trait]
 impl<'a, D: NgBlockDevice + 'a> BlockDevice for NgBlockDeviceWrapper<'a, D> {
+    async fn read_aligned(
+        &self,
+        pos: u64,
+        buf: &mut PageSlice<MaybeUninit<u8>>,
+    ) -> Result<(), Error> {
+        if pos % self.block_size as u64 != 0 || buf.len() % self.block_size != 0 {
+            // TODO fallback to unaligned read
+            todo!()
+        }
+
+        let range = 0..buf.len() / self.block_size;
+
+        for chunk in range.chunked_range(self.max_blocks_per_request) {
+            let lba = chunk.start as u64 + pos / self.block_size as u64;
+            let buffer_range = chunk.mul(self.block_size);
+
+            self.device
+                .read(lba, buf.subslice_mut(buffer_range))
+                .await
+                .map_err(Self::handle_drive_error)?;
+        }
+
+        Ok(())
+    }
+
+    async fn write_aligned(&self, pos: u64, buf: &PageSlice<u8>) -> Result<(), Error> {
+        if pos % self.block_size as u64 != 0 || buf.len() % self.block_size != 0 {
+            // TODO fallback to unaligned write
+            todo!()
+        }
+
+        let range = 0..buf.len() / self.block_size;
+
+        for chunk in range.chunked_range(self.max_blocks_per_request) {
+            let lba = chunk.start as u64 + pos / self.block_size as u64;
+            let buffer_range = chunk.mul(self.block_size);
+
+            self.device
+                .write(lba, buf.subslice(buffer_range))
+                .await
+                .map_err(Self::handle_drive_error)?;
+        }
+
+        Ok(())
+    }
+
     async fn read(&self, mut pos: u64, mut buf: &mut [u8]) -> Result<usize, Error> {
         let len = buf.len();
         let mut remaining = buf.len();
@@ -88,7 +137,7 @@ impl<'a, D: NgBlockDevice + 'a> BlockDevice for NgBlockDeviceWrapper<'a, D> {
             let amount = core::cmp::min(self.block_size - block_offset, buf.len());
 
             self.device
-                .read(lba, &mut block)
+                .read(lba, block.as_slice_mut())
                 .await
                 .map_err(Self::handle_drive_error)?;
 
diff --git a/kernel/libk/src/vfs/block/mod.rs b/kernel/libk/src/vfs/block/mod.rs
index c7f62356..ed59ba81 100644
--- a/kernel/libk/src/vfs/block/mod.rs
+++ b/kernel/libk/src/vfs/block/mod.rs
@@ -1,8 +1,10 @@
 #![allow(missing_docs)]
 
+use core::mem::MaybeUninit;
+
 use alloc::{boxed::Box, vec::Vec};
 use async_trait::async_trait;
-use libk_mm::{PageBox, PageProvider};
+use libk_mm::{PageProvider, PageSlice};
 use partition::Partition;
 use yggdrasil_abi::{error::Error, io::DeviceRequest};
 
@@ -68,6 +70,18 @@ pub fn probe_partitions<
 #[allow(unused)]
 #[async_trait]
 pub trait BlockDevice: PageProvider + Sync {
+    async fn read_aligned(
+        &self,
+        pos: u64,
+        buf: &mut PageSlice<MaybeUninit<u8>>,
+    ) -> Result<(), Error> {
+        Err(Error::NotImplemented)
+    }
+
+    async fn write_aligned(&self, pos: u64, buf: &PageSlice<u8>) -> Result<(), Error> {
+        Err(Error::NotImplemented)
+    }
+
     async fn read(&self, pos: u64, buf: &mut [u8]) -> Result<usize, Error> {
         Err(Error::NotImplemented)
     }
diff --git a/kernel/libk/src/vfs/block/partition.rs b/kernel/libk/src/vfs/block/partition.rs
index 684db0d0..e49e4add 100644
--- a/kernel/libk/src/vfs/block/partition.rs
+++ b/kernel/libk/src/vfs/block/partition.rs
@@ -3,7 +3,7 @@ use core::mem::{size_of, MaybeUninit};
 use alloc::{boxed::Box, vec, vec::Vec};
 use async_trait::async_trait;
 use bytemuck::{Pod, Zeroable};
-use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageBox, PageProvider};
+use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageBox, PageProvider, PageSlice};
 use static_assertions::const_assert_eq;
 use uuid::Uuid;
 use yggdrasil_abi::{error::Error, io::DeviceRequest};
@@ -59,26 +59,56 @@ impl<'a, D: NgBlockDevice + 'a> Partition<'a, D> {
 }
 
 impl<'a, D: NgBlockDevice + 'a> PageProvider for Partition<'a, D> {
-    fn get_page(&self, offset: u64) -> Result<PhysicalAddress, Error> {
+    fn get_page(&self, _offset: u64) -> Result<PhysicalAddress, Error> {
         todo!()
     }
 
     fn clone_page(
         &self,
-        offset: u64,
-        src_phys: PhysicalAddress,
-        src_attrs: MapAttributes,
+        _offset: u64,
+        _src_phys: PhysicalAddress,
+        _src_attrs: MapAttributes,
     ) -> Result<PhysicalAddress, Error> {
         todo!()
     }
 
-    fn release_page(&self, offset: u64, phys: PhysicalAddress) -> Result<(), Error> {
+    fn release_page(&self, _offset: u64, _phys: PhysicalAddress) -> Result<(), Error> {
         todo!()
     }
 }
 
 #[async_trait]
 impl<'a, D: NgBlockDevice + 'a> BlockDevice for Partition<'a, D> {
+    async fn read_aligned(
+        &self,
+        pos: u64,
+        buf: &mut PageSlice<MaybeUninit<u8>>,
+    ) -> Result<(), Error> {
+        let start = self.start_byte() + pos;
+        if start % self.device.block_size as u64 != 0
+            || start + buf.len() as u64 >= self.end_byte()
+            || buf.len() % self.device.block_size != 0
+        {
+            // TODO fallback to unaligned read
+            todo!()
+        }
+
+        self.device.read_aligned(start, buf).await
+    }
+
+    async fn write_aligned(&self, pos: u64, buf: &PageSlice<u8>) -> Result<(), Error> {
+        let start = self.start_byte() + pos;
+        if start % self.device.block_size as u64 != 0
+            || start + buf.len() as u64 >= self.end_byte()
+            || buf.len() % self.device.block_size != 0
+        {
+            // TODO fallback to unaligned write
+            todo!()
+        }
+
+        self.device.write_aligned(start, buf).await
+    }
+
     async fn read(&self, pos: u64, buf: &mut [u8]) -> Result<usize, Error> {
         if pos >= self.end_byte() {
             return Ok(0);
diff --git a/xtask/src/build/mod.rs b/xtask/src/build/mod.rs
index 649106f8..28b08018 100644
--- a/xtask/src/build/mod.rs
+++ b/xtask/src/build/mod.rs
@@ -116,7 +116,12 @@ pub fn check_all(env: BuildEnv, action: CheckAction) -> Result<(), Error> {
 }
 
 pub fn test_all(env: BuildEnv) -> Result<(), Error> {
-    for path in ["kernel/driver/fs/memfs", "lib/abi", "kernel/libk"] {
+    for path in [
+        "kernel/driver/fs/memfs",
+        "lib/abi",
+        "kernel/libk",
+        "kernel/libk-util",
+    ] {
         CargoBuilder::Host(env.verbose).run(env.workspace_root.join(path), "test")?;
     }
     Ok(())