|
| 1 | +{-# LANGUAGE ViewPatterns #-} |
| 2 | +{-# LANGUAGE LambdaCase #-} |
| 3 | + |
| 4 | +module Ide.Plugin.Tactic.KnownStrategies.QuickCheck where |
| 5 | + |
| 6 | +import Control.Monad.Except (MonadError(throwError)) |
| 7 | +import Data.Bool (bool) |
| 8 | +import Data.List (partition) |
| 9 | +import DataCon ( DataCon, dataConName ) |
| 10 | +import Development.IDE.GHC.Compat (HsExpr, GhcPs, noLoc) |
| 11 | +import GHC.Exts ( IsString(fromString) ) |
| 12 | +import GHC.List ( foldl' ) |
| 13 | +import GHC.SourceGen (int) |
| 14 | +import GHC.SourceGen.Binds ( match, valBind ) |
| 15 | +import GHC.SourceGen.Expr ( case', lambda, let' ) |
| 16 | +import GHC.SourceGen.Overloaded ( App((@@)), HasList(list) ) |
| 17 | +import GHC.SourceGen.Pat ( conP ) |
| 18 | +import Ide.Plugin.Tactic.CodeGen |
| 19 | +import Ide.Plugin.Tactic.Judgements (jGoal) |
| 20 | +import Ide.Plugin.Tactic.Machinery (tracePrim) |
| 21 | +import Ide.Plugin.Tactic.Types |
| 22 | +import OccName (occNameString, mkVarOcc, HasOccName(occName) ) |
| 23 | +import Refinery.Tactic (goal, rule ) |
| 24 | +import TyCon (tyConName, TyCon, tyConDataCons ) |
| 25 | +import Type ( splitTyConApp_maybe ) |
| 26 | +import Data.Generics (mkQ, everything) |
| 27 | + |
| 28 | + |
| 29 | +------------------------------------------------------------------------------ |
| 30 | +-- | Known tactic for deriving @arbitrary :: Gen a@. This tactic splits the |
| 31 | +-- type's data cons into terminal and inductive cases, and generates code that |
| 32 | +-- produces a terminal if the QuickCheck size parameter is <=1, or any data con |
| 33 | +-- otherwise. It correctly scales recursive parameters, ensuring termination. |
| 34 | +deriveArbitrary :: TacticsM () |
| 35 | +deriveArbitrary = do |
| 36 | + ty <- jGoal <$> goal |
| 37 | + case splitTyConApp_maybe $ unCType ty of |
| 38 | + Just (gen_tc, [splitTyConApp_maybe -> Just (tc, apps)]) |
| 39 | + | occNameString (occName $ tyConName gen_tc) == "Gen" -> do |
| 40 | + rule $ \_ -> do |
| 41 | + let dcs = tyConDataCons tc |
| 42 | + (terminal, big) = partition ((== 0) . genRecursiveCount) |
| 43 | + $ fmap (mkGenerator tc apps) dcs |
| 44 | + terminal_expr = mkVal "terminal" |
| 45 | + oneof_expr = mkVal "oneof" |
| 46 | + pure |
| 47 | + ( tracePrim "deriveArbitrary" |
| 48 | + , noLoc $ |
| 49 | + let' [valBind (fromString "terminal") $ list $ fmap genExpr terminal] $ |
| 50 | + appDollar (mkFunc "sized") $ lambda [bvar' (mkVarOcc "n")] $ |
| 51 | + case' (infixCall "<=" (mkVal "n") (int 1)) |
| 52 | + [ match [conP (fromString "True") []] $ |
| 53 | + oneof_expr @@ terminal_expr |
| 54 | + , match [conP (fromString "False") []] $ |
| 55 | + appDollar oneof_expr $ |
| 56 | + infixCall "<>" |
| 57 | + (list $ fmap genExpr big) |
| 58 | + terminal_expr |
| 59 | + ] |
| 60 | + ) |
| 61 | + _ -> throwError $ GoalMismatch "deriveArbitrary" ty |
| 62 | + |
| 63 | + |
| 64 | +------------------------------------------------------------------------------ |
| 65 | +-- | Helper data type for the generator of a specific data con. |
| 66 | +data Generator = Generator |
| 67 | + { genRecursiveCount :: Integer |
| 68 | + , genExpr :: HsExpr GhcPs |
| 69 | + } |
| 70 | + |
| 71 | + |
| 72 | +------------------------------------------------------------------------------ |
| 73 | +-- | Make a 'Generator' for a given tycon instantiated with the given @[Type]@. |
| 74 | +mkGenerator :: TyCon -> [Type] -> DataCon -> Generator |
| 75 | +mkGenerator tc apps dc = do |
| 76 | + let dc_expr = var' $ occName $ dataConName dc |
| 77 | + args = dataConInstOrigArgTys' dc apps |
| 78 | + num_recursive_calls = sum $ fmap (bool 0 1 . doesTypeContain tc) args |
| 79 | + mkArbitrary = mkArbitraryCall tc num_recursive_calls |
| 80 | + Generator num_recursive_calls $ case args of |
| 81 | + [] -> mkFunc "pure" @@ dc_expr |
| 82 | + (a : as) -> |
| 83 | + foldl' |
| 84 | + (infixCall "<*>") |
| 85 | + (infixCall "<$>" dc_expr $ mkArbitrary a) |
| 86 | + (fmap mkArbitrary as) |
| 87 | + |
| 88 | + |
| 89 | +------------------------------------------------------------------------------ |
| 90 | +-- | Check if the given 'TyCon' exists anywhere in the 'Type'. |
| 91 | +doesTypeContain :: TyCon -> Type -> Bool |
| 92 | +doesTypeContain recursive_tc = |
| 93 | + everything (||) $ mkQ False (== recursive_tc) |
| 94 | + |
| 95 | + |
| 96 | +------------------------------------------------------------------------------ |
| 97 | +-- | Generate the correct sort of call to @arbitrary@. For recursive calls, we |
| 98 | +-- need to scale down the size parameter, either by a constant factor of 1 if |
| 99 | +-- it's the only recursive parameter, or by @`div` n@ where n is the number of |
| 100 | +-- recursive parameters. For all other types, just call @arbitrary@ directly. |
| 101 | +mkArbitraryCall :: TyCon -> Integer -> Type -> HsExpr GhcPs |
| 102 | +mkArbitraryCall recursive_tc n ty = |
| 103 | + let arbitrary = mkFunc "arbitrary" |
| 104 | + in case doesTypeContain recursive_tc ty of |
| 105 | + True -> |
| 106 | + mkFunc "scale" |
| 107 | + @@ bool (mkFunc "flip" @@ mkFunc "div" @@ int n) |
| 108 | + (mkFunc "subtract" @@ int 1) |
| 109 | + (n == 1) |
| 110 | + @@ arbitrary |
| 111 | + False -> arbitrary |
| 112 | + |
0 commit comments