AoC 2021 Day 12 in Go

package main

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

func main() {
    g := NewGraph(utils.ReadLines())
    fmt.Println("part1: ", g.Part1())
    fmt.Println("part2: ", g.Part2())
}

type Graph struct {
    edges    [][]uint16
    bigCaves uint16
}

func (g *Graph) Part1() int {
    initCache()
    return g.countPathsPart1(startCave, 0)
}

func (g *Graph) Part2() int {
    initCache()
    return g.countPathsPart2(startCave, 0)
}

const (
    startCave = 1
    endCave   = 0
)

func caveMask(cave uint16) uint16 {
    return 1 << cave
}

func (g *Graph) countPathsPart1(from uint16, visited uint16) int {
    var count int
    if count, ok := getCachedResult(from, visited); ok {
        return count
    }
    for _, to := range g.edges[from] {
        if to == endCave {
            count++
            continue
        }
        toMask := caveMask(to) &^ g.bigCaves
        if visited&toMask == 0 {
            count += g.countPathsPart1(to, visited|toMask)
        }
    }
    putToCache(from, visited, count)
    return count
}

const doubleMask = 1 << 15

func (g *Graph) countPathsPart2(from uint16, visited uint16) int {
    var count int
    if count, ok := getCachedResult(from, visited); ok {
        return count
    }
    for _, to := range g.edges[from] {
        if to == endCave {
            count++
            continue
        }
        toMask := caveMask(to) &^ g.bigCaves
        if visited&toMask == 0 {
            count += g.countPathsPart2(to, visited|toMask)
        } else if visited&doubleMask == 0 {
            count += g.countPathsPart2(to, visited|doubleMask)
        }
    }
    putToCache(from, visited, count)
    return count
}

// Memoïzation

type key struct {
    cave    uint16
    visited uint16
}

type cache map[key]int

var memo cache

func initCache() {
    memo = make(cache, 500)
}

func getCachedResult(from uint16, visited uint16) (int, bool) {
    count, ok := memo[key{from, visited}]
    return count, ok
}

func putToCache(from uint16, visited uint16, count int) {
    memo[key{from, visited}] = count
}

// Input parsing

func NewGraph(input []string) *Graph {
    g := newBuilder()
    g.parse(input)
    return &g.Graph
}

func newBuilder() *graphBuilder {
    return &graphBuilder{
        Graph: Graph{edges: make([][]uint16, 2, 16)},
        caves: map[string]uint16{
            "start": startCave,
            "end":   endCave,
        },
    }
}

type graphBuilder struct {
    Graph
    caves map[string]uint16
}

func (g *graphBuilder) parse(input []string) {
    for _, line := range input {
        edge := strings.Split(line, "-")
        c1 := g.getCave(edge[0])
        c2 := g.getCave(edge[1])
        g.addEdge(c1, c2)
    }
}

func (g *graphBuilder) addEdge(c1, c2 uint16) {
    if c1 != endCave && c2 != startCave {
        g.edges[c1] = append(g.edges[c1], c2)
    }
    if c1 != startCave && c2 != endCave {
        g.edges[c2] = append(g.edges[c2], c1)
    }
}

func (g *graphBuilder) getCave(name string) uint16 {
    if cave, ok := g.caves[name]; ok {
        return cave
    }
    cave := uint16(len(g.caves))
    g.caves[name] = cave
    if 'A' <= name[0] && name[0] <= 'Z' {
        g.bigCaves |= caveMask(cave)
    }
    g.edges = append(g.edges, nil)
    return cave
}