rust-query 0.8.1

use std::marker::PhantomData;

use sea_query::ExprTrait;

use crate::{
    IntoSelect,
    select::{Cached, Cacher, ColumnImpl, Prepared, Row, Select, SelectImpl},
    value::{DynTypedExpr, EqTyp},
};

use super::{Expr, IntoExpr};

/// This is a combinator function that allows constructing single row optional queries.
///
/// ```
/// # use rust_query::IntoExpr;
/// # rust_query::private::doctest::get_txn(|txn| {
/// # use rust_query::optional;
/// let res = txn.query_one(optional(|row| {
///     let x = row.and(Some("test"));
///     let y = row.and(Some(42));
///     row.then_select((x, y))
/// }));
/// assert_eq!(res, Some(("test".to_owned(), 42)));
/// # });
/// ```
///
/// ```
/// # use rust_query::IntoExpr;
/// # rust_query::private::doctest::get_txn(|txn| {
/// # use rust_query::optional;
/// let res = txn.query_one(optional(|row| {
///     let x = row.and(Some("test"));
///     let y = row.and(None::<i64>);
///     row.then_select((x, y))
/// }));
/// assert_eq!(res, None);
/// # });
/// ```
pub fn optional<'outer, S, R>(
    f: impl for<'inner> FnOnce(&mut Optional<'outer, 'inner, S>) -> R,
) -> R {
    let mut optional = Optional {
        nulls: Vec::new(),
        _p: PhantomData,
        _p2: PhantomData,
    };
    f(&mut optional)
}

/// This is the argument type used by the [optional] combinator.
///
/// Joining more optional columns can be done with the [Optional::and] method.
/// Finally it is possible to return expressions or selections using [Optional::then] and [Optional::then_select].
pub struct Optional<'outer, 'inner, S> {
    nulls: Vec<DynTypedExpr>,
    _p: PhantomData<fn(&'inner ()) -> &'inner &'outer ()>,
    _p2: PhantomData<S>,
}

impl<'outer, 'inner, S> Optional<'outer, 'inner, S> {
    /// Join an optional column to the current row.
    ///
    /// If the joined column is [None], then the whole row
    /// will become [None], from that moment forward.
    ///
    /// ```
    /// # use rust_query::{private::doctest::*, optional};
    /// # get_txn(|txn| {
    /// let (total1, total2) = txn.query_one(optional(|row| {
    ///     let val1 = row.and(Some(100));
    ///     let val2 = row.and(Some(20));
    ///     let total = val1.add(val2);
    ///     let total1 = row.then(&total);
    ///
    ///     let _ = row.and(None::<i64>);
    ///     let total2 = row.then(total);
    ///
    ///      (total1, total2)
    /// }));
    /// assert_eq!(total1, Some(120));
    /// assert_eq!(total2, None);
    /// # });
    /// ```
    #[doc(alias = "join")]
    pub fn and<T: EqTyp>(
        &mut self,
        col: impl IntoExpr<'inner, S, Typ = Option<T>>,
    ) -> Expr<'inner, S, T> {
        let column = col.into_expr();
        self.nulls.push(DynTypedExpr::erase(column.is_none()));
        // `Expr::adhoc` is used here to reset `maybe_optional` to `true`.
        Expr::adhoc(move |b| column.inner.build_expr(b))
    }

    /// Return a [bool] column indicating whether the current row does not exists.
    ///
    /// ```
    /// # use rust_query::{private::doctest::*, optional};
    /// # get_txn(|txn| {
    /// let (none1, none2) = txn.query_one(optional(|row| {
    ///     let none1 = row.is_none();
    ///     let _ = row.and(None::<i64>);
    ///     let none2 = row.is_none();
    ///     (none1, none2)
    /// }));
    /// assert_eq!(none1, false);
    /// assert_eq!(none2, true);
    /// # });
    /// ```
    pub fn is_none(&self) -> Expr<'outer, S, bool> {
        let nulls = self.nulls.clone();
        Expr::adhoc(move |b| {
            nulls
                .iter()
                .map(|x| (x.func)(b))
                .reduce(|a, b| a.or(b))
                .unwrap_or(false.into())
        })
    }

    /// Return a [bool] column indicating whether the current row exists.
    ///
    /// ```
    /// # use rust_query::{private::doctest::*, optional};
    /// # get_txn(|txn| {
    /// let (some1, some2) = txn.query_one(optional(|row| {
    ///     let some1 = row.is_some();
    ///     let _ = row.and(None::<i64>);
    ///     let some2 = row.is_some();
    ///     (some1, some2)
    /// }));
    /// assert_eq!(some1, true);
    /// assert_eq!(some2, false);
    /// # });
    /// ```
    pub fn is_some(&self) -> Expr<'outer, S, bool> {
        self.is_none().not()
    }

    /// This is much like combining [Self::and] with [Self::then], but it
    /// allows returning an optional value without mutating the [Optional] row.
    ///
    ///  ```
    /// # use rust_query::{private::doctest::*, optional};
    /// # get_txn(|txn| {
    /// let (val1, val2) = txn.query_one(optional(|row| {
    ///     let val1 = row.and_then(None::<i64>);
    ///     let val2 = row.and_then(Some(1));
    ///     (val1, val2)
    /// }));
    /// assert_eq!(val1, None);
    /// assert_eq!(val2, Some(1));
    /// # });
    /// ```
    pub fn and_then<T: EqTyp>(
        &self,
        col: impl IntoExpr<'inner, S, Typ = Option<T>>,
    ) -> Expr<'outer, S, Option<T>> {
        const NULL: sea_query::Expr = sea_query::Expr::Keyword(sea_query::Keyword::Null);

        let col = col.into_expr().inner;
        let is_none = self.is_none().inner;
        Expr::adhoc(move |b| {
            sea_query::Expr::case(is_none.build_expr(b), NULL)
                .finally(col.build_expr(b))
                .into()
        })
    }

    /// Return [Some] column if the current row exists and [None] column otherwise.
    ///  ```
    /// # use rust_query::{private::doctest::*, optional};
    /// # get_txn(|txn| {
    /// let (val1, val2) = txn.query_one(optional(|row| {
    ///     let val1 = row.then(1);
    ///     let _ = row.and(None::<i64>);
    ///     let val2 = row.then(1);
    ///     (val1, val2)
    /// }));
    /// assert_eq!(val1, Some(1));
    /// assert_eq!(val2, None);
    /// # });
    /// ```
    pub fn then<T: EqTyp + 'outer>(
        &self,
        col: impl IntoExpr<'inner, S, Typ = T>,
    ) -> Expr<'outer, S, Option<T>> {
        self.and_then(Some(col))
    }

    /// Returns a [Select] with optional result. Useful for returning multiple values
    /// in a single [Option].
    ///
    ///  ```
    /// # use rust_query::{private::doctest::*, optional};
    /// # get_txn(|txn| {
    /// let pair = txn.query_one(optional(|row| {
    ///     let val1 = row.and(Some(1));
    ///     let val2 = row.and(Some(2));
    ///     row.then_select((val1, val2))
    /// }));
    /// assert_eq!(pair, Some((1, 2)));
    /// # });
    /// ```
    pub fn then_select<Out: 'static>(
        &self,
        d: impl IntoSelect<'inner, S, Out = Out>,
    ) -> Select<'outer, S, Option<Out>> {
        Select::new(OptionalImpl {
            inner: d.into_select().inner,
            is_some: ColumnImpl {
                expr: DynTypedExpr::erase(self.is_some()),
                _p: PhantomData,
            },
        })
    }
}

pub struct OptionalImpl<X> {
    inner: X,
    is_some: ColumnImpl<bool>,
}

impl<X: SelectImpl> SelectImpl for OptionalImpl<X> {
    type Out = Option<X::Out>;
    type Prepared = OptionalPrepared<X::Prepared>;

    fn prepare(self, cacher: &mut Cacher) -> Self::Prepared {
        OptionalPrepared {
            is_some: self.is_some.prepare(cacher),
            inner: self.inner.prepare(cacher),
        }
    }
}

pub struct OptionalPrepared<X> {
    inner: X,
    is_some: Cached<bool>,
}

impl<X: Prepared> Prepared for OptionalPrepared<X> {
    type Out = Option<X::Out>;

    fn call(&mut self, row: Row<'_>) -> Self::Out {
        if row.get(self.is_some) {
            Some(self.inner.call(row))
        } else {
            None
        }
    }
}