use.std::crypto::hashes::rpo

# ===== MEMORY FUNCTIONS ==========================================================================

#! Copies `n` words from `read_ptr` to `write_ptr`.
#!
#! Stack transition looks as follows:
#! [n, read_ptr, write_ptr, ...] -> [...]
#! cycles: 15 + 16n
export.memcopy
  # The loop variable is changed with an add instead of sub because the former
  # uses one fewer cycle. So here the counter is negated. (1 cycles)
  # stack: [-n, read_ptr, write_ptr, ...]
  neg

  # Pad the stack because mem_load overwrites it (4 cycles)
  # stack: [0, 0, 0, 0, -n, read_ptr, write_ptr, ...]
  padw

  # check loop condition (3 cycles)
  # stack: [b, 0, 0, 0, 0, -n, read_ptr, write_ptr, ...]
  dup.4 neq.0

  # LOOP: [0, 0, 0, 0, -n, read_ptr, write_ptr, ...]
  # while(n!=0) (16 cycles)
  while.true
    # perform read (2 cycles)
    # stack: [a3, a2, a1, a0, -n, read_ptr, write_ptr, ...]
    dup.5 mem_loadw

    # peform write (2 cycles)
    # stack: [a3, a2, a1, a0, -n, read_ptr, write_ptr, ...]
    dup.6 mem_storew

    # note: the values of `A` are no longer necessary, use 0 to signal its a
    # padding
    # stack: [-n, read_ptr, write_ptr, x, 0, 0, 0, 0, ...]
    swapw

    # stack: [-n+1, read_ptr+1, write_ptr+1, x, 0, 0, 0, 0, ...]
    # update counters (9 cycles)
    add.1 movup.3 movup.3 add.1 movup.3 add.1 movup.3

    # stack: [0, 0, 0, 0, -n+1, read_ptr+1, write_ptr+1, x, ...]
    swapw

    dup.4 neq.0 # while(n!=0) (3 cycles)
  end

  # clean stack (7 cycles)
  # stack: [...]
  dropw drop drop drop
end

#! Copies an even number of words from the advice_stack to memory.
#!
#! Input: [C, B, A, write_ptr, end_ptr, ...]
#! Output: [C, B, A, write_ptr, ...]
#!
#! Where:
#! - The words C, B, and A are the RPO hasher state
#!     - A is the capacity
#!     - C,B are the rate portion of the state
#! - The value `words = end_ptr - write_ptr` must be positive and even
#!
#! Cycles: 10 + 9 * word_pairs
export.pipe_double_words_to_memory.0
  dup.13 dup.13 neq # (4 cycles)

  # loop until write_ptr reaches end_ptr (9 cycles per iteration + 1)
  # LOOP: [b, C, B, A, write_ptr, end_ptr, ...]
  while.true
    adv_pipe hperm # (2 cycles)
    # => [C', B', A', write_ptr', end_ptr, ...]

    dup.13 dup.13 neq # (4 cycles)
    # LOOP: [b, C', B', A', write_ptr', end_ptr, ...]
  end

  movup.13 drop # (5 cycles)
  # [C', B', A', write_ptr', ...]
end

#! Copies an arbitrary number of words from the advice stack to memory
#!
#! Input: [num_words, write_ptr, ...]
#! Output: [C, B, A, write_ptr', ...]
#! Cycles:
#!  even num_words: 41 + 9 * num_words / 2
#!  odd num_words: 58 + 9 * round_down(num_words / 2)
export.pipe_words_to_memory.0
  # check if there is an odd number of words (6 cycles)
  dup is_odd
  # => [is_odd, num_words, write_ptr, ...]

  # copy is_odd, it defines if last last word requires padding (2 cycles)
  dup movdn.3
  # => [is_odd, num_words, write_ptr, needs_padding, ...]

  # compute `end_ptr` with an even number of words (5 cycles)
  sub dup.1 add swap
  # => [write_ptr, end_ptr, needs_padding, ...]

  # Prepare the capacity word. We use the padding rule which sets the first capacity
  # element to `len % 8` where `len` is the length of the hashed sequence. Since `len % 8`
  # is either equal to 0 or 4,  this is determined by the `needs_padding` flag multiplied
  # by 4. (6 cycles)
  dup.2 mul.4 push.0.0.0
  # => [A, write_ptr, end_ptr, needs_padding, ...]

  # set initial hasher state (8 cycles)
  padw padw
  # => [C, B, A, write_ptr, end_ptr, needs_padding, ...]

  # (9 + 6 * num_words cycles)
  exec.pipe_double_words_to_memory
  # => [C, B, A, write_ptr, needs_padding, ...]

  # (4 cycles)
  movup.13
  # => [needs_padding, C, B, A, write_ptr, ...]

  # if(needs_padding) (17 cycles)
  if.true
    # Rescue Prime Optimized uses overwrite mode, drop `C`. (4 cycles)
    dropw
    # => [B, A, write_ptr, ...]

    # Overwrite the `B` with the new data (1 cycles)
    adv_loadw
    # => [B', A, write_ptr, ...]

    # - get the memory address that B' should be saved to
    # - update the write_ptr to point to the next address (4 cycles)
    movup.8 dup.0 add.1 movdn.5
    # => [write_ptr, B', write_ptr+1, A, ...]

    # save data to memory (1 cycles)
    mem_storew
    # => [B', write_ptr+1, A, ...]

    # Fix write_ptr position (2 cycles)
    movup.4 movdn.8
    # => [B', A, write_ptr+1, ...]

    # Push padding word (4 cycles)
    padw
    # => [C, B', A, write_ptr+1, ...]

    # Run RPO permutation (1 cycles)
    hperm
    # => [C', B', A', write_ptr+1, ...]
  end
end

#! Moves an arbitrary number of words from the advice stack to memory and asserts it matches the commitment.
#!
#! Input: [num_words, write_ptr, COM, ...]
#! Output: [write_ptr', ...]
#! Cycles:
#!  even num_words: 62 + 9 * num_words / 2
#!  odd num_words: 79 + 9 * round_down(num_words / 2)
export.pipe_preimage_to_memory.0
  # Copies the advice stack data to memory
  exec.pipe_words_to_memory
  # => [C, B, A, write_ptr', COM, ...]

  # Leave only the digest on the stack
  exec.rpo::squeeze_digest
  # => [B, write_ptr', COM, ...]

  # Save the write_ptr (2 cycles)
  movup.4 movdn.8
  # => [HASH, COM, write_ptr', ...]

  # Check the COM (10 cycles)
  assert_eqw
  # => [write_ptr', ...]
end