Type annotations in arguments of function definition

2 files changed, 24 insertions, 9 deletions

diff --git a/src/Parser.hs b/src/Parser.hs
index 62549eb..5d545af 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -80,11 +80,12 @@ parseDefinition href = label "symbol definition" $ do
     def@( name, expr ) <- localState $ do
         wsymbol "def"
         name <- varName
-        argsDecl <- functionArguments (\off _ -> return . ( off, )) varName mzero (\_ -> return . VarName)
-        atypes <- forM argsDecl $ \( off, vname :: VarName ) -> do
-            tvar <- newTypeVar
-            modify $ \s -> s { testVars = ( vname, ( LocalVarName vname, ExprTypeVar tvar )) : testVars s }
-            return ( off, vname, tvar )
+        argsDecl <- functionArguments (\off _ -> return . ( off, ))
+            (typeAnnotated varName) mzero (\_ -> return . (, Nothing) . VarName)
+        atypes <- forM argsDecl $ \( off, ( vname :: VarName, mbstype :: Maybe SomeExprType ) ) -> do
+            stype <- maybe (ExprTypeVar <$> newTypeVar) return mbstype
+            modify $ \s -> s { testVars = ( vname, ( LocalVarName vname, stype )) : testVars s }
+            return ( off, vname, stype )
         SomeExpr expr <- choice
             [ do
                 osymbol ":"
@@ -102,8 +103,8 @@ parseDefinition href = label "symbol definition" $ do
     modify $ \s -> s { testVars = ( name, ( GlobalVarName (testCurrentModuleName s) name, someExprType expr )) : testVars s }
     return def
   where
-    getInferredTypes atypes = forM atypes $ \( _, vname, tvar ) -> do
-        ( vname, ) . SomeArgumentType OptionalArgument <$> typeClosure (ExprTypeVar tvar)
+    getInferredTypes atypes = forM atypes $ \( _, vname, stype ) -> do
+        ( vname, ) . SomeArgumentType OptionalArgument <$> typeClosure stype
 
     replaceDynArgs :: forall a. Expr a -> TestParser (Expr a)
     replaceDynArgs expr = do
@@ -120,6 +121,18 @@ parseDefinition href = label "symbol definition" $ do
                 replaceArgs (SomeExpr e) = SomeExpr (go unif e)
             e -> e
 
+    typeAnnotated p = do
+        x <- p
+        choice
+            [ do
+                void $ osymbol ":"
+                stype <- typeExpr
+                return ( x, Just stype )
+
+            , do
+                return ( x, Nothing )
+            ]
+
 parseAsset :: Pos -> TestParser ( VarName, SomeExpr )
 parseAsset href = label "asset definition" $ do
     wsymbol "asset"
diff --git a/test/asset/parser/type-annotation.et b/test/asset/parser/type-annotation.et
index 5fbf323..811b54b 100644
--- a/test/asset/parser/type-annotation.et
+++ b/test/asset/parser/type-annotation.et
@@ -1,6 +1,8 @@
-def fun (a) and b:
-    let c = [ a : String, "x" ]
+def fun (a : String) and (b : String):
+    let c = [ a, "x" ]
     let d = [ b : String, "y" : String ]
+    guard (c /= d)
 
 test Test:
     let x = "" : String
+    fun "a" and "b"