lux/src/ast.rs

//! Abstract Syntax Tree for the Lux language

#![allow(dead_code)]

use std::fmt;

/// Source location for error reporting
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}

impl Span {
    pub fn new(start: usize, end: usize) -> Self {
        Self { start, end }
    }

    pub fn merge(self, other: Span) -> Span {
        Span {
            start: self.start.min(other.start),
            end: self.end.max(other.end),
        }
    }
}

/// An identifier (variable or type name)
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Ident {
    pub name: String,
    pub span: Span,
}

impl Ident {
    pub fn new(name: impl Into<String>, span: Span) -> Self {
        Self {
            name: name.into(),
            span,
        }
    }
}

impl fmt::Display for Ident {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name)
    }
}

/// Visibility modifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Visibility {
    /// Public - exported from module
    Public,
    /// Private - only visible within module (default)
    #[default]
    Private,
}

// ============ Schema Evolution ============

/// A version number for schema evolution (e.g., @v1, @v2)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Version {
    pub number: u32,
    pub span: Span,
}

impl PartialOrd for Version {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Version {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.number.cmp(&other.number)
    }
}

impl Version {
    pub fn new(number: u32, span: Span) -> Self {
        Self { number, span }
    }
}

impl fmt::Display for Version {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "@v{}", self.number)
    }
}

/// Version constraint for type annotations
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VersionConstraint {
    /// Exactly this version: @v2
    Exact(Version),
    /// This version or later: @v2+
    AtLeast(Version),
    /// Latest version: @latest
    Latest(Span),
}

impl fmt::Display for VersionConstraint {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            VersionConstraint::Exact(v) => write!(f, "{}", v),
            VersionConstraint::AtLeast(v) => write!(f, "{}+", v),
            VersionConstraint::Latest(_) => write!(f, "@latest"),
        }
    }
}

/// Migration from one version to another
#[derive(Debug, Clone)]
pub struct Migration {
    /// Source version: from @v1
    pub from_version: Version,
    /// Migration body (expression that transforms old to new)
    pub body: Expr,
    pub span: Span,
}

// ============ Behavioral Types ============

/// A behavioral property that can be attached to functions
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BehavioralProperty {
    /// No side effects - function only depends on inputs
    Pure,
    /// Always terminates and doesn't throw exceptions
    Total,
    /// f(f(x)) == f(x) for all x
    Idempotent,
    /// Same inputs always produce same outputs (no randomness)
    Deterministic,
    /// Order of arguments doesn't matter: f(a, b) == f(b, a)
    Commutative,
}

impl fmt::Display for BehavioralProperty {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BehavioralProperty::Pure => write!(f, "pure"),
            BehavioralProperty::Total => write!(f, "total"),
            BehavioralProperty::Idempotent => write!(f, "idempotent"),
            BehavioralProperty::Deterministic => write!(f, "deterministic"),
            BehavioralProperty::Commutative => write!(f, "commutative"),
        }
    }
}

/// A where clause constraint on a function
#[derive(Debug, Clone)]
pub enum WhereClause {
    /// Type parameter has a property: where F is pure
    PropertyConstraint {
        type_param: Ident,
        property: BehavioralProperty,
        span: Span,
    },
    /// Result refinement: where result > 0
    ResultRefinement { predicate: Box<Expr>, span: Span },
    /// Trait constraint: where T: Show, where T: Eq + Ord
    TraitConstraint(TraitConstraint),
}

/// Module path: foo/bar/baz
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModulePath {
    pub segments: Vec<Ident>,
    pub span: Span,
}

impl ModulePath {
    pub fn to_string(&self) -> String {
        self.segments
            .iter()
            .map(|s| s.name.as_str())
            .collect::<Vec<_>>()
            .join("/")
    }
}

/// Import declaration
#[derive(Debug, Clone)]
pub struct ImportDecl {
    /// The module path being imported
    pub path: ModulePath,
    /// Optional alias: import foo/bar as baz
    pub alias: Option<Ident>,
    /// Specific items to import: import foo.{a, b, c}
    pub items: Option<Vec<Ident>>,
    /// Import all items: import foo.*
    pub wildcard: bool,
    pub span: Span,
}

/// A complete program (or module)
#[derive(Debug, Clone)]
pub struct Program {
    /// Module imports
    pub imports: Vec<ImportDecl>,
    /// Top-level declarations
    pub declarations: Vec<Declaration>,
}

/// Top-level declarations
#[derive(Debug, Clone)]
pub enum Declaration {
    /// Function definition: fn name(params): ReturnType with {Effects} = body
    Function(FunctionDecl),
    /// Effect declaration: effect Name { fn op1(...): T, ... }
    Effect(EffectDecl),
    /// Type alias or ADT: type Name = ...
    Type(TypeDecl),
    /// Handler definition: handler name: Effect { ... }
    Handler(HandlerDecl),
    /// Let binding at top level
    Let(LetDecl),
    /// Trait declaration: trait Name { fn method(...): T, ... }
    Trait(TraitDecl),
    /// Trait implementation: impl Trait for Type { ... }
    Impl(ImplDecl),
    /// Extern function declaration (FFI): extern fn name(params): ReturnType
    ExternFn(ExternFnDecl),
    /// Extern let declaration (FFI): extern let name: Type
    ExternLet(ExternLetDecl),
}

/// Function declaration
#[derive(Debug, Clone)]
pub struct FunctionDecl {
    pub visibility: Visibility,
    /// Documentation comment (from /// doc comments)
    pub doc: Option<String>,
    pub name: Ident,
    pub type_params: Vec<Ident>,
    pub params: Vec<Parameter>,
    pub return_type: TypeExpr,
    pub effects: Vec<Ident>,
    /// Behavioral properties: is pure, is total, etc.
    pub properties: Vec<BehavioralProperty>,
    /// Where clause constraints
    pub where_clauses: Vec<WhereClause>,
    pub body: Expr,
    pub span: Span,
}

/// Function parameter
#[derive(Debug, Clone)]
pub struct Parameter {
    pub name: Ident,
    pub typ: TypeExpr,
    pub span: Span,
}

/// Effect declaration
#[derive(Debug, Clone)]
pub struct EffectDecl {
    /// Documentation comment
    pub doc: Option<String>,
    pub name: Ident,
    pub type_params: Vec<Ident>,
    pub operations: Vec<EffectOp>,
    pub span: Span,
}

/// An operation within an effect
#[derive(Debug, Clone)]
pub struct EffectOp {
    pub name: Ident,
    pub params: Vec<Parameter>,
    pub return_type: TypeExpr,
    pub span: Span,
}

/// Type declaration (alias or ADT)
#[derive(Debug, Clone)]
pub struct TypeDecl {
    pub visibility: Visibility,
    /// Documentation comment
    pub doc: Option<String>,
    pub name: Ident,
    pub type_params: Vec<Ident>,
    /// Optional version annotation: type User @v2 { ... }
    pub version: Option<Version>,
    pub definition: TypeDef,
    /// Migrations from previous versions: from @v1 = { ... }
    pub migrations: Vec<Migration>,
    pub span: Span,
}

/// Type definition
#[derive(Debug, Clone)]
pub enum TypeDef {
    /// Type alias: type Foo = Bar
    Alias(TypeExpr),
    /// Record type: type Foo { field: Type, ... }
    Record(Vec<RecordField>),
    /// Enum/ADT: type Foo = A | B(Int) | C { x: Int }
    Enum(Vec<Variant>),
}

/// Record field
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordField {
    pub name: Ident,
    pub typ: TypeExpr,
    pub span: Span,
}

/// Enum variant
#[derive(Debug, Clone)]
pub struct Variant {
    pub name: Ident,
    pub fields: VariantFields,
    pub span: Span,
}

/// Variant field types
#[derive(Debug, Clone)]
pub enum VariantFields {
    /// Unit variant: A
    Unit,
    /// Tuple variant: A(Int, String)
    Tuple(Vec<TypeExpr>),
    /// Record variant: A { x: Int, y: String }
    Record(Vec<RecordField>),
}

/// Handler declaration
#[derive(Debug, Clone)]
pub struct HandlerDecl {
    pub name: Ident,
    pub params: Vec<Parameter>,
    pub effect: Ident,
    pub implementations: Vec<HandlerImpl>,
    pub span: Span,
}

/// Implementation of an effect operation in a handler
#[derive(Debug, Clone)]
pub struct HandlerImpl {
    pub op_name: Ident,
    pub params: Vec<Ident>,
    pub resume: Option<Ident>, // The continuation parameter
    pub body: Expr,
    pub span: Span,
}

/// Let declaration
#[derive(Debug, Clone)]
pub struct LetDecl {
    pub visibility: Visibility,
    /// Documentation comment
    pub doc: Option<String>,
    pub name: Ident,
    pub typ: Option<TypeExpr>,
    pub value: Expr,
    pub span: Span,
}

/// Trait declaration: trait Show { fn show(self): String }
#[derive(Debug, Clone)]
pub struct TraitDecl {
    pub visibility: Visibility,
    /// Documentation comment
    pub doc: Option<String>,
    pub name: Ident,
    /// Type parameters: trait Functor<F> { ... }
    pub type_params: Vec<Ident>,
    /// Super traits: trait Ord: Eq { ... }
    pub super_traits: Vec<TraitBound>,
    /// Method signatures
    pub methods: Vec<TraitMethod>,
    pub span: Span,
}

/// A trait method signature
#[derive(Debug, Clone)]
pub struct TraitMethod {
    pub name: Ident,
    pub type_params: Vec<Ident>,
    pub params: Vec<Parameter>,
    pub return_type: TypeExpr,
    /// Optional default implementation
    pub default_impl: Option<Expr>,
    pub span: Span,
}

/// A trait bound: Show, Eq, Ord<T>
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TraitBound {
    pub trait_name: Ident,
    pub type_args: Vec<TypeExpr>,
    pub span: Span,
}

/// Trait implementation: impl Show for Int { ... }
#[derive(Debug, Clone)]
pub struct ImplDecl {
    /// Type parameters: impl<T: Show> Show for List<T> { ... }
    pub type_params: Vec<Ident>,
    /// Trait constraints on type parameters
    pub constraints: Vec<TraitConstraint>,
    /// The trait being implemented
    pub trait_name: Ident,
    /// Type arguments for the trait: impl Functor<List> for ...
    pub trait_args: Vec<TypeExpr>,
    /// The type implementing the trait
    pub target_type: TypeExpr,
    /// Method implementations
    pub methods: Vec<ImplMethod>,
    pub span: Span,
}

/// A trait constraint: T: Show, T: Eq + Ord
#[derive(Debug, Clone)]
pub struct TraitConstraint {
    pub type_param: Ident,
    pub bounds: Vec<TraitBound>,
    pub span: Span,
}

/// A method implementation in an impl block
#[derive(Debug, Clone)]
pub struct ImplMethod {
    pub name: Ident,
    pub params: Vec<Parameter>,
    pub return_type: Option<TypeExpr>,
    pub body: Expr,
    pub span: Span,
}

/// Extern function declaration (FFI)
#[derive(Debug, Clone)]
pub struct ExternFnDecl {
    pub visibility: Visibility,
    /// Documentation comment
    pub doc: Option<String>,
    pub name: Ident,
    pub type_params: Vec<Ident>,
    pub params: Vec<Parameter>,
    pub return_type: TypeExpr,
    /// Optional JS name override: extern fn foo(...): T = "jsFoo"
    pub js_name: Option<String>,
    pub span: Span,
}

/// Extern let declaration (FFI)
#[derive(Debug, Clone)]
pub struct ExternLetDecl {
    pub visibility: Visibility,
    /// Documentation comment
    pub doc: Option<String>,
    pub name: Ident,
    pub typ: TypeExpr,
    /// Optional JS name override: extern let foo: T = "window.foo"
    pub js_name: Option<String>,
    pub span: Span,
}

/// Type expressions
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TypeExpr {
    /// Named type: Int, String, List
    Named(Ident),
    /// Generic type application: List<Int>, Map<String, Int>
    App(Box<TypeExpr>, Vec<TypeExpr>),
    /// Function type: fn(A, B): C
    Function {
        params: Vec<TypeExpr>,
        return_type: Box<TypeExpr>,
        effects: Vec<Ident>,
    },
    /// Tuple type: (A, B, C)
    Tuple(Vec<TypeExpr>),
    /// Record type: { name: String, age: Int }
    Record(Vec<RecordField>),
    /// Unit type
    Unit,
    /// Versioned type: User @v2, User @v2+, User @latest
    Versioned {
        base: Box<TypeExpr>,
        constraint: VersionConstraint,
    },
}

impl TypeExpr {
    pub fn named(name: &str) -> Self {
        TypeExpr::Named(Ident::new(name, Span::default()))
    }
}

/// Expressions
#[derive(Debug, Clone)]
pub enum Expr {
    /// Literal values
    Literal(Literal),
    /// Variable reference
    Var(Ident),
    /// Binary operation: a + b
    BinaryOp {
        op: BinaryOp,
        left: Box<Expr>,
        right: Box<Expr>,
        span: Span,
    },
    /// Unary operation: -a, !a
    UnaryOp {
        op: UnaryOp,
        operand: Box<Expr>,
        span: Span,
    },
    /// Function call: foo(a, b)
    Call {
        func: Box<Expr>,
        args: Vec<Expr>,
        span: Span,
    },
    /// Effect operation call: Effect.operation(args)
    EffectOp {
        effect: Ident,
        operation: Ident,
        args: Vec<Expr>,
        span: Span,
    },
    /// Field access: foo.bar
    Field {
        object: Box<Expr>,
        field: Ident,
        span: Span,
    },
    /// Tuple index access: tuple.0, tuple.1
    TupleIndex {
        object: Box<Expr>,
        index: usize,
        span: Span,
    },
    /// Lambda: fn(x, y) => x + y  or  fn(x: Int): Int => x + 1
    Lambda {
        params: Vec<Parameter>,
        return_type: Option<Box<TypeExpr>>,
        effects: Vec<Ident>,
        body: Box<Expr>,
        span: Span,
    },
    /// Let binding: let x = e1; e2
    Let {
        name: Ident,
        typ: Option<TypeExpr>,
        value: Box<Expr>,
        body: Box<Expr>,
        span: Span,
    },
    /// If expression: if cond then e1 else e2
    If {
        condition: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Box<Expr>,
        span: Span,
    },
    /// Match expression
    Match {
        scrutinee: Box<Expr>,
        arms: Vec<MatchArm>,
        span: Span,
    },
    /// Block: { e1; e2; e3 }
    Block {
        statements: Vec<Statement>,
        result: Box<Expr>,
        span: Span,
    },
    /// Record literal: { name: "Alice", age: 30 }
    /// With optional spread: { ...base, name: "Bob" }
    Record {
        spread: Option<Box<Expr>>,
        fields: Vec<(Ident, Expr)>,
        span: Span,
    },
    /// Tuple literal: (1, "hello", true)
    Tuple { elements: Vec<Expr>, span: Span },
    /// List literal: [1, 2, 3]
    List { elements: Vec<Expr>, span: Span },
    /// Run with handlers: run expr with { Effect = handler, ... }
    Run {
        expr: Box<Expr>,
        handlers: Vec<(Ident, Expr)>,
        span: Span,
    },
    /// Resume continuation in handler (like calling the continuation)
    Resume { value: Box<Expr>, span: Span },
}

impl Expr {
    pub fn span(&self) -> Span {
        match self {
            Expr::Literal(lit) => lit.span,
            Expr::Var(ident) => ident.span,
            Expr::BinaryOp { span, .. } => *span,
            Expr::UnaryOp { span, .. } => *span,
            Expr::Call { span, .. } => *span,
            Expr::EffectOp { span, .. } => *span,
            Expr::Field { span, .. } => *span,
            Expr::TupleIndex { span, .. } => *span,
            Expr::Lambda { span, .. } => *span,
            Expr::Let { span, .. } => *span,
            Expr::If { span, .. } => *span,
            Expr::Match { span, .. } => *span,
            Expr::Block { span, .. } => *span,
            Expr::Record { span, .. } => *span,
            Expr::Tuple { span, .. } => *span,
            Expr::List { span, .. } => *span,
            Expr::Run { span, .. } => *span,
            Expr::Resume { span, .. } => *span,
        }
    }
}

/// Literal values
#[derive(Debug, Clone)]
pub struct Literal {
    pub kind: LiteralKind,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub enum LiteralKind {
    Int(i64),
    Float(f64),
    String(String),
    Char(char),
    Bool(bool),
    Unit,
}

/// Binary operators
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinaryOp {
    // Arithmetic
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    // Comparison
    Eq,
    Ne,
    Lt,
    Le,
    Gt,
    Ge,
    // Logical
    And,
    Or,
    // Other
    Pipe,   // |>
    Concat, // ++
}

impl fmt::Display for BinaryOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BinaryOp::Add => write!(f, "+"),
            BinaryOp::Sub => write!(f, "-"),
            BinaryOp::Mul => write!(f, "*"),
            BinaryOp::Div => write!(f, "/"),
            BinaryOp::Mod => write!(f, "%"),
            BinaryOp::Eq => write!(f, "=="),
            BinaryOp::Ne => write!(f, "!="),
            BinaryOp::Lt => write!(f, "<"),
            BinaryOp::Le => write!(f, "<="),
            BinaryOp::Gt => write!(f, ">"),
            BinaryOp::Ge => write!(f, ">="),
            BinaryOp::And => write!(f, "&&"),
            BinaryOp::Or => write!(f, "||"),
            BinaryOp::Pipe => write!(f, "|>"),
            BinaryOp::Concat => write!(f, "++"),
        }
    }
}

/// Unary operators
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UnaryOp {
    Neg, // -
    Not, // !
}

impl fmt::Display for UnaryOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            UnaryOp::Neg => write!(f, "-"),
            UnaryOp::Not => write!(f, "!"),
        }
    }
}

/// Statement in a block
#[derive(Debug, Clone)]
pub enum Statement {
    /// Expression statement
    Expr(Expr),
    /// Let binding without body (in blocks)
    Let {
        name: Ident,
        typ: Option<TypeExpr>,
        value: Expr,
        span: Span,
    },
}

/// Match arm
#[derive(Debug, Clone)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub guard: Option<Expr>,
    pub body: Expr,
    pub span: Span,
}

/// Patterns for matching
#[derive(Debug, Clone)]
pub enum Pattern {
    /// Wildcard: _
    Wildcard(Span),
    /// Variable binding: x
    Var(Ident),
    /// Literal: 42, "hello", true
    Literal(Literal),
    /// Constructor: Some(x), None, Ok(v), module.Constructor(x)
    Constructor {
        module: Option<Ident>,
        name: Ident,
        fields: Vec<Pattern>,
        span: Span,
    },
    /// Record pattern: { name, age: a }
    Record {
        fields: Vec<(Ident, Pattern)>,
        span: Span,
    },
    /// Tuple pattern: (a, b, c)
    Tuple { elements: Vec<Pattern>, span: Span },
}

impl Pattern {
    pub fn span(&self) -> Span {
        match self {
            Pattern::Wildcard(span) => *span,
            Pattern::Var(ident) => ident.span,
            Pattern::Literal(lit) => lit.span,
            Pattern::Constructor { span, .. } => *span,
            Pattern::Record { span, .. } => *span,
            Pattern::Tuple { span, .. } => *span,
        }
    }
}