Partial Functions in Haskell

PartialFunc.hs

{-# LANGUAGE TypeOperators #-}

import Data.Maybe
import Control.Applicative
import Control.Category
import Prelude hiding ((.))

data a ~> b =
    Partial (a -> Maybe b) -- a partial function
  | Defaulted (a ~> b) b -- a partial function with a default value

instance Category (~>) where
  id = Partial Just
  (.) = flip andThen

partial :: (a -> Maybe b) -> a ~> b
partial f = Partial f

liftP :: (a -> b) -> a ~> b
liftP f = Partial (Just . f)

-- | Apply `a` to the partial function, returing a `Maybe` if not defined
apply :: (a ~> b) -> a -> Maybe b
apply (Partial f) a = f a
apply (Defaulted f d) a = apply f a <|> Just d

-- | Like `apply` but throws an error if not defined (and not defaulted)
unsafeApply :: (a ~> b) -> a -> b
unsafeApply f = fromJust . apply f

-- | Apply `a` and return `b` if the function isn't defined at `a`
applyOrElse :: (a ~> b) -> a -> b -> b
applyOrElse f a b = fromMaybe b $ apply f a

-- | Add a default value to a partial function. If the function was already
-- defaulted, override the value with the new default.
withDefault :: (a ~> b) -> b -> (a ~> b)
withDefault f@(Partial _) b = Defaulted f b
withDefault (Defaulted f _) b = Defaulted f b

-- | Whether `a` is defined in the function
isDefinedAt :: (a ~> b) -> a -> Bool
isDefinedAt (Partial f) a = isNothing $ f a
isDefinedAt (Defaulted f _) a = isDefinedAt f a

-- | Create a new partial function where the domain is the combination
-- of both partial functions. Priority is given to the first partial function
-- in case of conflict.
orElse :: (a ~> b) -> (a ~> b) -> a ~> b
orElse (Partial f) (Partial g) = Partial $ \a -> f a <|> g a
orElse (Defaulted f b) (Defaulted g _) = Defaulted (f `orElse` g) b
orElse f@(Partial _) (Defaulted g b) = Defaulted (f `orElse` g) b
orElse (Defaulted f b) g@(Partial _) = Defaulted (f `orElse` g) b

-- | Create a new partial function where the first function is called followed
-- by the second one. This is `(.)` flipped.
andThen :: (a ~> b) -> (b ~> c) -> a ~> c
andThen f g = Partial $ \a -> apply f a >>= apply g

data Msg = Timeout
         | Accepted
         | Error

-- This is where we hit the issue. It would be nice for the syntax to
-- understand that these things need to be lifted into `maybe` and that any
-- undefined case is "Nothing". Additionally, it would be nice to be able
-- do `f a` and have it call `apply` above.
recieve :: Msg ~> String
recieve = partial go
  where
    go Timeout = return "TM"
    go Accepted = return "Ac"