improve preamble to be faster and always allow a short

crates/hyperdrive-math/src/test_utils/preamble.rs

@@ -23,7 +23,7 @@ pub async fn initialize_pool_with_random_state(
     celine: &mut Agent<ChainClient<LocalWallet>, ChaCha8Rng>,
 ) -> Result<()> {
     // Get random pool liquidity & fixed rate.
-    let pool_initial_contribution = rng.gen_range(fixed!(1_000e18)..=fixed!(1_000_000_000e18));
+    let pool_initial_contribution = rng.gen_range(fixed!(10_000e18)..=fixed!(1_000_000_000e18));
     let fixed_rate = rng.gen_range(fixed!(0.01e18)..=fixed!(0.1e18));
 
     // Alice initializes the pool.
@@ -32,54 +32,72 @@ pub async fn initialize_pool_with_random_state(
         .initialize(fixed_rate, pool_initial_contribution, None)
         .await?;
 
-    // Advance the time for over a term and make trades in some of the checkpoints.
-    let mut time_remaining = alice.get_config().position_duration;
-    while time_remaining > uint256!(0) {
-        let mut state = alice.get_state().await?;
-        let min_txn = state.minimum_transaction_amount();
-
+    // Determine a bounded random amount of trade sequences.
+    let mut state = alice.get_state().await?;
+    let total_time = fixed!(1.1e18) * state.position_duration();
+    let num_steps = rng.gen_range(fixed!(1e18)..=fixed!(3e18));
+    let time_delta = total_time / num_steps;
+
+    // Advance the time for over a position term and make trades in some of the
+    // checkpoints.
+    let min_txn = state.minimum_transaction_amount();
+    let mut time_remaining = total_time.change_type::<I256>()?;
+    while time_remaining > fixed!(0) {
+        state = alice.get_state().await?;
         // Bob opens a long.
-        let max_long = get_max_long(state, None)?;
+        let max_long = get_max_long(&state, None)?;
         let long_amount = rng.gen_range(min_txn..=max_long);
-        bob.fund(long_amount + fixed!(10e18)).await?; // Fund a little extra for slippage.
+        // Fund extra for slippage.
+        bob.fund(long_amount + fixed!(10e18)).await?;
         bob.open_long(long_amount, None, None).await?;
 
         // Celine opens a short.
         state = alice.get_state().await?;
-
-        // Open a short.
-        let max_short = get_max_short(state, None)?;
-        let short_amount = rng.gen_range(min_txn..=max_short);
-        celine.fund(short_amount + fixed!(10e18)).await?; // Fund a little extra for slippage.
-        celine.open_short(short_amount, None, None).await?;
-
+        let max_short = state.absolute_max_short_guess()?; // Less accurate but faster to compute.
+        let short_bond_amount = if max_short == min_txn {
+            min_txn
+        } else {
+            rng.gen_range(min_txn..=max_short)
+        };
+        let user_deposit_shares =
+            state.calculate_open_short(short_bond_amount, state.vault_share_price())?;
+        // Fund extra for slippage.
+        celine
+            .fund(user_deposit_shares * state.vault_share_price() + fixed!(10e18))
+            .await?;
+        celine.open_short(short_bond_amount, None, None).await?;
         // Advance the time and mint all of the intermediate checkpoints.
-        let multiplier = rng.gen_range(fixed!(10e18)..=fixed!(100e18));
-        let delta = FixedPoint::from(time_remaining)
-            .min(FixedPoint::from(alice.get_config().checkpoint_duration) * multiplier);
-        time_remaining -= U256::from(delta);
         let variable_rate = rng.gen_range(fixed!(0.01e18)..=fixed!(0.1e18));
-        alice.advance_time(variable_rate, delta).await?;
+        alice.advance_time(variable_rate, time_delta).await?;
+        time_remaining -= time_delta.change_type::<I256>()?;
     }
 
-    // Mint a checkpoint to close any matured positions.
+    // Add some liquidity again to make sure future bots can make trades.
+    let state = alice.get_state().await?;
+    let min_share_reserves = state
+        .calculate_min_share_reserves_given_exposure()?
+        .change_type::<U256>()?;
+    let positive_share_adjustment =
+        FixedPoint::<U256>::try_from(state.share_adjustment().min(0.into()))?;
+    let min_base = (min_share_reserves + positive_share_adjustment) * state.vault_share_price()
+        + state.long_exposure();
+    let liquidity_amount = rng.gen_range(min_base..=min_base + fixed!(100_000e18));
+    alice.fund(liquidity_amount).await?;
+    alice.add_liquidity(liquidity_amount, None).await?;
+    // Advance time logner than a position duration & checkpoint one last time
+    // to ensure all positions are closed.
+    let variable_rate = rng.gen_range(fixed!(0.01e18)..=fixed!(0.1e18));
     alice
-        .checkpoint(alice.latest_checkpoint().await?, uint256!(0), None)
+        .advance_time(variable_rate, state.position_duration() * fixed!(2e18))
         .await?;
-
-    // Add some liquidity again to make sure future bots can make trades.
-    let liquidity_amount = rng.gen_range(fixed!(1_000e18)..=fixed!(100_000_000e18));
-    let exposure = alice.get_state().await?.long_exposure();
-    alice.fund(liquidity_amount + exposure).await?;
     alice
-        .add_liquidity(liquidity_amount + exposure, None)
+        .checkpoint(alice.latest_checkpoint().await?, uint256!(0), None)
         .await?;
-
     Ok(())
 }
 
 /// Conservative and safe estimate of the maximum long.
-fn get_max_long(state: State, maybe_max_num_tries: Option<usize>) -> Result<FixedPoint<U256>> {
+fn get_max_long(state: &State, maybe_max_num_tries: Option<usize>) -> Result<FixedPoint<U256>> {
     let max_num_tries = maybe_max_num_tries.unwrap_or(10);
     let checkpoint_exposure = I256::from(0);
 
@@ -127,78 +145,10 @@ fn get_max_long(state: State, maybe_max_num_tries: Option<usize>) -> Result<Fixe
     Ok(max_long)
 }
 
-/// Conservative and safe estimate of the maximum short.
-pub fn get_max_short(state: State, maybe_max_num_tries: Option<usize>) -> Result<FixedPoint<U256>> {
-    let max_num_tries = maybe_max_num_tries.unwrap_or(10);
-
-    // We linearly interpolate between the current spot price and the minimum
-    // price that the pool can support. This is a conservative estimate of
-    // the short's realized price.
-    let conservative_price = {
-        // We estimate the minimum price that short will pay by a
-        // weighted average of the spot price and the minimum possible
-        // spot price the pool can quote. We choose the weights so that this
-        // is an underestimate of the worst case realized price.
-        let spot_price = state.calculate_spot_price_down()?;
-        let min_price = state.calculate_min_spot_price()?;
-
-        // Calculate the linear interpolation.
-        let weight = fixed!(1e18).pow(fixed!(1e18) - state.time_stretch())?;
-        spot_price * (fixed!(1e18) - weight) + min_price * weight
-    };
-
-    // Compute the max short.
-    let mut max_short = match panic::catch_unwind(|| {
-        state.calculate_max_short(
-            U256::from(U128::MAX),
-            state.vault_share_price(),
-            Some(conservative_price),
-            Some(5),
-        )
-    }) {
-        Ok(max_short_no_panic) => match max_short_no_panic {
-            Ok(max_short) => max_short,
-            Err(_) => state.share_reserves() / state.vault_share_price() * fixed!(10e18),
-        },
-        Err(_) => state.share_reserves() / state.vault_share_price() * fixed!(10e18),
-    };
-    let mut num_tries = 0;
-    let mut success = false;
-    while !success {
-        max_short = match panic::catch_unwind(|| {
-            state.calculate_open_short(max_short, state.vault_share_price())
-        }) {
-            Ok(short_result_no_panic) => match short_result_no_panic {
-                Ok(_) => {
-                    success = true;
-                    max_short
-                }
-                Err(_) => max_short / fixed!(10e18),
-            },
-            Err(_) => max_short / fixed!(10e18),
-        };
-        if max_short < state.minimum_transaction_amount() {
-            return Err(eyre!(
-                "max_short={} was less than minimum_transaction_amount={}.",
-                max_short,
-                state.minimum_transaction_amount()
-            ));
-        }
-        num_tries += 1;
-        if num_tries > max_num_tries {
-            return Err(eyre!(
-                "Failed to find a max short. Last attempted value was {}",
-                max_short,
-            ));
-        }
-    }
-    Ok(max_short)
-}
-
 #[cfg(test)]
 mod tests {
     use fixedpointmath::fixed;
-    use hyperdrive_test_utils::{chain::TestChain, constants::FUZZ_RUNS};
+    use hyperdrive_test_utils::{chain::TestChain, constants::SLOW_FUZZ_RUNS};
     use rand::{thread_rng, Rng, SeedableRng};
     use rand_chacha::ChaCha8Rng;
 
@@ -222,7 +172,7 @@ mod tests {
         let mut bob = chain.bob().await?;
         let mut celine = chain.celine().await?;
 
-        for _ in 0..*FUZZ_RUNS {
+        for _ in 0..*SLOW_FUZZ_RUNS {
             // Snapshot the chain.
             let id = chain.snapshot().await?;
 
@@ -260,30 +210,22 @@ mod tests {
         let mut alice = chain.alice().await?;
         let mut bob = chain.bob().await?;
         let mut celine = chain.celine().await?;
-        for _ in 0..*FUZZ_RUNS {
+        for _ in 0..*SLOW_FUZZ_RUNS {
             // Snapshot the chain.
             let id = chain.snapshot().await?;
             // Run the preamble.
             initialize_pool_with_random_state(&mut rng, &mut alice, &mut bob, &mut celine).await?;
             // Get state.
             let state = alice.get_state().await?;
             // Check that a short is possible.
-            if state.effective_share_reserves()?
-                < state
-                    .calculate_min_share_reserves_given_exposure()?
-                    .change_type::<U256>()?
-            {
-                chain.revert(id).await?;
-                alice.reset(Default::default()).await?;
-                bob.reset(Default::default()).await?;
-                celine.reset(Default::default()).await?;
-                continue;
-            }
+            assert!(state
+                .solvency_after_short(state.minimum_transaction_amount())
+                .is_ok());
             // Open the max short.
-            let max_short = bob.calculate_max_short(None).await?;
+            let max_short = celine.calculate_max_short(None).await?;
             assert!(max_short >= state.minimum_transaction_amount());
-            bob.fund(max_short + fixed!(10e18)).await?;
-            bob.open_short(max_short, None, None).await?;
+            celine.fund(max_short + fixed!(10e18)).await?;
+            celine.open_short(max_short, None, None).await?;
             // Reset the chain & agents.
             chain.revert(id).await?;
             alice.reset(Default::default()).await?;