DSL.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
-- | The graded monad DSL: programs graded by their Knightian choice names.
module Imp.DSL
  ( Imp(..)
  , flip
  , knight
  , interval
  , observe
  , tag
  , IfThenElse(..)
  , IfR
  , Merge
  , Union
  , TagAll
  , Imp.DSL.return
  , (Imp.DSL.>>=)
  , (Imp.DSL.>>)
  , Imp.DSL.fmap
  , (Imp.DSL.<$>)
  ) where

import Prelude hiding (return, (>>=), (>>), flip, fmap, (<$>))
import Data.Proxy (Proxy(..))
import GHC.TypeLits (KnownSymbol, Symbol)
import Imp.DSL.Grade (Merge, Union, TagAll)

-- | A probabilistic program graded by @g@,
-- a type-level list of Symbols tracking which Knightian names this program uses.
data Imp (g :: [Symbol]) a where
  ImpReturn  :: a -> Imp '[] a
  ImpBind    :: Ord a => Imp g1 a -> (a -> Imp g2 b) -> Imp (Merge g1 g2) b
  ImpFlip    :: !Double -> Imp '[] Bool
  ImpKnight  :: KnownSymbol n => Proxy n -> Imp '[n] Bool
  ImpObserve :: Bool -> Imp '[] ()
  ImpTag     :: KnownSymbol t => Proxy t -> Imp g a -> Imp (TagAll t g) a
  ImpBranch  :: Bool -> Imp g1 a -> Imp g2 a -> Imp (Union g1 g2) a

-- | Graded monad return.
return :: a -> Imp '[] a
return = ImpReturn

-- | Graded monad bind.
infixl 1 >>=
(>>=) :: Ord a => Imp g1 a -> (a -> Imp g2 b) -> Imp (Merge g1 g2) b
(>>=) = ImpBind

-- | Graded monad sequence.
infixl 1 >>
(>>) :: Ord a => Imp g1 a -> Imp g2 b -> Imp (Merge g1 g2) b
m >> n = ImpBind m (const n)

-- | Graded functor map.
fmap :: Ord a => (a -> b) -> Imp g a -> Imp g b
fmap f m = ImpBind m (ImpReturn . f)

-- | Operator form of the graded functor map.
infixl 4 <$>
(<$>) :: Ord a => (a -> b) -> Imp g a -> Imp g b
(<$>) = fmap

-- | Probabilistic coin flip with probability @p@ of True; @p@ must lie in @[0,1]@.
flip :: Double -> Imp '[] Bool
flip p | p >= 0 && p <= 1 = ImpFlip p
       | otherwise        = error ("flip: probability " ++ show p ++ " outside [0,1]")

-- | Knightian (adversarial) binary choice, named at the type level.
--   Usage: @knight \@\"x\"@
knight :: forall n. KnownSymbol n => Imp '[n] Bool
knight = ImpKnight (Proxy @n)

-- | Interval-valued probability: @interval \@\"n\" lo hi@ returns a @Bool@
--   with @P(True) ∈ [lo, hi]@.  The bounds may be given in either order.
interval :: forall n. KnownSymbol n => Double -> Double -> Imp '[n] Bool
interval lo hi = ImpBind (ImpKnight (Proxy @n)) $ \x ->
  flip (if x then hi else lo)

-- | Condition on a Boolean predicate being True.
observe :: Bool -> Imp '[] ()
observe = ImpObserve

-- | Tag all Knightian choices in a subprogram enabling reusage.
tag :: forall t g a. KnownSymbol t => Imp g a -> Imp (TagAll t g) a
tag = ImpTag (Proxy @t)

-- ---------------------------------------------------------------------------
-- IfThenElse: grade-aware branching for RebindableSyntax
-- ---------------------------------------------------------------------------

-- | Compute the result type of an @if-then-else@ expression.
--   For @Imp@ branches with potentially different grades, the result grade is their 'Union'.
type family IfR t f where
  IfR (Imp g1 a) (Imp g2 a) = Imp (Union g1 g2) a
  IfR t          f           = t

-- | Overloaded @if-then-else@ for use with @RebindableSyntax@.
--
--   The pure instance is marked @INCOHERENT@ so that GHC can commit to
--   it when the branch types are not yet determined.
class IfThenElse t f where
  ifThenElse :: Bool -> t -> f -> IfR t f

instance {-# INCOHERENT #-} (a ~ b, IfR a b ~ a) => IfThenElse a b where
  ifThenElse True  t _ = t
  ifThenElse False _ f = f

instance IfThenElse (Imp g1 a) (Imp g2 a) where
  ifThenElse = ImpBranch