{-# LANGUAGE CPP
           , GADTs
           , KindSignatures
           , DataKinds
           , PolyKinds
           , TypeOperators
           , Rank2Types
           , FlexibleContexts
           , MultiParamTypeClasses
           , FlexibleInstances
           , UndecidableInstances
           , EmptyCase
           , ScopedTypeVariables
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                    2016.05.24
-- |
-- Module      :  Language.Hakaru.Evaluation.DisintegrationMonad
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :  wren@community.haskell.org
-- Stability   :  experimental
-- Portability :  GHC-only
--
-- The 'EvaluationMonad' for "Language.Hakaru.Disintegrate"
----------------------------------------------------------------
module Language.Hakaru.Evaluation.DisintegrationMonad
    (
    -- * The disintegration monad
    -- ** List-based version
      getStatements, putStatements
    , ListContext(..), Ans, Dis(..), runDis, runDisInCtx
    -- ** TODO: IntMap-based version
    
    -- * Operators on the disintegration monad
    -- ** The \"zero\" and \"one\"
    , bot
    --, reject
    -- ** Emitting code
    , emit
    , emitMBind , emitMBind2
    , emitLet
    , emitLet'
    , emitUnpair
    -- TODO: emitUneither
    -- emitCaseWith
    , emit_
    , emitMBind_
    , emitGuard
    , emitWeight
    , emitFork_
    , emitSuperpose
    , choose
    , pushWeight
    , pushGuard
    -- * Overrides for original in Evaluation.Types
    , pushPlate
    -- * For Arrays/Plate
    , getIndices
    , withIndices
    , extendIndices
    , selectMore
    , permutes
    , statementInds
    , sizeInnermostInd
    -- * Extras
    , Extra(..)
    , getExtras
    , putExtras
    , insertExtra
    , adjustExtra
    , mkLoc
    , freeLocError
    , zipInds
    , apply  
#ifdef __TRACE_DISINTEGRATE__
    , prettyExtra
    , prettyExtras
#endif    
    ) where

import           Prelude              hiding (id, (.))
import           Control.Category     (Category(..))
#if __GLASGOW_HASKELL__ < 710
import           Data.Monoid          (Monoid(..))
import           Data.Functor         ((<$>))
import           Control.Applicative  (Applicative(..))
#endif
import qualified Data.Set             as Set (fromList)
import           Data.Maybe
import qualified Data.Foldable        as F
import qualified Data.Traversable     as T
import           Data.List.NonEmpty   (NonEmpty(..))
import qualified Data.List.NonEmpty   as NE
import           Control.Applicative  (Alternative(..))
import           Control.Monad        (MonadPlus(..),foldM,guard)
#if __GLASGOW_HASKELL__ > 805
import           Control.Monad.Fail
#endif
import           Data.Text            (Text)
import qualified Data.Text            as Text
import           Data.Number.Nat

import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.Sing    (Sing(..), sUnMeasure, sUnPair, sUnit)
import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.DatumABT
import Language.Hakaru.Syntax.TypeOf
import Language.Hakaru.Syntax.Transform (TransformCtx(..), minimalCtx)
import Language.Hakaru.Syntax.ABT
import qualified Language.Hakaru.Syntax.Prelude as P
import Language.Hakaru.Evaluation.Types
import Language.Hakaru.Evaluation.PEvalMonad (ListContext(..))
import Language.Hakaru.Evaluation.Lazy (reifyPair)    

#ifdef __TRACE_DISINTEGRATE__
import Debug.Trace (trace, traceM)
import qualified Text.PrettyPrint     as PP
import Language.Hakaru.Pretty.Haskell (ppVariable, pretty)
#endif

getStatements :: Dis abt [Statement abt Location 'Impure]
getStatements :: Dis abt [Statement abt Location 'Impure]
getStatements = (forall (a :: Hakaru).
 [Index (abt '[])]
 -> ([Statement abt Location 'Impure] -> Ans abt a) -> Ans abt a)
-> Dis abt [Statement abt Location 'Impure]
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])]
  -> ([Statement abt Location 'Impure] -> Ans abt a) -> Ans abt a)
 -> Dis abt [Statement abt Location 'Impure])
-> (forall (a :: Hakaru).
    [Index (abt '[])]
    -> ([Statement abt Location 'Impure] -> Ans abt a) -> Ans abt a)
-> Dis abt [Statement abt Location 'Impure]
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ [Statement abt Location 'Impure] -> Ans abt a
c ListContext abt 'Impure
h -> [Statement abt Location 'Impure] -> Ans abt a
c (ListContext abt 'Impure -> [Statement abt Location 'Impure]
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
ListContext abt p -> [Statement abt Location p]
statements ListContext abt 'Impure
h) ListContext abt 'Impure
h

putStatements :: [Statement abt Location 'Impure] -> Dis abt ()
putStatements :: [Statement abt Location 'Impure] -> Dis abt ()
putStatements [Statement abt Location 'Impure]
ss =
    (forall (a :: Hakaru).
 [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
 -> Dis abt ())
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ () -> Ans abt a
c (ListContext Nat
i [Statement abt Location 'Impure]
_) Assocs (Extra (abt '[]))
extras ->
        () -> Ans abt a
c () (Nat -> [Statement abt Location 'Impure] -> ListContext abt 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Nat -> [Statement abt Location p] -> ListContext abt p
ListContext Nat
i [Statement abt Location 'Impure]
ss) Assocs (Extra (abt '[]))
extras

----------------------------------------------------------------
----------------------------------------------------------------

-- | Capturing substitution: 
plug :: forall abt a xs b
     .  (ABT Term abt)
     => Variable a
     -> abt '[] a
     -> abt xs b
     -> abt xs b
plug :: Variable a -> abt '[] a -> abt xs b -> abt xs b
plug Variable a
x abt '[] a
e = abt xs b -> abt xs b
forall (xs' :: [Hakaru]) (b' :: Hakaru). abt xs' b' -> abt xs' b'
start
    where
      start :: forall xs' b' . abt xs' b' -> abt xs' b'
      start :: abt xs' b' -> abt xs' b'
start abt xs' b'
f = abt xs' b' -> View (Term abt) xs' b' -> abt xs' b'
forall (xs' :: [Hakaru]) (b' :: Hakaru).
abt xs' b' -> View (Term abt) xs' b' -> abt xs' b'
loop abt xs' b'
f (abt xs' b' -> View (Term abt) xs' b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> View (syn abt) xs a
viewABT abt xs' b'
f)
      loop :: forall xs' b'. abt xs' b' -> View (Term abt) xs' b' -> abt xs' b'
      loop :: abt xs' b' -> View (Term abt) xs' b' -> abt xs' b'
loop abt xs' b'
_ (Syn Term abt b'
t) = Term abt b' -> abt '[] b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt b' -> abt '[] b') -> Term abt b' -> abt '[] b'
forall a b. (a -> b) -> a -> b
$! (forall (xs' :: [Hakaru]) (b' :: Hakaru). abt xs' b' -> abt xs' b')
-> Term abt b' -> Term abt b'
forall k1 k2 k3 (f :: (k1 -> k2 -> *) -> k3 -> *)
       (a :: k1 -> k2 -> *) (b :: k1 -> k2 -> *) (j :: k3).
Functor21 f =>
(forall (h :: k1) (i :: k2). a h i -> b h i) -> f a j -> f b j
fmap21 forall (xs' :: [Hakaru]) (b' :: Hakaru). abt xs' b' -> abt xs' b'
start Term abt b'
t
      loop abt xs' b'
f (Var Variable b'
z) = case Variable a -> Variable b' -> Maybe (TypeEq a b')
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Variable b -> Maybe (TypeEq a b)
varEq Variable a
x Variable b'
z of
                       Just TypeEq a b'
Refl -> abt xs' b'
abt '[] a
e
                       Maybe (TypeEq a b')
Nothing   -> abt xs' b'
f
      loop abt xs' b'
f (Bind Variable a
_ View (Term abt) xs b'
_) = abt (a : xs) b'
-> (Variable a -> abt xs b' -> abt (a : xs) b') -> abt (a : xs) b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt xs' b'
abt (a : xs) b'
f ((Variable a -> abt xs b' -> abt (a : xs) b') -> abt (a : xs) b')
-> (Variable a -> abt xs b' -> abt (a : xs) b') -> abt (a : xs) b'
forall a b. (a -> b) -> a -> b
$ \Variable a
z abt xs b'
f' -> 
                          Variable a -> abt xs b' -> abt (a : xs) b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
z (abt xs b' -> View (Term abt) xs b' -> abt xs b'
forall (xs' :: [Hakaru]) (b' :: Hakaru).
abt xs' b' -> View (Term abt) xs' b' -> abt xs' b'
loop abt xs b'
f' (abt xs b' -> View (Term abt) xs b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> View (syn abt) xs a
viewABT abt xs b'
f'))

-- | Perform multiple capturing substitutions
plugs :: forall abt xs a
      .  (ABT Term abt)
      => Assocs (abt '[]) 
      -> abt xs a
      -> abt xs a
plugs :: Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho0 abt xs a
e0 = (abt xs a -> Assoc (abt '[]) -> abt xs a)
-> abt xs a -> IntMap (Assoc (abt '[])) -> abt xs a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl (\abt xs a
e (Assoc Variable a
x abt '[] a
v) -> Variable a -> abt '[] a -> abt xs a -> abt xs a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (b :: Hakaru).
ABT Term abt =>
Variable a -> abt '[] a -> abt xs b -> abt xs b
plug Variable a
x abt '[] a
v abt xs a
e) abt xs a
e0 (Assocs (abt '[]) -> IntMap (Assoc (abt '[]))
forall k (ast :: k -> *). Assocs ast -> IntMap (Assoc ast)
unAssocs Assocs (abt '[])
rho0)

-- | Plug a term into a context. That is, the 'statements' of the
-- context specifies a program with a hole in it; so we plug the
-- given term into that hole, returning the complete program.
residualizeListContext
    :: forall abt a
    .  (ABT Term abt)
    => ListContext abt 'Impure
    -> Assocs (abt '[])
    -> abt '[] ('HMeasure a)
    -> abt '[] ('HMeasure a)
residualizeListContext :: ListContext abt 'Impure
-> Assocs (abt '[])
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure a)
residualizeListContext ListContext abt 'Impure
ss Assocs (abt '[])
rho abt '[] ('HMeasure a)
e0 =
    -- N.B., we use a left fold because the head of the list of
    -- statements is the one closest to the hole.
#ifdef __TRACE_DISINTEGRATE__
    trace ("e0: " ++ show (pretty e0) ++ "\n"
          ++ show (pretty_Statements (statements ss))) $
#endif
    (abt '[] ('HMeasure a)
 -> Statement abt Location 'Impure -> abt '[] ('HMeasure a))
-> abt '[] ('HMeasure a)
-> [Statement abt Location 'Impure]
-> abt '[] ('HMeasure a)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl abt '[] ('HMeasure a)
-> Statement abt Location 'Impure -> abt '[] ('HMeasure a)
step (Assocs (abt '[]) -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho abt '[] ('HMeasure a)
e0) (ListContext abt 'Impure -> [Statement abt Location 'Impure]
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
ListContext abt p -> [Statement abt Location p]
statements ListContext abt 'Impure
ss)
    where    
    step
        :: abt '[] ('HMeasure a)
        -> Statement abt Location 'Impure
        -> abt '[] ('HMeasure a)
    step :: abt '[] ('HMeasure a)
-> Statement abt Location 'Impure -> abt '[] ('HMeasure a)
step abt '[] ('HMeasure a)
e Statement abt Location 'Impure
s =        
#ifdef __TRACE_DISINTEGRATE__
        trace ("wrapping " ++ show (ppStatement 0 s) ++ "\n"
               ++ "around term " ++ show (pretty e)) $
#endif  
        case Statement abt Location 'Impure
s of       
        SBind (Location Variable a
x) Lazy abt ('HMeasure a)
body [Index (abt '[])]
_ ->
            -- TODO: if @body@ is dirac, then treat as 'SLet'
            Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC ('HMeasure a), '( '[a], 'HMeasure a)] ('HMeasure a)
forall (a :: Hakaru) (b :: Hakaru).
SCon '[LC ('HMeasure a), '( '[a], 'HMeasure b)] ('HMeasure b)
MBind SCon '[LC ('HMeasure a), '( '[a], 'HMeasure a)] ('HMeasure a)
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HMeasure a)]
-> Term abt ('HMeasure a)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ Assocs (abt '[]) -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho (Lazy abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt ('HMeasure a)
body) abt '[] ('HMeasure a)
-> SArgs abt '[ '( '[a], 'HMeasure a)]
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] ('HMeasure a) -> abt '[a] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x abt '[] ('HMeasure a)
e abt '[a] ('HMeasure a)
-> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)
        SLet (Location Variable a
x) Lazy abt a
body [Index (abt '[])]
_
            | Bool -> Bool
not (Variable a
x Variable a -> VarSet (KindOf ('HMeasure a)) -> Bool
forall k (a :: k) (kproxy :: KProxy k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> VarSet kproxy -> Bool
`memberVarSet` abt '[] ('HMeasure a) -> VarSet (KindOf ('HMeasure a))
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> VarSet (KindOf a)
freeVars abt '[] ('HMeasure a)
e) ->
#ifdef __TRACE_DISINTEGRATE__
               trace ("could not find location " ++ show x ++ "\n"
                     ++ "in term " ++ show (pretty e) ++ "\n"
                     ++ "given rho " ++ show (prettyAssocs rho)) $
#endif                
                abt '[] ('HMeasure a)
e
            -- TODO: if used exactly once in @e@, then inline.
            | Bool
otherwise ->
                case Lazy abt a -> Maybe (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> Maybe (Variable a)
getLazyVariable Lazy abt a
body of
                Just Variable a
y  -> Variable a
-> abt '[] a -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (b :: Hakaru).
ABT Term abt =>
Variable a -> abt '[] a -> abt xs b -> abt xs b
plug Variable a
x (Assocs (abt '[]) -> abt '[] a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
y)) abt '[] ('HMeasure a)
e
                Maybe (Variable a)
Nothing ->
                    case Lazy abt a -> Maybe (Literal a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> Maybe (Literal a)
getLazyLiteral Lazy abt a
body of
                    Just Literal a
v  -> Variable a
-> abt '[] a -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (b :: Hakaru).
ABT Term abt =>
Variable a -> abt '[] a -> abt xs b -> abt xs b
plug Variable a
x (Term abt a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt a -> abt '[] a) -> Term abt a -> abt '[] a
forall a b. (a -> b) -> a -> b
$ Literal a -> Term abt a
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
Literal a -> Term abt a
Literal_ Literal a
v) abt '[] ('HMeasure a)
e
                    Maybe (Literal a)
Nothing ->
                        Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC a, '( '[a], 'HMeasure a)] ('HMeasure a)
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], 'HMeasure a)] ('HMeasure a)
-> SArgs abt '[LC a, '( '[a], 'HMeasure a)]
-> Term abt ('HMeasure a)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ Assocs (abt '[]) -> abt '[] a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho (Lazy abt a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt a
body) abt '[] a
-> SArgs abt '[ '( '[a], 'HMeasure a)]
-> SArgs abt '[LC a, '( '[a], 'HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] ('HMeasure a) -> abt '[a] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x abt '[] ('HMeasure a)
e abt '[a] ('HMeasure a)
-> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)
        SGuard List1 Location xs
xs Pattern xs a
pat Lazy abt a
scrutinee [Index (abt '[])]
_ ->
            Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt ('HMeasure a) -> abt '[] ('HMeasure a))
-> Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall a b. (a -> b) -> a -> b
$ abt '[] a -> [Branch a abt ('HMeasure a)] -> Term abt ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
abt '[] a -> [Branch a abt b] -> Term abt b
Case_ (Assocs (abt '[]) -> abt '[] a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho (abt '[] a -> abt '[] a) -> abt '[] a -> abt '[] a
forall a b. (a -> b) -> a -> b
$ Lazy abt a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt a
scrutinee)
                [ Pattern xs a -> abt xs ('HMeasure a) -> Branch a abt ('HMeasure a)
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru)
       (xs :: [Hakaru]).
Pattern xs a -> abt xs b -> Branch a abt b
Branch Pattern xs a
pat   (List1 Variable xs -> abt '[] ('HMeasure a) -> abt xs ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (b :: k).
ABT syn abt =>
List1 Variable xs -> abt '[] b -> abt xs b
binds_ (List1 Location xs -> List1 Variable xs
forall k (a :: [k]). List1 Location a -> List1 Variable a
fromLocations1 List1 Location xs
xs) abt '[] ('HMeasure a)
e)
                , Pattern '[] a
-> abt '[] ('HMeasure a) -> Branch a abt ('HMeasure a)
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru)
       (xs :: [Hakaru]).
Pattern xs a -> abt xs b -> Branch a abt b
Branch Pattern '[] a
forall (a :: Hakaru). Pattern '[] a
PWild (Sing ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Sing ('HMeasure a) -> abt '[] ('HMeasure a)
P.reject (Sing ('HMeasure a) -> abt '[] ('HMeasure a))
-> Sing ('HMeasure a) -> abt '[] ('HMeasure a)
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] ('HMeasure a)
e)
                ]
        SWeight Lazy abt 'HProb
body [Index (abt '[])]
_ -> Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt ('HMeasure a) -> abt '[] ('HMeasure a))
-> Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall a b. (a -> b) -> a -> b
$ NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> Term abt ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> Term abt ('HMeasure a)
Superpose_ ((Assocs (abt '[]) -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Assocs (abt '[]) -> abt xs a -> abt xs a
plugs Assocs (abt '[])
rho (abt '[] 'HProb -> abt '[] 'HProb)
-> abt '[] 'HProb -> abt '[] 'HProb
forall a b. (a -> b) -> a -> b
$ Lazy abt 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt 'HProb
body, abt '[] ('HMeasure a)
e) (abt '[] 'HProb, abt '[] ('HMeasure a))
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
forall a. a -> [a] -> NonEmpty a
:| [])

----------------------------------------------------------------
-- An extra is a variable *use* instantiated at some list of indices.
data Extra :: (Hakaru -> *) -> Hakaru -> * where
     Loc :: Location a -> [ast 'HNat] -> Extra ast a

extrasInds :: Extra ast a -> [ast 'HNat]
extrasInds :: Extra ast a -> [ast 'HNat]
extrasInds (Loc       Location a
_ [ast 'HNat]
inds) = [ast 'HNat]
inds

selectMore :: [ast 'HNat] -> ast 'HNat -> [ast 'HNat]
selectMore :: [ast 'HNat] -> ast 'HNat -> [ast 'HNat]
selectMore = (ast 'HNat -> [ast 'HNat] -> [ast 'HNat])
-> [ast 'HNat] -> ast 'HNat -> [ast 'HNat]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)

-- Assumption: inds has no duplicates
permutes :: (ABT Term abt)
         => [abt '[] 'HNat]
         -> [Index (abt '[])]
         -> Bool
permutes :: [abt '[] 'HNat] -> [Index (abt '[])] -> Bool
permutes [abt '[] 'HNat]
ts [Index (abt '[])]
inds =
    (Maybe (Variable 'HNat) -> Bool)
-> [Maybe (Variable 'HNat)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe (Variable 'HNat) -> Bool
forall a. Maybe a -> Bool
isJust [Maybe (Variable 'HNat)]
ts' Bool -> Bool -> Bool
&&
    [Maybe (Variable 'HNat)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe (Variable 'HNat)]
ts' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Index (abt '[])] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (abt '[])]
inds Bool -> Bool -> Bool
&&
    [Variable 'HNat] -> Set (Variable 'HNat)
forall a. Ord a => [a] -> Set a
Set.fromList ((Maybe (Variable 'HNat) -> Variable 'HNat)
-> [Maybe (Variable 'HNat)] -> [Variable 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Maybe (Variable 'HNat) -> Variable 'HNat
forall a. HasCallStack => Maybe a -> a
fromJust [Maybe (Variable 'HNat)]
ts') Set (Variable 'HNat) -> Set (Variable 'HNat) -> Bool
forall a. Eq a => a -> a -> Bool
== [Variable 'HNat] -> Set (Variable 'HNat)
forall a. Ord a => [a] -> Set a
Set.fromList ((Index (abt '[]) -> Variable 'HNat)
-> [Index (abt '[])] -> [Variable 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> Variable 'HNat
forall (ast :: Hakaru -> *). Index ast -> Variable 'HNat
indVar [Index (abt '[])]
inds)
    where ts' :: [Maybe (Variable 'HNat)]
ts' = (abt '[] 'HNat -> Maybe (Variable 'HNat))
-> [abt '[] 'HNat] -> [Maybe (Variable 'HNat)]
forall a b. (a -> b) -> [a] -> [b]
map (\abt '[] 'HNat
t -> abt '[] 'HNat
-> (Variable 'HNat -> Maybe (Variable 'HNat))
-> (Term abt 'HNat -> Maybe (Variable 'HNat))
-> Maybe (Variable 'HNat)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] 'HNat
t Variable 'HNat -> Maybe (Variable 'HNat)
forall a. a -> Maybe a
Just (Maybe (Variable 'HNat) -> Term abt 'HNat -> Maybe (Variable 'HNat)
forall a b. a -> b -> a
const Maybe (Variable 'HNat)
forall a. Maybe a
Nothing)) [abt '[] 'HNat]
ts

#ifdef __TRACE_DISINTEGRATE__
      
prettyExtra :: (ABT Term abt) => Extra (abt '[]) a -> PP.Doc
prettyExtra (Loc (Location x) inds)      = PP.text "Loc" PP.<+> ppVariable x
                                           PP.<+> ppList (map pretty inds)

prettyExtras :: (ABT Term abt)
           => Assocs (Extra (abt '[]))
           -> PP.Doc
prettyExtras a = PP.vcat $ map go (fromAssocs a)
  where go (Assoc x l) = ppVariable x PP.<+>
                         PP.text "->" PP.<+>
                         prettyExtra l

#endif                           

-- In the paper we say that result must be a 'Whnf'; however, in
-- the paper it's also always @HMeasure a@ and everything of that
-- type is a WHNF (via 'WMeasure') so that's a trivial statement
-- to make. If we turn it back into some sort of normal form, then
-- it must be one preserved by 'residualizeContext'.
--
-- Also, we add the list in order to support "lub" without it living
-- in the AST.
-- TODO: really we should use LogicT...
type Ans abt a
  =  ListContext abt 'Impure
  -> Assocs (Extra (abt '[]))
  -> [abt '[] ('HMeasure a)]


----------------------------------------------------------------
-- TODO: defunctionalize the continuation. In particular, the only
-- heap modifications we need are 'push' and a variant of 'update'
-- for finding\/replacing a binding once we have the value in hand;
-- and the only 'freshNat' modifications are to allocate new 'Nat'.
-- We could defunctionalize the second arrow too by relying on the
-- @Codensity (ReaderT e m) ~= StateT e (Codensity m)@ isomorphism,
-- which makes explicit that the only thing other than 'ListContext'
-- updates is emitting something like @[Statement]@ to serve as the
-- beginning of the final result.
--
-- TODO: give this a better, more informative name!
--
-- N.B., This monad is used not only for both 'perform' and
-- 'constrainOutcome', but also for 'constrainValue'.
newtype Dis abt x =
    Dis { Dis abt x
-> forall (a :: Hakaru).
   [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a
unDis :: forall a. [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a }
    -- == @Codensity (Ans abt)@, assuming 'Codensity' is poly-kinded
    -- like it should be. If we don't want to allow continuations that
    -- can make nondeterministic choices, then we should use the right
    -- Kan extension itself, rather than the Codensity specialization
    -- of it.


-- | Run a computation in the 'Dis' monad, residualizing out all the
-- statements in the final evaluation context. The second argument
-- should include all the terms altered by the 'Dis' expression; this
-- is necessary to ensure proper hygiene; for example(s):
--
-- > runDis (perform e) [Some2 e]
-- > runDis (constrainOutcome e v) [Some2 e, Some2 v]
--
-- We use 'Some2' on the inputs because it doesn't matter what their
-- type or locally-bound variables are, so we want to allow @f@ to
-- contain terms with different indices.
runDisInCtx
    :: (ABT Term abt, F.Foldable f)
    => TransformCtx
    -> Dis abt (abt '[] a)
    -> f (Some2 abt)
    -> [abt '[] ('HMeasure a)]
runDisInCtx :: TransformCtx
-> Dis abt (abt '[] a) -> f (Some2 abt) -> [abt '[] ('HMeasure a)]
runDisInCtx TransformCtx
ctx Dis abt (abt '[] a)
d f (Some2 abt)
es =
    [Index (abt '[])]
-> ((abt '[] a, Assocs (abt '[])) -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])]
-> ((abt '[] a, Assocs (abt '[])) -> Ans abt a) -> Ans abt a
m0 [] (abt '[] a, Assocs (abt '[])) -> Ans abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru) p.
ABT Term abt =>
(abt '[] a, Assocs (abt '[]))
-> ListContext abt 'Impure -> p -> [abt '[] ('HMeasure a)]
c0 (Nat -> [Statement abt Location 'Impure] -> ListContext abt 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Nat -> [Statement abt Location p] -> ListContext abt p
ListContext Nat
i0 []) Assocs (Extra (abt '[]))
forall k (abt :: k -> *). Assocs abt
emptyAssocs
    where
    (Dis forall (a :: Hakaru).
[Index (abt '[])]
-> ((abt '[] a, Assocs (abt '[])) -> Ans abt a) -> Ans abt a
m0) = Dis abt (abt '[] a)
d Dis abt (abt '[] a)
-> (abt '[] a -> Dis abt (abt '[] a, Assocs (abt '[])))
-> Dis abt (abt '[] a, Assocs (abt '[]))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= abt '[] a -> Dis abt (abt '[] a, Assocs (abt '[]))
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] a -> Dis abt (abt '[] a, Assocs (abt '[]))
residualizeLocs
    -- TODO: we only use dirac because 'residualizeListContext'
    -- requires it to already be a measure; unfortunately this can
    -- result in an extraneous @(>>= \x -> dirac x)@ redex at the end
    -- of the program. In principle, we should be able to eliminate
    -- that redex by changing the type of 'residualizeListContext'...
    c0 :: (abt '[] a, Assocs (abt '[]))
-> ListContext abt 'Impure -> p -> [abt '[] ('HMeasure a)]
c0 (abt '[] a
e,Assocs (abt '[])
rho) ListContext abt 'Impure
ss p
_ = [ListContext abt 'Impure
-> Assocs (abt '[])
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
ListContext abt 'Impure
-> Assocs (abt '[])
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure a)
residualizeListContext ListContext abt 'Impure
ss Assocs (abt '[])
rho (Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn(SCon '[LC a] ('HMeasure a)
forall (a :: Hakaru). SCon '[LC a] ('HMeasure a)
Dirac SCon '[LC a] ('HMeasure a)
-> SArgs abt '[LC a] -> Term abt ('HMeasure a)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
e abt '[] a -> SArgs abt '[] -> SArgs abt '[LC a]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End))]
                  
    i0 :: Nat
i0 = f (Some2 abt) -> Nat
forall k2 (syn :: ([k2] -> k2 -> *) -> k2 -> *)
       (abt :: [k2] -> k2 -> *) (f :: * -> *).
(ABT syn abt, Foldable f) =>
f (Some2 abt) -> Nat
maxNextFree f (Some2 abt)
es Nat -> Nat -> Nat
forall a. Ord a => a -> a -> a
`max` TransformCtx -> Nat
nextFreeVar TransformCtx
ctx

runDis
    :: (ABT Term abt, F.Foldable f)
    => Dis abt (abt '[] a)
    -> f (Some2 abt)
    -> [abt '[] ('HMeasure a)]
runDis :: Dis abt (abt '[] a) -> f (Some2 abt) -> [abt '[] ('HMeasure a)]
runDis = TransformCtx
-> Dis abt (abt '[] a) -> f (Some2 abt) -> [abt '[] ('HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (f :: * -> *)
       (a :: Hakaru).
(ABT Term abt, Foldable f) =>
TransformCtx
-> Dis abt (abt '[] a) -> f (Some2 abt) -> [abt '[] ('HMeasure a)]
runDisInCtx TransformCtx
minimalCtx

{---------------------------------------------------------------------------------- 
 
 residualizeLocs does the following:
 1. update the heap by constructing plate/array around statements with nonempty indices
 2. use locations to construct terms out of var and "!" (for indexing into arrays)

For example, consider the state:

  list context (aka heap) =
  l1 <- lebesgue []
  l2 <- plate (normal 0 1) []
  l3 <- lebesgue [i]
  l4 <- dirac x3 []
  l5 <- normal 0 1 [j]
  
  assocs (aka locs) =
  x1 -> Loc l1 []
  x2 -> Loc l2 []
  x3 -> MultiLoc l3 []
  x4 -> Loc l4 []
  x5 -> Loc l5 [k]
  
Here the types of the above variables are:

  l1, x1 :: Real
  l2, x2 :: Array Real
  l3 :: Real
  x3 :: Array Real
  l4, x4 :: Array Real
  l5, x5 :: Real

Then residualizeLoc does two things:

1.Change the heap

  list context = 
  l1' <- lebesgue []
  l2' <- plate (normal 0 1) []
  l3' <- plate i (lebesgue)
  l4' <- dirac x3
  l5' <- plate j (normal 0 1)

2.Create new association table

  rho = 
  x1 -> var l1'
  x2 -> var l2'
  x3 -> array i' (l3' ! i')
  x4 -> var l4'
  x5 -> var l5' ! k 

----------------------------------------------------------------------------------}
residualizeLocs :: forall a abt. (ABT Term abt)
                => abt '[] a
                -> Dis abt (abt '[] a, Assocs (abt '[]))
residualizeLocs :: abt '[] a -> Dis abt (abt '[] a, Assocs (abt '[]))
residualizeLocs abt '[] a
e = do
  [Statement abt Location 'Impure]
ss <- Dis abt [Statement abt Location 'Impure]
forall (abt :: [Hakaru] -> Hakaru -> *).
Dis abt [Statement abt Location 'Impure]
getStatements
  ([Statement abt Location 'Impure]
ss', LAssocs Name
newlocs) <- (([Statement abt Location 'Impure], LAssocs Name)
 -> Statement abt Location 'Impure
 -> Dis abt ([Statement abt Location 'Impure], LAssocs Name))
-> ([Statement abt Location 'Impure], LAssocs Name)
-> [Statement abt Location 'Impure]
-> Dis abt ([Statement abt Location 'Impure], LAssocs Name)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Statement abt Location 'Impure], LAssocs Name)
-> Statement abt Location 'Impure
-> Dis abt ([Statement abt Location 'Impure], LAssocs Name)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
([Statement abt Location 'Impure], LAssocs Name)
-> Statement abt Location 'Impure
-> Dis abt ([Statement abt Location 'Impure], LAssocs Name)
step ([], LAssocs Name
forall k (abt :: k -> *). LAssocs abt
emptyLAssocs) [Statement abt Location 'Impure]
ss
  Assocs (abt '[])
rho <- LAssocs Name -> Dis abt (Assocs (abt '[]))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
LAssocs Name -> Dis abt (Assocs (abt '[]))
convertLocs LAssocs Name
newlocs
  [Statement abt Location 'Impure] -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
[Statement abt Location 'Impure] -> Dis abt ()
putStatements ([Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
forall a. [a] -> [a]
reverse [Statement abt Location 'Impure]
ss')
#ifdef __TRACE_DISINTEGRATE__
  trace ("residualizeLocs: old heap:\n" ++ show (pretty_Statements ss )) $ return ()
  trace ("residualizeLocs: new heap:\n" ++ show (pretty_Statements ss')) $ return ()
  extras <- getExtras
  traceM ("oldlocs:\n" ++ show (prettyExtras extras) ++ "\n")
  traceM ("new assoc for renaming:\n" ++ show (prettyAssocs rho))
#endif
  (abt '[] a, Assocs (abt '[]))
-> Dis abt (abt '[] a, Assocs (abt '[]))
forall (m :: * -> *) a. Monad m => a -> m a
return (abt '[] a
e, Assocs (abt '[])
rho)
    where step :: ([Statement abt Location 'Impure], LAssocs Name)
-> Statement abt Location 'Impure
-> Dis abt ([Statement abt Location 'Impure], LAssocs Name)
step ([Statement abt Location 'Impure]
ss',LAssocs Name
newlocs) Statement abt Location 'Impure
s = do (Statement abt Location 'Impure
s',LAssocs Name
newlocs') <- Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
residualizeLoc Statement abt Location 'Impure
s
                                    ([Statement abt Location 'Impure], LAssocs Name)
-> Dis abt ([Statement abt Location 'Impure], LAssocs Name)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s'Statement abt Location 'Impure
-> [Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
forall a. a -> [a] -> [a]
:[Statement abt Location 'Impure]
ss', LAssocs Name -> LAssocs Name -> LAssocs Name
forall k (ast :: k -> *). LAssocs ast -> LAssocs ast -> LAssocs ast
insertLAssocs LAssocs Name
newlocs' LAssocs Name
newlocs)

data Name (a :: Hakaru) = Name {Name a -> Text
nameHint :: Text, Name a -> Nat
nameID :: Nat}

locName :: Location a -> Name b
locName :: Location a -> Name b
locName (Location Variable a
x) = Text -> Nat -> Name b
forall (a :: Hakaru). Text -> Nat -> Name a
Name (Variable a -> Text
forall k (a :: k). Variable a -> Text
varHint Variable a
x) (Variable a -> Nat
forall k (a :: k). Variable a -> Nat
varID Variable a
x)

residualizeLoc :: (ABT Term abt)
               => Statement abt Location 'Impure
               -> Dis abt (Statement abt Location 'Impure, LAssocs Name)
residualizeLoc :: Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
residualizeLoc Statement abt Location 'Impure
s =
    case Statement abt Location 'Impure
s of
      SBind Location a
l Lazy abt ('HMeasure a)
_ [Index (abt '[])]
_ -> do 
             (Statement abt Location 'Impure
s', Name a
newname) <- Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement Statement abt Location 'Impure
s
             (Statement abt Location 'Impure, LAssocs Name)
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s', Location a -> Name a -> LAssocs Name
forall k (a :: k) (f :: k -> *). Location a -> f a -> LAssocs f
singletonLAssocs Location a
l Name a
newname)
      SLet  Location a
l Lazy abt a
_ [Index (abt '[])]
_ -> do
             (Statement abt Location 'Impure
s', Name a
newname) <- Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement Statement abt Location 'Impure
s
             (Statement abt Location 'Impure, LAssocs Name)
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s', Location a -> Name a -> LAssocs Name
forall k (a :: k) (f :: k -> *). Location a -> f a -> LAssocs f
singletonLAssocs Location a
l Name a
newname)
      SWeight Lazy abt 'HProb
w [Index (abt '[])]
inds    -> do
             Variable HUnit
x <- Text -> Sing HUnit -> Dis abt (Variable HUnit)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing HUnit
sUnit
             let bodyW :: Lazy abt ('HMeasure HUnit)
bodyW = abt '[] ('HMeasure HUnit) -> Lazy abt ('HMeasure HUnit)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk (abt '[] ('HMeasure HUnit) -> Lazy abt ('HMeasure HUnit))
-> abt '[] ('HMeasure HUnit) -> Lazy abt ('HMeasure HUnit)
forall a b. (a -> b) -> a -> b
$ abt '[] 'HProb -> abt '[] ('HMeasure HUnit)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HProb -> abt '[] ('HMeasure HUnit)
P.weight (Lazy abt 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt 'HProb
w)
             (Statement abt Location 'Impure
s', Name HUnit
newname) <- Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name HUnit)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement (Location HUnit
-> Lazy abt ('HMeasure HUnit)
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt v 'Impure
SBind (Variable HUnit -> Location HUnit
forall k (a :: k). Variable a -> Location a
Location Variable HUnit
x) Lazy abt ('HMeasure HUnit)
bodyW [Index (abt '[])]
inds)
             (Statement abt Location 'Impure, LAssocs Name)
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s', Location HUnit -> Name HUnit -> LAssocs Name
forall k (a :: k) (f :: k -> *). Location a -> f a -> LAssocs f
singletonLAssocs (Variable HUnit -> Location HUnit
forall k (a :: k). Variable a -> Location a
Location Variable HUnit
x) Name HUnit
newname)
      SGuard List1 Location xs
ls Pattern xs a
_ Lazy abt a
_ [Index (abt '[])]
ixs
        | [Index (abt '[])] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Index (abt '[])]
ixs  -> (Statement abt Location 'Impure, LAssocs Name)
-> Dis abt (Statement abt Location 'Impure, LAssocs Name)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s, List1 Location xs -> List1 Name xs -> LAssocs Name
forall k (xs :: [k]) (ast :: k -> *).
List1 Location xs -> List1 ast xs -> LAssocs ast
toLAssocs1 List1 Location xs
ls ((forall (i :: Hakaru). Location i -> Name i)
-> List1 Location xs -> List1 Name xs
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) (a :: k1 -> *)
       (b :: k1 -> *) (j :: k2).
Functor11 f =>
(forall (i :: k1). a i -> b i) -> f a j -> f b j
fmap11 forall k (a :: k) (b :: Hakaru). Location a -> Name b
forall (i :: Hakaru). Location i -> Name i
locName List1 Location xs
ls))
        | Bool
otherwise -> [Char] -> Dis abt (Statement abt Location 'Impure, LAssocs Name)
forall a. HasCallStack => [Char] -> a
error [Char]
"undefined: case statement under an array"

reifyStatement :: (ABT Term abt)
               => Statement abt Location 'Impure
               -> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement :: Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement Statement abt Location 'Impure
s =
    case Statement abt Location 'Impure
s of
      SBind Location a
l Lazy abt ('HMeasure a)
_    []     -> (Statement abt Location 'Impure, Name a)
-> Dis abt (Statement abt Location 'Impure, Name a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s, Location a -> Name a
forall k (a :: k) (b :: Hakaru). Location a -> Name b
locName Location a
l)
      SBind Location a
l Lazy abt ('HMeasure a)
body (Index (abt '[])
i:[Index (abt '[])]
is) -> do
        let plate :: abt '[] ('HMeasure a) -> Lazy abt ('HMeasure ('HArray a))
plate = abt '[] ('HMeasure ('HArray a)) -> Lazy abt ('HMeasure ('HArray a))
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk (abt '[] ('HMeasure ('HArray a))
 -> Lazy abt ('HMeasure ('HArray a)))
-> (abt '[] ('HMeasure a) -> abt '[] ('HMeasure ('HArray a)))
-> abt '[] ('HMeasure a)
-> Lazy abt ('HMeasure ('HArray a))
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. abt '[] 'HNat
-> Variable 'HNat
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure ('HArray a))
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] 'HNat
-> Variable 'HNat
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure ('HArray a))
P.plateWithVar (Index (abt '[]) -> abt '[] 'HNat
forall (ast :: Hakaru -> *). Index ast -> ast 'HNat
indSize Index (abt '[])
i) (Index (abt '[]) -> Variable 'HNat
forall (ast :: Hakaru -> *). Index ast -> Variable 'HNat
indVar Index (abt '[])
i)
        Variable ('HArray a)
x' <- Text -> Sing ('HArray a) -> Dis abt (Variable ('HArray a))
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar (Location a -> Text
forall k (a :: k). Location a -> Text
locHint Location a
l) (Sing a -> Sing ('HArray a)
forall (a :: Hakaru). Sing a -> Sing ('HArray a)
SArray (Location a -> Sing a
forall k (a :: k). Location a -> Sing a
locType Location a
l))
        Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement (Location ('HArray a)
-> Lazy abt ('HMeasure ('HArray a))
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt v 'Impure
SBind (Variable ('HArray a) -> Location ('HArray a)
forall k (a :: k). Variable a -> Location a
Location Variable ('HArray a)
x') (abt '[] ('HMeasure a) -> Lazy abt ('HMeasure ('HArray a))
plate (abt '[] ('HMeasure a) -> Lazy abt ('HMeasure ('HArray a)))
-> abt '[] ('HMeasure a) -> Lazy abt ('HMeasure ('HArray a))
forall a b. (a -> b) -> a -> b
$ Lazy abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt ('HMeasure a)
body) [Index (abt '[])]
is)
      SLet  Location a
l Lazy abt a
_    []     -> (Statement abt Location 'Impure, Name a)
-> Dis abt (Statement abt Location 'Impure, Name a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Statement abt Location 'Impure
s, Location a -> Name a
forall k (a :: k) (b :: Hakaru). Location a -> Name b
locName Location a
l)
      SLet  Location a
l Lazy abt a
body (Index (abt '[])
i:[Index (abt '[])]
is) -> do
        let array :: abt '[] a -> Lazy abt ('HArray a)
array = abt '[] ('HArray a) -> Lazy abt ('HArray a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk (abt '[] ('HArray a) -> Lazy abt ('HArray a))
-> (abt '[] a -> abt '[] ('HArray a))
-> abt '[] a
-> Lazy abt ('HArray a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. abt '[] 'HNat -> Variable 'HNat -> abt '[] a -> abt '[] ('HArray a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] 'HNat -> Variable 'HNat -> abt '[] a -> abt '[] ('HArray a)
P.arrayWithVar (Index (abt '[]) -> abt '[] 'HNat
forall (ast :: Hakaru -> *). Index ast -> ast 'HNat
indSize Index (abt '[])
i) (Index (abt '[]) -> Variable 'HNat
forall (ast :: Hakaru -> *). Index ast -> Variable 'HNat
indVar Index (abt '[])
i)
        Variable ('HArray a)
x' <- Text -> Sing ('HArray a) -> Dis abt (Variable ('HArray a))
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar (Location a -> Text
forall k (a :: k). Location a -> Text
locHint Location a
l) (Sing a -> Sing ('HArray a)
forall (a :: Hakaru). Sing a -> Sing ('HArray a)
SArray (Location a -> Sing a
forall k (a :: k). Location a -> Sing a
locType Location a
l))
        Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement (Location ('HArray a)
-> Lazy abt ('HArray a)
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity)
       (v :: Hakaru -> *) (a :: Hakaru).
v a -> Lazy abt a -> [Index (abt '[])] -> Statement abt v p
SLet  (Variable ('HArray a) -> Location ('HArray a)
forall k (a :: k). Variable a -> Location a
Location Variable ('HArray a)
x') (abt '[] a -> Lazy abt ('HArray a)
array (abt '[] a -> Lazy abt ('HArray a))
-> abt '[] a -> Lazy abt ('HArray a)
forall a b. (a -> b) -> a -> b
$ Lazy abt a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt a
body) [Index (abt '[])]
is)
      SWeight Lazy abt 'HProb
_    [Index (abt '[])]
_      -> [Char] -> Dis abt (Statement abt Location 'Impure, Name a)
forall a. HasCallStack => [Char] -> a
error [Char]
"reifyStatement called on SWeight"
      SGuard List1 Location xs
_ Pattern xs a
_ Lazy abt a
_ [Index (abt '[])]
_      -> [Char] -> Dis abt (Statement abt Location 'Impure, Name a)
forall a. HasCallStack => [Char] -> a
error [Char]
"reifyStatement called on SGuard"
                             
sizeInnermostInd :: (ABT Term abt)
                 => Location (a :: Hakaru)
                 -> Dis abt (abt '[] 'HNat)
sizeInnermostInd :: Location a -> Dis abt (abt '[] 'HNat)
sizeInnermostInd Location a
l =
    (Dis abt (abt '[] 'HNat)
-> (abt '[] 'HNat -> Dis abt (abt '[] 'HNat))
-> Maybe (abt '[] 'HNat)
-> Dis abt (abt '[] 'HNat)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Location a -> Dis abt (abt '[] 'HNat)
forall (a :: Hakaru) b. Location a -> b
freeLocError Location a
l) abt '[] 'HNat -> Dis abt (abt '[] 'HNat)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (abt '[] 'HNat) -> Dis abt (abt '[] 'HNat))
-> Dis abt (Maybe (abt '[] 'HNat)) -> Dis abt (abt '[] 'HNat)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (Dis abt (Maybe (abt '[] 'HNat)) -> Dis abt (abt '[] 'HNat))
-> ((Statement abt Location 'Impure
     -> Maybe (Dis abt (abt '[] 'HNat)))
    -> Dis abt (Maybe (abt '[] 'HNat)))
-> (Statement abt Location 'Impure
    -> Maybe (Dis abt (abt '[] 'HNat)))
-> Dis abt (abt '[] 'HNat)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Location a
-> (Statement abt Location 'Impure
    -> Maybe (Dis abt (abt '[] 'HNat)))
-> Dis abt (Maybe (abt '[] 'HNat))
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru) r.
EvaluationMonad abt m p =>
Location a
-> (Statement abt Location p -> Maybe (m r)) -> m (Maybe r)
select Location a
l ((Statement abt Location 'Impure
  -> Maybe (Dis abt (abt '[] 'HNat)))
 -> Dis abt (abt '[] 'HNat))
-> (Statement abt Location 'Impure
    -> Maybe (Dis abt (abt '[] 'HNat)))
-> Dis abt (abt '[] 'HNat)
forall a b. (a -> b) -> a -> b
$ \Statement abt Location 'Impure
s ->
        do Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard ([Index (abt '[])] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Statement abt Location 'Impure -> [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Statement abt Location p -> [Index (abt '[])]
statementInds Statement abt Location 'Impure
s) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1)
           case Statement abt Location 'Impure
s of
             SBind Location a
l' Lazy abt ('HMeasure a)
_ [Index (abt '[])]
ixs -> do TypeEq a a
Refl <- Location a -> Location a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Location a -> Location b -> Maybe (TypeEq a b)
locEq Location a
l Location a
l'
                                  Dis abt (abt '[] 'HNat) -> Maybe (Dis abt (abt '[] 'HNat))
forall a. a -> Maybe a
Just (Dis abt (abt '[] 'HNat) -> Maybe (Dis abt (abt '[] 'HNat)))
-> Dis abt (abt '[] 'HNat) -> Maybe (Dis abt (abt '[] 'HNat))
forall a b. (a -> b) -> a -> b
$ Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush Statement abt Location 'Impure
s Dis abt () -> Dis abt (abt '[] 'HNat) -> Dis abt (abt '[] 'HNat)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
                                         abt '[] 'HNat -> Dis abt (abt '[] 'HNat)
forall (m :: * -> *) a. Monad m => a -> m a
return (Index (abt '[]) -> abt '[] 'HNat
forall (ast :: Hakaru -> *). Index ast -> ast 'HNat
indSize ([Index (abt '[])] -> Index (abt '[])
forall a. [a] -> a
head [Index (abt '[])]
ixs))
             SLet  Location a
l' Lazy abt a
_ [Index (abt '[])]
ixs -> do TypeEq a a
Refl <- Location a -> Location a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Location a -> Location b -> Maybe (TypeEq a b)
locEq Location a
l Location a
l'
                                  Dis abt (abt '[] 'HNat) -> Maybe (Dis abt (abt '[] 'HNat))
forall a. a -> Maybe a
Just (Dis abt (abt '[] 'HNat) -> Maybe (Dis abt (abt '[] 'HNat)))
-> Dis abt (abt '[] 'HNat) -> Maybe (Dis abt (abt '[] 'HNat))
forall a b. (a -> b) -> a -> b
$ Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush Statement abt Location 'Impure
s Dis abt () -> Dis abt (abt '[] 'HNat) -> Dis abt (abt '[] 'HNat)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
                                         abt '[] 'HNat -> Dis abt (abt '[] 'HNat)
forall (m :: * -> *) a. Monad m => a -> m a
return (Index (abt '[]) -> abt '[] 'HNat
forall (ast :: Hakaru -> *). Index ast -> ast 'HNat
indSize ([Index (abt '[])] -> Index (abt '[])
forall a. [a] -> a
head [Index (abt '[])]
ixs))
             SWeight Lazy abt 'HProb
_ [Index (abt '[])]
_    -> Maybe (Dis abt (abt '[] 'HNat))
forall a. Maybe a
Nothing
             SGuard List1 Location xs
_ Pattern xs a
_ Lazy abt a
_ [Index (abt '[])]
_ -> [Char] -> Maybe (Dis abt (abt '[] 'HNat))
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: sizeInnermostInd{SGuard}"
                                         
fromName :: (ABT Term abt)
         => Name b
         -> Sing a
         -> [abt '[] 'HNat]
         -> abt '[] a
fromName :: Name b -> Sing a -> [abt '[] 'HNat] -> abt '[] a
fromName Name b
name Sing a
typ []     = Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a) -> Variable a -> abt '[] a
forall a b. (a -> b) -> a -> b
$ Variable :: forall k (a :: k). Text -> Nat -> Sing a -> Variable a
Variable { varHint :: Text
varHint = Name b -> Text
forall (a :: Hakaru). Name a -> Text
nameHint Name b
name
                                          , varID :: Nat
varID   = Name b -> Nat
forall (a :: Hakaru). Name a -> Nat
nameID Name b
name
                                          , varType :: Sing a
varType = Sing a
typ }
fromName Name b
name Sing a
typ (abt '[] 'HNat
i:[abt '[] 'HNat]
is) = Name b
-> Sing ('HArray a) -> [abt '[] 'HNat] -> abt '[] ('HArray a)
forall (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru)
       (a :: Hakaru).
ABT Term abt =>
Name b -> Sing a -> [abt '[] 'HNat] -> abt '[] a
fromName Name b
name (Sing a -> Sing ('HArray a)
forall (a :: Hakaru). Sing a -> Sing ('HArray a)
SArray Sing a
typ) [abt '[] 'HNat]
is abt '[] ('HArray a) -> abt '[] 'HNat -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HArray a) -> abt '[] 'HNat -> abt '[] a
P.! abt '[] 'HNat
i
                     
convertLocs :: (ABT Term abt)
            => LAssocs Name
            -> Dis abt (Assocs (abt '[]))
convertLocs :: LAssocs Name -> Dis abt (Assocs (abt '[]))
convertLocs LAssocs Name
newlocs = (Assoc (Extra (abt '[])) -> Assocs (abt '[]) -> Assocs (abt '[]))
-> Assocs (abt '[])
-> [Assoc (Extra (abt '[]))]
-> Assocs (abt '[])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr Assoc (Extra (abt '[])) -> Assocs (abt '[]) -> Assocs (abt '[])
step Assocs (abt '[])
forall k (abt :: k -> *). Assocs abt
emptyAssocs ([Assoc (Extra (abt '[]))] -> Assocs (abt '[]))
-> (Assocs (Extra (abt '[])) -> [Assoc (Extra (abt '[]))])
-> Assocs (Extra (abt '[]))
-> Assocs (abt '[])
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Assocs (Extra (abt '[])) -> [Assoc (Extra (abt '[]))]
forall k (ast :: k -> *). Assocs ast -> [Assoc ast]
fromAssocs (Assocs (Extra (abt '[])) -> Assocs (abt '[]))
-> Dis abt (Assocs (Extra (abt '[]))) -> Dis abt (Assocs (abt '[]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Dis abt (Assocs (Extra (abt '[])))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Dis abt (Assocs (Extra (abt '[])))
getExtras
    where
      build :: (ABT Term abt)
            => Assoc (Extra (abt '[]))
            -> Name a
            -> Assoc (abt '[])
      build :: Assoc (Extra (abt '[])) -> Name a -> Assoc (abt '[])
build (Assoc Variable a
x Extra (abt '[]) a
extra) Name a
name =
          Variable a -> abt '[] a -> Assoc (abt '[])
forall k (ast :: k -> *) (a :: k). Variable a -> ast a -> Assoc ast
Assoc Variable a
x (Name a -> Sing a -> [abt '[] 'HNat] -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru)
       (a :: Hakaru).
ABT Term abt =>
Name b -> Sing a -> [abt '[] 'HNat] -> abt '[] a
fromName Name a
name (Variable a -> Sing a
forall k (a :: k). Variable a -> Sing a
varType Variable a
x) (Extra (abt '[]) a -> [abt '[] 'HNat]
forall (ast :: Hakaru -> *) (a :: Hakaru).
Extra ast a -> [ast 'HNat]
extrasInds Extra (abt '[]) a
extra))
      step :: Assoc (Extra (abt '[])) -> Assocs (abt '[]) -> Assocs (abt '[])
step assoc :: Assoc (Extra (abt '[]))
assoc@(Assoc Variable a
_ Extra (abt '[]) a
extra) = Assoc (abt '[]) -> Assocs (abt '[]) -> Assocs (abt '[])
forall k (ast :: k -> *). Assoc ast -> Assocs ast -> Assocs ast
insertAssoc (Assoc (abt '[]) -> Assocs (abt '[]) -> Assocs (abt '[]))
-> Assoc (abt '[]) -> Assocs (abt '[]) -> Assocs (abt '[])
forall a b. (a -> b) -> a -> b
$
          case Extra (abt '[]) a
extra of
                 Loc      l _ -> Assoc (abt '[])
-> (Name a -> Assoc (abt '[])) -> Maybe (Name a) -> Assoc (abt '[])
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Location a -> Assoc (abt '[])
forall (a :: Hakaru) b. Location a -> b
freeLocError Location a
l)
                                       (Assoc (Extra (abt '[])) -> Name a -> Assoc (abt '[])
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Assoc (Extra (abt '[])) -> Name a -> Assoc (abt '[])
build Assoc (Extra (abt '[]))
assoc)
                                       (Location a -> LAssocs Name -> Maybe (Name a)
forall k (a :: k) (ast :: k -> *).
(Show1 Sing, JmEq1 Sing) =>
Location a -> LAssocs ast -> Maybe (ast a)
lookupLAssoc Location a
l LAssocs Name
newlocs)

freeLocError :: Location (a :: Hakaru) -> b
freeLocError :: Location a -> b
freeLocError Location a
l = [Char] -> b
forall a. HasCallStack => [Char] -> a
error ([Char] -> b) -> [Char] -> b
forall a b. (a -> b) -> a -> b
$ [Char]
"Found a free location " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Location a -> [Char]
forall a. Show a => a -> [Char]
show Location a
l

zipInds :: (ABT Term abt)
        => [Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
zipInds :: [Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
zipInds [Index (abt '[])]
inds [abt '[] 'HNat]
ts
    | [Index (abt '[])] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (abt '[])]
inds Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [abt '[] 'HNat] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [abt '[] 'HNat]
ts
        = [Char] -> Assocs (abt '[])
forall a. HasCallStack => [Char] -> a
error [Char]
"zipInds: argument lists must have the same length"
    | Bool
otherwise = [Assoc (abt '[])] -> Assocs (abt '[])
forall k (ast :: k -> *). [Assoc ast] -> Assocs ast
toAssocs ([Assoc (abt '[])] -> Assocs (abt '[]))
-> [Assoc (abt '[])] -> Assocs (abt '[])
forall a b. (a -> b) -> a -> b
$ (Variable 'HNat -> abt '[] 'HNat -> Assoc (abt '[]))
-> [Variable 'HNat] -> [abt '[] 'HNat] -> [Assoc (abt '[])]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Variable 'HNat -> abt '[] 'HNat -> Assoc (abt '[])
forall k (ast :: k -> *) (a :: k). Variable a -> ast a -> Assoc ast
Assoc ((Index (abt '[]) -> Variable 'HNat)
-> [Index (abt '[])] -> [Variable 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> Variable 'HNat
forall (ast :: Hakaru -> *). Index ast -> Variable 'HNat
indVar [Index (abt '[])]
inds) [abt '[] 'HNat]
ts

apply :: (ABT Term abt)
      => [Index (abt '[])]
      -> [Index (abt '[])]
      -> abt '[] a
      -> Dis abt (abt '[] a)
apply :: [Index (abt '[])]
-> [Index (abt '[])] -> abt '[] a -> Dis abt (abt '[] a)
apply [Index (abt '[])]
is [Index (abt '[])]
js abt '[] a
e = Assocs (abt '[]) -> abt '[] a -> Dis abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (m :: * -> *) (p :: Purity).
EvaluationMonad abt m p =>
Assocs (abt '[]) -> abt xs a -> m (abt xs a)
extSubsts ([Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
[Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
zipInds [Index (abt '[])]
is ((Index (abt '[]) -> abt '[] 'HNat)
-> [Index (abt '[])] -> [abt '[] 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> abt '[] 'HNat
fromIndex [Index (abt '[])]
js)) abt '[] a
e
                           
extendIndices
    :: (ABT Term abt)
    => Index (abt '[])
    -> [Index (abt '[])]
    -> [Index (abt '[])]
-- TODO: check all Indices are unique
extendIndices :: Index (abt '[]) -> [Index (abt '[])] -> [Index (abt '[])]
extendIndices Index (abt '[])
j [Index (abt '[])]
js | Index (abt '[])
j Index (abt '[]) -> [Index (abt '[])] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Index (abt '[])]
js
                   = [Char] -> [Index (abt '[])]
forall a. HasCallStack => [Char] -> a
error ([Char]
"Duplicate index between " )
                     -- TODO finish this error message by
                     -- defining Show for Index
                   | Bool
otherwise
                   = Index (abt '[])
j Index (abt '[]) -> [Index (abt '[])] -> [Index (abt '[])]
forall a. a -> [a] -> [a]
: [Index (abt '[])]
js

-- give better name
statementInds :: Statement abt Location p -> [Index (abt '[])]
statementInds :: Statement abt Location p -> [Index (abt '[])]
statementInds (SBind   Location a
_ Lazy abt ('HMeasure a)
_   [Index (abt '[])]
i) = [Index (abt '[])]
i
statementInds (SLet    Location a
_ Lazy abt a
_   [Index (abt '[])]
i) = [Index (abt '[])]
i
statementInds (SWeight Lazy abt 'HProb
_     [Index (abt '[])]
i) = [Index (abt '[])]
i
statementInds (SGuard  List1 Location xs
_ Pattern xs a
_ Lazy abt a
_ [Index (abt '[])]
i) = [Index (abt '[])]
i
statementInds (SStuff0 abt '[] 'HProb -> abt '[] 'HProb
_     [Index (abt '[])]
i) = [Index (abt '[])]
i
statementInds (SStuff1 Location a
_ abt '[] 'HProb -> abt '[] 'HProb
_   [Index (abt '[])]
i) = [Index (abt '[])]
i

getExtras :: (ABT Term abt)
          => Dis abt (Assocs (Extra (abt '[])))
getExtras :: Dis abt (Assocs (Extra (abt '[])))
getExtras = (forall (a :: Hakaru).
 [Index (abt '[])]
 -> (Assocs (Extra (abt '[])) -> Ans abt a) -> Ans abt a)
-> Dis abt (Assocs (Extra (abt '[])))
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])]
  -> (Assocs (Extra (abt '[])) -> Ans abt a) -> Ans abt a)
 -> Dis abt (Assocs (Extra (abt '[]))))
-> (forall (a :: Hakaru).
    [Index (abt '[])]
    -> (Assocs (Extra (abt '[])) -> Ans abt a) -> Ans abt a)
-> Dis abt (Assocs (Extra (abt '[])))
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ Assocs (Extra (abt '[])) -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l -> Assocs (Extra (abt '[])) -> Ans abt a
c Assocs (Extra (abt '[]))
l ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l

putExtras :: (ABT Term abt)
          => Assocs (Extra (abt '[]))
          -> Dis abt ()
putExtras :: Assocs (Extra (abt '[])) -> Dis abt ()
putExtras Assocs (Extra (abt '[]))
l = (forall (a :: Hakaru).
 [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
 -> Dis abt ())
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ () -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
_ -> () -> Ans abt a
c () ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l

insertExtra :: (ABT Term abt)
            => Variable a
            -> Extra (abt '[]) a
            -> Dis abt ()
insertExtra :: Variable a -> Extra (abt '[]) a -> Dis abt ()
insertExtra Variable a
v Extra (abt '[]) a
extra = 
  (forall (a :: Hakaru).
 [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
 -> Dis abt ())
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ () -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l -> () -> Ans abt a
c () ListContext abt 'Impure
h (Assocs (Extra (abt '[])) -> [abt '[] ('HMeasure a)])
-> Assocs (Extra (abt '[])) -> [abt '[] ('HMeasure a)]
forall a b. (a -> b) -> a -> b
$
    Assoc (Extra (abt '[]))
-> Assocs (Extra (abt '[])) -> Assocs (Extra (abt '[]))
forall k (ast :: k -> *). Assoc ast -> Assocs ast -> Assocs ast
insertAssoc (Variable a -> Extra (abt '[]) a -> Assoc (Extra (abt '[]))
forall k (ast :: k -> *) (a :: k). Variable a -> ast a -> Assoc ast
Assoc Variable a
v Extra (abt '[]) a
extra) Assocs (Extra (abt '[]))
l

adjustExtra :: (ABT Term abt)
            => Variable (a :: Hakaru)
            -> (Assoc (Extra (abt '[])) -> Assoc (Extra (abt '[])))
            -> Dis abt ()
adjustExtra :: Variable a
-> (Assoc (Extra (abt '[])) -> Assoc (Extra (abt '[])))
-> Dis abt ()
adjustExtra Variable a
x Assoc (Extra (abt '[])) -> Assoc (Extra (abt '[]))
f = do
    Assocs (Extra (abt '[]))
extras <- Dis abt (Assocs (Extra (abt '[])))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Dis abt (Assocs (Extra (abt '[])))
getExtras
    Assocs (Extra (abt '[])) -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Assocs (Extra (abt '[])) -> Dis abt ()
putExtras (Assocs (Extra (abt '[])) -> Dis abt ())
-> Assocs (Extra (abt '[])) -> Dis abt ()
forall a b. (a -> b) -> a -> b
$ Variable a
-> (Assoc (Extra (abt '[])) -> Assoc (Extra (abt '[])))
-> Assocs (Extra (abt '[]))
-> Assocs (Extra (abt '[]))
forall k1 k2 (a :: k2) (ast :: k1 -> *).
Variable a -> (Assoc ast -> Assoc ast) -> Assocs ast -> Assocs ast
adjustAssoc Variable a
x Assoc (Extra (abt '[])) -> Assoc (Extra (abt '[]))
f Assocs (Extra (abt '[]))
extras

mkLoc
    :: (ABT Term abt)
    => Text
    -> Location (a :: Hakaru)
    -> [abt '[] 'HNat]
    -> Dis abt (Variable a)
mkLoc :: Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
mkLoc Text
hint Location a
l [abt '[] 'HNat]
inds = do
  Variable a
x <- Text -> Sing a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
hint (Location a -> Sing a
forall k (a :: k). Location a -> Sing a
locType Location a
l)
  Variable a -> Extra (abt '[]) a -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Variable a -> Extra (abt '[]) a -> Dis abt ()
insertExtra Variable a
x (Location a -> [abt '[] 'HNat] -> Extra (abt '[]) a
forall (a :: Hakaru) (ast :: Hakaru -> *).
Location a -> [ast 'HNat] -> Extra ast a
Loc Location a
l [abt '[] 'HNat]
inds)
  Variable a -> Dis abt (Variable a)
forall (m :: * -> *) a. Monad m => a -> m a
return Variable a
x

mkLocs
    :: (ABT Term abt)
    => List1 Location (xs :: [Hakaru])
    -> [abt '[] 'HNat]
    -> Dis abt (List1 Variable xs)
mkLocs :: List1 Location xs -> [abt '[] 'HNat] -> Dis abt (List1 Variable xs)
mkLocs List1 Location xs
Nil1         [abt '[] 'HNat]
_    = List1 Variable '[] -> Dis abt (List1 Variable '[])
forall (m :: * -> *) a. Monad m => a -> m a
return List1 Variable '[]
forall k (a :: k -> *). List1 a '[]
Nil1
mkLocs (Cons1 Location x
l List1 Location xs
ls) [abt '[] 'HNat]
inds = Variable x -> List1 Variable xs -> List1 Variable (x : xs)
forall a (a :: a -> *) (x :: a) (xs :: [a]).
a x -> List1 a xs -> List1 a (x : xs)
Cons1
                           (Variable x -> List1 Variable xs -> List1 Variable (x : xs))
-> Dis abt (Variable x)
-> Dis abt (List1 Variable xs -> List1 Variable (x : xs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Location x -> [abt '[] 'HNat] -> Dis abt (Variable x)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
mkLoc Text
Text.empty Location x
l [abt '[] 'HNat]
inds
                           Dis abt (List1 Variable xs -> List1 Variable (x : xs))
-> Dis abt (List1 Variable xs) -> Dis abt (List1 Variable (x : xs))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> List1 Location xs -> [abt '[] 'HNat] -> Dis abt (List1 Variable xs)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru]).
ABT Term abt =>
List1 Location xs -> [abt '[] 'HNat] -> Dis abt (List1 Variable xs)
mkLocs List1 Location xs
ls [abt '[] 'HNat]
inds

instance Functor (Dis abt) where
    fmap :: (a -> b) -> Dis abt a -> Dis abt b
fmap a -> b
f (Dis forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m)  = (forall (a :: Hakaru).
 [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
-> Dis abt b
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
 -> Dis abt b)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
-> Dis abt b
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i b -> Ans abt a
c -> [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m [Index (abt '[])]
i (b -> Ans abt a
c (b -> Ans abt a) -> (a -> b) -> a -> Ans abt a
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> b
f)

instance Applicative (Dis abt) where
    pure :: a -> Dis abt a
pure a
x            = (forall (a :: Hakaru).
 [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
 -> Dis abt a)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ a -> Ans abt a
c -> a -> Ans abt a
c a
x
    Dis forall (a :: Hakaru).
[Index (abt '[])] -> ((a -> b) -> Ans abt a) -> Ans abt a
mf <*> :: Dis abt (a -> b) -> Dis abt a -> Dis abt b
<*> Dis forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
mx = (forall (a :: Hakaru).
 [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
-> Dis abt b
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
 -> Dis abt b)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
-> Dis abt b
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i b -> Ans abt a
c -> [Index (abt '[])] -> ((a -> b) -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> ((a -> b) -> Ans abt a) -> Ans abt a
mf [Index (abt '[])]
i (((a -> b) -> Ans abt a) -> Ans abt a)
-> ((a -> b) -> Ans abt a) -> Ans abt a
forall a b. (a -> b) -> a -> b
$ \a -> b
f -> [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
mx [Index (abt '[])]
i ((a -> Ans abt a) -> Ans abt a) -> (a -> Ans abt a) -> Ans abt a
forall a b. (a -> b) -> a -> b
$ \a
x -> b -> Ans abt a
c (a -> b
f a
x)

instance Monad (Dis abt) where
    return :: a -> Dis abt a
return      = a -> Dis abt a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    Dis forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m >>= :: Dis abt a -> (a -> Dis abt b) -> Dis abt b
>>= a -> Dis abt b
k = (forall (a :: Hakaru).
 [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
-> Dis abt b
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
 -> Dis abt b)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a)
-> Dis abt b
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i b -> Ans abt a
c -> [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m [Index (abt '[])]
i ((a -> Ans abt a) -> Ans abt a) -> (a -> Ans abt a) -> Ans abt a
forall a b. (a -> b) -> a -> b
$ \a
x -> Dis abt b -> [Index (abt '[])] -> (b -> Ans abt a) -> Ans abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
Dis abt x
-> forall (a :: Hakaru).
   [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a
unDis (a -> Dis abt b
k a
x) [Index (abt '[])]
i b -> Ans abt a
c

#if __GLASGOW_HASKELL__ > 805
instance MonadFail (Dis abt) where
    fail :: [Char] -> Dis abt a
fail        = [Char] -> Dis abt a
forall a. HasCallStack => [Char] -> a
error
#endif

instance Alternative (Dis abt) where
    empty :: Dis abt a
empty           = (forall (a :: Hakaru).
 [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
 -> Dis abt a)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ a -> Ans abt a
_ ListContext abt 'Impure
_ Assocs (Extra (abt '[]))
_ -> []
    Dis forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m <|> :: Dis abt a -> Dis abt a -> Dis abt a
<|> Dis forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
n = (forall (a :: Hakaru).
 [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
 -> Dis abt a)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l -> [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m [Index (abt '[])]
i a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l [abt '[] ('HMeasure a)]
-> [abt '[] ('HMeasure a)] -> [abt '[] ('HMeasure a)]
forall a. [a] -> [a] -> [a]
++ [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
n [Index (abt '[])]
i a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l

instance MonadPlus (Dis abt) where
    mzero :: Dis abt a
mzero = Dis abt a
forall (f :: * -> *) a. Alternative f => f a
empty -- aka "bot"
    mplus :: Dis abt a -> Dis abt a -> Dis abt a
mplus = Dis abt a -> Dis abt a -> Dis abt a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>) -- aka "lub"

instance (ABT Term abt) => EvaluationMonad abt (Dis abt) 'Impure where
    freshNat :: Dis abt Nat
freshNat =
        (forall (a :: Hakaru).
 [Index (abt '[])] -> (Nat -> Ans abt a) -> Ans abt a)
-> Dis abt Nat
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (Nat -> Ans abt a) -> Ans abt a)
 -> Dis abt Nat)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (Nat -> Ans abt a) -> Ans abt a)
-> Dis abt Nat
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ Nat -> Ans abt a
c (ListContext Nat
n [Statement abt Location 'Impure]
ss) ->
            Nat -> Ans abt a
c Nat
n (Nat -> [Statement abt Location 'Impure] -> ListContext abt 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Nat -> [Statement abt Location p] -> ListContext abt p
ListContext (Nat
nNat -> Nat -> Nat
forall a. Num a => a -> a -> a
+Nat
1) [Statement abt Location 'Impure]
ss)

    -- Note: we assume that freshLocStatement is never called on a
    -- statement already on the heap (list context)
    freshLocStatement :: Statement abt Variable 'Impure
-> Dis abt (Statement abt Location 'Impure, Assocs Variable)
freshLocStatement Statement abt Variable 'Impure
s =
        case Statement abt Variable 'Impure
s of
          SWeight Lazy abt 'HProb
w [Index (abt '[])]
e    -> (Statement abt Location 'Impure, Assocs Variable)
-> Dis abt (Statement abt Location 'Impure, Assocs Variable)
forall (m :: * -> *) a. Monad m => a -> m a
return (Lazy abt 'HProb
-> [Index (abt '[])] -> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
Lazy abt 'HProb -> [Index (abt '[])] -> Statement abt v 'Impure
SWeight Lazy abt 'HProb
w [Index (abt '[])]
e, Assocs Variable
forall a. Monoid a => a
mempty)
          SBind Variable a
x Lazy abt ('HMeasure a)
body [Index (abt '[])]
i -> do
               Variable a
x' <- Variable a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Variable a -> m (Variable a)
freshenVar Variable a
x
               let l :: Location a
l = Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location Variable a
x'
               Variable a
v  <- Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
mkLoc (Location a -> Text
forall k (a :: k). Location a -> Text
locHint Location a
l) Location a
l ((Index (abt '[]) -> abt '[] 'HNat)
-> [Index (abt '[])] -> [abt '[] 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> abt '[] 'HNat
fromIndex [Index (abt '[])]
i)
               (Statement abt Location 'Impure, Assocs Variable)
-> Dis abt (Statement abt Location 'Impure, Assocs Variable)
forall (m :: * -> *) a. Monad m => a -> m a
return (Location a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt v 'Impure
SBind Location a
l Lazy abt ('HMeasure a)
body [Index (abt '[])]
i, Variable a -> Variable a -> Assocs Variable
forall k (a :: k) (f :: k -> *). Variable a -> f a -> Assocs f
singletonAssocs Variable a
x Variable a
v)
          SLet  Variable a
x Lazy abt a
body [Index (abt '[])]
i -> do
               Variable a
x' <- Variable a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Variable a -> m (Variable a)
freshenVar Variable a
x
               let l :: Location a
l = Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location Variable a
x'
               Variable a
v  <- Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
mkLoc (Location a -> Text
forall k (a :: k). Location a -> Text
locHint Location a
l) Location a
l ((Index (abt '[]) -> abt '[] 'HNat)
-> [Index (abt '[])] -> [abt '[] 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> abt '[] 'HNat
fromIndex [Index (abt '[])]
i)
               (Statement abt Location 'Impure, Assocs Variable)
-> Dis abt (Statement abt Location 'Impure, Assocs Variable)
forall (m :: * -> *) a. Monad m => a -> m a
return (Location a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity)
       (v :: Hakaru -> *) (a :: Hakaru).
v a -> Lazy abt a -> [Index (abt '[])] -> Statement abt v p
SLet Location a
l Lazy abt a
body [Index (abt '[])]
i, Variable a -> Variable a -> Assocs Variable
forall k (a :: k) (f :: k -> *). Variable a -> f a -> Assocs f
singletonAssocs Variable a
x Variable a
v)
          SGuard List1 Variable xs
xs Pattern xs a
pat Lazy abt a
scrutinee [Index (abt '[])]
i -> do
               List1 Variable xs
xs'  <- List1 Variable xs -> Dis abt (List1 Variable xs)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (xs :: [Hakaru]).
EvaluationMonad abt m p =>
List1 Variable xs -> m (List1 Variable xs)
freshenVars List1 Variable xs
xs
               let ls :: List1 Location xs
ls = List1 Variable xs -> List1 Location xs
forall k (a :: [k]). List1 Variable a -> List1 Location a
locations1 List1 Variable xs
xs'
               List1 Variable xs
vs   <- List1 Location xs -> [abt '[] 'HNat] -> Dis abt (List1 Variable xs)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru]).
ABT Term abt =>
List1 Location xs -> [abt '[] 'HNat] -> Dis abt (List1 Variable xs)
mkLocs List1 Location xs
ls ((Index (abt '[]) -> abt '[] 'HNat)
-> [Index (abt '[])] -> [abt '[] 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> abt '[] 'HNat
fromIndex [Index (abt '[])]
i)
               (Statement abt Location 'Impure, Assocs Variable)
-> Dis abt (Statement abt Location 'Impure, Assocs Variable)
forall (m :: * -> *) a. Monad m => a -> m a
return (List1 Location xs
-> Pattern xs a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (xs :: [Hakaru]) (a :: Hakaru).
List1 v xs
-> Pattern xs a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt v 'Impure
SGuard List1 Location xs
ls Pattern xs a
pat Lazy abt a
scrutinee [Index (abt '[])]
i, List1 Variable xs -> List1 Variable xs -> Assocs Variable
forall k (xs :: [k]) (ast :: k -> *).
List1 Variable xs -> List1 ast xs -> Assocs ast
toAssocs1 List1 Variable xs
xs List1 Variable xs
vs)

    getIndices :: Dis abt [Index (abt '[])]
getIndices =  (forall (a :: Hakaru).
 [Index (abt '[])] -> ([Index (abt '[])] -> Ans abt a) -> Ans abt a)
-> Dis abt [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> ([Index (abt '[])] -> Ans abt a) -> Ans abt a)
 -> Dis abt [Index (abt '[])])
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> ([Index (abt '[])] -> Ans abt a) -> Ans abt a)
-> Dis abt [Index (abt '[])]
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i [Index (abt '[])] -> Ans abt a
c -> [Index (abt '[])] -> Ans abt a
c [Index (abt '[])]
i

    unsafePush :: Statement abt Location 'Impure -> Dis abt ()
unsafePush Statement abt Location 'Impure
s =
        (forall (a :: Hakaru).
 [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
 -> Dis abt ())
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ () -> Ans abt a
c (ListContext Nat
i [Statement abt Location 'Impure]
ss) ->
            () -> Ans abt a
c () (Nat -> [Statement abt Location 'Impure] -> ListContext abt 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Nat -> [Statement abt Location p] -> ListContext abt p
ListContext Nat
i (Statement abt Location 'Impure
sStatement abt Location 'Impure
-> [Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
forall a. a -> [a] -> [a]
:[Statement abt Location 'Impure]
ss))

    -- N.B., the use of 'reverse' is necessary so that the order
    -- of pushing matches that of 'pushes'
    unsafePushes :: [Statement abt Location 'Impure] -> Dis abt ()
unsafePushes [Statement abt Location 'Impure]
ss =
        (forall (a :: Hakaru).
 [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
 -> Dis abt ())
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ () -> Ans abt a
c (ListContext Nat
i [Statement abt Location 'Impure]
ss') ->
            () -> Ans abt a
c () (Nat -> [Statement abt Location 'Impure] -> ListContext abt 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Nat -> [Statement abt Location p] -> ListContext abt p
ListContext Nat
i ([Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
forall a. [a] -> [a]
reverse [Statement abt Location 'Impure]
ss [Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
forall a. [a] -> [a] -> [a]
++ [Statement abt Location 'Impure]
ss'))

    select :: Location a
-> (Statement abt Location 'Impure -> Maybe (Dis abt r))
-> Dis abt (Maybe r)
select Location a
l Statement abt Location 'Impure -> Maybe (Dis abt r)
p = [Statement abt Location 'Impure] -> Dis abt (Maybe r)
loop []
        where
        -- TODO: use a DList to avoid reversing inside 'unsafePushes'
        loop :: [Statement abt Location 'Impure] -> Dis abt (Maybe r)
loop [Statement abt Location 'Impure]
ss = do
            Maybe (Statement abt Location 'Impure)
ms <- Dis abt (Maybe (Statement abt Location 'Impure))
forall (abt :: [Hakaru] -> Hakaru -> *).
Dis abt (Maybe (Statement abt Location 'Impure))
unsafePop
            case Maybe (Statement abt Location 'Impure)
ms of
                Maybe (Statement abt Location 'Impure)
Nothing      -> do
                    [Statement abt Location 'Impure] -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
[Statement abt Location p] -> m ()
unsafePushes [Statement abt Location 'Impure]
ss
                    Maybe r -> Dis abt (Maybe r)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe r
forall a. Maybe a
Nothing
                Just Statement abt Location 'Impure
s ->
                    -- Alas, @p@ will have to recheck 'isBoundBy'
                    -- in order to grab the 'Refl' proof we erased;
                    -- but there's nothing to be done for it.
                    case Location a
l Location a -> Statement abt Location 'Impure -> Maybe ()
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *)
       (p :: Purity).
Location a -> Statement abt Location p -> Maybe ()
`isBoundBy` Statement abt Location 'Impure
s Maybe () -> Maybe (Dis abt r) -> Maybe (Dis abt r)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Statement abt Location 'Impure -> Maybe (Dis abt r)
p Statement abt Location 'Impure
s of
                    Maybe (Dis abt r)
Nothing -> [Statement abt Location 'Impure] -> Dis abt (Maybe r)
loop (Statement abt Location 'Impure
sStatement abt Location 'Impure
-> [Statement abt Location 'Impure]
-> [Statement abt Location 'Impure]
forall a. a -> [a] -> [a]
:[Statement abt Location 'Impure]
ss)
                    Just Dis abt r
mr -> do
                        r
r <- Dis abt r
mr
                        [Statement abt Location 'Impure] -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
[Statement abt Location p] -> m ()
unsafePushes [Statement abt Location 'Impure]
ss
                        Maybe r -> Dis abt (Maybe r)
forall (m :: * -> *) a. Monad m => a -> m a
return (r -> Maybe r
forall a. a -> Maybe a
Just r
r)

    substVar :: Variable a
-> abt '[] a
-> forall (b' :: Hakaru). Variable b' -> Dis abt (abt '[] b')
substVar Variable a
x abt '[] a
e Variable b'
z = do
      Assocs (Extra (abt '[]))
extras <- Dis abt (Assocs (Extra (abt '[])))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Dis abt (Assocs (Extra (abt '[])))
getExtras
      let defaultResult :: Dis abt (abt '[] b')
defaultResult = abt '[] b' -> Dis abt (abt '[] b')
forall (m :: * -> *) a. Monad m => a -> m a
return (Variable b' -> abt '[] b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable b'
z)
      case (Variable b'
-> Assocs (Extra (abt '[])) -> Maybe (Extra (abt '[]) b')
forall k (a :: k) (ast :: k -> *).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Assocs ast -> Maybe (ast a)
lookupAssoc Variable b'
z Assocs (Extra (abt '[]))
extras) of
        Maybe (Extra (abt '[]) b')
Nothing                -> Dis abt (abt '[] b')
defaultResult
        Just (Loc Location b'
l [abt '[] 'HNat]
inds)      ->
            if (abt '[] 'HNat -> Bool) -> [abt '[] 'HNat] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Variable a -> VarSet (KindOf ('HMeasure a)) -> Bool
forall k (a :: k) (kproxy :: KProxy k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> VarSet kproxy -> Bool
memberVarSet Variable a
x (VarSet (KindOf ('HMeasure a)) -> Bool)
-> (abt '[] 'HNat -> VarSet (KindOf ('HMeasure a)))
-> abt '[] 'HNat
-> Bool
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. abt '[] 'HNat -> VarSet (KindOf ('HMeasure a))
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> VarSet (KindOf a)
freeVars) [abt '[] 'HNat]
inds
            then do [abt '[] 'HNat]
inds' <- (abt '[] 'HNat -> Dis abt (abt '[] 'HNat))
-> [abt '[] 'HNat] -> Dis abt [abt '[] 'HNat]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Variable a -> abt '[] a -> abt '[] 'HNat -> Dis abt (abt '[] 'HNat)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (b :: Hakaru) (m :: * -> *) (p :: Purity).
EvaluationMonad abt m p =>
Variable a -> abt '[] a -> abt xs b -> m (abt xs b)
extSubst Variable a
x abt '[] a
e) [abt '[] 'HNat]
inds
                    Variable b' -> abt '[] b'
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable b' -> abt '[] b')
-> Dis abt (Variable b') -> Dis abt (abt '[] b')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Location b' -> [abt '[] 'HNat] -> Dis abt (Variable b')
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
mkLoc Text
Text.empty Location b'
l [abt '[] 'HNat]
inds'
            else Dis abt (abt '[] b')
defaultResult

    extFreeVars :: abt xs a -> Dis abt (VarSet (KindOf ('HMeasure a)))
extFreeVars abt xs a
e = do
      Assocs (Extra (abt '[]))
extras <- Dis abt (Assocs (Extra (abt '[])))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Dis abt (Assocs (Extra (abt '[])))
getExtras
      let fvs1 :: VarSet (KindOf ('HMeasure a))
fvs1 = abt xs a -> VarSet (KindOf ('HMeasure a))
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> VarSet (KindOf a)
freeVars abt xs a
e
          indFVs :: SomeVariable (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a))
indFVs (SomeVariable Variable a
v) =
              case (Variable a -> Assocs (Extra (abt '[])) -> Maybe (Extra (abt '[]) a)
forall k (a :: k) (ast :: k -> *).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Assocs ast -> Maybe (ast a)
lookupAssoc Variable a
v Assocs (Extra (abt '[]))
extras) of
                Maybe (Extra (abt '[]) a)
Nothing -> VarSet (KindOf ('HMeasure a))
forall k (kproxy :: KProxy k). VarSet kproxy
emptyVarSet
                Just (Loc _ is) -> (abt '[] 'HNat
 -> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a)))
-> VarSet (KindOf ('HMeasure a))
-> [abt '[] 'HNat]
-> VarSet (KindOf ('HMeasure a))
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VarSet (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a))
forall k (kproxy :: KProxy k).
(Show1 Sing, JmEq1 Sing) =>
VarSet kproxy -> VarSet kproxy -> VarSet kproxy
unionVarSet(VarSet (KindOf ('HMeasure a))
 -> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a)))
-> (abt '[] 'HNat -> VarSet (KindOf ('HMeasure a)))
-> abt '[] 'HNat
-> VarSet (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a))
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
.abt '[] 'HNat -> VarSet (KindOf ('HMeasure a))
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> VarSet (KindOf a)
freeVars) VarSet (KindOf ('HMeasure a))
forall k (kproxy :: KProxy k). VarSet kproxy
emptyVarSet [abt '[] 'HNat]
is
          locVars :: SomeVariable (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a))
locVars (SomeVariable Variable a
v) VarSet (KindOf ('HMeasure a))
b =
              case (Variable a -> Assocs (Extra (abt '[])) -> Maybe (Extra (abt '[]) a)
forall k (a :: k) (ast :: k -> *).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Assocs ast -> Maybe (ast a)
lookupAssoc Variable a
v Assocs (Extra (abt '[]))
extras) of
                Maybe (Extra (abt '[]) a)
Nothing -> VarSet (KindOf ('HMeasure a))
b
                Just (Loc l _) -> Variable a
-> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a))
forall k (a :: k).
Variable a -> VarSet (KindOf a) -> VarSet (KindOf a)
insertVarSet (Location a -> Variable a
forall k (a :: k). Location a -> Variable a
fromLocation Location a
l) VarSet (KindOf ('HMeasure a))
b
          fvs2 :: VarSet (KindOf ('HMeasure a))
fvs2 = (SomeVariable (KindOf ('HMeasure a))
 -> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a)))
-> VarSet (KindOf ('HMeasure a))
-> [SomeVariable (KindOf ('HMeasure a))]
-> VarSet (KindOf ('HMeasure a))
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VarSet (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a))
forall k (kproxy :: KProxy k).
(Show1 Sing, JmEq1 Sing) =>
VarSet kproxy -> VarSet kproxy -> VarSet kproxy
unionVarSet(VarSet (KindOf ('HMeasure a))
 -> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a)))
-> (SomeVariable (KindOf ('HMeasure a))
    -> VarSet (KindOf ('HMeasure a)))
-> SomeVariable (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a))
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
.SomeVariable (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a))
indFVs) VarSet (KindOf ('HMeasure a))
fvs1 (VarSet (KindOf ('HMeasure a))
-> [SomeVariable (KindOf ('HMeasure a))]
forall k (kproxy :: KProxy k).
VarSet kproxy -> [SomeVariable kproxy]
fromVarSet VarSet (KindOf ('HMeasure a))
fvs1)
      VarSet (KindOf ('HMeasure a))
-> Dis abt (VarSet (KindOf ('HMeasure a)))
forall (m :: * -> *) a. Monad m => a -> m a
return (VarSet (KindOf ('HMeasure a))
 -> Dis abt (VarSet (KindOf ('HMeasure a))))
-> VarSet (KindOf ('HMeasure a))
-> Dis abt (VarSet (KindOf ('HMeasure a)))
forall a b. (a -> b) -> a -> b
$ (SomeVariable (KindOf ('HMeasure a))
 -> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a)))
-> VarSet (KindOf ('HMeasure a))
-> [SomeVariable (KindOf ('HMeasure a))]
-> VarSet (KindOf ('HMeasure a))
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr SomeVariable (KindOf ('HMeasure a))
-> VarSet (KindOf ('HMeasure a)) -> VarSet (KindOf ('HMeasure a))
locVars VarSet (KindOf ('HMeasure a))
forall k (kproxy :: KProxy k). VarSet kproxy
emptyVarSet (VarSet (KindOf ('HMeasure a))
-> [SomeVariable (KindOf ('HMeasure a))]
forall k (kproxy :: KProxy k).
VarSet kproxy -> [SomeVariable kproxy]
fromVarSet VarSet (KindOf ('HMeasure a))
fvs2)

    -- | The forward disintegrator's function for evaluating case
    -- expressions. First we try calling 'defaultCaseEvaluator' which
    -- will evaluate the scrutinee and select the matching branch (if
    -- any). But that doesn't work out in general, since the scrutinee
    -- may contain heap-bound variables. So our fallback definition
    -- will push a 'SGuard' onto the heap and then continue evaluating
    -- each branch (thereby duplicating the continuation, calling it
    -- once on each branch).
    evaluateCase :: TermEvaluator abt (Dis abt) -> CaseEvaluator abt (Dis abt)
evaluateCase TermEvaluator abt (Dis abt)
evaluate_ = abt '[] a -> [Branch a abt b] -> Dis abt (Whnf abt b)
CaseEvaluator abt (Dis abt)
evaluateCase_
        where
          evaluateCase_ :: CaseEvaluator abt (Dis abt)
          evaluateCase_ :: abt '[] a -> [Branch a abt b] -> Dis abt (Whnf abt b)
evaluateCase_ abt '[] a
e [Branch a abt b]
bs =
              TermEvaluator abt (Dis abt)
-> abt '[] a -> [Branch a abt b] -> Dis abt (Whnf abt b)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
(ABT Term abt, EvaluationMonad abt m p) =>
TermEvaluator abt m -> CaseEvaluator abt m
defaultCaseEvaluator TermEvaluator abt (Dis abt)
evaluate_ abt '[] a
e [Branch a abt b]
bs
              Dis abt (Whnf abt b)
-> Dis abt (Whnf abt b) -> Dis abt (Whnf abt b)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> abt '[] a -> [Branch a abt b] -> Dis abt (Whnf abt b)
CaseEvaluator abt (Dis abt)
evaluateBranches abt '[] a
e [Branch a abt b]
bs

          evaluateBranches :: CaseEvaluator abt (Dis abt)
          evaluateBranches :: abt '[] a -> [Branch a abt b] -> Dis abt (Whnf abt b)
evaluateBranches abt '[] a
e = [Dis abt (Whnf abt b)] -> Dis abt (Whnf abt b)
forall (abt :: [Hakaru] -> Hakaru -> *) a.
ABT Term abt =>
[Dis abt a] -> Dis abt a
choose ([Dis abt (Whnf abt b)] -> Dis abt (Whnf abt b))
-> ([Branch a abt b] -> [Dis abt (Whnf abt b)])
-> [Branch a abt b]
-> Dis abt (Whnf abt b)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Branch a abt b -> Dis abt (Whnf abt b))
-> [Branch a abt b] -> [Dis abt (Whnf abt b)]
forall a b. (a -> b) -> [a] -> [b]
map Branch a abt b -> Dis abt (Whnf abt b)
evaluateBranch
              where
                evaluateBranch :: Branch a abt b -> Dis abt (Whnf abt b)
evaluateBranch (Branch Pattern xs a
pat abt xs b
body) =
                    let (List1 Variable xs
vars,abt '[] b
body') = abt xs b -> (List1 Variable xs, abt '[] b)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs b
body
                    in Dis abt [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
m [Index (abt '[])]
getIndices Dis abt [Index (abt '[])]
-> ([Index (abt '[])] -> Dis abt (Whnf abt b))
-> Dis abt (Whnf abt b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \[Index (abt '[])]
i ->
                        Statement abt Variable 'Impure -> abt '[] b -> Dis abt (abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (xs :: [Hakaru]) (a :: Hakaru).
(ABT Term abt, EvaluationMonad abt m p) =>
Statement abt Variable p -> abt xs a -> m (abt xs a)
push (List1 Variable xs
-> Pattern xs a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt Variable 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (xs :: [Hakaru]) (a :: Hakaru).
List1 v xs
-> Pattern xs a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt v 'Impure
SGuard List1 Variable xs
vars Pattern xs a
pat (abt '[] a -> Lazy abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk abt '[] a
e) [Index (abt '[])]
i) abt '[] b
body'
                           Dis abt (abt '[] b)
-> (abt '[] b -> Dis abt (Whnf abt b)) -> Dis abt (Whnf abt b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= abt '[] b -> Dis abt (Whnf abt b)
TermEvaluator abt (Dis abt)
evaluate_

    evaluateVar :: MeasureEvaluator abt (Dis abt)
-> TermEvaluator abt (Dis abt) -> VariableEvaluator abt (Dis abt)
evaluateVar MeasureEvaluator abt (Dis abt)
perform TermEvaluator abt (Dis abt)
evaluate_ Variable a
x =
      do Assocs (Extra (abt '[]))
extras <- Dis abt (Assocs (Extra (abt '[])))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Dis abt (Assocs (Extra (abt '[])))
getExtras
         -- If we get 'Nothing', then it turns out @x@ is a free variable
         Dis abt (Whnf abt a)
-> (Extra (abt '[]) a -> Dis abt (Whnf abt a))
-> Maybe (Extra (abt '[]) a)
-> Dis abt (Whnf abt a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Whnf abt a -> Dis abt (Whnf abt a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Whnf abt a -> Dis abt (Whnf abt a))
-> Whnf abt a -> Dis abt (Whnf abt a)
forall a b. (a -> b) -> a -> b
$ abt '[] a -> Whnf abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Whnf abt a
Neutral (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
x)) Extra (abt '[]) a -> Dis abt (Whnf abt a)
lookForLoc (Variable a -> Assocs (Extra (abt '[])) -> Maybe (Extra (abt '[]) a)
forall k (a :: k) (ast :: k -> *).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Assocs ast -> Maybe (ast a)
lookupAssoc Variable a
x Assocs (Extra (abt '[]))
extras)
      where lookForLoc :: Extra (abt '[]) a -> Dis abt (Whnf abt a)
lookForLoc (Loc      Location a
l [abt '[] 'HNat]
jxs) =
              (Dis abt (Whnf abt a)
-> (Whnf abt a -> Dis abt (Whnf abt a))
-> Maybe (Whnf abt a)
-> Dis abt (Whnf abt a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Location a -> Dis abt (Whnf abt a)
forall (a :: Hakaru) b. Location a -> b
freeLocError Location a
l) Whnf abt a -> Dis abt (Whnf abt a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Whnf abt a) -> Dis abt (Whnf abt a))
-> Dis abt (Maybe (Whnf abt a)) -> Dis abt (Whnf abt a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (Dis abt (Maybe (Whnf abt a)) -> Dis abt (Whnf abt a))
-> ((Statement abt Location 'Impure
     -> Maybe (Dis abt (Whnf abt a)))
    -> Dis abt (Maybe (Whnf abt a)))
-> (Statement abt Location 'Impure -> Maybe (Dis abt (Whnf abt a)))
-> Dis abt (Whnf abt a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Location a
-> (Statement abt Location 'Impure -> Maybe (Dis abt (Whnf abt a)))
-> Dis abt (Maybe (Whnf abt a))
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru) r.
EvaluationMonad abt m p =>
Location a
-> (Statement abt Location p -> Maybe (m r)) -> m (Maybe r)
select Location a
l ((Statement abt Location 'Impure -> Maybe (Dis abt (Whnf abt a)))
 -> Dis abt (Whnf abt a))
-> (Statement abt Location 'Impure -> Maybe (Dis abt (Whnf abt a)))
-> Dis abt (Whnf abt a)
forall a b. (a -> b) -> a -> b
$ \Statement abt Location 'Impure
s ->
                case Statement abt Location 'Impure
s of
                SBind Location a
l' Lazy abt ('HMeasure a)
e [Index (abt '[])]
ixs -> do
                  TypeEq a a
Refl <- Location a -> Location a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Location a -> Location b -> Maybe (TypeEq a b)
locEq Location a
l Location a
l'
                  Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a))
forall a. a -> Maybe a
Just (Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a)))
-> Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a))
forall a b. (a -> b) -> a -> b
$ do
                    Whnf abt a
w <- [Index (abt '[])] -> Dis abt (Whnf abt a) -> Dis abt (Whnf abt a)
forall (abt :: [Hakaru] -> Hakaru -> *) a.
[Index (abt '[])] -> Dis abt a -> Dis abt a
withIndices [Index (abt '[])]
ixs (Dis abt (Whnf abt a) -> Dis abt (Whnf abt a))
-> Dis abt (Whnf abt a) -> Dis abt (Whnf abt a)
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Dis abt (Whnf abt a)
MeasureEvaluator abt (Dis abt)
perform (Lazy abt ('HMeasure a)
-> (Whnf abt ('HMeasure a) -> abt '[] ('HMeasure a))
-> (abt '[] ('HMeasure a) -> abt '[] ('HMeasure a))
-> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru) r.
Lazy abt a -> (Whnf abt a -> r) -> (abt '[] a -> r) -> r
caseLazy Lazy abt ('HMeasure a)
e Whnf abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Whnf abt a -> abt '[] a
fromWhnf abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id)
                    Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity)
       (v :: Hakaru -> *) (a :: Hakaru).
v a -> Lazy abt a -> [Index (abt '[])] -> Statement abt v p
SLet Location a
l (Whnf abt a -> Lazy abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
Whnf abt a -> Lazy abt a
Whnf_ Whnf abt a
w) [Index (abt '[])]
ixs)
#ifdef __TRACE_DISINTEGRATE__
                    trace ("-- updated "
                           ++ show (ppStatement 11 s)
                           ++ " to "
                           ++ show (ppStatement 11 (SLet l (Whnf_ w) ixs))
                          ) $ return ()
#endif
                    Assocs (abt '[]) -> abt '[] a -> Dis abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (m :: * -> *) (p :: Purity).
EvaluationMonad abt m p =>
Assocs (abt '[]) -> abt xs a -> m (abt xs a)
extSubsts ([Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
[Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
zipInds [Index (abt '[])]
ixs [abt '[] 'HNat]
jxs) (Whnf abt a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Whnf abt a -> abt '[] a
fromWhnf Whnf abt a
w) Dis abt (abt '[] a)
-> (abt '[] a -> Dis abt (Whnf abt a)) -> Dis abt (Whnf abt a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= abt '[] a -> Dis abt (Whnf abt a)
TermEvaluator abt (Dis abt)
evaluate_
                SLet  Location a
l' Lazy abt a
e [Index (abt '[])]
ixs -> do
                  TypeEq a a
Refl <- Location a -> Location a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Location a -> Location b -> Maybe (TypeEq a b)
locEq Location a
l Location a
l'
                  Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a))
forall a. a -> Maybe a
Just (Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a)))
-> Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a))
forall a b. (a -> b) -> a -> b
$ do
                    Whnf abt a
w <- [Index (abt '[])] -> Dis abt (Whnf abt a) -> Dis abt (Whnf abt a)
forall (abt :: [Hakaru] -> Hakaru -> *) a.
[Index (abt '[])] -> Dis abt a -> Dis abt a
withIndices [Index (abt '[])]
ixs (Dis abt (Whnf abt a) -> Dis abt (Whnf abt a))
-> Dis abt (Whnf abt a) -> Dis abt (Whnf abt a)
forall a b. (a -> b) -> a -> b
$ Lazy abt a
-> (Whnf abt a -> Dis abt (Whnf abt a))
-> (abt '[] a -> Dis abt (Whnf abt a))
-> Dis abt (Whnf abt a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru) r.
Lazy abt a -> (Whnf abt a -> r) -> (abt '[] a -> r) -> r
caseLazy Lazy abt a
e Whnf abt a -> Dis abt (Whnf abt a)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] a -> Dis abt (Whnf abt a)
TermEvaluator abt (Dis abt)
evaluate_
                    Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity)
       (v :: Hakaru -> *) (a :: Hakaru).
v a -> Lazy abt a -> [Index (abt '[])] -> Statement abt v p
SLet Location a
l (Whnf abt a -> Lazy abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
Whnf abt a -> Lazy abt a
Whnf_ Whnf abt a
w) [Index (abt '[])]
ixs)
                    Assocs (abt '[]) -> abt '[] a -> Dis abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (m :: * -> *) (p :: Purity).
EvaluationMonad abt m p =>
Assocs (abt '[]) -> abt xs a -> m (abt xs a)
extSubsts ([Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
[Index (abt '[])] -> [abt '[] 'HNat] -> Assocs (abt '[])
zipInds [Index (abt '[])]
ixs [abt '[] 'HNat]
jxs) (Whnf abt a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Whnf abt a -> abt '[] a
fromWhnf Whnf abt a
w) Dis abt (abt '[] a)
-> (abt '[] a -> Dis abt (Whnf abt a)) -> Dis abt (Whnf abt a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= abt '[] a -> Dis abt (Whnf abt a)
TermEvaluator abt (Dis abt)
evaluate_
                -- This does not bind any variables, so it definitely can't match.
                SWeight   Lazy abt 'HProb
_ [Index (abt '[])]
_ -> Maybe (Dis abt (Whnf abt a))
forall a. Maybe a
Nothing
                -- This does bind variables,
                -- but there's no expression we can return for it
                -- because the variables are untouchable\/abstract.
                SGuard List1 Location xs
ls Pattern xs a
pat Lazy abt a
scrutinee [Index (abt '[])]
i -> Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a))
forall a. a -> Maybe a
Just (Dis abt (Whnf abt a) -> Maybe (Dis abt (Whnf abt a)))
-> (abt '[] a -> Dis abt (Whnf abt a))
-> abt '[] a
-> Maybe (Dis abt (Whnf abt a))
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Whnf abt a -> Dis abt (Whnf abt a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Whnf abt a -> Dis abt (Whnf abt a))
-> (abt '[] a -> Whnf abt a) -> abt '[] a -> Dis abt (Whnf abt a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. abt '[] a -> Whnf abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Whnf abt a
Neutral (abt '[] a -> Maybe (Dis abt (Whnf abt a)))
-> abt '[] a -> Maybe (Dis abt (Whnf abt a))
forall a b. (a -> b) -> a -> b
$ Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
x  
                 
        
withIndices :: [Index (abt '[])] -> Dis abt a -> Dis abt a
withIndices :: [Index (abt '[])] -> Dis abt a -> Dis abt a
withIndices [Index (abt '[])]
inds (Dis forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m) = (forall (a :: Hakaru).
 [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
 -> Dis abt a)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ a -> Ans abt a
c -> [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (a :: Hakaru).
[Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
m [Index (abt '[])]
inds a -> Ans abt a
c

-- | Not exported because we only need it for defining 'select' on 'Dis'.
unsafePop :: Dis abt (Maybe (Statement abt Location 'Impure))
unsafePop :: Dis abt (Maybe (Statement abt Location 'Impure))
unsafePop =
    (forall (a :: Hakaru).
 [Index (abt '[])]
 -> (Maybe (Statement abt Location 'Impure) -> Ans abt a)
 -> Ans abt a)
-> Dis abt (Maybe (Statement abt Location 'Impure))
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])]
  -> (Maybe (Statement abt Location 'Impure) -> Ans abt a)
  -> Ans abt a)
 -> Dis abt (Maybe (Statement abt Location 'Impure)))
-> (forall (a :: Hakaru).
    [Index (abt '[])]
    -> (Maybe (Statement abt Location 'Impure) -> Ans abt a)
    -> Ans abt a)
-> Dis abt (Maybe (Statement abt Location 'Impure))
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ Maybe (Statement abt Location 'Impure) -> Ans abt a
c h :: ListContext abt 'Impure
h@(ListContext Nat
i [Statement abt Location 'Impure]
ss) Assocs (Extra (abt '[]))
extras ->
        case [Statement abt Location 'Impure]
ss of
        []    -> Maybe (Statement abt Location 'Impure) -> Ans abt a
c Maybe (Statement abt Location 'Impure)
forall a. Maybe a
Nothing  ListContext abt 'Impure
h Assocs (Extra (abt '[]))
extras
        Statement abt Location 'Impure
s:[Statement abt Location 'Impure]
ss' -> Maybe (Statement abt Location 'Impure) -> Ans abt a
c (Statement abt Location 'Impure
-> Maybe (Statement abt Location 'Impure)
forall a. a -> Maybe a
Just Statement abt Location 'Impure
s) (Nat -> [Statement abt Location 'Impure] -> ListContext abt 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
Nat -> [Statement abt Location p] -> ListContext abt p
ListContext Nat
i [Statement abt Location 'Impure]
ss') Assocs (Extra (abt '[]))
extras

pushPlate
    :: (ABT Term abt)
    => abt '[] 'HNat
    -> abt '[ 'HNat ] ('HMeasure a)
    -> Dis abt (abt '[] ('HArray a))
pushPlate :: abt '[] 'HNat
-> abt '[ 'HNat] ('HMeasure a) -> Dis abt (abt '[] ('HArray a))
pushPlate abt '[] 'HNat
n abt '[ 'HNat] ('HMeasure a)
e =
  abt '[ 'HNat] ('HMeasure a)
-> (Variable 'HNat
    -> abt '[] ('HMeasure a) -> Dis abt (abt '[] ('HArray a)))
-> Dis abt (abt '[] ('HArray a))
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[ 'HNat] ('HMeasure a)
e ((Variable 'HNat
  -> abt '[] ('HMeasure a) -> Dis abt (abt '[] ('HArray a)))
 -> Dis abt (abt '[] ('HArray a)))
-> (Variable 'HNat
    -> abt '[] ('HMeasure a) -> Dis abt (abt '[] ('HArray a)))
-> Dis abt (abt '[] ('HArray a))
forall a b. (a -> b) -> a -> b
$ \Variable 'HNat
x abt '[] ('HMeasure a)
body -> do
    [Index (abt '[])]
inds <- Dis abt [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
m [Index (abt '[])]
getIndices
    Index (abt '[])
i    <- abt '[] 'HNat -> Dis abt (Index (abt '[]))
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
abt '[] 'HNat -> m (Index (abt '[]))
freshInd abt '[] 'HNat
n
    Location a
p    <- Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location (Variable a -> Location a)
-> Dis abt (Variable a) -> Dis abt (Location a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Sing a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty (Sing ('HMeasure a) -> Sing a
forall (a :: Hakaru). Sing ('HMeasure a) -> Sing a
sUnMeasure (Sing ('HMeasure a) -> Sing a) -> Sing ('HMeasure a) -> Sing a
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] ('HMeasure a)
body)
    let inds' :: [Index (abt '[])]
inds' = Index (abt '[]) -> [Index (abt '[])] -> [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> [Index (abt '[])] -> [Index (abt '[])]
extendIndices Index (abt '[])
i [Index (abt '[])]
inds
    Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt v 'Impure
SBind Location a
p (abt '[] ('HMeasure a) -> Lazy abt ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk (abt '[] ('HMeasure a) -> Lazy abt ('HMeasure a))
-> abt '[] ('HMeasure a) -> Lazy abt ('HMeasure a)
forall a b. (a -> b) -> a -> b
$ Variable 'HNat
-> Variable 'HNat -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
(JmEq1 Sing, Show1 Sing, Functor21 syn, ABT syn abt) =>
Variable a -> Variable a -> abt xs b -> abt xs b
rename Variable 'HNat
x (Index (abt '[]) -> Variable 'HNat
forall (ast :: Hakaru -> *). Index ast -> Variable 'HNat
indVar Index (abt '[])
i) abt '[] ('HMeasure a)
body) [Index (abt '[])]
inds')
    Variable a
v <- Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text -> Location a -> [abt '[] 'HNat] -> Dis abt (Variable a)
mkLoc Text
Text.empty Location a
p ([abt '[] 'HNat] -> Dis abt (Variable a))
-> [abt '[] 'HNat] -> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ (Index (abt '[]) -> abt '[] 'HNat)
-> [Index (abt '[])] -> [abt '[] 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> abt '[] 'HNat
fromIndex [Index (abt '[])]
inds'
    abt '[] ('HArray a) -> Dis abt (abt '[] ('HArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return (abt '[] ('HArray a) -> Dis abt (abt '[] ('HArray a)))
-> abt '[] ('HArray a) -> Dis abt (abt '[] ('HArray a))
forall a b. (a -> b) -> a -> b
$ abt '[] 'HNat -> Variable 'HNat -> abt '[] a -> abt '[] ('HArray a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] 'HNat -> Variable 'HNat -> abt '[] a -> abt '[] ('HArray a)
P.arrayWithVar abt '[] 'HNat
n (Index (abt '[]) -> Variable 'HNat
forall (ast :: Hakaru -> *). Index ast -> Variable 'HNat
indVar Index (abt '[])
i) (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
v)

-- | Calls unsafePush 
pushWeight :: (ABT Term abt) => abt '[] 'HProb -> Dis abt ()
pushWeight :: abt '[] 'HProb -> Dis abt ()
pushWeight abt '[] 'HProb
w = do
  [Index (abt '[])]
inds <- Dis abt [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
m [Index (abt '[])]
getIndices
  Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Statement abt Location 'Impure -> Dis abt ())
-> Statement abt Location 'Impure -> Dis abt ()
forall a b. (a -> b) -> a -> b
$ Lazy abt 'HProb
-> [Index (abt '[])] -> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
Lazy abt 'HProb -> [Index (abt '[])] -> Statement abt v 'Impure
SWeight (abt '[] 'HProb -> Lazy abt 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk abt '[] 'HProb
w) [Index (abt '[])]
inds

-- | Calls unsafePush 
pushGuard :: (ABT Term abt) => abt '[] HBool -> Dis abt ()
pushGuard :: abt '[] HBool -> Dis abt ()
pushGuard abt '[] HBool
b = do
  [Index (abt '[])]
inds <- Dis abt [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
m [Index (abt '[])]
getIndices
  Statement abt Location 'Impure -> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Statement abt Location 'Impure -> Dis abt ())
-> Statement abt Location 'Impure -> Dis abt ()
forall a b. (a -> b) -> a -> b
$ List1 Location '[]
-> Pattern '[] HBool
-> Lazy abt HBool
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (xs :: [Hakaru]) (a :: Hakaru).
List1 v xs
-> Pattern xs a
-> Lazy abt a
-> [Index (abt '[])]
-> Statement abt v 'Impure
SGuard List1 Location '[]
forall k (a :: k -> *). List1 a '[]
Nil1 Pattern '[] HBool
pTrue (abt '[] HBool -> Lazy abt HBool
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] a -> Lazy abt a
Thunk abt '[] HBool
b) [Index (abt '[])]
inds           

----------------------------------------------------------------
----------------------------------------------------------------

-- | It is impossible to satisfy the constraints, or at least we
-- give up on trying to do so. This function is identical to 'empty'
-- and 'mzero' for 'Dis'; we just give it its own name since this is
-- the name used in our papers.
--
-- TODO: add some sort of trace information so we can get a better
-- idea what caused a disintegration to fail.
bot :: (ABT Term abt) => Dis abt a
bot :: Dis abt a
bot = (forall (a :: Hakaru).
 [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
 -> Dis abt a)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ a -> Ans abt a
_ ListContext abt 'Impure
_ Assocs (Extra (abt '[]))
_ -> []


-- | The empty measure is a solution to the constraints.
-- reject :: (ABT Term abt) => Dis abt a
-- reject = Dis $ \_ _ -> [syn (Superpose_ [])]


-- Something essentially like this function was called @insert_@
-- in the finally-tagless code.
--
-- | Emit some code that binds a variable, and return the variable
-- thus bound. The function says what to wrap the result of the
-- continuation with; i.e., what we're actually emitting.
emit
    :: (ABT Term abt)
    => Text
    -> Sing a
    -> (forall r. abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
    -> Dis abt (Variable a)
emit :: Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
emit Text
hint Sing a
typ forall (r :: Hakaru).
abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r)
f = do
    Variable a
x <- Text -> Sing a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
hint Sing a
typ
    (forall (a :: Hakaru).
 [Index (abt '[])] -> (Variable a -> Ans abt a) -> Ans abt a)
-> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (Variable a -> Ans abt a) -> Ans abt a)
 -> Dis abt (Variable a))
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (Variable a -> Ans abt a) -> Ans abt a)
-> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ Variable a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l -> (abt '[a] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (r :: Hakaru).
abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r)
f (abt '[a] ('HMeasure a) -> abt '[] ('HMeasure a))
-> (abt '[] ('HMeasure a) -> abt '[a] ('HMeasure a))
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Variable a -> abt '[] ('HMeasure a) -> abt '[a] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x) (abt '[] ('HMeasure a) -> abt '[] ('HMeasure a))
-> [abt '[] ('HMeasure a)] -> [abt '[] ('HMeasure a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Variable a -> Ans abt a
c Variable a
x ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l

-- This function was called @lift@ in the finally-tagless code.
-- | Emit an 'MBind' (i.e., \"@m >>= \x ->@\") and return the
-- variable thus bound (i.e., @x@).
emitMBind :: (ABT Term abt) => abt '[] ('HMeasure a) -> Dis abt (Variable a)
emitMBind :: abt '[] ('HMeasure a) -> Dis abt (Variable a)
emitMBind abt '[] ('HMeasure a)
m =
    Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
emit Text
Text.empty (Sing ('HMeasure a) -> Sing a
forall (a :: Hakaru). Sing ('HMeasure a) -> Sing a
sUnMeasure (Sing ('HMeasure a) -> Sing a) -> Sing ('HMeasure a) -> Sing a
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] ('HMeasure a)
m) ((forall (r :: Hakaru).
  abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
 -> Dis abt (Variable a))
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ \abt '[a] ('HMeasure r)
e ->
        Term abt ('HMeasure r) -> abt '[] ('HMeasure r)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC ('HMeasure a), '( '[a], 'HMeasure r)] ('HMeasure r)
forall (a :: Hakaru) (b :: Hakaru).
SCon '[LC ('HMeasure a), '( '[a], 'HMeasure b)] ('HMeasure b)
MBind SCon '[LC ('HMeasure a), '( '[a], 'HMeasure r)] ('HMeasure r)
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HMeasure r)]
-> Term abt ('HMeasure r)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] ('HMeasure a)
m abt '[] ('HMeasure a)
-> SArgs abt '[ '( '[a], 'HMeasure r)]
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HMeasure r)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[a] ('HMeasure r)
e abt '[a] ('HMeasure r)
-> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HMeasure r)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)

emitMBind2 :: (ABT Term abt) => abt '[] ('HMeasure a) -> Dis abt (abt '[] a)
emitMBind2 :: abt '[] ('HMeasure a) -> Dis abt (abt '[] a)
emitMBind2 abt '[] ('HMeasure a)
m = do
  [Index (abt '[])]
inds <- Dis abt [Index (abt '[])]
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
m [Index (abt '[])]
getIndices
  let b :: Lazy abt ('HMeasure a)
b = Whnf abt ('HMeasure a) -> Lazy abt ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
Whnf abt a -> Lazy abt a
Whnf_ (Whnf abt ('HMeasure a) -> Lazy abt ('HMeasure a))
-> Whnf abt ('HMeasure a) -> Lazy abt ('HMeasure a)
forall a b. (a -> b) -> a -> b
$ Whnf abt ('HMeasure a)
-> Maybe (Whnf abt ('HMeasure a)) -> Whnf abt ('HMeasure a)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Whnf abt ('HMeasure a)
forall a. HasCallStack => [Char] -> a
error [Char]
"emitMBind2: non-hnf term") (abt '[] ('HMeasure a) -> Maybe (Whnf abt ('HMeasure a))
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Maybe (Whnf abt a)
toWhnf abt '[] ('HMeasure a)
m)
      typ :: Sing a
typ = Sing ('HMeasure a) -> Sing a
forall (a :: Hakaru). Sing ('HMeasure a) -> Sing a
sUnMeasure (Sing ('HMeasure a) -> Sing a) -> Sing ('HMeasure a) -> Sing a
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] ('HMeasure a)
m
  Location a
l <- Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location (Variable a -> Location a)
-> Dis abt (Variable a) -> Dis abt (Location a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Sing a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing a
typ
  (SBind Location a
l' Lazy abt ('HMeasure a)
b' [Index (abt '[])]
_, Name Any
name) <- Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name Any)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Statement abt Location 'Impure
-> Dis abt (Statement abt Location 'Impure, Name a)
reifyStatement (Location a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt Location 'Impure
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> Lazy abt ('HMeasure a)
-> [Index (abt '[])]
-> Statement abt v 'Impure
SBind Location a
l Lazy abt ('HMeasure a)
b [Index (abt '[])]
inds)
  let (abt '[] a
idx, abt '[] ('HMeasure a)
p) = (Name Any -> Sing a -> [abt '[] 'HNat] -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru)
       (a :: Hakaru).
ABT Term abt =>
Name b -> Sing a -> [abt '[] 'HNat] -> abt '[] a
fromName Name Any
name Sing a
typ ((Index (abt '[]) -> abt '[] 'HNat)
-> [Index (abt '[])] -> [abt '[] 'HNat]
forall a b. (a -> b) -> [a] -> [b]
map Index (abt '[]) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Index (abt '[]) -> abt '[] 'HNat
fromIndex [Index (abt '[])]
inds), Lazy abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt ('HMeasure a)
b')
  (forall (a :: Hakaru).
 [Index (abt '[])] -> (abt '[] a -> Ans abt a) -> Ans abt a)
-> Dis abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (abt '[] a -> Ans abt a) -> Ans abt a)
 -> Dis abt (abt '[] a))
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (abt '[] a -> Ans abt a) -> Ans abt a)
-> Dis abt (abt '[] a)
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ abt '[] a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
ex ->
      abt '[] a -> Ans abt a
c abt '[] a
idx ListContext abt 'Impure
h Assocs (Extra (abt '[]))
ex [abt '[] ('HMeasure a)]
-> (abt '[] ('HMeasure a) -> [abt '[] ('HMeasure a)])
-> [abt '[] ('HMeasure a)]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \abt '[] ('HMeasure a)
e ->
      abt '[] ('HMeasure a) -> [abt '[] ('HMeasure a)]
forall (m :: * -> *) a. Monad m => a -> m a
return (abt '[] ('HMeasure a) -> [abt '[] ('HMeasure a)])
-> abt '[] ('HMeasure a) -> [abt '[] ('HMeasure a)]
forall a b. (a -> b) -> a -> b
$ Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC ('HMeasure a), '( '[a], 'HMeasure a)] ('HMeasure a)
forall (a :: Hakaru) (b :: Hakaru).
SCon '[LC ('HMeasure a), '( '[a], 'HMeasure b)] ('HMeasure b)
MBind SCon '[LC ('HMeasure a), '( '[a], 'HMeasure a)] ('HMeasure a)
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HMeasure a)]
-> Term abt ('HMeasure a)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] ('HMeasure a)
p abt '[] ('HMeasure a)
-> SArgs abt '[ '( '[a], 'HMeasure a)]
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] ('HMeasure a) -> abt '[a] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind (Location a -> Variable a
forall k (a :: k). Location a -> Variable a
fromLocation Location a
l') abt '[] ('HMeasure a)
e abt '[a] ('HMeasure a)
-> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HMeasure a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)

-- | A smart constructor for emitting let-bindings. If the input
-- is already a variable then we just return it; otherwise we emit
-- the let-binding. N.B., this function provides the invariant that
-- the result is in fact a variable; whereas 'emitLet'' does not.
emitLet :: (ABT Term abt) => abt '[] a -> Dis abt (Variable a)
emitLet :: abt '[] a -> Dis abt (Variable a)
emitLet abt '[] a
e =
    abt '[] a
-> (Variable a -> Dis abt (Variable a))
-> (Term abt a -> Dis abt (Variable a))
-> Dis abt (Variable a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
e Variable a -> Dis abt (Variable a)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Term abt a -> Dis abt (Variable a)) -> Dis abt (Variable a))
-> (Term abt a -> Dis abt (Variable a)) -> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ \Term abt a
_ ->
        Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
emit Text
Text.empty (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
e) ((forall (r :: Hakaru).
  abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
 -> Dis abt (Variable a))
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ \abt '[a] ('HMeasure r)
m ->
            Term abt ('HMeasure r) -> abt '[] ('HMeasure r)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC a, '( '[a], 'HMeasure r)] ('HMeasure r)
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], 'HMeasure r)] ('HMeasure r)
-> SArgs abt '[LC a, '( '[a], 'HMeasure r)]
-> Term abt ('HMeasure r)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
e abt '[] a
-> SArgs abt '[ '( '[a], 'HMeasure r)]
-> SArgs abt '[LC a, '( '[a], 'HMeasure r)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[a] ('HMeasure r)
m abt '[a] ('HMeasure r)
-> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HMeasure r)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)

-- | A smart constructor for emitting let-bindings. If the input
-- is already a variable or a literal constant, then we just return
-- it; otherwise we emit the let-binding. N.B., this function
-- provides weaker guarantees on the type of the result; if you
-- require the result to always be a variable, then see 'emitLet'
-- instead.
emitLet' :: (ABT Term abt) => abt '[] a -> Dis abt (abt '[] a)
emitLet' :: abt '[] a -> Dis abt (abt '[] a)
emitLet' abt '[] a
e =
    abt '[] a
-> (Variable a -> Dis abt (abt '[] a))
-> (Term abt a -> Dis abt (abt '[] a))
-> Dis abt (abt '[] a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
e (Dis abt (abt '[] a) -> Variable a -> Dis abt (abt '[] a)
forall a b. a -> b -> a
const (Dis abt (abt '[] a) -> Variable a -> Dis abt (abt '[] a))
-> Dis abt (abt '[] a) -> Variable a -> Dis abt (abt '[] a)
forall a b. (a -> b) -> a -> b
$ abt '[] a -> Dis abt (abt '[] a)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] a
e) ((Term abt a -> Dis abt (abt '[] a)) -> Dis abt (abt '[] a))
-> (Term abt a -> Dis abt (abt '[] a)) -> Dis abt (abt '[] a)
forall a b. (a -> b) -> a -> b
$ \Term abt a
t ->
        case Term abt a
t of
        Literal_ Literal a
_ -> abt '[] a -> Dis abt (abt '[] a)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] a
e
        Term abt a
_          -> do
            Variable a
x <- Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Text
-> Sing a
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
emit Text
Text.empty (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
e) ((forall (r :: Hakaru).
  abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
 -> Dis abt (Variable a))
-> (forall (r :: Hakaru).
    abt '[a] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ \abt '[a] ('HMeasure r)
m ->
                Term abt ('HMeasure r) -> abt '[] ('HMeasure r)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC a, '( '[a], 'HMeasure r)] ('HMeasure r)
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], 'HMeasure r)] ('HMeasure r)
-> SArgs abt '[LC a, '( '[a], 'HMeasure r)]
-> Term abt ('HMeasure r)
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
e abt '[] a
-> SArgs abt '[ '( '[a], 'HMeasure r)]
-> SArgs abt '[LC a, '( '[a], 'HMeasure r)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[a] ('HMeasure r)
m abt '[a] ('HMeasure r)
-> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HMeasure r)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)
            abt '[] a -> Dis abt (abt '[] a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
x)

-- | A smart constructor for emitting \"unpair\". If the input
-- argument is actually a constructor then we project out the two
-- components; otherwise we emit the case-binding and return the
-- two variables.
emitUnpair
    :: (ABT Term abt)
    => Whnf abt (HPair a b)
    -> Dis abt (abt '[] a, abt '[] b)
emitUnpair :: Whnf abt (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
emitUnpair (Head_   Head abt (HPair a b)
w) = (abt '[] a, abt '[] b) -> Dis abt (abt '[] a, abt '[] b)
forall (m :: * -> *) a. Monad m => a -> m a
return ((abt '[] a, abt '[] b) -> Dis abt (abt '[] a, abt '[] b))
-> (abt '[] a, abt '[] b) -> Dis abt (abt '[] a, abt '[] b)
forall a b. (a -> b) -> a -> b
$ Head abt (HPair a b) -> (abt '[] a, abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
Head abt (HPair a b) -> (abt '[] a, abt '[] b)
reifyPair Head abt (HPair a b)
w
emitUnpair (Neutral abt '[] (HPair a b)
e) = do
    let (Sing a
a,Sing b
b) = Sing (HPair a b) -> (Sing a, Sing b)
forall (a :: Hakaru) (b :: Hakaru).
Sing (HPair a b) -> (Sing a, Sing b)
sUnPair (abt '[] (HPair a b) -> Sing (HPair a b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] (HPair a b)
e)
    Variable a
x <- Text -> Sing a -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing a
a
    Variable b
y <- Text -> Sing b -> Dis abt (Variable b)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing b
b
    Variable a
-> Variable b
-> abt '[] (HPair a b)
-> Dis abt (abt '[] a, abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
Variable a
-> Variable b
-> abt '[] (HPair a b)
-> Dis abt (abt '[] a, abt '[] b)
emitUnpair_ Variable a
x Variable b
y abt '[] (HPair a b)
e

emitUnpair_
    :: forall abt a b
    .  (ABT Term abt)
    => Variable a
    -> Variable b
    -> abt '[] (HPair a b)
    -> Dis abt (abt '[] a, abt '[] b)
emitUnpair_ :: Variable a
-> Variable b
-> abt '[] (HPair a b)
-> Dis abt (abt '[] a, abt '[] b)
emitUnpair_ Variable a
x Variable b
y = abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
loop
    where
    done :: abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
    done :: abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
done abt '[] (HPair a b)
e =
#ifdef __TRACE_DISINTEGRATE__
        trace "-- emitUnpair: done (term is not Datum_ nor Case_)" $
#endif
        (forall (a :: Hakaru).
 [Index (abt '[])]
 -> ((abt '[] a, abt '[] b) -> Ans abt a) -> Ans abt a)
-> Dis abt (abt '[] a, abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])]
  -> ((abt '[] a, abt '[] b) -> Ans abt a) -> Ans abt a)
 -> Dis abt (abt '[] a, abt '[] b))
-> (forall (a :: Hakaru).
    [Index (abt '[])]
    -> ((abt '[] a, abt '[] b) -> Ans abt a) -> Ans abt a)
-> Dis abt (abt '[] a, abt '[] b)
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ (abt '[] a, abt '[] b) -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l ->
            ( Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn
            (Term abt ('HMeasure a) -> abt '[] ('HMeasure a))
-> (abt '[] ('HMeasure a) -> Term abt ('HMeasure a))
-> abt '[] ('HMeasure a)
-> abt '[] ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. abt '[] (HPair a b)
-> [Branch (HPair a b) abt ('HMeasure a)] -> Term abt ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
abt '[] a -> [Branch a abt b] -> Term abt b
Case_ abt '[] (HPair a b)
e
            ([Branch (HPair a b) abt ('HMeasure a)] -> Term abt ('HMeasure a))
-> (abt '[] ('HMeasure a)
    -> [Branch (HPair a b) abt ('HMeasure a)])
-> abt '[] ('HMeasure a)
-> Term abt ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Branch (HPair a b) abt ('HMeasure a)
-> [Branch (HPair a b) abt ('HMeasure a)]
-> [Branch (HPair a b) abt ('HMeasure a)]
forall a. a -> [a] -> [a]
:[])
            (Branch (HPair a b) abt ('HMeasure a)
 -> [Branch (HPair a b) abt ('HMeasure a)])
-> (abt '[] ('HMeasure a) -> Branch (HPair a b) abt ('HMeasure a))
-> abt '[] ('HMeasure a)
-> [Branch (HPair a b) abt ('HMeasure a)]
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Pattern '[a, b] (HPair a b)
-> abt '[a, b] ('HMeasure a)
-> Branch (HPair a b) abt ('HMeasure a)
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru)
       (xs :: [Hakaru]).
Pattern xs a -> abt xs b -> Branch a abt b
Branch (Pattern '[a] a
-> Pattern '[b] b -> Pattern ('[a] ++ '[b]) (HPair a b)
forall (vars1 :: [Hakaru]) (a :: Hakaru) (vars2 :: [Hakaru])
       (b :: Hakaru).
Pattern vars1 a
-> Pattern vars2 b -> Pattern (vars1 ++ vars2) (HPair a b)
pPair Pattern '[a] a
forall (a :: Hakaru). Pattern '[a] a
PVar Pattern '[b] b
forall (a :: Hakaru). Pattern '[a] a
PVar)
            (abt '[a, b] ('HMeasure a) -> Branch (HPair a b) abt ('HMeasure a))
-> (abt '[] ('HMeasure a) -> abt '[a, b] ('HMeasure a))
-> abt '[] ('HMeasure a)
-> Branch (HPair a b) abt ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Variable a -> abt '[b] ('HMeasure a) -> abt '[a, b] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x
            (abt '[b] ('HMeasure a) -> abt '[a, b] ('HMeasure a))
-> (abt '[] ('HMeasure a) -> abt '[b] ('HMeasure a))
-> abt '[] ('HMeasure a)
-> abt '[a, b] ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Variable b -> abt '[] ('HMeasure a) -> abt '[b] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable b
y
            ) (abt '[] ('HMeasure a) -> abt '[] ('HMeasure a))
-> [abt '[] ('HMeasure a)] -> [abt '[] ('HMeasure a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (abt '[] a, abt '[] b) -> Ans abt a
c (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
x, Variable b -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable b
y) ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l

    loop :: abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
    loop :: abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
loop abt '[] (HPair a b)
e0 =
        abt '[] (HPair a b)
-> (Variable (HPair a b) -> Dis abt (abt '[] a, abt '[] b))
-> (Term abt (HPair a b) -> Dis abt (abt '[] a, abt '[] b))
-> Dis abt (abt '[] a, abt '[] b)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] (HPair a b)
e0 (abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
done (abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b))
-> (Variable (HPair a b) -> abt '[] (HPair a b))
-> Variable (HPair a b)
-> Dis abt (abt '[] a, abt '[] b)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Variable (HPair a b) -> abt '[] (HPair a b)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var) ((Term abt (HPair a b) -> Dis abt (abt '[] a, abt '[] b))
 -> Dis abt (abt '[] a, abt '[] b))
-> (Term abt (HPair a b) -> Dis abt (abt '[] a, abt '[] b))
-> Dis abt (abt '[] a, abt '[] b)
forall a b. (a -> b) -> a -> b
$ \Term abt (HPair a b)
t ->
            case Term abt (HPair a b)
t of
            Datum_ Datum (abt '[]) (HData' t)
d   -> do
#ifdef __TRACE_DISINTEGRATE__
                trace "-- emitUnpair: found Datum_" $ return ()
#endif
                (abt '[] a, abt '[] b) -> Dis abt (abt '[] a, abt '[] b)
forall (m :: * -> *) a. Monad m => a -> m a
return ((abt '[] a, abt '[] b) -> Dis abt (abt '[] a, abt '[] b))
-> (abt '[] a, abt '[] b) -> Dis abt (abt '[] a, abt '[] b)
forall a b. (a -> b) -> a -> b
$ Head abt (HPair a b) -> (abt '[] a, abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
Head abt (HPair a b) -> (abt '[] a, abt '[] b)
reifyPair (Datum (abt '[]) (HData' t) -> Head abt (HData' t)
forall (abt :: [Hakaru] -> Hakaru -> *) (t :: HakaruCon).
Datum (abt '[]) (HData' t) -> Head abt (HData' t)
WDatum Datum (abt '[]) (HData' t)
d)
            Case_ abt '[] a
e [Branch a abt (HPair a b)]
bs -> do
#ifdef __TRACE_DISINTEGRATE__
                trace "-- emitUnpair: going under Case_" $ return ()
#endif
                -- TODO: we want this to duplicate the current
                -- continuation for (the evaluation of @loop@ in)
                -- all branches. So far our traces all end up
                -- returning @bot@ on the first branch, and hence
                -- @bot@ for the whole case-expression, so we can't
                -- quite tell whether it does what is intended.
                --
                -- N.B., the only 'Dis'-effects in 'applyBranch'
                -- are to freshen variables; thus this use of
                -- 'traverse' is perfectly sound.
                (abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b))
-> abt '[] a
-> [Branch a abt (HPair a b)]
-> Dis abt (abt '[] a, abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (b :: Hakaru) r
       (a :: Hakaru).
ABT Term abt =>
(abt '[] b -> Dis abt r)
-> abt '[] a -> [Branch a abt b] -> Dis abt r
emitCaseWith abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
loop abt '[] a
e [Branch a abt (HPair a b)]
bs
            Term abt (HPair a b)
_ -> abt '[] (HPair a b) -> Dis abt (abt '[] a, abt '[] b)
done abt '[] (HPair a b)
e0


-- TODO: emitUneither


-- This function was called @insert_@ in the old finally-tagless code.
-- | Emit some code that doesn't bind any variables. This function
-- provides an optimisation over using 'emit' and then discarding
-- the generated variable.
emit_
    :: (ABT Term abt)
    => (forall r. abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
    -> Dis abt ()
emit_ :: (forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
emit_ forall (r :: Hakaru).
abt '[] ('HMeasure r) -> abt '[] ('HMeasure r)
f = (forall (a :: Hakaru).
 [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
 -> Dis abt ())
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (() -> Ans abt a) -> Ans abt a)
-> Dis abt ()
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
_ () -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (r :: Hakaru).
abt '[] ('HMeasure r) -> abt '[] ('HMeasure r)
f (abt '[] ('HMeasure a) -> abt '[] ('HMeasure a))
-> [abt '[] ('HMeasure a)] -> [abt '[] ('HMeasure a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> () -> Ans abt a
c () ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l


-- | Emit an 'MBind' that discards its result (i.e., \"@m >>@\").
-- We restrict the type of the argument to be 'HUnit' so as to avoid
-- accidentally dropping things.
emitMBind_ :: (ABT Term abt) => abt '[] ('HMeasure HUnit) -> Dis abt ()
emitMBind_ :: abt '[] ('HMeasure HUnit) -> Dis abt ()
emitMBind_ abt '[] ('HMeasure HUnit)
m = (forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
(forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
emit_ (abt '[] ('HMeasure HUnit)
m abt '[] ('HMeasure HUnit)
-> abt '[] ('HMeasure r) -> abt '[] ('HMeasure r)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
(ABT Term abt, SingI a) =>
abt '[] ('HMeasure a)
-> abt '[] ('HMeasure b) -> abt '[] ('HMeasure b)
P.>>)


-- TODO: if the argument is a value, then we can evaluate the 'P.if_'
-- immediately rather than emitting it.
-- | Emit an assertion that the condition is true.
emitGuard :: (ABT Term abt) => abt '[] HBool -> Dis abt ()
emitGuard :: abt '[] HBool -> Dis abt ()
emitGuard abt '[] HBool
b = (forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
(forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
emit_ (abt '[] HBool -> abt '[] ('HMeasure r) -> abt '[] ('HMeasure r)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] HBool -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
P.withGuard abt '[] HBool
b) -- == emit_ $ \m -> P.if_ b m P.reject

-- TODO: if the argument is the literal 1, then we can avoid emitting anything.
emitWeight :: (ABT Term abt) => abt '[] 'HProb -> Dis abt ()
emitWeight :: abt '[] 'HProb -> Dis abt ()
emitWeight abt '[] 'HProb
w = (forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
(forall (r :: Hakaru).
 abt '[] ('HMeasure r) -> abt '[] ('HMeasure r))
-> Dis abt ()
emit_ (abt '[] 'HProb -> abt '[] ('HMeasure r) -> abt '[] ('HMeasure r)
forall (abt :: [Hakaru] -> Hakaru -> *) (w :: Hakaru).
ABT Term abt =>
abt '[] 'HProb -> abt '[] ('HMeasure w) -> abt '[] ('HMeasure w)
P.withWeight abt '[] 'HProb
w)


-- N.B., this use of 'T.traverse' is definitely correct. It's
-- sequentializing @t [abt '[] ('HMeasure a)]@ into @[t (abt '[]
-- ('HMeasure a))]@ by chosing one of the possibilities at each
-- position in @t@. No heap\/context effects can escape to mess
-- things up. In contrast, using 'T.traverse' to sequentialize @t
-- (Dis abt a)@ as @Dis abt (t a)@ is /wrong/! Doing that would give
-- the conjunctive semantics where we have effects from one position
-- in @t@ escape to affect the other positions. This has to do with
-- the general issue in partial evaluation where we need to duplicate
-- downstream work (as we do by passing the same heap to everyone)
-- because there's no general way to combing the resulting heaps
-- for each branch.
--
-- | Run each of the elements of the traversable using the same
-- heap and continuation for each one, then pass the results to a
-- function for emitting code.
emitFork_
    :: (ABT Term abt, T.Traversable t)
    => (forall r. t (abt '[] ('HMeasure r)) -> abt '[] ('HMeasure r))
    -> t (Dis abt a)
    -> Dis abt a
emitFork_ :: (forall (r :: Hakaru).
 t (abt '[] ('HMeasure r)) -> abt '[] ('HMeasure r))
-> t (Dis abt a) -> Dis abt a
emitFork_ forall (r :: Hakaru).
t (abt '[] ('HMeasure r)) -> abt '[] ('HMeasure r)
f t (Dis abt a)
ms = (forall (a :: Hakaru).
 [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
 -> Dis abt a)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a)
-> Dis abt a
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l -> t (abt '[] ('HMeasure a)) -> abt '[] ('HMeasure a)
forall (r :: Hakaru).
t (abt '[] ('HMeasure r)) -> abt '[] ('HMeasure r)
f (t (abt '[] ('HMeasure a)) -> abt '[] ('HMeasure a))
-> [t (abt '[] ('HMeasure a))] -> [abt '[] ('HMeasure a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Dis abt a -> [abt '[] ('HMeasure a)])
-> t (Dis abt a) -> [t (abt '[] ('HMeasure a))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
T.traverse (\Dis abt a
m -> Dis abt a -> [Index (abt '[])] -> (a -> Ans abt a) -> Ans abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
Dis abt x
-> forall (a :: Hakaru).
   [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a
unDis Dis abt a
m [Index (abt '[])]
i a -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l) t (Dis abt a)
ms


-- | Emit a 'Superpose_' of the alternatives, each with unit weight.
emitSuperpose
    :: (ABT Term abt)
    => [abt '[] ('HMeasure a)]
    -> Dis abt (Variable a)
emitSuperpose :: [abt '[] ('HMeasure a)] -> Dis abt (Variable a)
emitSuperpose []  = [Char] -> Dis abt (Variable a)
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: emitSuperpose[]"
emitSuperpose [abt '[] ('HMeasure a)
e] = abt '[] ('HMeasure a) -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Dis abt (Variable a)
emitMBind abt '[] ('HMeasure a)
e
emitSuperpose [abt '[] ('HMeasure a)]
es  =
    abt '[] ('HMeasure a) -> Dis abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Dis abt (Variable a)
emitMBind (abt '[] ('HMeasure a) -> Dis abt (Variable a))
-> (NonEmpty (abt '[] ('HMeasure a)) -> abt '[] ('HMeasure a))
-> NonEmpty (abt '[] ('HMeasure a))
-> Dis abt (Variable a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> abt '[] ('HMeasure a)
P.superpose (NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
 -> abt '[] ('HMeasure a))
-> (NonEmpty (abt '[] ('HMeasure a))
    -> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a)))
-> NonEmpty (abt '[] ('HMeasure a))
-> abt '[] ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (abt '[] ('HMeasure a) -> (abt '[] 'HProb, abt '[] ('HMeasure a)))
-> NonEmpty (abt '[] ('HMeasure a))
-> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NE.map ((,) abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.one) (NonEmpty (abt '[] ('HMeasure a)) -> Dis abt (Variable a))
-> NonEmpty (abt '[] ('HMeasure a)) -> Dis abt (Variable a)
forall a b. (a -> b) -> a -> b
$ [abt '[] ('HMeasure a)] -> NonEmpty (abt '[] ('HMeasure a))
forall a. [a] -> NonEmpty a
NE.fromList [abt '[] ('HMeasure a)]
es


-- | Emit a 'Superpose_' of the alternatives, each with unit weight.
choose :: (ABT Term abt) => [Dis abt a] -> Dis abt a
choose :: [Dis abt a] -> Dis abt a
choose []  = [Char] -> Dis abt a
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: choose[]"
choose [Dis abt a
m] = Dis abt a
m
choose [Dis abt a]
ms  = (forall (r :: Hakaru).
 [abt '[] ('HMeasure r)] -> abt '[] ('HMeasure r))
-> [Dis abt a] -> Dis abt a
forall (abt :: [Hakaru] -> Hakaru -> *) (t :: * -> *) a.
(ABT Term abt, Traversable t) =>
(forall (r :: Hakaru).
 t (abt '[] ('HMeasure r)) -> abt '[] ('HMeasure r))
-> t (Dis abt a) -> Dis abt a
emitFork_ (NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure r))
-> abt '[] ('HMeasure r)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> abt '[] ('HMeasure a)
P.superpose (NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure r))
 -> abt '[] ('HMeasure r))
-> ([abt '[] ('HMeasure r)]
    -> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure r)))
-> [abt '[] ('HMeasure r)]
-> abt '[] ('HMeasure r)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (abt '[] ('HMeasure r) -> (abt '[] 'HProb, abt '[] ('HMeasure r)))
-> NonEmpty (abt '[] ('HMeasure r))
-> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure r))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NE.map ((,) abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.one) (NonEmpty (abt '[] ('HMeasure r))
 -> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure r)))
-> ([abt '[] ('HMeasure r)] -> NonEmpty (abt '[] ('HMeasure r)))
-> [abt '[] ('HMeasure r)]
-> NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure r))
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [abt '[] ('HMeasure r)] -> NonEmpty (abt '[] ('HMeasure r))
forall a. [a] -> NonEmpty a
NE.fromList) [Dis abt a]
ms


-- | Given some function we can call on the bodies of the branches,
-- freshen all the pattern-bound variables and then run the function
-- on all the branches in parallel (i.e., with the same continuation
-- and heap) and then emit a case-analysis expression with the
-- results of the continuations as the bodies of the branches. This
-- function is useful for when we really do want to emit a 'Case_'
-- expression, rather than doing the superpose of guard patterns
-- thing that 'constrainValue' does.
--
-- N.B., this function assumes (and does not verify) that the second
-- argument is emissible. So callers must guarantee this invariant,
-- by calling 'atomize' as necessary.
--
-- TODO: capture the emissibility requirement on the second argument
-- in the types.
emitCaseWith
    :: (ABT Term abt)
    => (abt '[] b -> Dis abt r)
    -> abt '[] a
    -> [Branch a abt b]
    -> Dis abt r
emitCaseWith :: (abt '[] b -> Dis abt r)
-> abt '[] a -> [Branch a abt b] -> Dis abt r
emitCaseWith abt '[] b -> Dis abt r
f abt '[] a
e [Branch a abt b]
bs = do
    [GBranch a (Dis abt r)]
gms <- [Branch a abt b]
-> (Branch a abt b -> Dis abt (GBranch a (Dis abt r)))
-> Dis abt [GBranch a (Dis abt r)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
T.for [Branch a abt b]
bs ((Branch a abt b -> Dis abt (GBranch a (Dis abt r)))
 -> Dis abt [GBranch a (Dis abt r)])
-> (Branch a abt b -> Dis abt (GBranch a (Dis abt r)))
-> Dis abt [GBranch a (Dis abt r)]
forall a b. (a -> b) -> a -> b
$ \(Branch Pattern xs a
pat abt xs b
body) ->
        let (List1 Variable xs
vars, abt '[] b
body') = abt xs b -> (List1 Variable xs, abt '[] b)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs b
body
        in  (\List1 Variable xs
vars' ->
                let rho :: Assocs Variable
rho = List1 Variable xs -> List1 Variable xs -> Assocs Variable
forall k (xs :: [k]) (ast :: k -> *).
List1 Variable xs -> List1 ast xs -> Assocs ast
toAssocs1 List1 Variable xs
vars List1 Variable xs
vars'
                in  Pattern xs a
-> List1 Variable xs -> Dis abt r -> GBranch a (Dis abt r)
forall (a :: Hakaru) r (xs :: [Hakaru]).
Pattern xs a -> List1 Variable xs -> r -> GBranch a r
GBranch Pattern xs a
pat List1 Variable xs
vars' (abt '[] b -> Dis abt r
f (abt '[] b -> Dis abt r) -> abt '[] b -> Dis abt r
forall a b. (a -> b) -> a -> b
$ Assocs Variable -> abt '[] b -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
(ABT syn abt, JmEq1 Sing, Show1 Sing, Functor21 syn) =>
Assocs Variable -> abt xs a -> abt xs a
renames Assocs Variable
rho abt '[] b
body')
            ) (List1 Variable xs -> GBranch a (Dis abt r))
-> Dis abt (List1 Variable xs) -> Dis abt (GBranch a (Dis abt r))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> List1 Variable xs -> Dis abt (List1 Variable xs)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (xs :: [Hakaru]).
EvaluationMonad abt m p =>
List1 Variable xs -> m (List1 Variable xs)
freshenVars List1 Variable xs
vars
    (forall (a :: Hakaru).
 [Index (abt '[])] -> (r -> Ans abt a) -> Ans abt a)
-> Dis abt r
forall (abt :: [Hakaru] -> Hakaru -> *) x.
(forall (a :: Hakaru).
 [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a)
-> Dis abt x
Dis ((forall (a :: Hakaru).
  [Index (abt '[])] -> (r -> Ans abt a) -> Ans abt a)
 -> Dis abt r)
-> (forall (a :: Hakaru).
    [Index (abt '[])] -> (r -> Ans abt a) -> Ans abt a)
-> Dis abt r
forall a b. (a -> b) -> a -> b
$ \[Index (abt '[])]
i r -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l ->
        (Term abt ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt ('HMeasure a) -> abt '[] ('HMeasure a))
-> ([Branch a abt ('HMeasure a)] -> Term abt ('HMeasure a))
-> [Branch a abt ('HMeasure a)]
-> abt '[] ('HMeasure a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. abt '[] a -> [Branch a abt ('HMeasure a)] -> Term abt ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
abt '[] a -> [Branch a abt b] -> Term abt b
Case_ abt '[] a
e) ([Branch a abt ('HMeasure a)] -> abt '[] ('HMeasure a))
-> [[Branch a abt ('HMeasure a)]] -> [abt '[] ('HMeasure a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [GBranch a (Dis abt r)]
-> (GBranch a (Dis abt r) -> [Branch a abt ('HMeasure a)])
-> [[Branch a abt ('HMeasure a)]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
T.for [GBranch a (Dis abt r)]
gms (\GBranch a (Dis abt r)
gm ->
            GBranch a (abt '[] ('HMeasure a)) -> Branch a abt ('HMeasure a)
forall (syn :: ([Hakaru] -> Hakaru -> *) -> Hakaru -> *)
       (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru) (b :: Hakaru).
ABT syn abt =>
GBranch a (abt '[] b) -> Branch a abt b
fromGBranch (GBranch a (abt '[] ('HMeasure a)) -> Branch a abt ('HMeasure a))
-> [GBranch a (abt '[] ('HMeasure a))]
-> [Branch a abt ('HMeasure a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GBranch a (Dis abt r)
-> (Dis abt r -> [abt '[] ('HMeasure a)])
-> [GBranch a (abt '[] ('HMeasure a))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
T.for GBranch a (Dis abt r)
gm (\Dis abt r
m ->
                Dis abt r -> [Index (abt '[])] -> (r -> Ans abt a) -> Ans abt a
forall (abt :: [Hakaru] -> Hakaru -> *) x.
Dis abt x
-> forall (a :: Hakaru).
   [Index (abt '[])] -> (x -> Ans abt a) -> Ans abt a
unDis Dis abt r
m [Index (abt '[])]
i r -> Ans abt a
c ListContext abt 'Impure
h Assocs (Extra (abt '[]))
l))
{-# INLINE emitCaseWith #-}


----------------------------------------------------------------
----------------------------------------------------------- fin.