Paste: aoc21

! read your input, modify the input to optimize the program and compute the idiv in 1 step instead of O(n) step with the new divi insn 

: 2ary-op-r ( regs inputs quot: ( a b -- c ) -- outputs )
  [ clone ] 2dip [ '[ [ second swap nth ] [ third swap nth ] 2bi @ ]
   [ fourth swap set-nth ] 2bi ] 2curry keep ; inline
: 2ary-op-i ( regs inputs quot: ( a b -- c ) -- outputs )
  [ clone ] 2dip [ '[ [ second swap nth ] [ third ] bi @ ]
  [ fourth swap set-nth ] 2bi ] 2curry keep ; inline
: 2ary-op-ir ( regs inputs quot: ( a b -- c ) -- outputs )
  [ clone ] 2dip [ '[ [ second swap ] [ third swap nth ] bi @ ]
  [ fourth swap set-nth ] 2bi ] 2curry keep ; inline


: addr ( regs inputs -- outputs ) [ + ] 2ary-op-r ;
: addi ( regs inputs -- outputs ) [ + ] 2ary-op-i ;
: mulr ( regs inputs -- outputs ) [ * ] 2ary-op-r ;
: muli ( regs inputs -- outputs ) [ * ] 2ary-op-i ;
: divr ( regs inputs -- outputs ) [ /i ] 2ary-op-r ;
: divi ( regs inputs -- outputs ) [ /i ] 2ary-op-i ;
: banr ( regs inputs -- outputs ) [ bitand ] 2ary-op-r ;
: bani ( regs inputs -- outputs ) [ bitand ] 2ary-op-i ;
: borr ( regs inputs -- outputs ) [ bitor ] 2ary-op-r ;
: bori ( regs inputs -- outputs ) [ bitor ] 2ary-op-i ;
: setr ( regs inputs -- outputs )
  [ clone ] dip [ [ second swap nth ] [ fourth swap set-nth ] 2bi ] curry keep ;
: seti ( regs inputs -- outputs )
  [ clone ] dip [ [ nip second ] [ fourth swap set-nth ] 2bi ] curry keep ;
: gtrr ( regs inputs -- outputs ) [ > 1 0 ? ] 2ary-op-r ;
: gtri ( regs inputs -- outputs ) [ > 1 0 ? ] 2ary-op-i ;
: gtir ( regs inputs -- outputs ) [ > 1 0 ? ] 2ary-op-ir ;
: eqrr ( regs inputs -- outputs ) [ = 1 0 ? ] 2ary-op-r ;
: eqri ( regs inputs -- outputs ) [ = 1 0 ? ] 2ary-op-i ;
: eqir ( regs inputs -- outputs ) [ = 1 0 ? ] 2ary-op-ir ;

: (execute-insn) ( regs inputs str -- regs' )
{
  { "addr" [ addr ] }
  { "addi" [ addi ] }
  { "mulr" [ mulr ] }
  { "muli" [ muli ] }
  { "divr" [ divr ] }
  { "divi" [ divi ] }
  { "banr" [ banr ] }
  { "bani" [ bani ] }
  { "borr" [ borr ] }
  { "bori" [ bori ] }
  { "setr" [ setr ] }
  { "seti" [ seti ] }
  { "gtrr" [ gtrr ] }
  { "gtri" [ gtri ] }
  { "gtir" [ gtir ] }
  { "eqrr" [ eqrr ] }
  { "eqri" [ eqri ] }
  { "eqir" [ eqir ] }
} case ;
: execute-insn ( regs str -- regs' )
" " split unclip [ [ string>number ] map f prefix ] dip (execute-insn) ;


:: aoc21p1 ( n -- )
<linked-hash> :> seen
0 :> count!
"/tmp/input" ascii file-lines unclip " " split second string>number :> ( program ip )
6 0 <array> :> regs!
n regs set-first
0 :> pc! [
[ 
  pc ip regs set-nth
  pc program nth :> insn
  pc 28 = [
    3 regs nth dup seen at
    [ drop seen >alist last first . "done" throw ]
    [ t swap seen set-at ] if
  ] when
  regs insn execute-insn regs!
  ip regs nth 1 + pc! t
  count 1 + count!
] loop
] [ . regs . ] recover count . ;