writing Arbitrary instances for monadic DSLs which generate references

ShrinkingScopedPrograms.hs

{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving, LambdaCase #-}
module ShrinkingScopedPrograms where
import Test.DocTest

import Control.Monad.State
import Data.Map (Map)
import Test.QuickCheck
import qualified Data.Map as Map

-- $setup
-- >>> import Data.Maybe


-- The challenge is to write an Arbitrary instance for a monadic DSL in which
-- some actions can only reference things which have already been generated.
--
-- First, let's create such a monadic DSL.

newtype Counter = UnsafeCounter
  { unCounter :: Int }

newtype CounterGenerator a = UnsafeCounterGenerator
  { unCounterGenerator :: StateT (Map Int Int) Maybe a }
  deriving (Functor, Applicative, Monad)

-- |
-- >>> :{
-- runCounterGenerator $ do
--   counter1 <- newCounter
--   counter2 <- newCounter
--   (,,,) <$> bumpCounter counter1
--         <*> bumpCounter counter2
--         <*> bumpCounter counter2
--         <*> bumpCounter counter1
-- :}
-- Just (0,0,1,1)
runCounterGenerator :: CounterGenerator a -> Maybe a
runCounterGenerator = flip evalStateT Map.empty . unCounterGenerator

newCounter :: CounterGenerator Counter
newCounter = UnsafeCounterGenerator $ do
  n <- Map.size <$> get
  modify $ Map.insert n 0
  pure (UnsafeCounter n)

-- Must be called on a 'Counter' generated by a previous call to 'newCounter' in
-- this 'CounterGenerator' computation. This restriction makes it difficult to
-- write an 'Arbitrary' instance: the set of valid 'Counter' values varies over
-- time, so we cannot simply call 'arbitrary' to pick a random one.
bumpCounter :: Counter -> CounterGenerator Int
bumpCounter (UnsafeCounter i) = UnsafeCounterGenerator $ do
  s <- get
  let s' = Map.adjust (+1) i s
  put s'

  -- an intentional defect which we'll find using 'quickCheck'
  when ( 0 `Map.member` s'
      && 1 `Map.member` s'
      && 2 `Map.member` s'
      && s' Map.! 0 < s' Map.! 1
      && s' Map.! 1 < s' Map.! 2
       ) $ do
    mzero

  pure (s Map.! i)

-- The trick is: don't write the 'arbitrary' method for 'CounterGenerator' in
-- terms of recursive calls to 'arbitrary', use a dedicated function which takes
-- the Counters which are in scope as an argument.
arbitraryCounterGenerator :: Int
                          -> Int
                          -> CounterGenerator [Counter]
                          -> Gen (CounterGenerator ())
arbitraryCounterGenerator remainingSteps numberOfCounters generateCounters =
  if remainingSteps == 0
  then terminate
  else oneof [addAnotherCounter, bumpArbitraryCounter]
  where
    terminate :: Gen (CounterGenerator ())
    terminate = pure $ do
      _ <- generateCounters
      pure ()

    bumpArbitraryCounter :: Gen (CounterGenerator ())
    bumpArbitraryCounter = do
      i <- choose (0, numberOfCounters - 1)
      arbitraryCounterGenerator (remainingSteps - 1) numberOfCounters $ do
        counters <- generateCounters
        _ <- bumpCounter (counters !! i)
        pure counters

    addAnotherCounter :: Gen (CounterGenerator ())
    addAnotherCounter = do
      arbitraryCounterGenerator (remainingSteps - 1) (numberOfCounters + 1) $ do
        counters <- generateCounters
        counter <- newCounter
        pure (counters ++ [counter])

instance Arbitrary (CounterGenerator ()) where
  arbitrary = sized $ \n -> do
    k <- choose (0, n)
    arbitraryCounterGenerator k 1 $ do
      counter <- newCounter
      pure [counter]

-- |
-- If we try enough programs, sooner or later we'll trigger the bug:
--
-- >>> :{
-- forever $ do
--   counterGenerator <- generate arbitrary :: IO (CounterGenerator ())
--   guard $ isJust $ runCounterGenerator counterGenerator
-- :}
-- *** Exception: user error (mzero)
-- ...
--
-- Unfortunately, that bug is going to be pretty hard to fix, because we cannot
-- examine nor shrink the program which triggers it. Using a Free Monad encoding
-- would help us to examine the program, but because that encoding contains
-- functions, it would be difficult to implement a shrink method.
--
-- Instead, let's split 'arbitraryCounterGenerator' into two pieces: one which
-- generates a value of a type we can easily examine and shrink, and one which
-- interprets that value into a 'CounterGenerator'.

data CounterAction
  = NewCounter Int
  | BumpCounter Int
  deriving (Eq, Show)

-- There is an implicit @NewCounter 0@ in front of the list.
-- The first 'NewCounter' must use 1, the second 2, etc.
-- If there are @n@ NewCounters to the left of a 'BumpCounter', that
-- 'BumpCounter' must be an element of @[0..n]@.
newtype CounterActions = CounterActions
  { unCounterActions :: [CounterAction] }
  deriving Show

-- precondition: @(actions ++ [BumpCounter (numberOfCounters-1)])@ must be a
-- valid 'CounterActions'
arbitraryCounterActions :: Int
                        -> Int
                        -> [CounterAction]
                        -> Gen CounterActions
arbitraryCounterActions remainingSteps numberOfCounters actions =
  if remainingSteps == 0
  then bumpFinalCounter
  else oneof [addAnotherCounter, bumpArbitraryCounter]
  where
    bumpArbitraryCounter :: Gen CounterActions
    bumpArbitraryCounter = do
      i <- choose (0, numberOfCounters - 1)
      let action = BumpCounter i
      arbitraryCounterActions (remainingSteps - 1)
                              numberOfCounters
                              (actions ++ [action])

    bumpFinalCounter :: Gen CounterActions
    bumpFinalCounter = do
      i <- choose (0, numberOfCounters - 1)
      let action = BumpCounter i
      pure $ CounterActions (actions ++ [action])

    addAnotherCounter :: Gen CounterActions
    addAnotherCounter = do
      let action = NewCounter numberOfCounters
      arbitraryCounterActions (remainingSteps - 1)
                              (numberOfCounters + 1)
                              (actions ++ [action])

instance Arbitrary CounterActions where
  arbitrary = sized $ \n -> do
    k <- choose (0, n)
    arbitraryCounterActions k 1 []
  shrink actions0 = fmap CounterActions . go . unCounterActions $ actions0
    where
      go :: [CounterAction] -> [[CounterAction]]
      go [] = []
      go (BumpCounter i : actions) = dropAction
                                   : shrinkAction
                                  ++ shrinkActions
        where
          dropAction :: [CounterAction]
          dropAction = actions

          shrinkAction :: [[CounterAction]]
          shrinkAction = do
            i' <- shrink i
            pure (BumpCounter i' : actions)

          shrinkActions :: [[CounterAction]]
          shrinkActions = do
            actions' <- go actions
            pure (BumpCounter i : actions')
      go (NewCounter i : actions) = dropAction
                                  : shrinkActions
        where
          dropAction :: [CounterAction]
          dropAction = map (\case
                              NewCounter  j | j > i -> NewCounter  (j - 1)
                              BumpCounter j | j > i -> BumpCounter (j - 1)
                              action                -> action)
                     . filter (/= BumpCounter i)
                     $ actions

          shrinkActions :: [[CounterAction]]
          shrinkActions = do
            actions' <- go actions
            pure (NewCounter i : actions')

interpretCounterActions :: CounterActions -> CounterGenerator ()
interpretCounterActions counterActions = do
  counter <- newCounter
  go [counter] (unCounterActions counterActions)
  where
    go :: [Counter] -> [CounterAction] -> CounterGenerator ()
    go _ [] = pure ()
    go counters (BumpCounter i : actions) = do
      _ <- bumpCounter (counters !! i)
      go counters actions
    go counters (NewCounter _ : actions) = do
      counter <- newCounter
      go (counters ++ [counter]) actions

-- We can now trigger the bug, shrink the 'CounterActions' which caused it, and
-- examine the problematic program:
--
-- >>> quickCheck (isJust . runCounterGenerator . interpretCounterActions)
-- *** Failed! Falsifiable (after ... tests and ... shrinks):
-- CounterActions {unCounterActions = [NewCounter 1,NewCounter 2,BumpCounter 2,BumpCounter 2,BumpCounter 1]}

-- | which corresponds to
--
-- >>> :{
-- runCounterGenerator $ do
--   counter0 <- newCounter  -- implicit @NewCounter 0@
--   counter1 <- newCounter  -- @NewCounter 1@
--   counter2 <- newCounter  -- @NewCounter 2@
--   bumpCounter counter2    -- @BumpCounter 2@
--   bumpCounter counter2    -- @BumpCounter 2@
--   bumpCounter counter1    -- @BumpCounter 1@
-- :}
-- Nothing


main :: IO ()
main = doctest ["ShrinkingScopedPrograms.hs"]