Haskell - Mixed stateful computations

Given the following structure:

data Computation a = Pure (CState -> (a, CState))
                   | Action (CState -> IO (a, CState))

(CState is some structure used for keeping a state, but is not of much interest now.)
Now I want to make it an instance of Monad, which is basically just a state monad and would be easily implemented with StateT. The only addition it has, is that I want to keep track, whether the resulting Computation is a Pure on or an Action and I want to be able to check whether a Computation contains any Actions, before executing an Action(so the IO in the Action doesn't get executed).

It should be also noted, that it doesn't really matter that Computation has two constructors. I just started implementing this with these constructors.

The rules of determining whether a >> b is Pure is simple: a >> b is Pure if a and b are both Pure, otherwise it's an action.

Now I started implementing the Monad instance:

instance Monad Computation where
    return x = Pure $ \s -> (x, s)
    (Action c) >>= f = Action . runStateT $
                         (StateT $ unpackAction oldComp) >>= (StateT . unpackAction . f)
    p@(Pure c) >>= f
      | givesPure f = Pure . runState $
                        state oldF >>= (state . unpackPure . f)
      | otherwise = liftComp p >>= f -- Just lift the first argument and recurse, to make an action

-- Helper functions used above:
unpackAction :: Computation a -> (CState -> IO (a, CState))
unpackAction (Pure c) = return . c 
unpackAction (Action c) = c

-- Make an Action out of a Pure
liftComp :: Computation a -> Computation a
liftComp (Pure c) = Action $ return . c
liftComp a@(Action _) = a

So the only missing part is the givesPurefunction and I'm not sure if it's even possible to implement it. I once had an implementation like this:

givesPure :: (a -> Computation b) -> Bool
givesPure f = isPure $ f undefined -- I actually used error with a custom message instead of undefined, but that shouldn't matter

isPure :: Computation a -> Bool
isPure (Pure _) = True
isPure (Action _) = False

This works, but makes the assumption that the function that I'm binding to, alway returns a Computation with the same purity, regardless of what its input is. This assumption appeared reasonably to me, since the purity of a Computation should be clearly stated and not depend on some calculations, until I noticed that a function in the following form dosn't work with this assumption:

baz :: Int -> Computation b
baz x = if x > 5 
        then foo
        else bar
-- foo and bar both have the type Computation b

So it seems to me, like it's impossible to do this, since I need the current state to apply the first Computation to, to get a proper input for the function to get the second Computation and test whether it's pure.

Does anyone see a solution to this or has a proof that it isn't possible?

Solution

You could try using a GADT:

data Pure
data Action

data Computation t a where
   Pure :: (CState -> (a, CState))      -> Computation t a
   Action :: (CState -> IO (a, CState)) -> Computation Action a

The idea is that a value x :: Computation Action a can do IO (but could also be pure), while a value y :: Computation Pure a can not do IO.

For instance,

liftComp :: Computation t a -> Computation Action a
liftComp (Pure c) = Pure c
liftComp x@(Action c) = x