Capture scry id in the main function and use it when compiling anomaGet (#2705)

The purpose of this PR is to wrap the compiled main function with Nockma
code that captures the argument tuple for use when compiling `anomaGet`
calls.


* The [Anoma system
expects](c7f2d69d1e/lib/anoma/node/executor/worker.ex (L20))
to receive a function of type `ScryId -> Transaction`
 
* The ScryId is only used to construct the argument to the Scry
operation (i.e the anomaGet builtin in the Juvix frontend),
 
* When the Juvix developer writes a function to submit to Anoma they use
type `() -> Transaction`, the main function wrapper is used to capture
the ScryId argument into the subject which is then used to construct
OpScry arguments when anomaGet is compiled.
 
* If the Anoma system expectation changes then the wrapper code must be
changed.

We could add a transformation that checks that the main function in the
Anoma target has no arguments. However it is convenient to be able to
write functions with arguments for testing and debugging (for example
compiling directly to a logic function).

---------

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

src/Juvix/Compiler/Nockma/Evaluator.hs

@@ -372,4 +372,4 @@ evalProfile inistack initerm =
               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))
+              fromMaybeM (throwKeyNotInStorage key) (HashMap.lookup (StorageKey key) <$> asks (^. storageKeyValueData))

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

@@ -146,7 +146,7 @@ instance (PrettyCode a, NockNatural a) => PrettyCode (KeyNotInStorage a) where
     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 :: (StorageKey a, Term a) -> Sem r (Doc Ann)
       combineKeyValue (t1, t2) = do
         pt1 <- ppCode t1
         pt2 <- ppCode t2

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

@@ -1,12 +1,55 @@
 module Juvix.Compiler.Nockma.Evaluator.Storage where
 
 import Juvix.Compiler.Nockma.Language
+import Juvix.Compiler.Nockma.Pretty.Base
 import Juvix.Prelude.Base
 
 newtype Storage a = Storage
-  {_storageKeyValueData :: HashMap (Term a) (Term a)}
+  {_storageKeyValueData :: HashMap (StorageKey a) (Term a)}
 
-emptyStorage :: (Hashable a) => Storage a
-emptyStorage = Storage {_storageKeyValueData = mempty}
+newtype StorageKey a = StorageKey {_storageKeyTerm :: Term a}
 
 makeLenses ''Storage
+makeLenses ''StorageKey
+
+stripMeta :: Term a -> Term a
+stripMeta = \case
+  TermAtom a ->
+    TermAtom
+      Atom
+        { _atom = a ^. atom,
+          _atomInfo = emptyAtomInfo
+        }
+  TermCell c ->
+    TermCell
+      Cell'
+        { _cellLeft = stripMeta (c ^. cellLeft),
+          _cellRight = stripMeta (c ^. cellRight),
+          _cellInfo = emptyCellInfo
+        }
+
+instance (NockmaEq a) => NockmaEq (StorageKey a) where
+  nockmaEq (StorageKey s1) (StorageKey s2) = nockmaEq s1 s2
+
+instance (NockmaEq a) => Eq (StorageKey a) where
+  (==) = nockmaEq
+
+instance (NockmaEq a, Hashable a) => Hashable (StorageKey a) where
+  hashWithSalt salt k = goTerm (k ^. storageKeyTerm)
+    where
+      goCell :: Cell a -> Int
+      goCell c = hashWithSalt salt (c ^. cellLeft, c ^. cellRight)
+
+      goAtom :: Atom a -> Int
+      goAtom a = hashWithSalt salt (a ^. atom)
+
+      goTerm :: Term a -> Int
+      goTerm = \case
+        TermAtom a -> goAtom a
+        TermCell c -> goCell c
+
+instance (PrettyCode a, NockNatural a) => PrettyCode (StorageKey a) where
+  ppCode k = ppCode (stripMeta (k ^. storageKeyTerm))
+
+emptyStorage :: (NockmaEq a, Hashable a) => Storage a
+emptyStorage = Storage {_storageKeyValueData = mempty}

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

@@ -133,6 +133,7 @@ data AnomaCallablePathId
   | ClosureTotalArgsNum
   | ClosureArgsNum
   | ClosureArgs
+  | AnomaGetOrder
   deriving stock (Enum, Bounded, Eq, Show)
 
 -- | A closure has the following structure:
@@ -301,6 +302,24 @@ anomaCallableClosureWrapper =
       adjustArgs = OpIf # closureArgsIsEmpty # (opAddress "wrapperSubject" emptyPath) # appendAndReplaceArgsTuple
    in opCall "closureWrapper" (closurePath RawCode) adjustArgs
 
+mainFunctionWrapper :: Term Natural
+mainFunctionWrapper =
+  -- 1. The Anoma system expects to receive a function of type `ScryId -> Transaction`
+  --
+  -- 2. The ScryId is only used to construct the argument to the Scry operation (i.e the anomaGet builtin in the Juvix frontend),
+  --
+  -- 3. When the Juvix developer writes a function to submit to Anoma they use
+  -- type `() -> Transaction`, this wrapper is used to capture the ScryId
+  -- argument into the subject which is then used to construct OpScry arguments
+  -- when anomaGet is compiled.
+  --
+  -- 4. If the Anoma system expectation changes then this code must be changed.
+  let captureAnomaGetOrder :: Term Natural
+      captureAnomaGetOrder = replaceSubject $ \case
+        AnomaGetOrder -> Just (getClosureFieldInSubject ArgsTuple)
+        _ -> Nothing
+   in opCall "mainFunctionWrapper" (closurePath RawCode) captureAnomaGetOrder
+
 compile :: forall r. (Members '[Reader FunctionCtx, Reader CompilerCtx] r) => Tree.Node -> Sem r (Term Natural)
 compile = \case
   Tree.Binop b -> goBinop b
@@ -413,7 +432,9 @@ compile = \case
       Tree.OpFieldToInt -> fieldErr
 
     goAnomaGet :: Term Natural -> Sem r (Term Natural)
-    goAnomaGet arg = return (OpScry # (OpQuote # nockNilTagged "OpScry-typehint") # arg)
+    goAnomaGet key = do
+      let arg = remakeList [getFieldInSubject AnomaGetOrder, key]
+      return (OpScry # (OpQuote # nockNilTagged "OpScry-typehint") # arg)
 
     goTrace :: Term Natural -> Sem r (Term Natural)
     goTrace arg = do
@@ -463,6 +484,7 @@ compile = \case
         ClosureTotalArgsNum -> nockNatLiteral farity
         ClosureArgsNum -> nockIntegralLiteral (length args)
         ClosureArgs -> remakeList args
+        AnomaGetOrder -> OpQuote # nockNilTagged "goAllocClosure-AnomaGetOrder"
 
     goExtendClosure :: Tree.NodeExtendClosure -> Sem r (Term Natural)
     goExtendClosure = extendClosure
@@ -487,6 +509,7 @@ compile = \case
                 ClosureArgs -> Nothing
                 ClosureTotalArgsNum -> Nothing
                 ClosureArgsNum -> Nothing
+                AnomaGetOrder -> Nothing
           return $ (opCall "callClosure" (closurePath WrapperCode) newSubject)
 
 isZero :: Term Natural -> Term Natural
@@ -575,6 +598,7 @@ extendClosure Tree.NodeExtendClosure {..} = do
     FunctionsLibrary -> getClosureField FunctionsLibrary closure
     TempStack -> getClosureField TempStack closure
     StandardLibrary -> getClosureField StandardLibrary closure
+    AnomaGetOrder -> getClosureField AnomaGetOrder closure
 
 -- Calling convention for Anoma stdlib
 --
@@ -673,10 +697,13 @@ remakeList :: (Foldable l) => l (Term Natural) -> Term Natural
 remakeList ts = foldTerms (toList ts `prependList` pure (OpQuote # nockNilTagged "remakeList"))
 
 runCompilerWith :: CompilerOptions -> ConstructorInfos -> [CompilerFunction] -> CompilerFunction -> AnomaResult
-runCompilerWith opts constrs libFuns mainFun = makeAnomaFun
+runCompilerWith opts constrs moduleFuns mainFun = makeAnomaFun
   where
+    libFuns :: [CompilerFunction]
+    libFuns = moduleFuns ++ (builtinFunction <$> allElements)
+
     allFuns :: NonEmpty CompilerFunction
-    allFuns = mainFun :| libFuns ++ (builtinFunction <$> allElements)
+    allFuns = mainFun :| libFuns
 
     compilerCtx :: CompilerCtx
     compilerCtx =
@@ -686,11 +713,16 @@ runCompilerWith opts constrs libFuns mainFun = makeAnomaFun
           _compilerOptions = opts
         }
 
+    mainClosure :: Term Natural
+    mainClosure = makeMainFunction (runCompilerFunction compilerCtx mainFun)
+
     compiledFuns :: NonEmpty (Term Natural)
     compiledFuns =
-      makeLibraryFunction
-        <$> ( runCompilerFunction compilerCtx <$> allFuns
-            )
+      mainClosure
+        :| ( makeLibraryFunction
+               <$> ( runCompilerFunction compilerCtx <$> libFuns
+                   )
+           )
 
     exportEnv :: Term Natural
     exportEnv = makeList compiledFuns
@@ -708,6 +740,22 @@ runCompilerWith opts constrs libFuns mainFun = makeAnomaFun
             ClosureTotalArgsNum -> nockNilHere
             ClosureArgsNum -> nockNilHere
             ClosureArgs -> nockNilHere
+            AnomaGetOrder -> nockNilHere
+
+    makeMainFunction :: Term Natural -> Term Natural
+    makeMainFunction c = makeClosure $ \p ->
+      let nockNilHere = nockNilTagged ("makeMainFunction-" <> show p)
+       in case p of
+            WrapperCode -> mainFunctionWrapper
+            ArgsTuple -> argsTuplePlaceholder "mainFunction"
+            FunctionsLibrary -> functionsLibraryPlaceHolder
+            RawCode -> c
+            TempStack -> nockNilHere
+            StandardLibrary -> stdlib
+            ClosureTotalArgsNum -> nockNilHere
+            ClosureArgsNum -> nockNilHere
+            ClosureArgs -> nockNilHere
+            AnomaGetOrder -> nockNilHere
 
     functionInfos :: HashMap FunctionId FunctionInfo
     functionInfos = hashMap (run (runInputNaturals (toList <$> userFunctions)))
@@ -725,11 +773,9 @@ runCompilerWith opts constrs libFuns mainFun = makeAnomaFun
 
     makeAnomaFun :: AnomaResult
     makeAnomaFun =
-      let mainClosure :: Term Natural
-          mainClosure = head compiledFuns
-       in AnomaResult
-            { _anomaClosure = substEnv mainClosure
-            }
+      AnomaResult
+        { _anomaClosure = substEnv mainClosure
+        }
       where
         -- Replaces all instances of functionsLibraryPlaceHolder by the actual
         -- functions library. Note that the functions library will have

test/Anoma/Compilation/Positive.hs

@@ -526,12 +526,16 @@ allTests =
           v1 :: Term Natural = [nock| 222 |]
           k2 :: Term Natural = [nock| [1 2 3 nil] |]
           v2 :: Term Natural = [nock| [4 5 6 nil] |]
+          -- The keys of the storage are of the form [id key nil].
+          -- The id is captured from the arguments tuple of the function.
+          sk1 :: Term Natural = [nock| [[333 1 2 3 nil] 333 nil] |]
+          sk2 :: Term Natural = [nock| [[333 1 2 3 nil] [1 2 3 nil] nil] |]
        in mkAnomaCallTest'
             True
             ( Storage
                 ( HashMap.fromList
-                    [ (k1, v1),
-                      (k2, v2)
+                    [ (StorageKey sk1, v1),
+                      (StorageKey sk2, v2)
                     ]
                 )
             )

test/Nockma/Eval/Positive.hs

@@ -123,7 +123,7 @@ anomaTest n mainFun args _testCheck _evalInterceptStdlibCalls =
    in Test {..}
 
 testWithStorage :: [(Term Natural, Term Natural)] -> Text -> Term Natural -> Term Natural -> Check () -> Test
-testWithStorage s = Test defaultEvalOptions Nothing (Storage (HashMap.fromList s))
+testWithStorage s = Test defaultEvalOptions Nothing (Storage (HashMap.fromList (first StorageKey <$> s)))
 
 test :: Text -> Term Natural -> Term Natural -> Check () -> Test
 test = testWithStorage []