qraft-core 0.1.2

//! Aggregate expression helpers such as `count`, `sum`, and `avg`.

use crate::{
    alias::Aliased,
    expression::Expression,
    lower::LowerCtx,
    query::{LowerHaving, LowerProject, Star},
};

/// Lowers a value into an aggregate argument list.
pub trait LowerAggregate {
    fn lower_aggregate(self, ctx: &mut LowerCtx);
}

impl<E> LowerAggregate for E
where
    E: Expression,
{
    fn lower_aggregate(self, ctx: &mut LowerCtx) {
        let _ = self.lower(ctx);
    }
}

/// Lowers a value into the argument list accepted by `count`.
pub trait LowerCount {
    fn lower_count(self, ctx: &mut LowerCtx);
}

impl<E> LowerCount for E
where
    E: LowerAggregate,
{
    fn lower_count(self, ctx: &mut LowerCtx) {
        E::lower_aggregate(self, ctx);
    }
}

impl LowerCount for Star {
    fn lower_count(self, ctx: &mut LowerCtx) {
        let _ = ctx.lower_star();
    }
}

/// A typed `count(...)` expression.
#[derive(Debug, Clone, Copy)]
pub struct Count<E> {
    inner: E,
}

impl<E: LowerCount> Count<E> {
    fn lower(self, ctx: &mut LowerCtx) -> usize {
        self.inner.lower_count(ctx);
        ctx.lower_fn("count", 1)
    }
}

/// A single-argument aggregate function such as `max` or `avg`.
#[derive(Debug, Clone, Copy)]
pub struct Aggregate<E> {
    inner: E,
    name: &'static str,
}

impl<E: LowerAggregate> Aggregate<E> {
    fn lower(self, ctx: &mut LowerCtx) -> usize {
        self.inner.lower_aggregate(ctx);
        ctx.lower_fn(self.name, 1)
    }
}

/// Builds a `count(...)` expression.
pub fn count<E>(e: E) -> Count<E>
where
    E: LowerCount,
{
    Count { inner: e }
}

/// Builds a `max(...)` expression.
pub fn max<E>(e: E) -> Aggregate<E>
where
    E: LowerAggregate,
{
    Aggregate {
        inner: e,
        name: "max",
    }
}

/// Builds a `min(...)` expression.
pub fn min<E>(e: E) -> Aggregate<E>
where
    E: LowerAggregate,
{
    Aggregate {
        inner: e,
        name: "min",
    }
}

/// Builds a `sum(...)` expression.
pub fn sum<E>(e: E) -> Aggregate<E>
where
    E: LowerAggregate,
{
    Aggregate {
        inner: e,
        name: "sum",
    }
}

/// Builds an `avg(...)` expression.
pub fn avg<E>(e: E) -> Aggregate<E>
where
    E: LowerAggregate,
{
    Aggregate {
        inner: e,
        name: "avg",
    }
}

impl<E: LowerCount> LowerProject for Count<E> {
    fn lower_project(self, ctx: &mut LowerCtx) {
        self.lower(ctx);
    }
}

impl<E: LowerCount> LowerHaving for Count<E> {
    fn lower_having(self, ctx: &mut LowerCtx) {
        self.lower(ctx);
    }
}

impl<E: LowerCount> LowerProject for Aliased<Count<E>> {
    fn lower_project(self, ctx: &mut LowerCtx) {
        let inner = self.inner.lower(ctx);
        let _ = ctx.lower_alias(self.alias, inner);
    }
}

impl<E: LowerAggregate> LowerProject for Aliased<Aggregate<E>> {
    fn lower_project(self, ctx: &mut LowerCtx) {
        let inner = self.inner.lower(ctx);
        let _ = ctx.lower_alias(self.alias, inner);
    }
}

impl<E: LowerAggregate> LowerProject for Aggregate<E> {
    fn lower_project(self, ctx: &mut LowerCtx) {
        self.lower(ctx);
    }
}