Type variables and simple unification

3 files changed, 153 insertions, 43 deletions

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 <> "'"