needs hal/instr hal/vreg lib/bit lib/psrs lib/struct
unit num/chacha

create CHACHA_CST ,"expand 32-byte k"
64 const BLKSZ
\ 64-bit counter variants are even
enum Legacy IETF XChaCha

struct ChaCha {
  +$00 [uint,16] wstate ; \ working state
  +$40 [uint,16] bstate ; \ base state
  +$70 uint cntlo cnthi ; \ counter
  +$74 [uint,3] nonce ;
  +$80 uchar rounds variant remaining void? ;
}

\ a += b ; c^=a ; c <<= n
: _regRound, ( n c b a )
  swap &) over (R R>) +,
  &) over (R R>) ^,
  swap i) swap (R R>) lrot, ;

\ a=S b=W c=A d=R0
\ b is W to remove some xchg W) Reg>)
: _rqround,
  16 R0) W) S) _regRound,
  12 W) R0) A) _regRound,
  08 R0) W) S) _regRound,
  07 W) R0) A) _regRound, ;

\ a=S b=W c=A d=R0
:~ W) swap 4* +) ;
: _load4, ( a b c d -- )
  R2) W>) !,
  ~ R0>) @, ~ A>) @,
  swap
  ~ S>) @, ~ W>) @, ;

:~ R1) swap 4* +) ;
: _store4, ( a b c d -- )
  R2) R1>) @,
  ~ R0>) !, ~ A>) !, ~ W>) !, ~ S>) !,
  R1) &) @, ;

: qround@, 4 ndup _load4, _rqround, _store4, ;

\ ---MAIN API---

: round ( chacha -- )
  bi wstate | rounds 2/ 0 do [
    0 4 08 12 qround@,
    1 5 09 13 qround@,
    2 6 10 14 qround@,
    3 7 11 15 qround@,
    0 5 10 15 qround@,
    1 6 11 12 qround@,
    2 7 08 13 qround@,
    3 4 09 14 qround@,
    ] loop drop ;

code _statesum ( src dst -- )
  PSP) A>) @+,
  BLKSZ 4/ i) R0>) @,
  begin
    A) S>) @+,
    W) S>) +,
    W) S>) dir) ale@,
    4 i) +,
    1 i) R0>) -, ?brnz,
  drop, exit,

: blockfn ( chacha -- ) dup round bi bstate | wstate _statesum ;

:~ W) offsetof cntlo +) ;
: counter+, ( chacha -- chacha run: hi lo -- )
  1 W) offsetof variant +) 8b) &nf,
  ifnz, over ~ +n,
  [compile] else rot>
    ~ S>) @, ~ 4 +) A>) @,
    i) S>) +c$,
    i) A>) +c,
    ~ S>) !, ~ 4 +) A>) !,
  [compile] then ;
code counter1+ 0 1 counter+, exit,

code loadstate ( chacha -- chacha )
  A) &) !, R0) &) !, BLKSZ 4/ i) S>) @,
  offsetof bstate i) W>) +,
  offsetof wstate i) A>) +,
  move[], R0) &) @, exit,

code _lemove ( src dst u -- src u )
  PSP) S>) @, PSP) !, PSP) 4 +) A>) @,
  begin
    A) R0>) le@,
    S) R0>) !+,
    4 i) A>) +,
    1 i) -, ?brnz,
  drop, exit,

: _?void ( a n chacha -- ) void? if 4* zfill else 2drop then ;
: ChaCha$ ( an ak chacha -- )
  tuck bstate 16 + 8 _lemove oover _?void
  dup variant case ( an chacha )
    IETF = of tuck nonce 3 _lemove rot _?void endof
    Legacy = of
      0 over to cnthi
      tuck nonce 4+ 2 _lemove rot _?void endof
    XChaCha = of
      2dup bstate 48 + 4 _lemove oover _?void
      dup loadstate round \ hchacha round
      dupbi wstate 00 + | bstate 16 + 4 move
      dupbi wstate 48 + | bstate 32 + 4 move
      dip 16 + | tuck nonce 4+ 2 _lemove oover _?void
      0 over to cnthi
      1 swap to cntlo
      endof
    drop abort"Unsupported variant"
  endcase ;

code _xor[] ( a u wstate -- a' )
  PSP) R0>) @,
  PSP) 4 +) A>) @,
  begin
    W) 8b) S>) @+,
    A) 8b) dir) S>) ^,
    1 i) A>) +,
    1 i) R0>) -, ?brnz,
  8 ps+, A) &) W>) @, exit,

: apply ( a u chacha -- a )
  r! remaining ?dup if ( a u n )
    over min tuck - ( a min u' ) rot> tuck
    r@ bi wstate | remaining BLKSZ swap- + _xor[] ( u' min a' )
    rot> r@ tuck doto remaining swap- dup |
    not if counter1+ then drop
  then ( a u )

  \ Main loop
  tuck 6 rshift 0 do ( u a )
    [ BLKSZ PSP) -n, ]
    V1 dup loadstate blockfn counter1+
    wstate BLKSZ swap _xor[]
  loop swap ( a u )

  ?dup if
    dup BLKSZ swap- r@ to remaining
    r@ dup loadstate blockfn
    wstate _xor[]
  then rdrop ;

: discard ( chacha -- )
  r! wstate 16 0 fill
  r@ bstate 16 + 12 0 fill
  0 1 1 r> >A to{ remaining void? cntlo } ;

: newChaCha20 ( -- chacha )
  here# ChaCha typesz allot >A
  20 IETF 0 1 1 to{ rounds variant remaining void? cntlo }
  A> CHACHA_CST over bstate 4 _lemove 2drop ;