Fix basic block calls in the Reg -> CASM translation

src/Juvix/Compiler/Casm/Translation/FromReg.hs

@@ -194,7 +194,9 @@ fromReg tab = mkResult $ run $ runLabelInfoBuilderWithNextId (Reg.getNextSymbolI
           Nothing -> do
             massert (isJust mout)
             massert (HashSet.member (fromJust mout) liveVars0)
-            goCallBlock False Nothing liveVars0
+            goAssignApBuiltins
+            whenJust mout saveLiveVar
+            goCallBlock mout liveVars0
       where
         output'' :: Instruction -> Sem r ()
         output'' i = do
@@ -206,22 +208,24 @@ fromReg tab = mkResult $ run $ runLabelInfoBuilderWithNextId (Reg.getNextSymbolI
           output'' i
           incAP apOff
 
-        goCallBlock :: Bool -> Maybe Reg.VarRef -> HashSet Reg.VarRef -> Sem r ()
-        goCallBlock updatedBuiltins outVar liveVars = do
-          let liveVars' = toList (maybe liveVars (`HashSet.delete` liveVars) outVar)
+        saveLiveVar :: Reg.VarRef -> Sem r ()
+        saveLiveVar var = do
+          ref <- mkMemRef var
+          let comment = Reg.ppPrint tab var
+          goAssignAp' (Just comment) (Val (Ref ref))
+
+        -- The `goCallBlock` function is used to switch to a new basic block.
+        -- Assumes that the builtins pointer and outVar (if present) were
+        -- already saved (in this order).
+        goCallBlock :: Maybe Reg.VarRef -> HashSet Reg.VarRef -> Sem r ()
+        goCallBlock outVar liveVars = do
+          let liveVars' = sort $ toList (maybe liveVars (`HashSet.delete` liveVars) outVar)
               n = length liveVars'
-              bltOff =
-                if
-                    | updatedBuiltins ->
-                        -argsOffset - n - fromEnum (isJust outVar)
-                    | otherwise ->
-                        -argsOffset - n
+              bltOff = -argsOffset - n - fromEnum (isJust outVar)
               vars =
                 HashMap.fromList $
-                  maybe [] (\var -> [(var, -argsOffset - n - if updatedBuiltins then 0 else 1)]) outVar
+                  maybe [] (\var -> [(var, -argsOffset - n)]) outVar
                     ++ zipWithExact (\var k -> (var, -argsOffset - k)) liveVars' [0 .. n - 1]
-          unless updatedBuiltins $
-            goAssignApBuiltins
           mapM_ saveLiveVar (reverse liveVars')
           output'' (mkCallRel $ Imm 3)
           output'' Return
@@ -232,12 +236,6 @@ fromReg tab = mkResult $ run $ runLabelInfoBuilderWithNextId (Reg.getNextSymbolI
           setAP 0
           setVars vars
           setBuiltinOffset bltOff
-          where
-            saveLiveVar :: Reg.VarRef -> Sem r ()
-            saveLiveVar var = do
-              ref <- mkMemRef var
-              let comment = Reg.ppPrint tab var
-              goAssignAp' (Just comment) (Val (Ref ref))
 
         goLocalBlock :: Int -> HashMap Reg.VarRef Int -> Int -> HashSet Reg.VarRef -> Maybe Reg.VarRef -> Reg.Block -> Sem r ()
         goLocalBlock ap0 vars bltOff liveVars mout' block = do
@@ -573,7 +571,11 @@ fromReg tab = mkResult $ run $ runLabelInfoBuilderWithNextId (Reg.getNextSymbolI
           val <- mkMemRef _instrExtendClosureValue
           goAssignAp (Val $ Ref val)
           output'' $ mkCallRel $ Lab $ LabelRef (blts ^. stdlibExtendClosure) (Just (blts ^. stdlibExtendClosureName))
-          goCallBlock False (Just _instrExtendClosureResult) liveVars
+          -- the `juvix_extend_closure` runtime function does not accept or
+          -- return the builtins pointer
+          goAssignApBuiltins
+          goAssignAp (Val $ Ref $ MemRef Ap (-2))
+          goCallBlock (Just _instrExtendClosureResult) liveVars
 
         goCall' :: Reg.CallType -> [Reg.Value] -> Sem r ()
         goCall' ct args = case ct of
@@ -593,7 +595,7 @@ fromReg tab = mkResult $ run $ runLabelInfoBuilderWithNextId (Reg.getNextSymbolI
         goCall :: HashSet Reg.VarRef -> Reg.InstrCall -> Sem r ()
         goCall liveVars Reg.InstrCall {..} = do
           goCall' _instrCallType _instrCallArgs
-          goCallBlock True (Just _instrCallResult) liveVars
+          goCallBlock (Just _instrCallResult) liveVars
 
         -- There is no way to make "proper" tail calls in Cairo, because
         -- the only way to set the `fp` register is via the `call` instruction.

src/Juvix/Compiler/Reg/Language/Instrs.hs

@@ -26,7 +26,7 @@ data ConstrField = ConstrField
 data VarGroup
   = VarGroupArgs
   | VarGroupLocal
-  deriving stock (Eq, Generic, Show)
+  deriving stock (Eq, Ord, Generic, Show)
 
 instance Hashable VarGroup
 
@@ -48,6 +48,9 @@ instance Eq VarRef where
     vr1 ^. varRefGroup == vr2 ^. varRefGroup
       && vr1 ^. varRefIndex == vr2 ^. varRefIndex
 
+instance Ord VarRef where
+  compare = compare `on` \vr -> (vr ^. varRefGroup, vr ^. varRefIndex)
+
 deriving stock instance (Eq ConstrField)
 
 deriving stock instance (Eq Value)

test/Casm/Compilation/Positive.hs

@@ -565,5 +565,13 @@ tests =
         $(mkRelDir ".")
         $(mkRelFile "test078.juvix")
         Nothing
-        $(mkRelFile "out/test078.out")
+        $(mkRelFile "out/test078.out"),
+      posTest
+        "Test079: Trivial resource logic"
+        False
+        True
+        $(mkRelDir ".")
+        $(mkRelFile "test079.juvix")
+        (Just $(mkRelFile "in/test079.json"))
+        $(mkRelFile "out/test079.out")
     ]

test/Casm/Reg/Cairo.hs

@@ -53,5 +53,11 @@ cairoTests =
       $(mkRelDir ".")
       $(mkRelFile "test042.jvr")
       $(mkRelFile "out/test042.out")
-      Nothing
+      Nothing,
+    P.PosTest
+      "Test048: Mock resource logic"
+      $(mkRelDir ".")
+      $(mkRelFile "test048.jvr")
+      $(mkRelFile "out/test048.out")
+      (Just $(mkRelFile "in/test048.json"))
   ]

test/Casm/Reg/Positive.hs

@@ -228,5 +228,29 @@ tests =
       $(mkRelDir ".")
       $(mkRelFile "test043.jvr")
       $(mkRelFile "out/test043.out")
-      Nothing
+      Nothing,
+    PosTest
+      "Test044: Extend closure in true branch"
+      $(mkRelDir ".")
+      $(mkRelFile "test044.jvr")
+      $(mkRelFile "out/test044.out")
+      (Just $(mkRelFile "in/test044.json")),
+    PosTest
+      "Test045: Extend closure in false branch"
+      $(mkRelDir ".")
+      $(mkRelFile "test045.jvr")
+      $(mkRelFile "out/test045.out")
+      (Just $(mkRelFile "in/test045.json")),
+    PosTest
+      "Test046: Call in true branch"
+      $(mkRelDir ".")
+      $(mkRelFile "test046.jvr")
+      $(mkRelFile "out/test046.out")
+      (Just $(mkRelFile "in/test046.json")),
+    PosTest
+      "Test047: Call in false branch"
+      $(mkRelDir ".")
+      $(mkRelFile "test047.jvr")
+      $(mkRelFile "out/test047.out")
+      (Just $(mkRelFile "in/test047.json"))
   ]

tests/Casm/Compilation/positive/in/test079.json

@@ -0,0 +1,31 @@
+{
+  "self_resource": {
+    "logic": "0x373bb1d37414c2edf111cf2f9f076517da99d38e44cdd716ca2ad00a07731e5",
+    "label": "0x12",
+    "quantity": "0x13",
+    "data": "0x14",
+    "eph": false,
+    "nonce": "0x26",
+    "npk": "0x7752582c54a42fe0fa35c40f07293bb7d8efe90e21d8d2c06a7db52d7d9b7a1",
+    "rseed": "0x48"
+  },
+  "resource_nf_key": "0x1",
+  "merkle_path": [
+    {
+      "fst": "0x33",
+      "snd": true
+    },
+    {
+      "fst": "0x83",
+      "snd": false
+    },
+    {
+      "fst": "0x73",
+      "snd": false
+    },
+    {
+      "fst": "0x23",
+      "snd": false
+    }
+  ]
+}

tests/Casm/Compilation/positive/out/test079.out

@@ -0,0 +1,16 @@
+-395556098205570544514862427675713135822454173064116032036066260837844667166
+1199411604616075448730293334298953634347373067899103883531383580500542532805
+741345583372708758677381254341695646215212764479558601822324496916305564860
+906223455619847457072543249061510257957686047823824674682567491887881744554
+-767849295212578246709801782763128612227201311770055823223370752694466982385
+-499605824796177847482548145602529914672118948474574930466002320448397702145
+-739783929947191450539981235446122797172963523649497524598736585887720470605
+11906781661858318111608384113371656853984100876730919753607956359054872835
+1184591657410410618259810478559756080508915282346219818107533089968423438529
+-1196756736539952528850324578540484369187211012447802934902211689810543298321
+-615406306260154438751417537243664316178100185892451009663053034794681290878
+-479989799644409226462391378855274180707944349292197032426785420935231903693
+1666094585867140022923253270503506891580143265539166904031542100935830693071
+874739451078007766457464989774322083649278607533249481151382481072868806602
+152666792071518830868575557812948353041420400780739481342941381225525861407
+1