Serialize Nockma output using nock jam (#3066)

The Anoma API accepts jammed nock terms as input. The benefit to this is
that jammed terms are greatly compressed compared to the original term.

* Closes https://github.com/anoma/juvix/issues/3017

Remaining tasks:

- [x] Deserialize input nockma file in `juvix dev nockma {run, eval}`
- [x] Support debug input nockma file in `juvix dev nockma {run, eval,
repl}` i.e there should be a way to pass the `*.debug.nockma` (output of
`juvix compile anoma --debug`) file to `juvix dev nockma {run, eval,
repl}`
- [x] Add proper JuvixErrors for deserialisation failures

---------

Co-authored-by: Jan Mas Rovira <janmasrovira@gmail.com>

app/Commands/Compile/Anoma.hs

@@ -3,6 +3,7 @@ module Commands.Compile.Anoma where
 import Commands.Base
 import Commands.Compile.Anoma.Options
 import Commands.Extra.NewCompile
+import Juvix.Compiler.Nockma.Encoding.Jam qualified as Encoding
 import Juvix.Compiler.Nockma.Pretty qualified as Nockma
 import Juvix.Compiler.Nockma.Translation.FromTree qualified as Nockma
 
@@ -25,9 +26,9 @@ runCommand opts = do
   res <- getRight r
   outputAnomaResult (opts' ^. compileDebug) nockmaFile res
 
-outputAnomaResult :: (Members '[EmbedIO, App] r) => Bool -> Path Abs File -> Nockma.AnomaResult -> Sem r ()
+outputAnomaResult :: (Members '[EmbedIO, App, Files] r) => Bool -> Path Abs File -> Nockma.AnomaResult -> Sem r ()
 outputAnomaResult debugOutput nockmaFile Nockma.AnomaResult {..} = do
-  let code = Nockma.ppSerialize _anomaClosure
-      prettyNockmaFile = replaceExtensions' [".debug", ".nockma"] nockmaFile
-  writeFileEnsureLn nockmaFile code
+  let code = Encoding.jamToByteString _anomaClosure
+      prettyNockmaFile = replaceExtensions' nockmaDebugFileExts nockmaFile
+  writeFileBS nockmaFile code
   when debugOutput (writeFileEnsureLn prettyNockmaFile (Nockma.ppPrint _anomaClosure))

app/Commands/Dev/Nockma/Eval.hs

@@ -10,10 +10,9 @@ import Juvix.Compiler.Nockma.Translation.FromSource qualified as Nockma
 runCommand :: forall r. (Members AppEffects r) => NockmaEvalOptions -> Sem r ()
 runCommand opts = do
   afile <- fromAppPathFile file
-  parsedTerm <- Nockma.parseTermFile afile
+  parsedTerm <- runAppError @JuvixError (Nockma.cueJammedFileOrPretty afile)
   case parsedTerm of
-    Left err -> exitJuvixError (JuvixError err)
-    Right (TermCell c) -> do
+    TermCell c -> do
       (counts, res) <-
         runOpCounts
           . runReader defaultEvalOptions
@@ -22,7 +21,7 @@ runCommand opts = do
       putStrLn (ppPrint res)
       let statsFile = replaceExtension' ".profile" afile
       writeFileEnsureLn statsFile (prettyText counts)
-    Right TermAtom {} -> exitFailMsg "Expected nockma input to be a cell"
+    TermAtom {} -> exitFailMsg "Expected nockma input to be a cell"
   where
     file :: AppPath File
     file = opts ^. nockmaEvalFile

app/Commands/Dev/Nockma/Format.hs

@@ -5,13 +5,11 @@ import Commands.Dev.Nockma.Format.Options
 import Juvix.Compiler.Nockma.Pretty
 import Juvix.Compiler.Nockma.Translation.FromSource qualified as Nockma
 
-runCommand :: forall r. (Members '[EmbedIO, App] r) => NockmaFormatOptions -> Sem r ()
+runCommand :: forall r. (Members AppEffects r) => NockmaFormatOptions -> Sem r ()
 runCommand opts = do
   afile <- fromAppPathFile file
-  parsedTerm <- Nockma.parseTermFile afile
-  case parsedTerm of
-    Left err -> exitJuvixError (JuvixError err)
-    Right t -> putStrLn (ppPrint t)
+  parsedTerm <- runAppError @JuvixError (Nockma.parseTermFile afile)
+  putStrLn (ppPrint parsedTerm)
   where
     file :: AppPath File
     file = opts ^. nockmaFormatFile

app/Commands/Dev/Nockma/Repl.hs

@@ -10,7 +10,7 @@ import Juvix.Compiler.Nockma.Evaluator.Options
 import Juvix.Compiler.Nockma.Language
 import Juvix.Compiler.Nockma.Pretty
 import Juvix.Compiler.Nockma.Pretty qualified as Nockma
-import Juvix.Compiler.Nockma.Translation.FromSource (parseProgramFile, parseReplStatement, parseReplText, parseText)
+import Juvix.Compiler.Nockma.Translation.FromSource (cueJammedFileOrPrettyProgram, parseReplStatement, parseReplText, parseText)
 import Juvix.Compiler.Nockma.Translation.FromSource qualified as Nockma
 import Juvix.Parser.Error
 import Juvix.Prelude qualified as Prelude
@@ -111,7 +111,8 @@ getProgram :: Repl (Maybe (Program Natural))
 getProgram = State.gets (^. replStateProgram)
 
 readProgram :: Prelude.Path Abs File -> Repl (Program Natural)
-readProgram s = fromMegaParsecError <$> parseProgramFile s
+readProgram s = runM . runFilesIO $ do
+  runErrorIO' @JuvixError (cueJammedFileOrPrettyProgram s)
 
 direction' :: String -> Repl ()
 direction' s = Repline.dontCrash $ do
@@ -181,7 +182,7 @@ replAction =
         banner
       }
 
-runCommand :: forall r. (Members '[EmbedIO, App] r) => NockmaReplOptions -> Sem r ()
+runCommand :: forall r. (Members '[Files, EmbedIO, App] r) => NockmaReplOptions -> Sem r ()
 runCommand opts = do
   mt :: Maybe (Term Natural) <- mapM iniStack (opts ^. nockmaReplOptionsStackFile)
   liftIO . (`State.evalStateT` (iniState mt)) $ replAction
@@ -189,10 +190,7 @@ runCommand opts = do
     iniStack :: AppPath File -> Sem r (Term Natural)
     iniStack af = do
       afile <- fromAppPathFile af
-      parsedTerm <- Nockma.parseTermFile afile
-      case parsedTerm of
-        Left err -> exitJuvixError (JuvixError err)
-        Right t -> return t
+      checkCued (Nockma.cueJammedFile afile)
 
     iniState :: Maybe (Term Natural) -> ReplState
     iniState mt =
@@ -202,3 +200,6 @@ runCommand opts = do
           _replStateLoadedFile = Nothing,
           _replStateLastResult = nockNilTagged "repl-result"
         }
+
+    checkCued :: Sem (Error JuvixError ': r) a -> Sem r a
+    checkCued = runErrorNoCallStackWith exitJuvixError

app/Commands/Dev/Nockma/Run.hs

@@ -7,14 +7,13 @@ import Juvix.Compiler.Nockma.EvalCompiled
 import Juvix.Compiler.Nockma.Evaluator
 import Juvix.Compiler.Nockma.Pretty
 import Juvix.Compiler.Nockma.Translation.FromSource qualified as Nockma
-import Juvix.Parser.Error
 
 runCommand :: forall r. (Members AppEffects r) => NockmaRunOptions -> Sem r ()
 runCommand opts = do
   afile <- fromAppPathFile inputFile
   argsFile <- mapM fromAppPathFile (opts ^. nockmaRunArgs)
-  parsedArgs <- mapM (Nockma.parseTermFile >=> checkParsed) argsFile
-  parsedTerm <- Nockma.parseTermFile afile >>= checkParsed
+  parsedArgs <- runAppError @JuvixError (mapM Nockma.cueJammedFileOrPretty argsFile)
+  parsedTerm <- checkCued (Nockma.cueJammedFileOrPretty afile)
   case parsedTerm of
     t@(TermCell {}) -> do
       let formula = anomaCallTuple parsedArgs
@@ -31,7 +30,5 @@ runCommand opts = do
     inputFile :: AppPath File
     inputFile = opts ^. nockmaRunFile
 
-    checkParsed :: Either MegaparsecError (Term Natural) -> Sem r (Term Natural)
-    checkParsed = \case
-      Left err -> exitJuvixError (JuvixError err)
-      Right tm -> return tm
+    checkCued :: Sem (Error JuvixError ': r) a -> Sem r a
+    checkCued = runErrorNoCallStackWith exitJuvixError

src/Juvix/Compiler/Nockma/Encoding/Cue.hs

@@ -10,7 +10,7 @@ import Juvix.Prelude.Base
 import VectorBuilder.Builder as Builder
 import VectorBuilder.Vector
 
-data CueState a = CueState
+newtype CueState a = CueState
   { _cueStateCache :: HashMap Int (Term a)
   }
 
@@ -20,7 +20,7 @@ initCueState =
     { _cueStateCache = mempty
     }
 
-data CueEnv = CueEnv
+newtype CueEnv = CueEnv
   {_cueEnvStartPos :: Int}
 
 initCueEnv :: CueEnv
@@ -38,14 +38,20 @@ data DecodingError
   | DecodingErrorInvalidBackref
   deriving stock (Show)
 
+instance Pretty DecodingError where
+  pretty = unAnnotate . ppCodeAnn
+
+instance PrettyCodeAnn DecodingError where
+  ppCodeAnn = \case
+    DecodingErrorInvalidTag -> "Invalid tag"
+    DecodingErrorCacheMiss -> "Cache miss"
+    DecodingErrorInvalidLength -> "Invalid length"
+    DecodingErrorExpectedAtom -> "Expected atom"
+    DecodingErrorInvalidAtom -> "Invalid atom"
+    DecodingErrorInvalidBackref -> "Invalid backref"
+
 instance PrettyCode DecodingError where
-  ppCode = \case
-    DecodingErrorInvalidTag -> return "Invalid tag"
-    DecodingErrorCacheMiss -> return "Cache miss"
-    DecodingErrorInvalidLength -> return "Invalid length"
-    DecodingErrorExpectedAtom -> return "Expected atom"
-    DecodingErrorInvalidAtom -> return "Invalid atom"
-    DecodingErrorInvalidBackref -> return "Invalid backref"
+  ppCode = return . ppCodeAnn
 
 -- | Register the start of processing a new entity
 registerElementStart ::
@@ -154,7 +160,7 @@ atomToBits a' = do
   n <- nockNatural' a'
   return (integerToVectorBits @Integer (fromIntegral n))
 
--- | Transfor a vector of bits to a decoded term
+-- | Transform a vector of bits to a decoded term
 cueFromBits ::
   forall a r.
   ( NockNatural a,
@@ -168,6 +174,19 @@ cueFromBits ::
   Sem r (Term a)
 cueFromBits v = evalBitReader v (evalState (initCueState @a) (runReader initCueEnv cueFromBitsSem))
 
+cueFromByteString' ::
+  forall a r.
+  ( NockNatural a,
+    Members
+      '[ Error DecodingError,
+         Error (ErrNockNatural' a)
+       ]
+      r
+  ) =>
+  ByteString ->
+  Sem r (Term a)
+cueFromByteString' = cueFromBits . cloneFromByteString
+
 cueFromBitsSem ::
   forall a r.
   ( NockNatural a,
@@ -275,6 +294,28 @@ cueEither =
     . runErrorNoCallStack @DecodingError
     . cue'
 
+cueFromByteString ::
+  -- NB: The signature returns the DecodingError in an Either to avoid
+  -- overlapping instances with `ErrNockNatural a` when errors are handled. See
+  -- the comment above `ErrNockNatural' a` for more explanation.
+  forall a r.
+  ( NockNatural a,
+    Member (Error (ErrNockNatural a)) r
+  ) =>
+  ByteString ->
+  Sem r (Either DecodingError (Term a))
+cueFromByteString =
+  runErrorNoCallStackWith @(ErrNockNatural' a) (\(ErrNockNatural' e) -> throw e)
+    . runErrorNoCallStack @DecodingError
+    . cueFromByteString'
+
+cueFromByteString'' ::
+  forall a.
+  (NockNatural a) =>
+  ByteString ->
+  Either (ErrNockNatural a) (Either DecodingError (Term a))
+cueFromByteString'' = run . runErrorNoCallStack . cueFromByteString
+
 {- `ErrNockNatural a` must be wrapped in a newtype to avoid overlapping instances
 with `DecodingError` when errors are handled before the type variable `a` is
 resolved.

src/Juvix/Compiler/Nockma/Encoding/Jam.hs

@@ -95,6 +95,10 @@ jamToBits =
     . evalState (initJamState @a)
     . jamSem
 
+-- | jam encode a Nock term to the bytes encoding of an atom
+jamToByteString :: forall a. (Integral a, Hashable a) => Term a -> ByteString
+jamToByteString = vectorBitsToByteString . jamToBits
+
 -- | jam encode a Nock term to an atom
 jam :: forall a r. (Integral a, Hashable a, NockNatural a, Member (Error (ErrNockNatural a)) r) => Term a -> Sem r (Atom a)
 jam t = do

src/Juvix/Compiler/Nockma/Evaluator/Error.hs

@@ -26,13 +26,6 @@ data NockEvalError a
   | ErrDecodingFailed (DecodingFailed a)
   | ErrVerificationFailed (VerificationFailed a)
 
-newtype GenericNockEvalError = GenericNockEvalError
-  { _genericNockEvalErrorMessage :: AnsiText
-  }
-
-class ToGenericNockEvalError a where
-  toGenericNockEvalError :: a -> GenericNockEvalError
-
 data ExpectedCell a = ExpectedCell
   { _expectedCellCtx :: EvalCtx,
     _expectedCellAtom :: Atom a

src/Juvix/Compiler/Nockma/Language.hs

@@ -11,8 +11,8 @@ import GHC.Base (Type)
 import Juvix.Compiler.Core.Language.Base (Symbol)
 import Juvix.Compiler.Nockma.AnomaLib.Base
 import Juvix.Compiler.Nockma.Language.Path
+import Juvix.Data.CodeAnn
 import Juvix.Prelude hiding (Atom, Path)
-import Juvix.Prelude.Pretty
 
 data ReplStatement a
   = ReplStatementExpression (ReplExpression a)
@@ -245,6 +245,9 @@ makeLenses ''WithStack
 makeLenses ''AtomInfo
 makeLenses ''CellInfo
 
+singletonProgram :: Term a -> Program a
+singletonProgram t = Program [StatementStandalone t]
+
 isCell :: Term a -> Bool
 isCell = \case
   TermCell {} -> True

src/Juvix/Compiler/Nockma/Pretty/Base.hs

@@ -1,3 +1,5 @@
+{-# OPTIONS_GHC -Wno-orphans #-}
+
 module Juvix.Compiler.Nockma.Pretty.Base
   ( module Juvix.Compiler.Nockma.Pretty.Base,
     module Juvix.Data.CodeAnn,
@@ -12,6 +14,9 @@ import Juvix.Data.CodeAnn
 import Juvix.Extra.Strings qualified as Str
 import Juvix.Prelude hiding (Atom, Path)
 
+docDefault :: (PrettyCode c) => c -> Doc Ann
+docDefault = doc defaultOptions
+
 doc :: (PrettyCode c) => Options -> c -> Doc Ann
 doc opts =
   run
@@ -24,6 +29,11 @@ class PrettyCode c where
 runPrettyCode :: (PrettyCode c) => Options -> c -> Doc Ann
 runPrettyCode opts = run . runReader opts . ppCode
 
+instance PrettyCodeAnn NockNaturalNaturalError where
+  ppCodeAnn = \case
+    NaturalInvalidPath a -> "Invalid path" <+> docDefault a
+    NaturalInvalidOp a -> "Invalid operator code" <+> docDefault a
+
 instance forall a. (PrettyCode a, NockNatural a) => PrettyCode (Atom a) where
   ppCode atm = do
     t <- runFail $ do

src/Juvix/Compiler/Nockma/Translation/FromSource/Base.hs

@@ -4,7 +4,9 @@ import Data.HashMap.Internal.Strict qualified as HashMap
 import Data.List.NonEmpty qualified as NonEmpty
 import Data.Text qualified as Text
 import Juvix.Compiler.Nockma.Encoding.ByteString (textToNatural)
+import Juvix.Compiler.Nockma.Encoding.Cue qualified as Cue
 import Juvix.Compiler.Nockma.Language
+import Juvix.Data.CodeAnn
 import Juvix.Extra.Paths
 import Juvix.Extra.Strings qualified as Str
 import Juvix.Parser.Error
@@ -23,15 +25,67 @@ parseText = runParser noFile
 parseReplText :: Text -> Either MegaparsecError (ReplTerm Natural)
 parseReplText = runParserFor replTerm noFile
 
-parseTermFile :: (MonadIO m) => Prelude.Path Abs File -> m (Either MegaparsecError (Term Natural))
-parseTermFile fp = do
-  txt <- readFile fp
-  return (runParser fp txt)
+-- | If the file ends in .debug.nockma it parses an annotated unjammed term. Otherwise
+-- it is equivalent to cueJammedFile
+cueJammedFileOrPretty ::
+  forall r.
+  (Members '[Files, Error JuvixError] r) =>
+  Prelude.Path Abs File ->
+  Sem r (Term Natural)
+cueJammedFileOrPretty f
+  | f `hasExtensions` nockmaDebugFileExts = parseTermFile f
+  | otherwise = cueJammedFile f
 
-parseProgramFile :: (MonadIO m) => Prelude.Path Abs File -> m (Either MegaparsecError (Program Natural))
+-- | If the file ends in .debug.nockma it parses an annotated unjammed program. Otherwise
+-- it parses program with a single jammed term
+cueJammedFileOrPrettyProgram ::
+  forall r.
+  (Members '[Files, Error JuvixError] r) =>
+  Prelude.Path Abs File ->
+  Sem r (Program Natural)
+cueJammedFileOrPrettyProgram f
+  | f `hasExtensions` nockmaDebugFileExts = parseProgramFile f
+  | otherwise = singletonProgram <$> cueJammedFile f
+
+cueJammedFile :: forall r. (Members '[Files, Error JuvixError] r) => Prelude.Path Abs File -> Sem r (Term Natural)
+cueJammedFile fp = do
+  bs <- readFileBS' fp
+  case Cue.cueFromByteString'' @Natural bs of
+    Left e -> natErr e
+    Right (Left e) -> decodingErr e
+    Right (Right t) -> return t
+  where
+    err :: AnsiText -> Sem r x
+    err msg =
+      throw $
+        JuvixError
+          GenericError
+            { _genericErrorLoc = i,
+              _genericErrorIntervals = [i],
+              _genericErrorMessage = msg
+            }
+
+    decodingErr :: Cue.DecodingError -> Sem r x
+    decodingErr e = err (mkAnsiText (ppCodeAnn e))
+
+    natErr :: NockNaturalNaturalError -> Sem r x
+    natErr e = err (mkAnsiText (ppCodeAnn e))
+
+    i :: Interval
+    i = mkInterval loc loc
+      where
+        loc :: Loc
+        loc = mkInitialLoc fp
+
+parseTermFile :: (Members '[Files, Error JuvixError] r) => Prelude.Path Abs File -> Sem r (Term Natural)
+parseTermFile fp = do
+  txt <- readFile' fp
+  either (throw . JuvixError) return (runParser fp txt)
+
+parseProgramFile :: (Members '[Files, Error JuvixError] r) => Prelude.Path Abs File -> Sem r (Program Natural)
 parseProgramFile fp = do
-  txt <- readFile fp
-  return (runParserProgram fp txt)
+  txt <- readFile' fp
+  either (throw . JuvixError) return (runParserProgram fp txt)
 
 parseReplStatement :: Text -> Either MegaparsecError (ReplStatement Natural)
 parseReplStatement = runParserFor replStatement noFile

src/Juvix/Data/FileExt.hs

@@ -32,9 +32,29 @@ data FileExt
 
 $(genSingletons [''FileExt])
 
+splitExtensions :: Path b File -> (Path b File, [String])
+splitExtensions =
+  swap
+    . run
+    . runAccumListReverse
+    . go
+  where
+    go :: (Members '[Accum String] r) => Path b File -> Sem r (Path b File)
+    go f = case splitExtension f of
+      Nothing -> return f
+      Just (f', ext) -> do
+        accum ext
+        go f'
+
+hasExtensions :: (Foldable l) => Path b File -> l String -> Bool
+hasExtensions f exts = toList exts == snd (splitExtensions f)
+
 juvixFileExt :: (IsString a) => a
 juvixFileExt = ".juvix"
 
+nockmaDebugFileExts :: (IsString a) => NonEmpty a
+nockmaDebugFileExts = ".debug" :| [".nockma"]
+
 juvixMarkdownFileExt :: (IsString a) => a
 juvixMarkdownFileExt = ".juvix.md"
 

src/Juvix/Prelude/Effects/Accum.hs

@@ -22,6 +22,7 @@ newtype instance StaticRep (Accum o) = Accum
 accum :: (Member (Accum o) r) => o -> Sem r ()
 accum o = overStaticRep (\(Accum l) -> Accum (o : l))
 
+-- | Accumulates in LIFO order
 runAccumListReverse :: Sem (Accum o ': r) a -> Sem r ([o], a)
 runAccumListReverse m = do
   (a, Accum s) <- runStaticRep (Accum mempty) m

src/Juvix/Prelude/Path.hs

@@ -88,18 +88,16 @@ removeExtension = fmap fst . splitExtension
 removeExtension' :: Path b File -> Path b File
 removeExtension' = fst . fromJust . splitExtension
 
-addExtensions :: (MonadThrow m) => [String] -> Path b File -> m (Path b File)
-addExtensions ext p = case ext of
-  [] -> return p
-  e : es -> addExtension e p >>= addExtensions es
+addExtensions :: forall m l b. (MonadThrow m, Foldable l) => l String -> Path b File -> m (Path b File)
+addExtensions exts p = foldM (flip addExtension) p exts
 
-replaceExtensions :: (MonadThrow m) => [String] -> Path b File -> m (Path b File)
+replaceExtensions :: (MonadThrow m, Foldable l) => l String -> Path b File -> m (Path b File)
 replaceExtensions ext = addExtensions ext . removeExtensions
 
-replaceExtensions' :: [String] -> Path b File -> Path b File
+replaceExtensions' :: (Foldable l) => l String -> Path b File -> Path b File
 replaceExtensions' ext = fromJust . replaceExtensions ext
 
-addExtensions' :: [String] -> Path b File -> Path b File
+addExtensions' :: (Foldable l) => l String -> Path b File -> Path b File
 addExtensions' ext = fromJust . addExtensions ext
 
 -- | TODO this is ugly. Please, fix it. FileExtJuvixMarkdown needs special

tests/Anoma/Compilation/positive/test001-args.debug.nockma

@@ -0,0 +1 @@
+[1 100]

tests/smoke/Commands/compile.smoke.yaml

@@ -144,6 +144,38 @@ tests:
     stdout: ""
     exit-status: 0
 
+  - name: target-anoma-run-debug
+    command:
+      shell:
+        - bash
+      script: |
+        temp=$(mktemp -d)
+        trap 'rm -rf -- "$temp"' EXIT
+        testdir=$PWD/tests/Anoma/Compilation/positive
+        cd $temp
+        juvix --log-level error compile anoma $testdir/test001.juvix --debug
+        [ -f test001.debug.nockma ]
+        juvix dev nockma run test001.debug.nockma --args $testdir/test001-args.debug.nockma
+    stdout: |
+      106
+    exit-status: 0
+
+  - name: target-anoma-run
+    command:
+      shell:
+        - bash
+      script: |
+        temp=$(mktemp -d)
+        trap 'rm -rf -- "$temp"' EXIT
+        testdir=$PWD/tests/Anoma/Compilation/positive
+        cd $temp
+        juvix --log-level error compile anoma $testdir/test001.juvix
+        [ -f test001.nockma ]
+        juvix dev nockma run test001.nockma --args $testdir/test001-args.debug.nockma
+    stdout: |
+      106
+    exit-status: 0
+
   - name: target-anoma
     command:
       shell: