main.rs

pub use ::std::io::prelude::*;
pub use ::std::io::Result;

pub trait Node
    : Copy + Ord + Eq
{
    fn add(self, b: Self) -> Self;
    fn sub(self, b: Self) -> Self;
    fn neg(self) -> Self;

    fn get_self() -> Self;
    fn get_assocs() -> Vec<Self>
    {
        let mut v = Self::get_all_special_tys();
        let i = v.iter().position(|ty| ty.is_self()).unwrap();
        v.remove(i);
        v
    }
    // Get the self ty and all associated types
    fn get_all_special_tys() -> Vec<Self>;

    fn classify(self) -> NodeClass {
        // NOTE: this default definition peforms an allocation and O(n) search that
        //       are avoidable; but I don't think this will be a big deal in practice.
        if self.is_self() {
            NodeClass::SelfTy
        } else {
            match Self::get_all_special_tys().iter().find(|x| x == &&self) {
                Some(_) => NodeClass::AssocTy,
                None => NodeClass::NotSpecial,
            }
        }
    }

    fn is_self(self) -> bool
    { self == Self::get_self() }

    fn output(self, edge: Edge<Self>) -> Self
    { match edge {
        Edge::BinOp { op: BinOpKind::Add, arg: b } => self.add(b),
        Edge::BinOp { op: BinOpKind::Sub, arg: b } => self.sub(b),
        Edge::UnOp { op: UnOpKind::Neg } => self.neg(),
    }}

    fn assoc_name(self, set: TraitSet) -> Option<&'static str>;
    fn format_trait(set: TraitSet) -> String;
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Ty1(i8);

impl Node for Ty1 {
    fn format_trait(set: TraitSet) -> String
    { match set {
        TraitSet::AddSub => format!("AddSub1"),
        TraitSet::MulDiv => format!("MulDiv1"),
    }}

    fn add(self, b: Ty1) -> Ty1 { Ty1(self.0 + b.0) }
    fn sub(self, b: Ty1) -> Ty1 { Ty1(self.0 - b.0) }
    fn neg(self) -> Ty1 { Ty1(-self.0) }

    fn get_self() -> Self { Ty1(1) }
    fn get_all_special_tys() -> Vec<Self> { (-3..3+1).map(Ty1).collect() }

    fn assoc_name(self, set: TraitSet) -> Option<&'static str> {
        match (set, self) {
            (TraitSet::AddSub, Ty1( 0)) => Some("Zero"),
            (TraitSet::AddSub, Ty1( 2)) => Some("Twice"),
            (TraitSet::AddSub, Ty1( 3)) => Some("Thrice"),
            (TraitSet::AddSub, Ty1(-1)) => Some("Neg"),
            (TraitSet::AddSub, Ty1(-2)) => Some("NTwice"),
            (TraitSet::AddSub, Ty1(-3)) => Some("NThrice"),

            (TraitSet::MulDiv, Ty1( 0)) => Some("One"),
            (TraitSet::MulDiv, Ty1( 2)) => Some("Square"),
            (TraitSet::MulDiv, Ty1( 3)) => Some("Cube"),
            (TraitSet::MulDiv, Ty1(-1)) => Some("Recip"),
            (TraitSet::MulDiv, Ty1(-2)) => Some("RSquare"),
            (TraitSet::MulDiv, Ty1(-3)) => Some("RCube"),

            _ => None,
        }
    }
}

/// Powers of Self and B
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Ty2(i8, i8);

impl Node for Ty2 {
    fn format_trait(set: TraitSet) -> String
    {
        // HACK: Hardcoding the type parameter name.
        //       This mostly works just due to luck.

        // This string gets used in the following places:
        //   * "pub trait {}"
        //   * "impl<...> {} for ..."
        //   * "<Self as {}>"
        //
        // In the trait definition, this just declares `OtherP` as a new type parameter.
        // In the impl, `OtherP` is already a type parameter, so it binds to that.
        // In assoc types, it works out fine simply thanks to the above two facts.
        match set {
            TraitSet::AddSub => format!("AddSub2<OtherP>"),
            TraitSet::MulDiv => format!("MulDiv2<OtherP>"),
        }
    }

    fn add(self, b: Ty2) -> Ty2 { Ty2(self.0 + b.0, self.1 + b.1) }
    fn sub(self, b: Ty2) -> Ty2 { Ty2(self.0 - b.0, self.1 - b.1) }
    fn neg(self) -> Ty2 { Ty2(-self.0, -self.1) }

    fn get_self() -> Self { Ty2(1, 0) }
    fn get_all_special_tys() -> Vec<Self> {
        [0, 1, -1].iter().flat_map(|&j| {
            [0, 1, -1].iter().map(move |&k| Ty2(j, k))
        }).collect()
    }

    fn assoc_name(self, set: TraitSet) -> Option<&'static str> {
        match (set, self) {
            (TraitSet::AddSub, Ty2( 0,  0)) => Some("Zero"),
            (TraitSet::AddSub, Ty2(-1,  0)) => Some("Neg"),
            (TraitSet::AddSub, Ty2( 0,  1)) => Some("Other"),
            (TraitSet::AddSub, Ty2( 0, -1)) => Some("NOther"),
            (TraitSet::AddSub, Ty2( 1,  1)) => Some("Add"),
            (TraitSet::AddSub, Ty2(-1, -1)) => Some("NAdd"),
            (TraitSet::AddSub, Ty2( 1, -1)) => Some("Sub"),
            (TraitSet::AddSub, Ty2(-1,  1)) => Some("NSub"),

            (TraitSet::MulDiv, Ty2( 0,  0)) => Some("One"),
            (TraitSet::MulDiv, Ty2(-1,  0)) => Some("Recip"),
            (TraitSet::MulDiv, Ty2( 0,  1)) => Some("Other"),
            (TraitSet::MulDiv, Ty2( 0, -1)) => Some("ROther"),
            (TraitSet::MulDiv, Ty2( 1,  1)) => Some("Mul"),
            (TraitSet::MulDiv, Ty2(-1, -1)) => Some("RMul"),
            (TraitSet::MulDiv, Ty2( 1, -1)) => Some("Div"),
            (TraitSet::MulDiv, Ty2(-1,  1)) => Some("RDiv"),

            _ => None,
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum NodeClass {
    SelfTy,
    AssocTy,
    NotSpecial,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum TraitSet { AddSub, MulDiv }

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum BinOpKind { Add, Sub }
impl BinOpKind {
    pub fn name(self, set: TraitSet) -> &'static str
    { match (set, self) {
        (TraitSet::AddSub, BinOpKind::Add) => "Add",
        (TraitSet::AddSub, BinOpKind::Sub) => "Sub",
        (TraitSet::MulDiv, BinOpKind::Add) => "Mul",
        (TraitSet::MulDiv, BinOpKind::Sub) => "Div",
    }}
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum UnOpKind { Neg }
impl UnOpKind {
    pub fn name(self, set: TraitSet) -> &'static str
    { match (set, self) {
        (TraitSet::AddSub, UnOpKind::Neg) => "Neg",
        (TraitSet::MulDiv, UnOpKind::Neg) => "Recip",
    }}
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum Edge<T> {
    BinOp { op: BinOpKind, arg: T },
    UnOp { op: UnOpKind, }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum TraitBound<T> {
    Edge {
        edge: Edge<T>,
        output: Option<T>
    },
    General(String),
}

pub type Bounds<T> = Vec<TraitBound<T>>;
pub type FullBoundSet<T> = Vec<(T, Bounds<T>)>;

pub fn compute_full_bound_set<T: Node>() -> FullBoundSet<T> {

    T::get_all_special_tys().into_iter().map(|a| {
        let mut bounds = vec![];

        if let NodeClass::SelfTy = a.classify() {
            bounds.push(TraitBound::General("Sized".to_owned()));
        }

        T::get_all_special_tys().into_iter().for_each(|arg| {
            for &op in &[BinOpKind::Add, BinOpKind::Sub] {
                let edge = Edge::BinOp { op, arg };

                // Only include equality bounds when the output type is special
                let output = a.output(edge);
                let output = match output.classify() {
                    NodeClass::SelfTy => Some(output),
                    NodeClass::AssocTy => Some(output),
                    NodeClass::NotSpecial => None,
                };

                bounds.push(TraitBound::Edge { edge, output })
            }
        });

        (a, bounds)
    }).collect()
}

mod format {
    use super::*;

    // abstracts over impl vs trait definition
    pub trait Format<T: Node> {
        fn format_type(&self, set: TraitSet, ty: T) -> Option<String>;
    }

    pub struct Formatter<F, T: Node> {
        pub fmt: F,
        pub set: TraitSet,
        // in the future Formatter might hold data dependent on T, so might as well
        // make it generic now
        pub _marker: ::std::marker::PhantomData<T>,
    }

    impl<F: Format<T>, T: Node> Formatter<F, T> {
        pub fn ty(&self, ty: T) -> String
        {
            self.fmt.format_type(self.set, ty)
                .expect("Could not format type!")
        }

        pub fn assoc_name(&self, ty: T) -> Option<&'static str>
        { ty.assoc_name(self.set) }

        pub fn main_trait_name(&self) -> String
        { T::format_trait(self.set) }

        pub fn bound(&self, b: TraitBound<T>) -> String
        {
            match b {
                TraitBound::General(s) => s,
                TraitBound::Edge { edge, output } => {
                    match (edge, output) {
                        (Edge::BinOp { op, arg }, None) => {
                            format!("{}<{}>", op.name(self.set), self.ty(arg))
                        },
                        (Edge::UnOp { op }, None) => {
                            format!("{}", op.name(self.set))
                        },
                        (Edge::BinOp { op, arg }, Some(out)) => {
                            format!("{}<{}, Output={}>", op.name(self.set), self.ty(arg), self.ty(out))
                        },
                        (Edge::UnOp { op }, Some(out)) => {
                            format!("{}<Output={}>", op.name(self.set), self.ty(out))
                        },
                    }
                }
            }
        }
    }

    /// Formats types as `Self` or `<Self as Trait>::AssocTy`
    #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
    pub struct DefnFormat;

    /// Formats types as dummy type parameter names (`A`, `AssocTyP`...)
    #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
    pub struct ImplFormat;

    impl<T: Node> Format<T> for DefnFormat {
        fn format_type(&self, set: TraitSet, ty: T) -> Option<String>
        { match ty.classify() {
            NodeClass::SelfTy => Some(format!("Self")),
            NodeClass::AssocTy => {
                let tr = T::format_trait(set);
                let assoc = ty.assoc_name(set).unwrap();
                Some(format!("<Self as {}>::{}", tr, assoc))
            },
            NodeClass::NotSpecial => None,
        }}
    }

    impl<T: Node> Format<T> for ImplFormat {
        fn format_type(&self, set: TraitSet, ty: T) -> Option<String>
        { match ty.classify() {
            NodeClass::SelfTy => Some(format!("A")),
            NodeClass::AssocTy => {
                let assoc = ty.assoc_name(set).unwrap();
                Some(format!("{}P", assoc))
            },
            NodeClass::NotSpecial => None,
        }}
    }
}

pub struct Printer<W> {
    writer: W,
    indent: String,
}

impl<W: Write> Printer<W> {
    pub fn writeln<S: AsRef<str>>(&mut self, s: S) -> Result<()> {
        writeln!(self.writer, "{}{}", self.indent, s.as_ref())
    }

    pub fn indented(&mut self) -> Printer<&mut W>
    {
        Printer {
            writer: &mut self.writer,
            indent: self.indent.clone() + "    ",
        }
    }

    pub fn format_bounds(&mut self, bounds: &[String]) -> Result<()>
    {Ok({
        let mut it = bounds.iter();
        if let Some(b) = it.next() {
            self.writeln(format!(": {}", b))?;
        }
        for b in it {
            self.writeln(format!("+ {}", b))?;
        }
    })}
}

fn generate<T: Node, W: Write>(p: &mut Printer<W>, set: TraitSet) -> Result<()>
{Ok({
    // bounds in an abstract form that doesn't format the types
    let bounds = compute_full_bound_set::<T>();

    // Trait definition
    {
        let fmt = format::Formatter::<_, T> {
            fmt: format::DefnFormat,
            set,
            _marker: Default::default(),
        };

        p.writeln("")?;
        p.writeln(format!("pub trait {}", fmt.main_trait_name()))?;

        let bounds: ::std::collections::BTreeMap<_,_>
            = bounds.iter().map(|&(ty, ref bounds)| {
                (ty, bounds.iter().map(|b| fmt.bound(b.clone())).collect::<Vec<_>>())
            }).collect();

        p.indented().format_bounds(&bounds[&T::get_self()])?;

        p.writeln("{")?;
        {
            let mut p = p.indented();
            for ty in T::get_assocs() {
                p.writeln(format!("type {}", fmt.assoc_name(ty).unwrap()))?;

                p.indented().format_bounds(&bounds[&ty])?;
                p.indented().writeln(";")?;
                p.writeln("")?;
            }
        }
        p.writeln("}")?;
    }

    {
        let fmt = format::Formatter::<_, T> {
            fmt: format::ImplFormat,
            set,
            _marker: Default::default(),
        };

        p.writeln("")?;
        p.writeln(format!(
            "impl<{}> {} for {}",
            {
                T::get_all_special_tys()
                    .into_iter()
                    .map(|ty| fmt.ty(ty))
                    .collect::<Vec<_>>()
                    .join(", ")
            },
            fmt.main_trait_name(),
            fmt.ty(T::get_self()),
        ))?;

        p.writeln("where")?;
        {
            let mut p = p.indented();

            for &(ty, ref bounds) in &bounds {
                let bounds = bounds.iter().map(|b| fmt.bound(b.clone())).collect::<Vec<_>>();

                p.writeln(fmt.ty(ty))?;
                p.indented().format_bounds(&bounds)?;
                p.indented().writeln(",")?;
                p.indented().writeln("")?;
            }
        }

        p.writeln("{")?;
        {
            let mut p = p.indented();
            for ty in T::get_assocs() {
                p.writeln(format!(
                    "type {} = {};",
                    fmt.assoc_name(ty).unwrap(),
                    fmt.ty(ty),
                ))?;
            }
        }
        p.writeln("}")?;
    }
})}


fn _main() -> Result<()> {Ok({

    let p = ::std::io::stdout();
    let p = p.lock();
    let mut p = Printer { writer: p, indent: String::new() };

    p.writeln("use ::std::ops::{Add, Sub, Neg, Mul, Div};")?;
    p.writeln("use ::dim::Recip;")?;

    generate::<Ty1, _>(&mut p, TraitSet::AddSub)?;
    generate::<Ty1, _>(&mut p, TraitSet::MulDiv)?;
    generate::<Ty2, _>(&mut p, TraitSet::AddSub)?;
    generate::<Ty2, _>(&mut p, TraitSet::MulDiv)?;
})}

fn main() -> () {
    _main().unwrap();
}