Module import parsing and type check

5 files changed, 114 insertions, 33 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 01bb766..73d8c02 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -148,8 +148,8 @@ main = do
         Nothing -> queryTerminal (Fd 1)
     out <- startOutput (optVerbose opts) useColor
 
-    tests <- forM files $ \(path, mbTestName) -> do
-        Module {..} <- parseTestFile path
+    modules <- parseTestFiles $ map fst files
+    tests <- forM (zip modules $ map snd files) $ \( Module {..}, mbTestName ) -> do
         return $ map ( , moduleDefinitions ) $ case mbTestName of
             Nothing -> moduleTests
             Just name -> filter ((==name) . testName) moduleTests
diff --git a/src/Parser.hs b/src/Parser.hs
index 00f6f3e..323f2cf 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -1,12 +1,13 @@
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Parser (
-    parseTestFile,
+    parseTestFiles,
 ) where
 
 import Control.Monad
 import Control.Monad.State
 
+import Data.IORef
 import Data.Map qualified as M
 import Data.Maybe
 import Data.Proxy
@@ -23,6 +24,7 @@ import Text.Megaparsec.Char.Lexer qualified as L
 import System.Directory
 import System.Exit
 import System.FilePath
+import System.IO.Error
 
 import Network
 import Parser.Core
@@ -103,50 +105,87 @@ parseExport = label "export declaration" $ toplevel id $ do
         return [ ToplevelDefinition def, ToplevelExport name ]
       , do
         names <- listOf varName
-        void eol
+        eol >> scn
         return $ map ToplevelExport names
       ]
 
+parseImport :: TestParser [ Toplevel ]
+parseImport = label "import declaration" $ toplevel (\() -> []) $ do
+    wsymbol "import"
+    name <- parseModuleName
+    importedModule <- getOrParseModule name
+    let importedDefs = filter ((`elem` moduleExports importedModule) . fst) (moduleDefinitions importedModule)
+    modify $ \s -> s { testVars = map (fmap someExprType) importedDefs ++ testVars s }
+    eol >> scn
+
 parseTestModule :: FilePath -> TestParser Module
 parseTestModule absPath = do
     moduleName <- choice
         [ label "module declaration" $ do
             wsymbol "module"
             off <- stateOffset <$> getParserState
-            x <- identifier
-            name <- (x:) <$> many (symbol "." >> identifier)
-            when (or (zipWith (/=) (reverse name) (reverse $ map T.pack $ splitDirectories $ dropExtension $ absPath))) $ do
+            name@(ModuleName tname) <- parseModuleName
+            when (or (zipWith (/=) (reverse tname) (reverse $ map T.pack $ splitDirectories $ dropExtension $ absPath))) $ do
                 registerParseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
                     "module name does not match file path"
             eol >> scn
             return name
         , do
-            return $ [ T.pack $ takeBaseName absPath ]
+            return $ ModuleName [ T.pack $ takeBaseName absPath ]
         ]
+    modify $ \s -> s { testCurrentModuleName = moduleName }
     toplevels <- fmap concat $ many $ choice
         [ (: []) <$> parseTestDefinition
         , (: []) <$> toplevel ToplevelDefinition parseDefinition
         , parseExport
+        , parseImport
         ]
     let moduleTests = catMaybes $ map (\case ToplevelTest x -> Just x; _ -> Nothing) toplevels
         moduleDefinitions = catMaybes $ map (\case ToplevelDefinition x -> Just x; _ -> Nothing) toplevels
+        moduleExports = catMaybes $ map (\case ToplevelExport x -> Just x; _ -> Nothing) toplevels
     eof
     return Module {..}
 
-parseTestFile :: FilePath -> IO Module
-parseTestFile path = do
-    content <- TL.readFile path
-    absPath <- makeAbsolute path
-    let initState = TestParserState
-            { testVars = concat
-                [ map (fmap someVarValueType) builtins
-                ]
-            , testContext = SomeExpr (Undefined "void" :: Expr Void)
-            , testNextTypeVar = 0
-            , testTypeUnif = M.empty
-            }
-        res = runTestParser path content initState $ parseTestModule absPath
+parseTestFiles :: [ FilePath ] -> IO [ Module ]
+parseTestFiles paths = do
+    parsedModules <- newIORef []
+    reverse <$> foldM (go parsedModules) [] paths
+  where
+    go parsedModules res path = do
+        let moduleName = error "current module name should be set at the beginning of parseTestModule"
+        parseTestFile parsedModules moduleName path >>= \case
+            Left err -> do
+                putStr (showErrorComponent err)
+                exitFailure
+            Right cur -> do
+                return $ cur : res
 
-    case res of
-         Left err -> putStr (errorBundlePretty err) >> exitFailure
-         Right testModule -> return testModule
+parseTestFile :: IORef [ ( FilePath, Module ) ] -> ModuleName -> FilePath -> IO (Either CustomTestError Module)
+parseTestFile parsedModules moduleName path = do
+    absPath <- makeAbsolute path
+    (lookup absPath <$> readIORef parsedModules) >>= \case
+        Just found -> return $ Right found
+        Nothing -> do
+            let initState = TestParserState
+                    { testVars = concat
+                        [ map (fmap someVarValueType) builtins
+                        ]
+                    , testContext = SomeExpr (Undefined "void" :: Expr Void)
+                    , testNextTypeVar = 0
+                    , testTypeUnif = M.empty
+                    , testCurrentModuleName = moduleName
+                    , testParseModule = \(ModuleName current) mname@(ModuleName imported) -> do
+                        let projectRoot = iterate takeDirectory absPath !! length current
+                        parseTestFile parsedModules mname $ projectRoot </> foldr (</>) "" (map T.unpack imported) <.> takeExtension absPath
+                    }
+            mbContent <- (Just <$> TL.readFile path) `catchIOError` \e ->
+                if isDoesNotExistError e then return Nothing else ioError e
+            case mbContent of
+                Just content -> do
+                    runTestParser path content initState (parseTestModule absPath) >>= \case
+                        Left bundle -> do
+                            return $ Left $ ImportModuleError bundle
+                        Right testModule -> do
+                            modifyIORef parsedModules (( absPath, testModule ) : )
+                            return $ Right testModule
+                Nothing -> return $ Left $ ModuleNotFound moduleName
diff --git a/src/Parser/Core.hs b/src/Parser/Core.hs
index 2b8837a..f964291 100644
--- a/src/Parser/Core.hs
+++ b/src/Parser/Core.hs
@@ -2,7 +2,6 @@ module Parser.Core where
 
 import Control.Applicative
 import Control.Monad
-import Control.Monad.Identity
 import Control.Monad.State
 
 import Data.Map (Map)
@@ -12,7 +11,6 @@ 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)
 import Text.Megaparsec.Char
@@ -21,32 +19,54 @@ import qualified Text.Megaparsec.Char.Lexer as L
 import Network ()
 import Test
 
-newtype TestParser a = TestParser (StateT TestParserState (ParsecT Void TestStream Identity) a)
+newtype TestParser a = TestParser (StateT TestParserState (ParsecT CustomTestError TestStream IO) a)
     deriving
         ( Functor, Applicative, Alternative, Monad
         , MonadState TestParserState
         , MonadPlus
         , MonadFail
-        , MonadParsec Void TestStream
+        , MonadParsec CustomTestError TestStream
         )
 
 type TestStream = TL.Text
 
-type TestParseError = ParseError TestStream Void
+type TestParseError = ParseError TestStream CustomTestError
 
-runTestParser :: String -> TestStream -> TestParserState -> TestParser a -> Either (ParseErrorBundle TestStream Void) a
-runTestParser path content initState (TestParser parser) = runIdentity . flip (flip runParserT path) content . flip evalStateT initState $ parser
+data CustomTestError
+    = ModuleNotFound ModuleName
+    | ImportModuleError (ParseErrorBundle TestStream CustomTestError)
+    deriving (Eq)
+
+instance Ord CustomTestError where
+    compare (ModuleNotFound a) (ModuleNotFound b) = compare a b
+    compare (ModuleNotFound _) _                  = LT
+    compare _                  (ModuleNotFound _) = GT
+
+    -- Ord instance is required to store errors in Set, but there shouldn't be
+    -- two ImportModuleErrors at the same possition, so "dummy" comparison
+    -- should be ok.
+    compare (ImportModuleError _) (ImportModuleError _) = EQ
+
+instance ShowErrorComponent CustomTestError where
+    showErrorComponent (ModuleNotFound name) = "module `" <> T.unpack (textModuleName name) <> "' not found"
+    showErrorComponent (ImportModuleError bundle) = "error parsing imported module:\n" <> errorBundlePretty bundle
+
+runTestParser :: String -> TestStream -> TestParserState -> TestParser a -> IO (Either (ParseErrorBundle TestStream CustomTestError) a)
+runTestParser path content initState (TestParser parser) = flip (flip runParserT path) content . flip evalStateT initState $ parser
 
 data Toplevel
     = ToplevelTest Test
     | ToplevelDefinition ( VarName, SomeExpr )
     | ToplevelExport VarName
+    | ToplevelImport ( ModuleName, VarName )
 
 data TestParserState = TestParserState
     { testVars :: [ ( VarName, SomeExprType ) ]
     , testContext :: SomeExpr
     , testNextTypeVar :: Int
     , testTypeUnif :: Map TypeVar SomeExprType
+    , testCurrentModuleName :: ModuleName
+    , testParseModule :: ModuleName -> ModuleName -> IO (Either CustomTestError Module)
     }
 
 newTypeVar :: TestParser TypeVar
@@ -231,3 +251,12 @@ getSourceLine = do
         , T.pack ": "
         , TL.toStrict $ TL.takeWhile (/='\n') $ pstateInput pstate
         ]
+
+
+getOrParseModule :: ModuleName -> TestParser Module
+getOrParseModule name = do
+    current <- gets testCurrentModuleName
+    parseModule <- gets testParseModule
+    (TestParser $ lift $ lift $ parseModule current name) >>= \case
+        Right parsed -> return parsed
+        Left err -> customFailure err
diff --git a/src/Parser/Expr.hs b/src/Parser/Expr.hs
index d8d96eb..41790bb 100644
--- a/src/Parser/Expr.hs
+++ b/src/Parser/Expr.hs
@@ -1,5 +1,6 @@
 module Parser.Expr (
     identifier,
+    parseModuleName,
 
     varName,
     newVarName,
@@ -58,6 +59,11 @@ identifier = label "identifier" $ do
             ]
         return ident
 
+parseModuleName :: TestParser ModuleName
+parseModuleName = do
+    x <- identifier
+    ModuleName . (x :) <$> many (symbol "." >> identifier)
+
 varName :: TestParser VarName
 varName = label "variable name" $ VarName <$> identifier
 
diff --git a/src/Test.hs b/src/Test.hs
index 3db7919..01b2d95 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -1,5 +1,5 @@
 module Test (
-    Module(..),
+    Module(..), ModuleName(..), textModuleName,
     Test(..),
     TestStep(..),
     TestBlock(..),
@@ -53,11 +53,18 @@ import {-# SOURCE #-} Process
 import Util
 
 data Module = Module
-    { moduleName :: [ Text ]
+    { moduleName :: ModuleName
     , moduleTests :: [ Test ]
     , moduleDefinitions :: [ ( VarName, SomeExpr ) ]
+    , moduleExports :: [ VarName ]
     }
 
+newtype ModuleName = ModuleName [ Text ]
+    deriving (Eq, Ord)
+
+textModuleName :: ModuleName -> Text
+textModuleName (ModuleName parts) = T.intercalate "." parts
+
 data Test = Test
     { testName :: Text
     , testSteps :: Expr TestBlock