AoC 2021 day14 (smalltalk)

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

Character extend [
    , chr  [ ^self asString, chr asString ]
]

SequenceableCollection extend [
    " Returns collection of results from calling binBlock for each link pair "
    chain: binBlock [
        | res |
        res := self copyEmptyForCollect.
        self fold: [:curr :next | res add: (binBlock value: curr value: next). next].
        ^res
    ]
]

Object subclass: Polymer [
    | start rules |

    Polymer class >> new: start withRules: lines [
        ^super new init: start rules: lines
    ]

    init: str rules: lines [
        start := str asOrderedCollection.
        rules := Dictionary from: (lines collect: [ :line |
                     | parts |
                     parts := line tokenize: ' -> '.
                     parts first -> parts second first
                 ]).
        ^self.
    ]

    elementsAfter: steps [
        | pairs elements |
        pairs := Bag from: (start chain: [:a :b | a,b]).

        " Track number of each adjacent pair across the steps "
        steps timesRepeat: [
            | new_pairs |
            new_pairs := Bag new.

            pairs contents keysAndValuesDo: [ :pair :num |
                | insert |
                insert := rules at: pair.

                new_pairs add: (pair first,  insert) withOccurrences: num.
                new_pairs add: (insert, pair second) withOccurrences: num.
            ].

            pairs := new_pairs.
        ].

        " Convert pairs into counts of each element "
        elements := Bag from: {start first. start last}.
        pairs contents keysAndValuesDo: [ :pair :num |
            pair do: [ :el | elements add: el withOccurrences: num ]
        ].

        " each element counted twice (once from each side) so /2 "
        ^elements sortedByCount collect: [:a | (a key / 2) -> a value].
    ]
]

"
| Mainline
"
input := stdin lines contents.

poly := Polymer new: input first withRules: (input allButFirst: 2).

{ 1 -> 10. 2 -> 40 } do: [ :part |
    elem := poly elementsAfter: part value.
    ('Part %1: %2' % {part key. elem first key - elem last key}) displayNl.
].