Multiple case: check consistency with leaf index

src/merkle/smt/full/error.rs

@@ -12,16 +12,20 @@ use crate::{
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum SmtLeafError {
-    InconsistentKeysInEntries {
+    InconsistentKeys {
         entries: Vec<(RpoDigest, Word)>,
         key_1: RpoDigest,
         key_2: RpoDigest,
     },
     InvalidNumEntriesForMultiple(usize),
-    KeyInconsistentWithLeafIndex {
+    SingleKeyInconsistentWithLeafIndex {
         key: RpoDigest,
         leaf_index: LeafIndex<SMT_DEPTH>,
     },
+    MultipleKeysInconsistentWithLeafIndex {
+        leaf_index_from_keys: LeafIndex<SMT_DEPTH>,
+        leaf_index_supplied: LeafIndex<SMT_DEPTH>,
+    },
 }
 
 #[cfg(feature = "std")]
@@ -34,16 +38,27 @@ impl fmt::Display for SmtLeafError {
             InvalidNumEntriesForMultiple(num_entries) => {
                 write!(f, "Multiple leaf requires 2 or more entries. Got: {num_entries}")
             }
-            InconsistentKeysInEntries { entries, key_1, key_2 } => {
+            InconsistentKeys { entries, key_1, key_2 } => {
                 write!(f, "Multiple leaf requires all keys to map to the same leaf index. Offending keys: {key_1} and {key_2}. Entries: {entries:?}.")
             }
-            KeyInconsistentWithLeafIndex { key, leaf_index } => {
+            SingleKeyInconsistentWithLeafIndex { key, leaf_index } => {
                 write!(
                     f,
                     "Single key in leaf inconsistent with leaf index. Key: {key}, leaf index: {}",
                     leaf_index.value()
                 )
             }
+            MultipleKeysInconsistentWithLeafIndex {
+                leaf_index_from_keys,
+                leaf_index_supplied,
+            } => {
+                write!(
+                    f,
+                    "Keys in entries map to leaf index {}, but leaf index {} was supplied",
+                    leaf_index_from_keys.value(),
+                    leaf_index_supplied.value()
+                )
+            }
         }
     }
 }

src/merkle/smt/full/leaf.rs

@@ -22,6 +22,8 @@ impl SmtLeaf {
     ///
     /// # Errors
     ///   - Returns an error if 2 keys in `entries` map to a different leaf index
+    ///   - Returns an error if 1 or more keys in `entries` map to a leaf index
+    ///     different from `leaf_index`
     pub fn new(
         entries: Vec<(RpoDigest, Word)>,
         leaf_index: LeafIndex<SMT_DEPTH>,
@@ -32,12 +34,28 @@ impl SmtLeaf {
                 let (key, value) = entries[0];
 
                 if LeafIndex::<SMT_DEPTH>::from(key) != leaf_index {
-                    return Err(SmtLeafError::KeyInconsistentWithLeafIndex { key, leaf_index });
+                    return Err(SmtLeafError::SingleKeyInconsistentWithLeafIndex {
+                        key,
+                        leaf_index,
+                    });
                 }
 
                 Ok(Self::new_single(key, value))
             }
-            _ => Self::new_multiple(entries),
+            _ => {
+                let leaf = Self::new_multiple(entries)?;
+
+                // `new_multiple()` checked that all keys map to the same leaf index. We still need
+                // to ensure that that leaf index is `leaf_index`.
+                if leaf.index() != leaf_index {
+                    Err(SmtLeafError::MultipleKeysInconsistentWithLeafIndex {
+                        leaf_index_from_keys: leaf.index(),
+                        leaf_index_supplied: leaf_index,
+                    })
+                } else {
+                    Ok(leaf)
+                }
+            }
         }
     }
 
@@ -73,7 +91,7 @@ impl SmtLeaf {
                 let next_leaf_index: LeafIndex<SMT_DEPTH> = next_key.into();
 
                 if next_leaf_index != first_leaf_index {
-                    return Err(SmtLeafError::InconsistentKeysInEntries {
+                    return Err(SmtLeafError::InconsistentKeys {
                         entries,
                         key_1: first_key,
                         key_2: next_key,