Scala Native Optimised Advent of Code 2022 Neal Wu Solution

aoc-2021-day22-opt.scala

// using scala 3.0.2

package day22Wu


import scala.util.Using
import scala.io.Source

import scala.collection.mutable

import scala.util.chaining.*

import Command.*

@main def part1(): Unit =
  println(s"The solution is ${part1(readInput())}")

@main def part2(): Unit =
  println(s"The solution is ${part2(readInput())}")

def readInput(): String =
  Using.resource(Source.fromFile("day22"))(_.mkString)

case class Dimension(min: Int, max: Int):
  require(min <= max)

  def isSubset(d: Dimension): Boolean =
    min >= d.min && max <= d.max

extension (x1: Int)
  infix def by (x2: Int): Dimension = Dimension(x1, x2)

case class Cuboid(xs: Dimension, ys: Dimension, zs: Dimension)

enum Command:
  case On, Off

case class Step(command: Command, cuboid: Cuboid)

object Indexed:
  opaque type Index = Int

  object Index:

    def apply(x: Int, y: Int, z: Int): Index =
      val x0 = x & 0x3ff
      val y0 = y & 0x3ff
      val z0 = z & 0x3ff
      (x0 << 20) | (y0 << 10) | z0

    def decode(i: Int): Index = i

    extension (i: Index)
      def x: Int =
        ((i >> 20) & 0x3ff)
      def y: Int =
        ((i >> 10) & 0x3ff)
      def z: Int =
        (i & 0x3ff)
      def encoded: Int =
        i

import Indexed.*

inline def fastFor(inline begin: Int, inline end: Int)(
  inline body: Int => Unit): Unit =
  var i = begin
  while i < end do
    body(i)
    i += 1

// class Grid3D(size: Int):
//   val grid: Array[Array[Array[Boolean]]] = new Array(size)
//   for (i <- 0 until size)
//     grid(i) = new Array(size)
//     for (j <- 0 until size)
//       grid(i)(j) = new Array(size)

//   def put(x: Int, y: Int, z: Int, b: Boolean): Unit =
//     grid(x)(y)(z) = b

//   def get(x: Int, y: Int, z: Int): Boolean =
//     grid(x)(y)(z)

//   inline def foreach(inline f: (Int, Int, Int) => Unit): Unit =
//     fastFor(0, size) { x =>
//       fastFor(0, size) { y =>
//         fastFor(0, size) { z =>
//           if (grid(x)(y)(z))
//             f(x, y, z)
//         }
//       }
//     }

def run(steps: IndexedSeq[Step]): Long = {

  val X0, Y0, Z0 = mutable.ArrayBuffer.empty[Int]

  for Step(_, Cuboid(xs, ys, zs)) <- steps do
    X0 += xs.min
    X0 += xs.max + 1
    Y0 += ys.min
    Y0 += ys.max + 1
    Z0 += zs.min
    Z0 += zs.max + 1

  // println("saw steps")

  val X = IArray.from(X0).sorted
  val Y = IArray.from(Y0).sorted
  val Z = IArray.from(Z0).sorted

  val XL = X.zipWithIndex.toMap
  val YL = Y.zipWithIndex.toMap
  val ZL = Z.zipWithIndex.toMap

  val N = X.size

  val g = mutable.BitSet.empty
  // val gg = Grid3D(N)

  // println("sorted steps")

  for Step(tpe, Cuboid(xs, ys, zs)) <- steps do
    val x0 = XL(xs.min)
    val x1 = XL(xs.max + 1)
    val y0 = YL(ys.min)
    val y1 = YL(ys.max + 1)
    val z0 = ZL(zs.min)
    val z1 = ZL(zs.max + 1)
    val value = tpe == Command.On
    fastFor(x0, x1) { x =>
      fastFor(y0, y1) { y =>
        fastFor(z0, z1) { z =>
          val i = Index(x, y, z).encoded
          if value then g += i else g -= i
          // gg.put(x, y, z, value)
        }
      }
    }

  // println("ran steps")

  val sum =
    class Sum extends (Int => Unit):
      import Indexed.Index.*
      private var acc = 0L

      def add(i: Index): Unit =
        add(i.x, i.y, i.z)

      def add(x: Int, y: Int, z: Int): Unit =
        acc += 1L * (X(x + 1) - X(x)) * (Y(y + 1) - Y(y)) * (Z(z + 1) - Z(z))

      def apply(enc: Int): Unit = add(Index.decode(enc))

      def run =
        g.foreach(this)
        acc

    Sum().run

  sum
}

def isInit(cuboid: Cuboid): Boolean =
  Seq(cuboid.xs, cuboid.ys, cuboid.zs).forall(_.isSubset(-50 by 50))

type Parser[A] = PartialFunction[String, A]

val NumOf: Parser[Int] =
  case s if s.matches(raw"-?\d+") => s.toInt

val DimensionOf: Parser[Dimension] =
  case s"${NumOf(begin)}..${NumOf(end)}" => begin by end

val CuboidOf: Parser[Cuboid] =
  case s"x=${DimensionOf(xs)},y=${DimensionOf(ys)},z=${DimensionOf(zs)}" => Cuboid(xs, ys, zs)

val CommandOf: Parser[Command] =
  case "on" => On
  case "off" => Off

val StepOf: Parser[Step] =
  case s"${CommandOf(command)} ${CuboidOf(cuboid)}" => Step(command, cuboid)

def part1(input: String) =
  run(input.linesIterator.map(StepOf).filter(s => isInit(s.cuboid)).toIndexedSeq)

def part2(input: String) =
  run(input.linesIterator.map(StepOf).toIndexedSeq)