remove branch reader / writer

crates/gitbutler-core/src/virtual_branches.rs

@@ -16,9 +16,6 @@ pub use base::*;
 pub mod controller;
 pub use controller::Controller;
 
-mod iterator;
-pub use iterator::BranchIterator as Iterator;
-
 mod r#virtual;
 pub use r#virtual::*;
 

crates/gitbutler-core/src/virtual_branches/branch.rs

@@ -1,16 +1,12 @@
 mod file_ownership;
 mod hunk;
 mod ownership;
-mod reader;
-mod writer;
 
 use anyhow::Result;
 pub use file_ownership::OwnershipClaim;
 pub use hunk::Hunk;
 pub use ownership::{reconcile_claims, BranchOwnershipClaims};
-pub use reader::BranchReader as Reader;
 use serde::{Deserialize, Serialize};
-pub use writer::BranchWriter as Writer;
 
 use crate::{git, id::Id};
 

crates/gitbutler-core/src/virtual_branches/branch/reader.rs

@@ -1,32 +0,0 @@
-use super::{Branch, BranchId};
-use crate::{reader, sessions, virtual_branches::VirtualBranchesHandle};
-
-pub struct BranchReader<'r> {
-    reader: &'r reader::Reader<'r>,
-    state_handle: VirtualBranchesHandle,
-    use_state_handle: bool,
-}
-
-impl<'r> BranchReader<'r> {
-    pub fn new(
-        reader: &'r sessions::Reader<'r>,
-        state_handle: VirtualBranchesHandle,
-        use_state_handle: bool,
-    ) -> Self {
-        Self {
-            reader: reader.reader(),
-            state_handle,
-            use_state_handle,
-        }
-    }
-
-    pub fn read(&self, id: &BranchId) -> Result<Branch, reader::Error> {
-        if self.use_state_handle && self.state_handle.file_exists() {
-            self.state_handle
-                .get_branch(id)
-                .map_err(|_| reader::Error::NotFound)
-        } else {
-            Branch::from_reader(&self.reader.sub(format!("branches/{}", id)))
-        }
-    }
-}

crates/gitbutler-core/src/virtual_branches/branch/writer.rs

@@ -1,156 +0,0 @@
-use anyhow::Result;
-
-use super::Branch;
-use crate::{gb_repository, reader, virtual_branches::state::VirtualBranchesHandle, writer};
-
-pub struct BranchWriter<'writer> {
-    repository: &'writer gb_repository::Repository,
-    writer: writer::DirWriter,
-    reader: reader::Reader<'writer>,
-    state_handle: VirtualBranchesHandle,
-}
-
-impl<'writer> BranchWriter<'writer> {
-    pub fn new(
-        repository: &'writer gb_repository::Repository,
-        state_handle: VirtualBranchesHandle,
-    ) -> Result<Self, std::io::Error> {
-        let reader = reader::Reader::open(repository.root())?;
-        let writer = writer::DirWriter::open(repository.root())?;
-        Ok(Self {
-            repository,
-            writer,
-            reader,
-            state_handle,
-        })
-    }
-
-    pub fn delete(&self, branch: &Branch) -> Result<()> {
-        match self
-            .reader
-            .sub(format!("branches/{}", branch.id))
-            .read("id")
-        {
-            Ok(_) => {
-                self.repository.mark_active_session()?;
-                let _lock = self.repository.lock();
-                self.writer.remove(format!("branches/{}", branch.id))?;
-                // Write in the state file as well
-                let _ = self.state_handle.remove_branch(branch.id);
-                Ok(())
-            }
-            Err(reader::Error::NotFound) => Ok(()),
-            Err(err) => Err(err.into()),
-        }
-    }
-
-    pub fn write(&self, branch: &mut Branch) -> Result<()> {
-        let reader = self.reader.sub(format!("branches/{}", branch.id));
-        match Branch::from_reader(&reader) {
-            Ok(existing) if existing.eq(branch) => return Ok(()),
-            Ok(_) | Err(reader::Error::NotFound) => {}
-            Err(err) => return Err(err.into()),
-        }
-
-        self.repository.mark_active_session()?;
-
-        branch.updated_timestamp_ms = std::time::SystemTime::now()
-            .duration_since(std::time::UNIX_EPOCH)?
-            .as_millis();
-
-        let mut batch = vec![];
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/id", branch.id),
-            branch.id.to_string(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/name", branch.id),
-            branch.name.clone(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/notes", branch.id),
-            branch.notes.clone(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/order", branch.id),
-            branch.order.to_string(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/applied", branch.id),
-            branch.applied.to_string(),
-        ));
-
-        if let Some(upstream) = &branch.upstream {
-            batch.push(writer::BatchTask::Write(
-                format!("branches/{}/meta/upstream", branch.id),
-                upstream.to_string(),
-            ));
-        } else {
-            batch.push(writer::BatchTask::Remove(format!(
-                "branches/{}/meta/upstream",
-                branch.id
-            )));
-        }
-
-        if let Some(upstream_head) = &branch.upstream_head {
-            batch.push(writer::BatchTask::Write(
-                format!("branches/{}/meta/upstream_head", branch.id),
-                upstream_head.to_string(),
-            ));
-        } else {
-            batch.push(writer::BatchTask::Remove(format!(
-                "branches/{}/meta/upstream_head",
-                branch.id
-            )));
-        }
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/tree", branch.id),
-            branch.tree.to_string(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/head", branch.id),
-            branch.head.to_string(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/created_timestamp_ms", branch.id),
-            branch.created_timestamp_ms.to_string(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/updated_timestamp_ms", branch.id),
-            branch.updated_timestamp_ms.to_string(),
-        ));
-
-        batch.push(writer::BatchTask::Write(
-            format!("branches/{}/meta/ownership", branch.id),
-            branch.ownership.to_string(),
-        ));
-
-        if let Some(selected_for_changes) = branch.selected_for_changes {
-            batch.push(writer::BatchTask::Write(
-                format!("branches/{}/meta/selected_for_changes", branch.id),
-                selected_for_changes.to_string(),
-            ));
-        } else {
-            batch.push(writer::BatchTask::Remove(format!(
-                "branches/{}/meta/selected_for_changes",
-                branch.id
-            )));
-        }
-
-        self.writer.batch(&batch)?;
-
-        // Write in the state file as well
-        self.state_handle.set_branch(branch.clone())?;
-
-        Ok(())
-    }
-}

crates/gitbutler-core/src/virtual_branches/iterator.rs

@@ -1,66 +0,0 @@
-use std::collections::HashSet;
-
-use anyhow::Result;
-
-use super::{
-    branch::{self, BranchId},
-    VirtualBranchesHandle,
-};
-use crate::sessions;
-
-pub struct BranchIterator<'i> {
-    branch_reader: branch::Reader<'i>,
-    ids: Vec<BranchId>,
-}
-
-impl<'i> BranchIterator<'i> {
-    pub fn new(
-        session_reader: &'i sessions::Reader<'i>,
-        state_handle: VirtualBranchesHandle,
-        use_state_handle: bool,
-    ) -> Result<Self> {
-        let unique_ids: HashSet<String> = if use_state_handle && state_handle.file_exists() {
-            state_handle.list_branch_ids()?
-        } else {
-            session_reader
-                .reader()
-                .list_files("branches")?
-                .into_iter()
-                .map(|file_path| {
-                    file_path
-                        .iter()
-                        .next()
-                        .unwrap()
-                        .to_string_lossy()
-                        .to_string()
-                })
-                .filter(|file_path| file_path != "selected")
-                .filter(|file_path| file_path != "target")
-                .collect()
-        };
-        let mut ids: Vec<BranchId> = unique_ids
-            .into_iter()
-            .map(|id| id.parse())
-            .filter_map(Result::ok)
-            .collect();
-        ids.sort();
-        Ok(Self {
-            branch_reader: branch::Reader::new(session_reader, state_handle, use_state_handle),
-            ids,
-        })
-    }
-}
-
-impl Iterator for BranchIterator<'_> {
-    type Item = Result<branch::Branch, crate::reader::Error>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.ids.is_empty() {
-            return None;
-        }
-
-        let id = self.ids.remove(0);
-        let branch = self.branch_reader.read(&id);
-        Some(branch)
-    }
-}

crates/gitbutler-core/tests/virtual_branches/branch/mod.rs

@@ -1,7 +1,3 @@
-use gitbutler_core::virtual_branches::Branch;
-
 mod file_ownership;
 mod hunk;
 mod ownership;
-mod reader;
-mod writer;

crates/gitbutler-core/tests/virtual_branches/branch/reader.rs

@@ -1,112 +0,0 @@
-use std::sync::atomic::{AtomicUsize, Ordering};
-
-use anyhow::Result;
-use gitbutler_core::virtual_branches::{
-    branch::{self, BranchOwnershipClaims},
-    Branch, BranchId, VirtualBranchesHandle,
-};
-use once_cell::sync::Lazy;
-
-use gitbutler_testsupport::{Case, Suite};
-
-static TEST_INDEX: Lazy<AtomicUsize> = Lazy::new(|| AtomicUsize::new(0));
-
-fn test_branch() -> Branch {
-    TEST_INDEX.fetch_add(1, Ordering::Relaxed);
-
-    Branch {
-        id: BranchId::generate(),
-        name: format!("branch_name_{}", TEST_INDEX.load(Ordering::Relaxed)),
-        notes: String::new(),
-        applied: true,
-        order: TEST_INDEX.load(Ordering::Relaxed),
-        upstream: Some(
-            format!(
-                "refs/remotes/origin/upstream_{}",
-                TEST_INDEX.load(Ordering::Relaxed)
-            )
-            .parse()
-            .unwrap(),
-        ),
-        upstream_head: Some(
-            format!(
-                "0123456789abcdef0123456789abcdef0123456{}",
-                TEST_INDEX.load(Ordering::Relaxed)
-            )
-            .parse()
-            .unwrap(),
-        ),
-        created_timestamp_ms: TEST_INDEX.load(Ordering::Relaxed) as u128,
-        updated_timestamp_ms: (TEST_INDEX.load(Ordering::Relaxed) + 100) as u128,
-        head: format!(
-            "0123456789abcdef0123456789abcdef0123456{}",
-            TEST_INDEX.load(Ordering::Relaxed)
-        )
-        .parse()
-        .unwrap(),
-        tree: format!(
-            "0123456789abcdef0123456789abcdef012345{}",
-            TEST_INDEX.load(Ordering::Relaxed) + 10
-        )
-        .parse()
-        .unwrap(),
-        ownership: BranchOwnershipClaims {
-            claims: vec![format!("file/{}:1-2", TEST_INDEX.load(Ordering::Relaxed))
-                .parse()
-                .unwrap()],
-        },
-        selected_for_changes: Some(1),
-    }
-}
-
-#[test]
-fn read_not_found() -> Result<()> {
-    let suite = Suite::default();
-    let Case {
-        gb_repository,
-        project,
-        ..
-    } = &suite.new_case();
-
-    let session = gb_repository.get_or_create_current_session()?;
-    let session_reader = gitbutler_core::sessions::Reader::open(gb_repository, &session)?;
-
-    let reader = branch::Reader::new(
-        &session_reader,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-        project.use_toml_vbranches_state(),
-    );
-    let result = reader.read(&BranchId::generate());
-    assert!(result.is_err());
-    assert_eq!(result.unwrap_err().to_string(), "file not found");
-
-    Ok(())
-}
-
-#[test]
-fn read_override() -> Result<()> {
-    let suite = Suite::default();
-    let Case {
-        gb_repository,
-        project,
-        ..
-    } = &suite.new_case();
-
-    let mut branch = test_branch();
-
-    let writer = branch::Writer::new(gb_repository, VirtualBranchesHandle::new(&project.gb_dir()))?;
-    writer.write(&mut branch)?;
-
-    let session = gb_repository.get_current_session()?.unwrap();
-    let session_reader = gitbutler_core::sessions::Reader::open(gb_repository, &session)?;
-
-    let reader = branch::Reader::new(
-        &session_reader,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-        project.use_toml_vbranches_state(),
-    );
-
-    assert_eq!(branch, reader.read(&branch.id).unwrap());
-
-    Ok(())
-}

crates/gitbutler-core/tests/virtual_branches/branch/writer.rs

@@ -1,218 +0,0 @@
-use std::{
-    fs,
-    sync::atomic::{AtomicUsize, Ordering},
-};
-
-use anyhow::Context;
-use gitbutler_core::virtual_branches::{branch, VirtualBranchesHandle};
-use once_cell::sync::Lazy;
-
-use self::branch::BranchId;
-use super::*;
-use gitbutler_testsupport::{Case, Suite};
-
-static TEST_INDEX: Lazy<AtomicUsize> = Lazy::new(|| AtomicUsize::new(0));
-
-fn new_test_branch() -> Branch {
-    TEST_INDEX.fetch_add(1, Ordering::Relaxed);
-
-    Branch {
-        id: BranchId::generate(),
-        name: format!("branch_name_{}", TEST_INDEX.load(Ordering::Relaxed)),
-        notes: String::new(),
-        applied: true,
-        upstream: Some(
-            format!(
-                "refs/remotes/origin/upstream_{}",
-                TEST_INDEX.load(Ordering::Relaxed)
-            )
-            .parse()
-            .unwrap(),
-        ),
-        upstream_head: None,
-        created_timestamp_ms: TEST_INDEX.load(Ordering::Relaxed) as u128,
-        updated_timestamp_ms: (TEST_INDEX.load(Ordering::Relaxed) + 100) as u128,
-        head: format!(
-            "0123456789abcdef0123456789abcdef0123456{}",
-            TEST_INDEX.load(Ordering::Relaxed)
-        )
-        .parse()
-        .unwrap(),
-        tree: format!(
-            "0123456789abcdef0123456789abcdef012345{}",
-            TEST_INDEX.load(Ordering::Relaxed) + 10
-        )
-        .parse()
-        .unwrap(),
-        ownership: gitbutler_core::virtual_branches::branch::BranchOwnershipClaims {
-            claims: vec![gitbutler_core::virtual_branches::branch::OwnershipClaim {
-                file_path: format!("file/{}:1-2", TEST_INDEX.load(Ordering::Relaxed)).into(),
-                hunks: vec![],
-            }],
-        },
-        order: TEST_INDEX.load(Ordering::Relaxed),
-        selected_for_changes: Some(1),
-    }
-}
-
-#[test]
-fn write_branch() -> anyhow::Result<()> {
-    let suite = Suite::default();
-    let Case {
-        gb_repository,
-        project,
-        ..
-    } = &suite.new_case();
-
-    let mut branch = new_test_branch();
-
-    let writer = branch::Writer::new(gb_repository, VirtualBranchesHandle::new(&project.gb_dir()))?;
-    writer.write(&mut branch)?;
-
-    let root = gb_repository
-        .root()
-        .join("branches")
-        .join(branch.id.to_string());
-
-    assert_eq!(
-        fs::read_to_string(root.join("meta").join("name").to_str().unwrap())
-            .context("Failed to read branch name")?,
-        branch.name
-    );
-    assert_eq!(
-        fs::read_to_string(root.join("meta").join("applied").to_str().unwrap())?
-            .parse::<bool>()
-            .context("Failed to read branch applied")?,
-        branch.applied
-    );
-    assert_eq!(
-        fs::read_to_string(root.join("meta").join("upstream").to_str().unwrap())
-            .context("Failed to read branch upstream")?,
-        branch.upstream.clone().unwrap().to_string()
-    );
-    assert_eq!(
-        fs::read_to_string(
-            root.join("meta")
-                .join("created_timestamp_ms")
-                .to_str()
-                .unwrap()
-        )
-        .context("Failed to read branch created timestamp")?
-        .parse::<u128>()
-        .context("Failed to parse branch created timestamp")?,
-        branch.created_timestamp_ms
-    );
-    assert_eq!(
-        fs::read_to_string(
-            root.join("meta")
-                .join("updated_timestamp_ms")
-                .to_str()
-                .unwrap()
-        )
-        .context("Failed to read branch updated timestamp")?
-        .parse::<u128>()
-        .context("Failed to parse branch updated timestamp")?,
-        branch.updated_timestamp_ms
-    );
-
-    writer.delete(&branch)?;
-    fs::read_dir(root).unwrap_err();
-
-    Ok(())
-}
-
-#[test]
-fn should_create_session() -> anyhow::Result<()> {
-    let suite = Suite::default();
-    let Case {
-        gb_repository,
-        project,
-        ..
-    } = &suite.new_case();
-
-    let mut branch = new_test_branch();
-
-    let writer = branch::Writer::new(gb_repository, VirtualBranchesHandle::new(&project.gb_dir()))?;
-    writer.write(&mut branch)?;
-
-    assert!(gb_repository.get_current_session()?.is_some());
-
-    Ok(())
-}
-
-#[test]
-fn should_update() -> anyhow::Result<()> {
-    let suite = Suite::default();
-    let Case {
-        gb_repository,
-        project,
-        ..
-    } = &suite.new_case();
-
-    let mut branch = new_test_branch();
-
-    let writer = branch::Writer::new(gb_repository, VirtualBranchesHandle::new(&project.gb_dir()))?;
-    writer.write(&mut branch)?;
-
-    let mut updated_branch = Branch {
-        name: "updated_name".to_string(),
-        applied: false,
-        upstream: Some("refs/remotes/origin/upstream_updated".parse().unwrap()),
-        created_timestamp_ms: 2,
-        updated_timestamp_ms: 3,
-        ownership: gitbutler_core::virtual_branches::branch::BranchOwnershipClaims {
-            claims: vec![],
-        },
-        ..branch.clone()
-    };
-
-    writer.write(&mut updated_branch)?;
-
-    let root = gb_repository
-        .root()
-        .join("branches")
-        .join(branch.id.to_string());
-
-    assert_eq!(
-        fs::read_to_string(root.join("meta").join("name").to_str().unwrap())
-            .context("Failed to read branch name")?,
-        updated_branch.name
-    );
-    assert_eq!(
-        fs::read_to_string(root.join("meta").join("applied").to_str().unwrap())?
-            .parse::<bool>()
-            .context("Failed to read branch applied")?,
-        updated_branch.applied
-    );
-    assert_eq!(
-        fs::read_to_string(root.join("meta").join("upstream").to_str().unwrap())
-            .context("Failed to read branch upstream")?,
-        updated_branch.upstream.unwrap().to_string()
-    );
-    assert_eq!(
-        fs::read_to_string(
-            root.join("meta")
-                .join("created_timestamp_ms")
-                .to_str()
-                .unwrap()
-        )
-        .context("Failed to read branch created timestamp")?
-        .parse::<u128>()
-        .context("Failed to parse branch created timestamp")?,
-        updated_branch.created_timestamp_ms
-    );
-    assert_eq!(
-        fs::read_to_string(
-            root.join("meta")
-                .join("updated_timestamp_ms")
-                .to_str()
-                .unwrap()
-        )
-        .context("Failed to read branch updated timestamp")?
-        .parse::<u128>()
-        .context("Failed to parse branch updated timestamp")?,
-        updated_branch.updated_timestamp_ms
-    );
-
-    Ok(())
-}

crates/gitbutler-core/tests/virtual_branches/iterator.rs

@@ -69,22 +69,12 @@ fn new_test_target() -> virtual_branches::target::Target {
 #[test]
 fn empty_iterator() -> Result<()> {
     let suite = Suite::default();
-    let Case {
+    let Case { project, .. } = &suite.new_case();
-        gb_repository,
-        project,
-        ..
-    } = &suite.new_case();
 
-    let session = gb_repository.get_or_create_current_session()?;
+    let vb_state = VirtualBranchesHandle::new(&project.gb_dir());
-    let session_reader = gitbutler_core::sessions::Reader::open(gb_repository, &session)?;
+    let iter = vb_state.list_branches()?;
 
-    let iter = virtual_branches::Iterator::new(
+    assert_eq!(iter.len(), 0);
-        &session_reader,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-        project.use_toml_vbranches_state(),
-    )?;
-
-    assert_eq!(iter.count(), 0);
 
     Ok(())
 }

crates/gitbutler-core/tests/virtual_branches/mod.rs

@@ -14,7 +14,7 @@ use std::{
 
 use anyhow::{Context, Result};
 use gitbutler_core::{
-    git, reader, sessions,
+    git,
     virtual_branches::{
         self, apply_branch,
         branch::{BranchCreateRequest, BranchOwnershipClaims},
@@ -262,7 +262,6 @@ fn create_branch_with_ownership() -> Result<()> {
     let Case {
         project,
         project_repository,
-        gb_repository,
         ..
     } = &suite.new_case();
 
@@ -276,14 +275,8 @@ fn create_branch_with_ownership() -> Result<()> {
 
     virtual_branches::get_status_by_branch(project_repository, None).expect("failed to get status");
 
-    let current_session = gb_repository.get_or_create_current_session().unwrap();
+    let vb_state = VirtualBranchesHandle::new(&project_repository.project().gb_dir());
-    let current_session_reader = sessions::Reader::open(gb_repository, &current_session).unwrap();
+    let branch0 = vb_state.get_branch(&branch0.id).unwrap();
-    let branch_reader = virtual_branches::branch::Reader::new(
-        &current_session_reader,
-        VirtualBranchesHandle::new(&project_repository.project().gb_dir()),
-        project_repository.project().use_toml_vbranches_state(),
-    );
-    let branch0 = branch_reader.read(&branch0.id).unwrap();
 
     let branch1 = create_virtual_branch(
         project_repository,
@@ -314,9 +307,7 @@ fn create_branch_with_ownership() -> Result<()> {
 fn create_branch_in_the_middle() -> Result<()> {
     let suite = Suite::default();
     let Case {
-        project_repository,
+        project_repository, ..
-        gb_repository,
-        ..
     } = &suite.new_case();
 
     set_test_target(project_repository)?;
@@ -334,16 +325,8 @@ fn create_branch_in_the_middle() -> Result<()> {
     )
     .expect("failed to create virtual branch");
 
-    let current_session = gb_repository.get_or_create_current_session()?;
+    let vb_state = VirtualBranchesHandle::new(&project_repository.project().gb_dir());
-    let current_session_reader = sessions::Reader::open(gb_repository, &current_session)?;
+    let mut branches = vb_state.list_branches().expect("failed to read branches");
-
-    let mut branches = virtual_branches::Iterator::new(
-        &current_session_reader,
-        VirtualBranchesHandle::new(&project_repository.project().gb_dir()),
-        project_repository.project().use_toml_vbranches_state(),
-    )?
-    .collect::<Result<Vec<virtual_branches::Branch>, reader::Error>>()
-    .expect("failed to read branches");
     branches.sort_by_key(|b| b.order);
     assert_eq!(branches.len(), 3);
     assert_eq!(branches[0].name, "Virtual branch");
@@ -357,9 +340,7 @@ fn create_branch_in_the_middle() -> Result<()> {
 fn create_branch_no_arguments() -> Result<()> {
     let suite = Suite::default();
     let Case {
-        project_repository,
+        project_repository, ..
-        gb_repository,
-        ..
     } = &suite.new_case();
 
     set_test_target(project_repository)?;
@@ -367,16 +348,8 @@ fn create_branch_no_arguments() -> Result<()> {
     create_virtual_branch(project_repository, &BranchCreateRequest::default())
         .expect("failed to create virtual branch");
 
-    let current_session = gb_repository.get_or_create_current_session()?;
+    let vb_state = VirtualBranchesHandle::new(&project_repository.project().gb_dir());
-    let current_session_reader = sessions::Reader::open(gb_repository, &current_session)?;
+    let branches = vb_state.list_branches().expect("failed to read branches");
-
-    let branches = virtual_branches::Iterator::new(
-        &current_session_reader,
-        VirtualBranchesHandle::new(&project_repository.project().gb_dir()),
-        project_repository.project().use_toml_vbranches_state(),
-    )?
-    .collect::<Result<Vec<virtual_branches::branch::Branch>, reader::Error>>()
-    .expect("failed to read branches");
     assert_eq!(branches.len(), 1);
     assert_eq!(branches[0].name, "Virtual branch");
     assert!(branches[0].applied);
@@ -519,7 +492,6 @@ fn move_hunks_multiple_sources() -> Result<()> {
     let Case {
         project_repository,
         project,
-        gb_repository,
         ..
     } = &suite.new_case_with_files(HashMap::from([(
         PathBuf::from("test.txt"),
@@ -543,27 +515,17 @@ fn move_hunks_multiple_sources() -> Result<()> {
         "line0\nline1\nline2\nline3\nline4\nline5\nline6\nline7\nline8\nline9\nline10\nline11\nline12\nline13\n",
     )?;
 
-    let current_session = gb_repository.get_or_create_current_session()?;
+    let vb_state = VirtualBranchesHandle::new(&project.gb_dir());
-    let current_session_reader = sessions::Reader::open(gb_repository, &current_session)?;
+    let mut branch2 = vb_state.get_branch(&branch2_id)?;
-    let branch_reader = virtual_branches::branch::Reader::new(
-        &current_session_reader,
-        VirtualBranchesHandle::new(&project_repository.project().gb_dir()),
-        project_repository.project().use_toml_vbranches_state(),
-    );
-    let branch_writer = virtual_branches::branch::Writer::new(
-        gb_repository,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-    )?;
-    let mut branch2 = branch_reader.read(&branch2_id)?;
     branch2.ownership = BranchOwnershipClaims {
         claims: vec!["test.txt:1-5".parse()?],
     };
-    branch_writer.write(&mut branch2)?;
+    vb_state.set_branch(branch2.clone())?;
-    let mut branch1 = branch_reader.read(&branch1_id)?;
+    let mut branch1 = vb_state.get_branch(&branch1_id)?;
     branch1.ownership = BranchOwnershipClaims {
         claims: vec!["test.txt:11-15".parse()?],
     };
-    branch_writer.write(&mut branch1)?;
+    vb_state.set_branch(branch1.clone())?;
 
     let statuses = virtual_branches::get_status_by_branch(project_repository, None)
         .expect("failed to get status")
@@ -757,7 +719,6 @@ fn merge_vbranch_upstream_clean_rebase() -> Result<()> {
     let Case {
         project_repository,
         project,
-        gb_repository,
         ..
     } = &suite.new_case();
 
@@ -829,16 +790,12 @@ fn merge_vbranch_upstream_clean_rebase() -> Result<()> {
     std::fs::write(Path::new(&project.path).join(file_path2), "file2\n")?;
 
     let remote_branch: git::RemoteRefname = "refs/remotes/origin/master".parse().unwrap();
-    let branch_writer = virtual_branches::branch::Writer::new(
-        gb_repository,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-    )?;
     let mut branch = create_virtual_branch(project_repository, &BranchCreateRequest::default())
         .expect("failed to create virtual branch");
     branch.upstream = Some(remote_branch.clone());
     branch.head = last_push;
-    branch_writer
+    vb_state
-        .write(&mut branch)
+        .set_branch(branch.clone())
         .context("failed to write target branch after push")?;
 
     // create the branch
@@ -878,7 +835,6 @@ fn merge_vbranch_upstream_conflict() -> Result<()> {
     let Case {
         project_repository,
         project,
-        gb_repository,
         ..
     } = &suite.new_case();
 
@@ -952,16 +908,12 @@ fn merge_vbranch_upstream_conflict() -> Result<()> {
     )?;
 
     let remote_branch: git::RemoteRefname = "refs/remotes/origin/master".parse().unwrap();
-    let branch_writer = virtual_branches::branch::Writer::new(
-        gb_repository,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-    )?;
     let mut branch = create_virtual_branch(project_repository, &BranchCreateRequest::default())
         .expect("failed to create virtual branch");
     branch.upstream = Some(remote_branch.clone());
     branch.head = last_push;
-    branch_writer
+    vb_state
-        .write(&mut branch)
+        .set_branch(branch.clone())
         .context("failed to write target branch after push")?;
 
     // create the branch
@@ -1255,18 +1207,6 @@ fn detect_mergeable_branch() -> Result<()> {
         .expect("failed to create virtual branch")
         .id;
 
-    let current_session = gb_repository.get_or_create_current_session()?;
-    let current_session_reader = sessions::Reader::open(gb_repository, &current_session)?;
-    let branch_reader = virtual_branches::branch::Reader::new(
-        &current_session_reader,
-        VirtualBranchesHandle::new(&project_repository.project().gb_dir()),
-        project_repository.project().use_toml_vbranches_state(),
-    );
-    let branch_writer = virtual_branches::branch::Writer::new(
-        gb_repository,
-        VirtualBranchesHandle::new(&project.gb_dir()),
-    )?;
-
     update_branch(
         project_repository,
         virtual_branches::branch::BranchUpdateRequest {
@@ -1357,11 +1297,13 @@ fn detect_mergeable_branch() -> Result<()> {
         "line1\nline2\nline3\nline4\nbranch4\n",
     )?;
 
-    let mut branch4 = branch_reader.read(&branch4_id)?;
+    let vb_state = VirtualBranchesHandle::new(&project.gb_dir());
+
+    let mut branch4 = vb_state.get_branch(&branch4_id)?;
     branch4.ownership = BranchOwnershipClaims {
         claims: vec!["test2.txt:1-6".parse()?],
     };
-    branch_writer.write(&mut branch4)?;
+    vb_state.set_branch(branch4.clone())?;
 
     let (branches, _) = virtual_branches::list_virtual_branches(project_repository)?;
     assert_eq!(branches.len(), 4);