advent-of-code/aoc2024/day03/main.go

package main

import (
	_ "embed"
	"fmt"
	"io"
	"iter"
	"strings"

	aoc "go.sour.is/advent-of-code"
)

// var log = aoc.Log

func main() { aoc.MustResult(aoc.RunnerReader(run)) }

type result struct {
	valuePT1 int
	valuePT2 int
}

func (r result) String() string { return fmt.Sprintf("%#v", r) }

func run(read io.Reader) (*result, error) {
	sum := 0
	sum2 := 0

	lexer := Lexer(read)
	active := true

	for lexer.NextTok() {
		switch lexer.token {
		case "TokMUL":
			if m := readMul(lexer); m != 0 {
				sum += m
				if active {
					sum2 += m
				}
			}
		case "TokDONT":
			if readDont(lexer) {
				active = false
			}

		case "TokDO":
			if readDo(lexer) {
				active = true
			}
		}
	}

	return &result{sum, sum2}, nil
}

func Lexer(in io.Reader) *lexer {
	seq := func(yield func(rune) bool) {
		buf := make([]byte, 256)
		s, _ := in.Read(buf)
		buf = buf[:s]

		for len(buf) > 0 {
			for _, r := range string(buf) {
				if !yield(r) {
					return
				}
			}
			s, _ := in.Read(buf)
			buf = buf[:s]
		}

		for !yield(-1) {
		}
	}
	next, stop := iter.Pull(seq)
	lex := lexer{iter: next, stop: stop}
	lex.readRune()
	lex.readRune()
	return &lex
}

type lexer struct {
	rune rune
	next rune

	iter func() (rune, bool)
	stop func()

	buf     []rune
	token   string
	literal []rune
}

func (l *lexer) Stop() { l.stop() }

func (l *lexer) readRune() {
	if l.rune == -1 {
		return
	}
	if r, ok := l.iter(); ok {
		l.rune, l.next = l.next, r
	} else {
		l.rune, l.next = l.next, -1
	}
}

func (l *lexer) loadRune(tok string) {
	l.token = tok
	l.literal = append(l.literal, l.rune)
	l.readRune()
}

func (l *lexer) loadNumber() {
	l.token = "TokNUMBER"
	for strings.ContainsRune("0123456789", l.rune) {
		l.literal = append(l.literal, l.rune)
		l.readRune()
	}
}

func (l *lexer) loadString(accept string) {
	l.token = "TokSTRING"
	for !(!strings.ContainsRune(accept, l.rune) || l.rune == 0 || l.rune == -1) {
		l.literal = append(l.literal, l.rune)
		l.readRune()
	}
}

func (l *lexer) NextTok() bool {
	l.literal = l.literal[:0]

	switch l.rune {
	case 'm':
		l.loadString("mul")
		l.token = "TokMUL"
		if string(l.literal) != "mul" {
			l.token = "TokILLEGAL"
		}
		return true
	case 'd':
		l.loadString("don't")
		switch string(l.literal) {
		case "don't":
			l.token = "TokDONT"
		case "do":
			l.token = "TokDO"
		default:
			l.token = "TokILLEGAL"
		}
		return true

	case '(':
		l.loadRune("TokLPAREN")
		return true
	case ')':
		l.loadRune("TokRPAREN")
		return true
	case ',':
		l.loadRune("TokCOMMA")
		return true

	case -1:
		l.loadRune("TokEOF")
		return false
	default:
		if '0' <= l.rune && l.rune <= '9' {
			l.loadNumber()
			return true
		}
	}

	l.loadRune("TokILLEGAL")
	return true
}

func (l *lexer) Iter() iter.Seq2[string, string] {
	return func(yield func(string, string) bool) {
		for l.NextTok() {
			if !yield(l.token, string(l.literal)) {
				return
			}
		}
	}
}

func readMul(lex *lexer) int {
	if lex.token != "TokMUL" || string(lex.literal) != "mul" {
		return 0
	}

	var a, b = -1, -1
	if !lex.NextTok() || lex.token != "TokLPAREN" {
		return 0
	}

	if !lex.NextTok() || lex.token != "TokNUMBER" {
		return 0
	}
	a = aoc.Atoi(string(lex.literal))

	if !lex.NextTok() || lex.token != "TokCOMMA" {
		return 0
	}

	if !lex.NextTok() || lex.token != "TokNUMBER" {
		return 0
	}
	b = aoc.Atoi(string(lex.literal))

	if !lex.NextTok() || lex.token != "TokRPAREN" {
		return 0
	}
	return a * b
}

func readDont(lex *lexer) bool {
	if lex.token != "TokDONT" || string(lex.literal) != "don't" {
		return false
	}
	if !lex.NextTok() || lex.token != "TokLPAREN" {
		return false
	}
	if !lex.NextTok() || lex.token != "TokRPAREN" {
		return false
	}
	return true
}

func readDo(lex *lexer) bool {
	if lex.token != "TokDO" || string(lex.literal) != "do" {
		return false
	}
	if !lex.NextTok() || lex.token != "TokLPAREN" {
		return false
	}
	if !lex.NextTok() || lex.token != "TokRPAREN" {
		return false
	}
	return true
}