From 1577b408a69f76a27d2bf69a19a2dede469f58da Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Fri, 25 Aug 2023 18:49:15 -0700
Subject: [PATCH] store: add function to look up `MergedTree` by `MergedTreeId`

We'll start seeing `MergedTreeId` in more places and we'll want it to
be easy to look up the tree.
---
 lib/src/commit.rs | 11 +----------
 lib/src/store.rs  | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/lib/src/commit.rs b/lib/src/commit.rs
index 391e67aa2..98d013741 100644
--- a/lib/src/commit.rs
+++ b/lib/src/commit.rs
@@ -110,16 +110,7 @@ impl Commit {
     }
 
     pub fn merged_tree(&self) -> Result<MergedTree, BackendError> {
-        match &self.data.root_tree {
-            MergedTreeId::Legacy(id) => {
-                let tree = self.store.get_tree(&RepoPath::root(), id)?;
-                Ok(MergedTree::Legacy(tree))
-            }
-            MergedTreeId::Merge(ids) => {
-                let trees = ids.try_map(|id| self.store.get_tree(&RepoPath::root(), id))?;
-                Ok(MergedTree::Merge(trees))
-            }
-        }
+        self.store.get_root_tree(&self.data.root_tree)
     }
 
     // TODO(#1624): delete when all callers have been updated to support tree-level
diff --git a/lib/src/store.rs b/lib/src/store.rs
index 5339041a5..020c21497 100644
--- a/lib/src/store.rs
+++ b/lib/src/store.rs
@@ -21,10 +21,12 @@ use std::sync::{Arc, RwLock};
 
 use crate::backend;
 use crate::backend::{
-    Backend, BackendResult, ChangeId, CommitId, ConflictId, FileId, SymlinkId, TreeId, TreeValue,
+    Backend, BackendResult, ChangeId, CommitId, ConflictId, FileId, MergedTreeId, SymlinkId,
+    TreeId, TreeValue,
 };
 use crate::commit::Commit;
 use crate::merge::Merge;
+use crate::merged_tree::MergedTree;
 use crate::repo_path::RepoPath;
 use crate::tree::Tree;
 use crate::tree_builder::TreeBuilder;
@@ -125,6 +127,19 @@ impl Store {
         Ok(data)
     }
 
+    pub fn get_root_tree(self: &Arc<Self>, id: &MergedTreeId) -> BackendResult<MergedTree> {
+        match &id {
+            MergedTreeId::Legacy(id) => {
+                let tree = self.get_tree(&RepoPath::root(), id)?;
+                Ok(MergedTree::Legacy(tree))
+            }
+            MergedTreeId::Merge(ids) => {
+                let trees = ids.try_map(|id| self.get_tree(&RepoPath::root(), id))?;
+                Ok(MergedTree::Merge(trees))
+            }
+        }
+    }
+
     pub fn write_tree(
         self: &Arc<Self>,
         path: &RepoPath,