Merge pull request tensorflow#160 from adamcrume/master

Add support for while loops

Author: adamcrume

src/graph.rs

@@ -35,6 +35,7 @@ struct GraphLifetime;
 #[derive(Debug)]
 struct GraphImpl {
     inner: *mut tf::TF_Graph,
+    owned: bool,
 }
 
 unsafe impl Send for GraphImpl {}
@@ -43,8 +44,10 @@ unsafe impl Sync for GraphImpl {}
 impl Drop for GraphImpl {
     /// Graph will be deleted once no more Sessions are referencing it.
     fn drop(&mut self) {
-        unsafe {
-            tf::TF_DeleteGraph(self.inner);
+        if self.owned {
+            unsafe {
+                tf::TF_DeleteGraph(self.inner);
+            }
         }
     }
 }
@@ -273,7 +276,10 @@ impl Graph {
     pub fn new() -> Graph {
         unsafe {
             Graph {
-                gimpl: Arc::new(GraphImpl { inner: tf::TF_NewGraph() }),
+                gimpl: Arc::new(GraphImpl {
+                    inner: tf::TF_NewGraph(),
+                    owned: true,
+                }),
                 lifetime: GraphLifetime,
             }
         }
@@ -333,6 +339,32 @@ impl Graph {
         }
     }
 
+    /// Finds a unique operation name.  The pattern must contain exactly one
+    /// '{}' placeholder to indicate where a unique ID can be inserted, e.g.
+    /// 'Add_{}' or 'while_loop_{}/Merge', and the function returns an integer
+    /// which, when inserted into the placeholder, yields an operation name
+    /// which does not appear in the graph.
+    pub(crate) fn generate_operation_name(&self, operation_name_pattern: &str) -> Result<i64> {
+        let parts: Vec<_> = operation_name_pattern.split("{}").collect();
+        if parts.len() != 2 {
+            return Err(invalid_arg!(
+                "operation_name_pattern must contain placeholder"
+            ));
+        }
+        // Can't use format! because its argument must be a string literal.
+        let mut i = 0;
+        loop {
+            let name = format!("{}{}{}", parts[0], i, parts[1]);
+            let c_name = CString::new(name)?;
+            unsafe {
+                if tf::TF_GraphOperationByName(self.gimpl.inner, c_name.as_ptr()).is_null() {
+                    return Ok(i);
+                }
+            }
+            i += 1;
+        }
+    }
+
     /// Iterates over the operations in the graph.
     pub fn operation_iter(&self) -> OperationIter {
         OperationIter {
@@ -717,6 +749,16 @@ impl GraphTrait for Graph {
     fn inner(&self) -> *mut tf::TF_Graph {
         self.gimpl.inner
     }
+
+    unsafe fn from_c(inner: *mut tf::TF_Graph) -> Self {
+        Graph {
+            gimpl: Arc::new(GraphImpl {
+                inner,
+                owned: false,
+            }),
+            lifetime: GraphLifetime,
+        }
+    }
 }
 
 ////////////////////////
@@ -1523,14 +1565,14 @@ pub struct Output {
 }
 
 impl Output {
-    fn to_c(&self) -> tf::TF_Output {
+    pub(crate) fn to_c(&self) -> tf::TF_Output {
         tf::TF_Output {
             oper: self.operation.inner,
             index: self.index,
         }
     }
 
-    fn from_c(graph: &Graph, output: &tf::TF_Output) -> Self {
+    pub(crate) fn from_c(graph: &Graph, output: &tf::TF_Output) -> Self {
         Output {
             operation: Operation {
                 inner: output.oper,
@@ -2479,4 +2521,17 @@ mod tests {
         // We don't want to compare the actual proto because it may change across releases.
         assert!(g.versions().unwrap().len() > 0);
     }
+
+    #[test]
+    fn graph_generate_operation_name() {
+        let mut g = Graph::new();
+        for i in 0..5 {
+            assert_eq!(i, g.generate_operation_name("foo_{}").unwrap());
+            let mut nd = g.new_operation("Placeholder", &format!("foo_{}", i))
+                .unwrap();
+            nd.set_attr_type("dtype", DataType::Float).unwrap();
+            nd.set_attr_shape("shape", &Shape(Some(vec![]))).unwrap();
+            nd.finish().unwrap();
+        }
+    }
 }

src/lib.rs

@@ -47,11 +47,11 @@ use std::str::Utf8Error;
 /// Will panic if `msg` contains an embedded 0 byte.
 macro_rules! invalid_arg {
   ($fmt:expr) => {
-    Status::new_set(Code::InvalidArgument, $fmt).unwrap()
+    ::Status::new_set(::Code::InvalidArgument, $fmt).unwrap()
   };
   ($fmt:expr, $($arg:tt)*) => ({
     let msg = format!($fmt, $($arg)*);
-    Status::new_set(Code::InvalidArgument, &msg).unwrap()
+    ::Status::new_set(::Code::InvalidArgument, &msg).unwrap()
   });
 }
 
@@ -1240,6 +1240,8 @@ impl Library {
 
 ////////////////////////
 
+// TODO: Replace these with pub(crate)
+
 /// This exposes Buffer behavior without making it public.
 trait BufferTrait {
     fn is_owned(&self) -> bool;
@@ -1251,6 +1253,7 @@ trait BufferTrait {
 /// This exposes Graph behavior without making it public.
 trait GraphTrait {
     fn inner(&self) -> *mut tf::TF_Graph;
+    unsafe fn from_c(inner: *mut tf::TF_Graph) -> Self;
 }
 
 
@@ -1323,6 +1326,11 @@ impl Display for Shape {
 
 ////////////////////////
 
+mod while_loop;
+pub use while_loop::*;
+
+////////////////////////
+
 #[cfg(test)]
 mod tests {
     use super::*;