aoc2024/day-09/internal/two/disk.go

package two

import (
	"cmp"
	"container/heap"
	"errors"
	"math"
	"slices"

	log "github.com/sirupsen/logrus"
)

const numEmptyBlocks int = 9

type disk struct {
	data        []int
	fileblocks  []*block
	emptyblocks [numEmptyBlocks]minHeap
}

func newDisk(raw []int) *disk {
	fileblocks := make([]*block, 0)
	emptyblockHeaps := make([]minHeap, numEmptyBlocks)
	for i := range numEmptyBlocks {
		heap.Init(&emptyblockHeaps[i])
	}

	var sz int
	for i, id := 0, 0; i < len(raw); id++ {
		sz = 0
		for j := i; j < len(raw) && raw[j] == id; j++ {
			sz++
		}
		fileblocks = append(fileblocks, newBlock(kindOfFile, id, i, sz, id))

		i += sz

		sz = 0
		for j := i; j < len(raw) && raw[j] == empty; j++ {
			sz++
		}
		if sz > 0 {
			heap.Push(&emptyblockHeaps[sz-1], newBlock(kindOfEmpty, id, i, sz, empty))
		}

		i += sz
	}

	log.Debugf("\n%v\n%v", fileblocks, emptyblockHeaps)

	return &disk{data: raw, fileblocks: fileblocks, emptyblocks: [numEmptyBlocks]minHeap(emptyblockHeaps)}
}

func (d *disk) defragment() {
	for i := len(d.fileblocks) - 1; i >= 0; i-- {
		log.Debugf("searching for space for fileblock %d: %v", i, d.fileblocks[i])

		emptyblock, err := d.getNextEmptyBlock(d.fileblocks[i])
		if err != nil {
			log.Debug(err)
			continue
		}

		oldStart := d.fileblocks[i].start
		oldLength := d.fileblocks[i].length

		// we've found an appropriate empty block, now swap the data
		d.fileblocks[i].start, emptyblock.start = emptyblock.start, emptyblock.start+d.fileblocks[i].length

		// push the resized empty block to a new heap
		emptyblock.length -= d.fileblocks[i].length
		if emptyblock.length > 0 {
			log.Debugf("emptyblock resized %d", emptyblock.length)
			heap.Push(&d.emptyblocks[emptyblock.length-1], emptyblock)
		}

		// now create a new empty block and push it to the appropriate heap
		heap.Push(
			&d.emptyblocks[oldLength-1],
			newBlock(kindOfEmpty, math.MaxInt, oldStart, oldLength, empty),
		)
	}
}

func (d *disk) getNextEmptyBlock(currentFile *block) (*block, error) {
	// collect all minblocks the same size as the current file or greater
	minBlocks := []*block{}
	for j := currentFile.length; j <= numEmptyBlocks; j++ {
		if d.emptyblocks[j-1].Len() == 0 {
			continue
		}

		currentblock := heap.Pop(&d.emptyblocks[j-1]).(*block)
		minBlocks = append(minBlocks, currentblock)
	}

	if len(minBlocks) == 0 {
		return nil, errors.New("no empty blocks found")
	}

	// sort the blocks by id
	slices.SortFunc(minBlocks, func(a, b *block) int {
		return cmp.Compare(a.id, b.id)
	})
	// push back the ones we won't be using
	for _, block := range minBlocks[1:] {
		heap.Push(&d.emptyblocks[block.length-1], block)
	}

	// if the lowest id minblock is positioned after the current file, push it back and return an error
	if minBlocks[0].start >= currentFile.start {
		heap.Push(&d.emptyblocks[minBlocks[0].length-1], minBlocks[0])
		return nil, errors.New("no empty blocks found")
	}

	log.Debugf("found empty space %v", minBlocks[0])

	return minBlocks[0], nil
}

func (d *disk) write() {
	allBlocks := make([]*block, 0)

	allBlocks = append(allBlocks, d.fileblocks...)

	for i := 0; i < numEmptyBlocks; i++ {
		for range d.emptyblocks[i].Len() {
			allBlocks = append(allBlocks, heap.Pop(&d.emptyblocks[i]).(*block))
		}
	}

	slices.SortFunc(allBlocks, func(a, b *block) int {
		return cmp.Compare(a.start, b.start)
	})

	for _, block := range allBlocks {
		for i := block.start; i < block.start+block.length; i++ {
			d.data[i] = block.value
		}
	}
}