Merge-sort in go

merge.go

package main

import "fmt"

func main() {
	a := createListNodes([]int{2, 3, 4, 0, 2, 8, 9})
	fmt.Println("starting with", a)
	fmt.Println("ending with", sortList(a))
}

type ListNode struct {
	Val  int
	Next *ListNode
}

func (n *ListNode) String() string {
	r := "["
	for n != nil {
		r += fmt.Sprintf("%d", n.Val)
		n = n.Next
		if n != nil {
			r += ", "
		}
	}
	r += "]"
	return r
}

func createListNodes(a []int) *ListNode {
	r := &ListNode{a[0], nil}
	c := r
	for i := 1; i < len(a); i++ {
		c.Next = &ListNode{a[i], nil}
		c = c.Next
	}
	return r
}

func sortList(head *ListNode) *ListNode {
	var a, b, last, rest *ListNode
	n := 1
	for {
		fmt.Printf("n=%d, current list is %v\n", n, head)
		a, rest = takeN(head, n)
		if rest == nil {
			break
		}
		b, rest = takeN(rest, n)
		head, last = merge(a, b)
		fmt.Printf("merged the first two sublists now have %v ", head)
		fmt.Printf("and the last value points to %v\n", last)
		for rest != nil {
			a, rest = takeN(rest, n)
			b, rest = takeN(rest, n)
			a, b = merge(a, b)
			fmt.Printf("merge returned head: %v and last %v\n", a, b)
			last.Next = a
			last = b
			fmt.Printf("Current list from head %v\n", head)
		}
		n = 2 * n
	}
	return head
}

func merge(a, b *ListNode) (head, last *ListNode) {
	fmt.Printf("Merging two sorted lists:\n")
	fmt.Printf("- a: %v\n", a)
	fmt.Printf("- b: %v\n", b)
	defer func() {
		fmt.Printf("-- merged: %v\n", head)
	}()

	defer func() {
		for last != nil && last.Next != nil {
			last = last.Next
		}
	}()

	if a == nil {
		head, last = b, b
		return
	}
	if b == nil {
		head, last = a, a
		return
	}
	if a.Val < b.Val {
		head, last, a = a, a, a.Next
	} else {
		head, last, b = b, b, b.Next
	}
	for {
		if a == nil {
			last.Next = b
			return
		}
		if b == nil {
			last.Next = a
			return
		}
		if a.Val < b.Val {
			last.Next, a = a, a.Next
		} else {
			last.Next, b = b, b.Next
		}
		last = last.Next
	}
}

// Returns a pointer to the first n elements of list disconnected from the
// remaining elements in rest.
func takeN(a *ListNode, n int) (first, rest *ListNode) {
	if a == nil {
		return
	}
	cur := a
	for i := 1; i < n; i++ {
		if cur.Next == nil {
			break
		}
		cur = cur.Next
	}
	first = a
	rest = cur.Next
	cur.Next = nil
	return
}