Use StorageKey newtype in Storage

In Storage we want key terms to be compared as if they had no tags. We
achieve this by introucing a newtype for keys and stripping out the tags
as part of the Hashable/Eq instance definitions

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,37 @@
 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
+
+stripTags :: Term a -> Term a
+stripTags = \case
+  TermAtom a -> TermAtom (set atomTag Nothing a)
+  TermCell c ->
+    TermCell
+      ( set cellTag Nothing
+          . over cellLeft stripTags
+          . over cellRight stripTags
+          . over (cellCall . _Just . stdlibCallArgs) stripTags
+          $ c
+      )
+
+instance (Eq a) => Eq (StorageKey a) where
+  s1 == s2 = stripTags (s1 ^. storageKeyTerm) == stripTags (s2 ^. storageKeyTerm)
+
+instance (Hashable a) => Hashable (StorageKey a) where
+  hashWithSalt s k = hashWithSalt s (stripTags (k ^. storageKeyTerm))
+
+instance (PrettyCode a, NockNatural a) => PrettyCode (StorageKey a) where
+  ppCode k = ppCode (stripTags (k ^. storageKeyTerm))
+
+emptyStorage :: (Hashable a) => Storage a
+emptyStorage = Storage {_storageKeyValueData = mempty}

test/Anoma/Compilation/Positive.hs

@@ -534,8 +534,8 @@ allTests =
             True
             ( Storage
                 ( HashMap.fromList
-                    [ (sk1, v1),
-                      (sk2, 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 []