Bartosz Milewski's neat defunctionalization talk

Bartosz kept a cool talk just recently. He also published a blogpost about it, "Defunctionalization and Freyd's Theorem".

In the video Bartosz tells how to defunctionalize arbitrary functions. This operation reveals the contents that these operations hold in their return stack during the evaluation. Bartosz also presents how he reasoned about this technique in category theory. I didn't understand much of the category-theory side, except that an infinite sum forming a state space for a function was brought up.

The video is part of the Haskell Love Conference, and what I can gather up from what's popping up in my Youtube feed there have been big names and interesting presentations. Outright I recognize and remember Wadler, Jones, Elliot, Milewski.

I've learned important things about category theory through the video lectures Bartosz Milewski has in his youtube channel. It took me few tries before I figured it out and Bartosz' lectures helped a lot.

Now, lets defunctionalize a program like Bartosz did in the video. We're going to write a program in Purescript, here's the import declarations.

import Data.Either (Either(..))
import Data.Foldable (for_)
import Data.Generic.Rep (class Generic)
import Data.Generic.Rep.Show (genericShow)
import Data.Maybe (Maybe(..))
import Data.Tuple (Tuple(..))
import Effect (Effect)
import Effect.Ref (Ref, new, read, write)
import Prelude (class Show, Unit, bind, discard, pure, show, unit, (+), (-), (<>), (>>=))
import Web.DOM.Document as D
import Web.DOM.Element as E
import Web.DOM.HTMLCollection as C
import Web.DOM.Node as N
import Web.Event.Event (Event, EventType(..))
import Web.Event.EventTarget (addEventListener, eventListener)
import Web.HTML.HTMLDocument (HTMLDocument, toDocument)
import Web.HTML (window)
import Web.HTML.Window as W

I take an example only slightly different from what Bartosz used. Here's a fibonacci function, written in Purescript. I'm using Purescript because it's been designed to use Javascript as a backend and it's easy to present the finished program this way.

fibonacci :: Int -> Int
fibonacci 0 = 0
fibonacci 1 = 1
fibonacci n = fibonacci (n-2) + fibonacci (n-1)

Bartosz perhaps picked a better example. Now you'll keep thinking about how there's a better algorithm to evaluating fibonacci than the recursive one presented here.

Like in Bartosz' example, we're recursively calling fibonacci and they are not tail calls. We'll transform the program such that they become tail calls.

fibonacci' :: Int -> Int
fibonacci' n = fibonacciCont n (\v -> v)

fibonacciCont :: forall k. Int -> (Int -> k) -> k
fibonacciCont 0 k = k 0
fibonacciCont 1 k = k 1
fibonacciCont n k
    = fibonacciCont (n-2) (\v ->
        fibonacciCont (n-1) (\w ->
            k (v + w)))

He additionally split it into functions like this:

fibonacci'' :: Int -> Int
fibonacci'' n = fibonacciCont' n fibonacciP0

fibonacciCont' :: forall k. Int -> (Int -> k) -> k
fibonacciCont' 0 k = k 0
fibonacciCont' 1 k = k 1
fibonacciCont' n k
    = fibonacciCont' (n-2) (fibonacciP1 {n, k})

fibonacciP0 :: forall k. k -> k
fibonacciP0 x = x

fibonacciP1 :: forall k. {n :: Int, k :: (Int -> k)} -> Int -> k
fibonacciP1 {n, k} v = fibonacciCont' (n-1)
                       (fibonacciP2 {v, k})

fibonacciP2 :: forall k. {v :: Int, k :: (Int -> k)} -> Int -> k
fibonacciP2 {v,k} w = k (v + w)

The split into functions allow determining what each function closes into itself. It allows to defunctionalize the whole thing.

data Kont = Done
          | Next {n :: Int, k :: Kont}
          | Conc {v :: Int, k :: Kont}

fibonacci''' :: Int -> Int
fibonacci''' n = fibonacciCont'' n Done

fibonacciCont'' :: Int -> Kont -> Int
fibonacciCont'' 0 k = apply k 0
fibonacciCont'' 1 k = apply k 1
fibonacciCont'' n k = fibonacciCont'' (n-2) (Next {n,k})

apply :: Kont -> Int -> Int
apply Done s = s
apply (Next {n, k}) v = fibonacciCont'' (n-1) (Conc {v,k})
apply (Conc {v, k}) w = apply k (v + w)

Now, what can you do with this? Well functions are generally a bit hard to serialize and this transformation turns them into simple data structures. Lets first derive the show instance, I did find it difficult to locate how to use Purescript's generics. The up-to-date documentation seem to be located in the generics-rep. The pursuit links to the repository of each module.

derive instance genericKont :: Generic Kont _

instance showKont :: Show Kont where
    show x = genericShow x

Now we're going to implement the fibonacci for the fourth time.

fibonacci'''' :: Int -> (Tuple Kont Int)
fibonacci'''' n = fibonacciCont''' n Done

fibonacciCont''' :: Int -> Kont -> Tuple Kont Int
fibonacciCont''' 0 k = Tuple k 0
fibonacciCont''' 1 k = Tuple k 1
fibonacciCont''' n k = fibonacciCont''' (n-2) (Next {n,k})

apply' :: Kont -> Int -> Either Int (Tuple Kont Int)
apply' Done s = Left s
apply' (Next {n, k}) v = Right (fibonacciCont''' (n-1) (Conc {v,k}))
apply' (Conc {v, k}) w = Right (Tuple k (v + w))

That's it. The only thing we need is bit of a driver code to show the results. I'd guess there's a better way to do this in Purescript, but I haven't learned it yet.

main :: Effect Unit
main = do
  w <- window
  d <- W.document w
  loadEvent <- eventListener (onLoadEvent d)
  addEventListener (EventType "load") loadEvent false (W.toEventTarget w)

onLoadEvent :: forall a. HTMLDocument -> a -> Effect Unit
onLoadEvent d e = do
  dialogs <- D.getElementsByClassName "fibonacci-dialog" (toDocument d) >>= C.toArray
  for_ dialogs setupDialog

setupDialog :: E.Element -> Effect Unit
setupDialog element = do
    maybe_output <- E.getElementsByClassName "output" element >>= C.item 0
    maybe_button <- E.getElementsByClassName "next" element >>= C.item 0
    case Tuple maybe_output maybe_button of
        Tuple (Just output) (Just button) -> do
            let setTexts a b = do
                                  N.setTextContent a (E.toNode output)
                                  N.setTextContent b (E.toNode button)
            setTexts "Program output goes here" "fibonacci 5"
            buttonRef <- new ((Tuple 5 Nothing) :: Tuple Int (Maybe (Either Int (Tuple Kont Int))))
            clickEvent <- eventListener (onButtonEvent buttonRef setTexts)
            addEventListener (EventType "click") clickEvent true (E.toEventTarget button)

            pure unit
        Tuple _ _ -> pure unit

onButtonEvent :: Ref (Tuple Int (Maybe (Either Int (Tuple Kont Int))))
              -> (String -> String -> Effect Unit) -> Event -> Effect Unit
onButtonEvent var update _ = do
    Tuple n state <- read var
    case state of
        Nothing -> do
            let res = fibonacci'''' n
            write (Tuple n (Just (Right res))) var
            update (show res) "next"
        Just (Right (Tuple k i)) -> do
            let res = apply' k i
            case res of
                Left _ -> do
                    let n' = n + 1
                    update (show res) ("fibonacci " <> show n')
                    write (Tuple n' (Just res)) var
                Right _ -> do
                    update (show res) "next"
                    write (Tuple n (Just res)) var
        Just (Left i) -> do
            let res = fibonacci'''' n
            write (Tuple n (Just (Right res))) var
            update (show res) "next"

2020-08-24 update: It may be Miles Frain just hinted how to do this better by using the signals library. He's got an example of it in purescript-cookbook.

So... that's it. Now we can show the intermediate computation steps of the fibonacci function.

Program output goes here

Here's the defunctionalized fibonacci again to illustrate that it remains to be a fairly small program.

data Kont = Done
          | Next {n :: Int, k :: Kont}
          | Conc {v :: Int, k :: Kont}

fibonacci'''' :: Int -> (Tuple Kont Int)
fibonacci'''' n = fibonacciCont''' n Done

fibonacciCont''' :: Int -> Kont -> Tuple Kont Int
fibonacciCont''' 0 k = Tuple k 0
fibonacciCont''' 1 k = Tuple k 1
fibonacciCont''' n k = fibonacciCont''' (n-2) (Next {n,k})

apply' :: Kont -> Int -> Either Int (Tuple Kont Int)
apply' Done s = Left s
apply' (Next {n, k}) v = Right (fibonacciCont''' (n-1) (Conc {v,k}))
apply' (Conc {v, k}) w = Right (Tuple k (v + w))

Here's the source code: fibonacci.zip

I got quite few questions about this method.

Could this kind of defunctionalization be automated?
Can you build a typeclass around the concept and vary how the program returns a value?
Can it be used for implementing memoization/dynamic programming algorithms?

On the typeclass, I guess it'd be something like:

class Defunctionalized k a r | k -> a where
  ap :: k -> a -> r

instance direct  :: Defunctionalized Kont Int Int
instance stepped :: Defunctionalized Kont Int (Tuple (Maybe Kont) Int)

2020-08-24 update: After reading the post, HakimSquirrel @ slack told me that he found a different post a bit easier to understand about defunctionalization itself. It can be found in James Koppel's blog, titled "The Best Refactoring You've Never Heard of". He kept his Compose 2019 talk from defunctionalization as well and the post is a transcript of that talk. It's a really long post but I think it may be helpful already simply because James explains what defunctionalization is. I also like the Hacker News example he gives, as an example of where he thinks defunctionalization can be useful.

Bartosz Milewski's neat defunctionalization talk

Similar posts