ProtoCompile.hs

-- | This is the Main module for the command line program 'hprotoc'
module Main where

import Control.Monad(when,forM_)
import qualified Data.ByteString.Lazy.Char8 as LC (writeFile)
import Data.List(break)
import qualified Data.Sequence as Seq (fromList,singleton)
import Data.Version(showVersion)
import Language.Haskell.Exts.Pretty(prettyPrintStyleMode,Style(..),Mode(..),PPHsMode(..),PPLayout(..))
import System.Console.GetOpt(OptDescr(Option),ArgDescr(NoArg,ReqArg)
                            ,usageInfo,getOpt,ArgOrder(ReturnInOrder))
import System.Directory(getCurrentDirectory,createDirectoryIfMissing)
import System.Environment(getArgs)
import System.FilePath(takeDirectory,combine,joinPath)
import qualified System.FilePath.Posix as Canon(takeBaseName)

import Text.ProtocolBuffers.Basic(defaultValue)
import Text.ProtocolBuffers.Identifiers(MName,checkDIString,mangle)
import Text.ProtocolBuffers.Reflections(ProtoInfo(..),EnumInfo(..))
import Text.ProtocolBuffers.WireMessage (messagePut)

import qualified Text.DescriptorProtos.FileDescriptorProto as D(FileDescriptorProto)
import qualified Text.DescriptorProtos.FileDescriptorProto as D.FileDescriptorProto(FileDescriptorProto(..))
import qualified Text.DescriptorProtos.FileDescriptorSet   as D(FileDescriptorSet)
import qualified Text.DescriptorProtos.FileDescriptorSet   as D.FileDescriptorSet(FileDescriptorSet(..))

import Text.ProtocolBuffers.ProtoCompile.BreakRecursion(makeResult)
import Text.ProtocolBuffers.ProtoCompile.Gen(protoModule,descriptorModules,enumModule)
import Text.ProtocolBuffers.ProtoCompile.MakeReflections(makeProtoInfo,serializeFDP)
import Text.ProtocolBuffers.ProtoCompile.Resolve(loadProto,makeNameMaps,getTLS
                                                ,LocalFP(..),CanonFP(..),TopLevel(..))

-- The Paths_hprotoc module is produced by cabal
import Paths_hprotoc(version)

data Options = Options { optPrefix :: [MName String]
                       , optAs :: [(CanonFP,[MName String])]
                       , optTarget :: LocalFP
                       , optInclude :: [LocalFP]
                       , optProto :: LocalFP
                       , optDesc :: Maybe (LocalFP)
                       , optImports,optVerbose,optUnknownFields,optDryRun :: Bool }
  deriving Show

setPrefix,setTarget,setInclude,setProto,setDesc :: String -> Options -> Options
setImports,setVerbose,setUnknown,setDryRun :: Options -> Options
setPrefix   s o = o { optPrefix = toPrefix s }
setTarget   s o = o { optTarget = (LocalFP s) }
setInclude  s o = o { optInclude = LocalFP s : optInclude o }
setProto    s o = o { optProto = LocalFP s }
setDesc     s o = o { optDesc = Just (LocalFP s) }
setImports    o = o { optImports = True }
setVerbose    o = o { optVerbose = True }
setUnknown    o = o { optUnknownFields = True }
setDryRun     o = o { optDryRun = True }

toPrefix :: String -> [MName String]
toPrefix s = case checkDIString s of
               Left msg -> error $ "Bad module name in options:"++show s++"\n"++msg
               Right (True,_) -> error $ "Bad module name in options (cannot start with '.'): "++show s
               Right (False,ms) -> map mangle ms

-- | 'setAs' puts both the full path and the basename as keys into the association list
setAs :: String -> Options -> Options
setAs s o =
   case break ('='==) s of
     (filepath,'=':rawPrefix) -> let value = toPrefix rawPrefix
                                 in o { optAs = (CanonFP filepath,value):
                                                (CanonFP (Canon.takeBaseName filepath),value):
                                                optAs o}
     _ -> error . unlines $ [ "Malformed -a or --as option "++show s
                            , "  Expected \"FILEPATH=MODULE\""
                            , "  where FILEPATH is the basename or relative path (using '/') of an imported file"
                            , "  where MODULE is a dotted Haskell name to use as a prefix"
                            , "  also MODULE can be empty to have no prefix" ]

data OptionAction = Mutate (Options->Options) | Run (Options->Options) | Switch Flag

data Flag = VersionInfo

optionList :: [OptDescr OptionAction]
optionList =
  [ Option ['a'] ["as"] (ReqArg (Mutate . setAs) "FILEPATH=MODULE")
               "assign prefix module to imported prot file: --as descriptor.proto=Text"
  , Option ['I'] ["proto_path"] (ReqArg (Mutate . setInclude) "DIR")
               "directory from which to search for imported proto files (default is pwd); all DIR searched"
  , Option ['d'] ["haskell_out"] (ReqArg (Mutate . setTarget) "DIR")
               "directory to use are root of generated files (default is pwd); last flag"
  , Option ['n'] ["dry_run"] (NoArg (Mutate setDryRun))
               "produce no output but still parse and check the proto file(s)"
  , Option ['o'] ["descriptor_set_out"] (ReqArg (Mutate . setDesc) "FILE")
               "filename to write binary FileDescriptorSet to"
  , Option [] ["include_imports"] (NoArg (Mutate setImports))
               "when writing descriptor_set_out include all imported files to be self-contained"
  , Option ['p'] ["prefix"] (ReqArg (Mutate . setPrefix) "MODULE")
               "dotted Haskell MODULE name to use as a prefix (default is none); last flag used"
  , Option ['u'] ["unknown_fields"] (NoArg (Mutate setUnknown))
               "generated messages and groups all support unknown fields"
  , Option ['v'] ["verbose"] (NoArg (Mutate  setVerbose))
               "increase amount of printed information"
  , Option [] ["version"]  (NoArg (Switch VersionInfo))
               "print out version information"
  ]

usageMsg,versionInfo :: String
usageMsg = usageInfo "Usage: protoCompile [OPTION..] path-to-file.proto ..." optionList

versionInfo = unlines $
  [ "Welcome to protocol-buffers version "++showVersion version
  , "Copyright (c) 2008, Christopher Kuklewicz."
  , "Released under BSD3 style license, see LICENSE file for details."
  , "Some proto files, such as descriptor.proto and unittest*.proto"
  , "are from google's code and are under an Apache 2.0 license."
  , ""
  , "Most command line arguments are similar to those of 'protoc'."
  , ""
  , "This program reads a .proto file and generates haskell code files."
  , "See http://code.google.com/apis/protocolbuffers/docs/overview.html for more."
  ]

processOptions :: [String] -> Either String [OptionAction]
processOptions argv =
    case getOpt (ReturnInOrder (Run . setProto)) optionList argv of
    (opts,_,[]) -> Right opts
    (_,_,errs) -> Left (unlines errs ++ usageMsg)

defaultOptions :: IO Options
defaultOptions = do
  pwd <- fmap LocalFP getCurrentDirectory
  return $ Options { optPrefix = []
                   , optAs = []
                   , optTarget = pwd
                   , optInclude = [pwd]
                   , optProto = LocalFP ""
                   , optDesc = Nothing
                   , optImports = False
                   , optVerbose = False
                   , optUnknownFields = False
                   , optDryRun = False }

main :: IO ()
main = do
  defs <- defaultOptions
  args <- getArgs
  case processOptions args of
    Left msg -> putStrLn msg
    Right todo -> process defs todo

process :: Options -> [OptionAction] -> IO ()
process options [] = if null (unLocalFP (optProto options))
                       then do putStrLn "No proto file specified (or empty proto file)"
                               putStrLn ""
                               putStrLn usageMsg
                       else putStrLn "Processing complete, have a nice day."
process options (Mutate f:rest) = process (f options) rest
process options (Run f:rest) = let options' = f options
                            in run options' >> process options' rest
process _options (Switch VersionInfo:_) = putStrLn versionInfo
  
mkdirFor :: FilePath -> IO ()
mkdirFor p = createDirectoryIfMissing True (takeDirectory p)

style :: Style
style = Style PageMode 132 0.6

myMode :: PPHsMode
myMode = PPHsMode 2 2 2 2 4 1 True PPOffsideRule False True

dump :: Bool -> Maybe LocalFP -> D.FileDescriptorProto -> [D.FileDescriptorProto] -> IO ()
dump _ Nothing _ _ = return ()
dump imports (Just (LocalFP dumpFile)) fdp fdps = do
  putStrLn $ "dumping to filename: "++show dumpFile
  let s = if imports then Seq.fromList fdps else Seq.singleton fdp
  LC.writeFile dumpFile (messagePut $ defaultValue { D.FileDescriptorSet.file = s })
  putStrLn $ "finished dumping FileDescriptorSet binary of: "++show (D.FileDescriptorProto.name fdp)

run :: Options -> IO ()
run options = do
  (env,fdps) <- loadProto (optInclude options) (optProto options)
  putStrLn "All proto files loaded"
  let fdp = either error id . top'FDP . fst . getTLS $ env
  when (not (optDryRun options)) $ dump (optImports options) (optDesc options) fdp fdps
  nameMap <- either error return $ makeNameMaps (optPrefix options) (optAs options) env
  putStrLn "Haskell name mangling done"
  let protoInfo = makeProtoInfo (optUnknownFields options) nameMap fdp
      result = makeResult protoInfo
  seq result (putStrLn "Recursive modules resolved")
  let produceMSG di = do
        when (not (optDryRun options)) $ do
          -- There might be several modules
          let fileModules = descriptorModules result di
          forM_ fileModules $ \ (relPath,modSyn) -> do
            let file = combine (unLocalFP . optTarget $ options) relPath
            print file
            mkdirFor file
            writeFile file (prettyPrintStyleMode style myMode modSyn)
      produceENM ei = do
        let file = combine (unLocalFP . optTarget $ options) . joinPath . enumFilePath $ ei
        print file
        when (not (optDryRun options)) $ do
          mkdirFor file
          writeFile file (prettyPrintStyleMode style myMode (enumModule ei))
  mapM_ produceMSG (messages protoInfo)
  mapM_ produceENM (enums protoInfo)
  let file = combine (unLocalFP . optTarget $ options) . joinPath . protoFilePath $ protoInfo
  print file
  writeFile file (prettyPrintStyleMode style myMode (protoModule result protoInfo (serializeFDP fdp)))