advent-of-code-2022/day17/common/tetris.go

package common

import (
	"bufio"
	"fmt"
)


type Rock byte

const (
  Minus Rock = iota
  Plus
  Angle
  Line
  Square
)

func GetRock(rockType int) Rock {
  switch rockType {
    case 0:
      return Minus
    case 1:
      return Plus
    case 2:
      return Angle
    case 3:
      return Line
    case 4:
      return Square
  }
  return Minus
}


type Tile byte

const (
  Empty Tile = iota
  MovingRock
  StoppedRock
)


type PreviousAction struct {
  endJet int
  heightDifference int
}


type Chamber struct {
  tiles [][]Tile
  highestPoint int
  offset int
  nRocks int
  jets []Jet
  currentJet int
  cache map[Rock]map[int]PreviousAction
}

func NewChamber(jets []Jet) *Chamber {
  return &Chamber{[][]Tile{}, -1, 0, 0, jets, 0, make(map[Rock]map[int]PreviousAction)}
}

func (chamber Chamber) Clone() *Chamber {
  return &Chamber{[][]Tile{}, -1, 0, 0, chamber.jets, 0, make(map[Rock]map[int]PreviousAction)}
}

func (chamber *Chamber) DropRock() {
  for i := len(chamber.tiles); i <= chamber.highestPoint + 7 - chamber.offset; i++ {
    chamber.tiles = append(chamber.tiles, []Tile{Empty, Empty, Empty, Empty, Empty, Empty, Empty})
  }

  newRock := GetRock(chamber.nRocks % 5)
  chamber.nRocks++
  startJet := chamber.currentJet
  currHighestPoint := chamber.highestPoint - chamber.offset

  switch newRock {
    case Minus:
      chamber.tiles[len(chamber.tiles) - 4][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 4][3] = MovingRock
      chamber.tiles[len(chamber.tiles) - 4][4] = MovingRock
      chamber.tiles[len(chamber.tiles) - 4][5] = MovingRock
    case Plus:
      chamber.tiles[len(chamber.tiles) - 4][3] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][3] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][4] = MovingRock
      chamber.tiles[len(chamber.tiles) - 2][3] = MovingRock
    case Angle:
      chamber.tiles[len(chamber.tiles) - 4][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 4][3] = MovingRock
      chamber.tiles[len(chamber.tiles) - 4][4] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][4] = MovingRock
      chamber.tiles[len(chamber.tiles) - 2][4] = MovingRock
    case Line:
      chamber.tiles[len(chamber.tiles) - 4][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 2][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 1][2] = MovingRock
    case Square:
      chamber.tiles[len(chamber.tiles) - 4][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 4][3] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][2] = MovingRock
      chamber.tiles[len(chamber.tiles) - 3][3] = MovingRock
  }

  moving := true
  for moving {
    canShift := true
    jet := chamber.jets[chamber.currentJet]
    chamber.currentJet = (chamber.currentJet + 1) % len(chamber.jets)

    rowsChecked := 0
    var lowestPoint int
    for i := len(chamber.tiles) - 1; i >= 0 && rowsChecked != 0b01111111; i-- {
      for j := 0; j < len(chamber.tiles[i]); j++ {
        if chamber.tiles[i][j] == MovingRock {
          if jet == LeftToRight {
            if j == 6 || chamber.tiles[i][j+1] == StoppedRock {
              canShift = false
            }
          } else if jet == RightToLeft {
            if j == 0 || chamber.tiles[i][j-1] == StoppedRock {
              canShift = false
            }
          }
        } else if chamber.tiles[i][j] == StoppedRock {
          rowsChecked |= 0b01 << j
        }
      }
      if rowsChecked == 0b01111111 {
        lowestPoint = i
      }
    }

    for i := lowestPoint; i < len(chamber.tiles); i++ {
      if jet == LeftToRight {
        for j := len(chamber.tiles[i]) - 1; j >= 0; j-- {
          if chamber.tiles[i][j] == MovingRock {
            if canShift {
              chamber.tiles[i][j] = Empty
              chamber.tiles[i][j+1] = MovingRock
              if i == 0 || chamber.tiles[i-1][j+1] == StoppedRock {
                moving = false
              }
            } else {
              if i == 0 || chamber.tiles[i-1][j] == StoppedRock {
                moving = false
              }
            }
          }
        }
      } else if jet == RightToLeft {
        for j := 0; j < len(chamber.tiles[i]); j++ {
          if chamber.tiles[i][j] == MovingRock {
            if canShift {
              chamber.tiles[i][j] = Empty
              chamber.tiles[i][j-1] = MovingRock
              if i == 0 || chamber.tiles[i-1][j-1] == StoppedRock {
                moving = false
              }
            } else {
              if i == 0 || chamber.tiles[i-1][j] == StoppedRock {
                moving = false
              }
            }
          }
        }  
      }
    }

    for i := lowestPoint; i < len(chamber.tiles); i++ {
      for j := 0; j < len(chamber.tiles[i]); j++ {
        if chamber.tiles[i][j] == MovingRock {
          if moving {
            chamber.tiles[i][j] = Empty
            chamber.tiles[i-1][j] = MovingRock
          } else {
            chamber.tiles[i][j] = StoppedRock
            if i > currHighestPoint {
              currHighestPoint = i
            }
          }
        }
      }
    }
  }

  if chamber.nRocks > len(chamber.jets) {
    if _, ok := chamber.cache[newRock]; !ok {
      chamber.cache[newRock] = make(map[int]PreviousAction)
    }
    chamber.cache[newRock][startJet] = PreviousAction{chamber.currentJet, currHighestPoint - chamber.highestPoint}
  }
  chamber.highestPoint = currHighestPoint + chamber.offset
}

func (chamber *Chamber) DropNRocks(nRocks int) {
  var loopNewRock Rock
  loopCurrentJet := -1
  loopEndJet := -1
  loopNRocks := 0
  loopHeightDifference := 0
  loopFound := false

  var i int
  for i = 0; i < nRocks && !loopFound; i++ {
    newRock := GetRock(chamber.nRocks % 5)
    found := false
    if jets, ok := chamber.cache[newRock]; ok {
      if result, ok := jets[chamber.currentJet]; ok {
        if loopCurrentJet < 0 {
          loopCurrentJet = chamber.currentJet
          loopNewRock = newRock
        }
        if loopNRocks > 0 && loopNewRock == newRock && loopCurrentJet == chamber.currentJet {
          loopFound = true
        } else {
          loopNRocks++
          loopHeightDifference += result.heightDifference
          loopEndJet = result.endJet
          chamber.nRocks++
          chamber.offset += result.heightDifference
          chamber.highestPoint += result.heightDifference
          chamber.currentJet = result.endJet
        }
        found = true
      }
    }
    if !found {
      chamber.DropRock()
    }
  }

  if loopFound {
    quotientRocks := int((nRocks - i) / loopNRocks)
    chamber.nRocks += quotientRocks * loopNRocks
    chamber.highestPoint += quotientRocks * loopHeightDifference
    chamber.offset += quotientRocks * loopHeightDifference
    chamber.currentJet = loopEndJet
    i += quotientRocks * loopNRocks

    for ; i <= nRocks; i++ {
      chamber.DropRock()
    }
  }
}

func (chamber *Chamber) GetHighestPoint() int {
  return chamber.highestPoint + 1
}

func (chamber *Chamber) Print() {
  for i := len(chamber.tiles) - 1; i >= 0; i-- {
    fmt.Print("|")
    for j := 0; j < len(chamber.tiles[i]); j++ {
      switch chamber.tiles[i][j] {
        case Empty:
          fmt.Print(".")
        case MovingRock:
          fmt.Print("@")
        case StoppedRock:
          fmt.Print("#")
      }
    }
    fmt.Println("|")
  }
  fmt.Println("+-------+")
}


type Jet byte

const (
  LeftToRight Jet = iota
  RightToLeft
)


func Parse(scanner bufio.Scanner) Chamber {
  jets := []Jet{}

  scanner.Scan()
  line := scanner.Text()
  for _, c := range(line) {
    if c == '>' {
      jets = append(jets, LeftToRight)
    } else if c == '<' {
      jets = append(jets, RightToLeft)
    }
  }

  return *NewChamber(jets)
}