add: variadic functions

lib/Ast/Parser/Type.hs

@@ -4,6 +4,7 @@ import qualified Ast.Parser.State as PS
 import qualified Ast.Parser.Utils as PU
 import qualified Ast.Types as AT
 import qualified Control.Monad.State as S
+import qualified Data.Maybe as M
 import qualified Text.Megaparsec as M
 import qualified Text.Megaparsec.Char as MC
 import qualified Text.Megaparsec.Char.Lexer as ML
@@ -78,9 +79,12 @@ arrayType = do
 -- TODO: find a way to do it without the parenthesis and avoid the infinite loop of parseType
 functionType :: PU.Parser AT.Type
 functionType = do
-  paramTypes <- M.some $ PU.lexeme (M.between (PU.symbol "(") (PU.symbol ")") functionType M.<|> parseTermType)
-  returnType <- PU.symbol "->" *> PU.lexeme parseTermType
-  return $ AT.TFunction {AT.returnType = returnType, AT.paramTypes = paramTypes, AT.isVariadic = False}
+  paramTypes <- M.some $ PU.lexeme $ functionParser M.<|> parseTermType
+  variadic <- M.fromMaybe False <$> M.optional (True <$ PU.symbol "...")
+  returnType <- PU.symbol "->" *> PU.lexeme (functionParser M.<|> parseTermType)
+  return $ AT.TFunction {AT.returnType = returnType, AT.paramTypes = paramTypes, AT.isVariadic = variadic}
+  where
+    functionParser = M.between (PU.symbol "(") (PU.symbol ")") functionType
 
 customType :: PU.Parser AT.Type
 customType = do

test/Ast/Parser/ExprSpec.hs

@@ -56,6 +56,19 @@ spec = do
       let result = normalizeExpr <$> parseWithEnv input
       result `shouldBe` expected
 
+    it "parses a variadic function" $ do
+      let input = "printf: *byte ... -> int = s { 1 }"
+      let expected =
+            Right $
+              AT.Function
+                normalizeLoc
+                "printf"
+                (AT.TFunction (AT.TInt 32) [AT.TPointer $ AT.TInt 8] True)
+                ["s"]
+                (AT.Block [AT.Return normalizeLoc (Just (AT.Lit normalizeLoc (AT.LInt 1)))])
+      let result = normalizeExpr <$> parseWithEnv input
+      result `shouldBe` expected
+
     it "parses a variable declaration with initialization" $
       do
         let input = "x : int = 42"

test/Ast/Parser/TypeSpec.hs

@@ -62,6 +62,24 @@ spec = do
     it "parses int int -> never" $ do
       parseWithEnv "int int -> never" `shouldBe` Right (AT.TFunction {AT.returnType = AT.TVoid, AT.paramTypes = [AT.TInt 32, AT.TInt 32], AT.isVariadic = False})
 
+    it "parses a HOF that takes a function" $ do
+      let input = "(int -> int) -> never"
+      let result = parseWithEnv input
+      let expected = Right $ AT.TFunction AT.TVoid [AT.TFunction (AT.TInt 32) [AT.TInt 32] False] False
+      result `shouldBe` expected
+
+    it "parses a HOF that returns a function" $ do
+      let input = "never -> (int -> int)"
+      let result = parseWithEnv input
+      let expected = Right $ AT.TFunction (AT.TFunction (AT.TInt 32) [AT.TInt 32] False) [AT.TVoid] False
+      result `shouldBe` expected
+
+    it "parses a variadic function" $ do
+      let input = "*byte ... -> int"
+      let result = parseWithEnv input
+      let expected = Right $ AT.TFunction (AT.TInt 32) [AT.TPointer $ AT.TInt 8] True
+      result `shouldBe` expected
+
   describe "Custom Types" $ do
     it "parses a defined custom struct type" $ do
       let env = PS.insertType "Point" (AT.TStruct "Point" [("x", AT.TInt 32), ("y", AT.TInt 32)]) PS.parserState