oftlisp-anf 0.1.3

use std::collections::{BTreeSet, HashMap};
use std::fmt::{Display, Formatter, Result as FmtResult, Write};

use gc::Gc;
use itertools::Itertools;
use oftlisp::{Context as OftLispContext, Symbol, Value};
use oftlisp::ast::{Args, Expr as AstExpr};
use oftlisp::collections::GcLinkedList;

use errors::RuntimeError;

/// The `Context` provided by this crate.
#[derive(Clone, Debug, Finalize, Trace)]
pub struct Context;

impl OftLispContext for Context {
    type BuiltinFunction = BuiltinFunction;
    type Expr = Expr;
    type ObjectVtable = HashMap<Symbol, Gc<Value<Context>>>;
    type UserFunction = UserFunction;
    type ValueMeta = ();

    fn from_expr(expr: Gc<AstExpr<Self>>) -> Gc<Expr> {
        Self::convert_expr(expr)
    }
}

/// The type of a built-in function.
pub type BuiltinFunction = fn(Vec<Gc<Value<Context>>>) -> Result<Gc<Value<Context>>, RuntimeError>;

/// The type of a user-defined function.
pub type UserFunction = (Gc<Expr>, GcLinkedList<(Symbol, Gc<Value<Context>>)>);

/// A complex expression, i.e. one that can fail to normalize.
#[derive(Clone, Debug, Finalize, PartialEq, Trace)]
pub enum Expr {
    /// A function call.
    Call(Gc<Prim>, Vec<Gc<Prim>>),

    /// A conditional expression.
    If(Gc<Prim>, Gc<Expr>, Gc<Expr>),

    /// A let binding.
    Let(Option<Symbol>, Gc<Expr>, Gc<Expr>),

    /// A primitive expression.
    Prim(Gc<Prim>),
}

impl Expr {
    /// Returns the free variables of the expression.
    pub fn freevars(&self) -> BTreeSet<Symbol> {
        match *self {
            Expr::Call(ref f, ref a) => {
                let mut fv = f.freevars();
                for a in a {
                    fv.extend(a.freevars());
                }
                fv
            },
            Expr::If(ref c, ref t, ref e) => {
                let mut fv = c.freevars();
                fv.extend(&t.freevars());
                fv.extend(&e.freevars());
                fv
            },
            Expr::Let(n, ref a, ref b) => {
                let mut fv = b.freevars();
                if let Some(n) = n {
                    fv.remove(&n);
                }
                fv.extend(&a.freevars());
                fv
            },
            Expr::Prim(ref p) => p.freevars(),
        }
    }
}

impl Display for Expr {
    fn fmt(&self, fmt: &mut Formatter) -> FmtResult {
        match *self {
            Expr::Call(ref f, ref a) => {
                write!(fmt, "{}({})", f, a.iter().join(" "))
            },
            Expr::If(ref c, ref t, ref e) => {
                write!(fmt, "if({}) {{ {} }} else {{ {} }}", c, t, e)
            },
            Expr::Let(None, ref x, ref y) => write!(fmt, "({}) >> ({})", x, y),
            Expr::Let(Some(n), ref x, ref y) => {
                write!(fmt, "let {} = ({}) in {}", n, x, y)
            },
            Expr::Prim(ref p) => Display::fmt(p, fmt),
        }
    }
}

/// A primitive expression, i.e. one that immediately normalizes.
#[derive(Clone, Debug, Finalize, PartialEq, Trace)]
pub enum Prim {
    /// A function literal, i.e. a lambda.
    Fn(Option<Symbol>, Gc<Args<Context>>, Gc<Expr>),

    /// A literal value.
    Lit(Gc<Value<Context>>),

    /// A variable.
    Var(Symbol),

    /// A literal vector.
    Vec(Vec<Gc<Prim>>),
}

impl Prim {
    /// Returns the free variables of the expression.
    pub fn freevars(&self) -> BTreeSet<Symbol> {
        match *self {
            Prim::Fn(n, ref a, ref b) => {
                let mut fv = b.freevars();
                for n in &a.required {
                    fv.remove(n);
                }
                for &(ref n, _) in &a.optional {
                    fv.remove(n);
                }
                if let Some(n) = a.rest {
                    fv.remove(&n);
                }
                if let Some(n) = n {
                    fv.remove(&n);
                }
                fv
            },
            Prim::Lit(_) => BTreeSet::new(),
            Prim::Var(s) => {
                let mut fv = BTreeSet::new();
                fv.insert(s);
                fv
            },
            Prim::Vec(ref v) => {
                let mut fv = BTreeSet::new();
                for v in v {
                    fv.extend(&v.freevars());
                }
                fv
            },
        }
    }
}

impl Display for Prim {
    fn fmt(&self, fmt: &mut Formatter) -> FmtResult {
        match *self {
            Prim::Fn(n, ref a, ref b) => {
                fmt.write_str("fn")?;
                if let Some(n) = n {
                    fmt.write_char(' ')?;
                    Display::fmt(&n, fmt)?;
                }
                Display::fmt(a, fmt)?;
                fmt.write_str(" -> ")?;
                Display::fmt(b, fmt)
            },
            Prim::Lit(ref v) => Display::fmt(v, fmt),
            Prim::Var(ref n) => Display::fmt(n, fmt),
            Prim::Vec(ref v) => {
                fmt.write_char('[')?;
                fmt.write_str(&v.iter().join(" "))?;
                fmt.write_char(']')
            },
        }
    }
}