From feab9b17a64969be4154365068f604052916460d Mon Sep 17 00:00:00 2001
From: Onyeka Obi <Onyeka.Obi@gmail.com>
Date: Thu, 16 Feb 2023 15:41:30 -0800
Subject: [PATCH] Return errors involving missing links instead of panicking

---
 src/error.rs                |  6 ++++++
 src/tree/mod.rs             | 13 +++++++------
 src/tree/walk/ref_walker.rs | 13 +++++--------
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 4f35da31..def66420 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -30,6 +30,12 @@ pub enum Error {
     KeyNotFound(String),
     #[error("Proof is missing data for query")]
     MissingData,
+    #[error("Expected link")]
+    MissingLink,
+    #[error("Expected Some(Link::Reference)")]
+    MissingLinkReference,
+    #[error("Cannot traverse Link::Modified")]
+    ModifiedLinkTraversal,
     #[error("Path Error: {0}")]
     Path(String),
     #[error("Proof Error: {0}")]
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index 0f38ee0a..071404c2 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -14,7 +14,7 @@ use std::cmp::max;
 
 use ed::{Decode, Encode};
 
-use super::error::Result;
+use super::error::{Error, Result};
 pub use commit::{Commit, NoopCommit};
 pub use hash::{kv_hash, node_hash, Hash, Hasher, HASH_LENGTH, NULL_HASH};
 use kv::KV;
@@ -380,16 +380,17 @@ impl Tree {
     /// `Link::Loaded`).
     #[inline]
     pub fn load<S: Fetch>(&mut self, left: bool, source: &S) -> Result<()> {
-        // TODO: return Err instead of panic?
-        let link = self.link(left).expect("Expected link");
+        let link = self.link(left).ok_or(Error::MissingLink)?;
         let (child_heights, hash) = match link {
             Link::Reference {
                 child_heights,
                 hash,
                 ..
-            } => (child_heights, hash),
-            _ => panic!("Expected Some(Link::Reference)"),
-        };
+            } => Ok((child_heights, hash)),
+            Link::Modified { .. } | Link::Uncommitted { .. } | Link::Loaded { .. } => {
+                Err(Error::MissingLinkReference)
+            }
+        }?;
 
         let tree = source.fetch(link)?;
         debug_assert_eq!(tree.key(), link.key());
diff --git a/src/tree/walk/ref_walker.rs b/src/tree/walk/ref_walker.rs
index 95602693..014654ab 100644
--- a/src/tree/walk/ref_walker.rs
+++ b/src/tree/walk/ref_walker.rs
@@ -1,6 +1,6 @@
 use super::super::{Link, Tree};
 use super::Fetch;
-use crate::error::Result;
+use crate::error::{Error::ModifiedLinkTraversal, Result};
 
 /// Allows read-only traversal of a `Tree`, fetching from the given source when
 /// traversing to a pruned node. The fetched nodes are then retained in memory
@@ -22,7 +22,6 @@ where
 {
     /// Creates a `RefWalker` with the given tree and source.
     pub fn new(tree: &'a mut Tree, source: S) -> Self {
-        // TODO: check if tree has modified links, panic if so
         RefWalker { tree, source }
     }
 
@@ -41,12 +40,10 @@ where
         };
 
         match link {
-            Link::Reference { .. } => {
-                self.tree.load(left, &self.source)?;
-            }
-            Link::Modified { .. } => panic!("Cannot traverse Link::Modified"),
-            Link::Uncommitted { .. } | Link::Loaded { .. } => {}
-        }
+            Link::Reference { .. } => self.tree.load(left, &self.source),
+            Link::Modified { .. } => Err(ModifiedLinkTraversal),
+            Link::Uncommitted { .. } | Link::Loaded { .. } => Ok(()),
+        }?;
 
         let child = self.tree.child_mut(left).unwrap();
         Ok(Some(RefWalker::new(child, self.source.clone())))