Probabilistic Effects. λθ

Fusion for Free

Free monads are at the heart of algebraic effect handlers, a functional approach for modelling side effects which allows the separation of the syntax and semantics of effectful operations, whilst enabling us to provide multiple different semantics for the same syntax.

The syntax of primitive side-effect operations is captured in a signature functor. The free monad over this functor assembles the syntax for the individual operations into an abstract syntax tree for an effectful program.

The semantics of the individual operations is captured in an algebra, and an effect handler folds the algebra over the syntax tree of the program to interpret it into a semantic domain.

2. Algebraic Effect Handlers

The idea of the algebraic effect handlers approach is to consider the free monad over a particular functor as an AST for an effectful computation. This functor is used to generate the nodes of a free structure whose leaves correspond to variables.

data Free f a where
    Var :: a -> Free f a
    Con :: f (Free f a) -> Free f a

Since a value of type Free f a is an inductive structure, we can define a fold for it by providing:

  • A function gen that deals with generation of values from Var x.
  • An algebra alg that is used to recursively collapse an operation Con op.
fold :: Functor f => (f b -> b) -> (a -> b) -> (Free f a -> b)
fold alg gen (Var x)  = gen x
fold alg gen (Con op) = alg (fmap (fold alg gen) op)

Algebraic effect handlers give a semantics to the syntax tree: one way of doing this is by using a fold. The behaviour of folds when composed with other functions is described by fusion laws.

  • The first law describes how certain functions that are precomposed with a fold can be incorporated into a new fold:
    fold alg gen · fmap h = fold alg (gen · h)
    
  • The second law shows how certain functions that are postcomposed with a fold can be incorporated into a new fold:
    k · fold alg gen = fold alg' (k · gen)
    

    This law is subject to the condition that:

    k · alg = alg' · fmap k
    

The monadic instance of the free monad given below.

instance Functor => Monad (Free f) where
    return x = Var x
    m >>= f  = fold Con f m

Variables are the way of providing a return for the monad. Extending a syntax tree by a means of a function f corresponds to applying that function to the variables found at the leaves of the tree.

2.1 Nondeterminism

A functor supplies the abstract syntax for the primitive effectful operations in the free monad: to demonsrate, the Nondet functor provides the Or k k syntax for a binary nondeterministic choice primitive. The parameter of type k marks the recursive site of syntax, which indicates where the continuation is after this syntactic fragment has been evaluated.

data Nondet k where
    Or :: k -> k -> Nondet k
instance Functor Nondet where
    fmap f (Or x y) = Or (f x) (f y)

This lets us express the syntax tree of a computation that nondeterministically returns True or False.

coin :: Free Nondet Bool
coin = Con (Or (Var True) (Var False))

The syntax is complemented by semantics in the form of effect handlers – functions which replace the syntax by values from a semantic domain. Using a fold for the free monad is a natural way of expressing such functions.

Here is an effect handler that interprets Nondet in terms of lists of possible outcomes:

handle_Nondet :: Free Nondet a -> [a]
handle_Nondet = fold alg_Nondet gen_Nondet

where alg_Nondet is the Nondet-algebra that interprets terms constructed by Or operations:

alg_Nondet :: Nondet [a] -> [a]
alg_Nondet (Or l1 l2) = l1 ++ l2

and gen_Nondet interprets variables:

gen_Nondet :: a -> [a]
gen_Nondet x = [x]

Generalising away from the details, handlers are usually presented in the following form, where F and H are arbitrary functors determined by the handler:

hdl :: ∀ a . Free F a -> H a

2.2 Handler Composition

We consider two classes of scenarios where effect handler composition is involved.

Effect Composition The first class of scenarios is where multiple effects are combined in the same program; to this end, we compose signatures and handlers.

The coproduct functor f + g makes it easy to compose functors.

data (+) f g a where
    Inl :: f a -> (f + g) a
    Inr :: g a -> (f + g) a

instance (Functor f, Functor g) => Functor (f + g) where
    fmap f (Inl s) = Inl (fmap f s)
    fmap f (Inr s) = Inr (fmap f s)

The free monad of a coproduct functor is then a tree where each node can be built from syntax from either f or g.

Composing handlers is easy too: if the handlers are written in a compositional style, then function composition does the trick.

Given that a normal effect handler for the functor F has the general form:

hdl ::  a . Free F a -> H a

A compositional handler for the functor F then has a signature of the form:

hdl ::  g a . Free (F + g) a -> H1 (Free g (G1 a))

This processes only the F-nodes in the AST and leaves the g-nodes as they are. Hence the result of the compositional handler is a new, typically smaller, AST with only g-nodes.

  • The variables of type G1 a in the resulting AST are derived from the variables of type a in the original AST as well as from the processed operations (the context G1 is due to the gen function which interprets variables).
  • Moreover, the new AST may be embedded in a context H1.

As an example of an effect signature, the compositional nondeterminism handler is defined as follows:

handle_Nondet :: Functor g => Free (Nondet + g) a -> Free g [a]
handle_Nondet = fold (alg_Nondet  Con) gen_Nondet

where given the general handler form hdl :: ∀ g a . Free (F + g) a -> H1 (Free g (G1 a))F is Nondet, G1 is [], and implicitly H1 is Id.

In here, the variables are handled with the monadified version of the previous gen_Nondet:

gen_Nondet :: Functor g => a -> Free g [a]
gen_Nondet x = Var [x]

The g nodes are handled by a Con algebra, which essentially leaves them untouched.

Con :: Functor g => g (Free g a) -> Free g a

The Nondet nodes are handled by the alg_Nondet algebra, which is a monadified version of the previous alg_Nondet:

alg_Nondet :: Functor g => Nondet (Free g [a]) -> Free g [a]
alg_Nondet (Or ml1 ml2) =
    do l1 <- ml1
       l2 <- ml2
       Var (l1 ++ l2)

The junction combinator (▽) composes the algebras for the two kinds of nodes.

() :: (f b -> b) -> (g b -> b) -> ((f + g) b -> b)
() alg_f alg_g (Inl s) = alg_f s
() alg_f alg_g (Inr s) = alg_g s

In the definition of handler_Nondet, we use alg_Nondet ▽ Con. Since the functor in question is Nondet + g, the values constructed by Nondet are handled by alg_Nondet, and the values constructed by g are left untouched (since the definition of fold unwraps one level of Con but using the Con as an algebra then replaces it with a Con again).

As another example of an effect signature is that of state.

data State s k where
    Put :: s -> k -> State s k
    Get :: (s -> k) -> State s k

instance Functor (State s) where
    fmap f (Put s k) = Put s (f k)
    fmap f (Get k)   = Get (f · k)

The compositional handler for state is defined as follows:

handle_State :: Functor g => Free (State s + g) a -> (s -> Free g a)
handle_State = fold (alg_State  con_State) gen_State

where given the general handler form hdl :: ∀ g a . Free (F + g) a -> H1 (Free g (G1 a))F is State s, H1 is s -> -, and implicitly G1 is Id.

In here, the variables are handled with gen_State. A variable x is replaced by a version that ignores any incoming state parameter s.

gen_State :: Functor g => a -> (s -> Free g a)
gen_State x s = Var x

The State nodes are handled by the alg_State algebra. This outputs a function which when given a state, the continuation k proceeds by using the appropriate state: if the syntax is a Put s' k, then the new state is s', otherwise the syntax is Get k, in which case the state is left unchanged and passed as s.

alg_State :: Functor g => State s (s -> Free g a) -> (s -> Free g a)
alg_State (Put s' k) = \s -> k s'
alg_State (Get (k :: s -> (s -> Free g a)))  = \s -> k s s

The g nodes are handled by the con_State algebra. This takes the extra state parameter s into account.

con_State :: Functor g => g (s -> Free g a) -> (s -> Free g a)
con_State op = \s -> Con (fmap (\(m :: s -> Free g a) -> m s) op)

To demonstrate effect composition, we can put Nondet and State together and handle them both. Before we do so, we also need a base case for the composition, which is the empty signature Void.

data Void k

instance Functor Void

The Void handler only provides a variable case since the signature has no constructors. In fact, a Free Void term can only be a Var x, so x is immediately output using the identity function (recall that fold alg gen (Var x) = gen x).

handle_Void :: Free Void a -> a
handle_Void = fold  id

Finally, we can put together a composite handler for programs that feature both non-determinism and state. The signature of such programs is the composition of the three basic signatures:

type Σ = Nondet + (State Int + Void)

The handler is the composition of the three handlers, working from the left-most functor in the signature:

handle_Σ :: Free Σ a -> Int [a]
handle_Σ prog = handle_Void . (handle_State . handle_Nondet) prog

Effect Delegation Another important class of applications are those where a handler expresses the complex semantics of particular operations in terms of more primitive effects.

For example, the syntax of the writer effect is captured by the folowing functor, where w is a parameter that represents the type of values that are written to the log.

data Writer w k where
    Tell :: w -> k -> Writer w k

instance Functor (Writer w) where
    fmap f (Tell w k) = Tell w (f k)

The following logging handler for state records every update of the state, by means of the Writer effect. For the Put cases, we take our previous definition \s -> k s' of alg_State, but wrap the result k s' in a Free (Writer String + Void) monad, by using Con (Inl (Tell "put" _)).

handle_LogState :: Free (State s) a -> s -> Free (Writer String + Void) a
handle_LogState = fold alg_LogState gen_State

alg_LogState :: State s (s -> Free (Writer String + Void) a)
                -> s -> Free (Writer String + Void) a
-- alg_State (Put s' k) = \s -> k s'
alg_LogState (Put s' k) = \s -> Con (Inl (Tell "put" (k s')))
-- alg_State (Get k)    = \s -> k s s
alg_LogState (Get k)    = \s -> k s s

A semantics for the writer effect can be given by the following handler, where w is constrained to be a member of the Monoid typeclass.

handle_Writer :: (Functor g, Monoid w) => Free (Writer w + g) a -> Free g (w, a)
handle_Writer = fold (alg_Writer  Con) gen_Writer

The variables are handled by gen_Writer, which pairs the variable with the unit of the writer’s log monoid, before embedding it into the monad m2 (which is the monad that we are logging e.g. the state monad).

gen_Writer :: (Monad m2, Monoid w) => a -> m2 (w, a)
gen_Writer x = return (mempty, x)

The Writer nodes are handled by the alg_Writer handler, which when encounters a Tell w1 k operation, the continuation k :: m2 (w, a) (which is already the correct return type) is >>=’d to a function which simply appends the log w1 to any generated logs.

alg_Writer :: (Monad m2, Monoid w) => Writer w (m2 (w, a)) -> m2 (w, a)
alg_Writer (Tell w1 k) = k >>= \(w2, x) -> return (w1 `mappend` w2, x)

To see this machinery in action, consider the following program that uses state (to add the sum of the sequence up to n onto the current state s):

program :: Int -> Free (State Int) Int
program n
    | n <= 0    = Con (Get Var)
    | otherwise = Con (Get (\s -> Con (Put (s + n) (program (n - 1)))))

This is then simply evaluated by running handlers in sequence.

example :: Int -> (String, Int)
example n = (handle_Void . handle_Writer . handle_LogState (program n)) 0

To fully interpret a stateful program, we must first run handle_LogState which interprets the Tell operations by generating a tree with Writer String syntax. This generated syntax is then handled with the handle_Writer handler.

3. Fusion

The composition of two handlers produces an intermediate abstract syntax tree. We aim to answer the question, “how can we use the two handlers into a single one that does not involve an intermediate tree?”.

More concretely, given two handlers of the form:

handler_1 :: Free F1 a -> H1 (Free F2 (G1 a))
handler_1 = fold alg_1 gen_1

handler_2 :: Free F2 a -> H2 a
handler_2 = fold alg_2 gen_2

F1 and F2 are signature functors. H1, G1, and H2 are arbitrary functors. Our goal is to obtain a combined handler:

pipeline_12 :: Free F1 a -> H1 (H2 (G1 a))
pipeline_12 = fold alg_12 gen_12

such that:

fmap handler_2 . handler_1 = pipeline_12
Last updated on 13 Nov 2020
Published on 13 Nov 2020