From e660dffe567fc12d831e4a3cb72382b4d7f4d305 Mon Sep 17 00:00:00 2001
From: Christopher Patton <cpatton@cloudflare.com>
Date: Mon, 6 Jan 2025 12:56:48 -0800
Subject: [PATCH] vidpf: Move `eval_prefix_tree_with_siblings()` to `impl<W:
 VidpfValue>`

This method is currently implemented for `Vidpf<VidpfWeight<F>>`, but it
applies to the more general `Vidpf<W>`.
---
 src/vidpf.rs | 130 +++++++++++++++++++++++++--------------------------
 1 file changed, 64 insertions(+), 66 deletions(-)

diff --git a/src/vidpf.rs b/src/vidpf.rs
index 9458a39b..d7c68b2a 100644
--- a/src/vidpf.rs
+++ b/src/vidpf.rs
@@ -298,6 +298,70 @@ impl<W: VidpfValue> Vidpf<W> {
         Ok(beta_share)
     }
 
+    /// Ensure `prefix_tree` contains the prefix tree for `prefixes`, as well as the sibling of
+    /// each node in the prefix tree. The return value is the weights for the prefixes
+    /// concatenated together.
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn eval_prefix_tree_with_siblings(
+        &self,
+        ctx: &[u8],
+        id: VidpfServerId,
+        public: &VidpfPublicShare<W>,
+        key: &VidpfKey,
+        nonce: &[u8],
+        prefixes: &[VidpfInput],
+        prefix_tree: &mut BinaryTree<VidpfEvalResult<W>>,
+    ) -> Result<Vec<W>, VidpfError> {
+        let mut out_shares = Vec::with_capacity(prefixes.len());
+
+        for prefix in prefixes {
+            if prefix.len() > public.cw.len() {
+                return Err(VidpfError::InvalidInputLength);
+            }
+
+            let mut sub_tree = prefix_tree.root.get_or_insert_with(|| {
+                Box::new(Node::new(VidpfEvalResult {
+                    state: VidpfEvalState::init_from_key(id, key),
+                    share: W::zero(&self.weight_parameter), // not used
+                }))
+            });
+
+            for (idx, cw) in self.index_iter(prefix)?.zip(public.cw.iter()) {
+                let left = sub_tree.left.get_or_insert_with(|| {
+                    Box::new(Node::new(self.eval_next(
+                        ctx,
+                        cw,
+                        idx.left_sibling(),
+                        &sub_tree.value.state,
+                        nonce,
+                    )))
+                });
+                let right = sub_tree.right.get_or_insert_with(|| {
+                    Box::new(Node::new(self.eval_next(
+                        ctx,
+                        cw,
+                        idx.right_sibling(),
+                        &sub_tree.value.state,
+                        nonce,
+                    )))
+                });
+
+                sub_tree = if idx.bit.unwrap_u8() == 0 {
+                    left
+                } else {
+                    right
+                };
+            }
+
+            out_shares.push(sub_tree.value.share.clone());
+        }
+
+        for out_share in out_shares.iter_mut() {
+            out_share.conditional_negate(Choice::from(id));
+        }
+        Ok(out_shares)
+    }
+
     fn extend(seed: VidpfSeed, ctx: &[u8], nonce: &[u8]) -> ExtendedSeed {
         let mut rng = XofFixedKeyAes128::seed_stream(
             &Seed(seed),
@@ -371,72 +435,6 @@ impl<W: VidpfValue> Vidpf<W> {
     }
 }
 
-impl<F: FieldElement> Vidpf<VidpfWeight<F>> {
-    /// Ensure `prefix_tree` contains the prefix tree for `prefixes`, as well as the sibling of
-    /// each node in the prefix tree. The return value is the weights for the prefixes
-    /// concatenated together.
-    #[allow(clippy::too_many_arguments)]
-    pub(crate) fn eval_prefix_tree_with_siblings(
-        &self,
-        ctx: &[u8],
-        id: VidpfServerId,
-        public: &VidpfPublicShare<VidpfWeight<F>>,
-        key: &VidpfKey,
-        nonce: &[u8],
-        prefixes: &[VidpfInput],
-        prefix_tree: &mut BinaryTree<VidpfEvalResult<VidpfWeight<F>>>,
-    ) -> Result<Vec<VidpfWeight<F>>, VidpfError> {
-        let mut out_shares = Vec::with_capacity(prefixes.len());
-
-        for prefix in prefixes {
-            if prefix.len() > public.cw.len() {
-                return Err(VidpfError::InvalidInputLength);
-            }
-
-            let mut sub_tree = prefix_tree.root.get_or_insert_with(|| {
-                Box::new(Node::new(VidpfEvalResult {
-                    state: VidpfEvalState::init_from_key(id, key),
-                    share: VidpfWeight::zero(&self.weight_parameter), // not used
-                }))
-            });
-
-            for (idx, cw) in self.index_iter(prefix)?.zip(public.cw.iter()) {
-                let left = sub_tree.left.get_or_insert_with(|| {
-                    Box::new(Node::new(self.eval_next(
-                        ctx,
-                        cw,
-                        idx.left_sibling(),
-                        &sub_tree.value.state,
-                        nonce,
-                    )))
-                });
-                let right = sub_tree.right.get_or_insert_with(|| {
-                    Box::new(Node::new(self.eval_next(
-                        ctx,
-                        cw,
-                        idx.right_sibling(),
-                        &sub_tree.value.state,
-                        nonce,
-                    )))
-                });
-
-                sub_tree = if idx.bit.unwrap_u8() == 0 {
-                    left
-                } else {
-                    right
-                };
-            }
-
-            out_shares.push(sub_tree.value.share.clone());
-        }
-
-        for out_share in out_shares.iter_mut() {
-            out_share.conditional_negate(Choice::from(id));
-        }
-        Ok(out_shares)
-    }
-}
-
 /// VIDPF key.
 ///
 /// Private key of an aggregation server.