JuvixReg recursors (#2641)

* Implements JuvixReg recursors, which will allow to implement JuvixReg
transformations more succinctly and effectively.
* Adds a transformation framework to JuvixReg.
* Adds identity transformation tests.
* Depends on #2635

app/Commands/Dev/Reg/Read.hs

@@ -3,7 +3,9 @@ module Commands.Dev.Reg.Read where
 import Commands.Base
 import Commands.Dev.Reg.Read.Options
 import Juvix.Compiler.Reg.Pretty qualified as Reg
+import Juvix.Compiler.Reg.Transformation qualified as Reg
 import Juvix.Compiler.Reg.Translation.FromSource qualified as Reg
+import RegInterpreter
 
 runCommand :: forall r. (Members '[EmbedIO, App] r) => RegReadOptions -> Sem r ()
 runCommand opts = do
@@ -12,8 +14,23 @@ runCommand opts = do
   case Reg.runParser (toFilePath afile) s of
     Left err ->
       exitJuvixError (JuvixError err)
-    Right tab ->
-      renderStdOut (Reg.ppOutDefault tab tab)
+    Right tab -> do
+      r <- runError @JuvixError (Reg.applyTransformations (project opts ^. regReadTransformations) tab)
+      case r of
+        Left err -> exitJuvixError (JuvixError err)
+        Right tab' -> do
+          unless (project opts ^. regReadNoPrint) $
+            renderStdOut (Reg.ppOutDefault tab' tab')
+          doRun tab'
   where
     file :: AppPath File
     file = opts ^. regReadInputFile
+
+    doRun :: Reg.InfoTable -> Sem r ()
+    doRun tab'
+      | project opts ^. regReadRun = do
+          putStrLn "--------------------------------"
+          putStrLn "|            Run               |"
+          putStrLn "--------------------------------"
+          runReg tab'
+      | otherwise = return ()

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

@@ -1,15 +1,30 @@
 module Commands.Dev.Reg.Read.Options where
 
 import CommonOptions
+import Juvix.Compiler.Reg.Data.TransformationId
 
-newtype RegReadOptions = RegReadOptions
-  { _regReadInputFile :: AppPath File
+data RegReadOptions = RegReadOptions
+  { _regReadTransformations :: [TransformationId],
+    _regReadRun :: Bool,
+    _regReadNoPrint :: Bool,
+    _regReadInputFile :: AppPath File
   }
   deriving stock (Data)
 
 makeLenses ''RegReadOptions
 
 parseRegReadOptions :: Parser RegReadOptions
 parseRegReadOptions = do
-  _regReadInputFile <- parseInputFile FileExtJuvixAsm
+  _regReadNoPrint <-
+    switch
+      ( long "no-print"
+          <> help "Do not print the transformed code"
+      )
+  _regReadRun <-
+    switch
+      ( long "run"
+          <> help "Run the code after the transformation"
+      )
+  _regReadTransformations <- optRegTransformationIds
+  _regReadInputFile <- parseInputFile FileExtJuvixReg
   pure RegReadOptions {..}

app/CommonOptions.hs

@@ -9,6 +9,7 @@ where
 import Control.Exception qualified as GHC
 import Data.List.NonEmpty qualified as NonEmpty
 import Juvix.Compiler.Core.Data.TransformationId.Parser qualified as Core
+import Juvix.Compiler.Reg.Data.TransformationId.Parser qualified as Reg
 import Juvix.Compiler.Tree.Data.TransformationId.Parser qualified as Tree
 import Juvix.Data.FileExt
 import Juvix.Prelude
@@ -217,32 +218,26 @@ optNoDisambiguate =
         <> help "Don't disambiguate the names of bound variables"
     )
 
-optCoreTransformationIds :: Parser [Core.TransformationId]
-optCoreTransformationIds =
+optTransformationIds :: forall a. (Text -> Either Text [a]) -> (String -> [String]) -> Parser [a]
+optTransformationIds parseIds completions =
   option
     (eitherReader parseTransf)
     ( long "transforms"
         <> short 't'
         <> value []
         <> metavar "[Transform]"
-        <> completer (mkCompleter (return . Core.completionsString))
+        <> completer (mkCompleter (return . completions))
         <> help "hint: use autocomplete"
     )
   where
-    parseTransf :: String -> Either String [Core.TransformationId]
-    parseTransf = mapLeft unpack . Core.parseTransformations . pack
+    parseTransf :: String -> Either String [a]
+    parseTransf = mapLeft unpack . parseIds . pack
+
+optCoreTransformationIds :: Parser [Core.TransformationId]
+optCoreTransformationIds = optTransformationIds Core.parseTransformations Core.completionsString
 
 optTreeTransformationIds :: Parser [Tree.TransformationId]
-optTreeTransformationIds =
-  option
-    (eitherReader parseTransf)
-    ( long "transforms"
-        <> short 't'
-        <> value []
-        <> metavar "[Transform]"
-        <> completer (mkCompleter (return . Tree.completionsString))
-        <> help "hint: use autocomplete"
-    )
-  where
-    parseTransf :: String -> Either String [Tree.TransformationId]
-    parseTransf = mapLeft unpack . Tree.parseTransformations . pack
+optTreeTransformationIds = optTransformationIds Tree.parseTransformations Tree.completionsString
+
+optRegTransformationIds :: Parser [Reg.TransformationId]
+optRegTransformationIds = optTransformationIds Reg.parseTransformations Reg.completionsString

src/Juvix/Compiler/Backend/C/Translation/FromReg.hs

@@ -8,7 +8,7 @@ import Juvix.Compiler.Backend.C.Extra.Serialization
 import Juvix.Compiler.Backend.C.Language as C
 import Juvix.Compiler.Backend.C.Translation.FromReg.Base
 import Juvix.Compiler.Reg.Data.InfoTable qualified as Reg
-import Juvix.Compiler.Reg.Extra qualified as Reg
+import Juvix.Compiler.Reg.Extra.Info qualified as Reg
 import Juvix.Compiler.Reg.Language qualified as Reg
 import Juvix.Data.Fixity qualified as Fixity
 import Juvix.Prelude

src/Juvix/Compiler/Backend/C/Translation/FromReg/Base.hs

@@ -4,7 +4,7 @@ import Data.HashMap.Strict qualified as HashMap
 import Data.Text qualified as T
 import Juvix.Compiler.Backend.C.Language
 import Juvix.Compiler.Reg.Data.InfoTable qualified as Reg
-import Juvix.Compiler.Reg.Extra qualified as Reg
+import Juvix.Compiler.Reg.Extra.Info qualified as Reg
 import Juvix.Compiler.Reg.Language qualified as Reg
 import Juvix.Prelude
 

src/Juvix/Compiler/Reg/Data/TransformationId.hs

@@ -0,0 +1,38 @@
+module Juvix.Compiler.Reg.Data.TransformationId where
+
+import Juvix.Compiler.Core.Data.TransformationId.Base
+import Juvix.Compiler.Reg.Data.TransformationId.Strings
+import Juvix.Prelude
+
+data TransformationId
+  = Identity
+  deriving stock (Data, Bounded, Enum, Show)
+
+data PipelineId
+  = PipelineC
+  | PipelineCairo
+  deriving stock (Data, Bounded, Enum)
+
+type TransformationLikeId = TransformationLikeId' TransformationId PipelineId
+
+toCTransformations :: [TransformationId]
+toCTransformations = []
+
+toCairoTransformations :: [TransformationId]
+toCairoTransformations = []
+
+instance TransformationId' TransformationId where
+  transformationText :: TransformationId -> Text
+  transformationText = \case
+    Identity -> strIdentity
+
+instance PipelineId' TransformationId PipelineId where
+  pipelineText :: PipelineId -> Text
+  pipelineText = \case
+    PipelineC -> strCPipeline
+    PipelineCairo -> strCairoPipeline
+
+  pipeline :: PipelineId -> [TransformationId]
+  pipeline = \case
+    PipelineC -> toCTransformations
+    PipelineCairo -> toCairoTransformations

src/Juvix/Compiler/Reg/Data/TransformationId/Parser.hs

@@ -0,0 +1,14 @@
+module Juvix.Compiler.Reg.Data.TransformationId.Parser (parseTransformations, TransformationId (..), completions, completionsString) where
+
+import Juvix.Compiler.Core.Data.TransformationId.Parser.Base
+import Juvix.Compiler.Reg.Data.TransformationId
+import Juvix.Prelude
+
+parseTransformations :: Text -> Either Text [TransformationId]
+parseTransformations = parseTransformations' @TransformationId @PipelineId
+
+completionsString :: String -> [String]
+completionsString = completionsString' @TransformationId @PipelineId
+
+completions :: Text -> [Text]
+completions = completions' @TransformationId @PipelineId

src/Juvix/Compiler/Reg/Data/TransformationId/Strings.hs

@@ -0,0 +1,12 @@
+module Juvix.Compiler.Reg.Data.TransformationId.Strings where
+
+import Juvix.Prelude
+
+strCPipeline :: Text
+strCPipeline = "pipeline-c"
+
+strCairoPipeline :: Text
+strCairoPipeline = "pipeline-cairo"
+
+strIdentity :: Text
+strIdentity = "identity"

src/Juvix/Compiler/Reg/Extra.hs → src/Juvix/Compiler/Reg/Extra/Info.hs

@@ -1,4 +1,4 @@
-module Juvix.Compiler.Reg.Extra where
+module Juvix.Compiler.Reg.Extra.Info where
 
 import Data.HashMap.Strict qualified as HashMap
 import Juvix.Compiler.Backend

src/Juvix/Compiler/Reg/Extra/Recursors.hs

@@ -0,0 +1,127 @@
+module Juvix.Compiler.Reg.Extra.Recursors where
+
+import Data.Functor.Identity
+import Juvix.Compiler.Reg.Language
+
+data ForwardRecursorSig m c = ForwardRecursorSig
+  { _forwardFun :: Instruction -> c -> m (c, Instruction),
+    _forwardCombine :: NonEmpty c -> c
+  }
+
+data BackwardRecursorSig m a = BackwardRecursorSig
+  { _backwardFun :: Code -> a -> [a] -> m (a, Code),
+    _backwardAdjust :: a -> a
+  }
+
+makeLenses ''ForwardRecursorSig
+makeLenses ''BackwardRecursorSig
+
+recurseF :: forall m c. (Monad m) => ForwardRecursorSig m c -> c -> Code -> m (c, Code)
+recurseF sig c = \case
+  i : instrs -> do
+    (c0, i0) <- (sig ^. forwardFun) i c
+    (c', i') <-
+      case i0 of
+        Branch x@InstrBranch {..} -> do
+          (c1, is1) <- recurseF sig c0 _instrBranchTrue
+          (c2, is2) <- recurseF sig c0 _instrBranchFalse
+          let c' = (sig ^. forwardCombine) (c1 :| [c2])
+          return (c', Branch x {_instrBranchTrue = is1, _instrBranchFalse = is2})
+        Case x@InstrCase {..} -> do
+          brs' <- mapM goBranch _instrCaseBranches
+          def' <- maybe (return Nothing) (\is -> Just <$> recurseF sig c0 is) _instrCaseDefault
+          let cs = map fst brs' ++ maybe [] (\md -> [fst md]) def'
+              brs = map snd brs'
+              def = maybe Nothing (Just . snd) def'
+              c' = (sig ^. forwardCombine) (nonEmpty' cs)
+          return (c', Case x {_instrCaseBranches = brs, _instrCaseDefault = def})
+          where
+            goBranch :: CaseBranch -> m (c, CaseBranch)
+            goBranch br@CaseBranch {..} = do
+              (c', is') <- recurseF sig c0 _caseBranchCode
+              return (c', br {_caseBranchCode = is'})
+        Block x@InstrBlock {..} -> do
+          (c', is) <- recurseF sig c0 _instrBlockCode
+          return (c', Block x {_instrBlockCode = is})
+        _ ->
+          return (c0, i0)
+    (c'', instrs') <- recurseF sig c' instrs
+    return (c'', i' : instrs')
+  [] ->
+    return (c, [])
+
+recurseB :: forall m a. (Monad m) => BackwardRecursorSig m a -> a -> Code -> m (a, Code)
+recurseB sig a = \case
+  i : instrs -> do
+    (a', instrs') <- recurseB sig a instrs
+    let a0 = (sig ^. backwardAdjust) a'
+    (as, i') <-
+      case i of
+        Branch x@InstrBranch {..} -> do
+          (a1, is1) <- recurseB sig a0 _instrBranchTrue
+          (a2, is2) <- recurseB sig a0 _instrBranchFalse
+          return ([a1, a2], Branch x {_instrBranchTrue = is1, _instrBranchFalse = is2})
+        Case x@InstrCase {..} -> do
+          brs' <- mapM goBranch _instrCaseBranches
+          def' <- maybe (return Nothing) (\is -> Just <$> recurseB sig a0 is) _instrCaseDefault
+          let as = map fst brs' ++ maybe [] (\md -> [fst md]) def'
+              brs = map snd brs'
+              def = maybe Nothing (Just . snd) def'
+          return (as, Case x {_instrCaseBranches = brs, _instrCaseDefault = def})
+          where
+            goBranch :: CaseBranch -> m (a, CaseBranch)
+            goBranch br@CaseBranch {..} = do
+              (aa, is') <- recurseB sig a0 _caseBranchCode
+              return (aa, br {_caseBranchCode = is'})
+        Block x@InstrBlock {..} -> do
+          (aa, is) <- recurseB sig a0 _instrBlockCode
+          return ([aa], Block x {_instrBlockCode = is})
+        _ ->
+          return ([], i)
+    (sig ^. backwardFun) (i' : instrs') a' as
+  [] ->
+    (sig ^. backwardFun) [] a []
+
+cmapM :: (Monad m) => (Code -> m Code) -> Code -> m Code
+cmapM f is0 = do
+  ((), is) <-
+    recurseB
+      BackwardRecursorSig
+        { _backwardFun = \is _ _ -> do
+            is' <- f is
+            return ((), is'),
+          _backwardAdjust = id
+        }
+      ()
+      is0
+  return is
+
+cmap :: (Code -> Code) -> Code -> Code
+cmap f is = runIdentity (cmapM (return . f) is)
+
+imapM :: (Monad m) => (Instruction -> m Instruction) -> Code -> m Code
+imapM f = cmapM $ \case
+  i : is -> do
+    i' <- f i
+    return (i' : is)
+  [] ->
+    return []
+
+imap :: (Instruction -> Instruction) -> Code -> Code
+imap f is = runIdentity (imapM (return . f) is)
+
+ifoldFM :: (Monad m, Monoid a) => (a -> Instruction -> m a) -> a -> Code -> m a
+ifoldFM f a0 is0 =
+  fst
+    <$> recurseF
+      ForwardRecursorSig
+        { _forwardFun = \i a -> do
+            a' <- f a i
+            return (a', i),
+          _forwardCombine = mconcat . toList
+        }
+      a0
+      is0
+
+ifoldF :: (Monoid a) => (a -> Instruction -> a) -> a -> Code -> a
+ifoldF f a is = runIdentity (ifoldFM (\a' -> return . f a') a is)

src/Juvix/Compiler/Reg/Interpreter.hs

@@ -11,7 +11,7 @@ import Data.Vector qualified as Vec
 import Data.Vector.Mutable qualified as MV
 import Juvix.Compiler.Reg.Data.InfoTable
 import Juvix.Compiler.Reg.Error
-import Juvix.Compiler.Reg.Extra
+import Juvix.Compiler.Reg.Extra.Info
 import Juvix.Compiler.Reg.Interpreter.Base
 import Juvix.Compiler.Reg.Interpreter.Error
 import Juvix.Compiler.Reg.Pretty

src/Juvix/Compiler/Reg/Transformation.hs

@@ -0,0 +1,17 @@
+module Juvix.Compiler.Reg.Transformation
+  ( module Juvix.Compiler.Reg.Transformation.Base,
+    module Juvix.Compiler.Reg.Transformation,
+    module Juvix.Compiler.Reg.Data.TransformationId,
+  )
+where
+
+import Juvix.Compiler.Reg.Data.TransformationId
+import Juvix.Compiler.Reg.Transformation.Base
+import Juvix.Compiler.Reg.Transformation.Identity
+
+applyTransformations :: forall r. [TransformationId] -> InfoTable -> Sem r InfoTable
+applyTransformations ts tbl = foldM (flip appTrans) tbl ts
+  where
+    appTrans :: TransformationId -> InfoTable -> Sem r InfoTable
+    appTrans = \case
+      Identity -> return . identity

src/Juvix/Compiler/Reg/Transformation/Base.hs

@@ -0,0 +1,10 @@
+module Juvix.Compiler.Reg.Transformation.Base
+  ( module Juvix.Compiler.Tree.Transformation.Generic.Base,
+    module Juvix.Compiler.Reg.Data.InfoTable,
+    module Juvix.Compiler.Reg.Language,
+  )
+where
+
+import Juvix.Compiler.Reg.Data.InfoTable
+import Juvix.Compiler.Reg.Language
+import Juvix.Compiler.Tree.Transformation.Generic.Base

src/Juvix/Compiler/Reg/Transformation/Identity.hs

@@ -0,0 +1,7 @@
+module Juvix.Compiler.Reg.Transformation.Identity where
+
+import Juvix.Compiler.Reg.Extra.Recursors
+import Juvix.Compiler.Reg.Transformation.Base
+
+identity :: InfoTable -> InfoTable
+identity = mapT (const (cmap id))

src/Juvix/Compiler/Tree/Transformation.hs

@@ -5,6 +5,7 @@ module Juvix.Compiler.Tree.Transformation
   )
 where
 
+import Juvix.Compiler.Tree.Data.InfoTable
 import Juvix.Compiler.Tree.Data.TransformationId
 import Juvix.Compiler.Tree.Error
 import Juvix.Compiler.Tree.Transformation.Apply