<commit_before>{-# LANGUAGE
  TupleSections,
  GeneralizedNewtypeDeriving,
  FlexibleInstances,
  FlexibleContexts
 #-}

module Particle (
    ParticleT,
    synchronize,
    flatten,
    advance,
    finished,
    mapMonad
                ) where

import Control.Monad.Trans.Class
import Control.Monad (liftM2)
import Control.Monad.Coroutine
import Control.Monad.Coroutine.SuspensionFunctors
import Data.Either

import Base

-- | Particle represents a computation that can be paused at certain points.
-- The intermediate monadic effects can be extracted, which is particularly useful
-- for implementation of SMC-related methods.
-- All the probabilistic effects are delegated to the transformed monad,
-- but also `synchronize` is inserted after each `factor`.
type ParticleT m a = Coroutine (Await ()) m a
extract (Await f) = f ()

-- | A synchronization barrier where computation is paused.
synchronize :: Monad m => ParticleT m ()
synchronize = await

-- | Removes the synchronization barriers.
flatten :: Monad m => ParticleT m a -> m a
flatten = pogoStick extract

-- | Run a particle to the next barrier.
-- If the computation is finished do nothing.
advance :: Monad m => ParticleT m a -> ParticleT m a
advance = bounce extract

-- | Checks if the particle is finished.
finished :: Monad m => ParticleT m a -> m Bool
finished = fmap isRight . resume

instance MonadDist m => MonadDist (Coroutine (Await ()) m) where
    primitive = lift . primitive

instance MonadBayes m => MonadBayes (Coroutine (Await ()) m) where
    factor w = lift (factor w) >> synchronize
<commit_msg>Fixed Particle.mapMonad not the apply the map recursively<commit_after>{-# LANGUAGE
  TupleSections,
  GeneralizedNewtypeDeriving,
  FlexibleInstances,
  FlexibleContexts,
  RankNTypes
 #-}

module Particle (
    ParticleT,
    synchronize,
    flatten,
    advance,
    finished,
    Particle.mapMonad
                ) where

import Control.Monad.Trans.Class
import Control.Monad (liftM2)
import Control.Monad.Coroutine
import Control.Monad.Coroutine.SuspensionFunctors
import Data.Either

import Base

-- | Particle represents a computation that can be paused at certain points.
-- The intermediate monadic effects can be extracted, which is particularly useful
-- for implementation of SMC-related methods.
-- All the probabilistic effects are delegated to the transformed monad,
-- but also `synchronize` is inserted after each `factor`.
type ParticleT m a = Coroutine (Await ()) m a
extract (Await f) = f ()

-- | A synchronization barrier where computation is paused.
synchronize :: Monad m => ParticleT m ()
synchronize = await

-- | Removes the synchronization barriers.
flatten :: Monad m => ParticleT m a -> m a
flatten = pogoStick extract

-- | Run a particle to the next barrier.
-- If the computation is finished do nothing.
advance :: Monad m => ParticleT m a -> ParticleT m a
advance = bounce extract

-- | Checks if the particle is finished.
finished :: Monad m => ParticleT m a -> m Bool
finished = fmap isRight . resume

mapMonad :: Monad m =>
            (forall a. m a -> m a) -> ParticleT m a -> ParticleT m a
mapMonad f cort = Coroutine {resume= f $ resume cort}

instance MonadDist m => MonadDist (Coroutine (Await ()) m) where
    primitive = lift . primitive

instance MonadBayes m => MonadBayes (Coroutine (Await ()) m) where
    factor w = lift (factor w) >> synchronize
{'fraction_non_alphanumeric': 0.08483594442802246, 'fraction_numerical': 0.0005911912503694946, 'mean_word_length': 3.542281879194631, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2433737', 'n_tokens_mistral': 1017, 'n_tokens_neox': 1005, 'n_words': 397}