oftlisp-anf 0.1.3

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

use byteorder::{ByteOrder, LE};
use gc::Trace;
use oftlisp::{CompilerContext, Symbol, Value};
use oftlisp::ast::Args;
use oftlisp::modules::ExportedModule;
use smallvec::SmallVec;

use errors::CompileError;
use types::{Context, Expr, Prim};

const ANFIR_VERSION: u32 = 0x001_00_000;

impl CompilerContext for Context {
    type CompileError = CompileError;

    fn compile(main_module: Symbol, mut exported_modules: HashMap<Symbol, ExportedModule<Self>>) -> Result<Vec<u8>, Self::CompileError> {
        // First, convert the modules to serialize into a Vec. We can also
        // detect some issues in the dependency graph here.
        let mut modules = vec![];
        let mut moved = HashSet::new();
        move_modules(main_module, &mut exported_modules, &mut modules, &mut moved)?;

        // Warn about any unused modules.
        for (name, _) in exported_modules {
            warn!("Unused module: {}", name);
        }

        // Convert `modules.len()` to a u32.
        if modules.len() > u32::MAX as usize {
            return Err(CompileError::TooManyModules);
        }
        let num_modules = modules.len() as u32;

        // Serialize the modules.
        let mut buf = vec![0; 8];
        LE::write_u32(&mut buf[0..4], ANFIR_VERSION);
        LE::write_u32(&mut buf[4..8], num_modules);
        for (name, m) in modules {
            serialize_module(name, &m, &mut buf)?;
        }
        Ok(buf)
    }
}

fn move_modules(name: Symbol, exported: &mut HashMap<Symbol, ExportedModule<Context>>, modules: &mut Vec<(Symbol, ExportedModule<Context>)>, moved: &mut HashSet<Symbol>) -> Result<(), CompileError> {
    if moved.contains(&name) {
        return Ok(());
    }

    if let Some(m) = exported.remove(&name) {
        let imports = m.imports.iter()
            .map(|(&n, _)| n)
            .collect::<SmallVec<[Symbol; 64]>>();
        modules.push((name, m));
        moved.insert(name);
        for name in imports {
            move_modules(name, exported, modules, moved)?;
        }
        Ok(())
    } else {
        Err(CompileError::MissingModule(name))
    }
}

fn serialize_module(name: Symbol, m: &ExportedModule<Context>, buf: &mut Vec<u8>) -> Result<(), CompileError> {
    let mut out = vec![];
    name.serialize(&mut out)?;

    if m.exports.len() > u32::MAX as usize {
        return Err(CompileError::TooManyExports(name));
    } else {
        let mut len = [0, 0, 0, 0];
        LE::write_u32(&mut len, m.exports.len() as u32);
        out.extend(&len);
    }
    for &ex in &m.exports {
        ex.serialize(&mut out)?;
    }

    if m.imports.len() > u32::MAX as usize {
        return Err(CompileError::TooManyImports(name));
    } else {
        let mut len = [0, 0, 0, 0];
        LE::write_u32(&mut len, m.imports.len() as u32);
        out.extend(&len);
    }
    for (&im, syms) in &m.imports {
        im.serialize(&mut out)?;

        if syms.len() > u32::MAX as usize {
            return Err(CompileError::TooManyImportedSymbols(name, im));
        } else {
            let mut len = [0, 0, 0, 0];
            LE::write_u32(&mut len, syms.len() as u32);
            out.extend(&len);
        }
        for &sym in syms {
            sym.serialize(&mut out)?;
        }
    }

    if m.body.len() > u32::MAX as usize {
        return Err(CompileError::TooManyDecls(name));
    } else {
        let mut len = [0, 0, 0, 0];
        LE::write_u32(&mut len, m.body.len() as u32);
        out.extend(&len);
    }
    for &(name, ref expr) in &m.body {
        name.serialize(&mut out)?;
        expr.serialize(&mut out)?;
    }

    buf.extend(out);
    Ok(())
}

trait Serializable: Trace {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError>;
}

impl Serializable for Symbol {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError> {
        if self.len() > u16::MAX as usize {
            Err(CompileError::SymbolTooLong(*self))
        } else {
            let s = self.as_str();
            let mut len = [0, 0];
            LE::write_u16(&mut len, s.len() as u16);
            buf.extend(&len);
            buf.extend(s.bytes());
            Ok(())
        }
    }
}

impl Serializable for Option<Symbol> {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError> {
        match *self {
            Some(sym) => sym.serialize(buf),
            None => {
                buf.push(0);
                buf.push(0);
                Ok(())
            },
        }
    }
}

impl Serializable for Args<Context> {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError> {
        if self.required.len() > u16::MAX as usize {
            return Err(CompileError::TooManyArgsInDefn);
        } else {
            let mut len = [0, 0];
            LE::write_u16(&mut len, self.required.len() as u16);
            buf.extend(&len);
            for sym in &self.required {
                sym.serialize(buf)?;
            }
        }

        if self.optional.len() > u16::MAX as usize {
            return Err(CompileError::TooManyArgsInDefn);
        } else {
            let mut len = [0, 0];
            LE::write_u16(&mut len, self.optional.len() as u16);
            buf.extend(&len);
            for &(sym, ref val) in &self.optional {
                sym.serialize(buf)?;
                val.serialize(buf)?;
            }
        }

        self.rest.serialize(buf)
    }
}

impl Serializable for Expr {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError> {
        match *self {
            Expr::Call(ref f, ref a) => {
                buf.push(0x80);
                f.serialize(buf)?;
                if a.len() > u16::MAX as usize {
                    Err(CompileError::TooManyArgsInCall)
                } else {
                    let mut len = [0, 0];
                    LE::write_u16(&mut len, a.len() as u16);
                    buf.extend(&len);
                    for a in a {
                        a.serialize(buf)?;
                    }
                    Ok(())
                }
            },
            Expr::If(ref c, ref t, ref e) => {
                buf.push(0xc0);
                c.serialize(buf)?;
                t.serialize(buf)?;
                e.serialize(buf)
            },
            Expr::Let(n, ref x, ref y) => {
                buf.push(0xf0);
                n.serialize(buf)?;
                x.serialize(buf)?;
                y.serialize(buf)
            },
            Expr::Prim(ref prim) => prim.serialize(buf),
        }
    }
}

impl Serializable for Prim {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError> {
        match *self {
            Prim::Fn(ref n, ref a, ref b) => {
                buf.push(0x10);
                n.serialize(buf)?;
                a.serialize(buf)?;
                b.serialize(buf)
            },
            Prim::Lit(ref v) => v.serialize(buf),
            Prim::Var(ref n) => {
                buf.push(0x20);
                n.serialize(buf)
            },
            Prim::Vec(ref ps) => {
                buf.push(0x30);
                if ps.len() > u16::MAX as usize {
                    Err(CompileError::TooManyArgsInCall)
                } else {
                    let mut len = [0, 0];
                    LE::write_u16(&mut len, ps.len() as u16);
                    buf.extend(&len);
                    for p in ps {
                        p.serialize(buf)?;
                    }
                    Ok(())
                }
            },
        }
    }
}

impl Serializable for Value<Context> {
    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CompileError> {
        match *self {
            Value::AtomicWord(..) => Err(CompileError::CannotSerialize(self.clone())),
            Value::BuiltinFunction(..) => Err(CompileError::CannotSerialize(self.clone())),
            Value::Byte(n, _) => {
                buf.push(0x01);
                buf.push(n);
                Ok(())
            },
            Value::Bytes(ref bs, _) => {
                buf.push(0x02);
                if bs.len() > u32::MAX as usize {
                    Err(CompileError::BytesTooLong(bs.clone()))
                } else {
                    let mut len = [0, 0, 0, 0];
                    LE::write_u32(&mut len, bs.len() as u32);
                    buf.extend(&len);
                    buf.extend(bs.as_ref() as &[u8]);
                    Ok(())
                }
            },
            Value::Cons(ref h, ref t, _) => {
                buf.push(0x03);
                h.serialize(buf)?;
                t.serialize(buf)
            },
            Value::Fixnum(n, _) => {
                buf.push(0x04);
                let mut n_buf = [0, 0, 0, 0, 0, 0, 0, 0];
                LE::write_i64(&mut n_buf, n as i64);
                buf.extend(&n_buf);
                Ok(())
            },
            Value::Func(..) => Err(CompileError::CannotSerialize(self.clone())),
            Value::Object(..) => Err(CompileError::CannotSerialize(self.clone())),
            Value::Nil(_) => {
                buf.push(0x00);
                Ok(())
            },
            Value::String(ref s, _) => {
                buf.push(0x05);
                if s.len() > u32::MAX as usize {
                    Err(CompileError::StringTooLong(s.clone()))
                } else {
                    let mut len = [0, 0, 0, 0];
                    LE::write_u32(&mut len, s.len() as u32);
                    buf.extend(&len);
                    buf.extend(s.bytes());
                    Ok(())
                }
            },
            Value::Symbol(s, _) => {
                buf.push(0x06);
                s.serialize(buf)
            },
            Value::Vector(ref vs, _) => {
                buf.push(0x07);
                if vs.len() > u32::MAX as usize {
                    Err(CompileError::VectorTooLong(vs.clone()))
                } else {
                    let mut len = [0, 0, 0, 0];
                    LE::write_u32(&mut len, vs.len() as u32);
                    buf.extend(&len);
                    for v in vs {
                        v.serialize(buf)?;
                    }
                    Ok(())
                }
            },
        }
    }
}