Merge rust-bitcoin/rust-bitcoin#1710: Support weight prediction in `const` context

00b46d6d9d Indent functions (Martin Habovstiak)
d56d202aeb Support weight prediction in `const` context (Martin Habovstiak)

Pull request description:

  **Notes for reviewers:**
  This is something that I want to use in my code and hopefully reasonably easy to review, so if this can get into 0.30 that'd be really nice. No hard feelings if it doesn't.
  I tried to put extra effort into making review easier by:
  * intentionally "mis-formatting" the first commit so diff is smaller and easy to understand - see individual commits.
  * copying patterns from non-const fn to const fn so it's obviously correct (includes same variable names)
  * not bothering with the array trick in `VarInt::len` and simply accepting the limitation of Rust 1.46+ (I use 1.48 BTW).

  **Description**

  Some smart contracts or simplified wallets statically know the sizes of
  transactions or inputs. The possible approaches to handling them so far
  were re-computing the values (and hoping the optimizer will const fold
  them) or using a simple constant which may be harder to understand and
  get right. It's much nicer to just use a `const` but our code didn't
  support it until now.

  This change adds methods that can compute the prediction in `const`
  context for Rust versions >= 1.46.0 which allow use of loops (and
  conditions but those could be workaround anyway).

  As a side effect of this, the change also adds `const` to `VarInt::len`
  in Rust 1.46+. While this one could be made unconditional using array
  trick it's probably not worth it because of the planned MSRV bump.

ACKs for top commit:
  apoelstra:
    ACK 00b46d6d9d
  tcharding:
    ACK 00b46d6d9d

Tree-SHA512: 5509886a68b4de5227db0e28d92a40be8de64592e0b189c519213db21bcfe98ca03d9a1936b1024729b97db69e8ec0b55fac870a7ce9bab0d0c9a47b2087990f

bitcoin/src/blockdata/transaction.rs

@@ -1247,23 +1247,64 @@ pub fn predict_weight<I, O>(inputs: I, output_script_lens: O) -> Weight
                 let script_size = script_len + VarInt(script_len as u64).len();
                 (output_count + 1, total_scripts_size + script_size)
         });
-    let input_weight = partial_input_weight + input_count * 4 * (32 + 4 + 4);
+    predict_weight_internal(input_count, partial_input_weight, inputs_with_witnesses, output_count, output_scripts_size)
-    let output_size = 8 * output_count + output_scripts_size;
+}
-    let non_input_size =
+
-        // version:
+crate::internal_macros::maybe_const_fn! {
-        4 +
+    fn predict_weight_internal(input_count: usize, partial_input_weight: usize, inputs_with_witnesses: usize, output_count: usize, output_scripts_size: usize) -> Weight {
-        // count varints:
+        let input_weight = partial_input_weight + input_count * 4 * (32 + 4 + 4);
-        VarInt(input_count as u64).len() +
+        let output_size = 8 * output_count + output_scripts_size;
-        VarInt(output_count as u64).len() +
+        let non_input_size =
-        output_size +
+            // version:
-        // lock_time
+            4 +
-        4;
+            // count varints:
-    let weight = if inputs_with_witnesses == 0 {
+            VarInt(input_count as u64).len() +
-        non_input_size * 4 + input_weight
+            VarInt(output_count as u64).len() +
-    } else {
+            output_size +
-        non_input_size * 4 + input_weight + input_count - inputs_with_witnesses + 2
+            // lock_time
-    };
+            4;
-    Weight::from_wu(weight as u64)
+        let weight = if inputs_with_witnesses == 0 {
+            non_input_size * 4 + input_weight
+        } else {
+            non_input_size * 4 + input_weight + input_count - inputs_with_witnesses + 2
+        };
+        Weight::from_wu(weight as u64)
+    }
+}
+
+/// Predicts the weight of a to-be-constructed transaction in const context.
+///
+/// *Important: only available in Rust 1.46+*
+///
+/// This is a `const` version of [`predict_weight`] which only allows slices due to current Rust
+/// limitations around `const fn`. Because of these limitations it may be less efficient than
+/// `predict_weight` and thus is intended to be only used in `const` context.
+///
+/// Please see the documentation of `predict_weight` to learn more about this function.
+#[cfg(rust_v_1_46)]
+pub const fn predict_weight_from_slices(inputs: &[InputWeightPrediction], output_script_lens: &[usize]) -> Weight {
+    let mut partial_input_weight = 0;
+    let mut inputs_with_witnesses = 0;
+
+    // for loops not supported in const fn
+    let mut i = 0;
+    while i < inputs.len() {
+        let prediction = inputs[i];
+        partial_input_weight += prediction.script_size * 4 + prediction.witness_size;
+        inputs_with_witnesses += (prediction.witness_size > 0) as usize;
+        i += 1;
+    }
+
+    let mut output_scripts_size = 0;
+
+    i = 0;
+    while i < output_script_lens.len() {
+        let script_len = output_script_lens[i];
+        output_scripts_size += script_len + VarInt(script_len as u64).len();
+        i += 1;
+    }
+
+    predict_weight_internal(inputs.len(), partial_input_weight, inputs_with_witnesses, output_script_lens.len(), output_scripts_size)
 }
 
 /// Weight prediction of an individual input.
@@ -1351,6 +1392,37 @@ impl InputWeightPrediction {
             witness_size,
         }
     }
+
+    /// Computes the prediction for a single input in `const` context.
+    ///
+    /// *Important: only available in Rust 1.46+*
+    ///
+    /// This is a `const` version of [`new`](Self::new) which only allows slices due to current Rust
+    /// limitations around `const fn`. Because of these limitations it may be less efficient than
+    /// `new` and thus is intended to be only used in `const` context.
+    #[cfg(rust_v_1_46)]
+    pub const fn from_slice(input_script_len: usize, witness_element_lengths: &[usize]) -> Self {
+        let mut i = 0;
+        let mut total_size = 0;
+        // for loops not supported in const fn
+        while i < witness_element_lengths.len() {
+            let elem_len = witness_element_lengths[i];
+            let elem_size = elem_len + VarInt(elem_len as u64).len();
+            total_size += elem_size;
+            i += 1;
+        }
+        let witness_size = if !witness_element_lengths.is_empty() {
+            total_size + VarInt(witness_element_lengths.len() as u64).len()
+        } else {
+            0
+        };
+        let script_size = input_script_len + VarInt(input_script_len as u64).len();
+
+        InputWeightPrediction {
+            script_size,
+            witness_size,
+        }
+    }
 }
 
 #[cfg(test)]

bitcoin/src/consensus/encode.rs

@@ -372,17 +372,21 @@ impl_int_encodable!(i64, read_i64, emit_i64);
 
 #[allow(clippy::len_without_is_empty)] // VarInt has on concept of 'is_empty'.
 impl VarInt {
-    /// Gets the length of this VarInt when encoded.
+    crate::internal_macros::maybe_const_fn! {
-    ///
+        /// Gets the length of this VarInt when encoded.
-    /// Returns 1 for 0..=0xFC, 3 for 0xFD..=(2^16-1), 5 for 0x10000..=(2^32-1),
+        ///
-    /// and 9 otherwise.
+        /// *Important: this method is only `const` in Rust 1.46 or higher!*
-    #[inline]
+        ///
-    pub fn len(&self) -> usize {
+        /// Returns 1 for 0..=0xFC, 3 for 0xFD..=(2^16-1), 5 for 0x10000..=(2^32-1),
-        match self.0 {
+        /// and 9 otherwise.
-            0..=0xFC             => { 1 }
+        #[inline]
-            0xFD..=0xFFFF        => { 3 }
+        pub fn len(&self) -> usize {
-            0x10000..=0xFFFFFFFF => { 5 }
+            match self.0 {
-            _                    => { 9 }
+                0..=0xFC             => { 1 }
+                0xFD..=0xFFFF        => { 3 }
+                0x10000..=0xFFFFFFFF => { 5 }
+                _                    => { 9 }
+            }
         }
     }
 }

bitcoin/src/internal_macros.rs

@@ -45,6 +45,20 @@ macro_rules! impl_consensus_encoding {
     );
 }
 pub(crate) use impl_consensus_encoding;
+
+/// Marks the function const in Rust 1.46.0
+macro_rules! maybe_const_fn {
+    ($(#[$attr:meta])* $vis:vis fn $name:ident($($args:tt)*) -> $ret:ty $body:block) => {
+        #[cfg(rust_v_1_46)]
+        $(#[$attr])*
+        $vis const fn $name($($args)*) -> $ret $body
+
+        #[cfg(not(rust_v_1_46))]
+        $(#[$attr])*
+        $vis fn $name($($args)*) -> $ret $body
+    }
+}
+pub(crate) use maybe_const_fn;
 // We use test_macros module to keep things organised, re-export everything for ease of use.
 #[cfg(test)]
 pub(crate) use test_macros::*;