Implemented Gen for ForeignTyRefs

Author: bladyjoker

lambda-buffers-compiler/test/Test/LambdaBuffers/Compiler/Gen.hs

@@ -2,8 +2,8 @@ module Test.LambdaBuffers.Compiler.Gen (genCompilerInput) where
 
 import Control.Lens ((&), (.~), (^.))
 import Control.Monad (foldM)
+import Data.Either (isRight)
 import Data.Foldable (Foldable (toList))
-import Data.List qualified as List
 import Data.List.NonEmpty (NonEmpty ((:|)))
 import Data.Map (Map)
 import Data.Map qualified as Map
@@ -20,14 +20,14 @@ import Data.Text qualified as Text
 import Data.Traversable (for)
 import GHC.Enum qualified as Int
 import Proto.Compiler (ClassName, CompilerInput, ConstrName, Kind, Kind'KindRef (Kind'KIND_REF_TYPE), Module, ModuleName, ModuleNamePart, SourceInfo, Sum, Sum'Constructor, Ty, TyAbs, TyArg, TyBody, TyDef, TyName, VarName)
-import Proto.Compiler_Fields (argKind, argName, column, constrName, constructors, fields, file, kindArrow, kindRef, left, localTyRef, moduleName, modules, name, ntuple, parts, posFrom, posTo, right, row, sourceInfo, tyAbs, tyApp, tyArgs, tyBody, tyFunc, tyName, tyRef, tyVar, typeDefs, varName)
+import Proto.Compiler_Fields (argKind, argName, column, constrName, constructors, fields, file, foreignTyRef, kindArrow, kindRef, left, localTyRef, moduleName, modules, name, ntuple, parts, posFrom, posTo, right, row, sourceInfo, tyAbs, tyApp, tyArgs, tyBody, tyFunc, tyName, tyRef, tyVar, typeDefs, varName)
 import Proto.Compiler_Fields qualified as P
 import Test.QuickCheck qualified as QC (arbitraryPrintableChar)
 import Test.QuickCheck.Gen qualified as QC
 
 -- | Upper bound on various generators
 limit :: Int
-limit = 10
+limit = 5
 
 -- | Names
 genAlphaNum :: QC.Gen Char
@@ -88,27 +88,33 @@ genTyArg vn = do
       & argName .~ vn
       & argKind .~ starKind -- TODO(bladyjoker): Gen arbitrary kinds.
 
-genSum :: Map TyName TyDef -> Set TyArg -> NESet ConstrName -> QC.Gen Sum
+genSum :: TyDefs -> Set TyArg -> NESet ConstrName -> QC.Gen Sum
 genSum tydefs args ctorNs = do
   let (ctorN :| ctorNs') = NESet.toList ctorNs
   ctorNs'' <- QC.sublistOf (toList ctorNs')
   ctors <- for (ctorN :| ctorNs'') (genConstructor tydefs args)
   return $ defMessage & constructors .~ toList ctors
 
-genTy :: Kind -> Map TyName TyDef -> Set TyArg -> QC.Gen Ty
+genTy :: Kind -> TyDefs -> Set TyArg -> QC.Gen Ty
 genTy kind tydefs tyargs =
   QC.oneof $
     NESet.withNonEmpty [] (genTyVar kind) tyargs
-      <> NEMap.withNonEmpty [] (genTyRef kind) tydefs
+      <> genTyRef kind tydefs
       <> genTyApp kind tydefs tyargs
 
-genTyRef :: Kind -> NEMap TyName TyDef -> [QC.Gen Ty]
-genTyRef kind tydefs = case [tyd | tyd <- toList tydefs, kindOf (tyd ^. tyAbs) == kind] of
+genTyRef :: Kind -> TyDefs -> [QC.Gen Ty]
+genTyRef kind tydefs = case [tyd | tyd <- Map.toList tydefs, kindOf (snd tyd ^. tyAbs) == kind] of
   [] -> []
   tyds ->
     [ do
         tydef <- QC.elements tyds
-        return $ defMessage & tyRef . localTyRef . tyName .~ (tydef ^. tyName)
+        case fst tydef of
+          Left (mn, tyn) ->
+            return $
+              defMessage
+                & tyRef . foreignTyRef . moduleName .~ mn
+                & tyRef . foreignTyRef . tyName .~ tyn
+          Right tyn -> return $ defMessage & tyRef . localTyRef . tyName .~ tyn
     ]
 
 genTyVar :: Kind -> NESet TyArg -> [QC.Gen Ty]
@@ -120,7 +126,7 @@ genTyVar kind args = case [tyarg | tyarg <- toList args, tyarg ^. argKind == kin
         return $ defMessage & tyVar . varName .~ (tyarg ^. argName)
     ]
 
-genTyApp :: Kind -> Map TyName TyDef -> Set TyArg -> [QC.Gen Ty]
+genTyApp :: Kind -> TyDefs -> Set TyArg -> [QC.Gen Ty]
 genTyApp kind tydefs args =
   let kindFunc =
         defMessage
@@ -138,31 +144,31 @@ genTyApp kind tydefs args =
                   & tyApp . tyArgs .~ [tyarg] -- TODO(bladyjoker): Generate list arguments
           ]
 
-genConstructor :: Map TyName TyDef -> Set TyArg -> ConstrName -> QC.Gen Sum'Constructor
+genConstructor :: TyDefs -> Set TyArg -> ConstrName -> QC.Gen Sum'Constructor
 genConstructor tydefs args cn = do
   tys <- QC.chooseInt (0, limit) >>= vecOf (genTy starKind tydefs args)
   return $
     defMessage
       & constrName .~ cn
       & P.product . ntuple . fields .~ tys
 
-genTyBodySum :: Map TyName TyDef -> Set TyArg -> NESet ConstrName -> QC.Gen TyBody
+genTyBodySum :: TyDefs -> Set TyArg -> NESet ConstrName -> QC.Gen TyBody
 genTyBodySum tydefs args ctors = do
   b <- genSum tydefs args ctors
   return $ defMessage & P.sum .~ b
 
 genTyBodyOpaque :: QC.Gen TyBody
 genTyBodyOpaque = return $ defMessage & P.opaque .~ defMessage
 
-genTyBody :: Map TyName TyDef -> Set TyArg -> NESet ConstrName -> QC.Gen TyBody
+genTyBody :: TyDefs -> Set TyArg -> NESet ConstrName -> QC.Gen TyBody
 genTyBody tydefs args ctorNs =
   QC.oneof $
     [ genTyBodyOpaque
     ]
       -- Gen TyBody'Sum only if there's some TyDefs and TyArgs available
       <> [genTyBodySum tydefs args ctorNs | not (tydefs == mempty && args == mempty)]
 
-genTyAbs :: Map TyName TyDef -> NESet ConstrName -> QC.Gen TyAbs
+genTyAbs :: TyDefs -> NESet ConstrName -> QC.Gen TyAbs
 genTyAbs tydefs ctorNs = do
   vns <-
     if tydefs == mempty
@@ -175,37 +181,54 @@ genTyAbs tydefs ctorNs = do
       & tyArgs .~ toList args
       & tyBody .~ body
 
-genTyDef :: Map TyName TyDef -> TyName -> NESet ConstrName -> QC.Gen TyDef
+type TyDefs = Map (Either (ModuleName, TyName) TyName) TyDef
+
+genTyDef :: TyDefs -> TyName -> NESet ConstrName -> QC.Gen TyDef
 genTyDef tydefs tyn ctors = do
   tyabs <- genTyAbs tydefs ctors
-  return $
+  withSourceInfo $
     defMessage
       & tyName .~ tyn
       & tyAbs .~ tyabs
 
-genModule :: ModuleName -> QC.Gen Module
-genModule mn = do
-  tyNs <- QC.chooseInt (1, limit) >>= nesetOf genTyName
+genModule :: Map ModuleName Module -> ModuleName -> QC.Gen Module
+genModule availableMods mn = do
+  tyNs <- QC.chooseInt (0, limit) >>= nesetOf genTyName
   ctorNs <- QC.chooseInt (length tyNs, length tyNs * limit) >>= nesetOf genConstrName
   tyNsWithCtorNs <- Map.map NESet.fromList <$> distribute (toList ctorNs) (NESet.toSet tyNs)
+  let foreignTyDefs = collectTyDefs availableMods
   tydefs <-
     foldM
-      ( \availableTyDefs (tyN, ctorNs') -> do
-          tydef <- genTyDef availableTyDefs tyN ctorNs'
-          return $ Map.insert tyN tydef availableTyDefs
+      ( \allTyDefs (tyN, ctorNs') -> do
+          tydef <- genTyDef allTyDefs tyN ctorNs'
+          return $ Map.insert (Right tyN) tydef allTyDefs
       )
-      mempty
+      foreignTyDefs
       (Map.toList tyNsWithCtorNs)
   return $
     defMessage
       & moduleName .~ mn
-      & typeDefs .~ toList tydefs
+      & typeDefs .~ ([tydef | (n, tydef) <- Map.toList tydefs, isRight n])
+  where
+    collectTyDefs :: Map ModuleName Module -> TyDefs
+    collectTyDefs mods =
+      snd
+        <$> indexBy
+          (\(m, tydef) -> Left (m ^. moduleName, tydef ^. tyName))
+          [(m, tydef) | m <- toList mods, tydef <- m ^. typeDefs]
 
 genCompilerInput :: QC.Gen CompilerInput
 genCompilerInput = do
-  mns <- QC.chooseInt (1, limit) >>= vecOf genModuleName
-  ms <- for (List.nub mns) genModule
-  return $ defMessage & modules .~ ms
+  mns <- QC.chooseInt (0, limit) >>= setOf genModuleName
+  ms <-
+    foldM
+      ( \availableMods mn -> do
+          m <- genModule availableMods mn
+          return $ Map.insert mn m availableMods
+      )
+      mempty
+      (toList mns)
+  return $ defMessage & modules .~ toList ms
 
 -- | Utils
 
@@ -246,8 +269,12 @@ partition xs = go xs []
 _indexBy :: Ord k => (a -> k) -> NonEmpty a -> NEMap k a
 _indexBy keyF (x :| xs) = foldl (\t x' -> NEMap.insert (keyF x') x' t) (NEMap.singleton (keyF x) x) xs
 
-_withSourceInfo :: HasField a "sourceInfo" SourceInfo => a -> QC.Gen a
-_withSourceInfo msg = do
+-- | Index a list given a key function.
+indexBy :: Foldable t => Ord k => (a -> k) -> t a -> Map k a
+indexBy keyF = foldl (\t x -> Map.insert (keyF x) x t) mempty
+
+withSourceInfo :: HasField a "sourceInfo" SourceInfo => a -> QC.Gen a
+withSourceInfo msg = do
   f <- Text.pack <$> vecOf QC.arbitraryPrintableChar 10
   i <- QC.chooseInt (0, Int.maxBound)
   let pos =