AoC 2021 day 18 (smalltalk)

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

Object subclass: SFNode [
    | parent value left right |
    " Constructors "
    SFNode class >> new: aStream parent: node [
        ^super new initFrom: aStream parent: node
    ]

    initFrom: aStream parent: node [
        parent := node.
        ^self parseNode: aStream.
    ]

    SFNode class >> left: lnode right: rnode [
        ^super new init: lnode with: rnode
    ]

    init: lnode with: rnode [
        parent := nil.
        left  := lnode parent: self; yourself.
        right := rnode parent: self; yourself.
        ^self
    ]

    SFNode class >> leaf: val parent: node [
        ^super new initLeaf: val parent: node
    ]

    initLeaf: val parent: node [
        parent := node.
        value  := val.
        left   := nil.
        right  := nil.
        ^self
    ]

    " Tree parser and builder "
    parseNode: aStream [
        | chr |
        [ aStream atEnd not ] whileTrue: [
            chr := aStream next.

            " NOTE: Assuming single digit numbers only in input! "
            (chr isDigit) ifTrue: [ ^SFNode leaf: chr digitValue parent: self parent ].
            (chr = $])    ifTrue: [ ^self ].

            (chr = $[) ifTrue: [
                left  := SFNode new: aStream parent: self.
            ] ifFalse: [ " chr = $, "
                right := SFNode new: aStream parent: self.
            ]
        ].
        ^self
    ]

    " Access methods "
    left    [ ^left   ]
    right   [ ^right  ]
    parent  [ ^parent ]
    value   [ ^value  ]

    left:  val  [ left   := val ]
    right: val  [ right  := val ]
    parent: val [ parent := val ]
    value: val  [ value  := val ]

    isLeaf  [ ^value notNil ]

    magnitude [
        (self isLeaf) ifTrue: [ ^value ].
        ^(3 * left magnitude) + (2 * right magnitude)
    ]

    depth [
        | dep curr |
        dep := -1.
        curr := self.
        [curr notNil] whileTrue: [
            dep := dep + 1.
            curr := curr parent.
        ].
        ^dep
    ]

    " For display "
    printOn: aStream [
        (self isLeaf) ifTrue: [
            aStream nextPutAll: value asString.
        ] ifFalse: [
            aStream nextPutAll: '['.
            left printOn: aStream.
            aStream nextPutAll: ','.
            right printOn: aStream.
            aStream nextPutAll: ']'.
        ]
    ]
]

Object subclass: SnailfishNumber [
    | root |
    SnailfishNumber class >> from: aString [
        ^super new init: (SFNode new: (ReadStream on: aString) parent: nil)
    ]

    SnailfishNumber class >> left: snLeft right: snRight [
        ^super new init: (SFNode left: snLeft right: snRight)
    ]

    init: node [
        root := node.
        ^self
    ]

    root      [ ^root ]
    magnitude [ ^root magnitude ]

    + sfNum [
        ^(SnailfishNumber left: root right: sfNum root) reduce
    ]

    " Walk tree looking for first explode.  Returns true if valid, false if explode done. "
    validatorExplode [
        | curr stack prevLeaf explode |
        stack   := OrderedCollection with: root.
        explode := nil.

        [stack notEmpty] whileTrue: [
            curr := stack removeLast.

            (curr isLeaf) ifTrue: [
                (explode) ifNotNil: [
                    " Exploding!  Finish and quit. "
                    curr value: (curr value + explode).
                    ^false
                ] ifNil: [
                    " Not exploding!  Track most recent leaf in case we do. "
                    prevLeaf := curr.
                ]

            ] ifFalse: [
                ((curr depth >= 4) and: [explode isNil]) ifTrue: [
                    " Exploding current node "
                    " Add to previous if we've seen one: "
                    (prevLeaf) ifNotNil: [
                        prevLeaf value: (prevLeaf value + curr left value)
                    ].

                    " Mark that we're in exploding state with value to add to next "
                    explode := curr right value.

                    " Replace current node with 0 leaf node "
                    curr value: 0; left: nil; right: nil.
                ]
            ].

            (curr right) ifNotNil: [stack addLast: curr right].
            (curr left)  ifNotNil: [stack addLast: curr left ].
        ].

        ^(explode isNil)
    ]

    " Walk tree looking for first split.  Returns true if valid, false if split done. "
    validatorSplit [
        | curr stack |
        stack   := OrderedCollection with: root.

        [stack notEmpty] whileTrue: [
            curr := stack removeLast.

            (curr isLeaf) ifTrue: [
                | val |
                val := curr value.
                (val > 9) ifTrue: [
                    " Splitting current node "
                    curr left:  (SFNode leaf: (val / 2) floor   parent: curr);
                         right: (SFNode leaf: (val / 2) ceiling parent: curr);
                         value: nil.
                    ^false
                ]
            ].

            (curr right) ifNotNil: [stack addLast: curr right].
            (curr left)  ifNotNil: [stack addLast: curr left ].
        ].

        ^true
    ]

    reduce [
        [
            (self validatorExplode) and: [self validatorSplit]
        ] whileFalse.
    ]
]

"
| Mainline
"
input := stdin lines contents.

" Part 1 "
sum := (input collect: [:line | SnailfishNumber from: line]) fold: [:a :b | a + b].

('Final sum: %1' % {sum root}) displayNl.
('Part 1: %1' % {sum magnitude}) displayNl.

" Part 2 "
count := 0.
maxMag := 0.
input do: [ :i |
    stderr nextPutAll: ('Working: %1' % {count := count + 1}); cr; flush.
    input do: [ :j |
        (i ~= j) ifTrue: [
            maxMag := maxMag max: ((SnailfishNumber from: i)
                                    + (SnailfishNumber from: j)) magnitude.
        ]
    ]
].

stdout nl.
('Part 2: %1' % {maxMag}) displayNl.