{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- The ProvM is a Monad State Transformer, designed for computing provenance
(or metadata) about an underlying computation. By running a computation with type
(ProvM meta m a) one gets the underlying computation of type (m a). Switches
control whether the metadata is computed or not.
-}
module Control.Provenance (
  ProvM (..),

  -- * Basic Provenance Monad Transformer
  PObject,

  -- * Dynamically typed Provenance objects
  Provenance,

  -- * Type constraint on data stored in PObject
  Prov,

  -- * ProvM instantiated (Map Text PObject)
  BlackBox,

  -- * Abstraction barrier to isolate the provenence type from the result type.

  -- * Operations in ProvM
  lift,
  putM,
  getM,
  modifyM,
  modifyWithBlackBox,
  runProv,
  runWithProv,
  runOtherProv,
  liftProv,
  dump,

  -- * Operations in Prov instantiation
  store,
  push,
  pull,
  update,
  updateWithBlackBox,
  pushOtherProv,
  runWithProvM,
  runProvM,

  -- * Operation on PObject
  find,
  observe,

  -- * For testing invariants
  preservesNothing,
  preservesJust,
) where

import Control.Monad.State.Strict (MonadState (..), MonadTrans (..), StateT (..))
import Data.Aeson (ToJSON (..))
import Data.Map.Strict (Map, empty, insert)
import qualified Data.Map.Strict as Map
import Data.Maybe.Strict
import Data.Text (Text, unpack)
import Data.Type.Equality (TestEquality (testEquality))
import NoThunks.Class (NoThunks (..), allNoThunks)
import Type.Reflection (TypeRep, Typeable, typeOf, typeRep, (:~:) (Refl))

-- ======================================================

-- |
-- The Provenance Monad is just the StateT transformer
-- wrapped in a newtype, where the state is (StrictMaybe t).
-- By running the state transformer with an initial state of SNothing
-- we indicate we are not collecting provenance. With an initial
-- state of (SJust t), we are collecting provenance. If we start off
-- as SNothing, then actions on the state are no-ops and are ignored.
-- We maintain several invariants. If we start with SNothing, we get
-- SNothing, and if we start with (SJust s) we end up with (SJust t)
--
-- invariant1 (ProvM m) = do (_,SNothing) <- runStateT m SNothing
-- invariant2 (ProvM m) = do (_,SJust t) <- runStateT m (SJust 1)
--
-- The second invariant is that regardless of whether we compute
-- provenance or not, the non-provenance part is the same.
-- Currently this is enforced by the BlackBox type and its API.
--
-- invariant3 (ProvM m) = do
--   (ans1,SNothing) <- runStateT m SNothing
--   (ans2,SJust p) <- runStateT m (SJust s)
--   pure(ans1 == ans2)
--
-- All operations that read the provenance (i.e the state) return
-- a (BlackBox t). BlackBoxes can only be used to modify provenance.
newtype ProvM t m a = ProvM (StateT (StrictMaybe t) m a)
  deriving ((forall a b. (a -> b) -> ProvM t m a -> ProvM t m b)
-> (forall a b. a -> ProvM t m b -> ProvM t m a)
-> Functor (ProvM t m)
forall a b. a -> ProvM t m b -> ProvM t m a
forall a b. (a -> b) -> ProvM t m a -> ProvM t m b
forall t (m :: * -> *) a b.
Functor m =>
a -> ProvM t m b -> ProvM t m a
forall t (m :: * -> *) a b.
Functor m =>
(a -> b) -> ProvM t m a -> ProvM t m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall t (m :: * -> *) a b.
Functor m =>
(a -> b) -> ProvM t m a -> ProvM t m b
fmap :: forall a b. (a -> b) -> ProvM t m a -> ProvM t m b
$c<$ :: forall t (m :: * -> *) a b.
Functor m =>
a -> ProvM t m b -> ProvM t m a
<$ :: forall a b. a -> ProvM t m b -> ProvM t m a
Functor, Functor (ProvM t m)
Functor (ProvM t m) =>
(forall a. a -> ProvM t m a)
-> (forall a b. ProvM t m (a -> b) -> ProvM t m a -> ProvM t m b)
-> (forall a b c.
    (a -> b -> c) -> ProvM t m a -> ProvM t m b -> ProvM t m c)
-> (forall a b. ProvM t m a -> ProvM t m b -> ProvM t m b)
-> (forall a b. ProvM t m a -> ProvM t m b -> ProvM t m a)
-> Applicative (ProvM t m)
forall a. a -> ProvM t m a
forall a b. ProvM t m a -> ProvM t m b -> ProvM t m a
forall a b. ProvM t m a -> ProvM t m b -> ProvM t m b
forall a b. ProvM t m (a -> b) -> ProvM t m a -> ProvM t m b
forall a b c.
(a -> b -> c) -> ProvM t m a -> ProvM t m b -> ProvM t m c
forall t (m :: * -> *). Monad m => Functor (ProvM t m)
forall t (m :: * -> *) a. Monad m => a -> ProvM t m a
forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> ProvM t m b -> ProvM t m a
forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> ProvM t m b -> ProvM t m b
forall t (m :: * -> *) a b.
Monad m =>
ProvM t m (a -> b) -> ProvM t m a -> ProvM t m b
forall t (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> ProvM t m a -> ProvM t m b -> ProvM t m c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall t (m :: * -> *) a. Monad m => a -> ProvM t m a
pure :: forall a. a -> ProvM t m a
$c<*> :: forall t (m :: * -> *) a b.
Monad m =>
ProvM t m (a -> b) -> ProvM t m a -> ProvM t m b
<*> :: forall a b. ProvM t m (a -> b) -> ProvM t m a -> ProvM t m b
$cliftA2 :: forall t (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> ProvM t m a -> ProvM t m b -> ProvM t m c
liftA2 :: forall a b c.
(a -> b -> c) -> ProvM t m a -> ProvM t m b -> ProvM t m c
$c*> :: forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> ProvM t m b -> ProvM t m b
*> :: forall a b. ProvM t m a -> ProvM t m b -> ProvM t m b
$c<* :: forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> ProvM t m b -> ProvM t m a
<* :: forall a b. ProvM t m a -> ProvM t m b -> ProvM t m a
Applicative, Applicative (ProvM t m)
Applicative (ProvM t m) =>
(forall a b. ProvM t m a -> (a -> ProvM t m b) -> ProvM t m b)
-> (forall a b. ProvM t m a -> ProvM t m b -> ProvM t m b)
-> (forall a. a -> ProvM t m a)
-> Monad (ProvM t m)
forall a. a -> ProvM t m a
forall a b. ProvM t m a -> ProvM t m b -> ProvM t m b
forall a b. ProvM t m a -> (a -> ProvM t m b) -> ProvM t m b
forall t (m :: * -> *). Monad m => Applicative (ProvM t m)
forall t (m :: * -> *) a. Monad m => a -> ProvM t m a
forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> ProvM t m b -> ProvM t m b
forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> (a -> ProvM t m b) -> ProvM t m b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> (a -> ProvM t m b) -> ProvM t m b
>>= :: forall a b. ProvM t m a -> (a -> ProvM t m b) -> ProvM t m b
$c>> :: forall t (m :: * -> *) a b.
Monad m =>
ProvM t m a -> ProvM t m b -> ProvM t m b
>> :: forall a b. ProvM t m a -> ProvM t m b -> ProvM t m b
$creturn :: forall t (m :: * -> *) a. Monad m => a -> ProvM t m a
return :: forall a. a -> ProvM t m a
Monad)

instance MonadTrans (ProvM t) where
  lift :: forall (m :: * -> *) a. Monad m => m a -> ProvM t m a
lift m a
x = StateT (StrictMaybe t) m a -> ProvM t m a
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM (m a -> StateT (StrictMaybe t) m a
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (StrictMaybe t) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m a
x)

-- | Run and compute the result as well as the provenance. Supply an initial value for the provenance.
runWithProvM :: Monad m => s -> ProvM s m a -> m (a, s)
runWithProvM :: forall (m :: * -> *) s a. Monad m => s -> ProvM s m a -> m (a, s)
runWithProvM s
s (ProvM StateT (StrictMaybe s) m a
m) = do
  (a, x) <- StateT (StrictMaybe s) m a -> StrictMaybe s -> m (a, StrictMaybe s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (StrictMaybe s) m a
m (s -> StrictMaybe s
forall a. a -> StrictMaybe a
SJust s
s)
  case x of
    StrictMaybe s
SNothing -> [Char] -> m (a, s)
forall a. HasCallStack => [Char] -> a
error ([Char]
"(SJust state) returns SNothing in runWithProvM")
    SJust s
st -> (a, s) -> m (a, s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, s
st)
{-# INLINE runWithProvM #-}

-- | Run the computation with SNothing. Expend no resources to compute provenance.
runProvM :: Monad m => ProvM s m b -> m b
runProvM :: forall (m :: * -> *) s b. Monad m => ProvM s m b -> m b
runProvM (ProvM StateT (StrictMaybe s) m b
m) = do
  pair <- StateT (StrictMaybe s) m b -> StrictMaybe s -> m (b, StrictMaybe s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (StrictMaybe s) m b
m StrictMaybe s
forall a. StrictMaybe a
SNothing
  case pair of
    (b
a, StrictMaybe s
SNothing) -> b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
a
    (b
_, SJust s
_) -> [Char] -> m b
forall a. HasCallStack => [Char] -> a
error ([Char]
"SNothing returns (SJust p) in runProvM")
{-# INLINE runProvM #-}

-- ======================================================================
-- Helper functions and types

-- | BlackBox is an abstraction barrier. Reading the provenance is always returned
-- in a BlackBox. The only way to open the BlackBox is to use one of the BlackBlox
-- eliminator operations: modifyWithBlackBox or runOtherProv, that merge the contents of the
-- BlackBox into the current provenance. This ensures that there is no easy way for the
-- provenance computation to have an effect on the result of the underlying computation.
data BlackBox t = Box !t | NoBox
  deriving (Int -> BlackBox t -> ShowS
[BlackBox t] -> ShowS
BlackBox t -> [Char]
(Int -> BlackBox t -> ShowS)
-> (BlackBox t -> [Char])
-> ([BlackBox t] -> ShowS)
-> Show (BlackBox t)
forall t. Show t => Int -> BlackBox t -> ShowS
forall t. Show t => [BlackBox t] -> ShowS
forall t. Show t => BlackBox t -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall t. Show t => Int -> BlackBox t -> ShowS
showsPrec :: Int -> BlackBox t -> ShowS
$cshow :: forall t. Show t => BlackBox t -> [Char]
show :: BlackBox t -> [Char]
$cshowList :: forall t. Show t => [BlackBox t] -> ShowS
showList :: [BlackBox t] -> ShowS
Show, (forall a b. (a -> b) -> BlackBox a -> BlackBox b)
-> (forall a b. a -> BlackBox b -> BlackBox a) -> Functor BlackBox
forall a b. a -> BlackBox b -> BlackBox a
forall a b. (a -> b) -> BlackBox a -> BlackBox b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> BlackBox a -> BlackBox b
fmap :: forall a b. (a -> b) -> BlackBox a -> BlackBox b
$c<$ :: forall a b. a -> BlackBox b -> BlackBox a
<$ :: forall a b. a -> BlackBox b -> BlackBox a
Functor)

-- Helper for modifyM. This runs in (STateT p m a), modifyM runs in (ProvM p m a)

modifyMState :: Monad m => (t -> t) -> StateT (StrictMaybe t) m ()
modifyMState :: forall (m :: * -> *) t.
Monad m =>
(t -> t) -> StateT (StrictMaybe t) m ()
modifyMState t -> t
delta = do
  mstore <- StateT (StrictMaybe t) m (StrictMaybe t)
forall s (m :: * -> *). MonadState s m => m s
get
  case mstore of
    StrictMaybe t
SNothing -> () -> StateT (StrictMaybe t) m ()
forall a. a -> StateT (StrictMaybe t) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    (SJust t
st) -> StrictMaybe t -> StateT (StrictMaybe t) m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (t -> StrictMaybe t
forall a. a -> StrictMaybe a
SJust (t -> t
delta t
st))
{-# INLINE modifyMState #-}

-- | Overwrite the current provenance with something new.
putM :: Monad m => s -> ProvM s m ()
putM :: forall (m :: * -> *) s. Monad m => s -> ProvM s m ()
putM s
s = StateT (StrictMaybe s) m () -> ProvM s m ()
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM ((s -> s) -> StateT (StrictMaybe s) m ()
forall (m :: * -> *) t.
Monad m =>
(t -> t) -> StateT (StrictMaybe t) m ()
modifyMState (s -> s -> s
forall a b. a -> b -> a
const s
s))
{-# INLINE putM #-}

-- | Extract the current provenance. The result is wrapped in
-- a BlackBox. This ensures that provenance cannot be used in the
-- non-provenance part of the computation.
getM :: Monad m => ProvM s m (BlackBox s)
getM :: forall (m :: * -> *) s. Monad m => ProvM s m (BlackBox s)
getM = StateT (StrictMaybe s) m (BlackBox s) -> ProvM s m (BlackBox s)
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM (do m <- StateT (StrictMaybe s) m (StrictMaybe s)
forall s (m :: * -> *). MonadState s m => m s
get; case m of StrictMaybe s
SNothing -> BlackBox s -> StateT (StrictMaybe s) m (BlackBox s)
forall a. a -> StateT (StrictMaybe s) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BlackBox s
forall t. BlackBox t
NoBox; SJust s
t -> BlackBox s -> StateT (StrictMaybe s) m (BlackBox s)
forall a. a -> StateT (StrictMaybe s) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s -> BlackBox s
forall t. t -> BlackBox t
Box s
t))
{-# INLINE getM #-}

-- | Modify the provenance if collecting provenance, otherwise do nothing.
modifyM :: Monad m => (t -> t) -> ProvM t m ()
modifyM :: forall (m :: * -> *) t. Monad m => (t -> t) -> ProvM t m ()
modifyM t -> t
delta = StateT (StrictMaybe t) m () -> ProvM t m ()
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM ((t -> t) -> StateT (StrictMaybe t) m ()
forall (m :: * -> *) t.
Monad m =>
(t -> t) -> StateT (StrictMaybe t) m ()
modifyMState t -> t
delta)
{-# INLINE modifyM #-}

modifyWithBlackBox :: Monad m => BlackBox p -> (p -> t -> t) -> ProvM t m ()
modifyWithBlackBox :: forall (m :: * -> *) p t.
Monad m =>
BlackBox p -> (p -> t -> t) -> ProvM t m ()
modifyWithBlackBox (Box p
x) p -> t -> t
delta = StateT (StrictMaybe t) m () -> ProvM t m ()
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM ((t -> t) -> StateT (StrictMaybe t) m ()
forall (m :: * -> *) t.
Monad m =>
(t -> t) -> StateT (StrictMaybe t) m ()
modifyMState (p -> t -> t
delta p
x))
modifyWithBlackBox BlackBox p
NoBox p -> t -> t
_ = StateT (StrictMaybe t) m () -> ProvM t m ()
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM (() -> StateT (StrictMaybe t) m ()
forall a. a -> StateT (StrictMaybe t) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
{-# INLINE modifyWithBlackBox #-}

-- Are we currently collecting provenance? Not to be exported.
active :: Monad m => ProvM s m Bool
active :: forall (m :: * -> *) s. Monad m => ProvM s m Bool
active = StateT (StrictMaybe s) m Bool -> ProvM s m Bool
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM (do m <- StateT (StrictMaybe s) m (StrictMaybe s)
forall s (m :: * -> *). MonadState s m => m s
get; pure (case m of StrictMaybe s
SNothing -> Bool
False; SJust s
_ -> Bool
True))
{-# INLINE active #-}

-- | Run a provenance computation, with provenance s1, and lift the
-- result to a provenance computation with provenance s2. Use the active
-- state of the s2 computation to decide if we actually compute the
-- provenance s1. The s1 result is returned in a BlackBox. This ensures that
-- provenance cannot be used in the non-provenance part of the computation.
runOtherProv :: Monad m => s1 -> ProvM s1 m a -> ProvM s2 m (a, BlackBox s1)
runOtherProv :: forall (m :: * -> *) s1 a s2.
Monad m =>
s1 -> ProvM s1 m a -> ProvM s2 m (a, BlackBox s1)
runOtherProv s1
initial ProvM s1 m a
other = do
  t <- ProvM s2 m Bool
forall (m :: * -> *) s. Monad m => ProvM s m Bool
active
  if t
    then ProvM (lift $ do (a, s) <- runWithProvM initial other; pure (a, Box s))
    else ProvM (lift $ do a <- runProvM other; pure (a, NoBox))
{-# INLINE runOtherProv #-}

-- | lift a provenenace computation from one provenance type (s1) to another (s2)
liftProv :: Monad m => ProvM s1 m a -> s1 -> (a -> s1 -> s2 -> s2) -> ProvM s2 m a
liftProv :: forall (m :: * -> *) s1 a s2.
Monad m =>
ProvM s1 m a -> s1 -> (a -> s1 -> s2 -> s2) -> ProvM s2 m a
liftProv ProvM s1 m a
computation s1
inits1 a -> s1 -> s2 -> s2
combine =
  do
    (a, blackbox) <- s1 -> ProvM s1 m a -> ProvM s2 m (a, BlackBox s1)
forall (m :: * -> *) s1 a s2.
Monad m =>
s1 -> ProvM s1 m a -> ProvM s2 m (a, BlackBox s1)
runOtherProv s1
inits1 ProvM s1 m a
computation
    modifyWithBlackBox blackbox (combine a)
    pure a
{-# INLINE liftProv #-}

-- =======================================================================

-- | A special case of the ProvM Monad, where the state type is Store
-- a (Map Text PObject), where PObject is a dynamically typed value. This
-- allows the collection of a Set of provenance values, indexed by keys of
-- type Text. As in the ProvM monad, if run with 'runProv' operations on
-- the Store are ignored.
type Prov m a = ProvM Store m a

-- | Run the (Prov m a) computation and ignore the provenance part
runProv :: Monad m => Prov m t -> m t
runProv :: forall (m :: * -> *) t. Monad m => Prov m t -> m t
runProv (ProvM StateT (StrictMaybe Store) m t
m) = do (a, _) <- StateT (StrictMaybe Store) m t
-> StrictMaybe Store -> m (t, StrictMaybe Store)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (StrictMaybe Store) m t
m StrictMaybe Store
forall a. StrictMaybe a
SNothing; pure a

-- | Run the (Prov m a) computation, compute and return the provenance as well as the result.
runWithProv :: Monad m => Prov m t -> m (t, Store)
runWithProv :: forall (m :: * -> *) t. Monad m => Prov m t -> m (t, Store)
runWithProv = Store -> ProvM Store m t -> m (t, Store)
forall (m :: * -> *) s a. Monad m => s -> ProvM s m a -> m (a, s)
runWithProvM Store
forall k a. Map k a
empty

-- | Run a computation in the underlying monad (m), return that value in the
-- (Prov m) monad. As a side effect store that value under the given key
-- equivalent to: store key m = do { a <- lift m; push key a; pure a}
store :: forall t m. (Provenance t, Monad m) => Text -> m t -> Prov m t
store :: forall t (m :: * -> *).
(Provenance t, Monad m) =>
Text -> m t -> Prov m t
store Text
key m t
m = StateT (StrictMaybe Store) m t -> ProvM Store m t
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM (do a <- m t -> StateT (StrictMaybe Store) m t
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (StrictMaybe Store) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m t
m; modifyMState (insert key (pobject a)); pure a)
{-# INLINE store #-}

-- | Push a key value pair into the provenance store. Overwrites any
-- existing POjects stored under that key.
push :: (Provenance t, Monad m) => Text -> t -> Prov m ()
push :: forall t (m :: * -> *).
(Provenance t, Monad m) =>
Text -> t -> Prov m ()
push Text
key t
t = StateT (StrictMaybe Store) m () -> ProvM Store m ()
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM ((Store -> Store) -> StateT (StrictMaybe Store) m ()
forall (m :: * -> *) t.
Monad m =>
(t -> t) -> StateT (StrictMaybe t) m ()
modifyMState (Text -> PObject -> Store -> Store
forall k a. Ord k => k -> a -> Map k a -> Map k a
insert Text
key (t -> PObject
forall t. Provenance t => t -> PObject
pobject t
t)))
{-# INLINE push #-}

-- | Modify the value stored at the given key. If the key isn't found
-- or the PObject at that key has the wrong type, do nothing.
update :: forall t m. (Provenance t, Monad m) => Text -> (t -> t) -> Prov m ()
update :: forall t (m :: * -> *).
(Provenance t, Monad m) =>
Text -> (t -> t) -> Prov m ()
update Text
key t -> t
delta = StateT (StrictMaybe Store) m () -> ProvM Store m ()
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM StateT (StrictMaybe Store) m ()
action2
  where
    action2 :: StateT (StrictMaybe Store) m ()
action2 = do
      m <- StateT (StrictMaybe Store) m (StrictMaybe Store)
forall s (m :: * -> *). MonadState s m => m s
get
      case findM @t key m of
        SJust t
t -> (Store -> Store) -> StateT (StrictMaybe Store) m ()
forall (m :: * -> *) t.
Monad m =>
(t -> t) -> StateT (StrictMaybe t) m ()
modifyMState (Text -> PObject -> Store -> Store
forall k a. Ord k => k -> a -> Map k a -> Map k a
insert Text
key (forall t. Provenance t => t -> PObject
pobject @t (t -> t
delta t
t)))
        StrictMaybe t
SNothing -> () -> StateT (StrictMaybe Store) m ()
forall a. a -> StateT (StrictMaybe Store) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE update #-}

-- | Modify the value stored at the given key using a value in a BlackBox.
-- if the key isn't found, or the PObject at that key has the wrong type, do nothing.
updateWithBlackBox ::
  forall t m s. (Provenance t, Monad m) => Text -> BlackBox s -> (s -> t -> t) -> Prov m ()
updateWithBlackBox :: forall t (m :: * -> *) s.
(Provenance t, Monad m) =>
Text -> BlackBox s -> (s -> t -> t) -> Prov m ()
updateWithBlackBox Text
key (Box s
s) s -> t -> t
delta = Text -> (t -> t) -> Prov m ()
forall t (m :: * -> *).
(Provenance t, Monad m) =>
Text -> (t -> t) -> Prov m ()
update Text
key (s -> t -> t
delta s
s)
updateWithBlackBox Text
_ BlackBox s
NoBox s -> t -> t
_ = () -> Prov m ()
forall a. a -> ProvM Store m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE updateWithBlackBox #-}

-- | Reads the provenance value at a key. The value is returned
-- in a BlackBox. There are 3 reasons the BlackBox may be empty.
-- 1) The computation is not collecting provenance.
-- 2) The map does not contain the key
-- 3) The value at the key has the wrong type.
pull :: forall t m. (Monad m, Typeable t) => Text -> Prov m (BlackBox t)
pull :: forall t (m :: * -> *).
(Monad m, Typeable t) =>
Text -> Prov m (BlackBox t)
pull Text
key = StateT (StrictMaybe Store) m (BlackBox t)
-> ProvM Store m (BlackBox t)
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM (do m <- StateT (StrictMaybe Store) m (StrictMaybe Store)
forall s (m :: * -> *). MonadState s m => m s
get; case findM key m of StrictMaybe t
SNothing -> BlackBox t -> StateT (StrictMaybe Store) m (BlackBox t)
forall a. a -> StateT (StrictMaybe Store) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BlackBox t
forall t. BlackBox t
NoBox; SJust t
t -> BlackBox t -> StateT (StrictMaybe Store) m (BlackBox t)
forall a. a -> StateT (StrictMaybe Store) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (t -> BlackBox t
forall t. t -> BlackBox t
Box t
t))
{-# INLINE pull #-}

-- | Return a String representation of the current provenance store.
dump :: Monad m => Prov m String
dump :: forall (m :: * -> *). Monad m => Prov m [Char]
dump =
  StateT (StrictMaybe Store) m [Char] -> ProvM Store m [Char]
forall t (m :: * -> *) a. StateT (StrictMaybe t) m a -> ProvM t m a
ProvM
    ( do
        mstore <- StateT (StrictMaybe Store) m (StrictMaybe Store)
forall s (m :: * -> *). MonadState s m => m s
get
        case mstore of
          SJust Store
m -> [Char] -> StateT (StrictMaybe Store) m [Char]
forall a. a -> StateT (StrictMaybe Store) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Store -> [Char]
observe Store
m)
          StrictMaybe Store
SNothing -> [Char] -> StateT (StrictMaybe Store) m [Char]
forall a. a -> StateT (StrictMaybe Store) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
"SNothing Store"
    )

-- | Push the provenance of a computation, under the given key. The
-- computation has provenance s1, and lift the result to a provenance
-- computation with provenance Store. Use the active state of the Store
-- computation to decide if we actually want to compute the
-- provenance s1, and push it, or simply ignore it.
pushOtherProv :: (Provenance s1, Monad m) => Text -> s1 -> ProvM s1 m a -> ProvM Store m a
pushOtherProv :: forall s1 (m :: * -> *) a.
(Provenance s1, Monad m) =>
Text -> s1 -> ProvM s1 m a -> ProvM Store m a
pushOtherProv Text
key s1
initial ProvM s1 m a
other = do
  t <- ProvM Store m Bool
forall (m :: * -> *) s. Monad m => ProvM s m Bool
active
  if t
    then
      ProvM
        ( do
            (a, v) <- lift (runWithProvM initial other)
            modifyMState (insert key (pobject v))
            pure a
        )
    else ProvM (lift $ runProvM other)

-- ================================================================
-- Helper functions to implement the operations on Prov
-- that manipulate PObject values.

-- | Since PObjects are dynamically typed, What operations are
-- required on a type to act as Provenance? We might want to
-- add or subtract some properties from this list.
type Provenance t = (Typeable t, ToJSON t, Show t, NoThunks t)

-- | Provenance Object
data PObject where
  PObject :: Provenance t => !(TypeRep t) -> !t -> PObject

instance NoThunks PObject where
  showTypeOf :: Proxy PObject -> [Char]
showTypeOf Proxy PObject
_ = [Char]
"PObject"
  wNoThunks :: Context -> PObject -> IO (Maybe ThunkInfo)
wNoThunks Context
ctxt (PObject TypeRep t
_ t
t) = [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks [Context -> t -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctxt t
t]

instance Show PObject where
  show :: PObject -> [Char]
show (PObject TypeRep t
ty t
t) = [Char]
"#" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ t -> [Char]
forall a. Show a => a -> [Char]
show t
t [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"::" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeRep t -> [Char]
forall a. Show a => a -> [Char]
show TypeRep t
ty

-- | extract a value with the given type from a PObject. If the PObject
--   does not have the right type, returns SNothing. If the type context
--   of the call does not fix the type, one can use a type application like
--   extract @[Int] pobject
extract :: forall t. Typeable t => PObject -> StrictMaybe t
extract :: forall t. Typeable t => PObject -> StrictMaybe t
extract (PObject TypeRep t
ty t
n) = case TypeRep t -> TypeRep t -> Maybe (t :~: t)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality TypeRep t
ty (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @t) of Just t :~: t
Refl -> t -> StrictMaybe t
forall a. a -> StrictMaybe a
SJust t
t
n; Maybe (t :~: t)
Nothing -> StrictMaybe t
forall a. StrictMaybe a
SNothing

-- | inject a type into the PObject type.
pobject :: Provenance t => t -> PObject
pobject :: forall t. Provenance t => t -> PObject
pobject !t
n = TypeRep t -> t -> PObject
forall t. Provenance t => TypeRep t -> t -> PObject
PObject (t -> TypeRep t
forall a. Typeable a => a -> TypeRep a
typeOf t
n) t
n

-- ==============================================================
-- Maps of PObjects is how we store Provenance in the ProvM Store monad.
-- Here we define some useful helper functions on these Maps

type Store = Map Text PObject

-- | Find a value for a given key from a Store (Map Text PObject). If
--  the Store does not have that key, or the PObject at that key
--  does not have the right type, returns SNothing. If the type context
--  of the call does not fix the type, one can use a type
--  application like:  find @Bool key map
find :: forall t k. (Ord k, Typeable t) => k -> Map k PObject -> StrictMaybe t
find :: forall t k.
(Ord k, Typeable t) =>
k -> Map k PObject -> StrictMaybe t
find k
key Map k PObject
m = case k -> Map k PObject -> Maybe PObject
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
key Map k PObject
m of Just PObject
obj -> PObject -> StrictMaybe t
forall t. Typeable t => PObject -> StrictMaybe t
extract PObject
obj; Maybe PObject
Nothing -> StrictMaybe t
forall a. StrictMaybe a
SNothing

-- A useful helper function
findM :: forall t k. (Ord k, Typeable t) => k -> StrictMaybe (Map k PObject) -> StrictMaybe t
findM :: forall t k.
(Ord k, Typeable t) =>
k -> StrictMaybe (Map k PObject) -> StrictMaybe t
findM k
_ StrictMaybe (Map k PObject)
SNothing = StrictMaybe t
forall a. StrictMaybe a
SNothing
findM k
key (SJust Map k PObject
m) = k -> Map k PObject -> StrictMaybe t
forall t k.
(Ord k, Typeable t) =>
k -> Map k PObject -> StrictMaybe t
find k
key Map k PObject
m

-- | Turn a Map of PObjects into a String, indicating its contents.
observe :: Store -> String
observe :: Store -> [Char]
observe Store
m = Context -> [Char]
unlines (((Text, PObject) -> [Char]) -> [(Text, PObject)] -> Context
forall a b. (a -> b) -> [a] -> [b]
map (Text, PObject) -> [Char]
f (Store -> [(Text, PObject)]
forall k a. Map k a -> [(k, a)]
Map.assocs Store
m))
  where
    f :: (Text, PObject) -> [Char]
f (Text
key, PObject TypeRep t
_ t
t) = Text -> [Char]
unpack Text
key [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" =\n   " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ t -> [Char]
forall a. Show a => a -> [Char]
show t
t

-- =======================================================================
-- useful for testing invariants, The type StrictMaybe is local to this
-- module, and is not exported, so these predicates are defined here to
-- support testing these invariants.

preservesNothing :: Monad m => ProvM t m a -> m Bool
preservesNothing :: forall (m :: * -> *) t a. Monad m => ProvM t m a -> m Bool
preservesNothing (ProvM StateT (StrictMaybe t) m a
m) = do
  (_, maybet) <- StateT (StrictMaybe t) m a -> StrictMaybe t -> m (a, StrictMaybe t)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (StrictMaybe t) m a
m StrictMaybe t
forall a. StrictMaybe a
SNothing
  case maybet of StrictMaybe t
SNothing -> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True; SJust t
_ -> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

preservesJust :: Monad m => t -> ProvM t m a -> m Bool
preservesJust :: forall (m :: * -> *) t a. Monad m => t -> ProvM t m a -> m Bool
preservesJust t
t (ProvM StateT (StrictMaybe t) m a
m) = do
  (_, maybet) <- StateT (StrictMaybe t) m a -> StrictMaybe t -> m (a, StrictMaybe t)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (StrictMaybe t) m a
m (t -> StrictMaybe t
forall a. a -> StrictMaybe a
SJust t
t)
  case maybet of StrictMaybe t
SNothing -> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False; SJust t
_ -> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True