{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DefaultSignatures     #-}
{-# LANGUAGE DerivingVia           #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module NoThunks.Class (
    -- * Check a value for unexpected thunks
    NoThunks(..)
  , ThunkInfo(..)
  , Context
  , unsafeNoThunks
    -- * Helpers for defining instances
  , allNoThunks
  , noThunksInValues
  , noThunksInKeysAndValues
    -- * Deriving-via wrappers
  , OnlyCheckWhnf(..)
  , OnlyCheckWhnfNamed(..)
  , InspectHeap(..)
  , InspectHeapNamed(..)
  , AllowThunk(..)
  , AllowThunksIn(..)
    -- * Generic class
  , GWNoThunks(..)
  ) where

import Data.Proxy
import Data.Typeable
import System.IO.Unsafe (unsafePerformIO)

import GHC.Exts.Heap
import GHC.Generics
import GHC.Records
import GHC.TypeLits

-- For instances

import Data.Foldable (toList)
import Data.Int
import Data.IntMap (IntMap)
import Data.Kind (Type)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map (Map)
import Data.Ratio
import Data.Sequence (Seq)
import Data.Set (Set)
import Data.Time
import Data.Void (Void)
import Data.Word
import GHC.Stack
-- base-4.16 exports 'Natural' from 'GHC.TypeLits'
#if !MIN_VERSION_base(4,16,0)
import Numeric.Natural
#endif

import qualified Control.Concurrent.MVar       as MVar
import qualified Control.Concurrent.STM.TVar   as TVar
import qualified Data.IntMap                   as IntMap
import qualified Data.IORef                    as IORef
import qualified Data.Map                      as Map
import qualified Data.Set                      as Set

#ifdef MIN_VERSION_bytestring
import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString               as BS.Strict
import qualified Data.ByteString.Lazy          as BS.Lazy
import qualified Data.ByteString.Lazy.Internal as BS.Lazy.Internal
#endif

#ifdef MIN_VERSION_text
import qualified Data.Text                     as Text.Strict
import qualified Data.Text.Internal.Lazy       as Text.Lazy.Internal
import qualified Data.Text.Lazy                as Text.Lazy
#endif

#ifdef MIN_VERSION_vector
import qualified Data.Vector                   as Vector.Boxed
import qualified Data.Vector.Unboxed           as Vector.Unboxed
#endif

{-------------------------------------------------------------------------------
  Check a value for unexpected thunks
-------------------------------------------------------------------------------}

-- | Check a value for unexpected thunks
class NoThunks a where
  -- | Check if the argument does not contain any unexpected thunks
  --
  -- For most datatypes, we should have that
  --
  -- > noThunks ctxt x == Nothing
  --
  -- if and only if
  --
  -- > checkContainsThunks x
  --
  -- For some datatypes however, some thunks are expected. For example, the
  -- internal fingertree 'Data.Sequence.Sequence' might contain thunks (this is
  -- important for the asymptotic complexity of this data structure). However,
  -- we should still check that the /values/ in the sequence don't contain any
  -- unexpected thunks.
  --
  -- This means that we need to traverse the sequence, which might force some of
  -- the thunks in the tree. In general, it is acceptable for
  -- 'noThunks' to force such "expected thunks", as long as it always
  -- reports the /unexpected/ thunks.
  --
  -- The default implementation of 'noThunks' checks that the argument is in
  -- WHNF, and if so, adds the type into the context (using 'showTypeOf'), and
  -- calls 'wNoThunks'. See 'ThunkInfo' for a detailed discussion of the type
  -- context.
  --
  -- See also discussion of caveats listed for 'checkContainsThunks'.
  noThunks :: Context -> a -> IO (Maybe ThunkInfo)
  noThunks Context
ctxt a
x = do
      Bool
isThunk <- forall a. a -> IO Bool
checkIsThunk a
x
      if Bool
isThunk
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just ThunkInfo { thunkContext :: Context
thunkContext = Context
ctxt' }
        else forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt' a
x
    where
      ctxt' :: Context
      ctxt' :: Context
ctxt' = forall a. NoThunks a => Proxy a -> String
showTypeOf (forall {k} (t :: k). Proxy t
Proxy @a) forall a. a -> [a] -> [a]
: Context
ctxt

  -- | Check that the argument is in normal form, assuming it is in WHNF.
  --
  -- The context will already have been extended with the type we're looking at,
  -- so all that's left is to look at the thunks /inside/ the type. The default
  -- implementation uses GHC Generics to do this.
  wNoThunks :: Context -> a -> IO (Maybe ThunkInfo)
  default wNoThunks :: (Generic a, GWNoThunks '[] (Rep a))
                    => Context -> a -> IO (Maybe ThunkInfo)
  wNoThunks Context
ctxt a
x = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks (forall {k} (t :: k). Proxy t
Proxy @'[]) Context
ctxt forall x. Rep a x
fp
    where
      -- Force the result of @from@ to WHNF: we are not interested in thunks
      -- that arise from the translation to the generic representation.
      fp :: Rep a x
      !fp :: forall x. Rep a x
fp = forall a x. Generic a => a -> Rep a x
from a
x

  -- | Show type @a@ (to add to the context)
  --
  -- We try hard to avoid 'Typeable' constraints in this module: there are types
  -- with no 'Typeable' instance but with a 'NoThunks' instance (most
  -- important example are types such as @ST s@ which rely on parametric
  -- polymorphism). By default we should therefore only show the "outer layer";
  -- for example, if we have a type
  --
  -- > Seq (ST s ())
  --
  -- then 'showTypeOf' should just give @Seq@, leaving it up to the instance for
  -- @ST@ to decide how to implement 'showTypeOf'; this keeps things
  -- compositional. The default implementation does precisely this using the
  -- metadata that GHC Generics provides.
  --
  -- For convenience, however, some of the @deriving via@ newtype wrappers we
  -- provide /do/ depend on @Typeable@; see below.
  showTypeOf :: Proxy a -> String
  default showTypeOf :: (Generic a, GShowTypeOf (Rep a)) => Proxy a -> String
  showTypeOf Proxy a
_ = forall (f :: * -> *) x. GShowTypeOf f => f x -> String
gShowTypeOf (forall a x. Generic a => a -> Rep a x
from a
x)
    where
      x :: a
      x :: a
x = a
x

-- | Context where a thunk was found
--
-- This is intended to give a hint about which thunk was found. For example,
-- a thunk might be reported with context
--
-- > ["Int", "(,)", "Map", "AppState"]
--
-- telling you that you have an @AppState@ containing a @Map@ containing a pair,
-- all of which weren't thunks (were in WHNF), but that pair contained an
-- @Int@ which was a thunk.
type Context = [String]

{-------------------------------------------------------------------------------
  Results of the check
-------------------------------------------------------------------------------}

-- | Information about unexpected thunks
--
-- TODO: The ghc-debug work by Matthew Pickering includes some work that allows
-- to get source spans from closures. If we could take advantage of that, we
-- could not only show the type of the unexpected thunk, but also where it got
-- allocated.
newtype ThunkInfo = ThunkInfo {
      -- The @Context@ argument is intended to give a clue to add debugging.
      -- For example, suppose we have something of type @(Int, [Int])@. The
      -- various contexts we might get are
      --
      -- > Context                  The thunk is..
      -- > ---------------------------------------------------------------------
      -- > ["(,)"]                  the pair itself
      -- > ["Int","(,)"]            the Int in the pair
      -- > ["List","(,)"]           the [Int] in the pair
      -- > ["Int","List","(,)"]     an Int in the [Int] in the pair
      --
      -- Note: prior to `ghc-9.6` a list was indicated by `[]`.
      ThunkInfo -> Context
thunkContext :: Context
    }
  deriving (Int -> ThunkInfo -> ShowS
[ThunkInfo] -> ShowS
ThunkInfo -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ThunkInfo] -> ShowS
$cshowList :: [ThunkInfo] -> ShowS
show :: ThunkInfo -> String
$cshow :: ThunkInfo -> String
showsPrec :: Int -> ThunkInfo -> ShowS
$cshowsPrec :: Int -> ThunkInfo -> ShowS
Show)

{-# NOINLINE unsafeNoThunks #-}
-- | Call 'noThunks' in a pure context (relies on 'unsafePerformIO').
unsafeNoThunks :: NoThunks a => a -> Maybe ThunkInfo
unsafeNoThunks :: forall a. NoThunks a => a -> Maybe ThunkInfo
unsafeNoThunks a
a = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks [] a
a

{-------------------------------------------------------------------------------
  Helpers for defining NoThunks instances
-------------------------------------------------------------------------------}

-- | Short-circuit a list of checks
allNoThunks :: [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks :: [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks = [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
go
  where
    go :: [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
    go :: [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
go []     = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go (IO (Maybe ThunkInfo)
a:[IO (Maybe ThunkInfo)]
as) = do
        Maybe ThunkInfo
nf <- IO (Maybe ThunkInfo)
a
        case Maybe ThunkInfo
nf of
          Maybe ThunkInfo
Nothing    -> [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
go [IO (Maybe ThunkInfo)]
as
          Just ThunkInfo
thunk -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just ThunkInfo
thunk

-- | Check that all elements in the list are thunk-free
--
-- Does not check the list itself. Useful for checking the elements of a
-- container.
--
-- See also 'noThunksInKeysAndValues'
noThunksInValues :: NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues :: forall a. NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues Context
ctxt = [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt)

-- | Variant on 'noThunksInValues' for keyed containers.
--
-- Neither the list nor the tuples are checked for thunks.
noThunksInKeysAndValues :: (NoThunks k, NoThunks v)
                        => Context -> [(k, v)] -> IO (Maybe ThunkInfo)
noThunksInKeysAndValues :: forall k v.
(NoThunks k, NoThunks v) =>
Context -> [(k, v)] -> IO (Maybe ThunkInfo)
noThunksInKeysAndValues Context
ctxt =
      [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(k
k, v
v) -> [ forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt k
k
                            , forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt v
v
                            ])

{-------------------------------------------------------------------------------
  Newtype wrappers for deriving via
-------------------------------------------------------------------------------}

-- | Newtype wrapper for use with @deriving via@ to check for WHNF only
--
-- For some types we don't want to check for nested thunks, and we only want
-- check if the argument is in WHNF, not in NF. A typical example are functions;
-- see the instance of @(a -> b)@ for detailed discussion. This should be used
-- sparingly.
--
-- Example:
--
-- > deriving via OnlyCheckWhnf T instance NoThunks T
newtype OnlyCheckWhnf a = OnlyCheckWhnf a

-- | Variant on 'OnlyCheckWhnf' that does not depend on 'Generic'
--
-- Example:
--
-- > deriving via OnlyCheckWhnfNamed "T" T instance NoThunks T
newtype OnlyCheckWhnfNamed (name :: Symbol) a = OnlyCheckWhnfNamed a

-- | Newtype wrapper for values that should be allowed to be a thunk
--
-- This should be used /VERY/ sparingly, and should /ONLY/ be used on values
-- (or, even rarer, types) which you are /SURE/ cannot retain any data that they
-- shouldn't. Bear in mind allowing a value of type @T@ to be a thunk might
-- cause a value of type @S@ to be retained if @T@ was computed from @S@.
newtype AllowThunk a = AllowThunk a

-- | Newtype wrapper for records where some of the fields are allowed to be
-- thunks.
--
-- Example:
--
-- > deriving via AllowThunksIn '["foo","bar"] T instance NoThunks T
--
-- This will create an instance that skips the thunk checks for the "foo" and
-- "bar" fields.
newtype AllowThunksIn (fields :: [Symbol]) a = AllowThunksIn a

-- | Newtype wrapper for use with @deriving via@ to inspect the heap directly
--
-- This bypasses the class instances altogether, and inspects the GHC heap
-- directly, checking that the value does not contain any thunks /anywhere/.
-- Since we can do this without any type classes instances, this is useful for
-- types that contain fields for which 'NoThunks' instances are not available.
--
-- Since the primary use case for 'InspectHeap' then is to give instances
-- for 'NoThunks' from third party libraries, we also don't want to
-- rely on a 'Generic' instance, which may likewise not be available. Instead,
-- we will rely on 'Typeable', which is available for /all/ types. However, as
-- 'showTypeOf' explains, requiring 'Typeable' may not always be suitable; if
-- it isn't, 'InspectHeapNamed' can be used.
--
-- Example:
--
-- > deriving via InspectHeap T instance NoThunks T
newtype InspectHeap a = InspectHeap a

-- | Variant on 'InspectHeap' that does not depend on 'Typeable'.
--
-- > deriving via InspectHeapNamed "T" T instance NoUnexpecedThunks T
newtype InspectHeapNamed (name :: Symbol) a = InspectHeapNamed a

{-------------------------------------------------------------------------------
  Internal: instances for the deriving-via wrappers
-------------------------------------------------------------------------------}

instance Typeable a => NoThunks (OnlyCheckWhnf a) where
  showTypeOf :: Proxy (OnlyCheckWhnf a) -> String
showTypeOf Proxy (OnlyCheckWhnf a)
_  = forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy @a)
  wNoThunks :: Context -> OnlyCheckWhnf a -> IO (Maybe ThunkInfo)
wNoThunks Context
_ OnlyCheckWhnf a
_ = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

instance KnownSymbol name => NoThunks (OnlyCheckWhnfNamed name a) where
  showTypeOf :: Proxy (OnlyCheckWhnfNamed name a) -> String
showTypeOf Proxy (OnlyCheckWhnfNamed name a)
_  = forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (forall {k} (t :: k). Proxy t
Proxy @name)
  wNoThunks :: Context -> OnlyCheckWhnfNamed name a -> IO (Maybe ThunkInfo)
wNoThunks Context
_ OnlyCheckWhnfNamed name a
_ = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

instance NoThunks (AllowThunk a) where
  showTypeOf :: Proxy (AllowThunk a) -> String
showTypeOf Proxy (AllowThunk a)
_ = String
"<never used since never fails>"
  noThunks :: Context -> AllowThunk a -> IO (Maybe ThunkInfo)
noThunks Context
_ AllowThunk a
_ = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
  wNoThunks :: Context -> AllowThunk a -> IO (Maybe ThunkInfo)
wNoThunks    = forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks

instance (HasFields s a, Generic a, Typeable a, GWNoThunks s (Rep a))
      => NoThunks (AllowThunksIn s a) where
  showTypeOf :: Proxy (AllowThunksIn s a) -> String
showTypeOf Proxy (AllowThunksIn s a)
_ = forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy @a)
  wNoThunks :: Context -> AllowThunksIn s a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt (AllowThunksIn a
x) = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks (forall {k} (t :: k). Proxy t
Proxy @s) Context
ctxt forall x. Rep a x
fp
    where
      fp :: Rep a x
      !fp :: forall x. Rep a x
fp = forall a x. Generic a => a -> Rep a x
from a
x

instance Typeable a => NoThunks (InspectHeap a) where
  showTypeOf :: Proxy (InspectHeap a) -> String
showTypeOf Proxy (InspectHeap a)
_ = forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy @a)
  wNoThunks :: Context -> InspectHeap a -> IO (Maybe ThunkInfo)
wNoThunks = forall a. Context -> a -> IO (Maybe ThunkInfo)
inspectHeap

instance KnownSymbol name => NoThunks (InspectHeapNamed name a) where
  showTypeOf :: Proxy (InspectHeapNamed name a) -> String
showTypeOf Proxy (InspectHeapNamed name a)
_ = forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (forall {k} (t :: k). Proxy t
Proxy @name)
  wNoThunks :: Context -> InspectHeapNamed name a -> IO (Maybe ThunkInfo)
wNoThunks = forall a. Context -> a -> IO (Maybe ThunkInfo)
inspectHeap

-- | Internal: implementation of 'wNoThunks' for 'InspectHeap'
-- and 'InspectHeapNamed'
inspectHeap :: Context -> a -> IO (Maybe ThunkInfo)
inspectHeap :: forall a. Context -> a -> IO (Maybe ThunkInfo)
inspectHeap Context
ctxt a
x = do
    Bool
containsThunks <- forall a. a -> IO Bool
checkContainsThunks a
x
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if Bool
containsThunks
               then forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ ThunkInfo { thunkContext :: Context
thunkContext = String
"..." forall a. a -> [a] -> [a]
: Context
ctxt }
               else forall a. Maybe a
Nothing

{-------------------------------------------------------------------------------
  Internal: generic infrastructure
-------------------------------------------------------------------------------}

-- | Generic infrastructure for checking for unexpected thunks
--
-- The @a@ argument records which record fields are allowed to contain thunks;
-- see 'AllowThunksIn' and 'GWRecordField', below.
class GWNoThunks (a :: [Symbol]) f where
  -- | Check that the argument does not contain any unexpected thunks
  --
  -- Precondition: the argument is in WHNF.
  gwNoThunks :: proxy a -> Context -> f x -> IO (Maybe ThunkInfo)

instance GWNoThunks a f => GWNoThunks a (D1 c f) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> D1 c f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt (M1 f x
fp) = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt f x
fp

instance GWNoThunks a f => GWNoThunks a (C1 c f) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> C1 c f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt (M1 f x
fp) = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt f x
fp

instance GWNoThunks a f => GWNoThunks a (S1 ('MetaSel ('Nothing) su ss ds) f) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a
-> Context
-> S1 ('MetaSel 'Nothing su ss ds) f x
-> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt (M1 f x
fp) = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt f x
fp

instance (GWNoThunks a f, GWNoThunks a g) => GWNoThunks a (f :*: g) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> (:*:) f g x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt (f x
fp :*: g x
gp) = [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks [
        forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt f x
fp
      , forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt g x
gp
      ]

instance (GWNoThunks a f, GWNoThunks a g) => GWNoThunks a (f :+: g) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> (:+:) f g x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt (L1 f x
fp) = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt f x
fp
  gwNoThunks proxy a
a Context
ctxt (R1 g x
gp) = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
a Context
ctxt g x
gp

instance NoThunks c => GWNoThunks a (K1 i c) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> K1 i c x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
_a Context
ctxt (K1 c
c) = forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt' c
c
    where
      -- If @c@ is a recursive occurrence of the type itself, we want to avoid
      -- accumulating context. For example, suppose we are dealing with @[Int]@,
      -- and we have an unexpected thunk as the third @Int@ in the list. If
      -- we use the generic instance, then without this correction, the final
      -- context will look something like
      --
      -- > ["Int", "[]", "[]", "[]"]
      --
      -- While that is more informative (it's the /third/ element that is a
      -- thunk), it's not that helpful (typically we just want /all/ elements
      -- to be in NF). We strip the context here so that we just get
      --
      -- > ["Int", "[]"]
      --
      -- which is a bit easier to interpret.
      ctxt' :: Context
ctxt' = case Context
ctxt of
                String
hd : Context
tl | String
hd forall a. Eq a => a -> a -> Bool
== forall a. NoThunks a => Proxy a -> String
showTypeOf (forall {k} (t :: k). Proxy t
Proxy @c) -> Context
tl
                Context
_otherwise                            -> Context
ctxt

instance GWNoThunks a U1 where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> U1 x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
_a Context
_ctxt U1 x
U1 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

instance GWNoThunks a V1 where
  -- By assumption, the argument is already in WHNF. Since every inhabitant of
  -- this type is bottom, this code is therefore unreachable.
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a -> Context -> V1 x -> IO (Maybe ThunkInfo)
gwNoThunks proxy a
_a Context
_ctxt V1 x
_ = forall a. HasCallStack => String -> a
error String
"unreachable gwNoThunks @V1"

{-------------------------------------------------------------------------------
  Skip fields with allowed thunks
-------------------------------------------------------------------------------}

-- | If @fieldName@ is allowed to contain thunks, skip it.
instance GWRecordField f (Elem fieldName a)
      => GWNoThunks a (S1 ('MetaSel ('Just fieldName) su ss ds) f) where
  gwNoThunks :: forall (proxy :: [Symbol] -> *) x.
proxy a
-> Context
-> S1 ('MetaSel ('Just fieldName) su ss ds) f x
-> IO (Maybe ThunkInfo)
gwNoThunks proxy a
_ Context
ctxt (M1 f x
fp) =
      forall (f :: * -> *) (b :: Bool) (proxy :: Bool -> *) x.
GWRecordField f b =>
proxy b -> Context -> f x -> IO (Maybe ThunkInfo)
gwRecordField (forall {k} (t :: k). Proxy t
Proxy @(Elem fieldName a)) Context
ctxt f x
fp

class GWRecordField f (b :: Bool) where
  gwRecordField :: proxy b -> Context -> f x -> IO (Maybe ThunkInfo)

-- | If the field is allowed to contain thunks, don't check anything.
instance GWRecordField f 'True where
  gwRecordField :: forall (proxy :: Bool -> *) x.
proxy 'True -> Context -> f x -> IO (Maybe ThunkInfo)
gwRecordField proxy 'True
_ Context
_ f x
_ = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

instance GWNoThunks '[] f => GWRecordField f 'False where
  gwRecordField :: forall (proxy :: Bool -> *) x.
proxy 'False -> Context -> f x -> IO (Maybe ThunkInfo)
gwRecordField proxy 'False
_ Context
ctxt f x
f = forall (a :: [Symbol]) (f :: * -> *) (proxy :: [Symbol] -> *) x.
GWNoThunks a f =>
proxy a -> Context -> f x -> IO (Maybe ThunkInfo)
gwNoThunks (forall {k} (t :: k). Proxy t
Proxy @'[]) Context
ctxt f x
f

{-------------------------------------------------------------------------------
  Internal: generic function to get name of a type
-------------------------------------------------------------------------------}

class GShowTypeOf f where
  gShowTypeOf :: f x -> String

instance Datatype c => GShowTypeOf (D1 c f) where
  gShowTypeOf :: forall x. D1 c f x -> String
gShowTypeOf = forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
datatypeName

{-------------------------------------------------------------------------------
  Instances for primitive types
-------------------------------------------------------------------------------}

deriving via OnlyCheckWhnf Bool    instance NoThunks Bool
deriving via OnlyCheckWhnf Natural instance NoThunks Natural
deriving via OnlyCheckWhnf Integer instance NoThunks Integer
deriving via OnlyCheckWhnf Float   instance NoThunks Float
deriving via OnlyCheckWhnf Double  instance NoThunks Double
deriving via OnlyCheckWhnf Char    instance NoThunks Char

deriving via OnlyCheckWhnf Int   instance NoThunks Int
deriving via OnlyCheckWhnf Int8  instance NoThunks Int8
deriving via OnlyCheckWhnf Int16 instance NoThunks Int16
deriving via OnlyCheckWhnf Int32 instance NoThunks Int32
deriving via OnlyCheckWhnf Int64 instance NoThunks Int64

deriving via OnlyCheckWhnf Word   instance NoThunks Word
deriving via OnlyCheckWhnf Word8  instance NoThunks Word8
deriving via OnlyCheckWhnf Word16 instance NoThunks Word16
deriving via OnlyCheckWhnf Word32 instance NoThunks Word32
deriving via OnlyCheckWhnf Word64 instance NoThunks Word64

{-------------------------------------------------------------------------------
  Mutable Vars
-------------------------------------------------------------------------------}

instance NoThunks a => NoThunks (IORef.IORef a) where
    showTypeOf :: Proxy (IORef a) -> String
showTypeOf Proxy (IORef a)
_ = String
"IORef"
    wNoThunks :: Context -> IORef a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctx IORef a
ref = do
        a
val <- forall a. IORef a -> IO a
IORef.readIORef IORef a
ref
        forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx a
val

instance NoThunks a => NoThunks (MVar.MVar a) where
    showTypeOf :: Proxy (MVar a) -> String
showTypeOf Proxy (MVar a)
_ = String
"MVar"
    wNoThunks :: Context -> MVar a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctx MVar a
ref = do
        Maybe a
val <- forall a. MVar a -> IO (Maybe a)
MVar.tryReadMVar MVar a
ref
        forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing) (forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx) Maybe a
val

instance NoThunks a => NoThunks (TVar.TVar a) where
    showTypeOf :: Proxy (TVar a) -> String
showTypeOf Proxy (TVar a)
_ = String
"TVar"
    wNoThunks :: Context -> TVar a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctx TVar a
ref = do
        -- An alternative is to use
        --
        --   val <- STM.atomically $ TVar.readTVar ref
        --
        -- but that would cause nested atomically failures with
        -- unsafeNoThunks. Fortunately, readTVarIO doesn't make a transaction.
        --
        -- See related tests.
        --
        a
val <- forall a. TVar a -> IO a
TVar.readTVarIO TVar a
ref
        forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx a
val

{-------------------------------------------------------------------------------
  Time
-------------------------------------------------------------------------------}

deriving via InspectHeap Day              instance NoThunks Day
deriving via InspectHeap DiffTime         instance NoThunks DiffTime
deriving via InspectHeap LocalTime        instance NoThunks LocalTime
deriving via InspectHeap NominalDiffTime  instance NoThunks NominalDiffTime
deriving via InspectHeap TimeLocale       instance NoThunks TimeLocale
deriving via InspectHeap TimeOfDay        instance NoThunks TimeOfDay
deriving via InspectHeap TimeZone         instance NoThunks TimeZone
deriving via InspectHeap UniversalTime    instance NoThunks UniversalTime
deriving via InspectHeap UTCTime          instance NoThunks UTCTime
deriving via InspectHeap ZonedTime        instance NoThunks ZonedTime

{-------------------------------------------------------------------------------
  ByteString
-------------------------------------------------------------------------------}

#ifdef MIN_VERSION_bytestring

-- | Instance for string bytestrings
--
-- Strict bytestrings /shouldn't/ contain any thunks, but could, due to
-- <https://gitlab.haskell.org/ghc/ghc/issues/17290>. However, such thunks can't
-- retain any data that they shouldn't, and so it's safe to ignore such thunks.
deriving via OnlyCheckWhnfNamed "Strict.ByteString" BS.Strict.ByteString
         instance NoThunks BS.Strict.ByteString

-- | Instance for short bytestrings
--
-- We have
--
-- > data ShortByteString = SBS ByteArray#
--
-- Values of this type consist of a tag followed by an _unboxed_ byte array,
-- which can't contain thunks. Therefore we only check WHNF.
deriving via OnlyCheckWhnfNamed "ShortByteString" ShortByteString
         instance NoThunks ShortByteString

-- | Instance for lazy bytestrings
--
-- Defined manually so that it piggy-backs on the one for strict bytestrings.
instance NoThunks BS.Lazy.ByteString where
  showTypeOf :: Proxy ByteString -> String
showTypeOf Proxy ByteString
_      = String
"Lazy.ByteString"
  wNoThunks :: Context -> ByteString -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt ByteString
bs =
      case ByteString
bs of
        ByteString
BS.Lazy.Internal.Empty           -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
        BS.Lazy.Internal.Chunk ByteString
chunk ByteString
bs' -> [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks [
              forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt ByteString
chunk
            , forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt ByteString
bs'
            ]

#endif

{-------------------------------------------------------------------------------
  Instances for text types

  For consistency, we follow the same pattern as for @ByteString@.
-------------------------------------------------------------------------------}

#ifdef MIN_VERSION_text

deriving via OnlyCheckWhnfNamed "Strict.Text" Text.Strict.Text
         instance NoThunks Text.Strict.Text

instance NoThunks Text.Lazy.Text where
  showTypeOf :: Proxy Text -> String
showTypeOf Proxy Text
_      = String
"Lazy.Text"
  wNoThunks :: Context -> Text -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt Text
bs =
      case Text
bs of
        Text
Text.Lazy.Internal.Empty           -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
        Text.Lazy.Internal.Chunk Text
chunk Text
bs' -> [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks [
              forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt Text
chunk
            , forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt Text
bs'
            ]

#endif

{-------------------------------------------------------------------------------
  Tuples
-------------------------------------------------------------------------------}

instance ( NoThunks a
         , NoThunks b
         ) => NoThunks (a, b)

instance ( NoThunks a
         , NoThunks b
         , NoThunks c
         ) => NoThunks (a, b, c)

instance ( NoThunks a
         , NoThunks b
         , NoThunks c
         , NoThunks d
         ) => NoThunks (a, b, c, d)

instance ( NoThunks a
         , NoThunks b
         , NoThunks c
         , NoThunks d
         , NoThunks e
         ) => NoThunks (a, b, c, d, e)

instance ( NoThunks a
         , NoThunks b
         , NoThunks c
         , NoThunks d
         , NoThunks e
         , NoThunks f
         ) => NoThunks (a, b, c, d, e, f)

instance ( NoThunks a
         , NoThunks b
         , NoThunks c
         , NoThunks d
         , NoThunks e
         , NoThunks f
         , NoThunks g
         ) => NoThunks (a, b, c, d, e, f, g)

{-------------------------------------------------------------------------------
  Base types (other than tuples)
-------------------------------------------------------------------------------}

instance NoThunks Void
instance NoThunks ()

instance NoThunks a => NoThunks [a]
instance NoThunks a => NoThunks (Maybe a)
instance NoThunks a => NoThunks (NonEmpty a)

instance (NoThunks a, NoThunks b) => NoThunks (Either a b)

{-------------------------------------------------------------------------------
  Spine-strict container types

  Such types can /only/ contain thunks in the values, so that's all we check.
  Note that containers using keys are typically strict in those keys, but that
  forces them to WHNF only, not NF; in /most/ cases the @Ord@ instance on those
  keys will force them to NF, but not /always/ (for example, when using lists
  as keys); this means we must check keys for thunks to be sure.
-------------------------------------------------------------------------------}

instance (NoThunks k, NoThunks v) => NoThunks (Map k v) where
  showTypeOf :: Proxy (Map k v) -> String
showTypeOf Proxy (Map k v)
_   = String
"Map"
  wNoThunks :: Context -> Map k v -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt = forall k v.
(NoThunks k, NoThunks v) =>
Context -> [(k, v)] -> IO (Maybe ThunkInfo)
noThunksInKeysAndValues Context
ctxt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toList

instance NoThunks a => NoThunks (Set a) where
  showTypeOf :: Proxy (Set a) -> String
showTypeOf Proxy (Set a)
_   = String
"Set"
  wNoThunks :: Context -> Set a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt = forall a. NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues Context
ctxt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
Set.toList

instance NoThunks a => NoThunks (IntMap a) where
  showTypeOf :: Proxy (IntMap a) -> String
showTypeOf Proxy (IntMap a)
_   = String
"IntMap"
  wNoThunks :: Context -> IntMap a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt = forall a. NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues Context
ctxt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntMap a -> [(Int, a)]
IntMap.toList

{-------------------------------------------------------------------------------
  Vector
-------------------------------------------------------------------------------}

#ifdef MIN_VERSION_vector

instance NoThunks a => NoThunks (Vector.Boxed.Vector a) where
  showTypeOf :: Proxy (Vector a) -> String
showTypeOf Proxy (Vector a)
_   = String
"Boxed.Vector"
  wNoThunks :: Context -> Vector a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt = forall a. NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues Context
ctxt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Vector a -> [a]
Vector.Boxed.toList

-- | Unboxed vectors can't contain thunks
--
-- Implementation note: defined manually rather than using 'OnlyCheckWhnf'
-- due to ghc limitation in deriving via, making it impossible to use with it
-- with data families.
instance NoThunks (Vector.Unboxed.Vector a) where
  showTypeOf :: Proxy (Vector a) -> String
showTypeOf Proxy (Vector a)
_  = String
"Unboxed.Vector"
  wNoThunks :: Context -> Vector a -> IO (Maybe ThunkInfo)
wNoThunks Context
_ Vector a
_ = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

#endif

{-------------------------------------------------------------------------------
  Function types
-------------------------------------------------------------------------------}

-- | We do NOT check function closures for captured thunks by default
--
-- Since we have no type information about the values captured in a thunk, the
-- only check we could possibly do is 'checkContainsThunks': we can't
-- recursively call 'noThunks' on those captured values, which is problematic if
-- any of those captured values /requires/ a custom instance (for example, data
-- types that depend on laziness, such as 'Seq').
--
-- By default we therefore /only/ check if the function is in WHNF, and don't
-- check the captured values at all. If you want a stronger check, you can
-- use @'InspectHeap' (a -> b)@ instead.
deriving via OnlyCheckWhnfNamed "->" (a -> b) instance NoThunks (a -> b)

-- | We do not check IO actions for captured thunks by default
--
-- See instance for @(a -> b)@ for detailed discussion.
deriving via OnlyCheckWhnfNamed "IO" (IO a) instance NoThunks (IO a)

{-------------------------------------------------------------------------------
  Special cases
-------------------------------------------------------------------------------}

-- | Since CallStacks can't retain application data, we don't want to check
-- them for thunks /at all/
deriving via AllowThunk CallStack instance NoThunks CallStack

-- | Instance for 'Seq' checks elements only
--
-- The internal fingertree in 'Seq' might have thunks, which is essential for
-- its asymptotic complexity.
instance NoThunks a => NoThunks (Seq a) where
  showTypeOf :: Proxy (Seq a) -> String
showTypeOf Proxy (Seq a)
_ = String
"Seq"
  wNoThunks :: Context -> Seq a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt = forall a. NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues Context
ctxt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

instance NoThunks a => NoThunks (Ratio a) where
  showTypeOf :: Proxy (Ratio a) -> String
showTypeOf Proxy (Ratio a)
_ = String
"Ratio"
  wNoThunks :: Context -> Ratio a -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt Ratio a
r = forall a. NoThunks a => Context -> [a] -> IO (Maybe ThunkInfo)
noThunksInValues Context
ctxt [a
n, a
d]
   where
     -- The 'Ratio' constructor is not exported: we only have two accessor
     -- functions. However, @numerator r@ is obviously trivially a trunk
     -- (due to the unevaluated call to @numerator@). By forcing the values of
     -- @n@ and @d@ where we get rid of these function calls, leaving only the
     -- values inside the @Ratio@. Note that @Ratio@ is strict in both of these
     -- fields, so forcing them to WHNF won't change them.
     !n :: a
n = forall a. Ratio a -> a
numerator   Ratio a
r
     !d :: a
d = forall a. Ratio a -> a
denominator Ratio a
r

{-------------------------------------------------------------------------------
  Type level symbol comparison logic
-------------------------------------------------------------------------------}

type family Same s t where
  Same s t = IsSame (CmpSymbol s t)

type family IsSame (o :: Ordering) where
  IsSame 'EQ = 'True
  IsSame _x  = 'False

type family Or (a :: Bool) (b :: Bool) where
  Or 'False 'False = 'False
  Or _a     _b     = 'True

type family Elem (s :: Symbol) (xs :: [Symbol]) where
  Elem s  (x ': xs) = Or (Same s x) (Elem s xs)
  Elem _s '[]       = 'False

{-------------------------------------------------------------------------------
  Check that all mentioned record fields are known fields
-------------------------------------------------------------------------------}

-- | Check that type @a@ has all record fields listed in @s@
--
-- This exists to catch mismatches between the arguments to `AllowThunksIn` and
-- the fields of a record. If any of the symbols is not the name of a field then
-- this constraint won't be satisfied.
class HasFields (s :: [Symbol]) (a :: Type)
instance HasFields '[] a
instance (HasField x a t, HasFields xs a) => HasFields (x ': xs) a

{-------------------------------------------------------------------------------
  Internal: low level magic
-------------------------------------------------------------------------------}

-- | Is the argument a (top-level thunk)?
checkIsThunk :: a -> IO Bool
checkIsThunk :: forall a. a -> IO Bool
checkIsThunk a
x = Closure -> Bool
closureIsThunk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Box -> IO Closure
getBoxedClosureData (forall a. a -> Box
asBox a
x)

-- | Is the argument a thunk, or does it (recursively) contain any?
checkContainsThunks :: a -> IO Bool
checkContainsThunks :: forall a. a -> IO Bool
checkContainsThunks a
x = Box -> IO Bool
go (forall a. a -> Box
asBox a
x)
  where
    go :: Box -> IO Bool
    go :: Box -> IO Bool
go Box
b = do
        Closure
c <- Box -> IO Closure
getBoxedClosureData Box
b
        if Closure -> Bool
closureIsThunk Closure
c then
          forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        else do
          Closure
c' <- Box -> IO Closure
getBoxedClosureData Box
b
          forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyM Box -> IO Bool
go (forall b. GenClosure b -> [b]
allClosures Closure
c')

-- | Check if the given 'Closure' is a thunk.
--
-- Indirections are not considered to be thunks.
closureIsThunk :: Closure -> Bool
closureIsThunk :: Closure -> Bool
closureIsThunk ThunkClosure{}    = Bool
True
closureIsThunk APClosure{}       = Bool
True
closureIsThunk SelectorClosure{} = Bool
True
closureIsThunk BCOClosure{}      = Bool
True
closureIsThunk Closure
_                 = Bool
False

anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool
anyM :: forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
_ []       = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
anyM a -> m Bool
p (a
x : [a]
xs) = do
    Bool
q <- a -> m Bool
p a
x
    if Bool
q then forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
         else forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
p [a]
xs