Implementing tail for a peculiar stream representation.

StreamTailCPS.hs

-- In response to https://twitter.com/totbwf/status/1565154284135534592?s=20&t=1rtY9RMVImoAopR-ia24zA
{-# LANGUAGE RankNTypes, ScopedTypeVariables, TypeApplications #-}
module Main where

import Prelude hiding (head, tail)
import Test.DocTest (doctest)

-- $setup
-- >>> import System.Timeout (timeout)


-- The challenge is to implement 'tail' for a stream defined as
-- @forall r. (a -> r -> r) -> r@.

newtype Stream a = Stream
  { runStream :: forall r. (a -> r -> r) -> r }


-- How does this type represent a stream? The @theirCons :: a -> r -> r@
-- function has the right shape to be a @cons@ function which prepends an @a@
-- element to a stream of type @r@, resulting in a bigger stream @r@. Since a
-- stream has infinitely-many elements, we need to call this @theirCons@
-- function infinitely-many times.

nats :: Stream Int
nats = Stream go
  where
    go :: forall r. (Int -> r -> r) -> r
    go theirCons = goFrom 0
      where
        goFrom :: Int -> r
        goFrom n = theirCons n (goFrom (n+1))


-- |
-- Since @r@ is universally-quantified, a function which wants to observe and
-- manipulate a value of type @Stream A@ must pick a concrete type @R@ at which
-- to instantiate @r@, and a concrete function @myCons :: A -> R -> R@ for the
-- chosen @R@. That function will be called infinitely-many times, so it is
-- important to pick a lazy function! With a function which is strict in @R@,
-- 'runStream' loops forever:
--
-- >>> timeout 100000 $ print $ badLength nats
-- Nothing
badLength :: forall a. Stream a -> Int
badLength xs = runStream xs @Int myStrictCons
  where
    myStrictCons :: a -> Int -> Int
    myStrictCons _ n = n + 1

-- |
-- One way to be lazy in @R@ is to completely ignore the @R@ argument. This is
-- how 'head' is implemented:
--
-- >>> head nats
-- 0
head :: forall a. Stream a -> a
head xs = runStream xs @a myLazyCons
  where
    myLazyCons :: a -> a -> a
    myLazyCons x _ = x

-- |
-- Another way is to construct a datatype. Here, the @R@ argument will not be
-- forced until the tail of the resulting list is forced. Thus, if we only
-- force the first 10 elements of the list, then only 10 of the infinitely-many
-- calls to 'myCons' will actually run:
--
-- >>> take 10 $ toList nats
-- [0,1,2,3,4,5,6,7,8,9]
toList :: forall a. Stream a -> [a]
toList xs = runStream xs @[a] myCons
  where
    myCons :: a -> [a] -> [a]
    myCons = (:)

-- |
-- Before we can take on the challenge of implementing 'tail', we need to
-- implement another common function: 'cons'. It simply needs to call
-- @theirCons@ one more time than its argument.
--
-- >>> take 10 $ toList $ cons (-1) nats
-- [-1,0,1,2,3,4,5,6,7,8]
cons :: forall a. a -> Stream a -> Stream a
cons x xs = Stream go
  where
    go :: forall r. (a -> r -> r) -> r
    go theirCons = theirCons x (runStream xs @r theirCons)

-- |
-- We are finally ready to take on the challenge! We need to use the common
-- trick of _returning_a_function_ [1]. That's it, that's the trick which makes
-- this a puzzle instead of a mere exercise.
-- 
-- [1] https://hackage.haskell.org/package/recursion-schemes-5.2.2.2/docs/Data-Functor-Foldable.html#v:cataA
--
-- >>> take 10 $ toList $ tail nats
-- [1,2,3,4,5,6,7,8,9,10]
tail :: forall a. Stream a -> Stream a
tail xs = runStream xs @(Bool -> Stream a) myCons False
  where
    -- The Bool specifies whether to include the head
    myCons :: a -> (Bool -> Stream a) -> Bool -> Stream a
    myCons _ f False = f True
    myCons x f True = cons x (f True)


main :: IO ()
main = do
  putStrLn "typechecks."
  doctest ["src/Main.hs"]