From a119a27ee715e4258964cddeea0f9e92ed10957a Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@electriccoin.co>
Date: Fri, 30 Apr 2021 07:59:46 -0600
Subject: [PATCH] Clean up value balance generation.

---
 src/bundle.rs | 21 ++++++++-------------
 src/value.rs  | 14 +++++++++++++-
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/bundle.rs b/src/bundle.rs
index 9a9fa063..5ebab773 100644
--- a/src/bundle.rs
+++ b/src/bundle.rs
@@ -358,7 +358,8 @@ pub mod testing {
             },
         },
         value::{
-            testing::arb_note_value, NoteValue, ValueCommitTrapdoor, ValueCommitment, ValueSum,
+            testing::{arb_positive_note_value, MAX_MONEY},
+            NoteValue, ValueCommitTrapdoor, ValueCommitment, ValueSum,
         },
         Anchor,
     };
@@ -450,7 +451,7 @@ pub mod testing {
         /// [`crate::builder::testing::arb_bundle`]
         pub fn arb_unauthorized_bundle()(
             acts in vec(
-                arb_note_value().prop_flat_map(|v|
+                arb_positive_note_value(MAX_MONEY as u64 / 10000).prop_flat_map(|v|
                     arb_unauthorized_action(v).prop_map(move |a| (v, a))
                 ),
                 1..10
@@ -458,15 +459,12 @@ pub mod testing {
             flags in arb_flags(),
             anchor in prop::array::uniform32(prop::num::u8::ANY).prop_map(Anchor)
         ) -> Bundle<Unauthorized, ValueSum> {
-            let (values, actions): (Vec<NoteValue>, Vec<Action<()>>) = acts.into_iter().unzip();
+            let (note_values, actions): (Vec<NoteValue>, Vec<Action<_>>) = acts.into_iter().unzip();
 
             Bundle::from_parts(
                 NonEmpty::from_vec(actions).unwrap(),
                 flags,
-                values.into_iter().fold(
-                    ValueSum::zero(),
-                    |acc, cv| (acc + (cv - NoteValue::zero()).unwrap()).unwrap()
-                ),
+                note_values.into_iter().map(ValueSum::from).sum::<Result<ValueSum, _>>().unwrap(),
                 anchor,
                 Unauthorized
             )
@@ -479,7 +477,7 @@ pub mod testing {
         /// [`crate::builder::testing::arb_bundle`]
         pub fn arb_bundle()(
             acts in vec(
-                arb_note_value().prop_flat_map(|v|
+                arb_positive_note_value(MAX_MONEY as u64 / 10000).prop_flat_map(|v|
                     arb_action(v).prop_map(move |a| (v, a))
                 ),
                 1..10
@@ -491,16 +489,13 @@ pub mod testing {
             fake_proof in vec(prop::num::u8::ANY, 1973),
             fake_sighash in prop::array::uniform32(prop::num::u8::ANY),
         ) -> Bundle<Authorized, ValueSum> {
-            let (values, actions): (Vec<NoteValue>, Vec<Action<redpallas::Signature<reddsa::orchard::SpendAuth>>>) = acts.into_iter().unzip();
+            let (note_values, actions): (Vec<NoteValue>, Vec<Action<_>>) = acts.into_iter().unzip();
             let rng = StdRng::from_seed(rng_seed);
 
             Bundle::from_parts(
                 NonEmpty::from_vec(actions).unwrap(),
                 flags,
-                values.into_iter().fold(
-                    ValueSum::zero(),
-                    |acc, cv| (acc + (cv - NoteValue::zero()).unwrap()).unwrap()
-                ),
+                note_values.into_iter().map(ValueSum::from).sum::<Result<ValueSum, _>>().unwrap(),
                 anchor,
                 Authorized {
                     proof: Proof(fake_proof),
diff --git a/src/value.rs b/src/value.rs
index 29712bbe..51d6f811 100644
--- a/src/value.rs
+++ b/src/value.rs
@@ -110,7 +110,13 @@ impl Add for ValueSum {
 
 impl<'a> Sum<&'a ValueSum> for Result<ValueSum, OverflowError> {
     fn sum<I: Iterator<Item = &'a ValueSum>>(iter: I) -> Self {
-        iter.fold(Ok(ValueSum(0)), |acc, cv| (acc? + *cv).ok_or(OverflowError))
+        iter.fold(Ok(ValueSum(0)), |acc, v| (acc? + *v).ok_or(OverflowError))
+    }
+}
+
+impl Sum<ValueSum> for Result<ValueSum, OverflowError> {
+    fn sum<I: Iterator<Item = ValueSum>>(iter: I) -> Self {
+        iter.fold(Ok(ValueSum(0)), |acc, v| (acc? + v).ok_or(OverflowError))
     }
 }
 
@@ -122,6 +128,12 @@ impl TryFrom<ValueSum> for i64 {
     }
 }
 
+impl From<NoteValue> for ValueSum {
+    fn from(v: NoteValue) -> Self {
+        ValueSum(v.0 as i128)
+    }
+}
+
 /// The blinding factor for a [`ValueCommitment`].
 #[derive(Clone, Debug)]
 pub struct ValueCommitTrapdoor(pallas::Scalar);