Generating data structures that need additional functions

There has been a lot of talk about generics lately in the Go community which as lead to me thinking about them a lot lately. In thinking about generics, my mind instinctively wandered to code generation because that has been my go-to tool when I do need something resembling a generic. In fact, I have written about using code generation to get by without generics in Go in the past.

While thinking about some of my code generation, it got me to thinking about a question that was posed to me but I never felt I answered adequately.

How do you generate data structures that require additional information, such as a comparison function?

Generators work well enough for linked lists and other data structures that are data agnostic, but what happens when you want to write a heap template?

Ideally we would like to create a template like we do for linked lists, but how do we do this if we don’t know how to compare elements ahead of time?

After toying around with it, I came up with something I figured I would share. It takes advantage of first class functions and essentially allows developers to write the comparison logic after the template is generated. The end result is code that feels somewhat reminiscent of sort.Slice.

An example - generating heaps

Let’s dig into an example to see how this would work. We will be writing a template that might be used to generate heap implementations for various data types. Heaps require a way to compare elements in the heap so that it can always provide you with the smallest or largest element in the heap (depending on how you compare them), so it clearly needs to know a little bit about the data we store inside of our heap data structure.

Note: I am going to assume you have used the container/heap package in the past and won’t be spending much time on the code that actually implements heap.Interface.

First let’s look at your template code.

package main

import "container/heap"

type Type struct{}
type TypeHeap struct {
	heap typeHeap
}

func (t *TypeHeap) Init(less func(i, j Type) bool) {
	t.heap.less = less
}

func (h *TypeHeap) Push(v Type) {
	heap.Push(&h.heap, v)
}

func (h *TypeHeap) Peek() Type {
	if h.Len() == 0 {
		panic("heap is empty")
	}
	return h.heap.slice[0]
}

func (h *TypeHeap) Pop() Type {
	val, ok := heap.Pop(&h.heap).(Type)
	if !ok {
		panic("invalid type in our heap - this shouldn't ever happen")
	}
	return val
}

func (h *TypeHeap) Len() int {
	return h.heap.Len()
}

type typeHeap struct {
	slice []Type
	less  func(i, j Type) bool
}

func (h typeHeap) Len() int {
	return len(h.slice)
}

func (h typeHeap) Less(i, j int) bool {
	return h.less(h.slice[i], h.slice[j])
}

func (h typeHeap) Swap(i, j int) {
	h.slice[i], h.slice[j] = h.slice[j], h.slice[i]
}

func (h *typeHeap) Push(x interface{}) {
	h.slice = append(h.slice, x.(Type))
}

func (h *typeHeap) Pop() interface{} {
	n := len(h.slice)
	ret := h.slice[n-1]
	h.slice = h.slice[0 : n-1]
	return ret
}

You can view this code on the Go Playground here: https://play.golang.org/p/8Zi-fIyePl

In the code above we have three types.

• Type - this is a filler meant to stub out whatever final data type our heap uses (eg int). It will eventually be removed from the generated code and won’t be exported.
• TypeHeap - this is the heap that will be exported when we generate our code. This is basically a wrapper around the functions provided by the container/heap package and is bundled with some type safety to make your life easier when using the generated heap. With this you will never need to import the container/heap package into your other code.
• typeHeap - this is an unexported type that implements the heap.Interface type and is hidden from our end users. This will actually store our slice and handles most of the heap logic and gets wrapped by the TypeHeap type.

One thing to note is that our typeHeap type has a Less method, but it defers that logic to a less field stored on the type. This allows us to define that function later.

Also, unlike the sort.Interface, our less function that needs defined will NOT pass in indices, but will actually pass in values that need compared. This can catch some people off guard at times.

When generating our heaps we would generate everything except for our less function, and then we would write that function manually. When you write the less function you could do so as part of the package with all your generated code, or you could leave that detail up to users of your heap. Either way is fine.

An example of this can be found below (or on the Go Playground: https://play.golang.org/p/dGr-2BQKq9).

In the example code we define our own less function (inside of the main function) and could compare our strings however we see fit.

package main

import (
	"container/heap"
	"fmt"
)

func main() {
	less := func(i, j string) bool {
		return i < j
	}
	var h String
	h.Init(less)
	h.Push("cat")
	h.Push("dog")
	h.Push("a")
	h.Push("bird with red")
	h.Push("bird")
	for h.Len() > 0 {
		fmt.Println(h.Pop())
	}
}

type String struct {
	heap stringHeap
}

func (t *String) Init(less func(i, j string) bool) {
	t.heap.less = less
}

func (h *String) Push(v string) {
	heap.Push(&h.heap, v)
}

func (h *String) Peek() string {
	if h.Len() == 0 {
		panic("heap is empty")
	}
	return h.heap.slice[0]
}

func (h *String) Pop() string {
	val, ok := heap.Pop(&h.heap).(string)
	if !ok {
		panic("invalid type in our heap - this shouldn't ever happen")
	}
	return val
}

func (h *String) Len() int {
	return h.heap.Len()
}

type stringHeap struct {
	slice []string
	less  func(i, j string) bool
}

func (h stringHeap) Len() int {
	return len(h.slice)
}

func (h stringHeap) Less(i, j int) bool {
	return h.less(h.slice[i], h.slice[j])
}

func (h stringHeap) Swap(i, j int) {
	h.slice[i], h.slice[j] = h.slice[j], h.slice[i]
}

func (h *stringHeap) Push(x interface{}) {
	h.slice = append(h.slice, x.(string))
}

func (h *stringHeap) Pop() interface{} {
	n := len(h.slice)
	ret := h.slice[n-1]
	h.slice = h.slice[0 : n-1]
	return ret
}

We could probably even reduce our code footprint further (specifically the stringHeap type likely has parts that could be shared across many different heap types), but for now this illustrates my point - you can generate typed data structures even when they require information about the underlying data.

Note: You could also likely use an interface like Comparer and require your data type have that method, but this feels less Go-ish to me so I opted to use this approach.

What about the generics?!?!

Generics could potentially make this much simpler, but before I even start to talk about them I wanted to first illustrate how you might create a template for a data structure like this without generics. All too often people jump to massive tools or features to solve a simple problem without at least exploring other options.

That doesn’t mean I wouldn’t like to see generics in Go - I would - but it is important to explore other options before adding such massive changes to a language.