Advanced rebalancing

def rebalance_once(position_token_id: str, tick_lower: int,
                   tick_upper: int, deadline_offset_blocks: int):
    position = dex.get_owned_position(position_token_id)
    if position is None or position.state is None or position.state.liquidity == 0:
        raise ValueError("position unavailable or empty")
    plan = dex.plan_rebalance(
        pool_key=position.pool_key,
        position_token_id=position_token_id,
        tick_lower=tick_lower,
        tick_upper=tick_upper,
        max_funding0=0,
        max_funding1=0,
    )
    return dex.rebalance_position(
        plan=plan,
        position_record=position.record,
        deadline_offset_blocks=deadline_offset_blocks,
    ).delegate(wait=True)

cargo add snarkvm --no-default-features --features console,ledger,synthesizer,utilities
cargo add reqwest --features blocking,json
cargo add serde_json

use reqwest::blocking::Client;
use snarkvm::{
    ledger::{query::QueryTrait, store::ConsensusStorage},
    prelude::*,
};
use std::{
    error::Error, fs::OpenOptions, path::Path,
    str::FromStr, time::Duration,
};

type SubmitResult = Result<String, Box<dyn Error>>;

pub struct Submitter<'a, N: Network, C: ConsensusStorage<N>> {
    pub vm: &'a VM<N, C>,
    pub key: &'a PrivateKey<N>,
    pub query: &'a dyn QueryTrait<N>,
    pub broadcast_url: &'a str,
    pub priority_fee_microcredits: u64,
}

impl<N: Network, C: ConsensusStorage<N>> Submitter<'_, N, C> {
    pub fn build_and_broadcast<R: Rng + CryptoRng>(
        &self,
        entrypoint: (&str, &str),
        encoded_inputs: &[String],
        fee_record: Record<N, Plaintext<N>>,
        checkpoint: &Path,
        rng: &mut R,
    ) -> SubmitResult {
        let http = Client::builder()
            .timeout(Duration::from_secs(60))
            .retry(reqwest::retry::never())
            .redirect(reqwest::redirect::Policy::none())
            .build()?;
        let inputs = encoded_inputs.iter()
            .map(|input| Value::<N>::from_str(input))
            .collect::<Result<Vec<_>, _>>()?;
        let transaction = self.vm.execute(
            self.key, entrypoint, inputs.into_iter(), Some(fee_record),
            self.priority_fee_microcredits, Some(self.query), rng,
        )?;
        let id = transaction.id().to_string();

        // Save before POST. An HTTP error leaves submission status unresolved.
        let mut file = OpenOptions::new().write(true).create_new(true)
            .open(checkpoint)?;
        serde_json::to_writer(&mut file, &transaction)?;
        file.sync_all()?;
        eprintln!("Broadcasting {id}; saved at {}", checkpoint.display());

        let returned_id: String = http.post(self.broadcast_url)
            .json(&transaction)
            .send()?
            .error_for_status()?
            .json()?;
        if returned_id != id {
            return Err(format!("broadcast ID mismatch; reconcile {id}").into());
        }
        Ok(id)
    }
}

pub fn rebalance_plain<N: Network, C: ConsensusStorage<N>, R: Rng + CryptoRng>(
    submitter: &Submitter<N, C>,
    encoded_inputs: [&str; 6],
    fee_record: Record<N, Plaintext<N>>,
    checkpoint: &Path,
    rng: &mut R,
) -> SubmitResult {
    let inputs = encoded_inputs.map(str::to_owned);
    submitter.build_and_broadcast(
        ("shield_swap_rebalance_router.aleo", "rebalance_plain_plain_none"),
        &inputs, fee_record, checkpoint, rng,
    )
}

from aleo_shield_swap.position_math import amounts_for_liquidity
from aleo_shield_swap.tick_math import get_sqrt_price_at_tick_x128

def decrease_for_rebalance(position_token_id: str,
                           acceptable_tick_low: int, acceptable_tick_high: int):
    if acceptable_tick_low > acceptable_tick_high:
        raise ValueError("inverted price bounds")
    p = dex.get_owned_position(position_token_id)
    if p is None or p.state is None or p.state.liquidity == 0:
        raise ValueError("position unavailable or empty")
    lower = get_sqrt_price_at_tick_x128(p.tick_lower)
    upper = get_sqrt_price_at_tick_x128(p.tick_upper)
    amount0_min, _ = amounts_for_liquidity(
        get_sqrt_price_at_tick_x128(acceptable_tick_high), lower, upper,
        p.state.liquidity, round_up=False,
    )
    _, amount1_min = amounts_for_liquidity(
        get_sqrt_price_at_tick_x128(acceptable_tick_low), lower, upper,
        p.state.liquidity, round_up=False,
    )
    return dex.decrease_liquidity(
        pool_key=p.pool_key,
        position_record=p.record,
        liquidity_to_remove=p.state.liquidity,
        amount0_min=amount0_min,
        amount1_min=amount1_min,
    ).delegate(wait=True)

pub fn decrease_for_rebalance<N: Network, C: ConsensusStorage<N>, R: Rng + CryptoRng>(
    submitter: &Submitter<N, C>,
    nft: &str,
    liquidity_to_remove: u128,
    amount0_min: u128,
    amount1_min: u128,
    fee_record: Record<N, Plaintext<N>>,
    checkpoint: &Path,
    rng: &mut R,
) -> SubmitResult {
    let inputs = [
        nft.to_owned(),
        format!("{liquidity_to_remove}u128"),
        format!("{amount0_min}u128"),
        format!("{amount1_min}u128"),
    ];

    submitter.build_and_broadcast(
        ("shield_swap.aleo", "decrease_liquidity"),
        &inputs, fee_record, checkpoint, rng,
    )
}

def collect_closed_position(position_token_id: str):
    p = dex.get_owned_position(position_token_id)
    if p is None or p.state is None or p.state.liquidity != 0:
        raise ValueError("full decrease not confirmed")
    if p.state.tokens_owed0 == 0 and p.state.tokens_owed1 == 0:
        return None
    return dex.collect(
        pool_key=p.pool_key,
        position_record=p.record,
        amount0_requested=p.state.tokens_owed0,
        amount1_requested=p.state.tokens_owed1,
    ).delegate(wait=True)

pub fn collect_closed_position<N: Network, C: ConsensusStorage<N>, R: Rng + CryptoRng>(
    submitter: &Submitter<N, C>,
    nft: &str,
    liquidity: u128,
    owed: [u128; 2],
    token_ids: [&str; 2],
    amm_proofs: [&str; 2], // owner, withdrawal
    receiver_proofs: [Option<&str>; 2], // underlying token0, token1
    fee_record: Record<N, Plaintext<N>>,
    checkpoint: &Path,
    rng: &mut R,
) -> Result<Option<String>, Box<dyn Error>> {
    if liquidity != 0 {
        return Err("full decrease not confirmed".into());
    }
    if owed == [0, 0] {
        return Ok(None);
    }
    let entrypoint = match receiver_proofs {
        [None, None] => ("shield_swap.aleo", "collect"),
        [Some(_), None] => ("shield_swap_lp_router.aleo", "collect_to_wrapped_arc20"),
        [None, Some(_)] => ("shield_swap_lp_router.aleo", "collect_to_arc20_wrapped"),
        [Some(_), Some(_)] => ("shield_swap_lp_router.aleo", "collect_to_wrapped_wrapped"),
    };
    let mut inputs = vec![
        nft.to_owned(), format!("{}u128", owed[0]), format!("{}u128", owed[1]),
        token_ids[0].to_owned(), token_ids[1].to_owned(),
        amm_proofs[0].to_owned(), amm_proofs[1].to_owned(),
    ];
    // Wrapped receiver bundles follow both AMM bundles, in token order.
    inputs.extend(receiver_proofs.into_iter().flatten().map(str::to_owned));
    let id = submitter.build_and_broadcast(
        entrypoint, &inputs, fee_record, checkpoint, rng,
    )?;
    Ok(Some(id))
}

def mint_replacement(pool_key: str, tick_lower: int, tick_upper: int,
                     amount0_desired: int, amount1_desired: int,
                     amount0_min: int, amount1_min: int,
                     recipient: str, withdrawal: str):
    return dex.mint(
        pool_key=pool_key, tick_lower=tick_lower, tick_upper=tick_upper,
        amount0_desired=amount0_desired, amount1_desired=amount1_desired,
        amount0_min=amount0_min, amount1_min=amount1_min,
        recipient=recipient, withdrawal=withdrawal,
    ).delegate(wait=True)

pub fn mint_replacement<N: Network, C: ConsensusStorage<N>, R: Rng + CryptoRng>(
    submitter: &Submitter<N, C>,
    nonce: &str,
    funding_records: [&str; 2],
    sender_proofs: [Option<&str>; 2], // underlying token0, token1
    recipient: &str,
    withdrawal: &str,
    request: &str,
    token_ids: [&str; 2],
    amm_proofs: [&str; 3], // signer, recipient, withdrawal
    fee_record: Record<N, Plaintext<N>>,
    checkpoint: &Path,
    rng: &mut R,
) -> SubmitResult {
    let entrypoint = match sender_proofs {
        [None, None] => ("shield_swap.aleo", "mint"),
        [Some(_), None] => ("shield_swap_lp_router.aleo", "mint_from_wrapped_arc20"),
        [None, Some(_)] => ("shield_swap_lp_router.aleo", "mint_from_arc20_wrapped"),
        [Some(_), Some(_)] => ("shield_swap_lp_router.aleo", "mint_from_wrapped_wrapped"),
    };
    let mut inputs = vec![nonce.to_owned()];
    for (record, proof) in funding_records.into_iter().zip(sender_proofs) {
        inputs.push(record.to_owned());
        // A wrapped side's sender bundle immediately follows its record.
        if let Some(bundle) = proof {
            inputs.push(bundle.to_owned());
        }
    }
    inputs.extend([
        recipient.to_owned(), withdrawal.to_owned(), request.to_owned(),
        token_ids[0].to_owned(), token_ids[1].to_owned(),
    ]);
    inputs.extend(amm_proofs.map(str::to_owned));
    submitter.build_and_broadcast(
        entrypoint, &inputs, fee_record, checkpoint, rng,
    )
}