Use Anoma compatible Nockma serialization of Bools and List-like data structures (#2591)

This PR changes the Nockma representation of builtin Bool and list-like
types to make them compatible with Anoma.

True and False are now compiled to the Nockma atoms 0 and 1
respectively.

For inductive types that have exactly two constructors, one of arity
zero and one of arity two, we compile the arity zero constructor to
Nockma zero, and the arity two constructor to a Nockma cell. In
particular a Juvix stdlib List will be compiled to an Anoma/Nockma list.
This is necessary for compatibility with the layout of resource and
resource logic types in Anoma.

In tests we avoid using the StackRef memory reference because it will be
removed as part of the JuvixTree work.

---------

Co-authored-by: Łukasz Czajka <62751+lukaszcz@users.noreply.github.com>

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

@@ -96,7 +96,7 @@ instance (PrettyCode a, NockNatural a) => PrettyCode (ExpectedCell a) where
     atm <- ppCode _expectedCellAtom
     ctx <- ppCtx _expectedCellCtx
     let cell = annotate AnnImportant "cell"
-    return (ctx <> "Expected an" <+> atm <+> "but got:" <> line <> cell)
+    return (ctx <> "Expected a" <+> cell <+> "but got:" <> line <> atm)
 
 instance (PrettyCode a, NockNatural a) => PrettyCode (NockEvalError a) where
   ppCode = \case

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

@@ -17,9 +17,27 @@ nockmaMemRep = \case
   MemRepUnit -> NockmaMemRepConstr
   MemRepUnpacked {} -> NockmaMemRepConstr
 
+data NockmaMemRepListConstr
+  = NockmaMemRepListConstrNil
+  | NockmaMemRepListConstrCons
+  deriving stock (Eq)
+
 data NockmaMemRep
   = NockmaMemRepConstr
   | NockmaMemRepTuple
+  | NockmaMemRepList NockmaMemRepListConstr
+
+newtype NockmaBuiltinTag
+  = NockmaBuiltinBool Bool
+
+nockmaBuiltinTag :: Asm.BuiltinDataTag -> NockmaBuiltinTag
+nockmaBuiltinTag = \case
+  Asm.TagTrue -> NockmaBuiltinBool True
+  Asm.TagFalse -> NockmaBuiltinBool False
+  Asm.TagReturn -> impossible
+  Asm.TagBind -> impossible
+  Asm.TagWrite -> impossible
+  Asm.TagReadLn -> impossible
 
 type UserFunctionId = Symbol
 
@@ -168,11 +186,13 @@ stackPath :: StackId -> Path
 stackPath s = indexStack (fromIntegral (fromEnum s))
 
 indexTuple :: Natural -> Natural -> Path
-indexTuple len idx =
-  let lastL
-        | idx == len - 1 = []
-        | otherwise = [L]
-   in replicate idx R ++ lastL
+indexTuple len idx
+  | idx >= len = impossible
+  | otherwise =
+      let lastL
+            | idx == len - 1 = []
+            | otherwise = [L]
+       in replicate idx R ++ lastL
 
 indexStack :: Natural -> Path
 indexStack idx = replicate idx R ++ [L]
@@ -221,23 +241,50 @@ makeFunction f = f FunctionCode # f FunctionArgs
 foldTerms :: NonEmpty (Term Natural) -> Term Natural
 foldTerms = foldr1 (#)
 
+allConstructors :: Asm.InfoTable -> Asm.ConstructorInfo -> NonEmpty Asm.ConstructorInfo
+allConstructors Asm.InfoTable {..} ci =
+  let indInfo = getInductiveInfo (ci ^. Asm.constructorInductive)
+   in nonEmpty' (getConstructorInfo'' <$> indInfo ^. Asm.inductiveConstructors)
+  where
+    getInductiveInfo :: Symbol -> Asm.InductiveInfo
+    getInductiveInfo s = _infoInductives ^?! at s . _Just
+
+    getConstructorInfo'' :: Asm.Tag -> Asm.ConstructorInfo
+    getConstructorInfo'' t = _infoConstrs ^?! at t . _Just
+
+supportsListNockmaRep :: Asm.InfoTable -> Asm.ConstructorInfo -> Maybe NockmaMemRepListConstr
+supportsListNockmaRep tab ci = case allConstructors tab ci of
+  c1 :| [c2]
+    | [0, 2] `elem` permutations ((^. Asm.constructorArgsNum) <$> [c1, c2]) -> Just $ case ci ^. Asm.constructorArgsNum of
+        0 -> NockmaMemRepListConstrNil
+        2 -> NockmaMemRepListConstrCons
+        _ -> impossible
+    | otherwise -> Nothing
+  _ -> Nothing
+
 -- | Use `Asm.toNockma` before calling this function
 fromAsmTable :: (Members '[Error JuvixError, Reader CompilerOptions] r) => Asm.InfoTable -> Sem r (Cell Natural)
 fromAsmTable t = case t ^. Asm.infoMainFunction of
   Just mainFun -> do
     opts <- ask
     return (fromAsm opts mainFun t)
-  Nothing -> throw @JuvixError (error "TODO")
+  Nothing -> throw @JuvixError (error "TODO missing main")
   where
     fromAsm :: CompilerOptions -> Asm.Symbol -> Asm.InfoTable -> Cell Natural
-    fromAsm opts mainSym Asm.InfoTable {..} =
+    fromAsm opts mainSym tab@Asm.InfoTable {..} =
       let funs = map compileFunction allFunctions
           mkConstructorInfo :: Asm.ConstructorInfo -> ConstructorInfo
           mkConstructorInfo ci@Asm.ConstructorInfo {..} =
             ConstructorInfo
               { _constructorInfoArity = fromIntegral _constructorArgsNum,
-                _constructorInfoMemRep = nockmaMemRep (memRep ci (getInductiveInfo (ci ^. Asm.constructorInductive)))
+                _constructorInfoMemRep = rep
               }
+            where
+              rep :: NockmaMemRep
+              rep = maybe r NockmaMemRepList (supportsListNockmaRep tab ci)
+                where
+                  r = nockmaMemRep (memRep ci (getInductiveInfo (ci ^. Asm.constructorInductive)))
+
           constrs :: ConstructorInfos
           constrs = mkConstructorInfo <$> _infoConstrs
 
@@ -285,15 +332,21 @@ fromOffsetRef = fromIntegral . (^. Asm.offsetRefOffset)
 -- nil terminated list.
 goConstructor :: NockmaMemRep -> Asm.Tag -> [Term Natural] -> Term Natural
 goConstructor mr t args = case t of
-  Asm.BuiltinTag b -> makeConstructor $ \case
-    ConstructorTag -> builtinTagToTerm b
-    ConstructorArgs -> remakeList []
+  Asm.BuiltinTag b -> case nockmaBuiltinTag b of
+    NockmaBuiltinBool v -> nockBoolLiteral v
   Asm.UserTag tag -> case mr of
     NockmaMemRepConstr ->
       makeConstructor $ \case
         ConstructorTag -> OpQuote # (fromIntegral (tag ^. Asm.tagUserWord) :: Natural)
         ConstructorArgs -> remakeList args
     NockmaMemRepTuple -> foldTerms (nonEmpty' args)
+    NockmaMemRepList constr -> case constr of
+      NockmaMemRepListConstrNil
+        | null args -> remakeList []
+        | otherwise -> impossible
+      NockmaMemRepListConstrCons -> case args of
+        [l, r] -> TCell l r
+        _ -> impossible
 
 compile :: forall r. (Members '[Compiler] r) => Asm.Code -> Sem r ()
 compile = mapM_ goCommand
@@ -460,6 +513,9 @@ pushConstructorFieldOnto s tag refToConstr argIx = do
         NockmaMemRepTuple ->
           directRefPath refToConstr
             ++ indexTuple arity argIx
+        NockmaMemRepList constr -> case constr of
+          NockmaMemRepListConstrNil -> impossible
+          NockmaMemRepListConstrCons -> directRefPath refToConstr ++ indexTuple 2 argIx
   pushOnto s (OpAddress # path)
 
 pushConstructorField :: (Members '[Compiler] r) => Asm.Tag -> Asm.DirectRef -> Natural -> Sem r ()
@@ -562,6 +618,12 @@ initStack defs = makeList (initSubStack <$> allElements)
 push :: (Members '[Compiler] r) => Term Natural -> Sem r ()
 push = pushOnto ValueStack
 
+dupOnto :: (Members '[Compiler] r) => StackId -> Sem r ()
+dupOnto stackId = pushOnto stackId (OpAddress # topOfStack stackId)
+
+dup :: (Members '[Compiler] r) => Sem r ()
+dup = dupOnto ValueStack
+
 execCompilerList :: (Member (Reader CompilerCtx) r) => Sem (Compiler ': r) a -> Sem r [Term Natural]
 execCompilerList = fmap fst . runCompilerList
 
@@ -696,6 +758,10 @@ replaceSubterm obj relPath newVal = OpReplace # (relPath # newVal) # obj
 evaluated :: Term Natural -> Term Natural
 evaluated t = OpApply # (OpAddress # emptyPath) # t
 
+-- | The IsCell op but the argument is evaluated first.
+isCell' :: Term Natural -> Term Natural
+isCell' t = evaluated $ (OpQuote # OpIsCell) # t
+
 -- | The same as replaceSubterm but the path is a cell that is evaluated.
 -- i.e. replaceSubterm a p b = replaceSubterm' a (quote p) b
 replaceSubterm' :: Term Natural -> Term Natural -> Term Natural -> Term Natural
@@ -859,19 +925,14 @@ save' isTail m = do
       | isTail -> pureT ()
       | otherwise -> popFromH TempStack
 
-builtinTagToTerm :: Asm.BuiltinDataTag -> Term Natural
+builtinTagToTerm :: NockmaBuiltinTag -> Term Natural
 builtinTagToTerm = \case
-  Asm.TagTrue -> nockBoolLiteral True
-  Asm.TagFalse -> nockBoolLiteral False
-  Asm.TagReturn -> impossible
-  Asm.TagBind -> impossible
-  Asm.TagWrite -> impossible
-  Asm.TagReadLn -> impossible
+  NockmaBuiltinBool v -> nockBoolLiteral v
 
 constructorTagToTerm :: Asm.Tag -> Term Natural
 constructorTagToTerm = \case
   Asm.UserTag t -> OpQuote # toNock (fromIntegral (t ^. Asm.tagUserWord) :: Natural)
-  Asm.BuiltinTag b -> builtinTagToTerm b
+  Asm.BuiltinTag b -> builtinTagToTerm (nockmaBuiltinTag b)
 
 caseCmd ::
   forall r.
@@ -881,13 +942,19 @@ caseCmd ::
   Sem r ()
 caseCmd defaultBranch = \case
   [] -> sequence_ defaultBranch
-  (tag, b) : bs -> do
-    rep <- getConstructorMemRep tag
-    case rep of
-      NockmaMemRepConstr -> goRepConstr tag b bs
-      NockmaMemRepTuple
-        | null bs, isNothing defaultBranch -> b
-        | otherwise -> error "redundant branch. Impossible?"
+  (tag, b) : bs -> case tag of
+    Asm.BuiltinTag t -> case nockmaBuiltinTag t of
+      NockmaBuiltinBool v -> goBoolTag v b bs
+    Asm.UserTag {} -> do
+      rep <- getConstructorMemRep tag
+      case rep of
+        NockmaMemRepConstr -> goRepConstr tag b bs
+        NockmaMemRepTuple
+          | null bs, isNothing defaultBranch -> b
+          | otherwise -> error "redundant branch. Impossible?"
+        NockmaMemRepList constr -> do
+          bs' <- mapM (firstM asNockmaMemRepListConstr) bs
+          goRepList ((constr, b) :| bs')
   where
     goRepConstr :: Asm.Tag -> Sem r () -> [(Asm.Tag, Sem r ())] -> Sem r ()
     goRepConstr tag b bs = do
@@ -897,6 +964,40 @@ caseCmd defaultBranch = \case
       testEq
       branch b (caseCmd defaultBranch bs)
 
+    asNockmaMemRepListConstr :: Asm.Tag -> Sem r NockmaMemRepListConstr
+    asNockmaMemRepListConstr tag = case tag of
+      Asm.UserTag {} -> do
+        rep <- getConstructorMemRep tag
+        case rep of
+          NockmaMemRepList constr -> return constr
+          _ -> impossible
+      Asm.BuiltinTag {} -> impossible
+
+    goBoolTag :: Bool -> Sem r () -> [(Asm.Tag, Sem r ())] -> Sem r ()
+    goBoolTag v b bs = do
+      let otherBranch = fromJust (firstJust f bs <|> defaultBranch)
+      dup
+      if
+          | v -> branch b otherBranch
+          | otherwise -> branch otherBranch b
+      where
+        f :: (Asm.Tag, Sem r ()) -> Maybe (Sem r ())
+        f (tag', br) = case tag' of
+          Asm.UserTag {} -> impossible
+          Asm.BuiltinTag tag -> case nockmaBuiltinTag tag of
+            NockmaBuiltinBool v' -> guard (v /= v') $> br
+
+    goRepList :: NonEmpty (NockmaMemRepListConstr, Sem r ()) -> Sem r ()
+    goRepList ((c, b) :| bs) = do
+      push (isCell' (OpAddress # topOfStack ValueStack))
+      let otherBranch = fromJust (firstJust f bs <|> defaultBranch)
+      case c of
+        NockmaMemRepListConstrCons -> branch b otherBranch
+        NockmaMemRepListConstrNil -> branch otherBranch b
+      where
+        f :: (NockmaMemRepListConstr, Sem r ()) -> Maybe (Sem r ())
+        f (c', br) = guard (c /= c') $> br
+
 branch' ::
   (Functor f, Members '[Output (Term Natural), Reader CompilerCtx] r) =>
   m () ->

test/Nockma/Compile/Positive.hs

@@ -94,11 +94,19 @@ data ConstructorName
   | ConstructorTagged
   | ConstructorPair
   | ConstructorTuple
+  | ConstructorListNil
+  | ConstructorListCons
   deriving stock (Eq, Bounded, Enum)
 
 constructorTag :: ConstructorName -> Asm.Tag
 constructorTag n = Asm.UserTag (Asm.TagUser defaultModuleId (fromIntegral (fromEnum n)))
 
+builtinTrue :: Asm.Tag
+builtinTrue = Asm.BuiltinTag Asm.TagTrue
+
+builtinFalse :: Asm.Tag
+builtinFalse = Asm.BuiltinTag Asm.TagFalse
+
 constructorInfo :: ConstructorName -> ConstructorInfo
 constructorInfo = \case
   ConstructorFalse -> defaultInfo 0
@@ -110,6 +118,16 @@ constructorInfo = \case
       { _constructorInfoArity = 2,
         _constructorInfoMemRep = NockmaMemRepTuple
       }
+  ConstructorListNil ->
+    ConstructorInfo
+      { _constructorInfoArity = 0,
+        _constructorInfoMemRep = NockmaMemRepList NockmaMemRepListConstrNil
+      }
+  ConstructorListCons ->
+    ConstructorInfo
+      { _constructorInfoArity = 2,
+        _constructorInfoMemRep = NockmaMemRepList NockmaMemRepListConstrCons
+      }
 
 defaultInfo :: Natural -> ConstructorInfo
 defaultInfo ari =
@@ -532,21 +550,56 @@ tests =
       "cmdCase: case on builtin true"
       (eqStack ValueStack [nock| [5 nil] |])
       $ do
-        allocConstr (Asm.BuiltinTag Asm.TagTrue)
+        allocConstr builtinTrue
         caseCmd
           (Just (pushNat 0))
-          [ (Asm.BuiltinTag Asm.TagTrue, pop >> pushNat 5),
-            (Asm.BuiltinTag Asm.TagFalse, pushNat 0)
+          [ (builtinTrue, pop >> pushNat 5),
+            (builtinTrue, crash),
+            (builtinFalse, pushNat 0)
           ],
     defTest
       "cmdCase: case on builtin false"
       (eqStack ValueStack [nock| [5 nil] |])
       $ do
-        allocConstr (Asm.BuiltinTag Asm.TagFalse)
+        allocConstr builtinFalse
         caseCmd
           (Just (pushNat 0))
-          [ (Asm.BuiltinTag Asm.TagTrue, pushNat 0),
-            (Asm.BuiltinTag Asm.TagFalse, pop >> pushNat 5)
+          [ (builtinTrue, pushNat 0),
+            (builtinFalse, pop >> pushNat 5),
+            (builtinFalse, crash)
+          ],
+    defTest
+      "cmdCase: case on listy nil"
+      (eqStack ValueStack [nock| [5 nil] |])
+      $ do
+        allocConstr (constructorTag ConstructorListNil)
+        caseCmd
+          (Just (pushNat 0))
+          [ (constructorTag ConstructorListCons, pushNat 0),
+            (constructorTag ConstructorListNil, pop >> pushNat 5),
+            (constructorTag ConstructorListCons, crash)
+          ],
+    defTest
+      "cmdCase: case on listy cons and field accessor"
+      (eqStack ValueStack [nock| [[60 30 nil] nil] |])
+      $ do
+        let tagCons = constructorTag ConstructorListCons
+            tagNil = constructorTag ConstructorListNil
+        allocConstr tagNil
+        pushNat 30
+        allocConstr tagCons
+        caseCmd
+          Nothing
+          [ ( tagCons,
+              do
+                copyTopFromTo ValueStack TempStack
+                pushConstructorFieldOnto ValueStack tagCons (Asm.mkTempRef' 1 0) 0
+                pushConstructorFieldOnto ValueStack tagCons (Asm.mkTempRef' 1 0) 0
+                add
+                allocConstr tagCons
+            ),
+            (tagNil, pop >> pushNat 5),
+            (tagCons, crash)
           ],
     defTest
       "push constructor field"
@@ -567,5 +620,12 @@ tests =
       )
       $ do
         pushNat 10
-        traceTerm (OpAddress # topOfStack ValueStack)
+        traceTerm (OpAddress # topOfStack ValueStack),
+    defTest
+      "allocate listy constructors"
+      (eqStack ValueStack [nock| [[500 nil] nil] |])
+      $ do
+        allocConstr (constructorTag ConstructorListNil)
+        pushNat 500
+        allocConstr (constructorTag ConstructorListCons)
   ]