riftenlabs-indexer/src/db/cauldron/pool.rs

// Copyright (C) 2024 Riften Labs AS
//
// This software is licensed under the GNU Affero General Public License (AGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.html

use std::{
    collections::{HashMap, VecDeque},
    sync::atomic::AtomicI64,
};

use anyhow::Result;
use bitcoin_hashes::hex::{FromHex, ToHex};
use log::{debug, info, warn};
use riftenlabs_defi::{cauldron::ParsedContract, chainutil::OutPointHash};
use rusqlite::{params, Connection, Row};
use rust_decimal::prelude::Zero;

use crate::rpc::apy::PoolSnapshot;

pub fn create_table(conn: &Connection) {
    conn.execute(
        "
        CREATE TABLE pool (
            creation_utxo TEXT PRIMARY KEY REFERENCES utxo_funding(new_utxo_hash) ON DELETE CASCADE,
            owner_pkh TEXT NOT NULL,
            token_id TEXT NOT NULL,
            withdrawn_in_utxo TEXT REFERENCES utxo_spending(spent_utxo_hash) ON DELETE SET NULL
        )",
        [],
    )
    .expect("failed to create table pool");

    conn.execute(
        "CREATE TABLE pool_history_entry (
            utxo TEXT PRIMARY KEY REFERENCES utxo_funding(new_utxo_hash) ON DELETE CASCADE,
            pool TEXT REFERENCES pool(creation_utxo) ON DELETE CASCADE,
            txid TEXT REFERENCES tx(txid) ON DELETE CASCADE,
            tx_pos TEXT NOT NULL,
            mtp_timestamp BIGINT,
            first_seen_timestamp BIGINT,
            sequence BIGINT NOT NULL,
            sats BIGINT,
            token_amount BIGINT
        )",
        [],
    )
    .expect("failed to create table pool_history_entry");
}

fn get_pool_by_utxo(conn: &Connection, utxo_hash: &OutPointHash) -> Result<Option<OutPointHash>> {
    let mut stmt = conn.prepare("SELECT pool FROM pool_history_entry WHERE utxo = ?")?;

    let mut row = stmt.query([utxo_hash.to_hex()])?;
    let utxo_hex: Option<String> = row.next()?.map(|r| r.get(0).unwrap());

    match utxo_hex {
        Some(utxo) => Ok(Some(
            OutPointHash::from_hex(&utxo).expect("invalid original_utxo utxo in db"),
        )),
        None => Ok(None),
    }
}

pub fn flag_as_withdrawn(
    conn: &Connection,
    pool_utxo: &OutPointHash,
    cauldron: &ParsedContract,
) -> Result<()> {
    conn.execute(
        "UPDATE pool SET withdrawn_in_utxo = ? WHERE creation_utxo = ?",
        params![cauldron.spent_utxo_hash.to_hex(), pool_utxo.to_hex()],
    )
    .map_err(|e| anyhow::anyhow!("failed flag pool as withdrawn. Original error: {:?}", e))?;

    Ok(())
}

pub fn insert_new_pool(conn: &Connection, cauldron: &ParsedContract) -> Result<()> {
    // or replace, as it could have been added in mempool, then block
    conn.execute(
        "INSERT OR REPLACE INTO pool (creation_utxo, owner_pkh, token_id, withdrawn_in_utxo) VALUES (?, ?, ?, ?)",
        params![
            cauldron.new_utxo_hash.expect("outpoint hash for new pool missing").to_hex(),
            cauldron.pkh.to_hex(),
            cauldron.token_id.expect("token id for new pool missing").to_hex(),
            None::<String>,
        ]
    ).map_err(|e| {
        anyhow::anyhow!(
            "failed to insert new pool. Original error: {:?}",
            e
        )
    })?;

    Ok(())
}

/// Next sequence number in the `pool_history_entry` table
static NEXT_SEQUENCE: AtomicI64 = AtomicI64::new(-10);

pub fn initialize_seq(conn: &Connection) {
    let mut s = conn
        .prepare("SELECT IFNULL(MAX(sequence), 0) + 1 FROM pool_history_entry")
        .unwrap();
    let mut q = s.query(params![]).unwrap();
    let seq: i64 = q.next().unwrap().unwrap().get(0).unwrap();
    NEXT_SEQUENCE.store(seq, std::sync::atomic::Ordering::SeqCst);
}

#[allow(dead_code)] // for unit tests
pub fn dummy_init_seq() {
    if NEXT_SEQUENCE.load(std::sync::atomic::Ordering::SeqCst) < 0 {
        NEXT_SEQUENCE.store(42, std::sync::atomic::Ordering::SeqCst);
    }
}

pub fn insert_pool_history_entry(
    conn: &Connection,
    pool: &OutPointHash,
    cauldron: &ParsedContract,
    mtp_timestamp: Option<u64>,
    first_seen_timestamp: Option<u64>,
) -> Result<()> {
    let next_seq = NEXT_SEQUENCE.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
    assert!(next_seq >= 0);

    conn.execute(
        "INSERT INTO pool_history_entry (utxo, pool, txid, tx_pos, mtp_timestamp, first_seen_timestamp, sequence, sats, token_amount)
         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
         ON CONFLICT(utxo) DO UPDATE SET
            pool = excluded.pool,
            txid = excluded.txid,
            tx_pos = excluded.tx_pos,
            mtp_timestamp = COALESCE(excluded.mtp_timestamp, pool_history_entry.mtp_timestamp),
            first_seen_timestamp = COALESCE(excluded.first_seen_timestamp, pool_history_entry.first_seen_timestamp),
            sequence = excluded.sequence",
        params![
            cauldron
                .new_utxo_hash
                .expect("utxo hash on new pool history entry")
                .to_hex(),
            pool.to_hex(),
            cauldron
                .new_utxo_txid
                .expect("txid of new pool history entry")
                .to_hex(),
            cauldron
                .new_utxo_n
                .expect("utxo index of new pool history entry"),
            mtp_timestamp,
            first_seen_timestamp,
            next_seq,
            cauldron.sats,
            cauldron.token_amount
        ],
    )
    .map_err(|e| {
        anyhow::anyhow!(
            "failed to insert pool history entry. Original error: {:?}",
            e
        )
    })?;
    Ok(())
}

pub fn update_pool_history(
    conn: &Connection,
    cauldrons: Vec<ParsedContract>,
    mtp_timestamp: Option<u64>,
    first_seen_timestamp: Option<u64>,
) -> Result<()> {
    if cauldrons.is_empty() {
        return Ok(());
    }

    let mut queue = VecDeque::from(cauldrons);

    while let Some(current) = queue.pop_front() {
        let (pool_utxo, is_new) = match get_pool_by_utxo(conn, &current.spent_utxo_hash)? {
            Some(c) => (c, false),
            None => {
                // Not in our database. Check if child of another interaction in current batch.
                let has_parent = queue
                    .iter()
                    .any(|parent| Some(current.spent_utxo_hash) == parent.new_utxo_hash);

                if has_parent {
                    // Was child of current batch.
                    // Process later (after parent).
                    queue.push_back(current);
                    continue;
                }

                if let Some(utxo) = current.new_utxo_hash {
                    // Not in existing pool or child of a cauldron in current batch.
                    // This is a new pool.
                    assert!(!current.is_withdrawn);
                    (utxo, true)
                } else {
                    // This is a new pool that is immediately withdrawn. Just ignore.
                    assert!(current.is_withdrawn);
                    continue;
                }
            }
        };

        if current.is_withdrawn {
            info!(
                "LP {} withdrawn in {}",
                pool_utxo.to_hex(),
                current.spent_utxo_hash.to_hex()
            );
            flag_as_withdrawn(conn, &pool_utxo, &current)?;
        } else if is_new {
            info!(
                "Cauldron LP created in tx {}",
                current
                    .new_utxo_txid
                    .expect("expected txid in new LP")
                    .to_hex()
            );
            insert_new_pool(conn, &current)?;
            insert_pool_history_entry(
                conn,
                &pool_utxo,
                &current,
                mtp_timestamp,
                first_seen_timestamp,
            )?;
        } else {
            debug!("New entry for pool {}", pool_utxo.to_hex());
            insert_pool_history_entry(
                conn,
                &pool_utxo,
                &current,
                mtp_timestamp,
                first_seen_timestamp,
            )?;
        }
    }

    Ok(())
}

/// Strategy for which pool state to select when two pool states are available for a timestamp
#[derive(PartialEq, Eq)]
enum SnapshotSelection {
    UseBefore,
    UseAfter,
}

/// Get pools and their state nearest to given timestamp.
fn get_nearest_entries(
    conn: &Connection,
    timestamp: i64,
    token_id: Option<&str>,
    owner_pkh: Option<&str>,
    resolution: SnapshotSelection,
) -> Result<HashMap<String, PoolSnapshot>> {
    let extra_filters = match (token_id.is_some(), owner_pkh.is_some()) {
        (true, true) => "token_id = ?2 AND owner_pkh = ?3",
        (true, false) => "token_id = ?2",
        (false, true) => "owner_pkh = ?2",
        (false, false) => "1=1",
    };

    let params = match (token_id.is_some(), owner_pkh.is_some()) {
        (true, true) => params![timestamp, token_id, owner_pkh],
        (true, false) => params![timestamp, token_id],
        (false, true) => params![timestamp, owner_pkh],
        (false, false) => params![timestamp],
    };

    let ts = "COALESCE(first_seen_timestamp, mtp_timestamp)";

    let preferred = match resolution {
        SnapshotSelection::UseBefore => "ts ASC", // Prefer the lower timestamp first
        SnapshotSelection::UseAfter => "ts DESC", // Prefer the higher timestamp first
    };

    let query = format!(
        "WITH NearestEntries AS (
            -- Get the closest lower or equal to the timestamp (MAX for <= timestamp)
            SELECT phe.pool, phe.sats, phe.token_amount, {ts} AS ts, phe.sequence
            FROM pool_history_entry phe
            WHERE {ts} <= ?1
            AND phe.pool IN (
                SELECT creation_utxo
                FROM pool
                WHERE {extra_filters}
            )
            GROUP BY phe.pool
            HAVING MAX({ts})

            UNION ALL

            -- Get the closest greater than the timestamp (MIN for > timestamp)
            SELECT phe.pool, phe.sats, phe.token_amount, {ts} AS ts, phe.sequence
            FROM pool_history_entry phe
            WHERE {ts} > ?1
            AND phe.pool IN (
                SELECT creation_utxo
                FROM pool
                WHERE {extra_filters}
            )
            GROUP BY phe.pool
            HAVING MIN({ts})
        )
        SELECT * FROM NearestEntries
        ORDER BY {preferred}",
    );

    let mut stmt = conn.prepare(&query)?;
    let mut rows = stmt.query(params)?;

    let from_row = |row: &Row<'_>| -> Result<PoolSnapshot> {
        Ok(PoolSnapshot {
            pool_id: row.get(0)?,
            sats: row.get(1)?,
            token_amount: row.get(2)?,
            timestamp: row.get(3)?,
        })
    };

    let mut pools: HashMap<String, PoolSnapshot> = HashMap::default();

    while let Some(row) = rows.next()? {
        let pool_snapshot = from_row(row)?;
        if !pools.contains_key(&pool_snapshot.pool_id) {
            let existed = pools.insert(pool_snapshot.pool_id.clone(), pool_snapshot.clone());
            assert!(existed.is_none());
        }
    }

    Ok(pools)
}

/// Returns
pub fn get_pool_period_snapshot(
    conn: &Connection,
    token_id: Option<&str>,
    owner_pkh: Option<&str>,
    start: i64,
    end: i64,
) -> Result<Vec<(PoolSnapshot, PoolSnapshot)>> {
    let pools_start = get_nearest_entries(
        conn,
        start,
        token_id,
        owner_pkh,
        SnapshotSelection::UseBefore,
    )?;
    let mut pools_end =
        get_nearest_entries(conn, end, token_id, owner_pkh, SnapshotSelection::UseAfter)?;

    let mut pools: Vec<(PoolSnapshot, PoolSnapshot)> = Vec::default();
    for (start_pool_id, start_pool) in pools_start {
        let end_pool = match pools_end.remove(&start_pool_id) {
            Some(end) => end,
            None => {
                warn!("Found no 'end pool' for {}", start_pool_id);
                continue;
            }
        };

        let duration = end_pool.timestamp.saturating_sub(start_pool.timestamp);
        if !duration.is_zero() {
            pools.push((start_pool, end_pool))
        }
    }

    if !pools_end.is_empty() {
        warn!("Found {} end pools not in start pools", pools_end.len());
    }

    Ok(pools)
}