Function arguments unification

3 files changed, 128 insertions, 55 deletions

diff --git a/src/Parser/Core.hs b/src/Parser/Core.hs
index 3a3450b..7831682 100644
--- a/src/Parser/Core.hs
+++ b/src/Parser/Core.hs
@@ -1,6 +1,7 @@
 module Parser.Core where
 
 import Control.Applicative
+import Control.Arrow
 import Control.Monad
 import Control.Monad.State
 
@@ -119,6 +120,26 @@ lookupScalarVarExpr off sline name = do
             SomeExpr <$> unifyExpr off pa (FunVariable args sline fqn :: Expr (FunctionType a))
         stype -> return $ SomeExpr $ DynVariable stype sline fqn
 
+
+resolveKnownTypeVars :: SomeExprType -> TestParser SomeExprType
+resolveKnownTypeVars stype = case stype of
+    ExprTypePrim {} -> return stype
+    ExprTypeConstr1 {} -> return stype
+    ExprTypeVar tvar -> do
+        gets (M.lookup tvar . testTypeUnif) >>= \case
+            Just stype' -> resolveKnownTypeVars stype'
+            Nothing -> return stype
+    ExprTypeFunction args body -> ExprTypeFunction <$> resolveKnownTypeVars args <*> resolveKnownTypeVars body
+    ExprTypeArguments args -> ExprTypeArguments <$> mapM (\(SomeArgumentType a t) -> SomeArgumentType a <$> resolveKnownTypeVars t) args
+    ExprTypeApp ctor params -> do
+        ctor' <- resolveKnownTypeVars ctor
+        params' <- mapM resolveKnownTypeVars params
+        return $ case ( ctor', params' ) of
+            ( ExprTypeConstr1 (Proxy :: Proxy c'), [ ExprTypePrim (Proxy :: Proxy p') ] )
+                -> ExprTypePrim (Proxy :: Proxy (c' p'))
+            _ -> ExprTypeApp ctor' params'
+    ExprTypeForall tvar inner -> ExprTypeForall tvar <$> resolveKnownTypeVars inner
+
 unify :: Int -> SomeExprType -> SomeExprType -> TestParser SomeExprType
 unify _ (ExprTypeVar aname) (ExprTypeVar bname) | aname == bname = do
     cur <- gets testTypeUnif
@@ -206,6 +227,23 @@ unify off a b = do
         "couldn't match expected type ‘" <> textSomeExprType a <> "’ with actual type ‘" <> textSomeExprType b <> "’"
 
 
+unifyArguments
+    :: FunctionArguments SomeArgumentType
+    -> FunctionArguments ( Int, SomeExpr )
+    -> TestParser ( FunctionArguments SomeExpr, ( FunctionArguments SomeArgumentType, FunctionArguments ( Int, SomeExpr ) ) )
+unifyArguments (FunctionArguments am) (FunctionArguments bm) = (toArgs *** (toArgs *** toArgs)) <$> go (M.toAscList am) (M.toAscList bm)
+  where
+    toArgs = FunctionArguments . M.fromAscList
+    go [] bs = return ( [], ( [], bs ) )
+    go as [] = return ( [], ( as, [] ) )
+    go (a@( ak, SomeArgumentType _ at ) : as) (b@( bk, ( off, expr ) ) : bs)
+        | ak < bk = second (first  (a :)) <$> go as (b : bs)
+        | bk < ak = second (second (b :)) <$> go (a : as) bs
+        | otherwise = do
+            expr' <- unifySomeExpr off at expr
+            first (( ak, expr' ) :) <$> go as bs
+
+
 unifyExpr :: forall a b proxy. (ExprType a, ExprType b) => Int -> proxy a -> Expr b -> TestParser (Expr a)
 unifyExpr off pa expr = if
     | Just (Refl :: a :~: b) <- eqT
@@ -237,23 +275,7 @@ unifyExpr off pa expr = if
         unifyExpr off pa (f $ ExprTypePrim pt)
 
     | Just (Refl :: FunctionType a :~: b) <- eqT
-    -> do
-        let FunctionArguments remaining = exprArgs expr
-            showType ( Nothing, SomeArgumentType _ stype ) = "‘<" <> textSomeExprType stype <> ">’"
-            showType ( Just (ArgumentKeyword kw), SomeArgumentType _ stype ) = "‘" <> kw <> " <" <> textSomeExprType stype <> ">’"
-            err = parseError . FancyError off . S.singleton . ErrorFail . T.unpack
-
-        defaults <- fmap catMaybes $ forM (M.toAscList remaining) $ \case
-            arg@( _, SomeArgumentType RequiredArgument _ ) -> err $ "missing " <> showType arg <> " argument"
-            ( _, SomeArgumentType OptionalArgument _ ) -> return Nothing
-            ( kw, SomeArgumentType (ExprDefault def) _ ) -> return $ Just ( kw, def )
-            ( kw, SomeArgumentType ContextDefault (ExprTypePrim atype) ) -> do
-                SomeExpr context <- gets testContext
-                context' <- unifyExpr off atype context
-                return $ Just ( kw, SomeExpr context' )
-            ( _, SomeArgumentType ContextDefault _ ) -> err "non-primitive context requirement"
-        sline <- getSourceLine
-        return (FunctionEval sline $ ArgsApp (FunctionArguments $ M.fromAscList defaults) expr)
+    -> evalRemainingArguments off (exprArgs expr) expr
 
     | Just (Refl :: DynamicType :~: b) <- eqT
     , Undefined msg <- expr
@@ -266,6 +288,25 @@ unifyExpr off pa expr = if
             "couldn't match expected type ‘" <> textExprType pa <> "’ with actual type ‘" <> textExprType expr <> "’"
 
 
+evalRemainingArguments :: ExprType a => Int -> FunctionArguments SomeArgumentType -> Expr (FunctionType a) -> TestParser (Expr a)
+evalRemainingArguments off (FunctionArguments remaining) expr = do
+    let showType ( Nothing, SomeArgumentType _ stype ) = "‘<" <> textSomeExprType stype <> ">’"
+        showType ( Just (ArgumentKeyword kw), SomeArgumentType _ stype ) = "‘" <> kw <> " <" <> textSomeExprType stype <> ">’"
+        err = parseError . FancyError off . S.singleton . ErrorFail . T.unpack
+
+    defaults <- fmap catMaybes $ forM (M.toAscList remaining) $ \case
+        arg@( _, SomeArgumentType RequiredArgument _ ) -> err $ "missing " <> showType arg <> " argument"
+        ( _, SomeArgumentType OptionalArgument _ ) -> return Nothing
+        ( kw, SomeArgumentType (ExprDefault def) _ ) -> return $ Just ( kw, def )
+        ( kw, SomeArgumentType ContextDefault (ExprTypePrim atype) ) -> do
+            SomeExpr context <- gets testContext
+            context' <- unifyExpr off atype context
+            return $ Just ( kw, SomeExpr context' )
+        ( _, SomeArgumentType ContextDefault _ ) -> err "non-primitive context requirement"
+    sline <- getSourceLine
+    return (FunctionEval sline $ ArgsApp (FunctionArguments $ M.fromAscList defaults) expr)
+
+
 unifySomeExpr :: Int -> SomeExprType -> SomeExpr -> TestParser SomeExpr
 unifySomeExpr off stype sexpr@(SomeExpr (expr :: Expr a))
     | ExprTypePrim pa <- stype
@@ -298,6 +339,23 @@ unifySomeExpr off stype sexpr@(SomeExpr (expr :: Expr a))
             SomeExpr expr' <- unifySomeExpr off res sexpr
             return $ SomeExpr $ FunctionAbstraction expr'
 
+    | ExprTypeApp _ _ <- stype
+    , ExprTypeFunction args' res' <- someExprType sexpr
+    = do
+        ( _, ( remaining, _ ) ) <- case args' of
+            ExprTypeArguments args'' -> do
+                unifyArguments args'' mempty
+            _ -> do
+                _ <- unify off (ExprTypeArguments mempty) args'
+                return ( mempty, ( mempty, mempty ) )
+        unify off stype res' >>= \case
+            ExprTypePrim (Proxy :: Proxy r) | Just (Refl :: a :~: FunctionType r) <- eqT ->
+                SomeExpr <$> evalRemainingArguments off remaining expr
+            _ | Just (Refl :: a :~: FunctionType DynamicType) <- eqT ->
+                SomeExpr <$> evalRemainingArguments off remaining expr
+            _ ->
+                error $ "expecting function type, got: " <> show (typeRep expr)
+
     | otherwise
     = do
         _ <- unify off stype (someExprType sexpr)
diff --git a/src/Parser/Expr.hs b/src/Parser/Expr.hs
index c8a9e85..7d57791 100644
--- a/src/Parser/Expr.hs
+++ b/src/Parser/Expr.hs
@@ -16,8 +16,8 @@ module Parser.Expr (
 
     stringExpansion,
 
-    checkFunctionArguments,
     functionArguments,
+    applyFunctionArguments,
 ) where
 
 import Control.Applicative (liftA2)
@@ -418,27 +418,9 @@ constructor = label "constructor" $ do
 functionCall :: TestParser SomeExpr
 functionCall = do
     sline <- getSourceLine
-    off <- stateOffset <$> getParserState
-
     fun <- variable <|> constructor
-    FunctionArguments margs <- functionArguments (\poff _ e -> return ( poff, e )) (someExpr FunctionTerm) literal (\poff -> lookupVarExpr poff sline . VarName)
-    if M.null margs
-      then return fun
-      else do
-        dict <- newTypeVar
-        res <- newTypeVar
-        SomeExpr (expr :: Expr fa) <- unifySomeExpr off (ExprTypeFunction (ExprTypeVar dict) (ExprTypeVar res)) fun
-        Just (ExprTypeArguments argTypes) <- M.lookup dict <$> gets testTypeUnif
-        args <- fmap (FunctionArguments . M.fromAscList) $ mapM (\( kw, ( poff, e ) ) -> ( kw, ) <$> checkFunctionArguments argTypes poff kw e) $ M.toAscList margs
-        M.lookup res <$> gets testTypeUnif >>= \case
-            Just (ExprTypePrim (_ :: Proxy a))
-                | Just (Refl :: FunctionType a :~: fa) <- eqT
-                -> return $ SomeExpr $ ArgsApp args expr
-                | otherwise -> error $ "type mismatch after function unification: " <> show ( typeRep (Proxy @(FunctionType a)), typeRep (Proxy @fa) )
-            _
-                | Just (Refl :: FunctionType DynamicType :~: fa) <- eqT
-                -> return $ SomeExpr $ ArgsApp args expr
-                | otherwise -> error $ "type mismatch after function unification: " <> show ( typeRep (Proxy @(FunctionType DynamicType)), typeRep (Proxy @fa) )
+    args <- functionArguments (\poff _ e -> return ( poff, e )) (someExpr FunctionTerm) literal (\poff -> lookupVarExpr poff sline . VarName)
+    applyFunctionArguments args fun
 
 recordSelector :: SomeExpr -> TestParser SomeExpr
 recordSelector (SomeExpr expr) = do
@@ -454,21 +436,6 @@ recordSelector (SomeExpr expr) = do
     applyRecordSelector m e (RecordSelector f) = SomeExpr $ App (AnnRecord m) (pure f) e
 
 
-checkFunctionArguments :: FunctionArguments SomeArgumentType
-                       -> Int -> Maybe ArgumentKeyword -> SomeExpr -> TestParser SomeExpr
-checkFunctionArguments (FunctionArguments argTypes) poff kw expr = do
-    case M.lookup kw argTypes of
-        Just (SomeArgumentType _ stype) -> do
-            withRecovery (\e -> registerParseError e >> return expr) $ do
-                unifySomeExpr poff stype expr
-        Nothing -> do
-            registerParseError $ FancyError poff $ S.singleton $ ErrorFail $ T.unpack $
-                case kw of
-                    Just (ArgumentKeyword tkw) -> "unexpected parameter with keyword ‘" <> tkw <> "’"
-                    Nothing                    -> "unexpected parameter"
-            return expr
-
-
 functionArguments :: (Int -> Maybe ArgumentKeyword -> a -> TestParser b) -> TestParser a -> TestParser a -> (Int -> Text -> TestParser a) -> TestParser (FunctionArguments b)
 functionArguments check param lit promote = do
     args <- parseArgs True
@@ -496,3 +463,51 @@ functionArguments check param lit promote = do
     pparam = between (symbol "(") (symbol ")") param <|> lit
 
     checkAndInsert off kw x cont = M.insert kw <$> check off kw x <*> cont
+
+
+applyFunctionArguments :: FunctionArguments ( Int, SomeExpr ) -> SomeExpr -> TestParser SomeExpr
+applyFunctionArguments (FunctionArguments margs) sexpr
+    | M.null margs = return sexpr
+applyFunctionArguments args sexpr@(SomeExpr (expr :: Expr a))
+    | Just (Refl :: a :~: DynamicType) <- eqT
+    , ExprTypeForall qvar itype <- someExprType sexpr
+    = do
+        tvar <- newTypeVar
+        case renameVarInType qvar tvar itype of
+            ExprTypeFunction (ExprTypeArguments args') res' -> do
+                ( used, ( _, unexpectedArgs ) ) <- unifyArguments args' args
+                unexpectedArguments unexpectedArgs
+                t <- fromMaybe (ExprTypeVar tvar) . M.lookup tvar <$> gets testTypeUnif
+                resolveKnownTypeVars res' >>= \case
+                    res''@(ExprTypePrim (Proxy :: Proxy r)) ->
+                        return $ SomeExpr (ArgsApp used (ExposeFunType args' (TypeApp res'' t expr) :: Expr (FunctionType r)))
+                    r ->
+                        return $ SomeExpr (ArgsApp used (ExposeFunType args' (TypeApp r t expr) :: Expr (FunctionType DynamicType)))
+            _ -> do
+                unexpectedArguments args
+                return sexpr
+
+    | otherwise
+    = case someExprType sexpr of
+        ExprTypeFunction (ExprTypeArguments args') res' -> do
+            ( used, ( _, unexpectedArgs ) ) <- unifyArguments args' args
+            unexpectedArguments unexpectedArgs
+            resolveKnownTypeVars res' >>= \case
+                ExprTypePrim (Proxy :: Proxy r)
+                    | Just (Refl :: a :~: FunctionType r) <- eqT
+                    -> return $ SomeExpr (ArgsApp used expr)
+                _
+                    | Just (Refl :: a :~: FunctionType DynamicType) <- eqT
+                    -> return $ SomeExpr (ArgsApp used expr)
+                _ ->
+                    error $ "expecting function type, got: " <> show (typeRep expr)
+        _ -> do
+            unexpectedArguments args
+            return sexpr
+  where
+    unexpectedArguments (FunctionArguments amap) = do
+        forM_ (M.toAscList amap) $ \( kw, ( poff, _ ) ) ->
+            registerParseError $ FancyError poff $ S.singleton $ ErrorFail $ T.unpack $
+                case kw of
+                    Just (ArgumentKeyword tkw) -> "unexpected parameter with keyword ‘" <> tkw <> "’"
+                    Nothing                    -> "unexpected parameter"
diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index f4f5b61..c0a85e5 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -136,8 +136,8 @@ exprStatement = do
         blockOf indent $ do
             coff <- stateOffset <$> getParserState
             sline <- getSourceLine
-            args <- functionArguments (checkFunctionArguments (exprArgs fun)) (someExpr FunctionTerm) literal (\poff -> lookupVarExpr poff sline . VarName)
-            let fun' = ArgsApp args fun
+            args <- functionArguments (\poff _ e -> return ( poff, e )) (someExpr FunctionTerm) literal (\poff -> lookupVarExpr poff sline . VarName)
+            SomeExpr fun' <- applyFunctionArguments args (SomeExpr fun)
             choice
                 [ continuePartial coff indent fun'
                 , unifyExpr coff Proxy fun'