Decodal/crates/decodal-core/src/constraints.rs

use alloc::{string::String, vec::Vec};

use crate::{
    Diagnostic, DiagnosticKind, Span,
    ast::CompareOp,
    runtime::{Constraint, LiteralValue, PrimitiveType},
};

pub fn normalize_constraints(
    constraints: Vec<Constraint>,
    span: Span,
) -> crate::Result<Vec<Constraint>> {
    let mut primitive = None;
    let mut lower: Option<Bound> = None;
    let mut upper: Option<Bound> = None;
    let mut rest = Vec::new();

    for constraint in constraints {
        match constraint {
            Constraint::Type(next) => match primitive {
                Some(current) if current != next => {
                    return Err(Diagnostic::new(
                        DiagnosticKind::Conflict,
                        span,
                        "primitive type constraints conflict",
                    ));
                }
                Some(_) => {}
                None => primitive = Some(next),
            },
            Constraint::Compare(op, value) => {
                let number = Number::from_literal(&value).ok_or_else(|| {
                    Diagnostic::new(
                        DiagnosticKind::TypeMismatch,
                        span,
                        "comparison constraints require numeric literals",
                    )
                })?;
                match op {
                    CompareOp::Gt => merge_lower(&mut lower, Bound::new(number, false)),
                    CompareOp::Gte => merge_lower(&mut lower, Bound::new(number, true)),
                    CompareOp::Lt => merge_upper(&mut upper, Bound::new(number, false)),
                    CompareOp::Lte => merge_upper(&mut upper, Bound::new(number, true)),
                    CompareOp::Eq => {
                        merge_lower(&mut lower, Bound::new(number, true));
                        merge_upper(&mut upper, Bound::new(number, true));
                    }
                }
            }
            Constraint::Regex(pattern) => rest.push(Constraint::Regex(pattern)),
            Constraint::BuiltinPredicate(name) => rest.push(Constraint::BuiltinPredicate(name)),
        }
    }

    if matches!(primitive, Some(PrimitiveType::String | PrimitiveType::Bool))
        && (lower.is_some() || upper.is_some())
    {
        return Err(Diagnostic::new(
            DiagnosticKind::Conflict,
            span,
            "numeric comparison constraints conflict with non-numeric primitive type",
        ));
    }

    if primitive == Some(PrimitiveType::Int)
        && lower
            .iter()
            .chain(upper.iter())
            .any(|bound| !matches!(bound.number, Number::Int(_)))
    {
        return Err(Diagnostic::new(
            DiagnosticKind::Conflict,
            span,
            "Int comparison constraints must use integer literals",
        ));
    }

    ensure_bounds_non_empty(primitive, lower, upper, span)?;

    let mut normalized = Vec::new();
    if let Some(primitive) = primitive {
        normalized.push(Constraint::Type(primitive));
    }
    if let Some(lower) = lower {
        normalized.push(Constraint::Compare(
            if lower.inclusive {
                CompareOp::Gte
            } else {
                CompareOp::Gt
            },
            lower.number.into_literal(),
        ));
    }
    if let Some(upper) = upper {
        normalized.push(Constraint::Compare(
            if upper.inclusive {
                CompareOp::Lte
            } else {
                CompareOp::Lt
            },
            upper.number.into_literal(),
        ));
    }
    normalized.extend(rest);
    Ok(normalized)
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct Bound {
    number: Number,
    inclusive: bool,
}

impl Bound {
    fn new(number: Number, inclusive: bool) -> Self {
        Self { number, inclusive }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum Number {
    Int(i64),
    Float(f64),
}

impl Number {
    fn from_literal(value: &LiteralValue) -> Option<Self> {
        match value {
            LiteralValue::Int(value) => Some(Self::Int(*value)),
            LiteralValue::Float(value) => Some(Self::Float(*value)),
            LiteralValue::String(_) | LiteralValue::Bool(_) => None,
        }
    }

    fn into_literal(self) -> LiteralValue {
        match self {
            Self::Int(value) => LiteralValue::Int(value),
            Self::Float(value) => LiteralValue::Float(value),
        }
    }

    fn as_f64(self) -> f64 {
        match self {
            Self::Int(value) => value as f64,
            Self::Float(value) => value,
        }
    }
}

fn merge_lower(current: &mut Option<Bound>, next: Bound) {
    match current {
        None => *current = Some(next),
        Some(current_bound) if is_stricter_lower(next, *current_bound) => *current_bound = next,
        Some(_) => {}
    }
}

fn merge_upper(current: &mut Option<Bound>, next: Bound) {
    match current {
        None => *current = Some(next),
        Some(current_bound) if is_stricter_upper(next, *current_bound) => *current_bound = next,
        Some(_) => {}
    }
}

fn is_stricter_lower(next: Bound, current: Bound) -> bool {
    let next_value = next.number.as_f64();
    let current_value = current.number.as_f64();
    next_value > current_value
        || (next_value == current_value && !next.inclusive && current.inclusive)
}

fn is_stricter_upper(next: Bound, current: Bound) -> bool {
    let next_value = next.number.as_f64();
    let current_value = current.number.as_f64();
    next_value < current_value
        || (next_value == current_value && !next.inclusive && current.inclusive)
}

fn ensure_bounds_non_empty(
    primitive: Option<PrimitiveType>,
    lower: Option<Bound>,
    upper: Option<Bound>,
    span: Span,
) -> crate::Result<()> {
    if primitive == Some(PrimitiveType::Int) {
        let min = lower.map(int_lower_bound).unwrap_or(i128::from(i64::MIN));
        let max = upper.map(int_upper_bound).unwrap_or(i128::from(i64::MAX));
        if min > max {
            return Err(empty_numeric_bounds(span));
        }
        return Ok(());
    }

    if let (Some(lower), Some(upper)) = (lower, upper) {
        let lower_value = lower.number.as_f64();
        let upper_value = upper.number.as_f64();
        if lower_value > upper_value {
            return Err(empty_numeric_bounds(span));
        }
        if lower_value == upper_value && !(lower.inclusive && upper.inclusive) {
            return Err(empty_numeric_bounds(span));
        }
    }
    Ok(())
}

fn int_lower_bound(bound: Bound) -> i128 {
    let Number::Int(value) = bound.number else {
        unreachable!()
    };
    if bound.inclusive {
        i128::from(value)
    } else {
        i128::from(value) + 1
    }
}

fn int_upper_bound(bound: Bound) -> i128 {
    let Number::Int(value) = bound.number else {
        unreachable!()
    };
    if bound.inclusive {
        i128::from(value)
    } else {
        i128::from(value) - 1
    }
}

fn empty_numeric_bounds(span: Span) -> Diagnostic {
    Diagnostic::new(
        DiagnosticKind::Conflict,
        span,
        String::from("numeric comparison constraints have an empty intersection"),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::runtime::LiteralValue;

    #[test]
    fn detects_primitive_conflict() {
        assert!(
            normalize_constraints(
                alloc::vec![
                    Constraint::Type(PrimitiveType::Int),
                    Constraint::Type(PrimitiveType::String),
                ],
                Span::default(),
            )
            .is_err()
        );
    }

    #[test]
    fn detects_empty_int_range() {
        assert!(
            normalize_constraints(
                alloc::vec![
                    Constraint::Type(PrimitiveType::Int),
                    Constraint::Compare(CompareOp::Gt, LiteralValue::Int(10)),
                    Constraint::Compare(CompareOp::Lt, LiteralValue::Int(11)),
                ],
                Span::default(),
            )
            .is_err()
        );
    }

    #[test]
    fn keeps_regex_constraints_without_intersection_check() {
        let constraints = normalize_constraints(
            alloc::vec![
                Constraint::Regex(String::from("^a$")),
                Constraint::Regex(String::from("^b$")),
            ],
            Span::default(),
        )
        .unwrap();
        assert_eq!(constraints.len(), 2);
    }
}