Ben / Nujel

#!/usr/bin/env nujel

[defn bits/hex->bits [s]
      [string/pad-start [int->string/binary [read/single [cat "#x" s]]] 4 "0"]]

[defn bits/read-value [s offset len]
      [read/single [cat "#b" [string/cut s offset [+ offset len]]]]]

[defn bits/mul vals
      [reduce vals * 1]]

[defn bits/> [a b] [if [> a b] 1 0]]
[defn bits/< [a b] [if [< a b] 1 0]]
[defn bits/= [a b] [if [= a b] 1 0]]
[defn bits/type->op [t]
      [case t
            [0 +]
	    [1 bits/mul]
	    [2 min]
	    [3 max]
	    [5 bits/>]
	    [6 bits/<]
	    [7 bits/=]]]
[defn bits/compile [p]
      [if [= 4 [ref p :type]]
	  [ref p :value]
	  [cons [bits/type->op [ref p :type]]
                [map [ref p :value] bits/compile]]]]

[defn packet/operator/packets [p s offset packets]
      [def off offset]
      [dotimes [i packets]
	       [def child [bits/parse/packet s off]]
	       [tree/set! p :value [cons [cdr child] [ref p :value]]]
	       [set! off [car child]]]
      [tree/set! p :value [nreverse [ref p :value]]]
      [cons off p]]

[defn packet/operator/bits [p s offset bits]
      [def off offset]
      [def goal-offset [+ offset bits]]

      [while [< off goal-offset]
	     [def child [bits/parse/packet s off]]
	     [tree/set! p :value [cons [cdr child] [ref p :value]]]
	     [set! off [car child]]]
      [tree/set! p :value [nreverse [ref p :value]]]
      [cons off p]]

[defn packet/operator [p s offset]
      [if [= [buffer/ref s offset] #\1]
          [packet/operator/packets p s [+ offset 12] [bits/read-value s [+ 1 offset] 11]]
          [packet/operator/bits p s [+ offset 16] [bits/read-value s [+ 1 offset] 15]]]]

[defn packet/literal [p s offset]
      [def parts #nil]
      [def off offset]
      [while [= [buffer/ref s off] #\1]
	     [inc! off]
	     [set! parts [cons [string/cut s off [+ 4 off]] parts]]
	     [+= off 4]]
      [inc! off]
      [set! parts [cons [string/cut s off [+ 4 off]] parts]]
      [inc! off 4]
      [cons off [tree/set! p :value [read/single [cat "#b" [join [nreverse parts]]]]]]]

[defn bits/parse/packet-contents [p s offset]
      [case [ref p :type]
	    [4 [packet/literal p s offset]]
	    [otherwise [packet/operator p s offset]]]]

[defn bits/parse/packet [s offset]
      [-> @[:version [bits/read-value s offset 3]
                     :type [bits/read-value s [+ 3 offset] 3]]
	  [bits/parse/packet-contents s [+ offset 6]]]]

[defn bits/parse/hex [s]
      [bits/parse/packet [join [map [split s ""] bits/hex->bits]] 0]]

[defn bits/version-sum [p]
      [if [pair? [ref p :value]]
          [reduce [ref p :value] [fn [a b] [+ a [bits/version-sum b]]] [ref p :version]]
	  [ref p :version]]]

[defn bits/print [p indent]
      [when-not indent [set! indent 0]]
      [case [ref p :type]
	    [4 [println [cat [string/pad-start "" indent]
			     [ansi-blue [ref p :value]]
			     " v:" [ansi-green [ref p :version]]]]]
	    [otherwise [println [cat [string/pad-start "" indent]
				     [ansi-yellow [ref p :type]]
				     " v:" [ansi-green [ref p :version]]]]
		       [for-each [ref p :value] [fn [a] [bits/print a [+ 4 indent]]]]]]]

[def res-p1 [bits/version-sum [cdr [bits/parse/hex [file/read "tests/fast/day16.dat"]]]]]
[def res-p2 [eval [bits/compile [cdr [bits/parse/hex [file/read "tests/fast/day16.dat"]]]]]]
[when [not= res-p1 951]
      [throw [list :wrong-result "Wrong result" res-p1]]]
[when [not= res-p2 902198718880]
      [throw [list :wrong-result "Wrong result" res-p2]]]
[return :success]