content: {-# LANGUAGE ScopedTypeVariables #-}
  
  module Test.SmartCheck.SmartGen
    ( iterateArbIdx
    , iterateArb
    , resultify
    , replace
    , iter
    ) where
  
  import Test.SmartCheck.Types
  import Test.SmartCheck.DataToTree
  
  import qualified Test.QuickCheck.Gen      as Q
  import qualified Test.QuickCheck.Random   as Q
  import qualified Test.QuickCheck          as Q hiding (Result)
  import qualified Test.QuickCheck.Property as P
  
  import Prelude hiding (max)
  import System.Random
  import Data.Tree hiding (levels)
  
  --------------------------------------------------------------------------------
  
  -- | Driver for iterateArb.
  iterateArbIdx :: SubTypes a
                => a -> (Idx, Maybe Int) -> Int -> Int
                -> (a -> P.Property) -> IO (Int, Result a)
  iterateArbIdx d (idx, max) tries sz prop =
    case getAtIdx d idx max of
      Nothing  -> errorMsg "iterateArb 0"
      Just ext -> case ext of
                    -- Don't analyze base types.
                    SubT e -> if baseType e then return (0, BaseType)
                                else iterateArb d ext idx tries sz prop
  
  -- | Replace the hole in d indexed by idx with a bunch of random values, and
  -- test the new d against the property.  Returns the first new d (the full d but
  -- with the hole replaced) that succeeds.  "Succeeds" is determined by the call
  -- to resultify---if we're expecting failure, then we succeed by getting a value
  -- that passes the precondition but fails the property; otherwise we succeed by
  -- getting a value that passes the precondition and satisfies the property.  If
  -- no value ever satisfies the precondition, then we return FailedPreCond.
  -- (Thus, there's an implied linear order on the Result type: FailedPreCond <
  -- FailedProp < Result a.)
  iterateArb :: forall a. SubTypes a
    => a                  -- ^ Counterexample.
    -> SubT               -- ^ Sub-value in the counterexample.
    -> Idx                -- ^ Index of sub-value.
    -> Int                -- ^ Maximum number of iterations.
    -> Int                -- ^ Maximum size of value to generate.
    -> (a -> P.Property)      -- ^ Property.
    -> IO (Int, Result a) -- ^ Number of times precondition is passed and returned
                          -- result.
  iterateArb d ext idx tries max prop = do
    g <- Q.newQCGen
    iterateArb' (0, FailedPreCond) g 0 0
    where
    newMax SubT { unSubT = v } = valDepth v
  
    -- Main loop.  We break out if we ever satisfy the property.  Otherwise, we
    -- return the latest value.
    iterateArb' :: (Int, Result a) -> Q.QCGen -> Int -> Int -> IO (Int, Result a)
    iterateArb' (i, res) g try currMax
      -- We've exhausted the number of iterations.
      | try >= tries = return (i, res)
      -- The generated random value is too big.  Start again sampling again with
      -- size at 0.
      | newMax s >= max = iterateArb' (i, res) g0 (try + 1) 0
      | otherwise =
          case replace d idx s of
            Nothing -> errorMsg "iterateArb 1"
            Just d' -> do
              res' <- resultify prop d'
              case res' of
                FailedPreCond -> rec    (i  , FailedPreCond)
                FailedProp    -> rec    (i+1, FailedProp)
                Result x      -> return (i+1, Result x)
                BaseType      -> errorMsg "baseType from resultify"
      where
      (size, g0) = randomR (0, currMax) g
      sample SubT { unSubT = v } = newVal v
      s = sample ext g size
      rec res' =
        iterateArb' res' g0 (try + 1)
          -- XXX what ratio is right to increase size of values?  This gives us
          -- exponentail growth, but remember we're randomly chosing within the
          -- range of [0, max], so many values are significantly smaller than the
          -- max.  Plus we reset the size whenever we get a value that's too big.
          -- Note the need for (+ 1), since we seed with 0.
          ((currMax + 1) * 2)
  
  --------------------------------------------------------------------------------
  
  -- | Make a new random value given a generator and a max size.  Based on the
  -- value's type's arbitrary instance.
  newVal :: forall a. (SubTypes a, Q.Arbitrary a)
         => a -> Q.QCGen -> Int -> SubT
  newVal _ g size =
    let Q.MkGen m = Q.resize size (Q.arbitrary :: Q.Gen a) in
    subT (m g size)
  
  --------------------------------------------------------------------------------
  
  -- | Put a value v into a another value d at a hole idx, if v is well-typed.
  -- Return Nothing if dynamic typing fails.
  replace :: SubTypes a => a -> Idx -> SubT -> Maybe a
  replace d idx SubT { unSubT = v } = replaceAtIdx d idx v
  
  --------------------------------------------------------------------------------
  
  -- | Make a QuickCheck Result by applying a property function to a value and
  -- then get out the Result using our result type.
  resultify :: (a -> P.Property) -> a -> IO (Result a)
  resultify prop a = do
    P.MkRose r _ <- res fs
    return $ maybe FailedPreCond -- Failed precondition (discard)
                   -- If failed because of an exception, just say we failed.
                   (\b -> if notExceptionFail r then get b r else FailedProp)
                   (P.ok r) -- result of test case (True ==> passed)
    where
    get b r
      |     b &&      P.expect r  = Result a -- expected to pass and we did
      | not b && not (P.expect r) = Result a -- expected failure and got it
      | otherwise                 = FailedProp -- We'll just discard it.
  
    P.MkProperty { P.unProperty = Q.MkGen { Q.unGen = f } }
        = prop a :: P.Property
    fs  = P.unProp $ f err err            :: P.Rose P.Result
    res = P.protectRose . P.reduceRose
  
    -- XXX A hack!  Means we failed the property because it failed, not because of
    -- an exception (i.e., with partial function tests).
    notExceptionFail r = let e = P.reason r in
                         e == "Falsifiable" || e  == ""
  
    err = errorMsg "resultify: should not evaluate."
  
  --------------------------------------------------------------------------------
  
  type Test a b = a -> Idx -> IO b
  type Next a b = a -> b -> Forest Bool -> Idx -> [Idx] -> IO (a, [Idx])
  
  -- Do a breadth-first traversal of the data.  First, we find the next valid
  -- index we can use.  Then we apply our test function, passing the result to our
  -- next function.
  iter :: SubTypes a
       => a                 -- ^ Failed value
       -> Test a b          -- ^ Test to use
       -> Next a b          -- ^ What to do after the test
       -> (a -> Q.Property) -- ^ Property
       -> Maybe Int         -- ^ Max depth to analyze
       -> Forest Bool       -- ^ Only evaluate at True indexes.
       -> Idx               -- ^ Starting index to extrapolate
       -> [Idx]             -- ^ List of generalized indices
       -> IO (a, [Idx])
  iter d test nxt prop maxLevel forest idx idxs
    | done      = return (d, idxs)
    | nextLevel = iter'
    | atFalse   = iter' -- Must be last check or !! index below may be out of
                        -- bounds!
    | otherwise = do tries <- test d idx
                     nxt d tries forest idx idxs
    where
    -- Location is w.r.t. the forest, not the original data value.
    l          = level idx
    levels     = breadthLevels forest
    done       = length levels <= l || tooDeep l maxLevel
    nextLevel  = length (levels !! l) <= column idx
    atFalse    = not $ (levels !! l) !! column idx
    iter'      = iter d test nxt prop maxLevel forest
                   idx { level = l + 1, column = 0 } idxs
  
  --------------------------------------------------------------------------------
  
  -- | Get the maximum depth of a value, where depth is measured in the maximum
  -- depth of the tree representation, not counting base types (defined in
  -- Types.hs).
  valDepth :: SubTypes a => a -> Int
  valDepth d = depth (mkSubstForest d True)
  
  --------------------------------------------------------------------------------
  