<commit_before>-- Language extensions {{{
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UnicodeSyntax #-}
-- }}}

module Control.Visitor -- {{{
    ( MonadVisitor(..)
    , MonadVisitorTrans(..)
    , VisitorTInstruction(..)
    , VisitorInstruction
    , Visitor
    , VisitorIO
    , VisitorT(..)
    , allFrom
    , allFromBalanced
    , allFromBalancedGreedy
    , between
    , endowVisitor
    , msumBalanced
    , msumBalancedGreedy
    , runVisitor
    , runVisitorT
    , runVisitorTAndIgnoreResults
    , runVisitorUntilFirst
    , runVisitorTUntilFirst
    , runVisitorUntilFound
    , runVisitorTUntilFound
    ) where -- }}}

-- Imports {{{
import Control.Applicative (Alternative(..),Applicative(..))
import Control.Monad (MonadPlus(..),(>=>),guard,liftM,liftM2,msum)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Operational (ProgramT,ProgramViewT(..),singleton,view,viewT)
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.List (ListT)
import Control.Monad.Trans.Maybe (MaybeT)

import Data.Functor.Identity (Identity(..),runIdentity)
import Data.Maybe (isJust)
import Data.Monoid ((<>),Monoid(..))
import Data.Serialize (Serialize(),encode)

import Control.Visitor.Utils.MonadStacks
-- }}}

-- Classes {{{

class MonadPlus m ⇒ MonadVisitor m where -- {{{
    cache :: Serialize x ⇒ x → m x
    cache = cacheMaybe . Just

    cacheGuard :: Bool → m ()
    cacheGuard = cacheMaybe . (\x → if x then Just () else Nothing)

    cacheMaybe :: Serialize x ⇒ Maybe x → m x
-- }}}

class (MonadPlus m, Monad (NestedMonadInVisitor m)) ⇒ MonadVisitorTrans m where -- {{{
    type NestedMonadInVisitor m :: * → *
    runAndCache :: Serialize x ⇒ (NestedMonadInVisitor m) x → m x
    runAndCache = runAndCacheMaybe . liftM Just

    runAndCacheGuard :: (NestedMonadInVisitor m) Bool → m ()
    runAndCacheGuard = runAndCacheMaybe . liftM (\x → if x then Just () else Nothing)

    runAndCacheMaybe :: Serialize x ⇒ (NestedMonadInVisitor m) (Maybe x) → m x
-- }}}

-- }}}

-- Types {{{

data VisitorTInstruction m α where -- {{{
    Cache :: Serialize α ⇒ m (Maybe α) → VisitorTInstruction m α
    Choice :: VisitorT m α → VisitorT m α → VisitorTInstruction m α
    Null :: VisitorTInstruction m α
-- }}}
type VisitorInstruction = VisitorTInstruction Identity

newtype VisitorT m α = VisitorT { unwrapVisitorT :: ProgramT (VisitorTInstruction m) m α }
    deriving (Applicative,Functor,Monad,MonadIO)
type Visitor = VisitorT Identity
type VisitorIO = VisitorT IO

-- }}}

-- Instances {{{

instance Monad m ⇒ Alternative (VisitorT m) where -- {{{
    empty = mzero
    (<|>) = mplus
-- }}}

instance Eq α ⇒ Eq (Visitor α) where -- {{{
    (VisitorT x) == (VisitorT y) = e x y
      where
        e x y = case (view x, view y) of
            (Return x, Return y) → x == y
            (Null :>>= _, Null :>>= _) → True
            (Cache cx :>>= kx, Cache cy :>>= ky) →
                case (runIdentity cx, runIdentity cy) of
                    (Nothing, Nothing) → True
                    (Just x, Just y) → e (kx x) (ky y)
                    _ → False
            (Choice (VisitorT ax) (VisitorT bx) :>>= kx, Choice (VisitorT ay) (VisitorT by) :>>= ky) →
                e (ax >>= kx) (ay >>= ky) && e (bx >>= kx) (by >>= ky)
            _  → False
-- }}}

instance Monad m ⇒ MonadPlus (VisitorT m) where -- {{{
    mzero = VisitorT . singleton $ Null
    left `mplus` right = VisitorT . singleton $ Choice left right
-- }}}

instance MonadVisitor [] where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

instance Monad m ⇒ MonadVisitor (ListT m) where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

instance Monad m ⇒ MonadVisitorTrans (ListT m) where -- {{{
    type NestedMonadInVisitor (ListT m) = m
-- }}}
    runAndCacheMaybe = lift >=> maybe mzero return

instance MonadVisitor Maybe where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

instance Monad m ⇒ MonadVisitor (MaybeT m) where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

instance Monad m ⇒ MonadVisitorTrans (MaybeT m) where -- {{{
    type NestedMonadInVisitor (MaybeT m) = m
    runAndCacheMaybe = maybe mzero lift
-- }}}

instance Monad m ⇒ MonadVisitor (VisitorT m) where -- {{{
    cache = runAndCache . return
    cacheGuard = runAndCacheGuard . return
    cacheMaybe = runAndCacheMaybe . return
-- }}}

instance Monad m ⇒ MonadVisitorTrans (VisitorT m) where -- {{{
    type NestedMonadInVisitor (VisitorT m) = m
    runAndCache = runAndCacheMaybe . liftM Just
    runAndCacheGuard = runAndCacheMaybe . liftM (\x → if x then Just () else Nothing)
    runAndCacheMaybe = VisitorT . singleton . Cache
-- }}}

instance (MonadTrans t, MonadVisitor m, MonadPlus (t m)) ⇒ MonadVisitor (t m) where -- {{{
    cache = lift . cache
    cacheGuard = lift . cacheGuard
    cacheMaybe = lift . cacheMaybe
-- }}}

instance MonadTrans VisitorT where -- {{{
    lift = VisitorT . lift
-- }}}

instance Monad m ⇒ Monoid (VisitorT m α) where -- {{{
    mempty = mzero
    mappend = mplus
    mconcat = msum
-- }}}

instance Show α ⇒ Show (Visitor α) where -- {{{
    show = s . unwrapVisitorT
      where<commit_msg>Added docs for the type classes and types in Visitor.<commit_after>-- Language extensions {{{
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UnicodeSyntax #-}
-- }}}

module Control.Visitor -- {{{
    (
    -- * Visitor features
    -- $type-classes
      MonadVisitor(..)
    , MonadVisitorTrans(..)
    -- * Visitor types
    -- $types
    , Visitor
    , VisitorIO
    , VisitorT(..)
    -- * Rest
    , VisitorTInstruction(..)
    , VisitorInstruction
    , allFrom
    , allFromBalanced
    , allFromBalancedGreedy
    , between
    , endowVisitor
    , msumBalanced
    , msumBalancedGreedy
    , runVisitor
    , runVisitorT
    , runVisitorTAndIgnoreResults
    , runVisitorUntilFirst
    , runVisitorTUntilFirst
    , runVisitorUntilFound
    , runVisitorTUntilFound
    ) where -- }}}

-- Imports {{{
import Control.Applicative (Alternative(..),Applicative(..))
import Control.Monad (MonadPlus(..),(>=>),guard,liftM,liftM2,msum)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Operational (ProgramT,ProgramViewT(..),singleton,view,viewT)
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.List (ListT)
import Control.Monad.Trans.Maybe (MaybeT)

import Data.Functor.Identity (Identity(..),runIdentity)
import Data.Maybe (isJust)
import Data.Monoid ((<>),Monoid(..))
import Data.Serialize (Serialize(),encode)

import Control.Visitor.Utils.MonadStacks
-- }}}

-- Classes {{{

-- Doc: Visitor Features {{{
{- $type-classes

Visitors are instances of 'MonadVisitor' and/or 'MonadVisitorTrans', which are
both subclasses of 'MonadPlus'.  The additional functionality offered by these
type-classes is the ability to cache results so that a computation does not need
to be repeated when a path through the visitor is taken a second time, which
happens either when resuming from a checkpoint or when a workload has been
stolen by another processor as the first step is to retrace the path through the
visitor that leads to the stolen workload.

These features could have been provided as functions, but there are two reasons
why they were subsumed into type-classes: first, because one might want to
add another layer above the 'Visitor' monad transformers in the monad stack
(as is the case in "Control.Visitor.Label"), and second, because one might want
to run a visitor using a simpler monad such as [] for testing purposes.

NOTE:  Caching a computation takes space in the 'Checkpoint', so it is something
       you should only do when the result is relatively small and the
       computation is very expensive and is high enough in the search tree that
       it is likely to be repeated often.  If the calculation is low enough in
       the search tree that it is unlikely to be repeated, is cheap enough so
       that repeating it is not a big deal, or produces a result with an
       incredibly large memory footprint, then you are probably better off not
       caching the result.
 -}
-- }}}

{-| The 'MonadVisitor' class provides caching functionality when running a pure
    visitor;  at minimum 'cacheMaybe' needs to be defined.
 -}
class MonadPlus m ⇒ MonadVisitor m where -- {{{
    {-| Cache a value in case we visit this node again. -}
    cache :: Serialize x ⇒ x → m x
    cache = cacheMaybe . Just

    {-| This does the same thing as 'guard' but it caches the result. -}
    cacheGuard :: Bool → m ()
    cacheGuard = cacheMaybe . (\x → if x then Just () else Nothing)

    {-| This function is a combination of the previous two;  it performs a
        computation which might fail by returning 'Nothing', and if that happens
        it aborts the visitor;  if it passes then the result is cached and
        returned.

        Note that the previous two methods are essentially specializations of
        this method.
     -}
    cacheMaybe :: Serialize x ⇒ Maybe x → m x
-- }}}

{-| This class is like 'MonadVisitor', but it is designed to work with monad
    stacks;  at minimum 'runAndCacheMaybe' needs to be defined.
 -}
class (MonadPlus m, Monad (NestedMonadInVisitor m)) ⇒ MonadVisitorTrans m where -- {{{
    {-| The next layer down in the monad transformer stack. -}
    type NestedMonadInVisitor m :: * → *

    {-| Runs the given action in the nested monad and caches the result. -}
    runAndCache :: Serialize x ⇒ (NestedMonadInVisitor m) x → m x
    runAndCache = runAndCacheMaybe . liftM Just

    {-| Runs the given action in the nested monad and then does the equivalent
        of feeding it into 'guard', caching the result.
     -}
    runAndCacheGuard :: (NestedMonadInVisitor m) Bool → m ()
    runAndCacheGuard = runAndCacheMaybe . liftM (\x → if x then Just () else Nothing)

    {-| Runs the given action in the nested monad;  if it returns 'Nothing',
        then it acts like 'mzero',  if it returns 'Just x', then it caches the
        result.
     -}
    runAndCacheMaybe :: Serialize x ⇒ (NestedMonadInVisitor m) (Maybe x) → m x
-- }}}

-- }}}

-- Types {{{

data VisitorTInstruction m α where -- {{{
    Cache :: Serialize α ⇒ m (Maybe α) → VisitorTInstruction m α
    Choice :: VisitorT m α → VisitorT m α → VisitorTInstruction m α
    Null :: VisitorTInstruction m α
-- }}}
type VisitorInstruction = VisitorTInstruction Identity

-- Doc: Visitor types {{{
{- $types
The following are the visitor types that are accepted by most of he functions in
this package.  You do not need to know the details of their definitions unless
you intend to write your own custom routines for running and transforming
visitors, in which case the relevant information is at the bottom of this page
in the Implementation` section.

There is one type of pure visitor and two types of impure visitors.  In general,
your visitor should nearly always be pure if you are planning to make use of
checkpointing or parallel runs, because in general parts of the visitor will be
run multiple times, some parts may not be run at all on a given processor, and
whenever a leaf is hit there will be a jump to a higher node, so if your visitor
is impure the effects need to be meaningful no matter how the visitor is run on
a given processor.

Having said that, there are a few times when an impure visitor can make sense:
first, if the inner monad is something like the `Reader` monad, which has no
side-effects;  second, for testing purposes (e.g., many of my tests of the
various visitor runners use `MVar`s and the like to ensure that visitors are
explored in a certain way to test certain code paths);  finally, if there is some
side-effectful action that you want to run on each result (such as storing a
result into a database), though in this case you will need to make sure that
your code is robust against being run multiple times as there is no guarantee in
an environment where the system might be shut down and resumed from a checkpoint
that your action will only have been run once on a given result (i.e., if the
system goes down after your action was run but before a checkpoint was made
marking that it was run).

If you need something like state in your visitor, then you should consider
nesting the visitor monad in the state monad rather than vice-versa, because
this will do things like automatically erasing the change in state that happened
between an inner node and a leaf when the visitor jumps back up from the leaf
to an inner node.
-}
-- }}}

{-| A pure visitor, which is what you should normally be using. -}
type Visitor = VisitorT Identity

{-| A visitor running in the I/O monad, which you should only be using for
    testing purposes or, say, if you are planning on storing each result in an
    external database, in which case you need to guard against the possibility
    that an action for a given result might be run twice in checkpointing and/or
    parallel settings.
-}
type VisitorIO = VisitorT IO

{-| A visitor run in an arbitrary monad. -}
newtype VisitorT m α = VisitorT { unwrapVisitorT :: ProgramT (VisitorTInstruction m) m α }
    deriving (Applicative,Functor,Monad,MonadIO)

-- }}}

-- Instances {{{

{-| The 'Alternative' instance functions like the 'MonadPlus' instance. -}
instance Monad m ⇒ Alternative (VisitorT m) where -- {{{
    empty = mzero
    (<|>) = mplus
-- }}}

{-| Two visitors are equal if the generate exactly the same tree. -}
instance Eq α ⇒ Eq (Visitor α) where -- {{{
    (VisitorT x) == (VisitorT y) = e x y
      where
        e x y = case (view x, view y) of
            (Return x, Return y) → x == y
            (Null :>>= _, Null :>>= _) → True
            (Cache cx :>>= kx, Cache cy :>>= ky) →
                case (runIdentity cx, runIdentity cy) of
                    (Nothing, Nothing) → True
                    (Just x, Just y) → e (kx x) (ky y)
                    _ → False
            (Choice (VisitorT ax) (VisitorT bx) :>>= kx, Choice (VisitorT ay) (VisitorT by) :>>= ky) →
                e (ax >>= kx) (ay >>= ky) && e (bx >>= kx) (by >>= ky)
            _  → False
-- }}}

{-| For this type, 'mplus' provides a branch node with a choice between two
    values and 'mzero' provides a node with no values.
 -}
instance Monad m ⇒ MonadPlus (VisitorT m) where -- {{{
    mzero = VisitorT . singleton $ Null
    left `mplus` right = VisitorT . singleton $ Choice left right
-- }}}

{-| This instance performs no caching but is provided to make it easier to test
    running a visitor using the List monad.
 -}
instance MonadVisitor [] where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

{-| This instance performs no caching but is provided to make it easier to test
    running a visitor using the 'ListT' monad.
 -}
instance Monad m ⇒ MonadVisitor (ListT m) where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

{-| Like the 'MonadVisitor' isntance, this instance does no caching. -}
instance Monad m ⇒ MonadVisitorTrans (ListT m) where -- {{{
    type NestedMonadInVisitor (ListT m) = m
-- }}}
    runAndCacheMaybe = lift >=> maybe mzero return

{-| This instance performs no caching but is provided to make it easier to test
    running a visitor using the 'Maybe' monad.
 -}
instance MonadVisitor Maybe where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

{-| This instance performs no caching but is provided to make it easier to test
    running a visitor using the 'MaybeT' monad.
 -}
instance Monad m ⇒ MonadVisitor (MaybeT m) where -- {{{
    cacheMaybe = maybe mzero return
-- }}}

{-| Like the 'MonadVisitor' isntance, this instance does no caching. -}
instance Monad m ⇒ MonadVisitorTrans (MaybeT m) where -- {{{
    type NestedMonadInVisitor (MaybeT m) = m
    runAndCacheMaybe = maybe mzero lift
-- }}}

instance Monad m ⇒ MonadVisitor (VisitorT m) where -- {{{
    cache = runAndCache . return
    cacheGuard = runAndCacheGuard . return
    cacheMaybe = runAndCacheMaybe . return
-- }}}

instance Monad m ⇒ MonadVisitorTrans (VisitorT m) where -- {{{
    type NestedMonadInVisitor (VisitorT m) = m
    runAndCache = runAndCacheMaybe . liftM Just
    runAndCacheGuard = runAndCacheMaybe . liftM (\x → if x then Just () else Nothing)
    runAndCacheMaybe = VisitorT . singleton . Cache
-- }}}

{-| This instance allows you to automatically get a MonadVisitor instance for
    any monad transformer that has `MonadPlus` defined.  (Unfortunately its
    presence requires OverlappingInstances because it overlaps with the instance
    for `VisitorT`, even though the constraints are such that it is impossible
    in practice for there to ever be a case where a given type is satisfied by
    both instances.)
 -}
instance (MonadTrans t, MonadVisitor m, MonadPlus (t m)) ⇒ MonadVisitor (t m) where -- {{{
    cache = lift . cache
    cacheGuard = lift . cacheGuard
    cacheMaybe = lift . cacheMaybe
-- }}}

instance MonadTrans VisitorT where -- {{{
    lift = VisitorT . lift
-- }}}

{-| The 'Monoid' instance acts like the 'MonadPlus' instance. -}
instance Monad m ⇒ Monoid (VisitorT m α) where -- {{{
    mempty = mzero
    mappend = mplus
    mconcat = msum
-- }}}

instance Show α ⇒ Show (Visitor α) where -- {{{
    show = s . unwrapVisitorT
      where
{'fraction_non_alphanumeric': 0.10195878623188406, 'fraction_numerical': 0.00011322463768115942, 'mean_word_length': 3.023001594169893, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 55, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1501596', 'n_tokens_mistral': 5422, 'n_tokens_neox': 5349, 'n_words': 2176}