agda2train-0.0.1.0: src/Main.hs
{-# LANGUAGE TypeApplications #-}
-- | Main driver for the @agda2train@ executable.
module Main where
import GHC.Generics
import System.Environment ( getArgs, withArgs )
import System.Directory ( createDirectoryIfMissing, doesFileExist )
import System.FilePath ( (</>) )
import System.Console.GetOpt ( OptDescr(..), ArgDescr(..) )
import qualified Data.List as L
import qualified Data.Set as S
import qualified Data.HashMap.Strict as HM
import qualified Data.ByteString.Lazy as BL
import Data.Aeson ( ToJSON )
import Data.Aeson.Encode.Pretty
( encodePretty', Config(..), defConfig, Indent(..), keyOrder )
import Control.Monad
import Control.Monad.IO.Class ( liftIO )
import Control.DeepSeq ( NFData )
import Agda.Main ( runAgda )
import Agda.Compiler.Backend hiding (Reduced)
import Agda.Compiler.Common ( curDefs )
import Agda.Utils.Pretty
import Agda.Utils.Maybe
import Agda.Utils.Monad
import Agda.Syntax.Scope.Base
import Agda.Syntax.Scope.Monad
import Agda.Syntax.TopLevelModuleName
import Agda.TypeChecking.Reduce
import qualified Agda.TypeChecking.Pretty as P
import Agda.Compiler.Common ( curIF )
import ToTrain
import Output
-- | The main entrypoint for the @agda2train@ executable.
main :: IO ()
main = do
as <- getArgs
let extraFlags = [] -- ["--no-projection-like"]
withArgs (extraFlags ++ as) $
runAgda [Backend $ mkBackend train]
-- | Make an Agda backend from a given training function (c.f. 'ToTrain.TrainF').
mkBackend :: TrainF ->
Backend' Options Options ([ScopeEntry], [ScopeEntry]) () (String, [Sample])
mkBackend trainF = Backend'
{ backendName = "agda2train"
, backendVersion = Nothing
, options = defaultOptions
, commandLineFlags =
[ Option ['r'] ["recurse"] (NoArg recOpt)
"Recurse into imports/dependencies."
, Option ['x'] ["no-json"] (NoArg noJsonOpt)
"Skip generation of JSON files. (just debug print)"
, Option [] ["ignore-existing-json"] (NoArg ignoreJsonOpt)
"Ignore existing JSON files. (i.e. always overwrite)"
, Option [] ["print-json"] (NoArg printJsonOpt)
"Print JSON output. (for debugging)"
, Option [] ["no-terms"] (NoArg includeDefsOpt)
"Do not include definitions of things in scope"
, Option [] ["no-privates"] (NoArg includePrivsOpt)
"Do not include private definitions"
, Option ['o'] ["out-dir"] (ReqArg outDirOpt "DIR")
"Generate data at DIR. (default: project root)"
]
, isEnabled = \ _ -> True
, scopeCheckingSuffices = False
, mayEraseType = \ _ -> return True
, preCompile = return
, postCompile = \ _ _ _ -> return ()
, preModule = \ opts isMain md _ ->
let
processScopeEntry :: QName -> TCM (Maybe ScopeEntry)
processScopeEntry qn =
caseMaybeM (tryMaybe $ do
def <- getConstInfo qn
case theDef def of
GeneralizableVar -> fail "not handling `variable` definitions"
DataOrRecSig _ -> fail "not handling `DataOrRecSig` definitions"
_ -> typeOfConst qn) (pure Nothing) $ \ty -> do
rty <- mkReduced ty
pty <- ppm ty
rty' <- traverse convert rty
def <- getConstInfo qn
pdef <- ppm def
def' <- convert def
report 20 $ ppm (pp qn) <> " : " <> ppm ty
report 30 $ " *pretty: " <> P.text (pp ty)
report 20 $ ppm (pp qn) <> " = " <> ppm def
report 30 $ " *pretty: " <> P.text (pp def)
reportReduced rty
report 20 ""
return $ Just $ ppName qn :~ ScopeEntry
{ _type = prender pty :> rty'
, definition = boolToMaybe (includeDefs opts)
$ prender pdef :> def'
, holes = Nothing
}
in
ifM (skip opts isMain md) (return $ Skip ()) $ do
report 10 $ "************************ "
<> ppm md <> " (" <> ppm (show isMain)
<> ") ***********************************"
-- printScope "" 1 ""
setScope . iInsideScope =<< curIF
-- ** public
NameSpace _ _ ns <- allThingsInScope <$> getCurrentScope
scopeEntries <- mapMaybeM processScopeEntry (S.toList ns)
-- ** private
privs <- HM.filterWithKey (\k _ -> k `S.notMember` ns) <$> liftTCM curDefs
privScopeEntries <- mapMaybeM processScopeEntry (HM.keys privs)
report 10 "******************************************************************"
return $ Recompile (scopeEntries, privScopeEntries)
, compileDef = \ _ _ _ def -> (ppName (defName def) ,) <$> runC (forEachHole trainF def)
, postModule = \ opts (scopeEntries, privScopeEntries) isMain md samples ->
let mn = pp md in
unlessM (skip opts isMain md) $ do
whenJust (outDir opts) (liftIO . createDirectoryIfMissing True)
unless (noJson opts) $ liftIO $ do
let
isGlobal = \(n :~ _) -> isNothing $ lookup n samples
(globals, locals') = L.partition isGlobal scopeEntries
locals = flip map locals' $ \(n :~ l) ->
n :~ l { holes = Just $ fromJust $ lookup n samples }
plocals = flip map privScopeEntries $ \(n :~ l) ->
n :~ l { holes = Just $ fromJust $ lookup n samples }
fileData = mn :~ TrainData
{ scopeGlobal = globals
, scopeLocal = locals
, scopePrivate = boolToMaybe (includePrivs opts) plocals
}
encodeFile (getOutFn opts mn) fileData
when (printJson opts) $ do
putStrLn "************************ JSON ************************"
BL.putStr (encode fileData)
putStrLn "******************************************************"
}
where
skip :: Options -> IsMain -> TopLevelModuleName -> TCM Bool
skip opts@Options{..} isMain md = do
jsonExists <- liftIO $ doesFileExist (getOutFn opts $ pp md)
return $ (not recurse && (isMain == NotMain))
|| (not noJson && jsonExists && not ignoreJson)
-- * Command-line flags
-- | Available options to configure how @agda2train@ generates training data:
--
-- ['recurse'] whether to recursively generate data from all transitive dependencies
--
-- ['noJson'] mock run without generating any actual JSON files
--
-- ['ignoreJson'] compile everything from scratch
--
-- ['includeDefs'] whether to include definitions in scope, or just their type
--
-- ['includePrivs'] whether to include private definitions as well
--
-- (run @agda2train -h/--help@ for a human-readable description of all options)
data Options = Options
{ recurse, noJson, ignoreJson, printJson, includeDefs, includePrivs :: Bool
, outDir :: Maybe FilePath
} deriving (Generic, NFData)
-- | The default options.
defaultOptions = Options
{ recurse = False
, noJson = False
, ignoreJson = False
, printJson = False
, includeDefs = True
, includePrivs = True
, outDir = Nothing
}
recOpt, noJsonOpt, ignoreJsonOpt, printJsonOpt, includeDefsOpt, includePrivsOpt
:: Monad m => Options -> m Options
recOpt opts = return $ opts { recurse = True }
noJsonOpt opts = return $ opts { noJson = True }
printJsonOpt opts = return $ opts { printJson = True }
ignoreJsonOpt opts = return $ opts { ignoreJson = True }
includeDefsOpt opts = return $ opts { includeDefs = False }
includePrivsOpt opts = return $ opts { includePrivs = False }
outDirOpt :: Monad m => FilePath -> Options -> m Options
outDirOpt fp opts = return $ opts { outDir = Just fp }
getOutDir :: Options -> FilePath
getOutDir opts = case outDir opts of
Just fp -> fp
Nothing -> "."
getOutFn :: Options -> String -> FilePath
getOutFn opts mn = getOutDir opts </> mn <> ".json"
-- * JSON encoding
-- | Uses "Aeson.Pretty" to order the JSON fields.
encode :: ToJSON a => a -> BL.ByteString
encode = encodePretty' $ defConfig
{ confIndent = Spaces 2
, confCompare = keyOrder
[ "pretty"
, "tag"
, "name"
, "original", "simplified", "reduced", "normalised"
, "telescope", "patterns", "fields"
, "domain", "codomain"
, "abstraction", "body"
, "sort", "level", "literal"
, "head", "arguments"
, "variants", "reference", "variant"
, "index"
, "scopeGlobal", "scopeLocal"
, "type", "definition", "holes"
, "ctx", "goal", "term", "premises"
]
}
encodeFile :: ToJSON a => FilePath -> a -> IO ()
encodeFile = \fn -> BL.writeFile fn . encode