! Nothing too elegant today, I'm afraid. But it gets the
! job done!

USING: arrays assocs assocs.extras formatting io.encodings.ascii
io.files kernel literals math.combinatorics math.parser
namespaces prettyprint sequences splitting ;
IN: aoc.2020.14

CONSTANT: input $[
    "input.txt" ascii file-lines
    [
        "maskmem[] =" split harvest dup length 2 =
        [ first2 [ dec> ] bi@ "%036b" sprintf 2array ] when
    ] map
]

SYMBOLS: mask memory ;

: initialize-memory ( -- ) H{ } clone memory set ;

: set-mask ( seq -- ) first mask set ;

: (mask) ( seq seq -- seq )
    [ dup CHAR: X = [ drop ] [ nip ] if ] 2map ;

: assign ( seq -- )
    first2 mask get (mask) bin> swap memory get set-at ;

: process-line ( seq -- )
    dup length 1 = [ set-mask ] [ assign ] if ;

: part1 ( -- )
    initialize-memory input [ process-line ] each
    memory get sum-values . ;


! ===== PART 2 =====

CONSTANT: input2 $[
    "input.txt" ascii file-lines
    [
        "maskmem[] =" split harvest dup length 2 =
        [ first2 [ dec> ] bi@ [ "%036b" sprintf ] dip 2array ]
        when
    ] map
]

: (mask2) ( seq seq -- seq )
    [ dup CHAR: 0 = [ drop ] [ nip ] if ] 2map ;

: slot ( seq seq -- seq )
    clone tuck [ CHAR: X swap [ index ] keep set-nth ] curry each ;

: assign2 ( seq -- )
    first2 swap mask get (mask2) dup [ CHAR: X = ] count "10"
    swap selections dup length rot <repetition> [ slot bin> ]
    2map [ memory get set-at ] with each ;

: process-line2 ( seq -- )
    dup length 1 = [ set-mask ] [ assign2 ] if ;

: part2 ( -- )
    initialize-memory input2 [ process-line2 ] each
    memory get sum-values . ;