lutra-compiler 0.5.1

use enum_as_inner::EnumAsInner;
use indexmap::IndexMap;
use itertools::Itertools;

use crate::Span;
use crate::pr::path::Path;
use crate::pr::{Expr, Ty};

/// Definition.
#[derive(Debug, Clone, PartialEq)]
pub struct Def {
    pub kind: DefKind,

    pub annotations: Vec<Annotation>,

    pub doc_comment: Option<DocComment>,

    /// Code span of the whole definition (including doc comments and annotations)
    pub span: Option<Span>,

    /// Code span of definition name
    pub span_name: Option<Span>,
}

#[derive(Debug, EnumAsInner, PartialEq, Clone)]
pub enum DefKind {
    Module(ModuleDef),
    Expr(ExprDef),
    Ty(TyDef),
    Import(ImportDef),

    /// A definition that has not yet been resolved.
    /// Created during the first pass of the AST, must not be present in
    /// a fully resolved module structure.
    Unresolved(Option<Box<DefKind>>),
}

#[derive(PartialEq, Clone, Default)]
pub struct ModuleDef {
    pub defs: IndexMap<String, Def>,

    // Self-annotations, defined within the module
    pub annotations: Vec<Annotation>,

    /// Span covering the content of this module, i.e. the region inside the
    /// braces for inline modules, or the whole file body for file-based
    /// submodules. This excludes inner doc comments or annotations.
    pub span_content: Option<Span>,
}

#[derive(Debug, PartialEq, Clone)]
pub struct ExprDef {
    pub value: Box<Expr>,

    pub constant: bool,
    pub ty: Option<Ty>,
}

#[derive(Debug, PartialEq, Clone)]
pub struct TyDef {
    pub ty: Ty,
    pub is_framed: bool,
    pub framed_label: Option<String>,
}

#[derive(Debug, PartialEq, Clone)]
pub struct ImportDef {
    pub kind: ImportKind,
    pub span: Span,
}

#[derive(Debug, PartialEq, Clone)]
pub enum ImportKind {
    /// Represents `std::sql::from as read_table`
    Single(Path, Option<String>),

    /// Represents `std::sql::{...}`
    Many(Path, Vec<ImportDef>),
}

impl ImportDef {
    pub fn new_simple(path: Path, span: Span) -> Self {
        Self {
            kind: ImportKind::Single(path, None),
            span,
        }
    }

    pub fn as_simple(&self) -> Option<&Path> {
        let ImportKind::Single(path, _) = &self.kind else {
            return None;
        };
        Some(path)
    }

    pub fn path(&self) -> &Path {
        match &self.kind {
            ImportKind::Single(path, _) => path,
            ImportKind::Many(path, _) => path,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Annotation {
    pub expr: Box<Expr>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct DocComment {
    pub content: String,
    pub span: Span,
}

impl Def {
    pub fn new<K: Into<DefKind>>(kind: K) -> Def {
        Def {
            kind: kind.into(),
            annotations: Vec::new(),
            doc_comment: None,

            span: None,
            span_name: None,
        }
    }
}

impl From<ModuleDef> for DefKind {
    fn from(value: ModuleDef) -> Self {
        DefKind::Module(value)
    }
}
impl From<Expr> for DefKind {
    fn from(expr: Expr) -> Self {
        DefKind::Expr(ExprDef {
            value: Box::new(expr),
            ty: None,
            constant: false,
        })
    }
}

impl std::fmt::Debug for ModuleDef {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut ds = f.debug_struct("ModuleDef");

        if !self.annotations.is_empty() {
            ds.field("annotations", &self.annotations);
        }

        if self.defs.len() < 15 {
            ds.field("defs", &DebugNames(&self.defs));
        } else {
            ds.field("defs", &format!("... {} entries ...", self.defs.len()));
        }
        ds.finish()
    }
}

struct DebugNames<'a>(&'a IndexMap<String, Def>);

impl<'a> std::fmt::Debug for DebugNames<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut dm = f.debug_map();
        for (n, def) in self.0.iter().sorted_by_key(|x| x.0) {
            dm.entry(n, def);
        }
        dm.finish()
    }
}