AoC 2021 Day 14 in Go

package main

import (
    "aoc/utils"
    "fmt"
    "math"
)

func main() {
    template, rules := parseInput(utils.ReadLines())
    fmt.Println("part 1: ", GrowPolymer(template, rules, 10))
    fmt.Println("part 2: ", GrowPolymer(template, rules, 40))
}

type Symbol uint8 // Single-letter symbol (range [0-25])
type Key uint16   // Two-letter key (range [0-675])

type Rule struct {
    Symbol Symbol // New byte generated by this rule
    Keys   []Key  // New "rule keys" generated by this rule
}
type RuleSet map[Key]Rule

const (
    MaxKey    = 26*26 - 1
    MaxSymbol = 25
)

func GrowPolymer(template string, rules RuleSet, gens int) int {
    symCounts := make([]int, MaxSymbol+1)
    keyCounts := make([]int, MaxKey+1)
    initCounters(symCounts, keyCounts, template)

    keyCountsNext := make([]int, MaxKey+1)
    swap := func() { keyCounts, keyCountsNext = keyCountsNext, keyCounts }

    var count int
    for i := 0; i < gens; i++ {
        for k, rule := range rules {
            count, keyCounts[k] = keyCounts[k], 0
            symCounts[rule.Symbol] += count
            for _, newKey := range rule.Keys {
                keyCountsNext[newKey] += count
            }
        }
        swap()
    }
    return computeResult(symCounts)
}

func initCounters(symCounts, keyCounts []int, template string) {
    for i := 0; i < len(template)-1; i++ {
        k := key(template[i:])
        s := symbol(template[i])
        keyCounts[k]++
        symCounts[s]++
    }
    s := symbol(template[len(template)-1])
    symCounts[s]++
}

func computeResult(symCounts []int) int {
    min, max := math.MaxInt, 0
    for _, count := range symCounts {
        if count == 0 {
            continue
        }
        if count > max {
            max = count
        }
        if count < min {
            min = count
        }
    }
    return max - min
}

func parseInput(input []string) (string, RuleSet) {
    return input[0], parseRules(input[2:])
}

func parseRules(specs []string) RuleSet {
    rules := make(RuleSet, len(specs))
    for _, spec := range specs {
        var keyStr string
        var b byte
        _, err := fmt.Sscanf(spec, "%s -> %c", &keyStr, &b)
        if err != nil {
            panic(err)
        }
        k := key(keyStr)
        rules[k] = Rule{Symbol: symbol(b)}
    }
    for k, rule := range rules {
        key1, key2 := getNewKeys(k, rule.Symbol)
        if _, ok := rules[key1]; ok {
            rule.Keys = append(rule.Keys, key1)
        }
        if _, ok := rules[key2]; ok {
            rule.Keys = append(rule.Keys, key2)
        }
        rules[k] = rule
    }
    return rules
}

func key(str string) Key {
    return Key(symbol(str[0]))*26 + Key(symbol(str[1]))
}

func symbol(b byte) Symbol {
    return Symbol(b - 'A')
}

func getNewKeys(k Key, s Symbol) (Key, Key) {
    left, right := k/26, k%26
    return left*26 + Key(s), Key(s)*26 + right
}