QuickCheck.hs

{-# LANGUAGE ViewPatterns #-}

module Ide.Plugin.Tactic.KnownStrategies.QuickCheck where

import           Control.Monad.Except         (MonadError (throwError))
import           Data.Bool                    (bool)
import           Data.Generics                (everything, mkQ)
import           Data.List                    (partition)
import           DataCon                      (DataCon, dataConName)
import           Development.IDE.GHC.Compat   (GhcPs, HsExpr, noLoc)
import           GHC.Exts                     (IsString (fromString))
import           GHC.List                     (foldl')
import           GHC.SourceGen                (int)
import           GHC.SourceGen.Binds          (match, valBind)
import           GHC.SourceGen.Expr           (case', lambda, let')
import           GHC.SourceGen.Overloaded     (App ((@@)), HasList (list))
import           GHC.SourceGen.Pat            (conP)
import           Ide.Plugin.Tactic.CodeGen
import           Ide.Plugin.Tactic.Judgements (jGoal)
import           Ide.Plugin.Tactic.Machinery  (tracePrim)
import           Ide.Plugin.Tactic.Types
import           OccName                      (HasOccName (occName), mkVarOcc,
                                               occNameString)
import           Refinery.Tactic              (goal, rule)
import           TyCon                        (TyCon, tyConDataCons, tyConName)
import           Type                         (splitTyConApp_maybe)


------------------------------------------------------------------------------
-- | Known tactic for deriving @arbitrary :: Gen a@. This tactic splits the
-- type's data cons into terminal and inductive cases, and generates code that
-- produces a terminal if the QuickCheck size parameter is <=1, or any data con
-- otherwise. It correctly scales recursive parameters, ensuring termination.
deriveArbitrary :: TacticsM ()
deriveArbitrary = do
  ty <- jGoal <$> goal
  case splitTyConApp_maybe $ unCType ty of
    Just (gen_tc, [splitTyConApp_maybe -> Just (tc, apps)])
        | occNameString (occName $ tyConName gen_tc) == "Gen" -> do
      rule $ \_ -> do
        let dcs = tyConDataCons tc
            (terminal, big) = partition ((== 0) . genRecursiveCount)
                        $ fmap (mkGenerator tc apps) dcs
            terminal_expr = mkVal "terminal"
            oneof_expr = mkVal "oneof"
        pure
          $ Synthesized (tracePrim "deriveArbitrary")
              -- TODO(sandy): This thing is not actually empty! We produced
              -- a bespoke binding "terminal", and a not-so-bespoke "n".
              -- But maybe it's fine for known rules?
              mempty
              mempty
          $ noLoc $
              let' [valBind (fromString "terminal") $ list $ fmap genExpr terminal] $
                appDollar (mkFunc "sized") $ lambda [bvar' (mkVarOcc "n")] $
                  case' (infixCall "<=" (mkVal "n") (int 1))
                    [ match [conP (fromString "True") []] $
                        oneof_expr @@ terminal_expr
                    , match [conP (fromString "False") []] $
                        appDollar oneof_expr $
                          infixCall "<>"
                            (list $ fmap genExpr big)
                            terminal_expr
                    ]
    _ -> throwError $ GoalMismatch "deriveArbitrary" ty


------------------------------------------------------------------------------
-- | Helper data type for the generator of a specific data con.
data Generator = Generator
  { genRecursiveCount :: Integer
  , genExpr           :: HsExpr GhcPs
  }


------------------------------------------------------------------------------
-- | Make a 'Generator' for a given tycon instantiated with the given @[Type]@.
mkGenerator :: TyCon -> [Type] -> DataCon -> Generator
mkGenerator tc apps dc = do
  let dc_expr   = var' $ occName $ dataConName dc
      args = dataConInstOrigArgTys' dc apps
      num_recursive_calls = sum $ fmap (bool 0 1 . doesTypeContain tc) args
      mkArbitrary = mkArbitraryCall tc num_recursive_calls
  Generator num_recursive_calls $ case args of
    []  -> mkFunc "pure" @@ dc_expr
    (a : as) ->
      foldl'
        (infixCall "<*>")
        (infixCall "<$>" dc_expr $ mkArbitrary a)
        (fmap mkArbitrary as)


------------------------------------------------------------------------------
-- | Check if the given 'TyCon' exists anywhere in the 'Type'.
doesTypeContain :: TyCon -> Type -> Bool
doesTypeContain recursive_tc =
  everything (||) $ mkQ False (== recursive_tc)


------------------------------------------------------------------------------
-- | Generate the correct sort of call to @arbitrary@. For recursive calls, we
-- need to scale down the size parameter, either by a constant factor of 1 if
-- it's the only recursive parameter, or by @`div` n@ where n is the number of
-- recursive parameters. For all other types, just call @arbitrary@ directly.
mkArbitraryCall :: TyCon -> Integer -> Type -> HsExpr GhcPs
mkArbitraryCall recursive_tc n ty =
  let arbitrary = mkFunc "arbitrary"
   in case doesTypeContain recursive_tc ty of
        True ->
          mkFunc "scale"
            @@ bool (mkFunc "flip" @@ mkFunc "div" @@ int n)
                    (mkFunc "subtract" @@ int 1)
                    (n == 1)
            @@ arbitrary
        False -> arbitrary