content: -- Copyright 2016 TensorFlow authors.
  --
  -- Licensed under the Apache License, Version 2.0 (the "License");
  -- you may not use this file except in compliance with the License.
  -- You may obtain a copy of the License at
  --
  --     http://www.apache.org/licenses/LICENSE-2.0
  --
  -- Unless required by applicable law or agreed to in writing, software
  -- distributed under the License is distributed on an "AS IS" BASIS,
  -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  -- See the License for the specific language governing permissions and
  -- limitations under the License.
  {-# LANGUAGE CPP #-}
  
  -- | Generates the wrappers for Ops shipped with tensorflow.
  module Main where
  
  import Distribution.PackageDescription
      ( PackageDescription(..)
      , libBuildInfo
      , hsSourceDirs
      )
  import qualified Distribution.Simple.BuildPaths as BuildPaths
  import Distribution.Simple.LocalBuildInfo (LocalBuildInfo)
  import Distribution.Simple
      ( defaultMainWithHooks
      , simpleUserHooks
      , UserHooks(..)
      )
  #if MIN_VERSION_Cabal(3,6,0)
  import Distribution.Utils.Path (unsafeMakeSymbolicPath)
  #endif
  import Data.List (intercalate)
  import Data.ProtoLens (decodeMessage)
  import System.Directory (createDirectoryIfMissing)
  import System.Exit (exitFailure)
  import System.FilePath ((</>))
  import System.IO (hPutStrLn, stderr)
  import TensorFlow.Internal.FFI (getAllOpList)
  import TensorFlow.OpGen (docOpList, OpGenFlags(..))
  import Text.PrettyPrint.Mainland (prettyLazyText)
  import qualified Data.Text.Lazy.IO as Text
  
  main = defaultMainWithHooks generatingOpsWrappers
  
  -- TODO: Generalize for user libraries by replacing getAllOpList with
  -- a wrapper around TF_LoadLibrary. The complicated part is interplay
  -- between bazel and Haskell build system.
  generatingOpsWrappers :: UserHooks
  generatingOpsWrappers = hooks
      { buildHook = \p l h f -> generateSources l >> buildHook hooks p l h f
      , haddockHook = \p l h f -> generateSources l >> haddockHook hooks p l h f
      , replHook = \p l h f args -> generateSources l
                                          >> replHook hooks p l h f args
      }
    where
      flagsBuilder dir = OpGenFlags
          { outputFile = dir </> "Core.hs"
          , prefix = "TensorFlow.GenOps"
          , excludeList = intercalate "," blackList
          }
      hooks = simpleUserHooks
      generateSources :: LocalBuildInfo -> IO ()
      generateSources l = do
          let dir = autogenModulesDir l </> "TensorFlow/GenOps"
          createDirectoryIfMissing True dir
          let flags = flagsBuilder dir
          pb <- getAllOpList
          case decodeMessage pb of
              Left e -> hPutStrLn stderr e >> exitFailure
              Right x -> Text.writeFile (outputFile flags)
                                        (prettyLazyText 80 $ docOpList flags x)
  
  -- | Add the autogen directory to the hs-source-dirs of all the targets in the
  -- .cabal file.  Used to fool 'sdist' by pointing it to the generated source
  -- files.
  fudgePackageDesc
      :: LocalBuildInfo -> PackageDescription -> PackageDescription
  fudgePackageDesc lbi p = p
      { library =
          (\lib -> lib { libBuildInfo = fudgeBuildInfo (libBuildInfo lib) })
              <$> library p
      }
    where
      fudgeBuildInfo bi =
  #if MIN_VERSION_Cabal(3,6,0)
          bi { hsSourceDirs = unsafeMakeSymbolicPath (autogenModulesDir lbi) : hsSourceDirs bi }
  #else    
          bi { hsSourceDirs = autogenModulesDir lbi : hsSourceDirs bi }
  #endif
  
  blackList =
      [ -- Requires the "func" type:
        "FilterDataset"
      , "BatchFunction"
      , "Case"
      , "ChooseFastestBranchDataset"
      , "ExperimentalGroupByReducerDataset"
      , "ExperimentalGroupByWindowDataset"
      , "ExperimentalMapAndBatchDataset"
      , "ExperimentalMapDataset"
      , "ExperimentalNumaMapAndBatchDataset"
      , "ExperimentalParallelInterleaveDataset"
      , "ExperimentalScanDataset"
      , "ExperimentalTakeWhileDataset"
      , "FilterDataset"
      , "FlatMapDataset"
      , "For"
      , "GeneratorDataset"
      , "GroupByReducerDataset"
      , "GroupByWindowDataset"
      , "If"
      , "InterleaveDataset"
      , "LegacyParallelInterleaveDatasetV2"
      , "LoadDataset"
      , "MapAndBatchDataset"
      , "MapAndBatchDatasetV2"
      , "MapDataset"
      , "MapDefun"
      , "OneShotIterator"
      , "ParallelInterleaveDataset"
      , "ParallelInterleaveDatasetV2"
      , "ParallelInterleaveDatasetV3"
      , "ParallelInterleaveDatasetV4"
      , "ParallelMapDataset"
      , "ParallelMapDatasetV2"
      , "ParseSequenceExample"
      , "ParseSequenceExampleV2"
      , "ParseSingleSequenceExample"
      , "PartitionedCall"
      , "ReduceDataset"
      , "RemoteCall"
      , "SaveDataset"
      , "ScanDataset"
      , "SnapshotDatasetV2"
      , "StatefulPartitionedCall"
      , "StatelessCase"
      , "StatelessIf"
      , "StatelessWhile"
      , "SymbolicGradient"
      , "TakeWhileDataset"
      , "TPUCompile"
      , "TPUPartitionedCall"
      , "TPUReplicate"
      , "While"
      , "XlaIf"
      , "XlaLaunch"
      , "XlaReduce"
      , "XlaReduceWindow"
      , "XlaSelectAndScatter"
      , "XlaScatter"
      , "XlaWhile"
      , "_If"
      , "_TPUReplicate"
      , "_While"
      , "_XlaCompile"
      -- Incorrectly generated:
      , "_FusedBatchNormGradEx"
      -- Could not deduce:
      , "_MklFusedBatchNorm"
      , "_MklFusedBatchNormEx"
      , "_MklFusedBatchNormGrad"
      , "_MklFusedBatchNormGradV2"
      , "_MklFusedBatchNormGradV3"
      , "_MklFusedBatchNormV2"
      , "_MklFusedBatchNormV3"
      ]
  
  autogenModulesDir :: LocalBuildInfo -> FilePath
  #if MIN_VERSION_Cabal(2,0,0)
  autogenModulesDir = BuildPaths.autogenPackageModulesDir
  #else
  autogenModulesDir = BuildPaths.autogenModulesDir
  #endif
  