graphcal-compiler 0.0.1-alpha.14

//! Data types and function classification constants used by the declaration-collection layer.
//!
//! These types have no dependency on the resolution logic itself, making them
//! suitable for use across all compilation phases.

use std::collections::{HashMap, HashSet};

use crate::dag_id::DagId;
use crate::desugar::desugared_ast::{AssertBody, Expr, FigureDecl, LayerDecl, PlotDecl};
use crate::registry::declared_type::IndexTypeRef;
use crate::syntax::names::{
    DeclName, IndexName, IndexVariantName, NamePath, ResolvedIndexVariant, ScopedName,
};
use crate::syntax::span::Span;

// ---------------------------------------------------------------------------
// Function classification enums
// ---------------------------------------------------------------------------
//
// Each special-function category has its own enum whose variants are the
// canonical source of truth for the function names in that category.
// The `classify_special_fn` function maps a string to one of these.

/// Aggregation functions: operate on indexed collections.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggregationFn {
    Sum,
    Min,
    Max,
    Mean,
    Count,
}

impl AggregationFn {
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        match name {
            "sum" => Some(Self::Sum),
            "min" => Some(Self::Min),
            "max" => Some(Self::Max),
            "mean" => Some(Self::Mean),
            "count" => Some(Self::Count),
            _ => None,
        }
    }

    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Sum => "sum",
            Self::Min => "min",
            Self::Max => "max",
            Self::Mean => "mean",
            Self::Count => "count",
        }
    }
}

/// Type conversion functions: `to_float`, `to_int`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TypeConversionFn {
    ToFloat,
    ToInt,
}

impl TypeConversionFn {
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        match name {
            "to_float" => Some(Self::ToFloat),
            "to_int" => Some(Self::ToInt),
            _ => None,
        }
    }

    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::ToFloat => "to_float",
            Self::ToInt => "to_int",
        }
    }
}

/// Datetime constructor functions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConstructorFn {
    Datetime,
    Epoch,
}

impl ConstructorFn {
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        match name {
            "datetime" => Some(Self::Datetime),
            "epoch" => Some(Self::Epoch),
            _ => None,
        }
    }

    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Datetime => "datetime",
            Self::Epoch => "epoch",
        }
    }
}

/// Datetime extraction functions: extract a component from a `Datetime`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatetimeExtractFn {
    Year,
    Month,
    Day,
    Hour,
    Minute,
    Second,
    Weekday,
    DayOfYear,
}

impl DatetimeExtractFn {
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        match name {
            "year" => Some(Self::Year),
            "month" => Some(Self::Month),
            "day" => Some(Self::Day),
            "hour" => Some(Self::Hour),
            "minute" => Some(Self::Minute),
            "second" => Some(Self::Second),
            "weekday" => Some(Self::Weekday),
            "day_of_year" => Some(Self::DayOfYear),
            _ => None,
        }
    }

    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Year => "year",
            Self::Month => "month",
            Self::Day => "day",
            Self::Hour => "hour",
            Self::Minute => "minute",
            Self::Second => "second",
            Self::Weekday => "weekday",
            Self::DayOfYear => "day_of_year",
        }
    }
}

/// Datetime-from-numeric constructors.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatetimeFromFn {
    FromJd,
    FromMjd,
    FromUnix,
}

impl DatetimeFromFn {
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        match name {
            "from_jd" => Some(Self::FromJd),
            "from_mjd" => Some(Self::FromMjd),
            "from_unix" => Some(Self::FromUnix),
            _ => None,
        }
    }

    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::FromJd => "from_jd",
            Self::FromMjd => "from_mjd",
            Self::FromUnix => "from_unix",
        }
    }
}

/// Datetime-to-numeric functions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatetimeToFn {
    ToJd,
    ToMjd,
    ToUnix,
}

impl DatetimeToFn {
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        match name {
            "to_jd" => Some(Self::ToJd),
            "to_mjd" => Some(Self::ToMjd),
            "to_unix" => Some(Self::ToUnix),
            _ => None,
        }
    }

    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::ToJd => "to_jd",
            Self::ToMjd => "to_mjd",
            Self::ToUnix => "to_unix",
        }
    }
}

/// Classification of special built-in functions.
///
/// Each variant carries a sub-enum identifying the specific function, so
/// downstream handlers can match on typed variants instead of raw strings.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpecialFnKind {
    /// Aggregation functions: `sum`, `min`, `max`, `mean`, `count`.
    Aggregation(AggregationFn),
    /// Type conversion functions: `to_float`, `to_int`.
    TypeConversion(TypeConversionFn),
    /// Time-scale conversion functions: `to_utc`, `to_tai`, etc.
    TimeScaleConversion(crate::registry::time_scale::TimeScale),
    /// Constructor functions: `datetime`, `epoch`.
    Constructor(ConstructorFn),
    /// Datetime extraction functions: `year`, `month`, `day`, etc.
    DatetimeExtract(DatetimeExtractFn),
    /// Datetime-from-numeric functions: `from_jd`, `from_mjd`, `from_unix`.
    DatetimeFrom(DatetimeFromFn),
    /// Datetime-to-numeric functions: `to_jd`, `to_mjd`, `to_unix`.
    DatetimeTo(DatetimeToFn),
}

/// Classify a function name as a special built-in function.
///
/// Returns `None` if the name is not a recognized special function.
#[must_use]
pub fn classify_special_fn(name: &str) -> Option<SpecialFnKind> {
    if let Some(f) = AggregationFn::parse(name) {
        return Some(SpecialFnKind::Aggregation(f));
    }
    if let Some(f) = TypeConversionFn::parse(name) {
        return Some(SpecialFnKind::TypeConversion(f));
    }
    if let Some(scale) = crate::registry::time_scale::time_scale_from_conversion_fn(name) {
        return Some(SpecialFnKind::TimeScaleConversion(scale));
    }
    if let Some(f) = ConstructorFn::parse(name) {
        return Some(SpecialFnKind::Constructor(f));
    }
    if let Some(f) = DatetimeExtractFn::parse(name) {
        return Some(SpecialFnKind::DatetimeExtract(f));
    }
    if let Some(f) = DatetimeFromFn::parse(name) {
        return Some(SpecialFnKind::DatetimeFrom(f));
    }
    if let Some(f) = DatetimeToFn::parse(name) {
        return Some(SpecialFnKind::DatetimeTo(f));
    }
    None
}

/// Returns `true` if `name` is a built-in aggregation function (`sum`, `min`, etc.).
#[must_use]
pub fn is_aggregation_fn(name: &str) -> bool {
    AggregationFn::parse(name).is_some()
}

/// Returns `true` if `name` is a time scale identifier (`UTC`, `TT`, `TAI`, etc.).
#[must_use]
pub fn is_time_scale_name(name: &str) -> bool {
    crate::registry::time_scale::TimeScale::ALL_NAMES.contains(&name)
}

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/// Pre-evaluated value bindings imported from already-evaluated dependency files.
///
/// Unlike `ImportedNames` which carries AST expressions, this carries
/// evaluated values. Used in per-file evaluation where each file is
/// compiled and evaluated independently.
#[derive(Debug, Default, Clone)]
pub struct ImportedValueNames {
    /// Imported const names (for scope checking only — actual values are in the exec plan).
    pub const_names: Vec<(ScopedName, Span)>,
    /// Imported param names.
    pub param_names: Vec<(ScopedName, Span)>,
    /// Imported node names.
    pub node_names: Vec<(ScopedName, Span)>,
    /// Imported assert names (for `#[assumes]` validation).
    pub assert_names: Vec<(DeclName, Span)>,
    /// Plot aliases requested by include brace lists (#847). Registered in
    /// the value namespace for collision checking and recorded on the DAG so
    /// figures/layers can reference them.
    pub plot_names: Vec<(ScopedName, Span)>,
}

/// The kind of a declaration (used for source-order tracking).
#[derive(Debug, Clone, Copy)]
pub enum DeclCategory {
    Const,
    Param,
    Node,
    Assert,
    Plot,
    Figure,
    Layer,
}

impl std::fmt::Display for DeclCategory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Const => write!(f, "const"),
            Self::Param => write!(f, "param"),
            Self::Node => write!(f, "node"),
            Self::Assert => write!(f, "assert"),
            Self::Plot => write!(f, "plot"),
            Self::Figure => write!(f, "figure"),
            Self::Layer => write!(f, "layer"),
        }
    }
}

// ---------------------------------------------------------------------------
// Entry types for resolved declarations
// ---------------------------------------------------------------------------

/// A resolved const declaration (before type annotation is added).
#[derive(Debug)]
pub struct ResolvedConstEntry {
    pub name: String,
    pub expr: Expr,
    pub span: Span,
}

/// A resolved param declaration (before type annotation is added).
#[derive(Debug)]
pub struct ResolvedParamEntry {
    pub name: String,
    pub default_expr: Option<Expr>,
    pub span: Span,
}

/// A resolved node declaration (before type annotation is added).
#[derive(Debug)]
pub struct ResolvedNodeEntry {
    pub name: String,
    pub expr: Expr,
    pub span: Span,
}

/// A resolved assert declaration.
#[derive(Debug)]
pub struct ResolvedAssertEntry {
    pub name: String,
    pub body: AssertBody,
    pub span: Span,
}

/// A resolved plot declaration.
#[derive(Debug)]
pub struct ResolvedPlotEntry {
    pub name: String,
    pub decl: PlotDecl,
    pub span: Span,
}

/// A resolved figure declaration.
#[derive(Debug)]
pub struct ResolvedFigureEntry {
    pub name: String,
    pub decl: FigureDecl,
    pub span: Span,
}

/// A resolved layer declaration.
#[derive(Debug)]
pub struct ResolvedLayerEntry {
    pub name: String,
    pub decl: LayerDecl,
    pub span: Span,
}

/// One axis segment in a per-variant `#[expected_fail(...)]` key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExpectedFailKeyPart {
    /// An `Index.Variant` / `module.Index.Variant` segment for a named axis.
    ///
    /// The `index` carrier is the semantic key used by runtime assertion
    /// checks. Before module-aware TIR resolution, `source_index_path`
    /// preserves the structured syntax path. After resolution, `index`
    /// carries the canonical owner used by runtime checks.
    Named {
        index: IndexTypeRef,
        variant: IndexVariantName,
        source_index_path: Option<NamePath>,
        span: Span,
    },
    /// A `#N` segment for a Nat range axis (#816).
    ///
    /// Range axes have no source-level index name, so the axis identity is
    /// positional: the segment binds to the assertion's axis at the same
    /// tuple position, validated at dim-check time.
    RangeStep { step: u64, span: Span },
}

impl ExpectedFailKeyPart {
    fn unresolved_owner() -> DagId {
        DagId::root("<expected-fail-unresolved>")
    }

    #[must_use]
    pub fn unresolved(index_path: NamePath, variant: IndexVariantName, span: Span) -> Self {
        Self::Named {
            index: IndexTypeRef::with_owner(
                Self::unresolved_owner(),
                IndexName::from_atom(index_path.leaf().clone()),
            ),
            variant,
            source_index_path: Some(index_path),
            span,
        }
    }

    #[must_use]
    pub fn with_owner(
        owner: DagId,
        index: IndexName,
        variant: IndexVariantName,
        span: Span,
    ) -> Self {
        Self::Named {
            index: IndexTypeRef::with_owner(owner, index),
            variant,
            source_index_path: None,
            span,
        }
    }

    #[must_use]
    pub fn resolved(resolved: ResolvedIndexVariant, span: Span) -> Self {
        let (index, variant) = resolved.into_parts();
        Self::Named {
            index: IndexTypeRef::from_resolved(index),
            variant,
            source_index_path: None,
            span,
        }
    }

    #[must_use]
    pub fn with_resolved_variant(&self, resolved: ResolvedIndexVariant) -> Self {
        Self::resolved(resolved, self.span())
    }

    /// The source span of this key segment.
    #[must_use]
    pub const fn span(&self) -> Span {
        match self {
            Self::Named { span, .. } | Self::RangeStep { span, .. } => *span,
        }
    }

    /// The named index reference, when this segment targets a named axis.
    #[must_use]
    pub const fn named_index(&self) -> Option<&IndexTypeRef> {
        match self {
            Self::Named { index, .. } => Some(index),
            Self::RangeStep { .. } => None,
        }
    }

    /// The structured syntax path awaiting resolution, for named segments.
    #[must_use]
    pub const fn source_index_path(&self) -> Option<&NamePath> {
        match self {
            Self::Named {
                source_index_path, ..
            } => source_index_path.as_ref(),
            Self::RangeStep { .. } => None,
        }
    }

    /// The variant key this segment selects within its axis.
    #[must_use]
    pub fn variant(&self) -> IndexVariantName {
        match self {
            Self::Named { variant, .. } => variant.clone(),
            Self::RangeStep { step, .. } => IndexVariantName::range_step(*step),
        }
    }

    /// Whether this segment selects the given entry of an indexed value.
    ///
    /// Named segments require the entry's index identity to match; `#N`
    /// segments match the `#N` entry of any Nat range axis (the axis itself
    /// was bound positionally at dim-check time).
    #[must_use]
    pub fn matches_entry(&self, index: &IndexTypeRef, variant: &IndexVariantName) -> bool {
        match self {
            Self::Named {
                index: expected,
                variant: expected_variant,
                ..
            } => expected.matches_ref(index) && variant == expected_variant,
            Self::RangeStep { step, .. } => {
                matches!(index, IndexTypeRef::NatRange(_))
                    && *variant == IndexVariantName::range_step(*step)
            }
        }
    }

    /// Render this segment for diagnostics: `Index.Variant` or `#N`.
    #[must_use]
    pub fn display(&self) -> String {
        match self {
            Self::Named { index, variant, .. } => format!("{}.{variant}", index.display_name()),
            Self::RangeStep { step, .. } => format!("#{step}"),
        }
    }
}

/// A single expected-fail key: a list of index/variant pairs.
///
/// - Length 1 for single-index assertions: `[Mode.Boost]`
/// - Length >1 for multi-index assertions: `[(Mode.Boost, Phase.Launch)]`
pub type ExpectedFailKey = Vec<ExpectedFailKeyPart>;

/// Describes how an assertion is expected to fail.
#[derive(Debug, Clone)]
pub enum ExpectedFail {
    /// The entire assertion is expected to fail: `#[expected_fail]`.
    All,
    /// Specific index keys are expected to fail: `#[expected_fail(Index.Variant, ...)]`.
    Variants(Vec<ExpectedFailKey>),
}

/// The result of declaration collection: declarations separated by category.
#[derive(Debug)]
pub struct ResolvedFile {
    /// Const declarations in source order.
    pub consts: Vec<ResolvedConstEntry>,
    /// Param declarations in source order.
    pub params: Vec<ResolvedParamEntry>,
    /// Node declarations in source order.
    pub nodes: Vec<ResolvedNodeEntry>,
    /// Assert declarations in source order.
    pub asserts: Vec<ResolvedAssertEntry>,
    /// Plot declarations in source order.
    pub plots: Vec<ResolvedPlotEntry>,
    /// Figure declarations in source order.
    pub figures: Vec<ResolvedFigureEntry>,
    /// Layer declarations in source order.
    pub layers: Vec<ResolvedLayerEntry>,
    /// All declaration names in source order with their category.
    pub source_order: Vec<(DeclName, DeclCategory)>,
    /// Set of all assert names (for checking `@assert_name` errors).
    pub assert_names: HashSet<DeclName>,
    /// Mapping from assert name to the list of declarations that assume it.
    /// Built from `#[assumes(...)]` attributes.
    pub assumes_map: HashMap<DeclName, Vec<DeclName>>,
    /// Mapping from assert name to its expected-fail configuration.
    /// Built from `#[expected_fail]` / `#[expected_fail(...)]` attributes.
    pub expected_fail: HashMap<DeclName, ExpectedFail>,
    /// Plot names carrying `#[hidden]`: evaluated and referenceable from
    /// figures/layers, but excluded from standalone output (#847).
    pub hidden_plots: HashSet<DeclName>,
    /// Names of all declarations marked `pub` in this file (values + type-system).
    pub pub_names: HashSet<DeclName>,
}