Add support for anoma-decode builtin (#2775)

This PR adds support for the `anoma-decode` builtin

```
builtin anoma-decode
axiom anomaDecode : {A : Type} -> Nat -> A
```

Adds:
* An implementation of the `cue` function in Haskell
* Unit tests for `cue`
* A benchmark for `cue` applied to the Anoma / nockma stdlib

Benchmark results:

```
      cue (jam stdlib): OK
        36.0 ms ± 2.0 ms
```

Closes:
*  #2764

bench2/Benchmark/Nockma/Encoding.hs

@@ -6,11 +6,16 @@ import Juvix.Compiler.Nockma.Stdlib (stdlib)
 import Juvix.Prelude.Base
 import Test.Tasty.Bench
 
+jamStdlib :: Natural
+jamStdlib = runJam stdlib
+
 bm :: Benchmark
 bm =
   bgroup
     "Jam"
-    [bench "jam stdlib" $ nf runJam stdlib]
+    [ bench "jam stdlib" $ nf runJam stdlib,
+      bench "cue (jam stdlib)" $ nf runCue jamStdlib
+    ]
 
 runJam :: Term Natural -> Natural
 runJam =
@@ -19,3 +24,18 @@ runJam =
     . run
     . runError @NockNaturalNaturalError
     . jam
+
+mkAtom :: Natural -> Atom Natural
+mkAtom n =
+  Atom
+    { _atomInfo = emptyAtomInfo,
+      _atom = n
+    }
+
+runCue :: Natural -> Term Natural
+runCue =
+  run
+    . runErrorNoCallStackWith @NockNaturalNaturalError (const (error "NockNaturalNaturalError"))
+    . runErrorNoCallStackWith @DecodingError (const (error "decoding failed"))
+    . cue
+    . mkAtom

package.yaml

@@ -55,6 +55,7 @@ dependencies:
   - cereal == 0.5.*
   - containers == 0.6.*
   - cryptohash-sha256 == 0.11.*
+  - deepseq == 1.5.*
   - directory == 1.3.*
   - dlist == 1.0.*
   - edit-distance == 0.2.*

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

@@ -235,6 +235,7 @@ genCode fi =
       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/Builtins/Anoma.hs

@@ -30,3 +30,16 @@ registerAnomaEncode f = do
     ((ftype ==% (u <>--> encodeT --> nat)) freeVars)
     (error "anomaEncode must be of type {A : Type} -> A -> Nat")
   registerBuiltin BuiltinAnomaEncode (f ^. axiomName)
+
+registerAnomaDecode :: (Members '[Builtins, NameIdGen] r) => AxiomDef -> Sem r ()
+registerAnomaDecode f = do
+  let ftype = f ^. axiomType
+      u = ExpressionUniverse smallUniverseNoLoc
+      l = getLoc f
+  decodeT <- freshVar l "decodeT"
+  nat <- getBuiltinName (getLoc f) BuiltinNat
+  let freeVars = HashSet.fromList [decodeT]
+  unless
+    ((ftype ==% (u <>--> nat --> decodeT)) freeVars)
+    (error "anomaEncode must be of type {A : Type} -> Nat -> A")
+  registerBuiltin BuiltinAnomaDecode (f ^. axiomName)

src/Juvix/Compiler/Concrete/Data/Builtins.hs

@@ -189,6 +189,7 @@ data BuiltinAxiom
   | BuiltinIntPrint
   | BuiltinAnomaGet
   | BuiltinAnomaEncode
+  | BuiltinAnomaDecode
   | BuiltinPoseidon
   | BuiltinEcOp
   | BuiltinRandomEcPoint
@@ -225,6 +226,7 @@ instance Pretty BuiltinAxiom where
     BuiltinIntPrint -> Str.intPrint
     BuiltinAnomaGet -> Str.anomaGet
     BuiltinAnomaEncode -> Str.anomaEncode
+    BuiltinAnomaDecode -> Str.anomaDecode
     BuiltinPoseidon -> Str.cairoPoseidon
     BuiltinEcOp -> Str.cairoEcOp
     BuiltinRandomEcPoint -> Str.cairoRandomEcPoint

src/Juvix/Compiler/Core/Evaluator.hs

@@ -192,6 +192,7 @@ geval opts herr ctx env0 = eval' env0
       OpTrace -> traceOp
       OpAnomaGet -> anomaGetOp
       OpAnomaEncode -> anomaEncodeOp
+      OpAnomaDecode -> anomaDecodeOp
       OpPoseidonHash -> poseidonHashOp
       OpEc -> ecOp
       OpRandomEcPoint -> randomEcPointOp
@@ -344,9 +345,18 @@ geval opts herr ctx env0 = eval' env0
               | opts ^. evalOptionsNormalize || opts ^. evalOptionsNoFailure ->
                   mkBuiltinApp' OpAnomaEncode [eval' env arg]
               | otherwise ->
-                  err "unsupported builtin operation: OpAnomaGet"
+                  err "unsupported builtin operation: OpAnomaEncode"
         {-# INLINE anomaEncodeOp #-}
 
+        anomaDecodeOp :: [Node] -> Node
+        anomaDecodeOp = unary $ \arg ->
+          if
+              | opts ^. evalOptionsNormalize || opts ^. evalOptionsNoFailure ->
+                  mkBuiltinApp' OpAnomaDecode [eval' env arg]
+              | otherwise ->
+                  err "unsupported builtin operation: OpAnomaDecode"
+        {-# INLINE anomaDecodeOp #-}
+
         poseidonHashOp :: [Node] -> Node
         poseidonHashOp = unary $ \arg ->
           if

src/Juvix/Compiler/Core/Extra/Utils.hs

@@ -425,6 +425,7 @@ builtinOpArgTypes = \case
   OpFail -> [mkTypeString']
   OpAnomaGet -> [mkDynamic']
   OpAnomaEncode -> [mkDynamic']
+  OpAnomaDecode -> [mkDynamic']
   OpPoseidonHash -> [mkDynamic']
   OpEc -> [mkDynamic', mkTypeField', mkDynamic']
   OpRandomEcPoint -> []

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

@@ -28,6 +28,7 @@ data BuiltinOp
   | OpFail
   | OpAnomaGet
   | OpAnomaEncode
+  | OpAnomaDecode
   | OpPoseidonHash
   | OpEc
   | OpRandomEcPoint
@@ -73,6 +74,7 @@ builtinOpArgsNum = \case
   OpFail -> 1
   OpAnomaGet -> 1
   OpAnomaEncode -> 1
+  OpAnomaDecode -> 1
   OpPoseidonHash -> 1
   OpEc -> 3
   OpRandomEcPoint -> 0
@@ -111,6 +113,7 @@ builtinIsFoldable = \case
   OpFail -> False
   OpAnomaGet -> False
   OpAnomaEncode -> False
+  OpAnomaDecode -> False
   OpPoseidonHash -> False
   OpEc -> False
   OpRandomEcPoint -> False
@@ -125,4 +128,4 @@ builtinsCairo :: [BuiltinOp]
 builtinsCairo = [OpPoseidonHash, OpEc, OpRandomEcPoint]
 
 builtinsAnoma :: [BuiltinOp]
-builtinsAnoma = [OpAnomaGet, OpAnomaEncode]
+builtinsAnoma = [OpAnomaGet, OpAnomaEncode, OpAnomaDecode]

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

@@ -54,6 +54,7 @@ instance PrettyCode BuiltinOp where
     OpFail -> return primFail
     OpAnomaGet -> return primAnomaGet
     OpAnomaEncode -> return primAnomaEncode
+    OpAnomaDecode -> return primAnomaDecode
     OpPoseidonHash -> return primPoseidonHash
     OpEc -> return primEc
     OpRandomEcPoint -> return primRandomEcPoint
@@ -805,6 +806,9 @@ primAnomaGet = primitive Str.anomaGet
 primAnomaEncode :: Doc Ann
 primAnomaEncode = primitive Str.anomaEncode
 
+primAnomaDecode :: Doc Ann
+primAnomaDecode = primitive Str.anomaDecode
+
 primPoseidonHash :: Doc Ann
 primPoseidonHash = primitive Str.cairoPoseidon
 

src/Juvix/Compiler/Core/Transformation/ComputeTypeInfo.hs

@@ -68,6 +68,7 @@ computeNodeTypeInfo md = umapL go
           OpFail -> Info.getNodeType node
           OpAnomaGet -> Info.getNodeType node
           OpAnomaEncode -> Info.getNodeType node
+          OpAnomaDecode -> Info.getNodeType node
           OpPoseidonHash -> case _builtinAppArgs of
             [arg] -> Info.getNodeType arg
             _ -> error "incorrect poseidon builtin application"

src/Juvix/Compiler/Core/Translation/FromInternal.hs

@@ -578,6 +578,7 @@ goAxiomInductive a = whenJust (a ^. Internal.axiomBuiltin) builtinInductive
       Internal.BuiltinFieldToNat -> return ()
       Internal.BuiltinAnomaGet -> return ()
       Internal.BuiltinAnomaEncode -> return ()
+      Internal.BuiltinAnomaDecode -> return ()
       Internal.BuiltinPoseidon -> return ()
       Internal.BuiltinEcOp -> return ()
       Internal.BuiltinRandomEcPoint -> return ()
@@ -704,6 +705,14 @@ goAxiomDef a = maybe goAxiomNotBuiltin builtinBody (a ^. Internal.axiomBuiltin)
               mkSmallUniv
               (mkLambda' (mkVar' 0) (mkBuiltinApp' OpAnomaEncode [mkVar' 0]))
           )
+      Internal.BuiltinAnomaDecode -> do
+        natName <- getNatName
+        natSym <- getNatSymbol
+        registerAxiomDef
+          ( mkLambda'
+              (mkTypeConstr (setInfoName natName mempty) natSym [])
+              (mkLambda' (mkVar' 0) (mkBuiltinApp' OpAnomaDecode [mkVar' 0]))
+          )
       Internal.BuiltinPoseidon -> do
         psName <- getPoseidonStateName
         psSym <- getPoseidonStateSymbol
@@ -1103,6 +1112,7 @@ goApplication a = do
         Just Internal.BuiltinFieldToNat -> app
         Just Internal.BuiltinAnomaGet -> app
         Just Internal.BuiltinAnomaEncode -> app
+        Just Internal.BuiltinAnomaDecode -> app
         Just Internal.BuiltinPoseidon -> app
         Just Internal.BuiltinEcOp -> app
         Just Internal.BuiltinRandomEcPoint -> app

src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs

@@ -97,6 +97,7 @@ fromCore fsize tab =
         BuiltinIntPrint -> False
         BuiltinAnomaGet -> False
         BuiltinAnomaEncode -> False
+        BuiltinAnomaDecode -> False
         BuiltinPoseidon -> False
         BuiltinEcOp -> False
         BuiltinRandomEcPoint -> False

src/Juvix/Compiler/Internal/Translation/FromConcrete.hs

@@ -570,6 +570,7 @@ registerBuiltinAxiom d = \case
   BuiltinIntPrint -> registerIntPrint d
   BuiltinAnomaGet -> registerAnomaGet d
   BuiltinAnomaEncode -> registerAnomaEncode d
+  BuiltinAnomaDecode -> registerAnomaDecode d
   BuiltinPoseidon -> registerPoseidon d
   BuiltinEcOp -> registerEcOp d
   BuiltinRandomEcPoint -> registerRandomEcPoint d

src/Juvix/Compiler/Nockma/Encoding.hs

@@ -1,6 +1,8 @@
 module Juvix.Compiler.Nockma.Encoding
   ( module Juvix.Compiler.Nockma.Encoding.Jam,
+    module Juvix.Compiler.Nockma.Encoding.Cue,
   )
 where
 
+import Juvix.Compiler.Nockma.Encoding.Cue
 import Juvix.Compiler.Nockma.Encoding.Jam

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

@@ -24,13 +24,11 @@ integerToVectorBits = run . execBitWriter . writeIntegral
 -- | Computes the number of bits required to store the argument in binary
 -- NB: 0 is encoded to the empty bit vector (as specified by the Hoon serialization spec), so 0 has bit length 0.
 bitLength :: forall a. (Integral a) => a -> Int
-bitLength = \case
-  0 -> 0
-  n -> go (fromIntegral n) 0
-    where
-      go :: Integer -> Int -> Int
-      go 0 acc = acc
-      go x acc = go (x `shiftR` 1) (acc + 1)
+bitLength =
+  length
+    . takeWhile (/= 0)
+    . iterate (`shiftR` 1)
+    . toInteger
 
 -- | Decode a vector of bits (ordered from least to most significant bits) to an integer
 vectorBitsToInteger :: Bit.Vector Bit -> Integer