refactor: Add Anoma Node in Tree language (#2784)

Similarly to how the Cairo operations are handled we add a separate Tree
language Node for Anoma operations instead of handling them as an Unop
Node.

This is necessary because we need to add support for new Anoma
operations that are not unary.

This PR also adds support for `anoma-encode` and `anoma-decode`
functions in `jvt` tree source files which was missed in the previous
PRs.

src/Juvix/Compiler/Asm/Translation/FromTree.hs

@@ -39,6 +39,7 @@ genCode fi =
       Tree.Binop x -> goBinop isTail x
       Tree.Unop x -> goUnop isTail x
       Tree.Cairo x -> goCairo isTail x
+      Tree.Anoma {} -> error "Anoma instructions are not supported in the Asm backend"
       Tree.Constant x -> goConstant isTail x
       Tree.MemRef x -> goMemRef isTail x
       Tree.AllocConstr x -> goAllocConstr isTail x
@@ -233,9 +234,6 @@ genCode fi =
       Tree.PrimUnop op' -> mkUnop op'
       Tree.OpTrace -> mkInstr Trace
       Tree.OpFail -> mkInstr Failure
-      Tree.OpAnomaGet -> impossible
-      Tree.OpAnomaEncode -> impossible
-      Tree.OpAnomaDecode -> impossible
 
     snocReturn :: Bool -> Code' -> Code'
     snocReturn True code = DL.snoc code (mkInstr Return)

src/Juvix/Compiler/Core/Language/Builtins.hs

@@ -121,6 +121,9 @@ builtinIsFoldable = \case
 builtinIsCairo :: BuiltinOp -> Bool
 builtinIsCairo op = op `elem` builtinsCairo
 
+builtinIsAnoma :: BuiltinOp -> Bool
+builtinIsAnoma op = op `elem` builtinsAnoma
+
 builtinsString :: [BuiltinOp]
 builtinsString = [OpStrConcat, OpStrToInt, OpShow]
 

src/Juvix/Compiler/Nockma/Translation/FromTree.hs

@@ -330,6 +330,7 @@ compile = \case
   Tree.Binop b -> goBinop b
   Tree.Unop b -> goUnop b
   Tree.Cairo {} -> cairoErr
+  Tree.Anoma b -> goAnomaOp b
   Tree.Constant c -> return (goConstant (c ^. Tree.nodeConstant))
   Tree.MemRef c -> goMemRef (c ^. Tree.nodeMemRef)
   Tree.AllocConstr c -> goAllocConstr c
@@ -418,16 +419,21 @@ compile = \case
       iffalse <- compile _nodeBranchFalse
       return (branch arg iftrue iffalse)
 
+    goAnomaOp :: Tree.NodeAnoma -> Sem r (Term Natural)
+    goAnomaOp Tree.NodeAnoma {..} = do
+      args <- mapM compile _nodeAnomaArgs
+      case _nodeAnomaOpcode of
+        Tree.OpAnomaGet -> goAnomaGet args
+        Tree.OpAnomaEncode -> goAnomaEncode args
+        Tree.OpAnomaDecode -> goAnomaDecode args
+
     goUnop :: Tree.NodeUnop -> Sem r (Term Natural)
     goUnop Tree.NodeUnop {..} = do
       arg <- compile _nodeUnopArg
       case _nodeUnopOpcode of
         Tree.PrimUnop op -> return $ goPrimUnop op arg
         Tree.OpFail -> return crash
         Tree.OpTrace -> goTrace arg
-        Tree.OpAnomaGet -> goAnomaGet arg
-        Tree.OpAnomaEncode -> goAnomaEncode arg
-        Tree.OpAnomaDecode -> goAnomaDecode arg
 
     goPrimUnop :: Tree.UnaryOp -> Term Natural -> Term Natural
     goPrimUnop op arg = case op of
@@ -439,16 +445,16 @@ compile = \case
       Tree.OpIntToField -> fieldErr
       Tree.OpFieldToInt -> fieldErr
 
-    goAnomaGet :: Term Natural -> Sem r (Term Natural)
+    goAnomaGet :: [Term Natural] -> Sem r (Term Natural)
     goAnomaGet key = do
-      let arg = remakeList [getFieldInSubject AnomaGetOrder, key]
+      let arg = remakeList [getFieldInSubject AnomaGetOrder, foldTermsOrNil key]
       return (OpScry # (OpQuote # nockNilTagged "OpScry-typehint") # arg)
 
-    goAnomaEncode :: Term Natural -> Sem r (Term Natural)
-    goAnomaEncode arg = return (callStdlib StdlibEncode [arg])
+    goAnomaEncode :: [Term Natural] -> Sem r (Term Natural)
+    goAnomaEncode args = return (callStdlib StdlibEncode args)
 
-    goAnomaDecode :: Term Natural -> Sem r (Term Natural)
-    goAnomaDecode arg = return (callStdlib StdlibDecode [arg])
+    goAnomaDecode :: [Term Natural] -> Sem r (Term Natural)
+    goAnomaDecode args = return (callStdlib StdlibDecode args)
 
     goTrace :: Term Natural -> Sem r (Term Natural)
     goTrace arg = do

src/Juvix/Compiler/Tree/Evaluator.hs

@@ -37,6 +37,7 @@ hEval hout tab = eval' [] mempty
     eval' args temps node = case node of
       Binop x -> goBinop x
       Unop x -> goUnop x
+      Anoma {} -> evalError "unsupported: Anoma builtin"
       Cairo {} -> evalError "unsupported: Cairo builtin"
       Constant c -> goConstant c
       MemRef x -> goMemRef x
@@ -74,9 +75,6 @@ hEval hout tab = eval' [] mempty
                 PrimUnop op -> eitherToError $ evalUnop tab op v
                 OpTrace -> goTrace v
                 OpFail -> goFail v
-                OpAnomaGet -> evalError "Unsupported op: OpAnomaGet"
-                OpAnomaEncode -> evalError "Unsupported op: OpAnomaEncode"
-                OpAnomaDecode -> evalError "Unsupported op: OpAnomaDecode"
 
         goFail :: Value -> Value
         goFail v = evalError ("failure: " <> printValue tab v)

src/Juvix/Compiler/Tree/EvaluatorEff.hs

@@ -33,6 +33,7 @@ eval tab = runReader emptyEvalCtx . eval'
     eval' node = case node of
       Binop x -> goBinop x
       Unop x -> goUnop x
+      Anoma {} -> evalError "unsupported: Anoma builtins"
       Cairo {} -> evalError "unsupported: Cairo builtins"
       Constant c -> return (goConstant c)
       MemRef x -> goMemRef x
@@ -69,9 +70,6 @@ eval tab = runReader emptyEvalCtx . eval'
             PrimUnop op -> eitherToError $ evalUnop tab op v
             OpTrace -> goTrace v
             OpFail -> goFail v
-            OpAnomaGet -> evalError "Unsupported op: OpAnomaGet"
-            OpAnomaEncode -> evalError "Unsupported op: OpAnomaEncode"
-            OpAnomaDecode -> evalError "Unsupported op: OpAnomaDecode"
 
         goFail :: Value -> Sem r' Value
         goFail v = evalError ("failure: " <> printValue tab v)

src/Juvix/Compiler/Tree/Extra/Base.hs

@@ -19,6 +19,7 @@ getNodeInfo = \case
   Binop NodeBinop {..} -> _nodeBinopInfo
   Unop NodeUnop {..} -> _nodeUnopInfo
   Cairo NodeCairo {..} -> _nodeCairoInfo
+  Anoma NodeAnoma {..} -> _nodeAnomaInfo
   Constant NodeConstant {..} -> _nodeConstantInfo
   MemRef NodeMemRef {..} -> _nodeMemRefInfo
   AllocConstr NodeAllocConstr {..} -> _nodeAllocConstrInfo
@@ -152,6 +153,17 @@ destruct = \case
                 _nodeCairoInfo
               }
       }
+  Anoma NodeAnoma {..} ->
+    NodeDetails
+      { _nodeChildren = map noTempVar _nodeAnomaArgs,
+        _nodeReassemble = manyChildren $ \args ->
+          Anoma
+            NodeAnoma
+              { _nodeAnomaArgs = args,
+                _nodeAnomaOpcode,
+                _nodeAnomaInfo
+              }
+      }
   Constant c ->
     NodeDetails
       { _nodeChildren = [],

src/Juvix/Compiler/Tree/Keywords.hs

@@ -9,6 +9,8 @@ import Juvix.Compiler.Tree.Keywords.Base
 import Juvix.Data.Keyword.All
   ( kwAdd_,
     kwAlloc,
+    kwAnomaDecode,
+    kwAnomaEncode,
     kwAnomaGet,
     kwArgsNum,
     kwAtoi,
@@ -74,6 +76,8 @@ allKeywords =
          kwCase,
          kwSave,
          kwAnomaGet,
+         kwAnomaDecode,
+         kwAnomaEncode,
          kwPoseidon,
          kwEcOp,
          kwRandomEcPoint

src/Juvix/Compiler/Tree/Language.hs

@@ -17,6 +17,7 @@ data Node
   = Binop NodeBinop
   | Unop NodeUnop
   | Cairo NodeCairo
+  | Anoma NodeAnoma
   | -- | A constant value.
     Constant NodeConstant
   | -- | A memory reference.
@@ -67,12 +68,6 @@ data UnaryOpcode
     OpTrace
   | -- | Interrupt execution with a runtime error printing the argument.
     OpFail
-  | -- | Get a value by key from Anoma storage
-    OpAnomaGet
-  | -- | Encode a value to an Anoma atom
-    OpAnomaEncode
-  | -- | Decode a value from an Anoma atom
-    OpAnomaDecode
 
 data NodeBinop = NodeBinop
   { _nodeBinopInfo :: NodeInfo,
@@ -93,6 +88,12 @@ data NodeCairo = NodeCairo
     _nodeCairoArgs :: [Node]
   }
 
+data NodeAnoma = NodeAnoma
+  { _nodeAnomaInfo :: NodeInfo,
+    _nodeAnomaOpcode :: AnomaOp,
+    _nodeAnomaArgs :: [Node]
+  }
+
 data NodeConstant = NodeConstant
   { _nodeConstantInfo :: NodeInfo,
     _nodeConstant :: Constant

src/Juvix/Compiler/Tree/Language/Builtins.hs

@@ -47,3 +47,13 @@ cairoOpArgsNum = \case
   OpCairoPoseidon -> 1
   OpCairoEc -> 3
   OpCairoRandomEcPoint -> 0
+
+-- | Builtin Anoma operations. Implemented only in the Anoma backend.
+data AnomaOp
+  = -- | Get a value by key from Anoma storage
+    OpAnomaGet
+  | -- | Encode a value to an Anoma atom
+    OpAnomaEncode
+  | -- | Decode a value from an Anoma atom
+    OpAnomaDecode
+  deriving stock (Eq)

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

@@ -240,14 +240,17 @@ instance PrettyCode CairoOp where
     OpCairoEc -> Str.instrEcOp
     OpCairoRandomEcPoint -> Str.cairoRandomEcPoint
 
+instance PrettyCode AnomaOp where
+  ppCode op = return . primitive $ case op of
+    OpAnomaGet -> Str.anomaGet
+    OpAnomaEncode -> Str.anomaEncode
+    OpAnomaDecode -> Str.anomaDecode
+
 instance PrettyCode UnaryOpcode where
   ppCode = \case
     PrimUnop x -> ppCode x
     OpTrace -> return $ primitive Str.instrTrace
     OpFail -> return $ primitive Str.instrFailure
-    OpAnomaGet -> return $ primitive Str.anomaGet
-    OpAnomaEncode -> return $ primitive Str.anomaEncode
-    OpAnomaDecode -> return $ primitive Str.anomaDecode
 
 instance PrettyCode NodeUnop where
   ppCode NodeUnop {..} = do
@@ -261,6 +264,12 @@ instance PrettyCode NodeCairo where
     args <- ppCodeArgs _nodeCairoArgs
     return $ op <> parens args
 
+instance PrettyCode NodeAnoma where
+  ppCode NodeAnoma {..} = do
+    op <- ppCode _nodeAnomaOpcode
+    args <- ppCodeArgs _nodeAnomaArgs
+    return (op <> parens args)
+
 instance PrettyCode NodeConstant where
   ppCode NodeConstant {..} = ppCode _nodeConstant
 
@@ -353,6 +362,7 @@ instance PrettyCode Node where
   ppCode = \case
     Binop x -> ppCode x
     Unop x -> ppCode x
+    Anoma x -> ppCode x
     Cairo x -> ppCode x
     Constant x -> ppCode x
     MemRef x -> ppCode x

src/Juvix/Compiler/Tree/Transformation/CheckNoAnoma.hs

@@ -10,19 +10,16 @@ checkNoAnoma = walkT checkNode
   where
     checkNode :: Symbol -> Node -> Sem r ()
     checkNode _ = \case
-      Unop NodeUnop {..} -> case _nodeUnopOpcode of
+      Anoma NodeAnoma {..} -> case _nodeAnomaOpcode of
         OpAnomaGet -> unsupportedErr "OpAnomaGet"
         OpAnomaEncode -> unsupportedErr "OpAnomaEncode"
         OpAnomaDecode -> unsupportedErr "OpAnomaDecode"
-        OpFail -> return ()
-        OpTrace -> return ()
-        PrimUnop {} -> return ()
         where
           unsupportedErr :: Text -> Sem r ()
           unsupportedErr opName =
             throw
               TreeError
                 { _treeErrorMsg = opName <> " is unsupported",
-                  _treeErrorLoc = _nodeUnopInfo ^. nodeInfoLocation
+                  _treeErrorLoc = _nodeAnomaInfo ^. nodeInfoLocation
                 }
       _ -> return ()

src/Juvix/Compiler/Tree/Transformation/Validate.hs

@@ -16,6 +16,7 @@ inferType tab funInfo = goInfer mempty
       Binop x -> goBinop bl x
       Unop x -> goUnop bl x
       Cairo x -> goCairo bl x
+      Anoma x -> goAnoma bl x
       Constant x -> goConst bl x
       MemRef x -> goMemRef bl x
       AllocConstr x -> goAllocConstr bl x
@@ -65,9 +66,6 @@ inferType tab funInfo = goInfer mempty
       PrimUnop x -> checkPrimUnop x
       OpTrace -> goInfer bl _nodeUnopArg
       OpFail -> checkUnop TyDynamic TyDynamic
-      OpAnomaGet -> checkUnop TyDynamic TyDynamic
-      OpAnomaEncode -> checkUnop TyDynamic TyDynamic
-      OpAnomaDecode -> checkUnop TyDynamic TyDynamic
       where
         loc = _nodeUnopInfo ^. nodeInfoLocation
 
@@ -90,6 +88,11 @@ inferType tab funInfo = goInfer mempty
       mapM_ (\arg -> checkType bl arg TyDynamic) _nodeCairoArgs
       return TyDynamic
 
+    goAnoma :: BinderList Type -> NodeAnoma -> Sem r Type
+    goAnoma bl NodeAnoma {..} = do
+      mapM_ (\arg -> checkType bl arg TyDynamic) _nodeAnomaArgs
+      return TyDynamic
+
     goConst :: BinderList Type -> NodeConstant -> Sem r Type
     goConst _ NodeConstant {..} = case _nodeConstant of
       ConstInt {} -> return mkTypeInteger

src/Juvix/Compiler/Tree/Translation/FromCore.hs

@@ -145,6 +145,13 @@ genCode infoTable fi =
                 _nodeCairoOpcode = genCairoOp _builtinAppOp,
                 _nodeCairoArgs = args
               }
+      | Core.builtinIsAnoma _builtinAppOp =
+          Anoma $
+            NodeAnoma
+              { _nodeAnomaInfo = mempty,
+                _nodeAnomaOpcode = genAnomaOp _builtinAppOp,
+                _nodeAnomaArgs = args
+              }
       | otherwise =
           case args of
             [arg] ->
@@ -295,9 +302,6 @@ genCode infoTable fi =
       Core.OpFieldToInt -> PrimUnop OpFieldToInt
       Core.OpTrace -> OpTrace
       Core.OpFail -> OpFail
-      Core.OpAnomaGet -> OpAnomaGet
-      Core.OpAnomaEncode -> OpAnomaEncode
-      Core.OpAnomaDecode -> OpAnomaDecode
       _ -> impossible
 
     genCairoOp :: Core.BuiltinOp -> CairoOp
@@ -307,6 +311,13 @@ genCode infoTable fi =
       Core.OpRandomEcPoint -> OpCairoRandomEcPoint
       _ -> impossible
 
+    genAnomaOp :: Core.BuiltinOp -> AnomaOp
+    genAnomaOp = \case
+      Core.OpAnomaGet -> OpAnomaGet
+      Core.OpAnomaEncode -> OpAnomaEncode
+      Core.OpAnomaDecode -> OpAnomaDecode
+      _ -> impossible
+
     getArgsNum :: Symbol -> Int
     getArgsNum sym =
       fromMaybe

src/Juvix/Compiler/Tree/Translation/FromSource.hs

@@ -53,6 +53,7 @@ parseNode ::
 parseNode =
   (Binop <$> parseBinop)
     <|> (Unop <$> parseUnop)
+    <|> (Anoma <$> parseAnoma)
     <|> (Cairo <$> parseCairo)
     <|> (Constant <$> parseConst)
     <|> (AllocConstr <$> parseAlloc)
@@ -107,7 +108,6 @@ parseUnop =
     <|> parseUnaryOp kwTrace OpTrace
     <|> parseUnaryOp kwFail OpFail
     <|> parseUnaryOp kwArgsNum (PrimUnop OpArgsNum)
-    <|> parseUnaryOp kwAnomaGet (OpAnomaGet)
 
 parseUnaryOp ::
   (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>
@@ -119,6 +119,24 @@ parseUnaryOp kwd op = do
   arg <- parens parseNode
   return $ NodeUnop (NodeInfo (Just loc)) op arg
 
+parseAnoma ::
+  (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>
+  ParsecS r NodeAnoma
+parseAnoma =
+  parseAnoma' kwAnomaGet OpAnomaGet
+    <|> parseAnoma' kwAnomaDecode OpAnomaDecode
+    <|> parseAnoma' kwAnomaEncode OpAnomaEncode
+
+parseAnoma' ::
+  (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>
+  Keyword ->
+  AnomaOp ->
+  ParsecS r NodeAnoma
+parseAnoma' kwd op = do
+  loc <- onlyInterval (kw kwd)
+  args <- parseArgs
+  return $ NodeAnoma (NodeInfo (Just loc)) op args
+
 parseCairo ::
   (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>
   ParsecS r NodeCairo

src/Juvix/Data/Keyword/All.hs

@@ -445,6 +445,12 @@ kwLive = asciiKw Str.live
 kwAnomaGet :: Keyword
 kwAnomaGet = asciiKw Str.anomaGet
 
+kwAnomaDecode :: Keyword
+kwAnomaDecode = asciiKw Str.anomaDecode
+
+kwAnomaEncode :: Keyword
+kwAnomaEncode = asciiKw Str.anomaEncode
+
 delimBraceL :: Keyword
 delimBraceL = mkDelim Str.braceL