From 371e8c0326bdfb918726bb22791f7fa640bcb4a0 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 22:15:56 +0100
Subject: [PATCH] storage/virtio-blkd: Use DiskWrapper

---
 Cargo.lock                        |  1 -
 storage/virtio-blkd/Cargo.toml    |  1 -
 storage/virtio-blkd/src/scheme.rs | 17 ++++++-----------
 3 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0bc2322768..537432f9cd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1672,7 +1672,6 @@ dependencies = [
  "futures",
  "libredox",
  "log",
- "partitionlib",
  "pcid",
  "redox-daemon",
  "redox-scheme 0.3.0",
diff --git a/storage/virtio-blkd/Cargo.toml b/storage/virtio-blkd/Cargo.toml
index de7aba04a1..2b2a663565 100644
--- a/storage/virtio-blkd/Cargo.toml
+++ b/storage/virtio-blkd/Cargo.toml
@@ -15,7 +15,6 @@ spin = "*"
 redox-daemon = "0.1"
 redox_syscall = { version = "0.5", features = ["std"] }
 redox-scheme = { git = "https://gitlab.redox-os.org/redox-os/redox-scheme.git" }
-partitionlib = { path = "../partitionlib" }
 
 common = { path = "../../common" }
 driver-block = { path = "../driver-block" }
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index 710a43ba4d..e498798556 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -5,7 +5,6 @@ use std::sync::Arc;
 
 use common::dma::Dma;
 use driver_block::{Disk, DiskWrapper};
-use partitionlib::PartitionTable;
 
 use redox_scheme::CallerCtx;
 use redox_scheme::OpenResult;
@@ -134,19 +133,15 @@ pub enum Handle {
 }
 
 pub struct DiskScheme<'a> {
-    disk: VirtioDisk<'a>,
+    disk: DiskWrapper<VirtioDisk<'a>>,
     next_id: usize,
     handles: BTreeMap<usize, Handle>,
-    part_table: Option<PartitionTable>,
 }
 
 impl<'a> DiskScheme<'a> {
     pub fn new(queue: Arc<Queue<'a>>, cfg: BlockDeviceConfig) -> Self {
-        let mut disk = VirtioDisk::new(queue, cfg);
-
         Self {
-            part_table: DiskWrapper::pt(&mut disk),
-            disk,
+            disk: DiskWrapper::new(VirtioDisk::new(queue, cfg)),
             next_id: 0,
             handles: BTreeMap::new(),
         }
@@ -170,7 +165,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
                 //            to the namespace id).
                 write!(list, "{}\n", 0).unwrap();
 
-                if let Some(part_table) = &self.part_table {
+                if let Some(part_table) = &self.disk.pt {
                     for part_num in 0..part_table.partitions.len() {
                         write!(list, "{}p{}\n", 0, part_num).unwrap();
                     }
@@ -201,7 +196,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
             let part_num_str = &path_str[p_pos + 1..];
             let part_num = part_num_str.parse::<u32>().unwrap();
 
-            let part_table = self.part_table.as_ref().unwrap();
+            let part_table = self.disk.pt.as_ref().unwrap();
             let _part = part_table.partitions.get(part_num as usize).unwrap();
 
             let id = self.next_id;
@@ -246,7 +241,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
             }
 
             Handle::Partition { number } => {
-                let part_table = self.part_table.as_ref().unwrap();
+                let part_table = self.disk.pt.as_ref().unwrap();
                 let part = part_table
                     .partitions
                     .get(number as usize)
@@ -298,7 +293,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
                 }
 
                 Handle::Partition { number } => {
-                    let part_table = self.part_table.as_ref().unwrap();
+                    let part_table = self.disk.pt.as_ref().unwrap();
                     let part = part_table
                         .partitions
                         .get(number as usize)