Support nockma scry (#2678)

This PR adds support for Anoma/Nockma scry OP. It is used for obtaining
values from the Anoma storage (key value store). See the [linked
issue](#2672) for details on scry.

This PR adds support for scry to the Nockma language and compilation
from the frontend via a builtin: `anoma-get`:

```
builtin anoma-get
axiom anomaGet : {Value Key : Type} -> Key -> Value
```

In the backend, the `Value` and `Key` types could be anything, they will
depend on choices of Anoma applications. The type of the returned
`Value` is unchecked. It's currently the responsibility of the user to
match the annotated type with the type of data in storage.

We will not put this builtin in the standard library. It will be exposed
in the anoma-juvix library. It's likely that the frontend `anomaGet`
function will evolve as we use it to write Anoma applications and learn
how they work.

* Closes #2672

---------

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

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

@@ -227,6 +227,7 @@ genCode fi =
       Tree.PrimUnop op' -> mkUnop op'
       Tree.OpTrace -> mkInstr Trace
       Tree.OpFail -> mkInstr Failure
+      Tree.OpAnomaGet -> impossible
 
     snocReturn :: Bool -> Code' -> Code'
     snocReturn True code = DL.snoc code (mkInstr Return)

src/Juvix/Compiler/Builtins.hs

@@ -9,9 +9,11 @@ module Juvix.Compiler.Builtins
     module Juvix.Compiler.Builtins.Field,
     module Juvix.Compiler.Builtins.Debug,
     module Juvix.Compiler.Builtins.Control,
+    module Juvix.Compiler.Builtins.Anoma,
   )
 where
 
+import Juvix.Compiler.Builtins.Anoma
 import Juvix.Compiler.Builtins.Bool
 import Juvix.Compiler.Builtins.Control
 import Juvix.Compiler.Builtins.Debug

src/Juvix/Compiler/Builtins/Anoma.hs

@@ -0,0 +1,19 @@
+module Juvix.Compiler.Builtins.Anoma where
+
+import Data.HashSet qualified as HashSet
+import Juvix.Compiler.Builtins.Effect
+import Juvix.Compiler.Internal.Extra
+import Juvix.Prelude
+
+registerAnomaGet :: (Members '[Builtins, NameIdGen] r) => AxiomDef -> Sem r ()
+registerAnomaGet f = do
+  let ftype = f ^. axiomType
+      u = ExpressionUniverse smallUniverseNoLoc
+      l = getLoc f
+  keyT <- freshVar l "keyT"
+  valueT <- freshVar l "valueT"
+  let freeVars = HashSet.fromList [keyT, valueT]
+  unless
+    ((ftype ==% (u <>--> u <>--> keyT --> valueT)) freeVars)
+    (error "anomaGet must be of type {Value Key : Type} -> Key -> Value")
+  registerBuiltin BuiltinAnomaGet (f ^. axiomName)

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

@@ -177,6 +177,7 @@ data BuiltinAxiom
   | BuiltinFail
   | BuiltinIntToString
   | BuiltinIntPrint
+  | BuiltinAnomaGet
   deriving stock (Show, Eq, Ord, Enum, Bounded, Generic, Data)
 
 instance Hashable BuiltinAxiom
@@ -208,6 +209,7 @@ instance Pretty BuiltinAxiom where
     BuiltinFail -> Str.fail_
     BuiltinIntToString -> Str.intToString
     BuiltinIntPrint -> Str.intPrint
+    BuiltinAnomaGet -> Str.anomaGet
 
 data BuiltinType
   = BuiltinTypeInductive BuiltinInductive

src/Juvix/Compiler/Core/Evaluator.hs

@@ -187,6 +187,7 @@ geval opts herr ctx env0 = eval' env0
       OpSeq -> seqOp
       OpFail -> failOp
       OpTrace -> traceOp
+      OpAnomaGet -> err "unsupported op: OpAnomaGet"
       where
         err :: Text -> a
         err msg = evalError msg n

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

@@ -423,6 +423,7 @@ builtinOpArgTypes = \case
   OpSeq -> [mkDynamic', mkDynamic']
   OpTrace -> [mkDynamic']
   OpFail -> [mkTypeString']
+  OpAnomaGet -> [mkDynamic']
 
 translateCase :: (Node -> Node -> Node -> a) -> a -> Case -> a
 translateCase translateIfFun dflt Case {..} = case _caseBranches of

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

@@ -26,6 +26,7 @@ data BuiltinOp
   | OpSeq
   | OpTrace
   | OpFail
+  | OpAnomaGet
   deriving stock (Eq, Generic)
 
 instance Serialize BuiltinOp
@@ -66,6 +67,7 @@ builtinOpArgsNum = \case
   OpSeq -> 2
   OpTrace -> 1
   OpFail -> 1
+  OpAnomaGet -> 1
 
 builtinConstrArgsNum :: BuiltinDataTag -> Int
 builtinConstrArgsNum = \case

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

@@ -52,6 +52,7 @@ instance PrettyCode BuiltinOp where
     OpSeq -> return primSeq
     OpTrace -> return primTrace
     OpFail -> return primFail
+    OpAnomaGet -> return primAnomaGet
 
 instance PrettyCode BuiltinDataTag where
   ppCode = \case
@@ -794,6 +795,9 @@ primSeq = primitive Str.seqq_
 primFail :: Doc Ann
 primFail = primitive Str.fail_
 
+primAnomaGet :: Doc Ann
+primAnomaGet = primitive Str.anomaGet
+
 primTrace :: Doc Ann
 primTrace = primitive Str.trace_
 

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

@@ -66,6 +66,7 @@ computeNodeTypeInfo md = umapL go
             [arg] -> Info.getNodeType arg
             _ -> error "incorrect trace builtin application"
           OpFail -> Info.getNodeType node
+          OpAnomaGet -> Info.getNodeType node
       NCtr Constr {..} ->
         let ci = lookupConstructorInfo md _constrTag
             ii = lookupInductiveInfo md (ci ^. constructorInductive)

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

@@ -577,6 +577,7 @@ goAxiomInductive a = whenJust (a ^. Internal.axiomBuiltin) builtinInductive
       Internal.BuiltinFieldDiv -> return ()
       Internal.BuiltinFieldFromInt -> return ()
       Internal.BuiltinFieldToNat -> return ()
+      Internal.BuiltinAnomaGet -> return ()
 
     registerInductiveAxiom :: Maybe BuiltinAxiom -> [(Tag, Text, Type -> Type, Maybe BuiltinConstructor)] -> Sem r ()
     registerInductiveAxiom ax ctrs = do
@@ -685,6 +686,15 @@ goAxiomDef a = maybe goAxiomNotBuiltin builtinBody (a ^. Internal.axiomBuiltin)
         intName <- getIntName
         intSym <- getIntSymbol
         registerAxiomDef $ writeLambda (mkTypeConstr (setInfoName intName mempty) intSym [])
+      Internal.BuiltinAnomaGet ->
+        registerAxiomDef
+          ( mkLambda'
+              mkSmallUniv
+              ( mkLambda'
+                  mkSmallUniv
+                  (mkLambda' (mkVar' 0) (mkBuiltinApp' OpAnomaGet [mkVar' 0]))
+              )
+          )
 
     axiomType' :: Sem r Type
     axiomType' = fromTopIndex (goType (a ^. Internal.axiomType))
@@ -1060,6 +1070,7 @@ goApplication a = do
             [x] -> return $ mkBuiltinApp' OpFieldFromInt [x]
             _ -> app
         Just Internal.BuiltinFieldToNat -> app
+        Just Internal.BuiltinAnomaGet -> app
         Nothing -> app
     Internal.ExpressionIden (Internal.IdenFunction n) -> do
       funInfoBuiltin <- Internal.getFunctionBuiltinInfo n

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

@@ -93,6 +93,7 @@ fromCore fsize tab =
         BuiltinFail -> False
         BuiltinIntToString -> False
         BuiltinIntPrint -> False
+        BuiltinAnomaGet -> False
       BuiltinTypeInductive i -> case i of
         BuiltinList -> True
         BuiltinNat -> False

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

@@ -587,6 +587,7 @@ registerBuiltinAxiom d = \case
   BuiltinFail -> registerFail d
   BuiltinIntToString -> registerIntToString d
   BuiltinIntPrint -> registerIntPrint d
+  BuiltinAnomaGet -> registerAnomaGet d
 
 goInductive ::
   (Members '[Reader EntryPoint, Reader DefaultArgsStack, NameIdGen, Reader Pragmas, Builtins, Error ScoperError, State ConstructorInfos, Reader S.InfoTable] r) =>

src/Juvix/Compiler/Nockma/EvalCompiled.hs

@@ -10,6 +10,7 @@ evalCompiledNock' stack mainTerm = do
   evalT <-
     runError @(ErrNockNatural Natural)
       . runError @(NockEvalError Natural)
+      . runReader @(Storage Natural) emptyStorage
       $ eval stack mainTerm
   case evalT of
     Left e -> error (show e)

src/Juvix/Compiler/Nockma/Evaluator.hs

@@ -2,11 +2,14 @@ module Juvix.Compiler.Nockma.Evaluator
   ( module Juvix.Compiler.Nockma.Evaluator,
     module Juvix.Compiler.Nockma.Evaluator.Error,
     module Juvix.Compiler.Nockma.Evaluator.Options,
+    module Juvix.Compiler.Nockma.Evaluator.Storage,
   )
 where
 
+import Data.HashMap.Strict qualified as HashMap
 import Juvix.Compiler.Nockma.Evaluator.Error
 import Juvix.Compiler.Nockma.Evaluator.Options
+import Juvix.Compiler.Nockma.Evaluator.Storage
 import Juvix.Compiler.Nockma.Language
 import Juvix.Prelude hiding (Atom, Path)
 
@@ -102,7 +105,7 @@ programAssignments mprog =
 -- | The stack provided in the replExpression has priority
 evalRepl ::
   forall r a.
-  (Integral a, Members '[Reader EvalOptions, Error (NockEvalError a), Error (ErrNockNatural a)] r, NockNatural a) =>
+  (Hashable a, Integral a, Members '[Reader EvalOptions, Error (NockEvalError a), Error (ErrNockNatural a)] r, NockNatural a) =>
   (Term a -> Sem r ()) ->
   Maybe (Program a) ->
   Maybe (Term a) ->
@@ -115,7 +118,7 @@ evalRepl handleTrace mprog defaultStack expr = do
       t' <- fromReplTerm namedTerms (w ^. withStackStack)
       return (Just t', w ^. withStackTerm)
   stack <- maybe errNoStack return mstack
-  fromReplTerm namedTerms t >>= runOutputSem @(Term a) handleTrace . eval stack
+  fromReplTerm namedTerms t >>= runOutputSem @(Term a) handleTrace . runReader @(Storage a) emptyStorage . eval stack
   where
     errNoStack :: Sem r x
     errNoStack = throw @(NockEvalError a) (ErrNoStack NoStack)
@@ -125,7 +128,7 @@ evalRepl handleTrace mprog defaultStack expr = do
 
 eval ::
   forall s a.
-  (Integral a, Members '[Reader EvalOptions, Output (Term a), Error (NockEvalError a), Error (ErrNockNatural a)] s, NockNatural a) =>
+  (Hashable a, Integral a, Members '[Reader (Storage a), Reader EvalOptions, Output (Term a), Error (NockEvalError a), Error (ErrNockNatural a)] s, NockNatural a) =>
   Term a ->
   Term a ->
   Sem s (Term a)
@@ -210,6 +213,7 @@ eval inistack initerm =
           OpCall -> goOpCall
           OpReplace -> goOpReplace
           OpHint -> goOpHint
+          OpScry -> goOpScry
           OpTrace -> goOpTrace
           where
             crumb crumbTag =
@@ -317,3 +321,10 @@ eval inistack initerm =
               cellTerm <- withCrumb (crumb crumbDecodeFirst) (asCell (c ^. operatorCellTerm))
               t1' <- evalArg crumbEvalFirst stack (cellTerm ^. cellLeft)
               evalArg crumbEvalSecond t1' (cellTerm ^. cellRight)
+
+            goOpScry :: Sem r (Term a)
+            goOpScry = do
+              Cell' typeFormula subFormula _ <- withCrumb (crumb crumbDecodeFirst) (asCell (c ^. operatorCellTerm))
+              void (evalArg crumbEvalFirst stack typeFormula)
+              key <- evalArg crumbEvalSecond stack subFormula
+              fromMaybeM (throwKeyNotInStorage key) (HashMap.lookup key <$> asks (^. storageKeyValueData))

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

@@ -1,10 +1,13 @@
 module Juvix.Compiler.Nockma.Evaluator.Error
   ( module Juvix.Compiler.Nockma.Evaluator.Error,
     module Juvix.Compiler.Nockma.Evaluator.Crumbs,
+    module Juvix.Compiler.Nockma.Evaluator.Storage,
   )
 where
 
+import Data.HashMap.Strict qualified as HashMap
 import Juvix.Compiler.Nockma.Evaluator.Crumbs
+import Juvix.Compiler.Nockma.Evaluator.Storage
 import Juvix.Compiler.Nockma.Language
 import Juvix.Compiler.Nockma.Pretty.Base
 import Juvix.Prelude hiding (Atom, Path)
@@ -17,6 +20,7 @@ data NockEvalError a
     ErrNoStack NoStack
   | -- TODO perhaps this should be a repl error type
     ErrAssignmentNotFound Text
+  | ErrKeyNotInStorage (KeyNotInStorage a)
 
 newtype GenericNockEvalError = GenericNockEvalError
   { _genericNockEvalErrorMessage :: AnsiText
@@ -41,6 +45,11 @@ data InvalidPath a = InvalidPath
     _invalidPathPath :: Path
   }
 
+data KeyNotInStorage a = KeyNotInStorage
+  { _keyNotInStorageKey :: Term a,
+    _keyNotInStorageStorage :: Storage a
+  }
+
 data NoStack = NoStack
 
 throwInvalidPath :: (Members '[Error (NockEvalError a), Reader EvalCtx] r) => Term a -> Path -> Sem r x
@@ -74,6 +83,16 @@ throwExpectedAtom a = do
           _expectedAtomCell = a
         }
 
+throwKeyNotInStorage :: (Members '[Reader (Storage a), Error (NockEvalError a)] r) => Term a -> Sem r x
+throwKeyNotInStorage k = do
+  s <- ask
+  throw $
+    ErrKeyNotInStorage
+      KeyNotInStorage
+        { _keyNotInStorageKey = k,
+          _keyNotInStorageStorage = s
+        }
+
 instance PrettyCode NoStack where
   ppCode _ = return "Missing stack"
 
@@ -98,10 +117,24 @@ instance (PrettyCode a, NockNatural a) => PrettyCode (ExpectedCell a) where
     let cell = annotate AnnImportant "cell"
     return (ctx <> "Expected a" <+> cell <+> "but got:" <> line <> atm)
 
+instance (PrettyCode a, NockNatural a) => PrettyCode (KeyNotInStorage a) where
+  ppCode :: forall r. (Member (Reader Options) r) => KeyNotInStorage a -> Sem r (Doc Ann)
+  ppCode KeyNotInStorage {..} = do
+    tm <- ppCode _keyNotInStorageKey
+    hashMapKvs <- vsep <$> mapM combineKeyValue (HashMap.toList (_keyNotInStorageStorage ^. storageKeyValueData))
+    return ("The key" <+> tm <+> "is not found in the storage." <> line <> "Storage contains the following key value pairs:" <> line <> hashMapKvs)
+    where
+      combineKeyValue :: (Term a, Term a) -> Sem r (Doc Ann)
+      combineKeyValue (t1, t2) = do
+        pt1 <- ppCode t1
+        pt2 <- ppCode t2
+        return (pt1 <+> ":=" <+> pt2)
+
 instance (PrettyCode a, NockNatural a) => PrettyCode (NockEvalError a) where
   ppCode = \case
     ErrInvalidPath e -> ppCode e
     ErrExpectedAtom e -> ppCode e
     ErrExpectedCell e -> ppCode e
     ErrNoStack e -> ppCode e
     ErrAssignmentNotFound e -> return (pretty e)
+    ErrKeyNotInStorage e -> ppCode e

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

@@ -0,0 +1,12 @@
+module Juvix.Compiler.Nockma.Evaluator.Storage where
+
+import Juvix.Compiler.Nockma.Language
+import Juvix.Prelude.Base
+
+newtype Storage a = Storage
+  {_storageKeyValueData :: HashMap (Term a) (Term a)}
+
+emptyStorage :: (Hashable a) => Storage a
+emptyStorage = Storage {_storageKeyValueData = mempty}
+
+makeLenses ''Storage