Разбор монады ContT, или как написать продвинутый goto, генераторы и файберы на хаскеле

В этой статье я разберу монаду ContT, и покажу как вернуть return и другие control-flow операторы из императивных языков программирования, которых мне так нехватало, когда я начинал изучать хаскель.

0. Введение в продолжения (continuations)

Рассмотрим выражение f (g x). Такой подход к передаче аргументов не масштабируется. Если имена громоздкие и вложенность большая, код становится совершенно нечитаемым:

largeFunctionFoo (largeFunctionBar (largeFunctionBaz x))

Есть множество способов борьбы с этим, но два наиболее известных - это пайпы x |> g |> f и передача продолжений (CPS - Continuation passing style) ($ x) g' f' id.

Продолжения очень похожи на колбэки. Функция с поддержкой продолжений проводит вычисления, вызывает продолжение и возвращает его результат. Например, функцию g' можно определить так:

g' x cont = cont (g x)

Сразу отметим одну интересную особенность: CPS функция знает про своё продолжение. В то время как в f (g x) функции g ничего не известно про функцию f, функции g' продолжение передаётся явно. Это позволяет ей управлять продолжением, вызывать его множество раз или не вызывать вовсе.

1. Добавляем продолжения в хаскель

Определим тип для продолжаемых значений:

data ContT r m a = ContT {runCont :: (a -> m r) -> m r}

Дальше по тексту функция a -> m r часто будет называться cc - current continuation

Чтобы привыкнуть к определению разберём пару примеров:

Попробуем определить продолжаемое число 5:

cont5 :: ContT r m Integer
cont5 = ContT $ \cc -> cc 5 -- ContT ($ 5), если кратко

cont5 - это функция, которая получает продолжение и вызывает его с числом 5.

Можно убедиться, что cont5 действительно хранит нужное число, вызвав cont5 с продолжением print:

ghci> runCont cont5 print
5

Определим функцию прибавления единицы к продолжаемому числу:

addOne :: ContT r m Integer -> ContT r m Integer
addOne contX = ContT $
  \cc -> runCont contX $
    \x -> cc (x + 1)

addOne получает продолжаемое число, и на его основе строит новое число. Внутри функция получает продолжение cc, вызывает старое число contX с продолжением \x -> cc (x + 1), которое, в свою очередь, вызывает продолжение cc со значением x + 1.

Проверим, что (addOne cont5) хранит число 6:

ghci> runCont (addOne cont5) print
6

Чтобы с продолжаемыми значен��ями можно было удобно работать, определим для них инстансы классов Functor, Applicative и Monad:

instance Functor (ContT r m) where
  fmap f contX = ContT $
    \cc -> runCont contX $
      \x -> cc (f x)

instance Applicative (ContT r m) where
  pure a = ContT ($ a)

  contF <*> contX = ContT $
    \cc -> runCont contF $
      \f -> runCont contX $
        \x -> cc (f x)

instance Monad (ContT r m) where
  return = pure

  contX >>= f = ContT $
    \cc -> runCont contX $
      \x -> runCont (f x) $
        \res -> cc res

Теперь вычисления с продолжениями можно писать гораздо проще. Функции из примеров можно переопределить так:

cont5 :: ContT r m Integer
cont5 = return 5

addOne :: ContT r m Integer -> ContT r m Integer
addOne contX = do
  x <- contX
  return (x + 1)

Пока что все примеры с использованием ContT были довольно прямолинейными. Продолжения вызывались ровно один раз, никаких пропусков не происходило. В следующей главе мы это исправим

2. Строим callCC

Посмотрим ещё раз на тип ContT:

data ContT r m a = ContT {runCont :: (a -> m r) -> m r}

Заметили? Если у нас есть значение типа m r, можно вернуть его, не вызывая cc. Тогда все последующие действия (которые выполнялись внутри cc) будут пропущены.

Временно зафиксируем r = (), m = IO и напишем skip3 - продолжаемое значение, которое пропускает действия и печатает 3:

skip3 :: ContT () IO a
skip3 = ContT $ \cc -> print 3 -- ContT (const $ print 3)

Отметим, что тип a не важен, так как cc не вызывается. Проверим, что skip3 работает:

test1 = do
  return 5

test2 = do
  skip3
  return 5

ghci> runCont test1 print
5
ghci> runCont test2 print
3

Это даёт возможность пропускать шаги, но из-за фиксации r и m, skip3 получился слишком ограниченным. Напишем более гибкий механизм.

Для этого посмотрим на ситуацию изнутри ContT:

cont :: ContT r m a
cont = ContT $ \cc -> _ -- ?

cc внутри cont имеет тип a -> m r. Это даёт возможность получить m r, который уже можно использовать для пропуска. Добавим определение skip внутрь cont:

cont :: ContT r m a
cont = ContT $
  \cc ->
    let skip a = ContT (const $ cc a)
     in _ -- ?

skip имеет тип a -> ContT r m b, skip a работает аналогично skip3. Таким образом внутри cont мы можем запустить некоторую "малую" функцию, которой передадим функцию skip для переключения во внешний ContT r m a.

Так мы приходим к определению функции callCC (call with current continuation):

type Switch r m a = forall b. a -> ContT r m b

callCC :: (Switch r m a -> ContT r m a) -> ContT r m a
callCC f = ContT $
  \cc ->
    let switchToCC = \a -> ContT (const $ cc a)
     in runCont (f switchToCC) cc

Эта функция будет основой почти для всех дальнейших конструкций.

Рассмотрим несколько примеров с callCC.

Начнём с возможности досрочно завершить функцию:

test = callCC $ \exit -> do
  lift $ putStrLn "Reachable"
  exit ()
  lift $ putStrLn "Unreachable"

ghci> runCont test (const $ return ())
Reachable

Можно вернуть операторы break и continue:

test :: ContT r IO ()
test = do
  forM_ [1 .. 10] $ \i -> do
    callCC $ \continue -> do
      when (i == 5) $ do
        continue ()

      lift $ print i

ghci> runCont test (const $ return ())
1
2
3
4
6
7
8
9
10

test :: ContT r IO ()
test = do
  callCC $ \break -> do
    forM_ [1 .. 10] $ \i -> do
      when (i == 5) $ do
        break ()

      lift $ print i

ghci> runCont test (const $ return ())
1
2
3
4

Отметим одну очень важную особенность: функция switchToCC не прерывает "малую" функцию. Она завершает текущее вычисление и переключается на то, из которого была вызвана функция callCC. Если каким-то образом switchToCC сможет выйти за пределы callCC, вызов этой функции выполнит это переключение:

test = do
  val {- label -} <- callCC $ \exit -> do
    exit 10
  -- rest

exit переносит нас к точке label, исполнение продолжается с присваивания val.

Перейдём к более продвинутым примерам использующим вызовы switchToCC снаружи callCC

3. Функция label (она же goto, она же аналог хука useState)

В этой главе мы напишем функцию label. Она получает начальное значение и возвращает пару (restart, value).

• restart позволяет перезапустить вычисление с другим значением

• value - текущее значение

На самом деле определение этой функции довольно простое:

label :: a -> ContT r m (a -> ContT r m b, a)
label init = callCC $
  \switch ->
    let restart val = switch (restart, val)
     in return (restart, init)

Теперь мы можем описывать некоторые вещи в императивном стиле

Например, циклы:

test = do
  (restart, counter) <- label 0

  lift $ print counter

  when (counter < 10) $ do
    restart $ counter + 1

ghci> runCont test (const $ return ())
0
1
2
3
4
5
6
7
8
9
10

Можно написать пару setjmp/longjmp. Это даже проще чем label потому что им не нужно дополнительно нести состояние:

setjmp = do
  (restart, _) <- label ()

  let longjmp = restart ()

  return longjmp

test = do
  longjmp <- setjmp
  lift $ print 10

  longjmp

ghci> runCont test (const $ return ())
10
10
10
10
10
10
...

4. Генер��торы, файберы и планировщик

Попробуем написать ещё более продвинутые механизмы для управления программой.

Порассуждаем о том, как может выглядеть управление генератором. Для этого можно посмотреть как они сделаны в других языках. Обычно для этого используют ключевое слово yield, однако есть несколько нюансов:

0. Возможности добавить ключевое слово у нас нет, поэтому yield будет функцией

1. Необходимо, чтобы из yield x можно было получить следующий yield. Если дважды использовать один и тот же yield, мы выйдем в одной и той же точке, что нам не подходит

2. Нужна возможность вызвать exit (его тоже нужно обновлять при вызове yield) чтобы вернуться в планировщик

Положим управляющие функции в отдельную структуру:

data Controls r m x = Controls
  { yield :: x -> ContT r m (Controls r m x)
  , exit :: forall b. ContT r m b
  }

Тогда код генератора мог бы выглядеть так:

test controls = do
  controls <- yield controls 1
  controls <- yield controls 2
  controls <- yield controls 3
  exit controls

Может показаться странным, что мы явно вызываем exit controls для выхода в конце генератора, но дальше будет понятно зачем это.
Пока просто на уровне типов запретим генератору завершаться без вызова exit:

type Generator r m x = Controls r m x -> ContT r m Void

Теперь напишем функцию для запуска генератора до следующего yield.
Она будет возвращать пару из значения и следующей части генератора или Nothing если генератор завершился:

runToYield :: Generator r m x -> ContT r m (Maybe (x, Generator r m x))
runToYield generator = callCC $ \exitContext -> do
  let exit = exitContext Nothing
      yield value = callCC $ \continueGenerator ->
        exitContext $ Just (value, continueGenerator)

      controls = Controls{yield, exit}

  generator controls
  -- Если генератор добрался сюда, происходит что-то похожее на провал стека в
  -- императивных языках. Дальнейшее исполнение, если убрать error, хоть и
  -- определено, но будет очень странным и контринтуитивным
  error "Generator exit invariant violated"

Разберём этот код построчно:

• callCC $ \exitContext -> do - создаём новый контекст в котором запустим генератор

• exit = exitContext Nothing - функция для возвращения в родительский контекст

• yield value - определяем функцию для передачи значения

• callCC $ \continueGenerator -> - изнутри захватываем состояние генератора

• exitContext $ Just (value, continueGenerator) - выходим в родительский контекст, возвращаем значение и следующую часть генератора

• controls = Controls{yield, exit} - просто определяем удобный синоним

• generator controls - запускаем генератор

• error "Generator exit invariant violated" - мы с помощью типа Void запретили генератору завершаться как-либо помимо вызова exit. Эта строчка кода должна быть unreachable

Файберы - это генераторы, у которых yield передаёт управление назад планировщику (без передачи каких-либо значений).
Определим пару удобных синонимов для них:

type Fiber r m = Generator r m ()

suspend :: Controls r m () -> ContT r m (Controls r m ())
suspend controls = yield controls ()

Напишем планировщик для удобного запуска множества файберов. Работать он будет очень просто: берём список файберов, запускаем все по одному разу, убираем те, которые завершились. Повторять пока остались незавершённые файберы:

scheduler :: [Fiber r m] -> ContT r m ()
scheduler threads = do
  let round threads = do
        nextThreads <- forM threads $ \thread -> do
          res <- runToYield thread
          return $ snd <$> res

        return $ catMaybes nextThreads -- Фильтрация файберов

  (loop, threads) <- label threads
  threadsLeft <- round threads

  when (length threadsLeft /= 0) $ do
    loop threadsLeft

Осталось лишь проверить, что файберы действительно работают:

debug :: (MonadIO m) => String -> m ()
debug str = liftIO $ putStrLn str

fiberA :: Fiber r IO
fiberA controls = do
  debug "Started fiber A"
  controls <- suspend controls
  debug "Running fiber A"
  controls <- suspend controls
  debug "Exiting fiber A"
  exit controls

fiberB :: Fiber r IO
fiberB controls = do
  debug "Started fiber B"
  controls <- suspend controls
  debug "Running fiber B"
  controls <- suspend controls
  debug "Running fiber B again"
  controls <- suspend controls
  debug "Exiting fiber B"
  exit controls

ghci> runCont (scheduler [fiberA, fiberB]) (const $ return ())
Started fiber A
Started fiber B
Running fiber A
Running fiber B
Exiting fiber A
Running fiber B again
Exiting fiber B

Ну вот и всё. Спасибо за внимание