datatypes polymorphic in strictness

PolymorphicStrictness2.hs

-- a better version of https://gist.github.com/gelisam/58495f7b996f77e09d80
-- based on http://www.reddit.com/r/haskell/comments/2chb2h/fantasy_world_haskell/cjgepwj
--
-- I implement the same example three times: once with a lazy datatype,
-- once with a strict datatype, and once with a strict-polymorphic datatype.
{-# LANGUAGE BangPatterns, KindSignatures, TypeOperators #-}

import Debug.Trace
import Text.Printf


-- trace calls, to make it clear whether the code is evaluated eagerly or lazily.

noisyId :: Int -> Int
noisyId x = trace msg x
  where
    msg = printf "<generating %d>" x

noisyPlus :: Int -> Int -> Int
noisyPlus x y = trace msg (x + y)
  where
    msg = printf "<adding %d + %d>" x y


-- first version: an ordinary, lazy datatype.
data LazyList a
  = LazyNil
  | LazyCons a (LazyList a)
  deriving Show

lazyFromList :: [a] -> LazyList a
lazyFromList []     = LazyNil
lazyFromList (x:xs) = LazyCons x (lazyFromList xs)

lazyFMap :: (a -> b) -> LazyList a -> LazyList b
lazyFMap _ LazyNil         = LazyNil
lazyFMap f (LazyCons x xs) = LazyCons (f x) (lazyFMap f xs)

-- | with a lazy datatype, we get this evaluation pattern:
-- >>> lazyTest
-- [<generating 0>
-- <adding 1 + 0>
-- <adding 10 + 1>
-- 11,<generating 1>
-- <adding 1 + 1>
-- <adding 10 + 2>
-- 12,<generating 2>
-- <adding 1 + 2>
-- <adding 10 + 3>
-- 13]
lazyTest :: [Int]
lazyTest = map (noisyPlus 10)
         $ map (noisyPlus 1)
         $ map noisyId
         $ [0..2]


-- second version: a strict datatype, with bang patterns everywhere
data StrictList a
  = StrictNil
  | StrictCons !a !(StrictList a)
  deriving Show

strictFromList :: [a] -> StrictList a
strictFromList []     = StrictNil
strictFromList (x:xs) = StrictCons `app` x `app` strictFromList xs

strictFMap :: (a -> b) -> StrictList a -> StrictList b
strictFMap _ StrictNil         = StrictNil
strictFMap f (StrictCons x xs) = StrictCons (f x) (strictFMap f xs)

-- | with a strict datatype, we get this evaluation pattern:
-- >>> strictTest
-- <generating 0>
-- <generating 1>
-- <generating 2>
-- <adding 1 + 0>
-- <adding 1 + 1>
-- <adding 1 + 2>
-- <adding 10 + 1>
-- <adding 10 + 2>
-- <adding 10 + 3>
-- StrictCons 11 (StrictCons 12 (StrictCons 13 StrictNil))
strictTest :: StrictList Int
strictTest = strictFMap (noisyPlus 10)
           $ strictFMap (noisyPlus 1)
           $ strictFromList
           $ map noisyId
           $ [0..2]


-- for the third version, we need to create a few strictness-polymorphic constructs first.

data a :-> b = PrivateStrict (a -> b)

strictly :: (a -> b) -> a :-> b
strictly f = PrivateStrict (\a -> a `seq` f a)

runStrictly :: (a :-> b) -> a -> b
PrivateStrict f `runStrictly` a = f a


class FunctionLike f
  where
    lambda :: (a -> b) -> a `f` b
    app :: a `f` b -> a -> b

instance FunctionLike (->)
  where
    lambda = id
    app = ($)

instance FunctionLike (:->)
  where
    lambda = strictly
    app = runStrictly


-- finally, our strictness-polymorphic version.
-- List (->) Int is lazy, while List (:->) Int is strict.
data List (f :: * -> * -> *) (a :: *)
  = PrivateNil
  | PrivateCons a (List f a)
  deriving Show

mkNil :: List f a
mkNil = PrivateNil

-- this is the key:
-- if f is (->), mkCons will behave like LazyCons, and
-- if f is (:->), mkCons will behave like StrictCons.
mkCons :: FunctionLike f => a `f` (List f a `f` List f a)
mkCons = lambda $ \x -> lambda $ \xs -> PrivateCons x xs


fromList :: FunctionLike f => [a] -> List f a
fromList []     = mkNil
fromList (x:xs) = mkCons `app` x `app` fromList xs

-- pattern-matching on private constructors is a bit weird,
-- but I don't think it affects strictness.
ffmap :: FunctionLike f => (a -> b) -> List f a -> List f b
ffmap _ PrivateNil         = mkNil
ffmap f (PrivateCons x xs) = mkCons `app` f x `app` ffmap f xs


polymorphicTest :: FunctionLike f => List f Int
polymorphicTest = ffmap (noisyPlus 10)
                $ ffmap (noisyPlus 1)
                $ fromList
                $ map noisyId
                $ [0..2]

-- |
-- >>> main
-- PrivateCons <generating 0>
-- <adding 1 + 0>
-- <adding 10 + 1>
-- 11 (PrivateCons <generating 1>
-- <adding 1 + 1>
-- <adding 10 + 2>
-- 12 (PrivateCons <generating 2>
-- <adding 1 + 2>
-- <adding 10 + 3>
-- 13 PrivateNil))
-- <generating 0>
-- <generating 1>
-- <generating 2>
-- <adding 1 + 0>
-- <adding 1 + 1>
-- <adding 1 + 2>
-- <adding 10 + 1>
-- <adding 10 + 2>
-- <adding 10 + 3>
-- PrivateCons 11 (PrivateCons 12 (PrivateCons 13 PrivateNil))
main :: IO ()
main = do print (polymorphicTest :: List (->) Int)
          print (polymorphicTest :: List (:->) Int)