Very basic implementation of symbolic algebra in Rust

symbolic.rs

#![feature(globs)]

use std::fmt;
use std::hash::Hash;
use Expr::*;
use MaybeOwned::*;

// TODO:
// * multiple dispatch
// * sorting of symbols
// * ax + bx = (a+b)x for numbers a, b
// * Pow
// * x^a * x^b = x^(a+b) for numbers a, b
// * bigint instead of i64

///////////
// Types //
///////////

type Integer = i64;

#[deriving(Hash)]
enum Expr<'a> {
    Number(Integer),
    Symbol(&'static str),
    Sum(MaybeOwned<'a>, MaybeOwned<'a>),
    Product(MaybeOwned<'a>, MaybeOwned<'a>),
}

#[deriving(Hash)]
enum MaybeOwned<'a> {
    View(&'a Expr<'a>),
    Owned(Box<Expr<'a>>),
}

///////////////
// Constants //
///////////////

const ZERO: &'static Expr<'static> = &Number(0);
const ONE: &'static Expr<'static> = &Number(1);
const TWO: &'static Expr<'static> = &Number(2);


/////////////////
// Show and Eq //
/////////////////

impl<'a> fmt::Show for Expr<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Number(i) => try!(write!(f, "{}", i)),
            &Symbol(s) => try!(write!(f, "{}", s)),
            &Sum(ref a, ref b) => try!(write!(f, "({} + {})", a, b)),
            &Product(ref a, ref b) => try!(write!(f, "({} * {})", a, b)),
        }
        Ok(())
    }
}

impl<'a> PartialEq for Expr<'a> {
    fn eq(&self, other: &Expr<'a>) -> bool {
        match (self, other) {
            (&Number(ref a), &Number(ref b)) => a == b,
            (&Symbol(ref a), &Symbol(ref b)) => a == b,
            (&Sum(ref a, ref b), &Sum(ref c, ref d)) => a == c && b == d,
            (&Product(ref a, ref b), &Product(ref c, ref d)) => a == c && b == d,
            _ => false,
        }
    }
}

impl<'a> Eq for Expr<'a> {}

impl<'a> MaybeOwned<'a> {
    fn as_view<'b>(&'b self) -> &'b Expr<'a> {
        match *self {
            View(e) => e,
            Owned(box ref e) => e,
        }
    }
}

impl<'a> fmt::Show for MaybeOwned<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            View(ref e) => try!(write!(f, "{}", e)),
            Owned(ref e) => try!(write!(f, "{}", e)),
        }
        Ok(())
    }
}

impl<'a> PartialEq for MaybeOwned<'a> {
    fn eq(&self, other: &MaybeOwned<'a>) -> bool {
        self.as_view() == other.as_view()
    }
}

impl<'a> Eq for MaybeOwned<'a> {}

/////////////
// Flatten //
/////////////

impl<'a> Expr<'a> {
    fn flatten(&'a self) -> MaybeOwned<'a> {
        match *self {
            Sum(ref a, ref b) => {
                match (a.as_view(), b.as_view()) {
                    // (a + b) + (c + d) = ((a + b) + c) + d
                    (&Sum(ref a, ref b), &Sum(ref c, ref d)) => {
                        let a = a.as_view();
                        let b = b.as_view();
                        let c = c.as_view();
                        let d = d.as_view();
                        Owned(box Sum(
                            Owned(box Sum(
                                Owned(box Sum(View(a), View(b))),
                            View(c))),
                        View(d)))
                    },
                    // x + (y + z) = (x + y) + z
                    (x, &Sum(ref y, ref z)) => {
                        let y = y.as_view();
                        let z = z.as_view();
                        Owned(box Sum(Owned(box Sum(View(x), View(y))), View(z)))
                    },
                    _ => View(self),
                }
            },
            // FIXME: DRY
            Product(ref a, ref b) => {
                match (a.as_view(), b.as_view()) {
                    // (a * b) * (c * d) = ((a * b) * c) * d
                    (&Product(ref a, ref b), &Product(ref c, ref d)) => {
                        let a = a.as_view();
                        let b = b.as_view();
                        let c = c.as_view();
                        let d = d.as_view();
                        Owned(box Product(
                            Owned(box Product(
                                Owned(box Product(View(a), View(b))),
                            View(c))),
                        View(d)))
                    },
                    // x * (y * z) = (x * y) * z
                    (x, &Product(ref y, ref z)) => {
                        let y = y.as_view();
                        let z = z.as_view();
                        Owned(box Product(Owned(box Product(View(x), View(y))), View(z)))
                    },
                    _ => View(self),
                }
            }
            _ => View(self),
        }
    }
}

impl<'a> MaybeOwned<'a> {
    fn flatten(&'a self) -> MaybeOwned<'a> {
        self.as_view().flatten()
    }
}


//////////////
// Addition //
//////////////

impl<'a> Expr<'a> {
    fn add(&'a self, rhs: &'a Expr<'a>) -> MaybeOwned<'a> {
        match (self, rhs) {
            (&Number(a), &Number(b))
                => Owned(box Number(a + b)),
            (a, b) if a == b
                => Owned(box Product(View(TWO), View(a))),
            (a, b) if b == ZERO
                => View(a),
            (a, b) if a == ZERO
                => View(b),
            (a, b)
                => Owned(box Sum(View(a), View(b))),
        }
    }
}

impl<'a> MaybeOwned<'a> {
    fn add(&'a self, rhs: &'a Expr<'a>) -> MaybeOwned<'a> {
        self.as_view().add(rhs)
    }
}


////////////////////
// Multiplication //
////////////////////

impl<'a> Expr<'a> {
    fn mul(&'a self, rhs: &'a Expr<'a>) -> MaybeOwned<'a> {
        match (self, rhs) {
            (&Number(a), &Number(b))
                => Owned(box Number(a * b)),
            (a, b) if b == ONE
                => View(a),
            (a, b) if a == ONE
                => View(b),
            (a, b)
                => Owned(box Product(View(a), View(b))),
        }
    }
}

impl<'a> MaybeOwned<'a> {
    fn mul(&'a self, rhs: &'a Expr<'a>) -> MaybeOwned<'a> {
        self.as_view().mul(rhs)
    }
}


////////////////////////
// Tests and Examples //
////////////////////////

fn main() {
    let zero = Number(0);
    let one = Number(1);
    let two = Number(2);
    let t = Symbol("t");
    let x = Symbol("x");
    let y = Symbol("y");
    let z = Symbol("z");

    // test equality
    assert!(zero != one);
    assert!(zero == zero);
    assert!(x != zero);
    assert!(x == x);

    // test addition
    assert!(one.add(&two).as_view() == &Number(3));
    assert!(x.add(&x).as_view() == &Product(View(TWO), View(&x)));

    assert!(x.add(&zero).as_view() == &x);
    assert!(zero.add(&x).as_view() == &x);

    // test multiplication
    assert!(one.mul(&two).as_view() == &two);
    assert!(two.mul(&one).as_view() == &two);
    assert!(two.mul(&x).as_view() == &Product(View(TWO), View(&x)));

    assert!(x.mul(&one).as_view() == &x);
    assert!(one.mul(&x).as_view() == &x);

    // work in progress
    println!("  x + x + x");
    println!("= {}", x.add(&x).add(&x));
    println!("");
    println!("  {}", x.add(y.add(&z).as_view()));
    println!("= {}", x.add(y.add(&z).as_view()).flatten());
    println!("");
    println!("  {}", (t.add(&x)).add(y.add(&z).as_view()));
    println!("= {}", (t.add(&x)).add(y.add(&z).as_view()).flatten());
    println!("");
    println!("  {}", x.mul(x.mul(&x).as_view()));
    println!("= {}", x.mul(x.mul(&x).as_view()).flatten());
    println!("");
    println!("  {}", (t.mul(&x)).mul(y.mul(&z).as_view()));
    println!("= {}", (t.mul(&x)).mul(y.mul(&z).as_view()).flatten());
}