Tobin Yehle

A Sunrise Alarm Clock

Mon, 08 Jun 2020 00:00:00 +0000

I built a sunrise alarm clock so I wouldn’t need to wake up in the dark during the winter time when the sun rises very late. This was the first PCB I had ever made and I had a lot of fun with it. The hardware and firmware is all on github.

The completed project

The board is powered by an Atmega328 and has, in no particular order:

Three power transistors to control the RGB channels of an LED strip
A power regulator so it can be powered by a 12V wall wart
A real time clock with a battery backup to keep track of the time
No surface mount parts because I wasn’t sure I could solder them
Dope sunrise art

Hardware Development
The Firmware
Wishlist

Hardware Development

In the Beginning

I started this project with a cheap clone of the Arduino pro mini driving an LED strip, and the first time I left it on overnight it burned out the power regulator. This probably would not have happened to a legit Arduino pro mini, but the cheap clones have to make compromises in the design to lower costs, and the power regulator is a good way to save on component costs.

I was also keeping track of the time using the Atmega’s clock, but this isn’t the best approach because if the board lost power for any reason then it would lose the time. Maybe for a microwave that is ok, but not for an alarm clock (Why don’t microwaves save the time? An RTC only costs a few dollars).

Test circuit from the datasheet of the L4931 voltage regulator

My second attempt used a real time clock (RTC) module I bought from Adafruit, and a better power regulator. Figuring out how to use this power regulator was straightforward because its datasheet has an example circuit that was exactly what I was trying to do. All this assembled on a breadboard worked great! The LEDs powered up at the scheduled time, but the barrel jack was really easy to accidentally pull out of the bread board, so I decided I needed to solder some components into a circuit board.

The Circuit Board Design

My first thought was that I should make a board with the barrel jack, output transistors, better power regulation, and sockets for the Arduino and the RTC module. The RTC module communicates over I²C, but the position of the I²C lines on the Arduino aren’t on the standard set of pins that fit in a breadboard, so any Arduino I could plug in wouldn’t be compatible with a breadboard anymore.

It was at this point I realized there were only 15 components on the Arduino board and I was going to need my own copies of 4 of them for power regulation. The schematic is also open source, so I thought it would be more fun to build my own PCB with the components I needed from the Arduino and RTC modules as well as all the output and power regulation circuitry.

I started with the Arduino pro mini schematic and made changes and additions until it had all the things I needed. I knew I wanted to use the L4931 regulator and a barrel connector to provide power to the board, so I subbed those into the power regulation section of the Arduino’s schematic. The stabilization caps on the Arduino design are larger than those suggested by the L4931 datasheet, so I used the larger ones just to be safe. I have since learned that the smaller ceramic capacitor is used to smooth out higher frequency variations in the output voltage. It has a much lower equivalent series resistance (ESR) than an electrolytic capacitor so it can respond faster, but large ceramic caps are expensive, so the electrolytic capacitor smooths out any low frequency noise on the cheap and the ceramic cap takes care of the high frequency noise. There is also an optional power indication LED I chose not to solder because I don’t like lights I can’t turn off on things in my bedroom.

Power regulation with an electrolytic and ceramic cap to filter output voltage

The design of the RTC breakout from Adafruit is not open source, but the datasheet for the DS1307 is excellent, and has an example on the very first page! I sure do appreciate good documentation. There is one additional component on the breakout board I got from Adafruit, but I think it is just a power stabilization capacitor. This clock will keep the time using the backup battery if main power is lost, so if I unplug the board, or the power goes out I don’t need to reset the clock. The two pull-up resistors on the SDA and SCL lines are required by the I²C spec. The 32.768kHz crystal is a requirement of the chip. This is a very popular speed for clocks because it makes counting seconds so easy because 32,768 is 2¹⁵.

DS1307 Real Time Clock

I replicated the setup I had on the original breadboard design for the high power output transistors. The gates are connected to PWM capable outputs on the microcontroller and provide a high capacity current sink for the LED strip when on. Most of the strands I’ve seen are powered off of 12V and use a common high voltage line, so this setup should be compatible with many LED strips. The pull down resistors on the outputs here prevent the LED strip from momentarily turning on when the microcontroller is reset. The Atmega328 sets all pins to input without pull up resistors when it is reset so the voltage at the mosfet gate would be left floating, potentially allowing it to turn on.

High current output transistors for the RGB channels of the LED strip

The rest of the schematic is the same as the Arduino board I was using. There is another power stabilization cap drawn very close to the microcontroller. In this case this means that the cap should be physically close to the microcontroller on the board to provide a very stable voltage for it to use. The microcontroller will probably use different amounts of current at high frequency, so this makes sense to me.

The other interesting bit of this part of the schematic is the reset hardware. There is a cap between the DTR pin on the programming header and the reset pin on the microcontroller with a pull-up resistor. This is so that if software fails to pull the DTR pin high again the board will still reboot after charging the capacitor through the pull-up resistor. The values of the capacitor and resistor here have been tuned so the board is only momentarily off and can talk to the computer properly even if there are bad drivers for the programming cable.

GPIO and reset for the Atmega328

Circuit Board Layout

I decided to use only through hole parts because I wasn’t sure I had the skill to solder surface mount components. I also wanted to make sure I could debug and expand the design, so I wanted all the standard Arduino pins to be accessible if I plugged it into a bread board. All this means that the board is very skinny so it would fit on a bread board and in order for there to be any room left on the bread board it has to hang way off the end.

Top of the board with many tiny suns

Bottom of the board with one huge rising sun

The Firmware

Color Profile

When I first started thinking about a sunrise alarm clock I thought it would be neat to have it run through the colors of the actual sky during a sunrise. Surely somebody has collected such data. As I tried to find a dataset I discovered that “the sky at sunrise” isn’t a very precise idea. What angle above the horizon? East, West, straight up, or some other azimuth? A single point or the average color over an area? Are there clouds? If so where? Is there haze in the air?

I decided this was too much and I should instead choose a color profile that seemed natural without any relation to The Real Sky. I settled on the color of a perfect black body radiator because it goes from red to “white” which seemed nice, and there is a nice physical metaphor with the light source warming up for the day. I got some data about hue for different temperatures here. The colors in this dataset were converted using a D65 whitepoint, which is why I think there is a kink in the plot around 6500K. I played around with this data a bit and decided I liked the look of a ramp from 1000K up to 4500K. This leaves the lights a little more yellow than a true white which I find nice in the morning.

Color channel values for different temperature blackbody radiators using a D65 whitepoint

The intensity of a true blackbody goes up way too fast to be a satisfying light, and since our eyes respond exponentially, I thought an exponential brightness curve would look nice. The lights need to start off and end at fully bright over some period of time, so the brightness function $B(t)$ will need satisfy $B(0)=0$ and $B(1)=1$. By shifting the exponential function down by one and then using horizontal and vertical stretching I can ensure that our needs are met while still leaving a parameter free for the steepness of the curve. The exponential function with a vertical shift of -1, a vertical stretch of $k$, and a horizontal stretch of $a$ is $B_k(t) = k \left( e^{a t} - 1 \right)$. Solving $B_k(1) = 1$ to gives $a = \ln\left(\frac{1}{k} + 1\right)$, leaving $k$ as the single steepness parameter. The full brightness function is then

\[B_k(t) = k \left( \exp \left( t \ln \left( \frac{1}{k} + 1 \right) \right) - 1 \right)\]

The brightness function $B_k(t)$ for some values of $k$. For the lights I have, $k=\frac{1}{8}$ looks the best.

Normalization

The blue LEDs on the strip I have seemed much brighter than the green or red channels, so I knew I needed to tone them down, but by how much? To determine this I took a picture of a white sheet of paper lit by the LEDs at full power and then looked at a histogram of each color channel.

LEDs at full power reflecting off a white sheet of paper

Normalizing the green and blue channels to the max output of the red LEDs means that “full power” for the green channel should be $\frac{87}{96} \cdot 255 = 231$, and for the blue channel $\frac{87}{179} \cdot 255 = 124$.

High precision PWM

The difference between the lights being off, and being on at 1/255 power is pretty jarring. At higher powers the steps become less noticeable, but the first few steps are very clearly visible. Fortunately, there is a 16 bit clock on the Atmega328 and it can control the PWM output on pins 9 and 10, providing steps too small to notice. I didn’t know this when I fabricated the boards, so I got really lucky that one of the LED channels is wired to pin 9! Unfortunately, the Arduino environment doesn’t make it easy to access this functionality (it looks like there is some work to make that better, but not for AVR based boards), so I had to directly set values of the control registers. This stack overflow answer got me started nicely and I filled in the rest of the code by replicating the Arduino source and the taking a look at the relevant parts of the Atmega328 datasheet.

void setupPWM16() {
  TCCR1A = _BV(COM1A1) | _BV(COM1B1)  /* non-inverting PWM */
    | _BV(WGM11);                     /* mode 14: fast PWM, TOP=ICR1 */
  TCCR1B = _BV(WGM13) | _BV(WGM12)    /* other bits for mode 14 */
    | _BV(CS10);                      /* no prescaling */
  ICR1 = 0xffff;                      /* TOP counter value */
}

This setup function sets a bunch of flags that control the output of pins 9 and 10. These pins are internally called PB1 and PB2, and can optionally be driven by output compare 1A (OC1A) and output compare 1B (OC1B). The 1A and 1B suffix is used throughout to refer to these two outputs. WGM is the Waveform Generation Mode, and it is being set to mode 14 which is fast PWM with the maximum counter value controlled by ICR1. We want the steps to be as small as possible, so we will use the full 16 bits of the counter by setting the max value to 0xffff. CS is the Clock Select settings, and it is set so there is no prescaling on the timer because there is no reason to not have the PWM cycles happen as fast as possible.

void analogWrite16(uint8_t pin, uint16_t val)
{
  pinMode(pin, OUTPUT);

  // PWM outputs cannot be fully on or off, so use digitalWrite instead
  if(val == 0) {
    digitalWrite(pin, LOW);
  } else if(val == 65535) {
    digitalWrite(pin, HIGH);
  } else {
    // Calls to digitalWrite overwrite some things, so make sure to reset them
    switch (pin) {
      case 9:
        _SFR_BYTE(TCCR1A) |= _BV(COM1A1);
        ICR1 = 0xffff;
        OCR1A = val;
        break;

      case 10:
        _SFR_BYTE(TCCR1A) |= _BV(COM1B1);
        ICR1 = 0xffff;
        OCR1B = val;
        break;

      default:
        return;
    }
  }
}

This is a 16 bit version of the standard Arduino analogWrite function. Setting OCR1A to 0 won’t turn the PWM all the way off because of the way it is implemented internally, so in order to turn off the pin we need to switch it to digital mode. When we do that we end up stomping on the value in ICR1 and TCCR1A, so we need to set those if we want to update the PWM output.

void writePWM(uint8_t pin, double value) {
  if(pin != 9 && pin != 10) {
    analogWrite(pin, ratioToRange(value, 256));
  } else {
    analogWrite16(pin, ratioToRange(value, 65536));
  }
}

I didn’t want to keep track of which output channel was on the high resolution PWM, so I wrote this wrapper function that writes the proper PWM value based to the pin based on a floating point value from 0 to 1.

With all this setup, the red channel, which comes on first, can use the high resolution timer, which makes everything look much better. There is still a noticeable step when the green light goes from off to its lowest brightness, but there isn’t anything I can do about that without redesigning the board.

Wishlist

After using this for a while these are some features I wish it had:

A better UI. Setting the wake up time or resetting the clock requires a firmware flash which … uhh, encourages you to have a healthy regular sleep schedule 🙃 A single button that could reset the time, or do a snooze, or turn off the lights would be nice.
Mounting holes on the board. There isn’t a great way to put this thing into a case if I had one.
A case. It is currently just sitting on my floor. Not a great spot for electronics to stay safe.
Better interfacing with the LED strip. Terminal blocks, or some other type of convenient connector.
High res PWM on the other channels. There is a pretty noticeable change when the green channel comes on, so it would be great to use a high res timer for that as well. To get high res timers for all three channels I would have to use something other than the Atmega328, which would probably be a good idea anyway.

speakdata

Mon, 04 May 2020 00:00:00 +0000

Have you ever looked for a quick and easy way to say 84:a3:fe:67:c5:01, or maybe 192.168.0.1?

Look no further than speakdata, a canonical way to translate bytes into easy pronounceable syllables!

>>> import speakdata
>>>
>>> speakdata.pronounce(b'\x84\xa3\xfe\x67\xc5\x01')
'Lifufing torleyking sijabing.'
>>>
>>> speakdata.pronounce([192, 168, 0, 1])
'Saful bababing.'

Are you now wonding how to pronounce “Saful bababing?” Just pipe to say

$ say "Saful bababing"

See its easy to pronounce!

It can also create beautiful poetry:

$ curl https://tobin.yehle.us/favicon.ico -s | gzip | speakdata
Bortaning bubab taikaling naibowt babafing kaseyb kabaj jalajing.
Bijak baifailing nabubing sufaib fafaik safaib nababing sibaib.
Fajak bisak luboil beybal faibas jeynub beytaib bowkub.
Baitaij saisaib boitus borkaij jorjeyb baifaib najab nainowking.
Fabaib jabaib fitib teysab baiforting nisoifing torlab tainoifing.
Koilub towtaf bakoib beyteyt bileys taneyjing jorkuling tortowking.
Koibowt teynorting soifiling faiseyb borbowt kibaiking lownorling forkoifing.
Tuteybing kinab seyfaling jowkowt biforjing joibowting sorbafing sowsaifing.
Lunorsing lusoiking joikorning torteyling saisuking faifaik kiloiting boifafing.
Torfak susas lorsus fasas towkowj nownorting tafibing torjoil.
Soijab seykafing noifik leyforb saleyking forsas fowsait foilaif.
Faikiking jeyforning jukaling kusors towlifing kusut forlorjing lujoijing.
Lorlorb taituning fituk boisoin teysoiking toibowling buniking nowtas.
Laikuf nutorsing faforning kaborf seyfeyj lukoijing fortuking fukaf.
Foiforning feyteyk lowfufing seyfoiting towtorl bownois towloining finun.
Kanowl boijain tutub norbabing jaisafing fulorting lufoijing tortowj.
Nijufing kaforbing nitaiting bortaifing sowkaj bowbown saneyjing jakork.
Fakowning norsorb towjais kaibowt bitailing lakowl seyforf fuborn.
Lornaiking kijoiting nailorl jorlorning bifows boikowfing lownoiking nakeyl.
Bitijing feytowl sulik kowforning jajors koiborbing lataining korlowl.
Neylaking seyjaning julaking tubors kowborfing bufowling boitorfing forsus.
Korforb seyjating failiting tilowting jaiforn siteyking joinorj nowfajing.
Lornaik jaforbing norsorsing soisuk laineyf borsaij fataking forjais.
Foikais joikoib kowbeyl borluling lajoifing torluling jukorting sufeyling.
Boilowling foijul toinibing saisubing faikeyling lortowf tailaking bortuk.
Jitorling jisuning teytibing susijing koifowl foijaising kutowl seylaiting.
Junorsing loisating baisowning torsowfing sainort sakak jujeysing tortab.
Faifoiling toijoiting kaneysing sitowf tornoining lafut loitit toijowting.
Noijab koilafing borjis toijorb jeyleyfing jaifowting batowj bisaf.
Neyjaibing feytib kisafing borkeysing korlab tunaking boijeyb kowsorfing.
Bowfas baloib babab lowkeyjing nitifing forlibing babab.

How does it work?

The system uses a one-to-one mapping with 8 consonants and 8 vowels to build words. Each two byte chunk of the input forms a single word with three consonants and two vowels in the output. These words are further grouped into sentences of 8 words.

But the sequence of vowels and consonants only encodes 15 bits of data! The last bit is encoded with the optional suffix “–ing”. If the last bit of the chunk is a 1 then the word ends with –ing, otherwise it is left as is.

 c    v    c    v    c  ing?
...  ...  ...  ...  ...  .

The list of consonants and vowels were carefully chosen with help from Gavin Yehle to be far away from each other (eg: only one of f/v and s/z are present). The vowels are:

["a", "i", "u", "oi", "ai", "ow", "ey", "or"]

and the consonants are

["b", "f", "j", "k", "l", "n", "s", "t"]

If there are an odd number of bytes in the input then the last 16 bit chunk is padded with a zero byte.

Example

To convert 0x60ff:

Start with the bits.
0110 0000 1111 1111

Break into 3 bit chunks.
011 000 001 111 111 1

Look up sylables.
"k" "a" "f" "or" "t" 1

Set the ending.
kaforting

Cavort with joy!

Algebraic Data Types in the Lambda Calculus

Sun, 08 Oct 2017 00:00:00 +0000

As I was writing a compiler for the lambda calculus, I noticed some similarities between booleans and lists. They both have two variants, and in order to use a boolean or a list value you must apply it two to functions that say what to do in either case. Like this

(some-bool
  value-if-true
  value-if-false)

(some-list
  (λ (fst rst) (do-if-cons fst rst))
  value-if-nil)

There is only one way to construct boolean and list values you can use this way. Everything must be a closure, so booleans and lists are closures that expect two arguments.

(define (Cons fst rst) (λ (do-cons nil-val) (do-cons fst rst)))
(define Empty          (λ (do-cons nil-val) nil-val))

(define True  (λ (t f) t))
(define False (λ (t f) f))

Using a function in this way seemed very similar to the case or match syntax in many languages. I wondered if I could create some syntactic sugar for algebraic data types in general for which booleans and lists would be special cases.

For example a struct definition

(struct Bool
  [True]
  [False])

(struct (List a)
  [Cons a (List a)]
  [Empty])

with a corresponding case keyword for destructuring

(case some-bool
  [True value-if-true]
  [False value-if-false])

(case some-list
  [(Cons fst rst) (do-if-cons fst rst)]
  [Empty value-if-nil])

To make this work the struct keyword will need to introduce constructors for the variants, and the case keyword will need allow us to destructure an instance of a struct.

Case Keyword

Lets start by looking at an example of the transform. The following function will find the depth of a tree structure.

(define (depth tree)
  (case tree
    [Empty 0]
    [(Leaf data) 1]
    [(Tree left right) (+ 1 (max (depth left) (depth right)))]))

Just like the boolean and list examples, it should desugar into an application of closure named tree that is passed into the depth function.

(define (depth tree)
  (tree
    0
    [λ (data) 1]
    [λ (left right) (+ 1 (max (depth left) (depth right)))]))

The bindings in the patterns become lambda parameters, and the expression becomes the body of the lambda. If the pattern we are matching against is a constant like Empty then there are no bound variables, and the expression is applied directly. It could also become an application of a lambda with no parameters, but the evaluation model I’m using is lazy, so there is no need to wrap the value in an additional lambda. ie: we just have 0 instead of [λ () 0].

The only complication for doing this transform in general is the ordering of the case clauses. We should have been able to order the clauses Tree, Leaf, Empty, instead of Empty, Leaf, Tree, and it should have desugared to the same code. To ensure this happens we need to apply the clauses in the same order they appear in the struct definition.

In general the statement

(case <exp>
  [<C2>                      <exp2>]
  ...
  [(<Cm> <a1> <a2> ... <an>) <expm>]
  ...
  [(<C1> <a1> <a2> ... <an>) <exp1>])

desugars to

(<exp>
  (λ (<a1> <a2> ... <an>) <exp1>)
  <exp2>
  ...
  (λ (<a1> <a2> ... <an>) <expm>))

Struct keyword

The definition of the tree structure we used before would look like

(struct (Tree a)
  [Empty]
  [Leaf a]
  [Node (Tree a) (Tree a)])

The compiler needs to keep track of the order of the variants in order to desugar any case statements correctly, but it also needs to generate a constructor for each variant. This tree structure has three variants, so the constructors would be:

Empty :: Tree a
Leaf :: a -> Tree a
Node :: Tree a -> Tree a -> Tree a.

We know that the constructors need to generate a closure that expects three arguments, one for each variant. Each of those arguments will be a function that operates on the data stored in that variant.

(define Empty             (λ empty leaf node) empty)
(define (Leaf data)       (λ empty leaf node) (leaf data))
(define (Node left right) (λ empty leaf node) (node left right))

But why did we pick the variable names left, right, and data? We have some intuition about the meaning of these variables, but the desugarer will need to generate these variables names. The lambda parameters also need to be generated, and we could reuse the name of the variant, but this may cause confusion to a human, so I will mangle them. With these changes the definitions would look like this

(define Empty        (λ _Empty _Leaf _Node) _Empty)
(define (Leaf a1)    (λ _Empty _Leaf _Node) (_Leaf a1))
(define (Node a1 a2) (λ _Empty _Leaf _Node) (_Node a1 a2))

In general the statement

(struct <Name>
  [<C1> <T1> <T2> ... <Tn>]
  [<C2>]
  ...
  [<Cm> <T1> <T2> ... <Tn>])

desugars to

(define (<C1> a1 a2 ... an) (λ (<_C1> <_C2> ... <_Cm>) (<_C1> a1 a2 ... an) ))
(define <C2>                (λ (<_C1> <_C2> ... <_Cm>) <_C2>))
...
(define (<Cm> a1 a2 ... an) (λ (<_C1> <_C2> ... <_Cm>) (<_Cm> a1 a2 ... an) ))

Examples

I started this project with the goal of generalizing the syntax of booleans and lists to something that will work for any algebraic data type.

Booleans

It is easy to see that the simple example of an if statement will desugar the same way as the standard encoding for Church booleans.

(struct Bool
  [True]
  [False])

(define (if c t f)
  (case c
    [True  t]
    [False f]))

becomes

(define True  (λ (_True_0 _False_0) _True_0))
(define False (λ (_True_0 _False_0) _False_0))

(define (if c t f)
  (c t f))

Lists

The isNil function compiles to exactly the same lambda expression it did before.

(struct (List a)
  [Empty]
  [Cons a (List a)])

(define (isNil list)
  (case list
    [Empty True]
    [(Cons _ _) False]))

becomes

(define Empty        (λ (_Cons _Empty) _Empty))
(define (Cons a1 a2) (λ (_Cons _Empty) (_Cons a1 a2)))

(define (isNil list)
  (list True (λ (_ _) False)))

Maybe

It is now a breeze to define other well known functions like the safe version of head.

(struct (Maybe a)
  [Just a]
  [Nothing])

(struct (List a)
  [Empty]
  [Cons a (List a)])

(define (head? list)
  (case list
    [Empty             Nothing]
    [(Cons first rest) (Just first)]))

becomes

(define (Just a1) (λ (_Just _Nothing) (_Just a1)))
(define Nothing   (λ (_Just _Nothing) _Nothing))

(define Empty        (λ (_Cons _Empty) _Empty))
(define (Cons a1 a2) (λ (_Cons _Empty) (_Cons a1 a2)))

(define (head? list)
  (list
    Nothing
    [λ (first rest) (Just first)]))

Super Key on Dell XPS Developer Edition

Sun, 05 Mar 2017 00:00:00 +0000

The windows/super key on my new laptop was disabled by the manufacturer.

This removes the default behavior of left maximize and right maximize and the search function.

Dell has to include a package to disable the super key called dell-super-key because of licensing for a button that has the windows logo on it.

I removed this package with

sudo apt remove dell-super-key

and then enabled its default behavior with the compiz config manager.

Desktop → Ubuntu Unity Plugin → Launcher → Key to show the Dash ⇒ <Super>
Window management → Grid → Bindings → Left/Right Maximize ⇒ <Control><Super>Left/Right

And now my laptop works as it should. Why Dell doesn’t just print a different logo on that key is beyond me.

Automatic Table of Contents in Jekyll

Sun, 26 Feb 2017 00:00:00 +0000

Edit from 2020:

Kramdown now makes this easy! Checkout these gitlab docs for a modern way to do this.

I recently wrote a very long post and I wanted a table of contents to appear at the top.

I also came up with a good test for the length of a post; if a table of contents would be helpful, its too long.

All the code for this project is on github. There is also an example page.

Results

Kramdown adds an id attribute to all headers automatically. The table of contents contains links to these headers.

The most basic interface is:

{% include toc.html %}

The user should also be able to specify which tags to link to and to manually ignore specific tags. The include takes arguments to achieve this.

{% include toc.html tags="H2,H3" ignore="#sections" %}

Details

The html include loads jQuery from Google’s CDN if it can otherwise it falls back to the local copy. I do this because the CDN is probably faster to load, but if it fails, then the page will still load. I like this because one of the big advantages of Jekyll is offline development. I shamelessly stole this idea from this question.

The html include file uses liquid tags translate the arguments into JavaScript.

<ul id="_toc">
  <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.1.1/jquery.slim.min.js"></script>
  <script type="text/javascript">
    if(!window.jQuery) {
      document.write('<script src="{{ "/js/jquery-3.1.1.slim.min.js" | prepend: site.baseurl }}"><\/script>')
    }
  </script>
  <script src='{{ "/js/toc.js" | prepend: site.baseurl }}'></script>
  <script>
    let tags = {%if include.tags%}"{{include.tags}}"{%else%}"H1,H2,H3,H4"{%endif%};
    let ignore = {%if include.ignore%}"{{include.ignore}}"{%else%}''{%endif%};
    $(document).ready(function () {
      toc($("#_toc"), tags, ignore);
    });
  </script>
</ul>

The id of the <ul> to put the table in is _toc because, as far as I can tell, Kramdown will not generate that id automatically. If the document contained a header called ToC, and it appeared before the include call, then the entire table would be inserted into the header tag. Not pretty. This is still technically possible if someone manually set an id to _toc anywhere before the include call.

JavaScript

The first step is to get the headers and what level those headers should be displayed at. The order of the input tags is used to determine the level.

function getHeaders(tagNames, ignore) {
  let tags = tagNames.split(",");
  return $(tagNames.toString()).filter("[id]").not(ignore).map(function() {
    return {
      level: tags.indexOf($(this).prop("tagName")),
      link: "#" + $(this).attr("id"),
      text: $(this).text()
    };
  });
}

Next I used a linked list of buffers to build a nested array header information.

function nestHeaders(headers) {
  let base = [];
  let buffer = {level: 0, list: base, prev: undefined};
  headers.each(function(_, header) {
    if(header.level >= 0) {
      while(header.level > buffer.level) {
        let newBuffer = {level: buffer.level + 1, list: [], prev: buffer};
        buffer = newBuffer;
      }

      while(header.level < buffer.level) {
        buffer.prev.list.push(buffer.list);
        buffer = buffer.prev;
      }

      buffer.list.push(header);
    }
  });

  while(buffer.prev != undefined) {
    buffer.prev.list.push(buffer.list);
    buffer = buffer.prev;
  }

  return base;
}

Using that nested array, I can recursively build the table off of a given base ul element. On each call it either generates an li element for that header or a ul element that is filled by a recursive call to genToc.

function genToc(base, nestedHeaders) {
  nestedHeaders.forEach(function(elem) {
    if(elem instanceof Array) {
      let ul = $("<ul/>");
      genToc(ul, elem);
      base.append(ul);
    } else {
      base.append($("<li/>").append($("<a/>", {
        href: elem.link,
        text: elem.text
      })));
    }
  });
}

These pieces are combined in the function toc which is called from the included html.

function toc(base, tagNames, ignore) {
  genToc(base, nestHeaders(getHeaders(tagNames, ignore)));
}

Compiling and Interpreting the Lambda Calculus in Haskell

Mon, 23 Jan 2017 00:00:00 +0000

This whole project is based on this post by Matt Might. I decided I would try to implement it in Haskell instead of Racket. There are several steps of the process that are missed because they are built into the Racket language. In my implementation I needed a parser, an interpreter, and an extractor, in addition to a reimplementation of the compiler.

Instead of defining all the functions in the language in the compiler, I thought it would be more fun to write them in the language itself. This greatly reduces the complexity of the parser and the compiler. Also it means you can add any library functions you see fit, and I ended up adding some comparisons, and list operations.

All code for the project can be found on github.

Sections

Parser
Base Definitions
Compiler
Interpreter
Tying it all Together
- The Pieces
- The Interface

Parser

Because many functions are implemented in the language itself, the grammar quite small.

<program> ::= <define> ... <exp>

<module>  ::= <define> ...

<define>  ::= (define <var> <exp>)
           |  (define (<var> <arg> ...) <exp>)

<exp>     ::= <var>

           |  <nat>

           |  (<lam> (<arg> ...) <exp>)
           |  (let ((<var> <exp>) ...) <exp>)
           |  (letrec (<var> <exp>) <exp>)

           |  (<exp> <exp> ...)

<arg>     ::= _ | <var>

<lam>     ::= λ | lambda

<arg> is its own form because I like using _ as an ignored argument.

I added program, module and define forms to make it possible to do bootstrapping.

I used parsec to get the grammar into this data structure.

data Program = Program [Definition] Exp deriving (Show)

data Definition = Def String Exp deriving (Show)

data Exp = Var String

         | Num Int

         | Lambda [String] Exp
         | Let [(String, Exp)] Exp
         | Letrec String Exp Exp

         | Application Exp Exp
         deriving (Show)

I also wrote a very basic “lexer” to remove comments.

This isn’t a parsec tutorial, so I won’t bore you with details of the parser. The interface to the parser is

parseProgram :: String -> String -> Either String Program
parseProgram = lexParse programP

parseModule :: String -> String -> Either String [Definition]
parseModule = lexParse moduleP

lexParse :: Parser a -> String -> String -> Either String a
lexParse p name input = first show $ parse commentP name input >>= parse p name

I had never used bifunctors before, but it sure made handling Either easier.

This will result in a Left with the shown version of the ParseError if the parser fails. Throughout the project I used Left String to represent error.

Every stage of the compiler has the type a -> Either String b so it can all be strung together with bind.

Base Definitions

The bootstrapped definitions of basic functions in the language are in resources/base.lc.

All the desugarings implemented in Matt Might’s post that are missing from the compiler are here.

There are also some new definitions that make actually writing programs a bit easier.

Useful Functions

(define (id x) x)
(define (const x _) x)
(define hang
  ((λ (u) (u u)) (λ (u) (u u))))

const will work the same as in Haskell because the interpreter is lazy. (const x hang) won’t evaluate the second argument, so it won’t hang.

hang is used to kill the program if head or tail is called on an empty list. Might seem a bit draconian, but I didn’t implement errors or exceptions cause they’re hard. So too bad.

Booleans

All booleans are church encoded. Church booleans are functions that take two arguments; what to do if its true, and what to do if its false.

true = (λ (t f) t)
false = (λ (t f) f)

Since the interpreter is lazy if can be safely represented as a function because the branch that is not taken will never be evaluated.

(define (if c t f)
  (c t f))

All other boolean operations can be defined in a similar way, eg. not.

(define (not a)
  (a #f #t))

Church Numerals

All things are done using Church Numerals.

A church numeral is a function that takes two arguments: an operation to perform f and an object to operate on x. For the church numeral n, the operation is done n times to x.

1 = (λ (f x) (f x))
0 = (λ (f x) x)
3 = (λ (f x) (f (f (f x))))

Numeric Operations

The new numeric operations are defined as follows.

(define (even? n) (n not #t))

even? is its own function because a general remainder operator using repeated subtraction is (I think) less efficient. When the the number is evaluated it will apply not to #t n times.

You may have noticed this will take at least linear time as opposed to constant time one would expect from a real programming language. This is because church numerals are slow and should never be used unless you like office chair battles.

(define (/ n m)
  (letrec (div1 (λ (n1 m1)
                  (let ([diff (- n1 m1)])
                    (if (zero? diff)
                        0
                        (succ (div1 diff m1))))))
    (div1 (succ n) m)))

The division algorithm is repeated subtraction, and is described on wikipedia.

div1 is fast (because it only does one subtraction each iteration), but it is not quite the division operator we want.

(div1 2 3) = 0 as expected
(div1 4 3) = 1 as expected
(div1 3 3) = 0, but should be 1

To fix this problem we just need to add one to the input n. This fix would create problems with negative numbers, but there are no negative numbers in this language, so its all fine.

(define (mod n m)
  (letrec (mod1 (λ (n1 m1)
                  (let ([diff (- n1 m1)])
                    (if (zero? diff)
                        n1
                        (mod1 diff m1)))))
    (prev (mod1 (succ n) m))))

mod is the same algorithm as division, but instead of counting the number of iterations, we just return n1.

The fix is much the same; just add 1 to the input. Since its the final value of the input we are interested in we then subtract one from the result.

(define (<= a b)
  (zero? (- a b)))
(define (>= a b)
  (zero? (- b a)))
(define (< a b)
  (<= (+ 1 a) b))
(define (> a b)
  (>= a (+ 1 b)))

All the new comparison operations check if the result of a subtraction is zero.

Lists

Church encoded lists are similar to booleans. They are functions that take a thing to do if the list is a pair, and a result if the list is empty. If the list is a pair then the function is given two arguments, the head and the tail. If the list is empty then the result is passed straight through.

(define (cons h t)
  (λ (f _) (f h t)))
(define empty
  (λ (_ e) e))

This encoding will only work with a lazy interpreter. The argument e must only be evaluated if the list is empty.

Using these definitions we can define common list operations, some of which I have listed below.

(define (head l)
  (l (λ (h _) h) hang))
(define (tail l)
  (l (λ (_ t) t) hang))
(define (pair? l)
  (l (λ (_ _) #t) #f))
(define (null? l)
  (l (λ (_ _) #f) #t))

As I said earlier, head and tail will hang if called on the empty list because there isn’t anything else we could reasonably do.

pair? and null? never evaluate the contents of the list, so a program like (pair? (cons hang hang)) will still terminate.

(define (from n)
  (cons n (from (succ n))))

from builds an infinite list starting at the given value. Since the interpreter is lazy this is not a problem.

(define (take n l)
  (if (or (zero? n) (null? l))
      empty
      (cons (head l) (take (prev n) (tail l)))))

take results in a list containing the first n elements of the given list. It will return fewer than n items if the given list is not long enough.

(define (foldl fn acc xs)
  (if (null? xs)
      acc
      (foldl fn (fn acc (head xs)) (tail xs))))

(define (foldr fn acc xs)
  (if (null? xs)
      acc
      (fn (head xs) (foldr fn acc (tail xs)))))

Folds work here like they do in every other language.

The lazy interpreter allows the right fold to exit early.

A right fold that gets the first element if it exists, (foldr (λ (e _) #t) #f (from 0)), will terminate, but the left fold version will not.

(define (range low high)
  (if (>= low high)
      empty
      (cons low
            (range (+ 1 low)
                   high))))

This definition is just more efficient than (take (- high low) (from low)) because church numerals suck.

All the other language features must be implemented by the compiler.

Compiler

The purpose of the compiler is to transform the high level AST into one that represents the pure lambda calculus. The target structure is

data Node = Lam String Node
          | Ref String
          | App Node Node
          deriving (Show, Eq)

Scope and Free Variables

There are a bunch of places in the compiler and interpreter where knowing which variables will be referenced later is important.

Any variable that has no binding is considered free.

(λ (x) (f x)) – f is free and x is bound
(λ (y) (λ (x) (y x))) – x and y are both bound
((λ (x) x) x) – x is bound inside the lambda but free outside

The compiler needs to know which variables are free to determine which definitions can be dropped, and the interpreter needs to know which variables are free to handle closures properly.

Both the Node and Exp structures need a freeVars function, so they implement a typeclass.

class Scope a where
  freeVars :: a -> Set String

This allows a convenient function for guards in the compiler and interpreter.

isFree :: Scope e => String -> e -> Bool
isFree name expr = name `Set.member` freeVars expr

For the Node structure, Lam is the only way to bind a variable, and Ref is the only way to reference one.

instance Scope Node where
  freeVars (Lam arg body) = Set.delete arg $ freeVars body
  freeVars (Ref name) = Set.singleton name
  freeVars (App f x) = freeVars f `Set.union` freeVars x

For Exp, variables can be bound by Lambda, Let or Letrec, and referenced by Var.

instance Scope Exp where
  freeVars (Var name) = Set.singleton name
  freeVars (Lambda args body) = freeVars body `removeAll` args
  freeVars (Let bindings body) = foldr addBindingVars bodyVars bindings
    where
      addBindingVars = Set.union . freeVars . snd
      bodyVars = freeVars body `removeAll` map fst bindings
  freeVars (Letrec name binding body) = Set.delete name allVars
    where
      allVars = freeVars binding `Set.union` freeVars body
  freeVars (Num _) = Set.empty
  freeVars (Application f x) = freeVars f `Set.union` freeVars x

Define

The define forms desugar to a series of let or letrec expressions.

desugarDefs :: [Definition] -> Program -> Exp
desugarDefs baseDefs (Program defs expr) = desugarProgram newProg
  where
    newProg = Program (baseDefs ++ defs) expr

desugarProgram :: Program -> Exp
desugarProgram (Program ds e) = foldr defToLet e ds
  where
    defToLet (Def name expr) body = if isFree name body
                                    then if isFree name expr
                                         then Letrec name expr body
                                         else Let [(name, expr)] body
                                    else body

desugarDefs integrates the definitions of a module into those of a program.

desugarProgram only let-binds a definition if it is used, and then only does a letrec if the name being defined is free in its definition. Otherwise it desugars to the much simpler let form.

After this stage the program is represented as a single expression.

Easy Prey: Variables, Application and Numbers

References and applications have a one-to-one mapping. Numbers become repeated function application.

compileExp (Var name) = Ref name
compileExp (Application a b) = compileExp a `App` compileExp b
compileExp (Num n) = churchNum n

churchNum :: Int -> Node
churchNum n = Lam "f" $ Lam "x" $ foldr App (Ref "x") $ replicate n (Ref "f")

Multi-Argument Lambdas

The lambda form of the high level AST accepts multiple arguments. The lambda calculus only has single argument lambdas, so they need to be curried.

compileExp (Lambda args body) = foldr Lam (compileExp body) args

This will change an expression like (λ (x y z) body) to (λ (x) (λ (y) (λ (z) body))).

Function application is already curried by the parser, so there is nothing to fix up there.

Let

A let is changed into an immediately applied lambda.

(let ([x 5]) (+ x 3)) becomes ((λ (x) (+ x 3)) 5)
(let ([x a] [y b]) body) becomes ((λ (x) ((λ (y) body) b)) a)

compileExp (Let [] body) = compileExp body
compileExp (Let ((n,v):rest) body) = Lam n inner `App` compileExp v
  where
    inner = compileExp (Let rest body)

Letrec

To understand how the Y combinator allows recursion, it is sufficient to know that it is a function with the property (Y G) = (G (Y G)).

I’ll run through a quick example of defining a function F as some an arbitrary expression exp that can call F. The only way to bring a variable into scope is using lambda, so we might as well start there.

(λ (F) exp)

We can then use Y to get a definition for F.

(Y (λ (F) exp)) expands to ((λ (F) exp) (Y (λ (F) exp))) in the same way that (Y G) expands to (G (Y G)). We can then go ahead and evaluate exp with F bound to (Y (λ (F) exp)). If ever F is called inside exp, the Y combinator will expand again to produce another definition of F to call “recursively”.

The feeling I get about the Y combinator is this. Every time you try to evaluate (Y G) it expands to (G (Y G)), allowing you to continue a recursive computation in G for as long as you wish. From the interpreter’s perspective its not really recursion, its just expanding a copy of the same function definition as many times as necessary.

In the terms of the (λ (F) exp) example (letrec (F exp) body) is equivalent to (let ([F (Y (λ (F) exp))]) body)

My first thought when I saw this was: Hang on! The F bound in that let and the F in the definition aren’t the same. That’s variable shadowing!

And this is true from the standpoint of the interpreter, but the property of the Y combinator means that the two definitions of F are equivalent. Which is why it is a definition of recursion without recursion.

Here is the implementation in the compiler. The definition of the Y combinator still seems a bit like devil magic to me, but I thought this answer and Matt Might’s post made a lot of sense.

compileExp (Letrec name binding body) = compileExp recursiveLet
  where
    recursiveLet = Let [(name, y `Application` Lambda [name] binding)] body

-- (λy.λF.F (λx.y y F x))(λy.λF.F (λx.y y F x))
y :: Exp
y = term `Application` term
  where
    term = Lambda ["y", "F"] $ Var "F" `Application` Lambda ["x"] innerApp
    innerApp = Var "y" `Application`
               Var "y" `Application`
               Var "F" `Application`
               Var "x"

Interpreter

The task of the interpreter is to reduce a lambda term. But what does it mean to be reduced? Presumably we would like a result that can be easily interpreted by a human. The most important reduction we can do is a beta reduction, which is the same as a function call. (λx.body) arg would reduce to body with a new variable in scope x = arg (with some caveats about what variables are in scope for arg).

Beta-reduction is good enough for some programs. (= (+ (* 3 3) (* 4 4)) (* 5 5)) gives λt.λf.t, which is the church encoding for true.

Boolean results are pretty easy to see, but numbers get out of hand really quickly. The simple program (+ 2 3) results in (λb.(λa.λf.λx.a f (b f x)) (λf.λx.f (f x))) (λf.λx.f (f (f x))). The interpreter didn’t even do any work here. It has essentially told us that the answer is 2+3.

This is clearly not a solution. Instead, the interpreter should give an answer as a Haskell primitive. With the knowledge that (* 6 7) should be a number we can un-church the result to get 42.

Call by Need

I said earlier that the interpreter is lazy. The goal is to make programs like ((λ (_) 42) hang) work. When a function call is evaluated the argument is stored in a thunk. If ever that argument is used, then its value is computed, and the thunk is replaced with the computed value.

This requires state.

Data stored in the environment is either a thunk, or a forced result.

data Box s = Forced (Result s) | Thunk Node (Env s)

I use ST for mutation of the environment. All of those s variables are references to the state thread used to make sure state does not escape the ST monad. The environment is a map of variable names to a box.

type Env s = Map String (STRef s (Box s))

Why Use ST?

All of the other examples of interpreters I could find online used IORef instead of STRef for mutable boxes to store thunk results.

The reason they gave was that the interpreter was already dealing with IO. That is totally reasonable and makes their code much more simple.

Unfortunately my interpreter is not coupled at all to IO, so I had to use ST. This led to some interesting issues with impredictive types, and I ended up needing the Rank2Types language extension. More on that later.

Result Type

The result of a program could be a number, boolean, closure, or a list of results. I also have a function result type to help un-churching numbers and lists.

data Result s = Clos String Node (Env s)
              | RFun (Result s -> ExceptT String (ST s) (Result s))
              | RNum Integer
              | RBool Bool
              | RPair (Result s) (Result s)
              | REmpty

The RFun type transforms a result, and can possibly fail. The increment operator needed to decode numbers is defined as an RFun.

plus1 :: Result s -> ExceptT String (ST s) (Result s)
plus1 (RNum n) = return $ RNum (n+1)
plus1 _ = throwE "Cannot increment non-number type"

The type ExceptT String (ST s) (Result s) is a bit long and I use it a lot, so I made an alias.

type EST s a = ExceptT String (ST s) a

Interpreting An expression

The meat of the interpreter is the function lazyInterp

lazyInterp :: Env s -> Node -> EST s (Result s)
lazyInterp env (Lam arg body) = return $ Clos arg body env
lazyInterp env (Ref x) = maybe (scopeExcept x) force (Map.lookup x env)
lazyInterp env (App f x) = lazyInterp env f >>= (`app` Thunk x env)

scopeExcept :: String -> EST s a
scopeExcept = ExceptT . return . scopeError

There are only three forms, so only three cases to handle.

Lambdas immediately build a closure
References force a box from the environment, possibly failing if there is no variable in scope
Applications force the function, build a thunk for the argument, and then call out to an application function.

The scopeExcept function just changes the type of scopeError to EST s a instead of Either String a like is used the the rest of the pipeline.

To force a box, we check if its a thunk, if so, call lazyInterp and write the result.

force :: STRef s (Box s) -> EST s (Result s)
force ref = lift (readSTRef ref) >>= handleBox
  where
    handleBox (Forced r) = return r
    handleBox (Thunk node env) = do
      result <- lazyInterp env node
      lift $ writeSTRef ref (Forced result)
      return result

The type of the app function may seem a bit strange. I wanted to use it in the interpreter, and also to un-church results. Having it take a box lets me do this.

app :: Result s -> Box s -> EST s (Result s)
app (RFun f) (Forced x) = f x
app (RFun f) (Thunk node env) = lazyInterp env node >>= f
app (Clos arg body env) box = do
  argRef <- lift $ newSTRef box
  let extendedEnv = Map.insert arg argRef env
  lazyInterp extendedEnv body
app _ _ = throwE "Cannot apply non-function type"

Unlike closures, applying an RFun is strict in its argument. Fortunately RFun always needs its argument so this behavior is fine.

Extracting a Result

Once an expression has been interpreted, it is a closure. I then use an extraction function to end up with a Haskell literal.

The type of all the extractors is Result s -> EST s a, where a is the type of the literal. There are extractors for integers, booleans, and lists.

Integers can be extracted by applying the increment function and zero to the church encoded number.

(n inc 0)

Operations in our interpreter can fail, so each application needs to use bind. Once we have the final result we can unwrap it from the Result data type.

intExtractor :: Result s -> EST s Integer
intExtractor res = res `app` Forced (RFun plus1)
               >>= (`app` Forced (RNum 0))
               >>= getNum
  where
    plus1 (RNum n) = return $ RNum (n+1)
    plus1 _ = throwE "Cannot increment non-number type"
    getNum (RNum n) = return n
    getNum _ = throwE "Non-number result"

Booleans have a similar story, but here we apply the values True then False (b True False).

boolExtractor :: Result s -> EST s Bool
boolExtractor res = res `app` Forced (RBool True)
                >>= (`app` Forced (RBool False))
                >>= getBool
  where
    getBool (RBool b) = return b
    getBool _ = throwE "Non-boolean result"

Lists are more complicated. The list extractor first needs an extractor for the elements in the list. This ensures you can extract arbitrarily nested lists.

Lists are encoded as a function that takes something to do with a head and a tail and a result if the list is empty. To extract a literal list from this structure we apply an onCons function and the REmpty result.

onCons is a function that takes two results and yields an RPair. It looks a little funny because it has to be curried to fit into the application model of the interpreter.

To extract a list from a pair we run the extractor on the first of the pair to get a literal, and run the list extractor on the tail.

listExtractor :: (Result s -> EST s a) -> Result s -> EST s [a]
listExtractor ex res = res `app` Forced (RFun onCons)
                   >>= (`app` Forced REmpty)
                   >>= pairToList
  where
    onCons hd = return . RFun $ return . RPair hd
    pairToList REmpty = return []
    pairToList (RPair hd tl) = do
      x <- ex hd
      xs <- listExtractor ex tl
      return (x:xs)
    pairToList _ = throwE "Non-list result"

I lied a little bit. There is one more extractor.

My unit tests for early versions of the interpreter expected the output to be a Node, not the Result type I ended up with. Instead of rewriting all of the tests to expect a closure, I extract a Node from a closure.

The only variables that need to be bound from the closure’s environment are the ones that are free in the body. Dropping the environment from the closure, and instead binding those variables with immediately applied lambdas gets the job done.

unwrap :: Result s -> EST s Node
unwrap (Clos arg body env) = bindAll toBind $ Lam arg body
  where
    toBind = Set.toList $ arg `Set.delete` freeVars body
    bindAll [] expr = return expr
    bindAll (x:xs) expr = do
      binding <- maybe (scopeExcept x) force $ Map.lookup x env
      rawBinding <- unwrap binding
      bindAll xs $ Lam x expr `App` rawBinding
unwrap _ = throwE "Cannot unwrap non-closure type"

These three extractors result in EST s lit, not Either String lit like the rest of the pipeline. The last step is to run the stateful computation using runST.

Running a Stateful Computation

The only drawback to ST is the type of runST.

runST :: forall a. (forall s. ST s a) -> a

Functions in Haskell normally only have type variables on the outside of a type.

func :: forall a b. a -> b

They type of runST ensures that the resulting type cannot have any reference to the state thread, s.

The type forall s a. ST s a -> a could unify to ST s (Result s) -> Result s. This would be bad because a reference to the state thread would escape the stateful computation. This could cause side effects, and if you want that then Haskell is not the right language for you.

The type forall a. (forall s. ST s a) -> a cannot unify to ST s (Result s) -> Result s because the type variable s is not in scope on the right side of the function. This ensures that no references to the state thread escape.

Types of this form are beyond Haskell’s type inference engine. There is a special type rule for $ that allows code like runST $ do {...}. Unfortunately, the function for extracting lists requires s in its type signature, and the additional powers of the $ operator are not enough to make the type checker succeed.

To make this code work, I enabled the {-# LANGUAGE Rank2Types #-} extension. This allows all the extraction functions to have types that match the rest of the pipeline.

extract :: (forall s. Result s -> EST s a) -> Node -> Either String a
extract ex input = runST $ runExceptT (lazyInterp Map.empty input >>= ex)

extractInt :: Node -> Either String Integer
extractInt = extract intExtractor

extractBool :: Node -> Either String Bool
extractBool = extract boolExtractor

extractList :: (forall s. Result s -> EST s a) -> Node -> Either String [a]
extractList ex = extract $ listExtractor ex

interp :: Node -> Either String Node
interp = extract unwrap

To interpret and extract an int from a compiled program you could now call extractInt prog and get either an int or an error message.

Tying it all Together

Each piece of the pipeline transforms the program until it is a Haskell literal. The following example shows how each piece fits together by following a simple program through each stage.

The Pieces

Start with the input program (cons #t empty).

Parse the program to get:

Program [] (Application (Application (Var "cons") (Var "#t")) (Var "empty"))

The internal data structures aren’t very easy to read, so the rest of this is in nicer notation.

Desugar the program form to a single expression.

(let ([#t (λ (t f) t)])
  (let ([cons (λ (h t)
                (λ (f _)
                  ((f h) t)))])
    (let ([empty (λ (_ e) e)])
      (cons #t empty))))

Compile the program to the lambda calculus.

(λ#t.(λcons.(λempty.cons #t empty)
            (λ_.λe.e))
     (λh.λt.λf.λ_.f h t))
(λt.λf.t)

Interpret the program to get a closure.

{λf.λ_.f h t, [t={λ_.λe.e,[]}, h={λt.λf.t,[]}]}

Extract a pair from the closure.

< {λt.λf.t,[]} x {λ_.λe.e,[]} >

Finally, extract a list from the pair.

[True]

The Interface

I wanted to run programs from strings, and from files in the repl. I also wanted to interpret a program and pretty print the resulting lambda. runProgram, runFile and prettyProgram are the functions that do these things. runDisplayProgram contains all the common code between them.

type Interp a = Node -> Either String a

runDisplayProgram :: (a -> IO ()) -> String -> String -> Interp a -> IO ()
runDisplayProgram display filename input extractor = do
  base <- readBase
  either putStrLn display $ do
    defs <- base
    prog <- parseProgram filename input
    compiled <- compile defs prog
    extractor compiled

prettyProgram :: String -> IO ()
prettyProgram prog = runDisplayProgram (putStrLn . pretty) "input" prog interp

runProgram :: Show a => String -> Interp a -> IO ()
runProgram = runDisplayProgram print "input"

runFile :: Show a => String -> Interp a -> IO ()
runFile filename extractor = do
  input <- readFile filename
  runDisplayProgram print filename input extractor

The outer do in runDisplayProgram is for the IO monad. The inner one is for Either. To run the whole pipeline, first grab the base definitions while in IO, then drop into Either. Grab the base definitions, then parse the input, then compile the base definitions and the program together, and finally run the interpreter.

My original goal when writing this was to run the collatz sequence.

(define (step n)
  (if (even? n)
      (/ n 2)
      (+ (* n 3) 1)))

(define (collatz n)
  (if (<= n 1)
      0
      (+ 1 (collatz (step n)))))

(map collatz (range 1 15))

Save that to resources/collatz.lc, then run with

λ> runFile "resources/collatz.lc" (extractList intExtractor)
[0,1,7,2,5,8,16,3,19,6,14,9,9,17]

Tada! Those are the stopping times for the numbers 1 to 15.

The Best Wordament Players Are Mind Bogglingly Good

Tue, 04 Oct 2016 00:00:00 +0000

When I visit my family at Christmas boggle is always a popular choice of game, but a year ago the physical boggle board was often abandoned for wordament. Many of my family members were quite good at the game (myself not among them), but there were always a few players that seemed to find an impossible number of words. We collectively speculated that they were cheating to save our egos, so I wrote a simple solver to find every word on a Boggle board contained in a dictionary. Source code for the project can be found here.

The short answer is, either they are not not cheating, or if they are they are doing a very bad job of it. At least as of last Christmas. I haven’t looked at any scores since then.

Implementation

The approach I used is to build a trie of words in a dictionary file, and then do an exhaustive depth first search of the board to find every valid word on the board.

The trie is structured so every node has a flag saying if it is a valid word, and a list of children.

WordTree(children: Map[Char, WordTree], val isWord: Boolean)

With this representation, simultaneous recursion over the board and the trie will find every word. The board is small enough that overflowing the stack during the search will not be an issue.

def findAllWords(prevPath: Seq[Location], subDict: WordTree): Seq[String] = {
  val subwords = pathStep(prevPath).flatMap {
    step =>
      subDict.subTree(board.pieces(step)) match {
        case Some(subSubDict) => findAllWords(prevPath :+ step, subSubDict)
        case None => Nil
      }
  }

  val thisLetter = board.pieces(prevPath.last)
  val fixedSubwords = subwords.map(word => thisLetter + word)

  if(subDict.isWord)
    thisLetter +: fixedSubwords
  else
    fixedSubwords
}

Here pathStep is a function that returns all possible next steps in a valid path, and subTree recurses down the dictionary trie. Everything at the end of the function rebuilds words from the recursive calls to findAllWords.

Performance

Solving a 5x5 board with a list of over 100 000 words took less than a second on my laptop. Assuming this list could be fed directly to the game, there is no reason an automated solver should ever miss any word. Apparently there are some people who are just mind-bogglingly good at wordament.

Dominion AI

Fri, 29 Jul 2016 00:00:00 +0000

One of the projects in Matt Flatt’s functional programming studio class was to make an AI for the deck building game Dominion. We were each tasked with building an executable that would read what other players were doing on standard in, and communicate moves on standard out. The server that controlled all of the state of the game, and specified the protocol by which the processes would communicate was fixed. I chose to implement my client in Haskell.

There are three distinct pieces of the program following an architecture similar to the standard model-view-controller architecture. In this case there is a model responsible for representing the world state in the program, a parser similar to the standard view, but in this case it is responsible for I/O from the server instead of a user, and a two part controller. The controller is split into a driver side that contains all the code for dealing with I/O, and an agent side that is responsible for deciding how to play.

All source code is publicly available on github

Model

In order for the AI to make decisions it needs to know what the state of the game is. Every turn the server would send all of the known information about the game state to the client, so a client could be a simple reflex agent. This removes the difficulties of passing information along to the AI in the future in a purely functional way. I chose to represent the information received from the server as follows.

The most basic item in the game is probably the cards. I chose to represent them as a single union type. To simplify things we only dealt with the 10 suggested starting cards.

data Card = Cellar | Moat | ... | Gold | Province deriving (Eq, Show)

This representation makes no distinction between the victory, treasure, and kingdom cards, but such a system can get very complex for not much benefit.

I set up the Action type to represent any valid action a player can perform on their turn. They can Add treasures to their pool, Clean their hand, ending their turn, possibly showing one of any remaining cards in their hand, Buy a card from the supply, or Act by playing a kingdom card.

data Action = Add Card
            | Clean (Maybe Card)
            | Buy Card
            | Act Card [Card] deriving (Eq)

With the basic set of cards, a player can only be attacked by a Militia, and the only options they have are to discard, or to reveal a Moat.

data Defense = Reveal Card | Discard [Card] deriving (Eq)

Finally there is a GameState record that holds all the information about the current state of the game, and a Notification type. A Notification represents anything the server could send to the client. The server can Request an action from the client, Update the client about the actions of other players, Ask the client to defend itself if Attacked, or update the client that another player Defended them self.

data Notification = Request GameState
                  | Update String Action
                  | Attacked Action String GameState
                  | Defended String Defense deriving (Eq, Show)

This internal representation is the target of the parser that reads from standard in.

Parser and Output Format

The client uses a parser to read data from standard in, and Show instances defined on Action and Defense to build responses to put on standard out.

The show instances simply put the data into an s-expression to be written to standard out.

-- Wraps a list of strings in parens, and lower cases all of them
wrap :: [String] -> String
wrap ss = "(" ++ intercalate " " (map (map toLower) ss) ++ ")"

instance Show Action where
    show (Add c) = wrap ["add", show c]
    show (Clean (Just c)) = wrap ["clean", show c]
    show (Clean Nothing) = "(clean)"
    show (Buy c) = wrap ["buy", show c]
    show (Act c cs) = wrap $ "act" : show c : map show cs

instance Show Defense where
    show (Reveal c) = wrap [show c]
    show (Discard cs) = wrap $ "discard":(map show cs)

The parser module is responsible for taking valid input from standard in, and transforming it into a Notification. The interface it presents to the rest of the program is the function parseNotification which takes an input string, parses a notification and returns the pair of the parsed notification and any remaining input.

The details of the parser are long and boring, and all I will say here is that I quite like parser combinators. They make building a parser for a protocol like this much easier. My favorite explanation of what they are is here.

parseNotification :: String -> (Notification, String)
parseNotification input
    | Left e       <- parsed = error . show $ e
    | Right result <- parsed = result
    where
        parsed = parse (withRemaining notification) "stdin" input

withRemaining :: Parser a -> Parser (a, String)
withRemaining p = do
    result <- p
    rest <- getInput
    return (result, rest)

The first iterations of the server did not put any spaces between the notifications. It just wrote a single s-expression to the client, so in order to parse correctly it was necessary to read the input character by character, instead of using more common line reading functions. This created an interesting problem because if the client attempted to read a line the whole system would hang. The server would write the request to the client as a single s-expression "(move (...))", and then wait for the client to respond. Since there was no new line at the end of the s-expression the client, who had attempted to read a line, would be waiting for more input.

There were two ways to solve this problem. In an eager language the natural solution is to read standard in character by character, instead of waiting for a line. The reader would keep a counter of parentheses to determine where the end of the s-expression was. This solution is certainly possible in Haskell, but the lazy evaluation of strings presents a much more elegant solution.

In Haskell it is possible to get a lazy pointer to an entire buffer. This operation consumes the buffer, and in the case of standard in results in a string. This string, representing all of standard in over the lifetime of the program, is the string passed to parseNotification. The parser will remove all the characters containing the notification from the beginning of the string, and then return the parsed structure and any remaining input. In this case the remaining input is a lazy pointer to any remaining characters from standard in. Attempting to read any more characters from the string at this point would cause the program to hang because there aren’t more characters yet. The server hasn’t written them because it is waiting for a response from the client. Instead of reading the remainder of the string the parser simply returns it to be used later.

Driver

The driver code is responsible for all I/O and connects the parser to the agent. For now it is enough to know that there is a typeclass Agent, and it defines the function respond :: (Agent a) => a -> Notification -> String. I wanted to test my code out on different agents, so most of the code that would normally go in main got its own function drive. drive takes an agent, and runs the program, leaving the only code in main the choice of agent.

main :: IO ()
main = drive Miner49er

drive :: (Agent a) => a -> IO ()
drive agent = do
    hSetBuffering stdin NoBuffering
    hSetBuffering stdout NoBuffering
    hGetContents stdin >>= runClient agent

drive ensures there is no buffering on the I/O channels, and then gets a lazy pointer to all of standard in with hGetContents and hands that off to the runClient function which deals with looping responses from the client.

runClient :: (Agent a) => a -> String -> IO ()
runClient agent input
    | null trimmed = return ()
    | otherwise    = (putStr . respond agent) notification
                     >> runClient agent moreInput
    where
        trimmed = dropWhile isSpace input
        (notification, moreInput) = parseNotification trimmed

runClient will return if the server has closed the input stream, otherwise it prints the agents response, and loops on the remaining input from the server. Trimming the input is necessary because eventually the server was changed to emit new lines between each request, making it much easier to interface with, but by that time I had already built my parser, so I just remove all whitespace between the server’s messages.

Agent

The agent is responsible for deciding how to play dominion. The strategy I used is called big money smithy, and details about it can be found here, and how to beat it here. The basic strategy is to buy some small number of smithies, and treasure cards until you can buy provinces to end the game. It is surprisingly effective.

There are three distinct phases of a turn in dominion, during the first a player uses their actions to play kingdom cards, then they play their treasure cards, and then they can buy new cards from the supply. When their turn is over they discard any remaining cards in their hand, revealing one of them to the other players because it goes face up on their discard pile.

I wanted to make it easy to quickly define a new strategy, so I tried to generalize all of the decisions an agent would have to make. I came up with the following typeclass that I designed to be general enough for any strategy.

class Agent a where
    tryAction :: a -> GameState -> Either Action GameState
    tryAdd :: a -> GameState -> Either Action GameState
    tryBuy :: a -> GameState -> Either Action GameState
    tryDefend :: a -> GameState -> Either Defense GameState
    discardTo :: a -> GameState -> Int -> Defense

This requires an agent to define how to discard if necessary, and how to play actions, treasures, buys, and defenses using Either as a monad with “failure” being a valid response. This use of Either allows the implementation of playing actions to be very simple.

act :: (Agent a) => a -> GameState -> Action
act a state = case tryAction a state >>= tryAdd a >>= tryBuy a of
    Left action  -> action
    Right _      -> Clean $ find (const True) (hand state)

If none of tryAction, tryAdd, or tryBuy “fail” with a valid response the turn is over and the discard message is returned.

Playing Kingdom Cards

The simplest of the kingdom cards do not affect any other cards. There are more complex cards, like the mine or the cellar that require the player to do something with their hand. This function cannot deal with cards like that. For my strategy this is good enough because the only action card I will play is the smithy.

trySimplePlay :: Card -> GameState -> Either Action GameState
trySimplePlay card state
    | canAct state card = Left $ Act card []
    | otherwise = Right state

inHand :: GameState -> Card -> Bool
inHand state card = card `elem` hand state

canAct :: GameState -> Card -> Bool
canAct state card = actions state > 0 && inHand state card

Playing Treasures

I chose to always play all the treasure cards in my hand, which made this phase very easy

playAllTreasures :: GameState -> Either Action GameState
playAllTreasures state = case find isTreasure (hand state) of
    Just t -> Left $ Add t
    _ -> Right state

Buying

This is arguably where most of the strategy of the game is, and is also where my code is the ugliest. I decided to use a system of card priority with a function to determine if a card should be bought based on the game state.

buyPriority :: (GameState -> Card -> Bool) -> [Card] -> GameState -> Either Action GameState
buyPriority should priority state
    | buys state > 0 = case find (\c -> canBuy c && should state c) priority of
        Just c -> Left (Buy c)
        Nothing -> Right state
    | otherwise = Right state
    where
        canBuy c = cost c <= coins state && c `elem` supply state

If the agent has no buys left then it just passes on to the next phase, otherwise it searches through the prioritized list of cards and buys the first one that it can buy and should buy. The should buy function is provided by the agent instance, and allows it to control when to buy certain cards. I have found the ability to pass a function to a function like this to define part of its behavior an extremely useful pattern that allows for much more flexible implementations.

Discarding

The last interesting question to answer is how to discard cards if the player is attacked, and they cannot defend with a moat. I settled on a similar strategy as buying, allowing the agent instance to specify a priority of cards to discard.

discardPriority :: [Card] -> GameState -> Int -> Defense
discardPriority priority state n = Discard (findDiscards (length (hand state) - n)
                                                         (hand state)
                                                         priority)

findDiscards :: Int -> [Card] -> [Card] -> [Card]
findDiscards n hand toTry
    | n <= 0     = []
    | null toTry = take n hand
    | otherwise  = (take n toRemove) ++ findDiscards (n - length toRemove) toKeep (tail toTry)
    where
        (toRemove, toKeep) = partition (== (head toTry)) hand

This is a function that recursively removes cards in the order of the list toTry if it runs out of cards to try it takes from the beginning of the hand.

Defining Agent Instances

With all of these default implementations in place my smithy bot now only took a few lines to define. My goal was to build an agent that would:

always buy a province if it can
only buy smithies if less than 5% of its cards are smithies
avoid buying gold and silver in the late game
only buy duchies and estates at the end of the game
discard its least valuable cards first

data SmithyMoney = SmithyMoney

instance Agent SmithyMoney where
    tryAction _ = trySimplePlay Smithy

    tryAdd _ = playAllTreasures

    tryBuy _ = buyPriority shouldBuy priority
        where
            priority = [Province, Gold, Smithy, Duchy, Silver, Estate]
            shouldBuy _ Province = True
            shouldBuy state Gold = numInSupply state Province > 4
            shouldBuy state Smithy = probDraw (allMyCards state) Smithy < 0.05
            shouldBuy state Duchy = numInSupply state Province < 6
            shouldBuy state Silver = numInSupply state Province > 4
            shouldBuy state Estate = numInSupply state Province < 4

    tryDefend _ = revealMoat

    discardTo _ state = discardPriority ([Province, Duchy, Estate] ++ valueSorted) state
        where
            valueSorted = sortOn value [Copper, Silver, Gold, Smithy]
            value Smithy = expectedDrawValue state 3
            value t = treasureWorth t

And thats it.

I also made a different agent that uses some more of the kingdom cards just to see how it would do. I won’t show the whole definition here, but I think it is interesting to see how actions can be chained together even in the agent instance using bind.

    tryAction _ state = trySimplePlay Village state >>= trySimplePlay Militia >>= trySimplePlay Smithy >>= tryMine

To fine tune my smithy bot I ran it in batches of 50 games against this agent, and got it to the point where it would win about 95% of the time, and lose about 1% of the time.

Testing

I used HUnit and Tasty as a testing framework, which all fits in very nicely with stack. I wrote many tests, and they were invaluable in helping me write correct code, but they are also not very interesting. I found it quite annoying that I could not access a function for testing if it was not exposed by a module, and this lead to me exposing nearly every piece of code I wrote. I feel like there must be a better way.

Edit: The Internal/External module pattern seems to be the way to do this. This answer has a great explanation of how that should work.

Further Thoughts

After defining my agent instances it seems that most of the strategy of the game is not in how to play cards from a hand, but knowing what deck composition to aim for. A possible improvement to the project would be to write an agent interface that requires only an early game, mid game, and late game deck composition and attempts to satisfy it. This would make the internals of the agent much more complected, but would make tuning and testing a strategy much easier.

Python Compiler

Sat, 31 Oct 2015 00:00:00 +0000

Every assignment in Matt Might’s compilers class at the University of Utah is a piece of a compiler for Python. In an ideal world at the end of the semester all of the students would have made an end to end system that can fully compile Python into a low level language such as Java bytecode or x86. During the semester Matt Might was called to Washington to help set up the Precision Medicine Initiative which caused the class schedule to lag. There were only four projects in the end, and the output of the system was a normalized subset of Python ready for a CPS transform.

Lexing

This project took an file of valid Python code and produced a sequence of tokens. The tokens were one of eight s-expressions:

(NEWLINE)

(INDENT)

(DEDENT)

(ID name)

(LIT value)

(KEYWORD symbol)

(PUNCT text)

(ENDMARKER)

For example if the lexer got the input

def f(x):
    return 7*x

It would generate

(KEYWORD def)
(ID "f")
(PUNCT "(")
(ID "x")
(PUNCT ")")
(PUNCT ":")
(NEWLINE)
(INDENT)
(KEYWORD return)
(LIT 7)
(PUNCT "*")
(ID "x")
(NEWLINE)
(DEDENT)
(ENDMARKER)

The hardest part of the stage was dealing with Unicode in variable names, and strings. Variables were required to be NFKC normalized, and string literals had to be properly represented in Racket. The lexer requires some state to deal with the whitespace correctly, and there are some interesting corner cases with the ellipsis literal and a regular period.

Parsing

The parser we built was really just a bunch of reduction rules on top of the racket parsing library and DERP2 built by Matt Might (if you want to know more about how this works, the code is available, and there is also this fantastic post about how parser combinators work).

We started with the Python grammar specification, and inserted reductions to give us the parse tree matching Python’s built in AST module. As an example of what this entails I will walk through the process for dealing with an if statement.

From the Python grammar, an if statement looks like this (shown here in lisp-like syntax):

(if_stmt (seq "if"
              test
              ":"
              suite
              (rep (seq "elif" test ":" suite))
              (opt (seq "else" ":" suite))))

So an if statement must have a test expression and a body, and then any number of else if clauses, followed by an optional else clause.

In my code I use the function ($--> input output), which generates a reduction rule that matches the first argument and results in the second argument. Inside the reduction, parts of the input can be extracted with the ($ n) function, and the symbol $$ gets expanded to the entire input.

An example of a simple reduction is the following application.

($--> (seq "if" test ":" suite)

      (list ($ 2) ($ 4)))

This will match the literal "if", any expansion of the non-terminal test, the literal ":", and finally any expansion of the non-terminal suite. It will return a list of just the second and fourth things it matched, in this case the expansion of test and suite.

The following is the full set of reductions for an if statement.

(if_stmt ($--> (seq ($--> (seq "if" test ":" suite)
                          (list ($ 2) ($ 4)))
                    (rep ($--> (seq "elif" test ":" suite)
                               (list ($ 2) ($ 4))))
                    ($--> (opt (seq "else" ":" suite))
                          (match $$
                            [#f '()]
                            [`("else" ":" ,block) block])))
               (process-ifs (cons ($ 1) ($ 2)) ($ 3))))

The inner reductions filter out the unneeded keywords, and the outer reduction is a simple call to a helper function that builds the correct output syntax. The helper function takes a list of if/elif statements, and a possibly empty else statement. It recursively packages them up into the correct if-else syntax.

(define (process-ifs ifs tail)
  (match ifs
    [`((,test ,body))
     `(If (test ,test)
          (body ,@body)
          (orelse ,@tail))]
    [(cons `(,test ,body) others)
     `(If (test ,test)
          (body ,@body)
          (orelse ,(process-ifs others tail)))]))

The final system would take the following if statement

if t1:
  r1
elif t2:
  r2
else:
  r3

and produce this syntax tree

(If
 (test (Name t1))
 (body (Expr (Name r1)))
 (orelse (If (test (Name t2))
             (body (Expr (Name r2)))
             (orelse (Expr (Name r3))))))

Desugaring

The next two projects were desugaring and normalization steps. They both used a tree walker built for the output of the parser. The implementation of the tree walker simply does a match on a structure in the parse tree, and recursively descends until it hits a leaf, performing a transformation along the way. The interface provided by the tree walker is the ability to walk the tree performing some transformation on every subtree. It also provided the functions to perform a transformation until a fixed point was hit, and the option to transform top down or bottom up.

Basically all of the code that I wrote for this part was transformation functions that take a statement, and return a list of possible replacement statements. The most basic transform is fixing return statements so that a lone return call becomes return None.

(define (canonicalize-return stmt)
  (match stmt
    [`(Return)  (list '(Return (NameConstant None)))]
    [else       (list stmt)]))

Some of the other interesting transforms we did in the phase are lifting decorators, and removing classes. Removing decorators from functions is a straightforward operation. Generally code of the form

@<decorator>
def <name>(<parameters>):
  <body>

should be turned into code of the form

def <name>(<parameters>):
  <body>
<name> = <decorator>(<name>)

After the end of this phase of the compiler we needed to normalize before applying any more desugarings.

Normalizing

The main goal of normalization is breaking arbitrarily complex statements into bits that can be easily translated to machine code. This step is also where the order of evaluation is fixed. For example after this phase the statement f(print(1), print(2)) could not possibly do anything other than print 1 then print 2. In order to ensure the entire program was normalized we ran a normalizing tree transform until nothing happened. The normalizing transform matched every complex statement in the grammar and checked if all its sub-expressions were atomic. If they were not it bound them to a variable and then made the complex call with the variable instead.

For example the statement

result = 3*x + 4*y

would be internally represented after parsing and desugaring as

'(Assign (targets (Name result))
         (BinOp (BinOp (Num 3) Mult (Name x))
                Add
                (BinOp (Num 4) Mult (Name y))))

This internal representation would trigger the following match clause in the body of normalize/stmt.

[`(Assign (targets ,target) (value ,rhs))
 #:when (not (normalized? rhs))
 `((Assign (targets ,target)
           (value ,(normalize-expr! rhs))))]

Then the following match clause would do the real work in the body of normalize-expr!.

``racket [(BinOp ,lhs ,op ,rhs)(BinOp ,(atomize! lhs) ,op ,(atomize! rhs))]

The `atomize!` function would just push the complex statements into their own assign statements and return the name of the new variable they were assigned to. The original statement would then look like this

```python
_tmp_1 = 3*x
_tmp_2 = 4*y
result = _tmp_1 + _tmp_2

Finishing

There were still a couple of nasty things to take care of before the Python structures we had at this point could be translated into machine code. In order to deal with the yield construct we were going to do a double CPS transform. There was also a fair amount of other steps to take before we arrived at a true low level language.

Tail Call Optimized Memoization in Scala

Sat, 15 Aug 2015 00:00:00 +0000

Memoization is a technique where one saves the result of a function so if the function is ever called again with the same inputs it does not have to compute the result a second time. The example I will use here is the Fibonacci sequence. The obvious recursive implementation given below has comically terrible performance.

def fib(n: Int):Long = {
  case _ if n < 0 =>
    throw new Exception("negative input")
  case 0 => 1
  case 1 => 1
  case _ => fib(n-1) + fib(n-2)
}

I like this definition. I think it is easier to understand than an optimized version with loops and lists. Unfortunately this definition is unusably slow. Computing the value of the 10th Fibonacci number using this implementation requires 177 calls to fib! Many of these calls are calling fib with the same arguments. fib(1) was called 55 times.

Bummer, but memoization can save it! Since fib is a pure function, calls with the same arguments will always return the same value, so there is no need to compute the result again. The implementation below uses a map to store the results of previous computations.

val cache = mutable.Map.empty[Int, Long]

def memoizedFib(n: Int): Long = {
  cache.get(n) match {
    case Some(answer) => answer
    case None =>
      val answer = n match {
        case _ if n < 0 =>
          throw new Exception("negative input")
        case 0 => 1
        case 1 => 1
        case _ => memoizedFib(n-1) + memoizedFib(n-2)
      }
      cache.update(n, answer)
      answer
  }
}

This version has the desired linear performance, requiring only 11 calls to fib to compute the value of the 10th number in the sequence. Great, but that code is ugly and no fun to write.

Generalizing to Any Function

What I want is a small change I can make to the original function to get the power of memoization. The first step is to make a more general memoize function.

Due to limitations in scala generics the functions I will be able to memoize will all have the type signature I=>O, that is they accept only one input of type I and produce an output of type O.

If I am going to write a more general memoize function, in what way should it be more general? I decided that forcing a specific caching system was not necessary. Checking a cache may be able to give an answer very quickly, but not always, but may sometimes fail. I will call this function check and its type is I=>Option[O]. After I compute a result the hard way I should save it away so check can return it in the future. I will call this saving function update and its type is (I,O)=>Unit. Using this framework I implemented a generic memoize function.

def memoize[I, O](fn: I => O,
                  check: I => Option[O],
                  update: (I, O) => Unit): I => O = {
  (input: I) => check(input) match {
    case Some(answer) => answer
    case None =>
      val answer = fn(input)
      update(input, answer)
      answer
  }
}

Cool! I also implemented a memoizeMap function that memoizes some function, I=>O, using a map for convenience.

def memoizeMap[I, O](fn: I => O): I => O = {
  val cache = mutable.Map.empty[I, O]
  memoize(fn, cache.get, cache.update)
}

Now I can finally get the syntax I want. Unfortunately def is somewhat special in scala, so the best syntax I could wrangle is the following.

lazy val fib: Int => Long = memoizeMap {
  case n if n < 0 =>
    throw new Exception("negative input")
  case 0 => 1
  case 1 => 1
  case n => fib(n-1) + fib(n-2)
}

The lazy keyword is the only way I could figure out to get recursive binding without using def.

Tail Recursion

What if I want to compute the 100 000th Fibonacci number though? Things get a little more tricky. The syntax I ended up with is not very readable, so iteration is probably a more useful (and faster) solution for most functions. There are many functions that are not easily represented as iteration for which the method I use here is useful.

First of all: what is tail recursion? A tail recursive function only has calls to itself in tail position. In a language like Java with an explicit return keyword this means the only recursive calls are of the form return myself(...);. More generally, the result of the recursive function is equal to the recursive call, or some atomic value.

The Fibonacci function branches into two sub-calls, fib(n-1) and fib(n-2), at every step which introduces its own set of problems. For starters I implemented factorial.

lazy val fact: Int => BigInt = memoizeMap {
  case n if n < 0 =>
    throw new Exception("negative input")
  case 0 => 1
  case n => n * fact(n - 1)
}

This implementation is not tail recursive. Notice the value of a call to fact is n * fact(...). The traditional way to make this function tail recursive is include an accumulator as an argument to the function. Unfortunately this will make the memoization useless because every intermediate value is unique to the original call.

Continuation Passing Style

I will leave a full explanation of continuation passing style to others. Some good explanations are here and here. In this specific case the continuation acts as a return for the recursive function. The recursive function must, at some point, call the continuation it is passed with the result it is to return. Note the type of the continuation here is O=>O. The following is the same factorial function rewritten in continuation passing style. This change lets me write all recursive calls in tail position.

def cpsFact(input: Int, kont: BigInt=>BigInt): BigInt = {
  input match {
    case n if n < 0 =>
      throw new Exception("negative input")
    case 0 => kont(1)
    case n => cpsFact(n-1,
                      (result: BigInt) => kont(n * result))
  }
}

I wrote a new memoizing function that accepts CPS functions. Unlike the first one this memoizer is properly tail recursive. The update step is moved into the continuation of the recursive call, ensuring the stack frame of the lambda can be torn down.

def memoizeRec[I, O](fn: (I, O=>O) => O,
                     check: I => Option[O],
                     update: (I, O) => Unit):
                    (I, O=>O) => O = {
  (input: I, kont: O=>O) => check(input) match {
    case Some(answer) => kont(answer)
    case None => fn(input,
                    (result: O) => {
                      update(input, result)
                      kont(result)
                    })
  }
}

Unfortunately the scala compiler will not tail call optimize lazy vals, so I needed a different solution. The scala standard library had all I needed to fix the problem tucked away in the TailCalls library. The implementation containing the tail call optimization is given below.

def memoizeRec[I, O](fn: (I, O=>TailRec[O]) => TailRec[O],
                     check: I => Option[O],
                     update: (I, O) => Unit):
                    (I, O=>TailRec[O]) => TailRec[O] = {
  (input: I, kont: O=>TailRec[O]) => check(input) match {
    case Some(answer) => tailcall(kont(answer))
    case None =>
      tailcall(fn(input, (result: O) => {
        update(input, result)
        tailcall(kont(result))
      }))
  }
}

def memoizeRecMap[I, O](fn: (I, O=>TailRec[O]) => TailRec[O]):
    (I, O=>TailRec[O]) => TailRec[O] = {
  val cache = mutable.Map.empty[I, O]
  memoizeRec(fn, cache.get, cache.update)
}

lazy val cpsFactTail: (Int, BigInt=>TailRec[BigInt]) =>
                      TailRec[BigInt] = memoizeRecMap {
  (input, kont) => input match {
    case n if n < 0 =>
      throw new Exception("negative input")
    case 0 => tailcall(kont(1))
    case n =>
      tailcall(cpsFactTail(n-1,
        (result: BigInt) => tailcall(kont(n * result))))
  }
}

I said earlier that the Fibonacci sequence was harder to express using only tail calls. You are only allowed to make a single recursive call, so the continuation must contain the other recursive call. Remember continuations must also only contain tail calls, so the continuation passed to the recursive call will be a tail recursive call with a final continuation calling the outer continuation with the final result of this function call. Not very intuitive, I know.

lazy val cpsFibTail: (Int, BigInt=>TailRec[BigInt]) =>
                     TailRec[BigInt] = memoizeRecMap {
  (input, kont) => input match {
    case n if n < 0 =>
      throw new Exception("negative input")
    case 0 => tailcall(kont(1))
    case 1 => tailcall(kont(1))
    case n => tailcall(cpsFibTail(n-1, (res1: BigInt) =>
      tailcall(cpsFibTail(n-2, (res2: BigInt) =>
        tailcall(kont(res1 + res2))))))
  }
}

With this implementation I was able to compute the 100 000th Fibonacci number in just 2.2 seconds.

Limitations

Functions must be pure. Memoizing Random.nextInt is not very useful.
Continuations are not easy to read. All of the simplicity of the initial Fibonacci function has been lost. In most cases arbitrary depth recursion is not a requirement, so the more elegant memoizeMap version could be used instead.
Input must be a single type. Without macros accepting an arbitrary number inputs in a type-safe manner is simply not possible. Yes you could write a different version that accepts two input functions, or put your arguments into a tuple. In my mind these are not acceptable solutions.