1 files changed, 43 insertions, 4 deletions

diff --git a/modelica-parser-lalrpop/src/ast.rs b/modelica-parser-lalrpop/src/ast.rs
index 99d443c..467c518 100644
--- a/modelica-parser-lalrpop/src/ast.rs
+++ b/modelica-parser-lalrpop/src/ast.rs
@@ -6,7 +6,8 @@ use std::collections::HashMap;
 #[derive(Clone, PartialEq)]
 pub struct ModelicaModel {
     pub name: String,
-    pub components: Vec<Component>,
+    pub description: Option<String>,
+    pub component_clauses: Vec<ComponentClause>,
     pub equations: Vec<SimpleEquation>,
     pub connections: Vec<Connection>,
     pub extends: Vec<String>,
@@ -25,6 +26,22 @@ pub enum ComponentPrefix {
 }
 
 #[derive(Clone, PartialEq)]
+pub struct ComponentClause {
+    pub prefix: Option<ComponentPrefix>,
+    pub specifier: String,
+    pub declarations: Vec<ComponentDeclaration>,
+}
+
+#[derive(Clone, PartialEq)]
+pub struct ComponentDeclaration {
+    pub name: String,
+    pub value: Option<Expr>,
+    pub units: Option<String>,
+    pub description: Option<String>,
+}
+
+// This isn't part of the parse AST; it's a flattened helper type
+#[derive(Clone, PartialEq)]
 pub struct Component {
     pub prefix: Option<ComponentPrefix>,
     pub specifier: String,
@@ -68,10 +85,30 @@ pub enum BinOperator {
 
 impl ModelicaModel {
 
+    // This flattens the ComponentClause/ComponentDeclaration tree into a flat
+    // list of Components
+    // TODO: refactor this into an iterator `components()`?
+    pub fn get_components(&self) -> Vec<Component> {
+
+        let mut vars: Vec<Component> = vec![];
+        for clause in &self.component_clauses {
+            vars.extend(clause.declarations.iter().map(|ref dec|
+                Component {
+                    prefix: clause.prefix.clone(),
+                    specifier: clause.specifier.clone(),
+                    name: dec.name.clone(),
+                    value: dec.value.clone(),
+                    units: dec.units.clone(),
+                    description: dec.description.clone(),
+            }))
+        }
+        vars
+    }
+
     pub fn get_constant_vars(&self) -> HashMap<String,Option<Expr>> {
         let mut binds = HashMap::new();
         // XXX: actually implement this...
-        for c in &self.components {
+        for c in &self.get_components() {
             match c.prefix {
                 Some(ComponentPrefix::Constant) => { binds.insert(c.name.clone(), Some(Expr::Integer(123))); },
                 Some(ComponentPrefix::Parameter) => { binds.insert(c.name.clone(), Some(Expr::Float(4.56))); },
@@ -87,11 +124,13 @@ impl ModelicaModel {
     //   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 {
+        let components = self.get_components();
+        let vars = 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() {
@@ -159,7 +198,7 @@ impl Debug for ModelicaModel {
         for e in self.extends.iter() {
             try!(write!(fmt, "  extends {};\n", e));
         }
-        for v in self.components.iter() {
+        for v in self.get_components().iter() {
             try!(write!(fmt, "  {:?};\n", v));
         }
         try!(write!(fmt, "equation\n"));