main.rs

use std::env;
use std::fmt;
use std::ops::{Add, Mul};
use std::rc::Rc;

thread_local! {
    static MINUS_ONE_EXPR: Expr = Expr::int(-1);
    static ONE_EXPR: Expr = Expr::int(1);
    static ZERO_EXPR: Expr = Expr::int(0);
}

#[derive(Clone, Debug)]
enum _Expr {
    Int(i64),
    Var(String),
    Add(Expr, Expr),
    Mul(Expr, Expr),
    Pow(Expr, Expr),
    Ln(Expr),
}

#[derive(Clone, Debug)]
pub struct Expr(Rc<_Expr>);

impl From<_Expr> for Expr {
    fn from(e: _Expr) -> Self {
        Expr(e.into())
    }
}

fn pown(base: i64, expt: i64) -> i64 {
    match expt {
        0 => 1,
        1 => base,
        _ => {
            let b = pown(base, expt / 2);
            b * b * if expt % 2 == 0 { 1 } else { base }
        }
    }
}

fn decompose_add(_expr: &Expr) -> (i64, Option<&Expr>) {
    match &*_expr.0 {
        _Expr::Int(n) => (*n, None),
        _Expr::Add(f, g) => match *f.0 {
            _Expr::Int(n) => (n, Some(g)),
            _ => (0, Some(_expr)),
        },
        _ => (0, Some(_expr)),
    }
}

fn decompose_mul(_expr: &Expr) -> (i64, Option<&Expr>) {
    match &*_expr.0 {
        _Expr::Int(n) => (*n, None),
        _Expr::Mul(f, g) => match *f.0 {
            _Expr::Int(n) => (n, Some(g)),
            _ => (1, Some(_expr)),
        },
        _ => (1, Some(_expr)),
    }
}

impl Add<&Expr> for &Expr {
    type Output = Expr;

    fn add(self, other: &Expr) -> Expr {
        use crate::_Expr::*;

        let (m, fo) = decompose_add(self);
        let (n, go) = decompose_add(other);

        let mn = m + n;

        let fg = match (fo, go) {
            (None, None) => return Int(mn).into(),
            (Some(f), None) => f.clone(),
            (None, Some(g)) => g.clone(),
            (Some(f), Some(g)) => Add(f.clone(), g.clone()).into(),
        };

        if mn == 0 {
            fg
        } else {
            Add(Int(mn).into(), fg).into()
        }
    }
}

impl Mul<&Expr> for &Expr {
    type Output = Expr;

    fn mul(self, other: &Expr) -> Expr {
        use crate::_Expr::*;

        let (m, fo) = decompose_mul(self);
        let (n, go) = decompose_mul(other);

        let mn = m * n;

        if mn == 0 {
            return ZERO_EXPR.with(Clone::clone);
        }

        let fg = match (fo, go) {
            (None, None) => return Int(mn).into(),
            (Some(f), None) => f.clone(),
            (None, Some(g)) => g.clone(),
            (Some(f), Some(g)) => Mul(f.clone(), g.clone()).into(),
        };

        if mn == 1 {
            fg
        } else {
            Mul(Int(mn).into(), fg).into()
        }
    }
}

impl Expr {
    pub fn int(i: i64) -> Self {
        Expr(_Expr::Int(i).into())
    }

    pub fn var(s: &str) -> Self {
        Expr(_Expr::Var(s.to_owned()).into())
    }

    pub fn pow(&self, other: &Expr) -> Expr {
        use crate::_Expr::*;

        match (&*self.0, &*other.0) {
            (Int(m), Int(n)) => Int(pown(*m, *n)).into(),
            (_, Int(0)) => ONE_EXPR.with(Clone::clone),
            (_a, Int(1)) => self.clone(),
            (Int(0), _) => ZERO_EXPR.with(Clone::clone),
            (_a, _b) => Pow(self.clone(), other.clone()).into(),
        }
    }

    pub fn ln(&self) -> Expr {
        use crate::_Expr::*;

        match &*self.0 {
            Int(1) => ZERO_EXPR.with(Clone::clone),
            _a => Ln(self.clone()).into(),
        }
    }

    pub fn d(&self, x: &str) -> Expr {
        use crate::_Expr::*;

        match &*self.0 {
            Var(y) => {
                if *x == **y {
                    ONE_EXPR.with(Clone::clone)
                } else {
                    ZERO_EXPR.with(Clone::clone)
                }
            }
            Int(_) => ZERO_EXPR.with(Clone::clone),
            Add(f, g) => &f.d(x) + &g.d(x),
            Mul(f, g) => &(f * &g.d(x)) + &(g * &f.d(x)),
            Pow(f, g) => {
                &(self * &(g * &(&f.d(x) * &MINUS_ONE_EXPR.with(|minus_one| f.pow(minus_one)))))
                    + &(&f.ln() * &g.d(x))
            }
            Ln(f) => &f.d(x) * &MINUS_ONE_EXPR.with(|minus_one| f.pow(minus_one)),
        }
    }

    pub fn count(&self) -> usize {
        use crate::_Expr::*;

        match &*self.0 {
            Var(_) => 1,
            Int(_) => 1,
            Add(f, g) => f.count() + g.count(),
            Mul(f, g) => f.count() + g.count(),
            Pow(f, g) => f.count() + g.count(),
            Ln(f) => f.count(),
        }
    }

    fn format(&self, f: &mut fmt::Formatter, old_prec: usize) -> fmt::Result {
        use crate::_Expr::*;

        fn bracket<F>(
            f: &mut fmt::Formatter,
            old_prec: usize,
            new_prec: usize,
            body: F,
        ) -> fmt::Result
        where
            F: FnOnce(&mut fmt::Formatter) -> fmt::Result,
        {
            if new_prec < old_prec {
                f.write_str("(")?;
            }
            body(f)?;
            if new_prec < old_prec {
                f.write_str(")")?;
            }
            Ok(())
        }

        match &*self.0 {
            Var(name) => f.write_str(&*name),

            Int(n) => write!(f, "{}", n),

            Add(g, h) => bracket(f, old_prec, 1, |f| {
                g.format(f, 1)?;
                f.write_str(" + ")?;
                h.format(f, 1)
            }),

            Mul(g, h) => bracket(f, old_prec, 2, |f| {
                g.format(f, 2)?;
                f.write_str("*")?;
                h.format(f, 2)
            }),

            Pow(g, h) => bracket(f, old_prec, 3, |f| {
                g.format(f, 2)?;
                f.write_str("^")?;
                h.format(f, 3)
            }),

            Ln(g) => {
                f.write_str("ln(")?;
                g.format(f, 1)?;
                f.write_str(")")
            }
        }
    }
}

impl fmt::Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let n = self.count();
        if n > 100 {
            write!(f, "<<{}>>", n)
        } else {
            self.format(f, 1)
        }
    }
}

fn nest<A, F: Fn(A) -> A>(n: usize, f: F, mut x: A) -> A {
    for _ in 0..n {
        x = f(x);
    }

    x
}

fn deriv(f: Expr) -> Expr {
    let d = f.d("x");
    println!("D({}) = {}", f, d);
    d
}

fn main() {
    let x: Expr = _Expr::Var("x".into()).into();
    let f = x.clone().pow(&x);
    let n = env::args().nth(1).expect("n").parse().unwrap();
    nest(n, deriv, f);
}