From 6b5544f33518b78ab2b59897d4a28300e5fd9cb2 Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Wed, 23 Aug 2023 18:14:05 -0700
Subject: [PATCH] tree_builder: add a `set_or_remove()` and simplify callers

Many of the `TreeBuilder` users have an `Option<TreeValue>` and call
either `set()` or `remove()` or the builder depending on whether the
value is present. Let's centralize this logic in a new
`TreeBuilder::set_or_remove()`.
---
 cli/src/cli_util.rs     |  9 +--------
 cli/src/commands/mod.rs |  9 +--------
 lib/src/merged_tree.rs  | 10 ++--------
 lib/src/tree_builder.rs |  9 +++++++++
 4 files changed, 13 insertions(+), 24 deletions(-)

diff --git a/cli/src/cli_util.rs b/cli/src/cli_util.rs
index 15ed981d8..93ebd10ab 100644
--- a/cli/src/cli_util.rs
+++ b/cli/src/cli_util.rs
@@ -1498,14 +1498,7 @@ impl WorkspaceCommandTransaction<'_> {
         } else {
             let mut tree_builder = self.repo().store().tree_builder(left_tree.id().clone());
             for (repo_path, diff) in left_tree.diff(right_tree, matcher) {
-                match diff.into_options().1 {
-                    Some(value) => {
-                        tree_builder.set(repo_path, value);
-                    }
-                    None => {
-                        tree_builder.remove(repo_path);
-                    }
-                }
+                tree_builder.set_or_remove(repo_path, diff.into_options().1);
             }
             Ok(tree_builder.write_tree())
         }
diff --git a/cli/src/commands/mod.rs b/cli/src/commands/mod.rs
index 2db9fd1b6..7754e1e0a 100644
--- a/cli/src/commands/mod.rs
+++ b/cli/src/commands/mod.rs
@@ -2896,14 +2896,7 @@ fn cmd_restore(
             .store()
             .tree_builder(to_commit.tree_id().clone());
         for (repo_path, diff) in from_tree.diff(&to_commit.tree(), matcher.as_ref()) {
-            match diff.into_options().0 {
-                Some(value) => {
-                    tree_builder.set(repo_path, value);
-                }
-                None => {
-                    tree_builder.remove(repo_path);
-                }
-            }
+            tree_builder.set_or_remove(repo_path, diff.into_options().0);
         }
         tree_builder.write_tree()
     };
diff --git a/lib/src/merged_tree.rs b/lib/src/merged_tree.rs
index afcee3951..563caccf9 100644
--- a/lib/src/merged_tree.rs
+++ b/lib/src/merged_tree.rs
@@ -117,16 +117,10 @@ impl MergedTree {
             }
             // Now add the terms that were present in the conflict to the appropriate trees.
             for (i, term) in conflict.removes().iter().enumerate() {
-                match term {
-                    Some(value) => removes[i].set(path.clone(), value.clone()),
-                    None => removes[i].remove(path.clone()),
-                }
+                removes[i].set_or_remove(path.clone(), term.clone());
             }
             for (i, term) in conflict.adds().iter().enumerate() {
-                match term {
-                    Some(value) => adds[i].set(path.clone(), value.clone()),
-                    None => adds[i].remove(path.clone()),
-                }
+                adds[i].set_or_remove(path.clone(), term.clone());
             }
         }
 
diff --git a/lib/src/tree_builder.rs b/lib/src/tree_builder.rs
index 59e0e067a..66c1582a6 100644
--- a/lib/src/tree_builder.rs
+++ b/lib/src/tree_builder.rs
@@ -64,6 +64,15 @@ impl TreeBuilder {
         self.overrides.insert(path, Override::Tombstone);
     }
 
+    pub fn set_or_remove(&mut self, path: RepoPath, value: Option<TreeValue>) {
+        assert!(!path.is_root());
+        if let Some(value) = value {
+            self.overrides.insert(path, Override::Replace(value));
+        } else {
+            self.overrides.insert(path, Override::Tombstone);
+        }
+    }
+
     pub fn write_tree(self) -> TreeId {
         if self.overrides.is_empty() {
             return self.base_tree_id;