From 23adb593eafea31823330cde2b13a43c5f670d64 Mon Sep 17 00:00:00 2001
From: Christopher Patton <cpatton@cloudflare.com>
Date: Wed, 18 Dec 2024 13:14:36 -0800
Subject: [PATCH] Prio3Count: Don't branch when converting to bool

Avoid branching on the value of `DapMeasurement::U64(measurement)` when
converting to boolean.
---
 crates/daphne/src/vdaf/prio3.rs         | 14 ++------------
 crates/daphne/src/vdaf/prio3_draft09.rs | 14 ++------------
 2 files changed, 4 insertions(+), 24 deletions(-)

diff --git a/crates/daphne/src/vdaf/prio3.rs b/crates/daphne/src/vdaf/prio3.rs
index cb162ed3..949ca61e 100644
--- 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
index 0dd5251a..bd13f30b 100644
--- 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 {