Some REPL niceties

src/Main.hs

@@ -17,11 +17,13 @@ import FileEval
   )
 import IODriver (IOPermissions(..), runIO)
 import Parser (parseTricu)
-import REPL (repl)
+import REPL (repl, replWithStore)
 import Research (T, EvaluatedForm(..), Env, formatT, exportDag)
 import Wire (encodeBundle, defaultExportNames, Bundle(..))
 
-import Control.Monad (foldM, unless, when)
+import Control.Monad (foldM, forM, unless, when)
+import Data.Char (isAlphaNum)
+import Data.List (sortOn)
 import qualified Data.Text as T
 import Data.Version (showVersion)
 import Paths_tricu (version)
@@ -31,13 +33,18 @@ import qualified Data.ByteString as BS
 import qualified Data.ByteString.Lazy as BL
 import qualified Data.Sequence as Seq
 import qualified Data.Map as Map
-import System.Directory (getHomeDirectory)
+import System.Directory (createDirectoryIfMissing, getHomeDirectory)
 import System.FilePath (takeBaseName, (</>))
 
 -- ---------------------------------------------------------------------------
 -- CLI argument types
 -- ---------------------------------------------------------------------------
 
+data AppArgs = AppArgs
+  { globalStore :: Maybe FilePath
+  , appCommand   :: TricuArgs
+  } deriving (Show)
+
 data TricuArgs
   = Repl
   | Check
@@ -74,6 +81,8 @@ data TricuArgs
       , exportOutput  :: FilePath
       , exportNames   :: [String]
       , exportStore   :: Maybe FilePath
+      , exportAll     :: Bool
+      , exportSplit   :: Bool
       , dag           :: Bool
       }
   | StoreAliasList
@@ -251,6 +260,14 @@ exportParser = ArboricxExport
      <> metavar "PATH"
      <> help "Content-addressed store path"
       ))
+  <*> switch
+      ( long "all"
+     <> help "Export all name aliases that point at tree-term objects"
+      )
+  <*> switch
+      ( long "split"
+     <> help "Write one single-export bundle per export; --output is treated as a directory"
+      )
   <*> switch
       ( long "dag"
      <> help "Export as a topologically-sorted DAG node table instead of a bundle"
@@ -297,9 +314,15 @@ storeAliasGetParser = StoreAliasGet
 versionStr :: String
 versionStr = "tricu " ++ showVersion version
 
-tricuParser :: Parser TricuArgs
-tricuParser = (subparser topCommands <|> pure Repl)
-  <**> infoOption versionStr (long "version" <> help "Show version")
+tricuParser :: Parser AppArgs
+tricuParser = AppArgs
+  <$> optional (option str
+      ( long "store"
+     <> metavar "PATH"
+     <> help "Global content-addressed store path used by commands and the REPL unless a subcommand overrides it"
+      ))
+  <*> ((subparser topCommands <|> pure Repl)
+      <**> infoOption versionStr (long "version" <> help "Show version"))
   where
     topCommands = mconcat
       [ command "check" (info (checkParser <**> helper)
@@ -342,13 +365,15 @@ storeAliasParser = subparser $ mconcat
 
 main :: IO ()
 main = do
-  args <- execParser $ info (tricuParser <**> helper)
+  appArgs <- execParser $ info (tricuParser <**> helper)
     ( fullDesc
    <> progDesc "Exploring Tree Calculus"
    <> header versionStr
     )
+  let mGlobalStore = globalStore appArgs
+      args = applyGlobalStore mGlobalStore (appCommand appArgs)
   case args of
-    Repl               -> runRepl
+    Repl               -> runReplWithStore mGlobalStore
     Check {}           -> runCheck args
     Eval {}            -> runEval args
     ArboricxCompile {} -> runCompile args
@@ -362,11 +387,31 @@ main = do
 -- Command runners
 -- ---------------------------------------------------------------------------
 
+applyGlobalStore :: Maybe FilePath -> TricuArgs -> TricuArgs
+applyGlobalStore mGlobal args = case args of
+  Repl -> Repl
+  Check {} -> args { checkStore = preferLocal (checkStore args) }
+  Eval {} -> args { evalStore = preferLocal (evalStore args) }
+  ArboricxCompile {} -> args { compileStore = preferLocal (compileStore args) }
+  ArboricxImport {} -> args { importStore = preferLocal (importStore args) }
+  ArboricxExport {} -> args { exportStore = preferLocal (exportStore args) }
+  StoreAliasList {} -> args { storePathOpt = preferLocal (storePathOpt args) }
+  StoreAliasGet {} -> args { storePathOpt = preferLocal (storePathOpt args) }
+  where
+    preferLocal local = case local of
+      Just _  -> local
+      Nothing -> mGlobal
+
 runRepl :: IO ()
-runRepl = do
+runRepl = runReplWithStore Nothing
+
+runReplWithStore :: Maybe FilePath -> IO ()
+runReplWithStore mStore = do
   putStrLn "Welcome to the tricu REPL"
   putStrLn "You may exit with `CTRL+D` or the `!exit` command."
-  repl
+  case mStore of
+    Nothing -> repl
+    Just store -> replWithStore (StorePath store)
 
 runCheck :: TricuArgs -> IO ()
 runCheck opts = do
@@ -466,23 +511,46 @@ runExportBundle opts = do
       modules = exportModules opts
       out     = exportOutput opts
       names   = exportNames opts
+      allFlag = exportAll opts
+      splitFlag = exportSplit opts
   when (null out) $ die "tricu arboricx export: --output is required"
-  when (null targets && null modules) $
-    die "tricu arboricx export: at least one --target or --module is required"
+  when (null targets && null modules && not allFlag) $
+    die "tricu arboricx export: at least one --target, --module, or --all is required"
+  when (splitFlag && not (null names)) $
+    die "tricu arboricx export --split: --name is not supported; split bundles use their export names"
   store <- resolveStorePath (exportStore opts)
+  allEntries <- if allFlag then resolveAllNameExports store else pure []
   targetRoots <- mapM (resolveStoreTarget store) targets
   moduleRoots <- concat <$> mapM (resolveModuleExports store) modules
   let targetEntries = zip (defaultExportNames (length targetRoots)) targetRoots
-      entries = targetEntries ++ moduleRoots
+      entries = allEntries ++ targetEntries ++ moduleRoots
       expNames = if null names then map fst entries else map T.pack names
+  when (null entries) $
+    die "tricu arboricx export: no tree-term exports found"
   when (length expNames /= length entries) $
     die "tricu arboricx export: number of --name values must match number of exported roots"
-  bundle <- packBundleFromStore store (zip expNames (map snd entries))
-  let bundleData = encodeBundle bundle
-  BL.writeFile out (BL.fromStrict bundleData)
-  putStrLn $ "Exported bundle with " ++ show (length entries) ++ " export(s) to " ++ out
-  putStrLn $ "  nodes: " ++ show (Seq.length (bundleNodes bundle))
-  putStrLn $ "  size: " ++ show (BS.length bundleData) ++ " bytes"
+  if splitFlag
+    then runExportBundleSplit store out (zip expNames (map snd entries))
+    else do
+      bundle <- packBundleFromStore store (zip expNames (map snd entries))
+      let bundleData = encodeBundle bundle
+      BL.writeFile out (BL.fromStrict bundleData)
+      putStrLn $ "Exported bundle with " ++ show (length entries) ++ " export(s) to " ++ out
+      putStrLn $ "  nodes: " ++ show (Seq.length (bundleNodes bundle))
+      putStrLn $ "  size: " ++ show (BS.length bundleData) ++ " bytes"
+
+runExportBundleSplit :: StorePath -> FilePath -> [(T.Text, ObjectHash)] -> IO ()
+runExportBundleSplit store outDir entries = do
+  createDirectoryIfMissing True outDir
+  written <- forM (zip [0 :: Int ..] entries) $ \(i, (name, root)) -> do
+    bundle <- packBundleFromStore store [(name, root)]
+    let bundleData = encodeBundle bundle
+        path = outDir </> splitBundleFileName i name
+    BL.writeFile path (BL.fromStrict bundleData)
+    pure (path, Seq.length (bundleNodes bundle), BS.length bundleData)
+  putStrLn $ "Exported " ++ show (length written) ++ " split bundle(s) to " ++ outDir
+  mapM_ (\(path, nodeCount, byteCount) ->
+    putStrLn $ "  " ++ path ++ " (nodes: " ++ show nodeCount ++ ", size: " ++ show byteCount ++ " bytes)") written
 
 runStoreAliasList :: TricuArgs -> IO ()
 runStoreAliasList opts = do
@@ -543,6 +611,19 @@ resolveStoreTarget store target = do
     Just _ -> return root
     Nothing -> die $ "Term not found in store: " ++ target
 
+resolveAllNameExports :: StorePath -> IO [(T.Text, ObjectHash)]
+resolveAllNameExports store = do
+  aliases <- sortOn fst <$> listAliases store NameAlias
+  fmap concat $ mapM exportAlias aliases
+  where
+    exportAlias (name, ref)
+      | objectRefKind ref /= unDomain treeTermDomain = pure []
+      | otherwise = do
+          mTree <- getTreeTerm store (objectRefHash ref)
+          case mTree of
+            Nothing -> die $ "Name alias tree term not found: " ++ T.unpack name
+            Just _  -> pure [(name, objectRefHash ref)]
+
 resolveModuleExports :: StorePath -> String -> IO [(T.Text, ObjectHash)]
 resolveModuleExports store moduleTarget = do
   manifestHash <- resolveModuleManifestHash store moduleTarget
@@ -574,6 +655,17 @@ resolveModuleManifestHash store moduleTarget = do
 formatObjectRef :: ObjectRef -> String
 formatObjectRef ref = T.unpack (objectRefKind ref) ++ " " ++ T.unpack (objectRefHash ref)
 
+splitBundleFileName :: Int -> T.Text -> FilePath
+splitBundleFileName i name = show i ++ "-" ++ sanitize (T.unpack name) ++ ".arboricx"
+  where
+    sanitize [] = "export"
+    sanitize xs = case map safeChar xs of
+      [] -> "export"
+      ys -> ys
+    safeChar c
+      | isAlphaNum c || c == '-' || c == '_' || c == '.' = c
+      | otherwise = '_'
+
 writeOutput :: FilePath -> String -> IO ()
 writeOutput path content
   | null path = putStr content