rain-lang 0.0.1

/*!
An abstract syntax tree for `rain` programs
*/

//TODO: consider using a non collision-resistant hashmap
use std::fmt::{self, Debug, Display, Formatter};

use crate::value::primitive::binary::{Natural, Integer};
use crate::value::primitive::logical::{LogicalOp, Bool};

/// A scope in a `rain` program
#[derive(Clone, Eq, PartialEq)]
pub struct Scope<'a> {
    /// The definitions in this scope
    pub definitions: Vec<Let<'a>>,
    /// The value in this scope, if any
    pub value: Option<Box<Expr<'a>>>
}

impl Debug for Scope<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Scope<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        if self.definitions.len() == 0 {
            return if let Some(value) = &self.value { write!(fmt, "{{ {} }}", value) }
            else { write!(fmt, "{{}}")}
        }
        write!(fmt, "{{\n")?;
        for definition in self.definitions.iter() {
            write!(fmt, "{}\n", definition)?;
        }
        if let Some(value) = &self.value { write!(fmt, "{}\n", value)? }
        write!(fmt, "}}")
    }
}

/// A phi node
#[derive(Clone, Eq, PartialEq)]
pub struct Phi<'a>(pub Scope<'a>);

impl Debug for Phi<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Phi<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> { write!(fmt, "#phi {}", self.0) }
}


/// A `rain` path, denoting a value in a namespace
#[derive(Clone, Eq, PartialEq)]
pub struct Path<'a> {
    /// The base of this path, if any (TODO: this)
    pub base: Option<()>,
    /// The names in this path
    pub names: Vec<&'a str>
}

impl Debug for Path<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Path<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        let tail_names = if let Some(base) = self.base {
            write!(fmt, "{:?}", base)?; // TODO: fixme
            &self.names[..]
        } else if self.names.len() > 0 {
            write!(fmt, "{}", self.names[0])?;
            &self.names[1..]
        }
        else {
            // Invalid path. For now, do nothing. TODO: think about this
            return Ok(())
        };
        for name in tail_names { write!(fmt, ".{}", name)?; }
        Ok(())
    }
}

impl<'a> Path<'a> {
    /// Create a new path consisting of just an identifier
    pub fn ident<T>(ident: T) -> Path<'a> where T: Into<&'a str> {
        Path { base: None, names: vec![ident.into()] }
    }
}

impl<'a> From<Vec<&'a str>> for Path<'a> {
    fn from(names: Vec<&'a str>) -> Path<'a> { Path { names, base: None } }
}

/// An AST node for a `rain` expression
#[derive(Clone, Eq, PartialEq)]
pub enum Expr<'a> {
    /// A path
    Path(Path<'a>),
    /// A scope
    Scope(Scope<'a>),
    /// An S-expression
    Sexpr(Sexpr<'a>),
    /// A lambda expression
    Lambda(Lambda<'a>),
    /// A pi type
    Pi(Pi<'a>),
    /// A phi node
    Phi(Phi<'a>),
    /// A gamma node
    Gamma(Gamma<'a>),
    /// A boolean constant
    Bool(bool),
    /// The boolean type
    BoolTy(Bool),
    /// A logical operation
    LogicalOp(LogicalOp),
    /// A natural number
    Natural(Natural),
    /// An integer
    Integer(Integer)
}

impl<'a> Expr<'a> {
    /// Create a new path consisting of just an identifier
    pub fn ident<T>(ident: T) -> Expr<'a> where T: Into<&'a str> {
        Expr::Path(Path::ident(ident))
    }
    /// Attempt to build an S-expression from this value, with the other value at the front.
    /// For example, `5 + 5 = (+ 5 5)` can be written `let sum = 5`, `sum.push(5)`, `sum.push(+)`
    pub fn push<T>(self, other: T) -> Expr<'a> where T: Into<Box<Expr<'a>>> {
        match self {
            Expr::Sexpr(mut s) => { s.push(other); Expr::Sexpr(s) }
            atom => {
                Expr::Sexpr(Sexpr { ops: vec![Box::new(atom), other.into()] })
            },
        }
    }
}

impl Debug for Expr<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Expr<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        match self {
            Expr::Path(p) => std::fmt::Display::fmt(p, fmt),
            Expr::Scope(s) => std::fmt::Display::fmt(s, fmt),
            Expr::Sexpr(s) => std::fmt::Display::fmt(s, fmt),
            Expr::Lambda(l) => std::fmt::Display::fmt(l, fmt),
            Expr::Pi(p) => std::fmt::Display::fmt(p, fmt),
            Expr::Phi(p) => std::fmt::Display::fmt(p, fmt),
            Expr::Gamma(g) => std::fmt::Display::fmt(g, fmt),
            Expr::Bool(b) => write!(fmt, "#{}", b),
            Expr::BoolTy(b) => std::fmt::Display::fmt(b, fmt),
            Expr::LogicalOp(l) => std::fmt::Display::fmt(l, fmt),
            Expr::Natural(n) => std::fmt::Display::fmt(n, fmt),
            Expr::Integer(z) => std::fmt::Display::fmt(z, fmt)
        }
    }
}

/// An AST node for an S-expression
#[derive(Clone, Eq, PartialEq)]
pub struct Sexpr<'a> {
    /// The operands of this S-expression. Stored in *reverse* order, i.e. `f(5)` is represented as
    /// the list `[5, f]`.
    pub ops: Vec<Box<Expr<'a>>>
}

impl<'a> From<Vec<Box<Expr<'a>>>> for Sexpr<'a> {
    fn from(ops: Vec<Box<Expr<'a>>>) -> Sexpr<'a> { Sexpr { ops } }
}

impl<'a> Sexpr<'a> {
    /// Push an operand to this S-expression
    pub fn push<T>(&mut self, op: T) where T: Into<Box<Expr<'a>>> {
        self.ops.push(op.into())
    }
}

impl Debug for Sexpr<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Sexpr<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "(")?;
        let mut first: bool = true;
        for op in self.ops.iter().rev() {
            write!(fmt, "{}{}", if first {""} else {" "}, op)?;
            first = false;
        }
        write!(fmt, ")")
    }
}

/// An AST-node for a lambda expression
#[derive(Clone, Eq, PartialEq)]
pub struct Lambda<'a>(pub Parametrized<'a>);

impl Debug for Lambda<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Lambda<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "#lambda {}", self.0)
    }
}

/// An AST-node for a pi expression
#[derive(Clone, Eq, PartialEq)]
pub struct Pi<'a>(pub Parametrized<'a>);

impl Debug for Pi<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Pi<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "#pi {}", self.0)
    }
}

/// A a parametrized expression
#[derive(Clone, Eq, PartialEq)]
pub struct Parametrized<'a> {
    /// The arguments of this parametrized expression, along with their types
    pub args: Vec<(Option<&'a str>, Expr<'a>)>,
    /// The optional return type of this parametrized expression
    pub ret_ty: Option<Box<Expr<'a>>>,
    /// The result of this parametrized expression
    pub result: Box<Expr<'a>>
}

impl Debug for Parametrized<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Parametrized<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "|")?;
        let mut first: bool = true;
        for (arg, ty) in self.args.iter() {
            let sep = if first {""} else {", "};
            let arg = if let Some(arg) = arg { arg } else { "_" };
            write!(fmt, "{}{} : {}", sep, arg, ty)?;
            first = false;
        }
        write!(fmt, "| ")?;
        if let Some(ty) = &self.ret_ty {
            write!(fmt, "=> {} ", ty)?;
        }
        write!(fmt, "{}", self.result)
    }
}

/// A simple assignment (i.e. a single variable with an optional type)
#[derive(Clone, Eq, PartialEq)]
pub struct SimpleAssignment<'a> {
    /// The name of the variable being assigned to
    pub name: &'a str,
    /// The type constraint on the variable, if any
    pub ty: Option<Expr<'a>>
}

impl Debug for SimpleAssignment<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for SimpleAssignment<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", self.name)?;
        if let Some(ty) = &self.ty { write!(fmt, " : {}", ty) } else { Ok(()) }
    }
}


/// A pattern for assignment
#[derive(Clone, Eq, PartialEq)]
pub enum Pattern<'a> {
    /// A simple assignment
    Simple(SimpleAssignment<'a>)
}

impl Debug for Pattern<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Pattern<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        match self {
            Pattern::Simple(s) => write!(fmt, "{}", s)
        }
    }
}

/// A `let` statement, defining a set of variables
#[derive(Clone, Eq, PartialEq)]
pub struct Let<'a> {
    /// The pattern being assigned to
    pub pattern: Pattern<'a>,
    /// The value being assigned
    pub expr: Expr<'a>
}

impl Debug for Let<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Let<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "let {} = {};", self.pattern, self.expr) //TODO: this
    }
}

/// A `gamma` node, corresponding to a `match` statement
#[derive(Clone, Eq, PartialEq)]
pub struct Gamma<'a> {
    /// The branches of this `gamma` node
    pub branches: Vec<(Pattern<'a>, Expr<'a>)>
}

impl Debug for Gamma<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        <Self as Display>::fmt(self, fmt)
    }
}

impl Display for Gamma<'_> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "#match {{\n")?;
        let mut first = true;
        for (pattern, value) in self.branches.iter() {
            write!(fmt, "{}{} => {}", if first { "" } else { ",\n" }, pattern, value)?;
            first = false;
        }
        write!(fmt, "{}}}", if first { "" } else { "\n" })
    }
}

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

    #[test]
    fn expr_push_works_as_expected() {
        assert_eq!(
            Expr::ident("x").push(Expr::ident("f")),
            Expr::Sexpr(vec![Expr::ident("x").into(), Expr::ident("f").into()].into())
        );
        assert_eq!(
            Expr::ident("x").push(Expr::ident("y")).push(Expr::ident("f")),
            Expr::Sexpr(vec![
                Expr::ident("x").into(),
                Expr::ident("y").into(),
                Expr::ident("f").into()]
            .into())
        )
    }
}