day 14: better code (maybe)

14.rkt

@@ -1,60 +1,46 @@
 #lang curly-fn racket
 
 (require "scripts/aoc.rkt")
+;; data/bit-vector is hella nerfed so.... rosette uwu
+;; rosette is slow af for this but whatever
+(require (only-in rosette bv bvand bvor bvxor rotate-left concat bitvector zero-extend
+                  bitvector->natural))
 
 ;; solution for day 14
+(define (mask->bv bits select)
+  (apply concat (map #{if (equal? % select) (bv select 1) (bv (- 1 select) 1)} bits)))
+
+(define (pos->bv pos)
+  (rotate-left (- 35 pos) (zero-extend (bv 1 1) (bitvector 36))))
 
 (define (part1 input)
-  (define (parse-mask msk)
-    (define orval (for/fold ([val 0]) ([chr (in-string msk)])
-                    (+ (* val 2) (if (equal? chr #\1) 1 0))))
-    (define andval (for/fold ([val 0]) ([chr (in-string msk)])
-                     (+ (* val 2) (if (equal? chr #\0) 0 1))))
-    (values orval andval))
-  (define orval 0)
-  (define andval (sub1 (expt 2 37)))
+  (define orval #f)
+  (define andval #f)
   (define mem (make-hash))
   (for ([ins (in-list input)])
     (match ins
-      [(cons 'mask mask)
-       (define-values [ov av] (parse-mask mask))
-       (set! orval ov) (set! andval av)]
-      [(cons 'set (cons addr val))
-       (define new-val (bitwise-ior (bitwise-and val andval) orval))
-       (hash-set! mem addr new-val)]))
+      [(list 'mask bits ...)
+       (set! orval (mask->bv bits 1))
+       (set! andval (mask->bv bits 0))]
+      [(list 'set addr val)
+       (hash-set! mem addr
+                  (bitvector->natural (bvor (bvand (bv val 36) andval) orval)))]))
   (for/sum ([(k v) (in-hash mem)])
     v))
 
-(define (addr->bits addr)
-  (vector-copy 
-  (for/vector ([i (in-range 36)])
-    (bitwise-and 1 (arithmetic-shift addr (+ -35 i))))))
-(define (bits->addr bits)
-  (for/fold ([val 0]) ([bit (in-vector bits)])
-    (+ (* val 2) bit)))
-
 (define (part2 input)
   (define mem (make-hash))
   (define mask #f)
   (for ([ins (in-list input)])
     (match ins
-      [(cons 'mask m) (set! mask m)]
-      [(cons 'set (cons addr val))
-       (define addr-bits (addr->bits addr))
-       (for ([chr (in-string mask)] [i (in-naturals)])
-         (when (equal? chr #\1)
-           (vector-set! addr-bits i 1)))
+      [(list 'mask bits ...) (set! mask bits)]
+      [(list 'set addr val)
+       (define addr-bits (bvor (bv addr 36) (mask->bv mask 1)))
        (define addrs-set (list addr-bits))
-       (for ([chr (in-string mask)] [i (in-naturals)])
-         (when (equal? chr #\X)
-           (define ones (map (lambda (x) (define y (vector-copy x)) (vector-set! y i 1) y)
-                             addrs-set))
-           (define zeros (map (lambda (x) (define y (vector-copy x)) (vector-set! y i 0) y)
-                             addrs-set))
-           (set! addrs-set (append ones zeros))))
+       (for ([item (in-list mask)] [i (in-naturals)] #:when (equal? 'x item))
+         (set! addrs-set (append addrs-set (map #{bvxor (pos->bv i) %} addrs-set))))
        (for ([addr (in-list addrs-set)])
-         (hash-set! mem (bits->addr addr) val))]))
-
+         (hash-set! mem (bitvector->natural addr) val))]))
   (for/sum ([(k v) (in-hash mem)])
     v))
 
@@ -63,9 +49,14 @@
   (define input (file->lines fname))
   (for/list ([i (in-list input)])
     (match i
-      [(pregexp #px"^mask = (.*?)$" (list _ mask)) (cons 'mask mask)]
+      [(pregexp #px"^mask = (.*?)$" (list _ mask))
+       (cons 'mask (for/list ([chr (in-string mask)])
+                     (match chr
+                       [#\1 1]
+                       [#\0 0]
+                       [#\X 'x])))]
       [(pregexp #px"^mem\\[([0-9]+)\\] = ([0-9]+)$" (list _ addr val))
-       (cons 'set (cons (string->number addr) (string->number val)))]
+       (list 'set (string->number addr) (string->number val))]
       [_ (error "not shonks i")])))
 
 (module+ test