Use Anoma compatible Nockma serialization of Bools and List-like data structures

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 [email protected] {..} =
             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)
   ]