needs lib/struct hal/vreg hal/instr lib/func
unit num/poly1305

16 const BLKSZ
1 24 lshift const SETHIGH

struct Poly1305 {
  [uchar,BLKSZ] pbuffer ;
  \ _tmp is before polyr in order to simply muladd operation
  [uint,5] _tmp polyr polyh ;
  [uint,10] polyd ;
  [uint,8] polypad ;
  \ highbit is either 1<<24 or 0
  uint leftover highbit ;
}

\ clear from polyr to highbit included
offsetof highbit 4+ offsetof polyr - 4/ const _CLEARSZ

\ Poly1305 struct is stored in R1
:~ ( run: n -- op ) R1) swap 4* [compile] offsetof + +) ;
: _h) word"polyh" ~ ; : _r) word"polyr" ~ ; : _tmp) word"_tmp" ~ ;
: _dl) 2* word"polyd" ~ ; : _dh) 2* 1+ word"polyd" ~ ;
: _high) 0 word"highbit" ~ ; : _pad) word"polypad" ~ ;

\ ----- Used for blocks -----
\ Aligned version of hn+
: _mut 1 and if REGA else REGS then ; : _cst 1 xor _mut ;
: _ahn+, ( n -- )
  dup 4 = if drop
    8 i) S>) >>, _high) S>) |, 4 _h) S>) dir) +,
    exit then

  \ A) below are placeholders, see [doc/hal]
  W) over 4* +) over _cst R>) ale@,
  A) over _cst R) &) R0>) @,
  dup 6 * ?dup if i) R0>) <<, then
  $3ffffff i) R0>) &,
  dup 6 * ?dup if 32 swap-
      i) over _mut R>) >>,
      A) over _mut R) &) R0>) |, then
  _h) R0>) dir) +, ;

: _hn+, ( n -- ) ( run: a -- a )
  W) over 3 * +) S>) le@,
  dup 2* ?dup if i) S>) >>, then
  dup 4 <> if $3ffffff i) S>) &, else
    _high) S>) |, then
  _h) S>) dir) +, ;

\ result in dn = R0 A, S & W are temporaries
: _dn$, ( n -- )
  0 i) R0>) @, 0 i) A>) @,
  5 0 do i _h) @, dup i - _r) d*, W) &) R0>) +c$, S) &) A>) +c, loop
  dup _dl) R0>) !, _dh) A>) !, ;

\ dn = R0 A, S & W are temporaries
\ "carry" in S
: _hn%p, ( n -- )
  dup _dl) R0>) @, dup _dh) A>) @,
  dup if S) &) R0>) +c$, 0 i) A>) +c, then
  R0) &) S>) @, 26 i) S>) >>, 6 i) A>) <<, A) &) S>) |,
  $3ffffff i) R0>) &, _h) R0>) !, ;

\ ----- Used for both -----
\ end carry in h0 and h1
:~ A) &) S>) @, 26 i) S>) >>, $3ffffff i) A>) &, ;
: _hcend, ( -- )
  S) &) A>) @, 2 i) A>) <<, 0 _h) S>) +, S) &) A>) +,
  ~ 0 _h) A>) !, 1 _h) S>) dir) +, ;

\ ----- Used for finish -----
\ "carry" in S
: _hnc, ( n -- ) 1+ dup _h) A>) r! @, 1 <> if S) &) A>) +, then ~ r> !, ;
: _gn$, dup _h) A>) @, dup if S) &) else 5 i) then A>) +, ~ _tmp) A>) !, ;
: _g4$, ( -- ) 4 _h) A>) @, S) &) A>) +, 1 26 lshift i) A>) -, 4 _tmp) A>) !, ;
\ Select between h and g depending on A=mask
: _hn|gn, ( n -- )
  dup _tmp) S>) @, dup _h) R0>) @,
  S) &) R0>) ^, A) &) R0>) &,
  S) &) R0>) ^, _h) R0>) !, ;
: _hn%, ( n -- )
  dup _h) A>) r! @,
  dup 6 * r! if r@ i) A>) >>, then
  1+ _h) S>) @, 26 r> - i) S>) <<,
  S) &) A>) |, r> !, ;
: _hn+pad, r! _h) S>) @, r@ _pad) S>) r@ if +c, else +c$, then
  r> _tmp) dir) S>) ale@, ( _tmp is aligned thanks to here# ) ;

\ ----- High level implementations -----
:~ ( a u poly1305 -- )
  W) &) R1>) @,
  [compile] begin PSP) S>) @, BLKSZ i) S>) >=) if,
    PSP) 4 +) @, ' 5 timesi
    ['] _dn$, 5 timesi
    ['] _hn%p, 5 timesi _hcend,

    BLKSZ PSP) -n,
    BLKSZ PSP) 4 +) +n,
  swap bbr, [compile] then

  8 ps+, drop, exit, ;
code blocks ~ _hn+, code ablocks ~ _ahn+,

: update ( a u poly1305 -- )
  r! leftover if
    BLKSZ r@ leftover - over min rot ( u n a )
    2dup ( u n a n a ) r@ bi pbuffer | leftover + rot cmove
    over + rot> ( a u n ) dup r@ doto leftover + | -
    r@ leftover BLKSZ < if 2drop rdrop exit then
    r@ dup pbuffer BLKSZ swap ablocks \ pbuffer is always aligned
    0 r@ to leftover
  then ( a u )

  dup BLKSZ >= if
    dup [ BLKSZ 1- inv litn ] and ( a u n )
    oover over r@
    oover 3 and if blocks else ablocks then ( a u n )
    [ PSP) 4 +) dir) +, PSP) dir) -, ] drop
  then ( a u )

  ?dup if
    tuck r@ bi pbuffer | leftover + swap cmove ( u )
    r@ doto leftover + |
    else drop then rdrop ;

: _clear ( poly1305 -- ) dup BLKSZ zfill offsetof polyr + _CLEARSZ 0 fill ;
: finish ( poly1305 -- a )
  r! leftover if
    r@ dupbi pbuffer | leftover tuck + swap BLKSZ swap- zfill
    dupbi pbuffer | leftover + 1 swap c!
    0 over to highbit dup pbuffer BLKSZ rot ablocks
  then r> [
    W) &) R1>) @,
    ' _hnc, 4 timesi _hcend,
    ' _gn$, 4 timesi _g4$,
    31 i) A>) signed) >>, \ g4 -> mask
    ' _hn|gn, 5 timesi
    ' _hn%, 4 timesi
    ' _hn+pad, 4 timesi
  ] dup _clear _tmp ;

code verify ( a a -- f ) \ constant time equality
  PSP) A>) @, 16 PSP) !n, 0 i) S>) @,
  begin
    W) 8b) R0>) @+,
    A) 8b) R1>) @+,
    R0) &) R1>) ^,
    R1) &) S>) |,
    1 PSP) -n, ?brnz,
  S) &) @, 1 i) -, 31 i) >>,
  4 ps+, exit,

\ r &= 0xffffffc0ffffffc0ffffffc0fffffff
create KEY_MASK map< , $3ffffff $3ffff03 $3ffc0ff $3f03fff $00fffff
: sle@+ ( s for sesqui ) bi 3 + | le@ ; : le@+ bi 4+ | le@ ; :~ 4* + ;
: Poly1305$ r! Poly1305 typesz 4/ 0 fill
  5 0 do sle@+ i ?dup if 2* rshift then KEY_MASK i ~ @ and V1 polyr i ~ ! loop
  1+ 4 0 do le@+ V1 polypad i ~ ! loop drop
  r@ 4+ dupbi polyr | _tmp 4 move _tmp 4 [: 5 * :] map[]
  SETHIGH r> to highbit ;
: newPoly1305 here# tuck Poly1305 typesz allot Poly1305$ ;