Ben / Nujel

#!/usr/bin/env nujel

[defun snail/read-line [line]
       [read/single [join [split line ","] " . "]]]

[defun snail/parse-list [l d]
       [cond [[nil? l] #nil]
             [[int? l] @[:v l]]
             [#t @[:r [snail/parse-list [cdr l]] :l [snail/parse-list [car l]]]]]]

[defun snail/read [line]
       [snail/parse-list [snail/read-line line]]]

[defun snail/int? [l]
       [and l [tree/get l :v]]]

[defun snail->str [l]
       [cond [[snail/int? l] [string [tree/get l :v]]]
             [#t [cat "[" [snail->str [tree/get l :l]]  "," [snail->str [tree/get l :r]] "]"]]]]

[defun snail/find [root needle as bs]
       [def last #nil]
       [defun rec [t root needle]

              [cond [[snail/int? t] [set! last t] #f]
                    [[== needle t] last]
                    [#t [or [rec [tree/get t as] root needle]
                            [rec [tree/get t bs] root needle]]]]]
       [rec root root needle]]

[defun snail/find/previous [root needle]
       [snail/find root needle :l :r]]

[defun snail/find/next [root needle]
       [snail/find root needle :r :l]]

[defun snail/explode [t root]
       [def p [snail/find/previous root t]]
       [when p [tree/+= p :v [tree/get [tree/get t :l] :v]]]
       [def n [snail/find/next root t]]
       [when n [tree/+= n :v [tree/get [tree/get t :r] :v]]]
       [-> t
           [tree/set! :v 0]
           [tree/set! :l #nil]
           [tree/set! :r #nil]]]

[defun snail/split [t root]
       [def lv [/ [tree/get t :v] 2]]
       [-> t
           [tree/set! :l @[:v lv]]
           [tree/set! :r @[:v [- [tree/get t :v] lv]]]
           [tree/set! :v #nil]]]

[defun snail/explode/maybe [t root d]
       [cond [[snail/int? t] #f]
             [[>= d 4] [snail/explode t root]]
             [[snail/explode/maybe [tree/get t :l] root [+ 1 d]] #t]
             [#t [snail/explode/maybe [tree/get t :r] root [+ 1 d]]]]]

[defun snail/split/maybe [t root d]
       [cond [[snail/int? t] [and [>= [tree/get t :v] 10]
                                  [snail/split t root] #t]]
             [[snail/split/maybe [tree/get t :l] root [+ 1 d]] #t]
             [#t [snail/split/maybe [tree/get t :r] root [+ 1 d]]]]]

[defun snail/reduce [t]
       [while [or [snail/explode/maybe t t 0]
                  [snail/split/maybe t t 0]]] t]

[defun snail/add [a b]
       [snail/reduce @[:l a :r b]]]

[defun snail/sum [lines]
       [def vals [map [split lines "\n"] snail/read]]
       [reduce [cdr vals] snail/add [car vals]]]

[defun snail/mag [t]
       [cond [[snail/int? t] [tree/get t :v]]
             [#t [+ [* [snail/mag [tree/get t :l]] 3]
                    [* [snail/mag [tree/get t :r]] 2]]]]]

[defun snail/find-biggest [lines]
       [def vals [apply array/new [split lines "\n"]]]
       [def cmax 0]
       [for [a 0 100]
       [for [b 0 100]
            [when [!= a b]
                  [def cs [snail/add [snail/read [array/ref vals a]] [snail/read [array/ref vals b]]]]
                  [set! cmax [max cmax [snail/mag cs]]]
                  [def cs [snail/add [snail/read [array/ref vals b]] [snail/read [array/ref vals a]]]]
                  [set! cmax [max cmax [snail/mag cs]]]]]]
       cmax]
[optimize-all!]


[def res-p1 [snail/mag [snail/sum [file/read "tests/slow/day18.input"]]]]
[when [!= res-p1 3725]
      [throw [list :wrong-result "Wrong result" res-p1]]]
[def res-p2 [snail/find-biggest [file/read "tests/slow/day18.input"]]]
[when [!= res-p2 4832]
      [throw [list :wrong-result "Wrong result" res-p2]]]