Hide implementation details of extending a graph in FFI module

Author: Bart Schuurmans

tensorflow/src/TensorFlow/Internal/FFI.hs

@@ -27,20 +27,7 @@ module TensorFlow.Internal.FFI
     , Raw.SessionOptions
 
     , Raw.Graph
-    , Raw.graphOperationByName
-    , resolveOutput
-    , resolveOperation
-    , graphImportGraphDef
-    , forGraphOperations_
-
-    , Raw.Operation
-    , Raw.operationName
-    , Raw.operationNumOutputs
-
-    , Raw.ImportGraphDefOptions
-    , Raw.importGraphDefOptionsAddInputMapping
-
-    , Raw.Output(..)
+    , extendGraph
 
     , TensorData(..)
     , setSessionConfig
@@ -60,6 +47,7 @@ import Control.Monad.Catch (MonadMask, Exception, throwM, bracket, finally, mask
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.Bits (Bits, toIntegralSized)
 import Data.Int (Int64)
+import Data.Foldable (for_)
 import Data.Maybe (fromMaybe)
 import Data.Typeable (Typeable)
 import Data.Word (Word8)
@@ -174,6 +162,21 @@ forGraphOperations_ graph f = with 0 go
           Raw.Operation ptr | ptr == nullPtr -> return ()
           _ -> f op >> go indexPtr -- indexPtr is modified by Raw.graphNextOperation.
 
+extendGraph :: Raw.Graph -> GraphDef -> IO ()
+extendGraph graph graphDef =
+    graphImportGraphDef graph graphDef $ \opts ->
+        -- All inputs of the nodes in the GraphDef should either refer to
+        -- other nodes in the GraphDef, or be mapped to nodes already in
+        -- the Graph by adding an input mapping.
+        -- We add an input mapping for all existing nodes in the Graph in
+        -- case they are referenced in the GraphDef.
+        forGraphOperations_ graph $ \op -> do
+            srcName <- Raw.operationName op
+            numOutputs <- Raw.operationNumOutputs op
+            for_ [0..numOutputs] $ \srcIndex -> do
+                let dst = Raw.Output op (safeConvert srcIndex)
+                with dst $ Raw.importGraphDefOptionsAddInputMapping opts srcName srcIndex
+
 run :: Raw.Session
     -> Raw.Graph
     -> [(String, TensorData)] -- ^ Inputs.

tensorflow/src/TensorFlow/Internal/Raw.chs

@@ -76,7 +76,7 @@ instance Storable Operation where
 -- Output.
 data Output = Output
     { outputOperation :: Operation
-    , outputIndex     :: Int
+    , outputIndex     :: CInt
     }
 {# pointer *TF_Output as OutputPtr -> Output #}
 

tensorflow/src/TensorFlow/Session.hs

@@ -45,10 +45,8 @@ import Control.Monad.Trans.Class (MonadTrans, lift)
 import Control.Monad.Trans.Reader (ReaderT(..), ask, asks)
 import Data.ByteString (ByteString)
 import Data.Default (Default, def)
-import Data.Foldable (for_)
 import Data.ProtoLens (showMessage)
 import Data.Set (Set)
-import Foreign.Marshal.Utils (with)
 import Lens.Family2 (Lens', (^.), (&), (.~))
 import Lens.Family2.Unchecked (lens)
 import Proto.Tensorflow.Core.Framework.Graph (GraphDef)
@@ -156,18 +154,7 @@ extend = do
     unless (null nodesToExtend) $ liftIO $ do
         let graphDef = (defMessage :: GraphDef) & node .~ nodesToExtend
         trace ("Session.extend " <> Builder.string8 (showMessage graphDef))
-        FFI.graphImportGraphDef graph graphDef $ \opts ->
-            -- All inputs of the nodes in the GraphDef should either refer to
-            -- other nodes in the GraphDef, or be mapped to nodes already in
-            -- the Graph by adding an input mapping.
-            -- We add an input mapping for all existing nodes in the Graph in
-            -- case they are referenced in the GraphDef.
-            FFI.forGraphOperations_ graph $ \op -> do
-                srcName <- FFI.operationName op
-                numOutputs <- FFI.operationNumOutputs op
-                for_ [0..numOutputs] $ \srcIndex -> do
-                    let dst = FFI.Output op srcIndex
-                    with dst $ FFI.importGraphDefOptionsAddInputMapping opts srcName srcIndex
+        FFI.extendGraph graph graphDef
     -- Now that all the nodes are created, run the initializers.
     initializers <- build flushInitializers
     unless (null initializers) $