Add Imports to Module type, parser, and prettyprinter

cardano-foundation · Aug 29, 2023 · 15d014b · 15d014b
1 parent a18bb93
commit 15d014b
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 6 deletions.
diff --git a/lib/fine-types/src/Language/FineTypes/Module.hs b/lib/fine-types/src/Language/FineTypes/Module.hs
@@ -1,6 +1,9 @@
 -- | A 'Module' is a collection of 'Typ' definitions.
 module Language.FineTypes.Module
-    ( Module (..)
+    ( ModuleName
+    , Module (..)
+    , Import (..)
+    , Imports
     , Declarations
     , collectNotInScope
     , resolveVars
@@ -27,16 +30,25 @@ import qualified Data.Set as Set
 {-----------------------------------------------------------------------------
     Module type
 ------------------------------------------------------------------------------}
+type ModuleName = String
 
 -- | A 'Module' is a collection of 'Typ' definitions and documentation.
 data Module = Module
-    { moduleName :: String
+    { moduleName :: ModuleName
+    , moduleImports :: Imports
     , moduleDeclarations :: Declarations
     }
     deriving (Eq, Show)
 
 type Declarations = Map TypName Typ
 
+type Imports = Map ModuleName Import
+
+data Import
+    = AllNames
+    | SomeNames (Set TypName)
+    deriving (Eq, Show)
+
 {-----------------------------------------------------------------------------
     Module functions
 ------------------------------------------------------------------------------}

diff --git a/lib/fine-types/src/Language/FineTypes/Module/Gen.hs b/lib/fine-types/src/Language/FineTypes/Module/Gen.hs
@@ -5,7 +5,13 @@ module Language.FineTypes.Module.Gen where
 
 import Prelude
 
-import Language.FineTypes.Module (Declarations, Module (..))
+import Language.FineTypes.Module
+    ( Declarations
+    , Module (..)
+    , ModuleName
+    , Import (..)
+    , Imports
+    )
 import Language.FineTypes.Typ (Typ, TypName)
 import Language.FineTypes.Typ.Gen
     ( Mode (Complete)
@@ -17,16 +23,31 @@ import Language.FineTypes.Typ.Gen
 import Test.QuickCheck
     ( Gen
     , listOf
+    , oneof
     )
 
 import qualified Data.Map as Map
+import qualified Data.Set as Set
 
 genModule :: Gen Module
 genModule = do
     moduleName <- genModuleName
+    moduleImports <- genImports
     moduleDeclarations <- genDeclarations
     pure Module{..}
 
+genImports :: Gen Imports
+genImports = Map.fromList <$> listOf genImport
+
+genImport :: Gen (ModuleName, Import)
+genImport = do
+    moduleName <- genModuleName
+    imports <- oneof
+        [ pure AllNames
+        , SomeNames . Set.fromList <$> listOf genTypName
+        ]
+    pure (moduleName, imports)
+
 genDeclarations :: Gen Declarations
 genDeclarations = Map.fromList <$> listOf genDeclaration
 
@@ -47,6 +68,8 @@ shrinkModule m = do
     moduleDeclarations <- shrinkDeclarations (moduleDeclarations m)
     pure $ m{moduleDeclarations}
 
+-- TODO: Shrink import list
+
 shrinkDeclarations :: Declarations -> [Declarations]
 shrinkDeclarations xs = do
     (k, v) <- Map.toList xs

diff --git a/lib/fine-types/src/Language/FineTypes/Module/PrettyPrinter.hs b/lib/fine-types/src/Language/FineTypes/Module/PrettyPrinter.hs
@@ -10,7 +10,13 @@ module Language.FineTypes.Module.PrettyPrinter
 
 import Prelude
 
-import Language.FineTypes.Module (Declarations, Module (..))
+import Language.FineTypes.Module
+    ( Declarations
+    , Import (..)
+    , Imports
+    , Module (..)
+    , ModuleName
+    )
 import Language.FineTypes.Typ
     ( ConstructorName
     , FieldName
@@ -23,19 +29,22 @@ import Language.FineTypes.Typ
 import Prettyprinter
     ( Doc
     , Pretty (pretty)
+    , comma
     , concatWith
     , defaultLayoutOptions
     , encloseSep
     , indent
     , layoutPretty
     , line
+    , nest
     , vsep
     , (<+>)
     )
 import Prettyprinter.Render.Text (renderStrict)
 
 import Data.Functor ((<&>))
 import qualified Data.Map as Map
+import qualified Data.Set as Set
 import qualified Data.Text as T
 
 prettyText :: T.Text -> Doc ann
@@ -66,13 +75,29 @@ prettyPrintModule =
 
 -- | Pretty print a 'Module'.
 prettyModule :: Module -> Doc ann
-prettyModule Module{moduleName, moduleDeclarations} =
+prettyModule Module{moduleName, moduleDeclarations, moduleImports} =
     vsep
         [ prettyText "module" <+> pretty moduleName <+> prettyText "where"
         , prettyText ""
+        , prettyImports moduleImports
         , prettyDeclarations moduleDeclarations
         ]
 
+prettyImports :: Imports -> Doc ann
+prettyImports = vsep . map prettyImport . Map.toList
+
+prettyImport :: (ModuleName, Import) -> Doc ann
+prettyImport (name, scope) =
+    prettyText "import" <+> case scope of
+        AllNames ->
+            pretty name <> prettyText ";"
+        SomeNames names ->
+            pretty name
+                <> nest 4 (prettyList $ Set.toList names)
+                <> prettyText ";"
+  where
+    prettyList = encloseSep "(" ")" comma . map pretty
+
 prettyDeclarations :: Declarations -> Doc ann
 prettyDeclarations = vsep . map prettyDeclaration . Map.toList
 

diff --git a/lib/fine-types/src/Language/FineTypes/Parser.hs b/lib/fine-types/src/Language/FineTypes/Parser.hs
@@ -12,6 +12,9 @@ import Data.Void
 import Language.FineTypes.Module
     ( Declarations
     , Module (Module)
+    , ModuleName
+    , Import (..)
+    , Imports
     )
 import Language.FineTypes.Typ
     ( ConstructorName
@@ -28,6 +31,7 @@ import Text.Megaparsec
     , between
     , endBy
     , many
+    , optional
     , parse
     , parseMaybe
     , satisfy
@@ -39,6 +43,7 @@ import Text.Megaparsec
 
 import qualified Control.Monad.Combinators.Expr as Parser.Expr
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 import qualified Text.Megaparsec.Char as Parser.Char
 import qualified Text.Megaparsec.Char.Lexer as L
 
@@ -74,8 +79,23 @@ module' =
         <$ symbol "module"
         <*> moduleName
         <* symbol "where"
+        <*> imports
         <*> declarations
 
+imports :: Parser Imports
+imports = mconcat <$> (import' `endBy` symbol ";")
+
+import' :: Parser Imports
+import' =
+    Map.singleton
+        <$ symbol "import"
+        <*> moduleName
+        <*> (toImport <$> importedNames)
+  where
+    importedNames = optional $ parens (typName `sepBy` symbol ",")
+    toImport Nothing = AllNames
+    toImport (Just xs) = SomeNames $ Set.fromList xs
+
 declarations :: Parser Declarations
 declarations = mconcat <$> (declaration `endBy` symbol ";")
 
@@ -168,7 +188,7 @@ space = L.space Parser.Char.space1 lineComment blockComment
 symbol :: String -> Parser String
 symbol = L.symbol space
 
-moduleName :: Parser String
+moduleName :: Parser ModuleName
 moduleName = typName
 
 typName :: Parser TypName