VampIR pipeline: handle booleans in the type of main (#2137)

* Closes #2132

---------

Co-authored-by: Jonathan Cubides <jonathan.cubides@uib.no>

src/Juvix/Compiler/Backend/VampIR/Language.hs

@@ -48,7 +48,9 @@ data Function = Function
   { _functionName :: Text,
     _functionArguments :: [Text],
     _functionLocalDefs :: [LocalDef],
-    _functionExpression :: Expression
+    _functionExpression :: Expression,
+    _functionInputs :: [Text],
+    _functionOutput :: Text
   }
 
 newtype Program = Program

src/Juvix/Compiler/Backend/VampIR/Pretty/Base.hs

@@ -70,10 +70,10 @@ instance PrettyCode Function where
 
 ppEquation :: Function -> Sem r (Doc Ann)
 ppEquation Function {..} = do
-  let n = length _functionArguments
-      args = if n == 1 then ["(in + 0)"] else map (\k -> "(in" <> show k <> " + 0)") [1 .. n]
+  let args = map (\arg -> "(" <> pretty arg <> " + 0)") _functionInputs
+      out = pretty _functionOutput
   fn <- ppName KNameFunction _functionName
-  return $ fn <+> hsep args <+> kwEq <+> "(out + 0)" <> semi
+  return $ fn <+> hsep args <+> kwEq <+> "(" <> out <> " + 0)" <> semi
 
 instance PrettyCode Program where
   ppCode Program {..} = do

src/Juvix/Compiler/Backend/VampIR/Translation/FromCore.hs

@@ -9,19 +9,28 @@ import Juvix.Compiler.Core.Language as Core
 import Juvix.Compiler.Core.Transformation.DisambiguateNames (disambiguateNodeNames')
 
 fromCore :: InfoTable -> Program
-fromCore tab = fromCoreNode $ lookupIdentifierNode tab (fromJust (tab ^. infoMain))
+fromCore tab = fromCoreNode (ii ^. identifierType) node
+  where
+    sym = fromJust (tab ^. infoMain)
+    node = lookupIdentifierNode tab sym
+    ii = lookupIdentifierInfo tab sym
 
-fromCoreNode :: Node -> Program
-fromCoreNode node =
+fromCoreNode :: Type -> Node -> Program
+fromCoreNode ty node =
   let (lams, body) = unfoldLambdas (disambiguateNodeNames' disambiguate emptyInfoTable node)
       (defs, expr) = convertLets body
+      n = length lams
+      args = map (\k -> "in" <> if n == 1 then "" else show k) [1 .. n]
+      isBoolTarget = isTypeBool (typeTarget ty)
    in Program
         { _programFunctions =
             [ Function
                 { _functionName = "main",
                   _functionArguments = map (^. lambdaLhsBinder . binderName) lams,
                   _functionLocalDefs = defs,
-                  _functionExpression = expr
+                  _functionExpression = expr,
+                  _functionInputs = args,
+                  _functionOutput = if isBoolTarget then "1" else "out"
                 }
             ]
         }

src/Juvix/Compiler/Core/Evaluator.hs

@@ -157,7 +157,8 @@ geval opts herr ctx env0 = eval' env0
     evalBody i bi env v body
       | opts ^. evalOptionsNormalize
           && isTypePrim (bi ^. binderType)
-          && not (isImmediate v) =
+          && not (isImmediate v)
+          && not (isFailNode v) =
           Closure env (mkLet i bi v body)
       | otherwise =
           eval' (v : env) body

src/Juvix/Compiler/Core/Extra/Base.hs

@@ -236,6 +236,16 @@ isTypePrim = \case
   NPrim {} -> True
   _ -> False
 
+isTypeInteger :: Type -> Bool
+isTypeInteger = \case
+  NPrim (TypePrim _ (PrimInteger _)) -> True
+  _ -> False
+
+isTypeBool :: Type -> Bool
+isTypeBool = \case
+  NPrim (TypePrim _ (PrimBool _)) -> True
+  _ -> False
+
 -- | `expandType argtys ty` expands the dynamic target of `ty` to match the
 -- number of arguments with types specified by `argstys`. For example,
 -- `expandType [int, string] (int -> any) = int -> string -> any`.

src/Juvix/Compiler/Core/Extra/Utils.hs

@@ -102,6 +102,11 @@ isDataValue = \case
   NCtr Constr {..} -> all isDataValue _constrArgs
   _ -> False
 
+isFailNode :: Node -> Bool
+isFailNode = \case
+  NBlt (BuiltinApp {..}) | _builtinAppOp == OpFail -> True
+  _ -> False
+
 freeVarsSortedMany :: [Node] -> Set Var
 freeVarsSortedMany n = Set.fromList (n ^.. each . freeVars)
 

src/Juvix/Compiler/Core/Transformation/Check/Base.hs

@@ -72,7 +72,7 @@ checkBuiltins allowUntypedFail = dmapRM go
 
 -- | Checks that the root of the node is not `Bottom`. Currently the only way we
 -- create `Bottom` is when translating axioms that are not builtin. Hence it is
--- enought to check the root only.
+-- enough to check the root only.
 checkNoAxioms :: forall r. Member (Error CoreError) r => InfoTable -> Sem r ()
 checkNoAxioms = void . mapT' checkNodeNoAxiom
   where
@@ -127,7 +127,7 @@ checkNoRecursiveTypes tab =
   when (isCyclic (createTypeDependencyInfo tab)) $
     throw
       CoreError
-        { _coreErrorMsg = ppOutput "recursive types not supported for the GEB target",
+        { _coreErrorMsg = ppOutput "recursive types not supported for this target",
           _coreErrorNode = Nothing,
           _coreErrorLoc = defaultLoc
         }

src/Juvix/Compiler/Core/Transformation/Check/VampIR.hs

@@ -19,7 +19,7 @@ checkVampIR tab =
       unless (checkType (ii ^. identifierType)) $
         throw
           CoreError
-            { _coreErrorMsg = ppOutput "for this target the arguments and the result of the `main` function must be numbers",
+            { _coreErrorMsg = ppOutput "for this target the arguments and the result of the `main` function must be numbers or booleans",
               _coreErrorLoc = fromMaybe defaultLoc (ii ^. identifierLocation),
               _coreErrorNode = Nothing
             }
@@ -29,9 +29,7 @@ checkVampIR tab =
     checkType :: Node -> Bool
     checkType ty =
       let (tyargs, tgt) = unfoldPi' ty
-       in all isPrimInteger (tgt : tyargs)
+       in all isPrimIntegerOrBool (tgt : tyargs)
       where
-        isPrimInteger ty' = case ty' of
-          NPrim (TypePrim _ (PrimInteger _)) -> True
-          NDyn _ -> True
-          _ -> False
+        isPrimIntegerOrBool ty' =
+          isTypeInteger ty' || isTypeBool ty' || isDynamic ty'

test/Core/Eval/Base.hs

@@ -40,7 +40,8 @@ instance FromJSON EvalData where
       parseEvalData :: Parse JSONError EvalData
       parseEvalData = do
         _evalDataInput <- parseInputs
-        _evalDataOutput <- key "out" asText
+        mout <- keyMay "out" asText
+        let _evalDataOutput = fromMaybe "true" mout
         return EvalData {..}
 
       parseInputs :: Parse JSONError [Text]
@@ -66,7 +67,7 @@ coreEvalAssertion' ::
   Assertion
 coreEvalAssertion' mode tab mainFile expectedFile step =
   length (fromText (ppPrint tab) :: String) `seq`
-    case (tab ^. infoMain) >>= ((tab ^. identContext) HashMap.!?) of
+    case HashMap.lookup sym (tab ^. identContext) of
       Just node -> do
         d <- readEvalData
         case d of
@@ -77,7 +78,8 @@ coreEvalAssertion' mode tab mainFile expectedFile step =
                   let outputFile = dirPath <//> $(mkRelFile "out.out")
                   hout <- openFile (toFilePath outputFile) WriteMode
                   step "Evaluate"
-                  let args = map (mkConstant' . ConstInteger . fst . fromRight' . decimal) _evalDataInput
+                  let tyargs = typeArgs (lookupIdentifierInfo tab sym ^. identifierType)
+                      args = zipWith mkArg (tyargs ++ repeat mkDynamic') _evalDataInput
                       node' = mkApps' node args
                   r' <- doEval mainFile hout tab node'
                   case r' of
@@ -95,6 +97,16 @@ coreEvalAssertion' mode tab mainFile expectedFile step =
               )
       Nothing -> assertFailure ("No main function registered in: " <> toFilePath mainFile)
   where
+    sym = fromJust (tab ^. infoMain)
+
+    mkArg :: Type -> Text -> Node
+    mkArg ty arg =
+      let n = fst $ fromRight' $ decimal arg
+       in case (isTypeBool ty, n) of
+            (True, 0) -> mkConstr' (BuiltinTag TagFalse) []
+            (True, _) -> mkConstr' (BuiltinTag TagTrue) []
+            (False, _) -> mkConstant' (ConstInteger n)
+
     readEvalData :: IO (Either String EvalData)
     readEvalData = case mode of
       EvalModePlain -> do

test/VampIR/Compilation/Positive.hs

@@ -158,5 +158,11 @@ tests =
       "Test019: polymorphism"
       $(mkRelDir ".")
       $(mkRelFile "test019.juvix")
-      $(mkRelFile "data/test019.json")
+      $(mkRelFile "data/test019.json"),
+    posTest
+      8
+      "Test020: boolean target"
+      $(mkRelDir ".")
+      $(mkRelFile "test020.juvix")
+      $(mkRelFile "data/test020.json")
   ]

tests/VampIR/positive/Compilation/data/test020.json

@@ -0,0 +1,4 @@
+{
+    "in1": "9",
+    "in2": "0"
+}

tests/VampIR/positive/Compilation/test020.juvix

@@ -0,0 +1,12 @@
+-- boolean target
+module test020;
+
+import Stdlib.Prelude open;
+
+{-# unroll: 10 #-}
+compose : {A : Type} → Nat → (A → A) → A → A;
+compose zero _ := id;
+compose (suc n) f := f ∘ compose n f;
+
+main : Nat → Bool → Bool;
+main n b := compose n not b;