path: root/src/modelica_ast.rs

use std::clone::Clone;
use std::fmt::{Debug, Formatter, Error};
use std::collections::HashMap;

#[derive(PartialEq)]
pub struct ModelicaModel {
    pub name: String,
    pub components: Vec<Component>,
    pub equations: Vec<SimpleEquation>,
    pub connections: Vec<Connection>,
    pub extends: Vec<String>,
}

#[derive(Copy, Clone, PartialEq)]
pub enum ComponentPrefix {
    // incomplete: eg, can be parameter and input
    Flow,
    Stream,
    Input,
    Output,
    Discrete,
    Parameter,
    Constant,
}

#[derive(Clone, PartialEq)]
pub struct Component {
    pub prefix: Option<ComponentPrefix>,
    pub specifier: String,
    pub name: String,
    pub value: Option<Expr>,
    pub units: Option<String>,
    pub description: Option<String>,
}

#[derive(Clone, PartialEq)]
pub struct Connection {
    pub a: String,
    pub b: String,
}

#[derive(Clone, PartialEq)]
pub struct SimpleEquation {
    pub lhs: Expr,
    pub rhs: Expr,
}

#[derive(Clone, PartialEq)]
pub enum Expr {
    Integer(i64),
    Float(f64),
    Ident(String),
    Der(Box<Expr>),
    Abs(Box<Expr>),
    BinExpr(BinOperator, Box<Expr>, Box<Expr>),
}

#[derive(Copy, Clone, PartialEq)]
pub enum BinOperator {
    Multiply,
    Divide,
    Add,
    Subtract,
}

//// Helpers

impl ModelicaModel {

    pub fn get_constant_vars(&self) -> HashMap<String,Expr> {
        let mut binds = HashMap::new();
        for c in &self.components {
            match c.prefix {
                Some(ComponentPrefix::Constant) => { binds.insert(c.name.clone(), Expr::Integer(123)); },
                Some(ComponentPrefix::Parameter) => { binds.insert(c.name.clone(), Expr::Float(4.56)); },
                _ => (),
            }
        }
        binds
    }

    // This crude function finds "unbound" variables: those which are not constants, parameters, or
    // the sole element on the LHS of an equation.
    // Bugs:
    //   if a var is on LHS and RHS of same equation
    pub fn get_free_vars(&self) -> Vec<String> {
        // Start with components, and remove constants and parameters
        let vars = self.components.iter().filter(|v| match v.prefix {
                Some(ComponentPrefix::Constant) | Some(ComponentPrefix::Parameter) => false,
                _ => true,
            });

        // Remove LHS (bound) vars
        let mut outputs = vec![];
        for eq in self.equations.iter() {
            // TODO:
            if let Expr::Ident(ref symb) = eq.lhs {
                outputs.push(symb.to_string());
            }
        }
        let vars = vars.filter(|v| !outputs.contains(&v.name));

        vars.map(|c| c.name.clone()).collect()
    }

fn union_vecs(a: &Vec<String>, b: &Vec<String>) -> Vec<String> {
    let mut u = a.clone();
    for e in b {
        if !(u.contains(&e)) {
            u.push(e.clone());
        }
    }
    u
}

impl Expr {

    // Order is undefined
    pub fn identifiers(&self) -> Vec<String> {
        use self::Expr::*;
        match *self {
            Integer(_) | Float(_) => vec![],
            Ident(ref s) => vec![s.clone()],
            Der(ref e) | Abs(ref e) => e.identifiers(),
            BinExpr(_, ref e1, ref e2) => {
                union_vecs(&e1.identifiers(), &e2.identifiers())
            },
        }
    }

    pub fn contains(&self, ident: &str) -> bool{
        self.identifiers().contains(&ident.to_string())
    }
}

impl SimpleEquation {

    // Order is undefined
    pub fn identifiers(&self) -> Vec<String> {
        union_vecs(&self.lhs.identifiers(), &self.rhs.identifiers())
    }

    pub fn contains(&self, ident: &str) -> bool{
        let s = &ident.to_string();
        self.lhs.identifiers().contains(s) || self.rhs.identifiers().contains(s)
    }

    pub fn rebalance_for(&self, ident: String) -> Result<SimpleEquation,String> {
        let lvars = self.lhs.identifiers();
        let rvars = self.rhs.identifiers();

        let ret = match (lvars.contains(&ident), rvars.contains(&ident)) {
            (true, true) => Err("SymbolicError: NaiveImplementation".to_string()),
            (false, false) => Err("SymbolicError: VariableNotFound".to_string()),
            (true, false) => self.simplify_lhs(&ident),
            (false, true) =>
                SimpleEquation{lhs: self.rhs.clone(),
                               rhs: self.lhs.clone()}.simplify_lhs(&ident),
        };
        match ret {
            Ok(eqn) => {
                if eqn.rhs.contains(&ident) {
                    Err("SymbolicError: NaiveImplementation".to_string())
                } else {
                    Ok(eqn)
                }},
            Err(_) => ret,
        }
    }

    pub fn simplify_lhs(&self, ident: &str) -> Result<SimpleEquation,String> {
        use self::Expr::*;
        use self::BinOperator::*;
        match self.lhs {
            Ident(ref s) if s == ident => Ok((*self).clone()),
            Ident(_) | Integer(_) | Float(_) =>
                Err("SymbolicError: InternalError: expected var on LHS".to_string()),
            Der(_) | Abs(_) =>
                Err("SymbolicError: NaiveImplementation: can't simplify der() or abs()".to_string()),
            // TODO: create a macro for the below...
            BinExpr(Multiply, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Divide, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Multiply, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Divide, Box::new(self.rhs.clone()), a.clone())}.simplify_lhs(&ident) },
            BinExpr(Divide, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Multiply, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Divide, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Divide, a.clone(), Box::new(self.rhs.clone()))}.simplify_lhs(&ident) },
            BinExpr(Add, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Subtract, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Add, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Subtract, Box::new(self.rhs.clone()), a.clone())}.simplify_lhs(&ident) },
            BinExpr(Subtract, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Add, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Subtract, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Subtract, a.clone(), Box::new(self.rhs.clone()))}.simplify_lhs(&ident) },
            BinExpr(_, _, _) => Err("SymbolicError: NotImplemented BinOperator (or else couldn't find var...)".to_string()),
            // in case we add opers: _ => Err("NotImplemented".to_string()),
        }
    }
}


//// Debug Implementations

impl Debug for ModelicaModel {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        try!(write!(fmt, "model {}\n", self.name));
        for e in self.extends.iter() {
            try!(write!(fmt, "  extends {};\n", e));
        }
        for v in self.components.iter() {
            try!(write!(fmt, "  {:?};\n", v));
        }
        try!(write!(fmt, "equation\n"));
        for c in self.connections.iter() {
            try!(write!(fmt, "  {:?};\n", c));
        }
        for e in self.equations.iter() {
            try!(write!(fmt, "  {:?};\n", e));
        }
        write!(fmt, "end {};\n", self.name)
    }
}

impl Debug for ComponentPrefix {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        use self::ComponentPrefix::*;
        write!(fmt, "{}", 
            match *self {
                Flow => "flow",
                Stream => "stream",
                Input => "input",
                Output => "output",
                Discrete => "discrete",
                Parameter => "parameter",
                Constant => "constant",
            })
    }
}

impl Debug for Component {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "{}{} {}",
            match self.prefix {
                Some(p) => format!("{:?} ", p),
                None => "".to_string(),
            },
            self.specifier,
            self.name,
        )
    }
}

impl Debug for Connection {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "connect({}, {})", self.a, self.b)
    }
}

impl Debug for SimpleEquation {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "{:?} = {:?}", self.lhs, self.rhs)
    }
}

impl Debug for Expr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        use self::Expr::*;
        match *self {
            Integer(e) => write!(fmt, "{}", e),
            Float(e) => write!(fmt, "{}", e),
            Ident(ref e) => write!(fmt, "{}", e),
            Der(ref e) => write!(fmt, "der({:?})", e),
            Abs(ref e) => write!(fmt, "abs({:?})", e),
            BinExpr(op, ref l, ref r) => write!(fmt, "({:?} {:?} {:?})", l, op, r),
        }
    }
}

impl Debug for BinOperator {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        use self::BinOperator::*;
        match *self {
            Multiply => write!(fmt, "*"),
            Divide => write!(fmt, "/"),
            Add => write!(fmt, "+"),
            Subtract => write!(fmt, "-"),
        }
    }
}