AoC 2021 day 16 (Smalltalk)

#!/usr/local/bin/gst -q

Collection extend [
    sum   [ ^self inject: 0 into: [:a : b | a + b] ]
    prod  [ ^self inject: 1 into: [:a : b | a * b] ]

    min   [ ^self inject: self first into: [:a :b | a min: b] ]
    max   [ ^self inject: self first into: [:a :b | a max: b] ]

    greater  [ ^(self first > self second) asCBooleanValue ]
    less     [ ^(self first < self second) asCBooleanValue ]
    equal    [ ^(self first = self second) asCBooleanValue ]
]

Integer extend [
    timesCollect: aBlock  [ ^(1 to: self) collect: [:x | aBlock value] ]
]

String extend [
    asRadix: rad [ ^Number readFrom: (ReadStream on: self) radix: rad ]
]

Object subclass: BitStream [
    | bits bitStream |

    BitStream class >> on: bitString [
        ^super new init: bitString
    ]

    init: bitString [
        bits      := bitString.
        bitStream := ReadStream on: bits.
        ^self
    ]

    " Streamlike access methods "
    atEnd  [ ^bitStream atEnd ]

    next             [ ^bitStream next digitValue        ]
    next: num        [ ^(bitStream next: num) asRadix: 2 ]
    nextString: num  [ ^bitStream next: num              ]

    " Read an extended literal value from the stream: "
    readExtended [
        | ret extend |
        ret := 0.
        [
            extend := self next.
            ret := ret * 16 + (self next: 4).
            (extend = 1)
        ] whileTrue.

        ^ret
    ]
]

Object subclass: PacketReader [
    | bits verSum |

    operators := Dictionary from: {
                     0 -> #sum.
                     1 -> #prod.
                     2 -> #min.
                     3 -> #max.
                     5 -> #greater.
                     6 -> #less.
                     7 -> #equal.
                 }.

    PacketReader class >> on: bitStream [
        ^super new init: bitStream
    ]

    init: bitStream [
        bits   := bitStream.
        verSum := 0.
        ^self
    ]

    nextPacket [
        | version type ret |
        " read header "
        verSum := verSum + (bits next: 3).
        type   := bits next: 3.

        (type = 4) ifTrue: [
            " read literal "
            ret := bits readExtended.
        ] ifFalse: [
            | lenType vals num |
            " read subpackets "
            lenType := bits next.

            (lenType = 0) ifTrue: [
                | sub |
                num := bits next: 15.  " num bits to process packets from "

                sub    := PacketReader on: (BitStream on: (bits nextString: num)).
                vals   := sub parsePackets.
                verSum := verSum + sub versionSum.

            ] ifFalse: [
                num  := bits next: 11. " num packets to collect "
                vals := num timesCollect: [self nextPacket].
            ].

            ret := vals perform: (operators at: type).
        ].
        ^ret
    ]

    parsePackets [
        | vals |
        vals := OrderedCollection new.
        [ bits atEnd not ] whileTrue: [
            vals add: self nextPacket.
        ].
        ^vals
    ]

    versionSum  [ ^verSum ]
]

"
| Mainline
"
input := stdin nextLine.

" Convert hex string into a bit string: "
bitString := ((input asRadix: 16) radix: 2) allButFirst: 2.

" Add padding if needed: "
padding := ((input findFirst: [:chr | chr ~= $0]) - 1) * 4.  " lead zeros "
padding := padding + ((4 - bitString size) \\ 4).            " lead zeros in first nybble "
bitString := (String new: padding withAll: $0), bitString.

reader := PacketReader on: (BitStream on: bitString).
value  := reader nextPacket.

('Part 1: %1' % {reader versionSum}) displayNl.
('Part 2: %1' % {value}) displayNl.