Skip to content

Commit a136f38

Browse files
bladyjokerbladyjoker
committed
Improve robustness of Gens
1 parent 521093f commit a136f38

File tree

2 files changed

+188
-55
lines changed

2 files changed

+188
-55
lines changed

‎lambda-buffers-compiler/lambda-buffers-compiler.cabal‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,7 @@ test-suite tests
157157
, generic-lens
158158
, lambda-buffers-compiler
159159
, lambda-buffers-compiler-pb >=0.1
160+
, nonempty-containers
160161
, proto-lens >=0.7
161162
, QuickCheck >=2.14
162163
, tasty >=1.4

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

Lines changed: 187 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,33 @@
11
module Test.LambdaBuffers.Compiler.Gen (genCompilerInput) where
22

33
import Control.Lens ((&), (.~), (^.))
4+
import Control.Monad (foldM)
45
import Data.Foldable (Foldable (toList))
56
import Data.List qualified as List
67
import Data.List.NonEmpty (NonEmpty ((:|)))
8+
import Data.Map (Map)
9+
import Data.Map qualified as Map
10+
import Data.Map.NonEmpty (NEMap)
11+
import Data.Map.NonEmpty qualified as NEMap
712
import Data.ProtoLens (Message (defMessage))
13+
import Data.ProtoLens.Field (HasField)
14+
import Data.Set (Set)
15+
import Data.Set qualified as Set
16+
import Data.Set.NonEmpty (NESet)
17+
import Data.Set.NonEmpty qualified as NESet
818
import Data.Text (Text)
919
import Data.Text qualified as Text
1020
import Data.Traversable (for)
11-
import Proto.Compiler (ClassName, CompilerInput, ConstrName, Kind'KindRef (Kind'KIND_REF_TYPE), Module, ModuleName, ModuleNamePart, Sum, Sum'Constructor, Ty, TyAbs, TyArg, TyBody, TyDef, TyName, TyRef, VarName)
12-
import Proto.Compiler_Fields (argKind, argName, constrName, constructors, fields, kindRef, moduleName, modules, name, ntuple, parts, tyAbs, tyArgs, tyBody, tyName, tyRef, tyVar, typeDefs, varName)
21+
import GHC.Enum qualified as Int
22+
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)
23+
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)
1324
import Proto.Compiler_Fields qualified as P
14-
import Test.QuickCheck qualified as QC (Gen, chooseEnum, chooseInt, elements, oneof, vectorOf)
15-
16-
vecOf :: forall {a}. QC.Gen a -> Int -> QC.Gen [a]
17-
vecOf = flip QC.vectorOf
25+
import Test.QuickCheck qualified as QC (arbitraryPrintableChar)
26+
import Test.QuickCheck.Gen qualified as QC
1827

28+
-- | Upper bound on various generators
1929
limit :: Int
20-
limit = 4
30+
limit = 10
2131

2232
-- | Names
2333
genAlphaNum :: QC.Gen Char
@@ -60,85 +70,207 @@ genVarName = do
6070
t <- QC.vectorOf 4 (QC.chooseEnum ('a', 'z'))
6171
return $ defMessage & name .~ Text.pack (h : t)
6272

73+
starKind :: Kind
74+
starKind = defMessage & kindRef .~ Kind'KIND_REF_TYPE
75+
76+
kindOf :: TyAbs -> Kind
77+
kindOf tyabs = case tyabs ^. tyArgs of
78+
[] -> starKind
79+
(a : args) ->
80+
defMessage
81+
& kindArrow . left .~ (a ^. argKind)
82+
& kindArrow . right .~ kindOf (tyabs & tyArgs .~ args)
83+
6384
genTyArg :: VarName -> QC.Gen TyArg
6485
genTyArg vn = do
6586
return $
6687
defMessage
6788
& argName .~ vn
68-
& argKind . kindRef .~ Kind'KIND_REF_TYPE -- TODO(bladyjoker): QC.Gen arbitrary kinds.
69-
70-
genSum :: [TyArg] -> QC.Gen Sum
71-
genSum args = do
72-
cns <- QC.chooseInt (1, limit) >>= vecOf genConstrName
73-
ctors <- for (List.nub cns) (genConstructor args)
74-
return $ defMessage & constructors .~ ctors
75-
76-
-- TODO(bladyjoker): Add TyRef, TyApp etc.
77-
genTy :: [TyRef] -> [TyArg] -> QC.Gen Ty
78-
genTy (r : refs) (a : args) = QC.oneof [genTyVar (a :| args), genTyRef (r :| refs)]
79-
genTy [] (a : args) = QC.oneof [genTyVar (a :| args)]
80-
genTy (r : refs) [] = QC.oneof [genTyRef (r :| refs)]
81-
genTy _ _ = error "TODO(bladyjoker): Not yet implemented"
82-
83-
genTyRef :: NonEmpty TyRef -> QC.Gen Ty
84-
genTyRef refs = do
85-
r <- QC.elements (toList refs)
86-
return $ defMessage & tyRef .~ r
87-
88-
genTyVar :: NonEmpty TyArg -> QC.Gen Ty
89-
genTyVar args = do
90-
ar <- QC.elements (toList args)
91-
return $ defMessage & tyVar . varName .~ (ar ^. argName)
92-
93-
genConstructor :: [TyArg] -> ConstrName -> QC.Gen Sum'Constructor
94-
genConstructor args cn = do
95-
tys <- QC.chooseInt (1, limit) >>= vecOf (genTy mempty args)
89+
& argKind .~ starKind -- TODO(bladyjoker): Gen arbitrary kinds.
90+
91+
genSum :: Map TyName TyDef -> Set TyArg -> NESet ConstrName -> QC.Gen Sum
92+
genSum tydefs args ctorNs = do
93+
let (ctorN :| ctorNs') = NESet.toList ctorNs
94+
ctorNs'' <- QC.sublistOf (toList ctorNs')
95+
ctors <- for (ctorN :| ctorNs'') (genConstructor tydefs args)
96+
return $ defMessage & constructors .~ toList ctors
97+
98+
genTy :: Kind -> Map TyName TyDef -> Set TyArg -> QC.Gen Ty
99+
genTy kind tydefs tyargs =
100+
QC.oneof $
101+
NESet.withNonEmpty [] (genTyVar kind) tyargs
102+
<> NEMap.withNonEmpty [] (genTyRef kind) tydefs
103+
<> genTyApp kind tydefs tyargs
104+
105+
genTyRef :: Kind -> NEMap TyName TyDef -> [QC.Gen Ty]
106+
genTyRef kind tydefs = case [tyd | tyd <- toList tydefs, kindOf (tyd ^. tyAbs) == kind] of
107+
[] -> []
108+
tyds ->
109+
[ do
110+
tydef <- QC.elements tyds
111+
return $ defMessage & tyRef . localTyRef . tyName .~ (tydef ^. tyName)
112+
]
113+
114+
genTyVar :: Kind -> NESet TyArg -> [QC.Gen Ty]
115+
genTyVar kind args = case [tyarg | tyarg <- toList args, tyarg ^. argKind == kind] of
116+
[] -> []
117+
tyargs ->
118+
[ do
119+
tyarg <- QC.elements tyargs
120+
return $ defMessage & tyVar . varName .~ (tyarg ^. argName)
121+
]
122+
123+
genTyApp :: Kind -> Map TyName TyDef -> Set TyArg -> [QC.Gen Ty]
124+
genTyApp kind tydefs args =
125+
let kindFunc =
126+
defMessage
127+
& kindArrow . left .~ starKind -- TODO(bladyjoker): Generalize
128+
& kindArrow . right .~ kind
129+
in case [tyd | tyd <- toList tydefs, kindOf (tyd ^. tyAbs) == kindFunc] of
130+
[] -> []
131+
_ ->
132+
[ do
133+
tyfunc <- genTy kindFunc tydefs args
134+
tyarg <- genTy starKind tydefs args -- TODO(bladyjoker): Generalize
135+
return $
136+
defMessage
137+
& tyApp . tyFunc .~ tyfunc
138+
& tyApp . tyArgs .~ [tyarg] -- TODO(bladyjoker): Generate list arguments
139+
]
140+
141+
genConstructor :: Map TyName TyDef -> Set TyArg -> ConstrName -> QC.Gen Sum'Constructor
142+
genConstructor tydefs args cn = do
143+
tys <- QC.chooseInt (0, limit) >>= vecOf (genTy starKind tydefs args)
96144
return $
97145
defMessage
98146
& constrName .~ cn
99147
& P.product . ntuple . fields .~ tys
100148

101-
genTyBody :: [TyArg] -> QC.Gen TyBody
102-
genTyBody args = QC.oneof [genTyBodyOpaque, genTyBodySum args]
103-
104-
genTyBodySum :: [TyArg] -> QC.Gen TyBody
105-
genTyBodySum args = do
106-
b <- genSum args
149+
genTyBodySum :: Map TyName TyDef -> Set TyArg -> NESet ConstrName -> QC.Gen TyBody
150+
genTyBodySum tydefs args ctors = do
151+
b <- genSum tydefs args ctors
107152
return $ defMessage & P.sum .~ b
108153

109154
genTyBodyOpaque :: QC.Gen TyBody
110155
genTyBodyOpaque = return $ defMessage & P.opaque .~ defMessage
111156

112-
genTyAbs :: QC.Gen TyAbs
113-
genTyAbs = do
114-
-- TODO(bladyjoker): Allow empty args
115-
vns <- QC.chooseInt (1, limit) >>= vecOf genVarName
116-
args <- for (List.nub vns) genTyArg
117-
body <- genTyBody args
157+
genTyBody :: Map TyName TyDef -> Set TyArg -> NESet ConstrName -> QC.Gen TyBody
158+
genTyBody tydefs args ctorNs =
159+
QC.oneof $
160+
[ genTyBodyOpaque
161+
]
162+
-- Gen TyBody'Sum only if there's some TyDefs and TyArgs available
163+
<> [genTyBodySum tydefs args ctorNs | not (tydefs == mempty && args == mempty)]
164+
165+
genTyAbs :: Map TyName TyDef -> NESet ConstrName -> QC.Gen TyAbs
166+
genTyAbs tydefs ctorNs = do
167+
vns <-
168+
if tydefs == mempty
169+
then return mempty
170+
else QC.chooseInt (0, limit) >>= setOf genVarName
171+
args <- for (Set.toList vns) genTyArg
172+
body <- genTyBody tydefs (Set.fromList args) ctorNs
118173
return $
119174
defMessage
120-
& tyArgs .~ args
175+
& tyArgs .~ toList args
121176
& tyBody .~ body
122177

123-
genTyDef :: TyName -> QC.Gen TyDef
124-
genTyDef tn = do
125-
tyabs <- genTyAbs
178+
genTyDef :: Map TyName TyDef -> TyName -> NESet ConstrName -> QC.Gen TyDef
179+
genTyDef tydefs tyn ctors = do
180+
tyabs <- genTyAbs tydefs ctors
126181
return $
127182
defMessage
128-
& tyName .~ tn
183+
& tyName .~ tyn
129184
& tyAbs .~ tyabs
130185

131186
genModule :: ModuleName -> QC.Gen Module
132187
genModule mn = do
133-
tns <- QC.chooseInt (1, limit) >>= vecOf genTyName
134-
tydefs <- for (List.nub tns) genTyDef
188+
tyNs <- QC.chooseInt (1, limit) >>= nesetOf genTyName
189+
ctorNs <- QC.chooseInt (length tyNs, length tyNs * limit) >>= nesetOf genConstrName
190+
tyNsWithCtorNs <- Map.map NESet.fromList <$> distribute (toList ctorNs) (NESet.toSet tyNs)
191+
tydefs <-
192+
foldM
193+
( \availableTyDefs (tyN, ctorNs') -> do
194+
tydef <- genTyDef availableTyDefs tyN ctorNs'
195+
return $ Map.insert tyN tydef availableTyDefs
196+
)
197+
mempty
198+
(Map.toList tyNsWithCtorNs)
135199
return $
136200
defMessage
137201
& moduleName .~ mn
138-
& typeDefs .~ tydefs
202+
& typeDefs .~ toList tydefs
139203

140204
genCompilerInput :: QC.Gen CompilerInput
141205
genCompilerInput = do
142206
mns <- QC.chooseInt (1, limit) >>= vecOf genModuleName
143207
ms <- for (List.nub mns) genModule
144208
return $ defMessage & modules .~ ms
209+
210+
-- | Utils
211+
212+
-- | Distributes values (first argument) over the keys (second) randomly.
213+
distribute :: Foldable t => Ord k => t v -> Set k -> QC.Gen (Map k (NonEmpty v))
214+
distribute vals keys = do
215+
(leftover, distributed) <- distributeSingle vals keys
216+
if null leftover
217+
then return distributed
218+
else do
219+
distributed' <- distribute leftover keys
220+
return $ Map.unionWith (<>) distributed distributed'
221+
222+
distributeSingle :: Foldable t => Ord k => t v -> Set k -> QC.Gen ([v], Map k (NonEmpty v))
223+
distributeSingle vals =
224+
foldM
225+
( \(vals', dist) key ->
226+
case vals' of
227+
[] -> return (vals', dist)
228+
(v : vals'') -> do
229+
(chosenVals, leftoverVals) <- partition vals''
230+
return (leftoverVals, Map.insert key (v :| chosenVals) dist)
231+
)
232+
(toList vals, mempty)
233+
234+
-- | Partition a list randomly.
235+
partition :: forall {a}. [a] -> QC.Gen ([a], [a])
236+
partition xs = go xs []
237+
where
238+
go :: [a] -> [a] -> QC.Gen ([a], [a])
239+
go [] outs = return (outs, [])
240+
go (i : ins) outs = do
241+
b <- QC.chooseAny
242+
if b
243+
then go ins (i : outs)
244+
else return (outs, i : ins)
245+
246+
_indexBy :: Ord k => (a -> k) -> NonEmpty a -> NEMap k a
247+
_indexBy keyF (x :| xs) = foldl (\t x' -> NEMap.insert (keyF x') x' t) (NEMap.singleton (keyF x) x) xs
248+
249+
_withSourceInfo :: HasField a "sourceInfo" SourceInfo => a -> QC.Gen a
250+
_withSourceInfo msg = do
251+
f <- Text.pack <$> vecOf QC.arbitraryPrintableChar 10
252+
i <- QC.chooseInt (0, Int.maxBound)
253+
let pos =
254+
defMessage
255+
& row .~ fromIntegral i
256+
& column .~ fromIntegral i
257+
return $
258+
msg
259+
& sourceInfo . file .~ f
260+
& sourceInfo . posFrom .~ pos
261+
& sourceInfo . posTo .~ pos
262+
263+
vecOf :: forall {a}. QC.Gen a -> Int -> QC.Gen [a]
264+
vecOf = flip QC.vectorOf
265+
266+
nevecOf :: forall {a}. QC.Gen a -> Int -> QC.Gen (NonEmpty a)
267+
nevecOf g n =
268+
g >>= \x -> do
269+
xs <- QC.vectorOf (n - 1) g
270+
return $ x :| xs
271+
272+
nesetOf :: forall {a}. Ord a => QC.Gen a -> Int -> QC.Gen (NESet a)
273+
nesetOf g n = NESet.fromList <$> nevecOf g n
274+
275+
setOf :: forall {a}. Ord a => QC.Gen a -> Int -> QC.Gen (Set a)
276+
setOf g n = Set.fromList <$> vecOf g n

0 commit comments

Comments
 (0)