Stateful Functions

Motivating Example: Stacks

Let's define the following to represent and manipulate stacks, a data structure that allows adding and removing element only from the top end:

type Stack = [Int]

pop   :: Stack -> (Int, Stack)
push_ :: Int -> Stack -> Stack      -- first version

pop []     = (0, [])       -- default element, rather than failing
pop (i:is) = (i, is)

push_ i is = i:is

We can test our Stack with a sequence of operations:

testStack_ :: Stack -> (Int, Stack)
testStack_ s0 =
  let s1 = push_ 0 s0 in
  let s2 = push_ 1 s1 in
  let s3 = push_ 2 s2 in
  let s4 = push_ 3 s3 in
  let (a,s5) = pop s4 in
  let (b,s6) = pop s5 in
  let s7 = push_ (a+b) s6 in
  pop s7

> testStack_ []
(5,[1,0])

The stack is "threaded through" a series of computations. Let's make this even more explicit:

push :: Int -> Stack -> ((), Stack)
push i is = ((), i:is)

testStack :: Stack -> (Int, Stack)
testStack s0 =
  let (_,s1) = push 0 s0 in
  let (_,s2) = push 1 s1 in
  let (_,s3) = push 2 s2 in
  let (_,s4) = push 3 s3 in
  let (a,s5) = pop s4 in
  let (b,s6) = pop s5 in
  let (_,s7) = push (a+b) s6 in
  pop s7

What's wrong with this? First of all, there's a lot of syntactic "noise." Worse yet, it's easy to make a mistake; the intention is that each version of the stack si is referred to once and "discarded" in favor of the new stack returned by the operation.

Sequenced Computations

There is a common idiom in the code above: the threading of some object through a series of computations.

computation :: StateObject -> (T7, StateObject)
computation s0 =
  let (a1,s1) = f0 s0 in      -- f0 :: StateObject -> (T1, StateObject)
  let (a2,s2) = f1 s1 in      -- f1 :: StateObject -> (T2, StateObject)
  let (a3,s3) = f2 s2 in      -- f2 :: StateObject -> (T3, StateObject)
  let (a4,s4) = f3 s3 in      -- f3 :: StateObject -> (T4, StateObject)
  let (a5,s5) = f4 s4 in      -- f4 :: StateObject -> (T5, StateObject)
  let (a6,s6) = f5 s5 in      -- f5 :: StateObject -> (T6, StateObject)
  let (a7,s7) = f6 s6 in      -- f6 :: StateObject -> (T7, StateObject)
    (a7,s7)

We refer to this object as "the state" even though, if you're familiar with other languages with "mutable" or "stateful" features, there's nothing like that here. Just ordinary pure functions, with a pattern of use that feels like we're manipulating state.

type StatefulFunction s a =
  s -> (a,s)                        -- name for function type idiom

data StatefulFunction s a =
  StatefulFunction (s -> (a,s))     -- new datatype to define instances

newtype StatefulFunction s a =
  StatefulFunction (s -> (a,s))     -- newtype b/c one, one-arg constructor

newtype StatefulFunction s a =
  StatefulFunction { runStatefulFunction :: s -> (a,s) }  -- ... field for unboxing

StatefulFunction s a means a computation that, starting with an input state of type s, produces a value of type a and an updated state of type s.

When you read StatefulFunction s a, think "function of type s -> (a,s)" (but boxed up in a newtype). Or think "stateful computation that produces an a" keeping in mind that there is input and output state "in the background."

To preview the benefits that this abstraction will provide, we are going to define

pop'  :: State Stack Int
push' :: Int -> State Stack ()

and, then, because StatefulFunction s is a Monad, we will write the previous sequence of stack operations as:

testStack' = do
  push' 0
  push' 1
  push' 2
  push' 3
  a <- pop'
  b <- pop'
  push' (a+b)
  pop'

Sequencing Stateful Functions

We can think of sequencing StatefulFunctions as function composition, with the appropriate plumbing to thread the state objects through the component functions.

 s0   ------  s1   ------  s2   ------  s3
----> | f0 | ----> | f1 | ----> | f2 | ---->
      ------       ------       ------
          \_________/  \_________/  \------>
            a1           a2           a3

   f0                               :: s -> (a1, s)
   f1                               :: s -> (a2, s)
   f2                               :: s -> (a3, s)

   f0 >>= \a1 -> f1 >>= \a2 -> f2   :: s -> (a3, s)

So, let's define how StatefulFunction s forms a monadic, applicative functor.

fmap         :: (a -> b) -> StatefulFunction s a -> StatefulFunction s b
(<*>)        :: StatefulFunction s (a -> b) -> StatefulFunction s a -> StatefulFunction s b
(>>=)        :: StatefulFunction s a -> (a -> StatefulFunction s b) -> StatefulFunction s b
pure, return :: a -> StatefulFunction s a

(Recall that we could define return and (>>=) first, and then get the Functor and Applicative instances "for free" with boilerplate definitions. Below, we'll develop fmap and (<*>) separately to build an understanding of how they work.)

Let's first do things out of order, independent of instances.

st_pure x = StatefulFunction $ \s -> (x, s)

st_fmap g (StatefulFunction f) = StatefulFunction $ \s ->
  let (a,t) = f s in
  (g a, t)

Alternatively, runState to avoid pattern matching:

st_fmap g sa = StatefulFunction $ \s ->
  let (a,t) = runStatefulFunction sa s in
  (g a, t)
st_ap sab sa = StatefulFunction $ \s ->
  let (f, t) = runStatefulFunction sab s in
  let (a, u) = runStatefulFunction sa t in
  (f a, u)

st_bind sa f = StatefulFunction $ \s ->
  let (a, t) = runStatefulFunction sa s in
  let (b, u) = runStatefulFunction (f a) t in
  (b, u)

Now, put these definitions into instance declarations.

Helper functions

statefulFunction     :: (s -> (a,s)) -> StatefulFunction s a  -- create a stateful comp.
get                  :: StatefulFunction s s                  -- get state out
put                  :: s -> StatefulFunction s ()            -- set "current" state
modify               :: (s -> s) -> StatefulFunction s ()     -- modify the state
evalStatefulFunction :: StatefulFunction s a -> s -> a        -- run and return final value
execStatefulFunction :: StatefulFunction s a -> s -> s        -- run and return final state

statefulFunction f        = StatefulFunction $ f
get                       = StatefulFunction $ \s -> (s, s)
put t                     = StatefulFunction $ \s -> ((), t)
modify f                  = StatefulFunction $ \s -> ((), f s)
evalStatefulFunction sa s = fst $ runStatefulFunction sa s
execStatefulFunction sa s = snd $ runStatefulFunction sa s

Programming with StatefulFunction Stack

pop'  :: StatefulFunction Stack Int
push' :: Int -> State Stack ()

pop'    = StatefulFunction $ \stk -> case stk of {[] -> (0,[]); i:is -> (i,is)}
push' i = StatefulFunction $ \stk -> ((), i:stk)

If we want to, we can redefine pop' and push' using do-notation and the helpers for "reading" and "writing" the state (Stack) in the background.

push' i = do                -- push' i =
  is <- get                 --   get >>= \is ->
  put (i:is)                --   put (i:is)

pop' = do                   -- pop' =
  stk <- get                --   get >>= \stk ->
  case stk of               --   case stk of
    []   -> return 0        --     []   -> return 0
    i:is -> do              --     i:is ->
      put is                --       put is >>
      return i              --       return i

Now, let's go back to our long sequence of stack operations.

testStack' :: StatefulFunction Stack Int
testStack' = do
  push' 0
  push' 1
  push' 2
  push' 3
  a <- pop'
  b <- pop'
  push' (a+b)
  pop'

> runStatefulFunction testStack' []
(5,[1,0])

Cool!

Source Files

results matching ""

    No results matching ""