adventofcode/2016/day09/main.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strconv"
)

var re *regexp.Regexp

func main() {
	var input []byte
	var noDecompress bool
	var printSize bool
	if len(os.Args) > 1 {
		for i := 1; i < len(os.Args); i++ {
			if os.Args[i][0] != '-' {
				// Assume it's a filename
				// Read from the given filename
				input = fileToByteSlice(os.Args[1])
			} else {
				switch os.Args[i] {
				case "-size":
					printSize = true
				case "-dry":
					// -dry implies printSize
					printSize = true
					noDecompress = true
				}
			}
		}
	}
	if len(input) == 0 {
		// Must be stdin
		input = stdinToByteSlice()
	}
	input = bytes.TrimSpace(input)
	// re matches compression markers and submatches the values we care about
	re = regexp.MustCompile(`\((\d*)x(\d*)\)`)

	if noDecompress {
		cnt := depthDecompressLength(input)
		fmt.Println("Total Decompressed Bytes:", cnt)
		os.Exit(0)
	}
	output := depthDecompress(input)
	fmt.Println(string(output))
	if printSize {
		fmt.Println("Total Decompressed Bytes:", len(output))
	}
}

// depthDecompressLength counts how many uncompressed bytes are
// in the byte slice by unravelling the compression levels
// recursively
func depthDecompressLength(cmp []byte) uint64 {
	mrkParts := re.FindStringSubmatch(string(cmp))
	if len(mrkParts) < 3 {
		// No compressionmarker, just return cmp length
		return uint64(len(cmp))
	}
	marker, mrkBytes, mrkDupe := []byte(mrkParts[0]), atoi(mrkParts[1]), atoi(mrkParts[2])
	mrkPos := bytes.Index(cmp, marker)
	if mrkPos > 0 {
		cmp = cmp[mrkPos:]
	}
	recurBytes := bytes.TrimPrefix(cmp, marker)[:mrkBytes]
	remainder := bytes.TrimPrefix(cmp, append(marker, recurBytes...))
	return uint64(mrkPos) + (depthDecompressLength(recurBytes) * uint64(mrkDupe)) + depthDecompressLength(remainder)
}

func depthDecompress(cmp []byte) []byte {
	mrkParts := re.FindStringSubmatch(string(cmp))
	if len(mrkParts) < 3 {
		// No compression marker, just return cmp length
		return cmp
	}
	marker, mrkBytes, mrkDupe := []byte(mrkParts[0]), atoi(mrkParts[1]), atoi(mrkParts[2])
	mrkPos := bytes.Index(cmp, marker)
	var ret []byte
	if mrkPos > 0 {
		ret = cmp[:mrkPos]
		cmp = cmp[mrkPos:]
	}
	recurBytes := bytes.TrimPrefix(cmp, marker)[:mrkBytes]
	remainder := bytes.TrimPrefix(cmp, append(marker, recurBytes...))

	ret = append(ret, bytes.Repeat(depthDecompress(recurBytes), mrkDupe)...)
	return append(ret, depthDecompress(remainder)...)
}

func fileToByteSlice(fn string) []byte {
	var c []byte
	var err error
	c, err = ioutil.ReadFile(fn)
	if err != nil {
		fmt.Println("Unable to read file: " + fn)
		os.Exit(1)
	}
	return c
}

func stdinToByteSlice() []byte {
	var input string
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		input += scanner.Text()
	}
	return []byte(input)
}

func atoi(i string) int {
	var ret int
	var err error
	if ret, err = strconv.Atoi(i); err != nil {
		log.Fatal("Invalid Atoi")
	}
	return ret
}