Type variables and simple unification

6 files changed, 183 insertions, 49 deletions

diff --git a/erebos-tester.cabal b/erebos-tester.cabal
index d894395..c3a49e6 100644
--- a/erebos-tester.cabal
+++ b/erebos-tester.cabal
@@ -103,6 +103,7 @@ executable erebos-tester-core
                        ImportQualifiedPost
                        LambdaCase
                        MultiParamTypeClasses
+                       MultiWayIf
                        OverloadedStrings
                        RankNTypes
                        RecordWildCards
diff --git a/src/Parser.hs b/src/Parser.hs
index 3c43a69..6d6809b 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -8,6 +8,7 @@ import Control.Monad
 import Control.Monad.State
 import Control.Monad.Writer
 
+import Data.Map qualified as M
 import Data.Maybe
 import Data.Set qualified as S
 import Data.Text qualified as T
@@ -66,8 +67,10 @@ parseTestFile path = do
                 [ map (fmap someVarValueType) builtins
                 ]
             , testContext = SomeExpr RootNetwork
+            , testNextTypeVar = 0
+            , testTypeUnif = M.empty
             }
-        (res, _) = flip evalState initState $ runWriterT $ runParserT (parseTestModule absPath) path content
+        (res, _) = runWriter $ flip (flip runParserT path) content $ flip evalStateT initState $ parseTestModule absPath
 
     case res of
          Left err -> putStr (errorBundlePretty err) >> exitFailure
diff --git a/src/Parser/Core.hs b/src/Parser/Core.hs
index 2a74d3d..2a2fc89 100644
--- a/src/Parser/Core.hs
+++ b/src/Parser/Core.hs
@@ -4,8 +4,13 @@ import Control.Monad
 import Control.Monad.State
 import Control.Monad.Writer
 
-import Data.Text (Text)
-import qualified Data.Text.Lazy as TL
+import Data.Map (Map)
+import Data.Map qualified as M
+import Data.Maybe
+import Data.Set qualified as S
+import Data.Text qualified as T
+import Data.Text.Lazy qualified as TL
+import Data.Typeable
 import Data.Void
 
 import Text.Megaparsec hiding (State)
@@ -15,23 +20,127 @@ import qualified Text.Megaparsec.Char.Lexer as L
 import Network ()
 import Test
 
-type TestParser = ParsecT Void TestStream (WriterT [ Toplevel ] (State TestParserState))
+type TestParser = StateT TestParserState (ParsecT Void TestStream (Writer [ Toplevel ]))
 
 type TestStream = TL.Text
 
+type TestParseError = ParseError TestStream Void
+
 data Toplevel
     = ToplevelTest Test
 
 data TestParserState = TestParserState
-    { testVars :: [(VarName, SomeExprType)]
+    { testVars :: [ ( VarName, SomeExprType ) ]
     , testContext :: SomeExpr
+    , testNextTypeVar :: Int
+    , testTypeUnif :: Map TypeVar SomeExprType
     }
 
-textSomeExprType :: SomeExprType -> Text
-textSomeExprType (SomeExprType p) = textExprType p
+newTypeVar :: TestParser TypeVar
+newTypeVar = do
+    idx <- gets testNextTypeVar
+    modify $ \s -> s { testNextTypeVar = idx + 1 }
+    return $ TypeVar $ T.pack $ 'a' : show idx
+
+lookupVarType :: Int -> VarName -> TestParser SomeExprType
+lookupVarType off name = do
+    gets (lookup name . testVars) >>= \case
+        Nothing -> do
+            registerParseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
+                "variable not in scope: `" <> textVarName name <> "'"
+            vtype <- ExprTypeVar <$> newTypeVar
+            modify $ \s -> s { testVars = ( name, vtype ) : testVars s }
+            return vtype
+        Just t@(ExprTypeVar tvar) -> do
+            gets (fromMaybe t . M.lookup tvar . testTypeUnif)
+        Just x -> return x
+
+lookupVarExpr :: Int -> VarName -> TestParser SomeExpr
+lookupVarExpr off name = do
+    lookupVarType off name >>= \case
+        ExprTypePrim (Proxy :: Proxy a) -> return $ SomeExpr $ (Variable name :: Expr a)
+        ExprTypeVar tvar -> return $ SomeExpr $ DynVariable tvar name
+
+unify :: Int -> SomeExprType -> SomeExprType -> TestParser SomeExprType
+unify _ (ExprTypeVar aname) (ExprTypeVar bname) | aname == bname = do
+    cur <- gets testTypeUnif
+    case M.lookup aname cur of
+        Just a -> return a
+        Nothing -> return (ExprTypeVar aname)
+
+unify off (ExprTypeVar aname) (ExprTypeVar bname) = do
+    cur <- gets testTypeUnif
+    case ( M.lookup aname cur, M.lookup bname cur ) of
+        ( Just a, Just b ) -> do
+            c <- unify off a b
+            modify $ \s -> s { testTypeUnif = M.insert aname c $ M.insert bname c $ cur }
+            return c
+
+        ( Just a, Nothing ) -> do
+            modify $ \s -> s { testTypeUnif = M.insert bname a $ cur }
+            return a
+
+        ( Nothing, Just b ) -> do
+            modify $ \s -> s { testTypeUnif = M.insert aname b $ cur }
+            return b
+
+        ( Nothing, Nothing ) -> do
+            let b = ExprTypeVar bname
+            modify $ \s -> s { testTypeUnif = M.insert aname b $ cur }
+            return b
+
+unify off (ExprTypeVar aname) b = do
+    cur <- gets testTypeUnif
+    case M.lookup aname cur of
+        Just a -> do
+            c <- unify off a b
+            modify $ \s -> s { testTypeUnif = M.insert aname c $ cur }
+            return c
+        Nothing -> do
+            modify $ \s -> s { testTypeUnif = M.insert aname b $ cur }
+            return b
+
+unify off a (ExprTypeVar bname) = do
+    cur <- gets testTypeUnif
+    case M.lookup bname cur of
+        Just b -> do
+            c <- unify off a b
+            modify $ \s -> s { testTypeUnif = M.insert bname c $ cur }
+            return c
+
+        Nothing -> do
+            modify $ \s -> s { testTypeUnif = M.insert bname a $ cur }
+            return a
+
+unify _ res@(ExprTypePrim (Proxy :: Proxy a)) (ExprTypePrim (Proxy :: Proxy b))
+    | Just (Refl :: a :~: b) <- eqT
+    = return res
+
+unify off a b = do
+    parseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
+        "couldn't match expected type `" <> textSomeExprType a <> "' with actual type `" <> textSomeExprType b <> "'"
+
+
+unifyExpr :: forall a b proxy. (ExprType a, ExprType b) => Int -> proxy a -> Expr b -> TestParser (Expr a)
+unifyExpr off pa x = if
+    | Just (Refl :: a :~: b) <- eqT
+    -> return x
+
+    | DynVariable tvar name <- x
+    -> do
+        _ <- unify off (ExprTypePrim (Proxy :: Proxy a)) (ExprTypeVar tvar)
+        return $ Variable name
+
+    | Just (Refl :: DynamicType :~: b) <- eqT
+    , Undefined msg <- x
+    -> do
+        return $ Undefined msg
+
+    | otherwise
+    -> do
+        parseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
+            "couldn't match expected type `" <> textExprType pa <> "' with actual type `" <> textExprType x <> "'"
 
-lookupVarType :: VarName -> TestParser SomeExprType
-lookupVarType name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return =<< gets (lookup name . testVars)
 
 skipLineComment :: TestParser ()
 skipLineComment = L.skipLineComment $ TL.pack "#"
diff --git a/src/Parser/Expr.hs b/src/Parser/Expr.hs
index 8ea3ace..a228ad0 100644
--- a/src/Parser/Expr.hs
+++ b/src/Parser/Expr.hs
@@ -53,15 +53,15 @@ addVarName off (TypedVarName name) = do
         Just _ -> registerParseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
             T.pack "variable '" <> textVarName name <> T.pack "' already exists"
         Nothing -> return ()
-    modify $ \s -> s { testVars = (name, SomeExprType @a Proxy) : testVars s }
+    modify $ \s -> s { testVars = ( name, ExprTypePrim @a Proxy ) : testVars s }
 
 someExpansion :: TestParser SomeExpr
 someExpansion = do
     void $ char '$'
     choice
-        [do name <- VarName . TL.toStrict <$> takeWhile1P Nothing (\x -> isAlphaNum x || x == '_')
-            SomeExprType (_ :: Proxy a) <- lookupVarType name
-            return $ SomeExpr $ Variable @a name
+        [do off <- stateOffset <$> getParserState
+            name <- VarName . TL.toStrict <$> takeWhile1P Nothing (\x -> isAlphaNum x || x == '_')
+            lookupVarExpr off name
         , between (char '{') (char '}') someExpr
         ]
 
@@ -186,20 +186,20 @@ data SomeUnOp = forall a b. (ExprType a, ExprType b) => SomeUnOp (a -> b)
 
 applyUnOp :: forall a b sa.
     (ExprType a, ExprType b, ExprType sa) =>
-    (a -> b) -> Expr sa -> Maybe (Expr b)
-applyUnOp op x = do
-    Refl :: a :~: sa <- eqT
-    return $ op <$> x
+    Int -> (a -> b) -> Expr sa -> TestParser (Expr b)
+applyUnOp off op x = do
+    x' <- unifyExpr off (Proxy @a) x
+    return $ op <$> x'
 
 data SomeBinOp = forall a b c. (ExprType a, ExprType b, ExprType c) => SomeBinOp (a -> b -> c)
 
 applyBinOp :: forall a b c sa sb.
     (ExprType a, ExprType b, ExprType c, ExprType sa, ExprType sb) =>
-    (a -> b -> c) -> Expr sa -> Expr sb -> Maybe (Expr c)
-applyBinOp op x y = do
-    Refl :: a :~: sa <- eqT
-    Refl :: b :~: sb <- eqT
-    return $ op <$> x <*> y
+    Int -> (a -> b -> c) -> Expr sa -> Expr sb -> TestParser (Expr c)
+applyBinOp off op x y = do
+    x' <- unifyExpr off (Proxy @a) x
+    y' <- unifyExpr off (Proxy @b) y
+    return $ op <$> x' <*> y'
 
 someExpr :: TestParser SomeExpr
 someExpr = join inner <?> "expression"
@@ -251,9 +251,11 @@ someExpr = join inner <?> "expression"
         void $ osymbol name
         return $ \p -> do
             SomeExpr e <- p
-            let err = parseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $ T.concat
+            let err = FancyError off $ S.singleton $ ErrorFail $ T.unpack $ T.concat
                     [T.pack "operator '", T.pack name, T.pack "' not defined for '", textExprType e, T.pack "'"]
-            maybe err return $ listToMaybe $ catMaybes $ map (\(SomeUnOp op) -> SomeExpr <$> applyUnOp op e) ops
+            region (const err) $
+                choice $ map (\(SomeUnOp op) -> SomeExpr <$> applyUnOp off op e) ops
+
 
     binary :: String -> [SomeBinOp] -> Operator TestParser (TestParser SomeExpr)
     binary name = binary' name (undefined :: forall a b. (a -> b -> Void) -> [a] -> [b] -> Integer)
@@ -278,20 +280,22 @@ someExpr = join inner <?> "expression"
             let proxyOf :: proxy a -> Proxy a
                 proxyOf _ = Proxy
 
+            let err = FancyError off $ S.singleton $ ErrorFail $ T.unpack $ T.concat
+                    [T.pack "operator '", T.pack name, T.pack "' not defined for '", textExprType e, T.pack "' and '", textExprType f, T.pack "'"]
+
             let tryop :: forall a b d sa sb.
                     (ExprType a, ExprType b, ExprType d, ExprType sa, ExprType sb) =>
-                    (a -> b -> d) -> Proxy sa -> Proxy sb -> Maybe SomeExpr
-                tryop op pe pf = msum
-                    [ SomeExpr <$> applyBinOp op e f
-                    , do Refl <- eqT' op
-                         ExprListUnpacker _ une <- exprListUnpacker pe
-                         ExprListUnpacker _ unf <- exprListUnpacker pf
+                    (a -> b -> d) -> Proxy sa -> Proxy sb -> TestParser SomeExpr
+                tryop op pe pf = foldl1 (<|>) $
+                    [ SomeExpr <$> applyBinOp off op e f
+                    , do Refl <- maybe (parseError err) return $ eqT' op
+                         ExprListUnpacker _ une <- maybe (parseError err) return $ exprListUnpacker pe
+                         ExprListUnpacker _ unf <- maybe (parseError err) return $ exprListUnpacker pf
                          tryop (listmap op) (une pe) (unf pf)
                     ]
 
-            let err = parseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $ T.concat
-                    [T.pack "operator '", T.pack name, T.pack "' not defined for '", textExprType e, T.pack "' and '", textExprType f, T.pack "'"]
-            maybe err return $ listToMaybe $ catMaybes $ map (\(SomeBinOp op) -> tryop op (proxyOf e) (proxyOf f)) ops
+            region (const err) $
+                foldl1 (<|>) $ map (\(SomeBinOp op) -> tryop op (proxyOf e) (proxyOf f)) ops
 
     recordSelector :: Operator TestParser (TestParser SomeExpr)
     recordSelector = Postfix $ fmap (foldl1 (flip (.))) $ some $ do
@@ -315,16 +319,13 @@ someExpr = join inner <?> "expression"
         ]
 
     variable = label "variable" $ do
+        off <- stateOffset <$> getParserState
         name <- varName
-        SomeExprType (_ :: Proxy a) <- lookupVarType name
-        return $ return $ SomeExpr $ Variable @a name
+        e <- lookupVarExpr off name
+        return $ return e
 
 typedExpr :: forall a. ExprType a => TestParser (Expr a)
 typedExpr = do
     off <- stateOffset <$> getParserState
     SomeExpr e <- someExpr
-    let err = do
-            registerParseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $ T.concat
-                [ T.pack "expected '", textExprType @a Proxy, T.pack "', expression has type '", textExprType e, T.pack "'" ]
-            return $ Undefined "unexpected type"
-    maybe err return $ cast e
+    unifyExpr off Proxy e
diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index b2f3cd6..94a5583 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -259,12 +259,12 @@ testWith = do
     off <- stateOffset <$> getParserState
     ctx@(SomeExpr (_ :: Expr ctxe)) <- someExpr
     let expected =
-            [ SomeExprType @Network Proxy
-            , SomeExprType @Node Proxy
-            , SomeExprType @Process Proxy
+            [ ExprTypePrim @Network Proxy
+            , ExprTypePrim @Node Proxy
+            , ExprTypePrim @Process Proxy
             ]
     notAllowed <- flip allM expected $ \case
-        SomeExprType (Proxy :: Proxy a) | Just (Refl :: ctxe :~: a) <- eqT -> return False
+        ExprTypePrim (Proxy :: Proxy a) | Just (Refl :: ctxe :~: a) <- eqT -> return False
         _ -> return True
     when notAllowed $ registerParseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
         "expected " <> T.intercalate ", " (map (("'"<>) . (<>"'") . textSomeExprType) expected) <> ", expression has type '" <> textExprType @ctxe Proxy <> "'"
diff --git a/src/Test.hs b/src/Test.hs
index ba27153..bb65b81 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -7,7 +7,9 @@ module Test (
 
     MonadEval(..),
     VarName(..), TypedVarName(..), textVarName, unpackVarName,
-    ExprType(..), SomeExpr(..), SomeExprType(..), someExprType,
+    ExprType(..), SomeExpr(..),
+    TypeVar(..), SomeExprType(..), someExprType, textSomeExprType,
+    DynamicType,
     SomeVarValue(..), fromSomeVarValue, textSomeVarValue, someVarValueType,
     RecordSelector(..),
     ExprListUnpacker(..),
@@ -69,7 +71,7 @@ class MonadFail m => MonadEval m where
 
 
 newtype VarName = VarName Text
-    deriving (Eq, Ord)
+    deriving (Eq, Ord, Show)
 
 newtype TypedVarName a = TypedVarName { fromTypedVarName :: VarName }
     deriving (Eq, Ord)
@@ -128,12 +130,27 @@ instance ExprType TestBlock where
     textExprValue _ = "<test block>"
 
 
+data DynamicType
+
+instance ExprType DynamicType where
+    textExprType _ = "ambiguous type"
+    textExprValue _ = "<dynamic type>"
+
 data SomeExpr = forall a. ExprType a => SomeExpr (Expr a)
 
-data SomeExprType = forall a. ExprType a => SomeExprType (Proxy a)
+newtype TypeVar = TypeVar Text
+    deriving (Eq, Ord)
+
+data SomeExprType
+    = forall a. ExprType a => ExprTypePrim (Proxy a)
+    | ExprTypeVar TypeVar
 
 someExprType :: SomeExpr -> SomeExprType
-someExprType (SomeExpr (_ :: Expr a)) = SomeExprType (Proxy @a)
+someExprType (SomeExpr (_ :: Expr a)) = ExprTypePrim (Proxy @a)
+
+textSomeExprType :: SomeExprType -> Text
+textSomeExprType (ExprTypePrim p) = textExprType p
+textSomeExprType (ExprTypeVar (TypeVar name)) = name
 
 
 data SomeVarValue = forall a. ExprType a => SomeVarValue a
@@ -146,7 +163,7 @@ textSomeVarValue :: SomeVarValue -> Text
 textSomeVarValue (SomeVarValue value) = textExprValue value
 
 someVarValueType :: SomeVarValue -> SomeExprType
-someVarValueType (SomeVarValue (_ :: a)) = SomeExprType (Proxy @a)
+someVarValueType (SomeVarValue (_ :: a)) = ExprTypePrim (Proxy @a)
 
 
 data RecordSelector a = forall b. ExprType b => RecordSelector (a -> b)
@@ -158,6 +175,7 @@ data ExprEnumerator a = ExprEnumerator (a -> a -> [a]) (a -> a -> a -> [a])
 
 data Expr a where
     Variable :: ExprType a => VarName -> Expr a
+    DynVariable :: TypeVar -> VarName -> Expr DynamicType
     Pure :: a -> Expr a
     App :: AppAnnotation b -> Expr (a -> b) -> Expr a -> Expr b
     Concat :: [Expr Text] -> Expr Text
@@ -177,6 +195,7 @@ instance Applicative Expr where
 
 eval :: MonadEval m => Expr a -> m a
 eval (Variable name) = fromSomeVarValue name =<< lookupVar name
+eval (DynVariable _ _) = fail "ambiguous type"
 eval (Pure value) = return value
 eval (App _ f x) = eval f <*> eval x
 eval (Concat xs) = T.concat <$> mapM eval xs
@@ -193,6 +212,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
   where
     helper :: forall b. Expr b -> m [((VarName, [Text]), SomeVarValue)]
     helper (Variable var) = (:[]) . ((var, []),) <$> lookupVar var
+    helper (DynVariable _ var) = (:[]) . ((var, []),) <$> lookupVar var
     helper (Pure _) = return []
     helper e@(App (AnnRecord sel) _ x)
         | Just (var, sels) <- gatherSelectors x