JuvixTree parser and pretty printer (#2583)

This PR implements:
* JuvixTree parser.
* JuvixTree pretty printer.
* `juvix dev tree read file.jvt` command which reads and pretty prints a
JuvixTree file.
* The `tree` target in the `compile` command.
* Removal of `StackRef` in JuvixAsm. This makes JuvixAsm more consistent
with JuvixTree and simplifies the data structures. `StackRef` is not
needed for compilation from Core.

Tests for the parser will appear in a separate PR, when I implement an
automatic translation of JuvixAsm to JuvixTree files.

---------

Co-authored-by: Paul Cadman <[email protected]>
Co-authored-by: Jan Mas Rovira <[email protected]>

app/Commands/Compile.hs

@@ -24,6 +24,7 @@ runCommand opts@CompileOptions {..} = do
     TargetGeb -> Compile.runGebPipeline arg
     TargetVampIR -> Compile.runVampIRPipeline arg
     TargetCore -> writeCoreFile arg
+    TargetTree -> Compile.runTreePipeline arg
     TargetAsm -> Compile.runAsmPipeline arg
     TargetNockma -> Compile.runNockmaPipeline arg
 

app/Commands/Dev.hs

@@ -19,6 +19,7 @@ import Commands.Dev.Parse qualified as Parse
 import Commands.Dev.Runtime qualified as Runtime
 import Commands.Dev.Scope qualified as Scope
 import Commands.Dev.Termination qualified as Termination
+import Commands.Dev.Tree qualified as Tree
 import Commands.Repl qualified as Repl
 
 runCommand :: (Members '[Embed IO, App, TaggedLock] r) => DevCommand -> Sem r ()
@@ -31,6 +32,7 @@ runCommand = \case
   Core opts -> Core.runCommand opts
   Geb opts -> Geb.runCommand opts
   Asm opts -> Asm.runCommand opts
+  Tree opts -> Tree.runCommand opts
   Casm opts -> Casm.runCommand opts
   Runtime opts -> Runtime.runCommand opts
   DisplayRoot opts -> DisplayRoot.runCommand opts

app/Commands/Dev/Asm/Compile.hs

@@ -58,10 +58,14 @@ runCommand opts = do
       TargetWasm32Wasi -> return Backend.TargetCWasm32Wasi
       TargetNative64 -> return Backend.TargetCNative64
       TargetNockma -> return Backend.TargetNockma
-      TargetGeb -> exitMsg (ExitFailure 1) "error: GEB target not supported for JuvixAsm"
-      TargetVampIR -> exitMsg (ExitFailure 1) "error: VampIR target not supported for JuvixAsm"
-      TargetCore -> exitMsg (ExitFailure 1) "error: JuvixCore target not supported for JuvixAsm"
-      TargetAsm -> exitMsg (ExitFailure 1) "error: JuvixAsm target not supported for JuvixAsm"
+      TargetTree -> err "JuvixTree"
+      TargetGeb -> err "GEB"
+      TargetVampIR -> err "VampIR"
+      TargetCore -> err "JuvixCore"
+      TargetAsm -> err "JuvixAsm"
+      where
+        err :: Text -> Sem r a
+        err tgt = exitMsg (ExitFailure 1) ("error: " <> tgt <> " target not supported for JuvixAsm")
 
 inputCFile :: (Members '[App] r) => Path Abs File -> Sem r (Path Abs File)
 inputCFile inputFileCompile = do

app/Commands/Dev/Core/Compile.hs

@@ -19,6 +19,7 @@ runCommand opts = do
     TargetVampIR -> runVampIRPipeline arg
     TargetCore -> return ()
     TargetAsm -> runAsmPipeline arg
+    TargetTree -> runTreePipeline arg
     TargetNockma -> runNockmaPipeline arg
   where
     getFile :: Sem r (Path Abs File)

app/Commands/Dev/Core/Compile/Base.hs

@@ -10,6 +10,7 @@ import Juvix.Compiler.Backend.Geb qualified as Geb
 import Juvix.Compiler.Backend.VampIR.Translation qualified as VampIR
 import Juvix.Compiler.Core.Data.Module qualified as Core
 import Juvix.Compiler.Nockma.Pretty qualified as Nockma
+import Juvix.Compiler.Tree.Pretty qualified as Tree
 import System.FilePath (takeBaseName)
 
 data PipelineArg = PipelineArg
@@ -38,6 +39,7 @@ getEntry PipelineArg {..} = do
       TargetVampIR -> Backend.TargetVampIR
       TargetCore -> Backend.TargetCore
       TargetAsm -> Backend.TargetAsm
+      TargetTree -> Backend.TargetTree
       TargetNockma -> Backend.TargetNockma
 
     defaultOptLevel :: Int
@@ -111,6 +113,19 @@ runAsmPipeline pa@PipelineArg {..} = do
   let code = Asm.ppPrint tab' tab'
   embed @IO (writeFile (toFilePath asmFile) code)
 
+runTreePipeline :: (Members '[Embed IO, App, TaggedLock] r) => PipelineArg -> Sem r ()
+runTreePipeline pa@PipelineArg {..} = do
+  entryPoint <- getEntry pa
+  treeFile <- Compile.outputFile _pipelineArgOptions _pipelineArgFile
+  r <-
+    runReader entryPoint
+      . runError @JuvixError
+      . coreToTree
+      $ _pipelineArgModule
+  tab' <- getRight r
+  let code = Tree.ppPrint tab' tab'
+  embed @IO (writeFile (toFilePath treeFile) code)
+
 runNockmaPipeline :: (Members '[Embed IO, App, TaggedLock] r) => PipelineArg -> Sem r ()
 runNockmaPipeline pa@PipelineArg {..} = do
   entryPoint <- getEntry pa

app/Commands/Dev/Core/Compile/Options.hs

@@ -17,7 +17,8 @@ coreSupportedTargets =
       TargetNative64,
       TargetGeb,
       TargetVampIR,
-      TargetAsm
+      TargetAsm,
+      TargetTree
     ]
 
 parseCoreCompileOptions :: Parser CoreCompileOptions

app/Commands/Dev/Options.hs

@@ -26,6 +26,7 @@ import Commands.Dev.Repl.Options
 import Commands.Dev.Runtime.Options
 import Commands.Dev.Scope.Options
 import Commands.Dev.Termination.Options
+import Commands.Dev.Tree.Options
 import Commands.Repl.Options
 import CommonOptions
 
@@ -36,6 +37,7 @@ data DevCommand
   | Core CoreCommand
   | Geb GebCommand
   | Asm AsmCommand
+  | Tree TreeCommand
   | Casm CasmCommand
   | Runtime RuntimeCommand
   | Parse ParseOptions
@@ -55,6 +57,7 @@ parseDevCommand =
           commandCore,
           commandGeb,
           commandAsm,
+          commandTree,
           commandCasm,
           commandRuntime,
           commandParse,
@@ -102,6 +105,13 @@ commandAsm =
       (Asm <$> parseAsmCommand)
       (progDesc "Subcommands related to JuvixAsm")
 
+commandTree :: Mod CommandFields DevCommand
+commandTree =
+  command "tree" $
+    info
+      (Tree <$> parseTreeCommand)
+      (progDesc "Subcommands related to JuvixTree")
+
 commandCasm :: Mod CommandFields DevCommand
 commandCasm =
   command "casm" $

app/Commands/Dev/Tree.hs

@@ -0,0 +1,9 @@
+module Commands.Dev.Tree where
+
+import Commands.Base
+import Commands.Dev.Tree.Options
+import Commands.Dev.Tree.Read as Read
+
+runCommand :: forall r. (Members '[Embed IO, App, TaggedLock] r) => TreeCommand -> Sem r ()
+runCommand = \case
+  Read opts -> Read.runCommand opts

app/Commands/Dev/Tree/Options.hs

@@ -0,0 +1,24 @@
+module Commands.Dev.Tree.Options where
+
+import Commands.Dev.Tree.Read.Options
+import CommonOptions
+
+newtype TreeCommand
+  = Read TreeReadOptions
+  deriving stock (Data)
+
+parseTreeCommand :: Parser TreeCommand
+parseTreeCommand =
+  hsubparser $
+    mconcat
+      [ commandRead
+      ]
+  where
+    commandRead :: Mod CommandFields TreeCommand
+    commandRead = command "read" readInfo
+
+    readInfo :: ParserInfo TreeCommand
+    readInfo =
+      info
+        (Read <$> parseTreeReadOptions)
+        (progDesc "Parse a JuvixTree file and pretty print it")

app/Commands/Dev/Tree/Read.hs

@@ -0,0 +1,17 @@
+module Commands.Dev.Tree.Read where
+
+import Commands.Base
+import Commands.Dev.Tree.Read.Options
+import Juvix.Compiler.Tree.Pretty qualified as Tree
+import Juvix.Compiler.Tree.Translation.FromSource qualified as Tree
+
+runCommand :: forall r. (Members '[Embed IO, App] r) => TreeReadOptions -> Sem r ()
+runCommand opts = do
+  afile :: Path Abs File <- fromAppPathFile file
+  s <- readFile (toFilePath afile)
+  case Tree.runParser (toFilePath afile) s of
+    Left err -> exitJuvixError (JuvixError err)
+    Right tab -> renderStdOut (Tree.ppOutDefault tab tab)
+  where
+    file :: AppPath File
+    file = opts ^. treeReadInputFile

app/Commands/Dev/Tree/Read/Options.hs

@@ -0,0 +1,15 @@
+module Commands.Dev.Tree.Read.Options where
+
+import CommonOptions
+
+newtype TreeReadOptions = TreeReadOptions
+  { _treeReadInputFile :: AppPath File
+  }
+  deriving stock (Data)
+
+makeLenses ''TreeReadOptions
+
+parseTreeReadOptions :: Parser TreeReadOptions
+parseTreeReadOptions = do
+  _treeReadInputFile <- parseInputFile FileExtJuvixTree
+  pure TreeReadOptions {..}

app/Commands/Extra/Compile.hs

@@ -33,6 +33,7 @@ runCompile inputFile o = do
     TargetVampIR -> return (Right ())
     TargetCore -> return (Right ())
     TargetAsm -> return (Right ())
+    TargetTree -> return (Right ())
     TargetNockma -> return (Right ())
 
 prepareRuntime :: forall r. (Members '[App, Embed IO] r) => Path Abs Dir -> CompileOptions -> Sem r ()
@@ -49,6 +50,7 @@ prepareRuntime buildDir o = do
     TargetVampIR -> return ()
     TargetCore -> return ()
     TargetAsm -> return ()
+    TargetTree -> return ()
     TargetNockma -> return ()
   where
     wasiReleaseRuntime :: BS.ByteString
@@ -107,6 +109,8 @@ outputFile opts inputFile =
           replaceExtension' juvixCoreFileExt baseOutputFile
         TargetAsm ->
           replaceExtension' juvixAsmFileExt baseOutputFile
+        TargetTree ->
+          replaceExtension' juvixTreeFileExt baseOutputFile
         TargetNockma ->
           replaceExtension' nockmaFileExt baseOutputFile
 

app/Commands/Extra/Compile/Options.hs

@@ -11,6 +11,7 @@ data CompileTarget
   | TargetVampIR
   | TargetCore
   | TargetAsm
+  | TargetTree
   | TargetNockma
   deriving stock (Eq, Data, Bounded, Enum)
 
@@ -22,6 +23,7 @@ instance Show CompileTarget where
     TargetVampIR -> "vampir"
     TargetCore -> "core"
     TargetAsm -> "asm"
+    TargetTree -> "tree"
     TargetNockma -> "nockma"
 
 data CompileOptions = CompileOptions

cntlines.sh

@@ -22,6 +22,7 @@ CASM=$(count src/Juvix/Compiler/Casm/)
 NOCK=$(count src/Juvix/Compiler/Nockma)
 REG=$(count src/Juvix/Compiler/Reg/)
 ASM=$(count src/Juvix/Compiler/Asm/)
+TREE=$(count src/Juvix/Compiler/Tree/)
 CORE=$(count src/Juvix/Compiler/Core/)
 
 CONCRETE=$(count src/Juvix/Compiler/Concrete/)
@@ -37,7 +38,7 @@ PRELUDE=$(count src/Juvix/Prelude/)
 STORE=$(count src/Juvix/Compiler/Store/)
 
 FRONT=$((CONCRETE + INTERNAL + BUILTINS + PIPELINE))
-BACK=$((BACKENDC + GEB + VAMPIR + NOCK + REG + ASM + CORE + CASM + CAIRO))
+BACK=$((BACKENDC + GEB + VAMPIR + NOCK + REG + ASM + TREE + CORE + CASM + CAIRO))
 OTHER=$((APP + STORE + HTML + EXTRA + DATA + PRELUDE))
 TESTS=$(count test/)
 
@@ -56,6 +57,7 @@ echo "   Cairo backend: $((CASM + CAIRO)) LOC"
 echo "   Nockma backend: $NOCK LOC"
 echo "   JuvixReg: $REG LOC"
 echo "   JuvixAsm: $ASM LOC"
+echo "   JuvixTree: $TREE LOC"
 echo "   JuvixCore: $CORE LOC"
 echo "Runtime: $RUNTIME LOC"
 echo "   C runtime: $RUNTIME_C LOC"

src/Juvix/Compiler/Asm/Data/InfoTable.hs

@@ -17,3 +17,10 @@ data FunctionInfoExtra = FunctionInfoExtra
   }
 
 makeLenses ''FunctionInfoExtra
+
+instance Semigroup FunctionInfoExtra where
+  fi1 <> fi2 =
+    FunctionInfoExtra
+      { _functionMaxValueStackHeight = max (fi1 ^. functionMaxValueStackHeight) (fi2 ^. functionMaxValueStackHeight),
+        _functionMaxTempStackHeight = max (fi1 ^. functionMaxTempStackHeight) (fi2 ^. functionMaxTempStackHeight)
+      }

src/Juvix/Compiler/Asm/Data/InfoTableBuilder.hs

@@ -1,80 +1,40 @@
-module Juvix.Compiler.Asm.Data.InfoTableBuilder where
+module Juvix.Compiler.Asm.Data.InfoTableBuilder
+  ( module Juvix.Compiler.Asm.Data.InfoTableBuilder,
+    module Juvix.Compiler.Tree.Data.InfoTableBuilder.Base,
+  )
+where
 
-import Data.HashMap.Strict qualified as HashMap
 import Juvix.Compiler.Asm.Data.InfoTable
 import Juvix.Compiler.Asm.Language
+import Juvix.Compiler.Tree.Data.InfoTableBuilder.Base
 
-data IdentKind
-  = IdentFun Symbol
-  | IdentFwd Symbol
-  | IdentInd Symbol
-  | IdentConstr Tag
+type InfoTableBuilder = InfoTableBuilder' Code (Maybe FunctionInfoExtra)
 
-data InfoTableBuilder m a where
-  FreshSymbol :: InfoTableBuilder m Symbol
-  FreshTag :: InfoTableBuilder m Tag
-  RegisterFunction :: FunctionInfo -> InfoTableBuilder m ()
-  RegisterConstr :: ConstructorInfo -> InfoTableBuilder m ()
-  RegisterInductive :: InductiveInfo -> InfoTableBuilder m ()
-  RegisterForward :: Text -> Symbol -> InfoTableBuilder m ()
-  RegisterMain :: Symbol -> InfoTableBuilder m ()
-  GetIdent :: Text -> InfoTableBuilder m (Maybe IdentKind)
-  GetFunctionInfo :: Symbol -> InfoTableBuilder m FunctionInfo
+type BuilderState = BuilderState' Code (Maybe FunctionInfoExtra)
 
-makeSem ''InfoTableBuilder
+freshSymbol :: (Member InfoTableBuilder r) => Sem r Symbol
+freshSymbol = freshSymbol' @Code @(Maybe FunctionInfoExtra)
 
-data BuilderState = BuilderState
-  { _stateNextSymbolId :: Word,
-    _stateNextUserTag :: Word,
-    _stateInfoTable :: InfoTable,
-    _stateIdents :: HashMap Text IdentKind
-  }
+freshTag :: (Member InfoTableBuilder r) => Sem r Tag
+freshTag = freshTag' @Code @(Maybe FunctionInfoExtra)
 
-makeLenses ''BuilderState
+registerFunction :: (Member InfoTableBuilder r) => FunctionInfo -> Sem r ()
+registerFunction = registerFunction' @Code @(Maybe FunctionInfoExtra)
 
-emptyBuilderState :: BuilderState
-emptyBuilderState =
-  BuilderState
-    { _stateNextSymbolId = 0,
-      _stateNextUserTag = 0,
-      _stateInfoTable = emptyInfoTable,
-      _stateIdents = mempty
-    }
+registerConstr :: (Member InfoTableBuilder r) => ConstructorInfo -> Sem r ()
+registerConstr = registerConstr' @Code @(Maybe FunctionInfoExtra)
 
-runInfoTableBuilder :: Sem (InfoTableBuilder ': r) a -> Sem r (InfoTable, a)
-runInfoTableBuilder = fmap (first (^. stateInfoTable)) . runInfoTableBuilder' emptyBuilderState
+registerInductive :: (Member InfoTableBuilder r) => InductiveInfo -> Sem r ()
+registerInductive = registerInductive' @Code @(Maybe FunctionInfoExtra)
 
-runInfoTableBuilder' :: BuilderState -> Sem (InfoTableBuilder ': r) a -> Sem r (BuilderState, a)
-runInfoTableBuilder' bs =
-  runState bs
-    . reinterpret interp
-  where
-    interp :: InfoTableBuilder m a -> Sem (State BuilderState ': r) a
-    interp = \case
-      FreshSymbol -> do
-        s <- get
-        modify' (over stateNextSymbolId (+ 1))
-        return (Symbol defaultModuleId (s ^. stateNextSymbolId))
-      FreshTag -> do
-        modify' (over stateNextUserTag (+ 1))
-        s <- get
-        return (UserTag (TagUser defaultModuleId (s ^. stateNextUserTag - 1)))
-      RegisterFunction fi -> do
-        modify' (over (stateInfoTable . infoFunctions) (HashMap.insert (fi ^. functionSymbol) fi))
-        modify' (over stateIdents (HashMap.insert (fi ^. functionName) (IdentFun (fi ^. functionSymbol))))
-      RegisterConstr ci -> do
-        modify' (over (stateInfoTable . infoConstrs) (HashMap.insert (ci ^. constructorTag) ci))
-        modify' (over stateIdents (HashMap.insert (ci ^. constructorName) (IdentConstr (ci ^. constructorTag))))
-      RegisterInductive ii -> do
-        modify' (over (stateInfoTable . infoInductives) (HashMap.insert (ii ^. inductiveSymbol) ii))
-        modify' (over stateIdents (HashMap.insert (ii ^. inductiveName) (IdentInd (ii ^. inductiveSymbol))))
-      RegisterForward txt sym ->
-        modify' (over stateIdents (HashMap.insert txt (IdentFwd sym)))
-      RegisterMain sym ->
-        modify' (over stateInfoTable (set infoMainFunction (Just sym)))
-      GetIdent txt -> do
-        s <- get
-        return $ HashMap.lookup txt (s ^. stateIdents)
-      GetFunctionInfo sym -> do
-        s <- get
-        return (lookupFunInfo (s ^. stateInfoTable) sym)
+registerForward :: (Member InfoTableBuilder r) => Text -> Symbol -> Sem r ()
+registerForward = registerForward' @Code @(Maybe FunctionInfoExtra)
+
+registerMain :: (Member InfoTableBuilder r) => Symbol -> Sem r ()
+registerMain = registerMain' @Code @(Maybe FunctionInfoExtra)
+
+getIdent :: (Member InfoTableBuilder r) => Text -> Sem r (Maybe IdentKind)
+getIdent = getIdent' @Code @(Maybe FunctionInfoExtra)
+
+getFunctionInfo :: (Member InfoTableBuilder r) => Symbol -> Sem r FunctionInfo
+getFunctionInfo = getFunctionInfo' @Code @(Maybe FunctionInfoExtra)