Prio3Count: Don't branch when converting to bool

Avoid branching on the value of `DapMeasurement::U64(measurement)` when converting to boolean.
cloudflare · Dec 19, 2024 · 23adb59 · 23adb59
1 parent 16d61c2
commit 23adb59
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 24 deletions.
diff --git a/crates/daphne/src/vdaf/prio3.rs b/crates/daphne/src/vdaf/prio3.rs
@@ -33,23 +33,13 @@ pub(crate) fn prio3_shard(
     task_id: TaskId,
 ) -> Result<(Vec<u8>, [Vec<u8>; 2]), VdafError> {
     match (config, measurement) {
-        (Prio3Config::Count, DapMeasurement::U64(measurement)) => {
+        (Prio3Config::Count, DapMeasurement::U64(measurement)) if measurement < 2 => {
             let vdaf = Prio3::new_count(2).map_err(|e| {
                 VdafError::Dap(
                     fatal_error!(err = ?e, "failed to create prio3 count from num_aggregators(2)"),
                 )
             })?;
-            // TODO(cjpatton) Make this constant time.
-            let measurement = match measurement {
-                0 => false,
-                1 => true,
-                _ => {
-                    return Err(VdafError::Dap(fatal_error!(
-                        err = "cannot represent measurement as a 0 or 1"
-                    )))
-                }
-            };
-            shard_then_encode(&vdaf, task_id, &measurement, nonce)
+            shard_then_encode(&vdaf, task_id, &(measurement != 0), nonce)
         }
         (
             Prio3Config::Histogram {

diff --git a/crates/daphne/src/vdaf/prio3_draft09.rs b/crates/daphne/src/vdaf/prio3_draft09.rs
@@ -63,23 +63,13 @@ pub(crate) fn prio3_draft09_shard(
     nonce: &[u8; 16],
 ) -> Result<(Vec<u8>, [Vec<u8>; 2]), VdafError> {
     match (config, measurement) {
-        (Prio3Config::Count, DapMeasurement::U64(measurement)) => {
+        (Prio3Config::Count, DapMeasurement::U64(measurement)) if measurement < 2 => {
             let vdaf = Prio3::new_count(2).map_err(|e| {
                 VdafError::Dap(
                     fatal_error!(err = ?e, "failed to create prio3 count from num_aggregators(2)"),
                 )
             })?;
-            // TODO(cjpatton) Make this constant time.
-            let measurement = match measurement {
-                0 => false,
-                1 => true,
-                _ => {
-                    return Err(VdafError::Dap(fatal_error!(
-                        err = "cannot represent measurement as a 0 or 1"
-                    )))
-                }
-            };
-            shard_then_encode_draft09(&vdaf, &measurement, nonce)
+            shard_then_encode_draft09(&vdaf, &(measurement != 0), nonce)
         }
         (
             Prio3Config::Histogram {