use.std::mem
use.std::crypto::hashes::rpo
use.std::math::u64

#! Loads the leaf at the absolute `pos` in the MMR.
#!
#! This MMR implementation supports only u32 positions.
#!
#! Stack transition:
#! Input: [pos, mmr_ptr, ...]
#! Output: [N, ...] where `N` is the leaf and `R` is the MMR peak that owns the leaf.
#!
#! Cycles: 115
export.get
  # load the num_leaves of the MMR (2 cycles)
  dup.1 mem_load
  # stack: [num_leaves, pos, mmr_ptr, ...]

  # compute `num_leaves & pos`, this contains all peaks before `pos` (and maybe some after the owning peak) (3 cycles)
  dup.1 dup.1 u32and
  # stack: [before_candidates, num_leaves, pos, mmr_ptr, ...]

  # compute `num_leaves - before_candidates`, this removes every peak before the owner (result may include peaks after owner) (4 cycles)
  dup.1 swap sub
  # stack: [owner_candidates, num_leaves, pos, mmr_ptr, ...]

  # compute `ilog2(owner_candidates)` and `2**depth`, it corresponds to the owner peak and its depth (61 cycles)
  ilog2 dup.0 pow2
  # stack: [owner_peak, depth, num_leaves, pos, mmr_ptr, ...]

  # compute `owner_peak - 1`, this mask corresponds to every peak after the owner (3 cycles)
  dup.0 sub.1
  # stack: [after_mask, owner_peak, depth, num_leaves, pos, mmr_ptr, ...]

  # compute `num_leaves & after_mask`, uses the mask to compute the actual after peaks (2 cycles)
  dup.3 u32and
  # stack: [after_peaks, owner_peak, depth, num_leaves, pos, mmr_ptr, ...]

  # compute `num_leaves - (after_peaks + owner_peak)`, this computes the before peaks (5 cycles)
  add movup.2 swap sub
  # stack: [peaks_before, depth, pos, mmr_ptr, ...]

  # compute `pos - peaks_before`, this computes the relative_pos of the leaf w.r.t. the owner peak. (4 cycles)
  movup.2 dup.1 sub
  # stack: [relative_pos, peaks_before, depth, mmr_ptr, ...]

  # compute `popcount(peaks_before)`, the count peaks before the target to be skipped when loading from mem (2 cycles)
  swap u32assert u32popcnt
  # stack: [peak_count, relative_pos, depth, mmr_ptr, ...]

  # compute `mmr_ptr + peak_count + 1` the target tree index (3 cycles)
  movup.3 add add.1
  # stack: [mmr_ptr, relative_pos, depth, ...]

  # load the target peak (6 cycles)
  padw movup.4 mem_loadw
  # stack: [P, relative_pos, depth, ...]

  # find the tree depth (2 cycles)
  movup.4 movup.5
  # stack: [depth, relative_pos, P, ...]

  # corner case, leaf values are not supported in the VM's Merkle store, so the
  # `mtree_get` instruction will fail for the single leaf case of the MMR. (2 cycles)
  dup.0 eq.0
  if.true
    drop drop # (2 cycles)
    # stack: [leaf, ...]
  else
    # verify and get the leaf (9 cycles)
    mtree_get
    # stack: [leaf, root, ...]

    # drop the root (5 cycles)
    swapw dropw
    # stack: [leaf, ...]
  end
end

#! Given the num_leaves of a MMR returns the num_peaks.
#!
#! Input: [num_leaves, ...]
#! Output: [num_peaks, ...]
#! Cycles: 69
export.num_leaves_to_num_peaks
  # count number of peaks (69 cycles)
  u32split u32popcnt swap u32popcnt add
  # => [count, ...]
end

#! Given the num_peaks of a MMR, returns the hasher state size after accounting
#! for the required padding.
#!
#! Input: [num_peaks, ...]
#! Output: [len, ...]
#! Cycles: 17
export.num_peaks_to_message_size
  # the peaks are padded to a minimum length of 16 (10 cycles)
  push.16 u32max
  # => [count_min, ...]

  # when the number of peaks is greater than 16, then they are padded to an even number (7 cycles)
  dup is_odd add
  # => [even_count_min, ...]
end

#! Load the MMR peak data based on its hash.
#!
#! Input: [HASH, mmr_ptr, ...]
#! Output: [...]
#!
#! Where:
#!  - HASH: is the MMR peak hash, the hash is expected to be padded to an even
#!    length and to have a minimum size of 16 elements
#!  - The advice map must contain a key with HASH, and its value is
#!    `num_leaves || hash_data`, and hash_data is the data used to computed `HASH`
#!  - mmt_ptr: the memory location where the MMR data will be written to,
#!    starting with the MMR forest (its total leaves count) followed by its peaks
#!
#! Cycles: 162 + 9 * extra_peak_pair cycles
#!    where `extra_peak` is the number of peak pairs in addition to the first
#!    16, i.e. `round_up((num_of_peaks - 16) / 2)`
export.unpack
  # load the num_leaves and peaks to the advice_stack (0 cycles)
  adv.push_mapval
  # operand_stack => [HASH, mmr_ptr, ...]
  # advice_stack => [NUM_LEAVES, peaks*, ...]

  # load the size from the advice stack (7 cycles)
  adv_push.4 drop drop drop
  # operand_stack => [num_leaves, HASH, mmr_ptr, ...]
  # advice_stack => [peaks*, ...]

  # save the forest to memory (4 cycles)
  dup dup.6 mem_store
  # => [num_leaves, HASH, mmr_ptr, ...]

  # find the hasher state size, this is how many words will be read from the advice_stack
  exec.num_leaves_to_num_peaks
  exec.num_peaks_to_message_size
  # => [state_size, HASH, mmr_ptr, ...]

  # compute the end address including the padding data and forest (3 cycles)
  dup.5 add add.1
  # => [mmt_ptr_end, HASH, mmr_ptr, ...]

  # update the mmr_ptr to account for the size (2 cycles)
  movup.5 add.1
  # => [mmr_ptr+1, mmt_ptr_end, HASH, ...]

  # hash the first 16 words (28 cycles)
  padw padw padw
  adv_pipe hperm
  adv_pipe hperm
  adv_pipe hperm
  adv_pipe hperm
  adv_pipe hperm
  adv_pipe hperm
  adv_pipe hperm
  adv_pipe hperm
  # => [C, B, A, mmr_ptr+17, mmt_ptr_end, HASH, ...]

  # handle MMR with more than 16 elements (10 + 9 * words cycles)
  exec.mem::pipe_double_words_to_memory
  # => [C, B, A, mmr_ptr+17, HASH, ...]

  # drop anything but the hash result, word B (11 cycles)
  exec.rpo::squeeze_digest movup.4 drop
  # => [B, HASH, ...]

  # assert on the resulting hash (11 cycles)
  assert_eqw
end

#! Computes the hash of the given MMR and copies it to the Advice Map using its hash as a key.
#!
#! Input: [mmr_ptr, ...]
#! Output: [HASH, ...]
#! Cycles: 128 + 3 * num_peaks
export.pack
  # load num_leaves (2 cycles)
  dup mem_load
  # => [num_leaves, mmr_ptr, ...]

  # compute num_peaks (69 cycles)
  exec.num_leaves_to_num_peaks
  # => [num_peaks, mmr_ptr, ...]

  # compute the message size (18 cycles)
  exec.num_peaks_to_message_size
  # => [message_size, mmr_ptr, ...]

  # compute peaks_start and peaks_end (6 cycles)
  dup.1 add.1 swap dup.1 add swap
  # => [peaks_start, peaks_end, mmr_ptr, ...]

  # hash the memory contents (25 + 3 * num_peaks)
  padw padw padw
  exec.rpo::absorb_double_words_from_memory
  exec.rpo::squeeze_digest
  # => [HASH, peaks_end, peaks_end, mmr_ptr, ...]

  # prepare stack for adv.insert_mem (4 cycles)
  movup.4 drop movup.4 movdn.5
  # => [HASH, mmr_ptr, peaks_end, ...]

  # copy the data to advice map
  adv.insert_mem

  # drop the extra addresses (4 cycles)
  movup.4 drop movup.4 drop
  # => [HASH, ...]
end

#! Adds a new element to the MMR.
#!
#! This will update the MMR peaks in the VM's memory and the advice provider
#! with any merged nodes.
#!
#! Input: [EL, mmr_ptr, ...]
#! Output: [...]
#! Cycles: 144 + 39 * peak_merges
export.add
  # get num_leaves (2 cycles)
  dup.4 mem_load
  # => [num_leaves, EL, mmr_ptr]

  # update the num_leaves (5 cycles)
  dup add.1 dup.6 mem_store
  # => [num_leaves, EL, mmr_ptr]

  dup exec.num_leaves_to_num_peaks
  # [num_peaks, num_leaves, EL, mmr_ptr] (70 cycles)

  # compute peaks_end (3 cycles)
  movup.6 add add.1
  # [mmr_end, num_leaves, EL]

  # find how many MMR peaks will be merged (41 cycles)
  swap u32split exec.u64::cto
  # => [num_merges, mmr_end, EL]

  # optimization: negate num_merges to use add.1 instead of sub.1 (1 cycles)
  neg
  # => [-num_merges, mmr_end, EL]

  # move the control data after the working data (2 cycles)
  movdn.5 movdn.5
  # => [EL, -num_merges, mmr_end]

  # add a word of padding to load the peak from memory (4 cycles)
  padw
  # => [PAD, EL, -num_merges, mmr_end]

  # loop while there are merges left to be done (5 cycles)
  dup.8 neq.0
  # LOOP: [b, PAD, EL, -num_merges, mmr_end]

  while.true # (39 cycles)
    # load peak (4 cycles)
    dup.9 sub.1 mem_loadw
    # => [PEAK, EL, -num_merges, mmr_end]

    # merge the nodes (17 cycles)
    swapw mtree_merge
    # => [EL', -num_merges, mmr_end]

    # erase existing peak (6 cycles)
    padw dup.9 mem_storew
    # => [PAD, EL', -num_merges, mmr_end]

    # update control (7 cycles)
    swapw.2 add.1 swap sub.1 swap swapw.2
    # => [PAD, EL', -num_merges+1, mmr_end-1]

    # check loop condition (5 cycles)
    dup.8 neq.0
    # LOOP: [b, PAD, EL', -num_merges+1, mmr_end-1]
  end

  # drop padding (4 cycles)
  dropw
  # =>: [EL, -num_merges+1, mmr_end-1]

  # save the new peak (2 cycles)
  movup.5 mem_storew
  # =>: [EL, -num_merges+1]

  # clean stack (5 cycles)
  dropw drop
end