From 14154fd272d0f3bde96fe77bfe2aeb6825345f54 Mon Sep 17 00:00:00 2001
From: Tad Hardesty <tad@platymuus.com>
Date: Tue, 7 Nov 2017 16:51:50 -0800
Subject: [PATCH] Add more traversal options to the object tree

---
 src/dreammaker/lib.rs     |  2 +-
 src/dreammaker/objtree.rs | 68 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/src/dreammaker/lib.rs b/src/dreammaker/lib.rs
index 0450ad4e..bc1a6840 100644
--- a/src/dreammaker/lib.rs
+++ b/src/dreammaker/lib.rs
@@ -29,7 +29,7 @@ pub mod constants;
 ///
 /// Errors are automatically pretty-printed to stdout before they are returned.
 pub fn parse_environment(dme: &Path) -> Result<objtree::ObjectTree, DMError> {
-    let mut preprocessor = preprocessor::Preprocessor::new(dme.to_owned()).unwrap();
+    let mut preprocessor = preprocessor::Preprocessor::new(dme.to_owned())?;
     parser::parse(indents::IndentProcessor::new(&mut preprocessor)).map_err(|e| {
         pretty_print_error(&mut io::stdout(), &preprocessor, &e).unwrap();
         e
diff --git a/src/dreammaker/objtree.rs b/src/dreammaker/objtree.rs
index cb87e5b2..d8ff49b9 100644
--- a/src/dreammaker/objtree.rs
+++ b/src/dreammaker/objtree.rs
@@ -6,12 +6,16 @@ use std::cell::Cell;
 pub use petgraph::graph::NodeIndex;
 use petgraph::graph::Graph;
 use petgraph::visit::EdgeRef;
+use petgraph::Direction;
 use linked_hash_map::LinkedHashMap;
 
 use super::ast::{Expression, TypePath, PathOp, Prefab};
 use super::constants::Constant;
 use super::{DMError, Location};
 
+// ----------------------------------------------------------------------------
+// Variables
+
 pub type Vars = LinkedHashMap<String, Constant>;
 
 #[derive(Debug, Clone)]
@@ -46,9 +50,12 @@ pub struct TypeVar {
     pub declaration: Option<VarDeclaration>,
 }
 
+// ----------------------------------------------------------------------------
+// Types
+
 #[derive(Debug, Default)]
 pub struct Type {
-    name: String,
+    pub name: String,
     pub path: String,
     pub vars: LinkedHashMap<String, TypeVar>,
     parent_type: Cell<NodeIndex>,
@@ -132,6 +139,61 @@ pub fn subpath(path: &str, parent: &str) -> bool {
     path == &parent[..parent.len() - 1] || path.starts_with(parent)
 }
 
+// ----------------------------------------------------------------------------
+// Type references
+
+#[derive(Debug, Copy, Clone)]
+pub struct TypeRef<'a> {
+    idx: NodeIndex,
+    ty: &'a Type,
+}
+
+impl<'a> TypeRef<'a> {
+    fn new(tree: &'a ObjectTree, idx: NodeIndex) -> TypeRef<'a> {
+        TypeRef {
+            idx,
+            ty: tree.graph.node_weight(idx).unwrap(),
+        }
+    }
+
+    pub fn parent(&self, objtree: &'a ObjectTree) -> Option<TypeRef<'a>> {
+        objtree.graph.neighbors_directed(self.idx, Direction::Incoming).next().map(|i| TypeRef::new(objtree, i))
+    }
+
+    pub fn parent_type(&self, objtree: &'a ObjectTree) -> Option<TypeRef<'a>> {
+        let idx = self.ty.parent_type.get();
+        objtree.graph.node_weight(idx).map(|ty| TypeRef { idx, ty })
+    }
+
+    pub fn child(&self, name: &str, objtree: &'a ObjectTree) -> Option<TypeRef<'a>> {
+        for idx in objtree.graph.neighbors(self.idx) {
+            let ty = objtree.graph.node_weight(idx).unwrap();
+            if ty.name == name {
+                return Some(TypeRef { idx, ty });
+            }
+        }
+        None
+    }
+
+    pub fn children(&self, objtree: &'a ObjectTree) -> Vec<TypeRef<'a>> {
+        let mut output = Vec::new();
+        for idx in objtree.graph.neighbors(self.idx) {
+            output.push(TypeRef::new(objtree, idx));
+        }
+        output
+    }
+}
+
+impl<'a> ::std::ops::Deref for TypeRef<'a> {
+    type Target = Type;
+    fn deref(&self) -> &Type {
+        self.ty
+    }
+}
+
+// ----------------------------------------------------------------------------
+// The object tree itself
+
 #[derive(Debug)]
 pub struct ObjectTree {
     pub graph: Graph<Type, ()>,
@@ -164,6 +226,10 @@ impl ObjectTree {
     // ------------------------------------------------------------------------
     // Access
 
+    pub fn root(&self) -> TypeRef {
+        TypeRef::new(self, NodeIndex::new(0))
+    }
+
     pub fn find(&self, path: &str) -> Option<&Type> {
         self.types.get(path).and_then(|&ix| self.graph.node_weight(ix))
     }