CI: Stick with ubuntu-22.04 for CodeQL analysis

[zfs.git] / module / icp / algs / blake3 / blake3.c

/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "License").
 * You may not use this file except in compliance with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or https://opensource.org/licenses/CDDL-1.0.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */

/*
 * Based on BLAKE3 v1.3.1, https://github.com/BLAKE3-team/BLAKE3
 * Copyright (c) 2019-2020 Samuel Neves and Jack O'Connor
 * Copyright (c) 2021-2022 Tino Reichardt <milky-zfs@mcmilk.de>
 */

#include <sys/simd.h>
#include <sys/zfs_context.h>
#include <sys/blake3.h>

#include "blake3_impl.h"

/*
 * We need 1056 byte stack for blake3_compress_subtree_wide()
 * - we define this pragma to make gcc happy
 */
#if defined(__GNUC__)
#pragma GCC diagnostic ignored "-Wframe-larger-than="
#endif

/* internal used */
typedef struct {
        uint32_t input_cv[8];
        uint64_t counter;
        uint8_t block[BLAKE3_BLOCK_LEN];
        uint8_t block_len;
        uint8_t flags;
} output_t;

/* internal flags */
enum blake3_flags {
        CHUNK_START             = 1 << 0,
        CHUNK_END               = 1 << 1,
        PARENT                  = 1 << 2,
        ROOT                    = 1 << 3,
        KEYED_HASH              = 1 << 4,
        DERIVE_KEY_CONTEXT      = 1 << 5,
        DERIVE_KEY_MATERIAL     = 1 << 6,
};

/* internal start */
static void chunk_state_init(blake3_chunk_state_t *ctx,
    const uint32_t key[8], uint8_t flags)
{
        memcpy(ctx->cv, key, BLAKE3_KEY_LEN);
        ctx->chunk_counter = 0;
        memset(ctx->buf, 0, BLAKE3_BLOCK_LEN);
        ctx->buf_len = 0;
        ctx->blocks_compressed = 0;
        ctx->flags = flags;
}

static void chunk_state_reset(blake3_chunk_state_t *ctx,
    const uint32_t key[8], uint64_t chunk_counter)
{
        memcpy(ctx->cv, key, BLAKE3_KEY_LEN);
        ctx->chunk_counter = chunk_counter;
        ctx->blocks_compressed = 0;
        memset(ctx->buf, 0, BLAKE3_BLOCK_LEN);
        ctx->buf_len = 0;
}

static size_t chunk_state_len(const blake3_chunk_state_t *ctx)
{
        return (BLAKE3_BLOCK_LEN * (size_t)ctx->blocks_compressed) +
            ((size_t)ctx->buf_len);
}

static size_t chunk_state_fill_buf(blake3_chunk_state_t *ctx,
    const uint8_t *input, size_t input_len)
{
        size_t take = BLAKE3_BLOCK_LEN - ((size_t)ctx->buf_len);
        if (take > input_len) {
                take = input_len;
        }
        uint8_t *dest = ctx->buf + ((size_t)ctx->buf_len);
        memcpy(dest, input, take);
        ctx->buf_len += (uint8_t)take;
        return (take);
}

static uint8_t chunk_state_maybe_start_flag(const blake3_chunk_state_t *ctx)
{
        if (ctx->blocks_compressed == 0) {
                return (CHUNK_START);
        } else {
                return (0);
        }
}

static output_t make_output(const uint32_t input_cv[8],
    const uint8_t *block, uint8_t block_len,
    uint64_t counter, uint8_t flags)
{
        output_t ret;
        memcpy(ret.input_cv, input_cv, 32);
        memcpy(ret.block, block, BLAKE3_BLOCK_LEN);
        ret.block_len = block_len;
        ret.counter = counter;
        ret.flags = flags;
        return (ret);
}

/*
 * Chaining values within a given chunk (specifically the compress_in_place
 * interface) are represented as words. This avoids unnecessary bytes<->words
 * conversion overhead in the portable implementation. However, the hash_many
 * interface handles both user input and parent node blocks, so it accepts
 * bytes. For that reason, chaining values in the CV stack are represented as
 * bytes.
 */
static void output_chaining_value(const blake3_ops_t *ops,
    const output_t *ctx, uint8_t cv[32])
{
        uint32_t cv_words[8];
        memcpy(cv_words, ctx->input_cv, 32);
        ops->compress_in_place(cv_words, ctx->block, ctx->block_len,
            ctx->counter, ctx->flags);
        store_cv_words(cv, cv_words);
}

static void output_root_bytes(const blake3_ops_t *ops, const output_t *ctx,
    uint64_t seek, uint8_t *out, size_t out_len)
{
        uint64_t output_block_counter = seek / 64;
        size_t offset_within_block = seek % 64;
        uint8_t wide_buf[64];
        while (out_len > 0) {
                ops->compress_xof(ctx->input_cv, ctx->block, ctx->block_len,
                    output_block_counter, ctx->flags | ROOT, wide_buf);
                size_t available_bytes = 64 - offset_within_block;
                size_t memcpy_len;
                if (out_len > available_bytes) {
                        memcpy_len = available_bytes;
                } else {
                        memcpy_len = out_len;
                }
                memcpy(out, wide_buf + offset_within_block, memcpy_len);
                out += memcpy_len;
                out_len -= memcpy_len;
                output_block_counter += 1;
                offset_within_block = 0;
        }
}

static void chunk_state_update(const blake3_ops_t *ops,
    blake3_chunk_state_t *ctx, const uint8_t *input, size_t input_len)
{
        if (ctx->buf_len > 0) {
                size_t take = chunk_state_fill_buf(ctx, input, input_len);
                input += take;
                input_len -= take;
                if (input_len > 0) {
                        ops->compress_in_place(ctx->cv, ctx->buf,
                            BLAKE3_BLOCK_LEN, ctx->chunk_counter,
                            ctx->flags|chunk_state_maybe_start_flag(ctx));
                        ctx->blocks_compressed += 1;
                        ctx->buf_len = 0;
                        memset(ctx->buf, 0, BLAKE3_BLOCK_LEN);
                }
        }

        while (input_len > BLAKE3_BLOCK_LEN) {
                ops->compress_in_place(ctx->cv, input, BLAKE3_BLOCK_LEN,
                    ctx->chunk_counter,
                    ctx->flags|chunk_state_maybe_start_flag(ctx));
                ctx->blocks_compressed += 1;
                input += BLAKE3_BLOCK_LEN;
                input_len -= BLAKE3_BLOCK_LEN;
        }

        chunk_state_fill_buf(ctx, input, input_len);
}

static output_t chunk_state_output(const blake3_chunk_state_t *ctx)
{
        uint8_t block_flags =
            ctx->flags | chunk_state_maybe_start_flag(ctx) | CHUNK_END;
        return (make_output(ctx->cv, ctx->buf, ctx->buf_len, ctx->chunk_counter,
            block_flags));
}

static output_t parent_output(const uint8_t block[BLAKE3_BLOCK_LEN],
    const uint32_t key[8], uint8_t flags)
{
        return (make_output(key, block, BLAKE3_BLOCK_LEN, 0, flags | PARENT));
}

/*
 * Given some input larger than one chunk, return the number of bytes that
 * should go in the left subtree. This is the largest power-of-2 number of
 * chunks that leaves at least 1 byte for the right subtree.
 */
static size_t left_len(size_t content_len)
{
        /*
         * Subtract 1 to reserve at least one byte for the right side.
         * content_len
         * should always be greater than BLAKE3_CHUNK_LEN.
         */
        size_t full_chunks = (content_len - 1) / BLAKE3_CHUNK_LEN;
        return (round_down_to_power_of_2(full_chunks) * BLAKE3_CHUNK_LEN);
}

/*
 * Use SIMD parallelism to hash up to MAX_SIMD_DEGREE chunks at the same time
 * on a single thread. Write out the chunk chaining values and return the
 * number of chunks hashed. These chunks are never the root and never empty;
 * those cases use a different codepath.
 */
static size_t compress_chunks_parallel(const blake3_ops_t *ops,
    const uint8_t *input, size_t input_len, const uint32_t key[8],
    uint64_t chunk_counter, uint8_t flags, uint8_t *out)
{
        const uint8_t *chunks_array[MAX_SIMD_DEGREE];
        size_t input_position = 0;
        size_t chunks_array_len = 0;
        while (input_len - input_position >= BLAKE3_CHUNK_LEN) {
                chunks_array[chunks_array_len] = &input[input_position];
                input_position += BLAKE3_CHUNK_LEN;
                chunks_array_len += 1;
        }

        ops->hash_many(chunks_array, chunks_array_len, BLAKE3_CHUNK_LEN /
            BLAKE3_BLOCK_LEN, key, chunk_counter, B_TRUE, flags, CHUNK_START,
            CHUNK_END, out);

        /*
         * Hash the remaining partial chunk, if there is one. Note that the
         * empty chunk (meaning the empty message) is a different codepath.
         */
        if (input_len > input_position) {
                uint64_t counter = chunk_counter + (uint64_t)chunks_array_len;
                blake3_chunk_state_t chunk_state;
                chunk_state_init(&chunk_state, key, flags);
                chunk_state.chunk_counter = counter;
                chunk_state_update(ops, &chunk_state, &input[input_position],
                    input_len - input_position);
                output_t output = chunk_state_output(&chunk_state);
                output_chaining_value(ops, &output, &out[chunks_array_len *
                    BLAKE3_OUT_LEN]);
                return (chunks_array_len + 1);
        } else {
                return (chunks_array_len);
        }
}

/*
 * Use SIMD parallelism to hash up to MAX_SIMD_DEGREE parents at the same time
 * on a single thread. Write out the parent chaining values and return the
 * number of parents hashed. (If there's an odd input chaining value left over,
 * return it as an additional output.) These parents are never the root and
 * never empty; those cases use a different codepath.
 */
static size_t compress_parents_parallel(const blake3_ops_t *ops,
    const uint8_t *child_chaining_values, size_t num_chaining_values,
    const uint32_t key[8], uint8_t flags, uint8_t *out)
{
        const uint8_t *parents_array[MAX_SIMD_DEGREE_OR_2] = {0};
        size_t parents_array_len = 0;

        while (num_chaining_values - (2 * parents_array_len) >= 2) {
                parents_array[parents_array_len] = &child_chaining_values[2 *
                    parents_array_len * BLAKE3_OUT_LEN];
                parents_array_len += 1;
        }

        ops->hash_many(parents_array, parents_array_len, 1, key, 0, B_FALSE,
            flags | PARENT, 0, 0, out);

        /* If there's an odd child left over, it becomes an output. */
        if (num_chaining_values > 2 * parents_array_len) {
                memcpy(&out[parents_array_len * BLAKE3_OUT_LEN],
                    &child_chaining_values[2 * parents_array_len *
                    BLAKE3_OUT_LEN], BLAKE3_OUT_LEN);
                return (parents_array_len + 1);
        } else {
                return (parents_array_len);
        }
}

/*
 * The wide helper function returns (writes out) an array of chaining values
 * and returns the length of that array. The number of chaining values returned
 * is the dyanmically detected SIMD degree, at most MAX_SIMD_DEGREE. Or fewer,
 * if the input is shorter than that many chunks. The reason for maintaining a
 * wide array of chaining values going back up the tree, is to allow the
 * implementation to hash as many parents in parallel as possible.
 *
 * As a special case when the SIMD degree is 1, this function will still return
 * at least 2 outputs. This guarantees that this function doesn't perform the
 * root compression. (If it did, it would use the wrong flags, and also we
 * wouldn't be able to implement exendable ouput.) Note that this function is
 * not used when the whole input is only 1 chunk long; that's a different
 * codepath.
 *
 * Why not just have the caller split the input on the first update(), instead
 * of implementing this special rule? Because we don't want to limit SIMD or
 * multi-threading parallelism for that update().
 */
static size_t blake3_compress_subtree_wide(const blake3_ops_t *ops,
    const uint8_t *input, size_t input_len, const uint32_t key[8],
    uint64_t chunk_counter, uint8_t flags, uint8_t *out)
{
        /*
         * Note that the single chunk case does *not* bump the SIMD degree up
         * to 2 when it is 1. If this implementation adds multi-threading in
         * the future, this gives us the option of multi-threading even the
         * 2-chunk case, which can help performance on smaller platforms.
         */
        if (input_len <= (size_t)(ops->degree * BLAKE3_CHUNK_LEN)) {
                return (compress_chunks_parallel(ops, input, input_len, key,
                    chunk_counter, flags, out));
        }


        /*
         * With more than simd_degree chunks, we need to recurse. Start by
         * dividing the input into left and right subtrees. (Note that this is
         * only optimal as long as the SIMD degree is a power of 2. If we ever
         * get a SIMD degree of 3 or something, we'll need a more complicated
         * strategy.)
         */
        size_t left_input_len = left_len(input_len);
        size_t right_input_len = input_len - left_input_len;
        const uint8_t *right_input = &input[left_input_len];
        uint64_t right_chunk_counter = chunk_counter +
            (uint64_t)(left_input_len / BLAKE3_CHUNK_LEN);

        /*
         * Make space for the child outputs. Here we use MAX_SIMD_DEGREE_OR_2
         * to account for the special case of returning 2 outputs when the
         * SIMD degree is 1.
         */
        uint8_t cv_array[2 * MAX_SIMD_DEGREE_OR_2 * BLAKE3_OUT_LEN];
        size_t degree = ops->degree;
        if (left_input_len > BLAKE3_CHUNK_LEN && degree == 1) {

                /*
                 * The special case: We always use a degree of at least two,
                 * to make sure there are two outputs. Except, as noted above,
                 * at the chunk level, where we allow degree=1. (Note that the
                 * 1-chunk-input case is a different codepath.)
                 */
                degree = 2;
        }
        uint8_t *right_cvs = &cv_array[degree * BLAKE3_OUT_LEN];

        /*
         * Recurse! If this implementation adds multi-threading support in the
         * future, this is where it will go.
         */
        size_t left_n = blake3_compress_subtree_wide(ops, input, left_input_len,
            key, chunk_counter, flags, cv_array);
        size_t right_n = blake3_compress_subtree_wide(ops, right_input,
            right_input_len, key, right_chunk_counter, flags, right_cvs);

        /*
         * The special case again. If simd_degree=1, then we'll have left_n=1
         * and right_n=1. Rather than compressing them into a single output,
         * return them directly, to make sure we always have at least two
         * outputs.
         */
        if (left_n == 1) {
                memcpy(out, cv_array, 2 * BLAKE3_OUT_LEN);
                return (2);
        }

        /* Otherwise, do one layer of parent node compression. */
        size_t num_chaining_values = left_n + right_n;
        return compress_parents_parallel(ops, cv_array,
            num_chaining_values, key, flags, out);
}

/*
 * Hash a subtree with compress_subtree_wide(), and then condense the resulting
 * list of chaining values down to a single parent node. Don't compress that
 * last parent node, however. Instead, return its message bytes (the
 * concatenated chaining values of its children). This is necessary when the
 * first call to update() supplies a complete subtree, because the topmost
 * parent node of that subtree could end up being the root. It's also necessary
 * for extended output in the general case.
 *
 * As with compress_subtree_wide(), this function is not used on inputs of 1
 * chunk or less. That's a different codepath.
 */
static void compress_subtree_to_parent_node(const blake3_ops_t *ops,
    const uint8_t *input, size_t input_len, const uint32_t key[8],
    uint64_t chunk_counter, uint8_t flags, uint8_t out[2 * BLAKE3_OUT_LEN])
{
        uint8_t cv_array[MAX_SIMD_DEGREE_OR_2 * BLAKE3_OUT_LEN];
        size_t num_cvs = blake3_compress_subtree_wide(ops, input, input_len,
            key, chunk_counter, flags, cv_array);

        /*
         * If MAX_SIMD_DEGREE is greater than 2 and there's enough input,
         * compress_subtree_wide() returns more than 2 chaining values. Condense
         * them into 2 by forming parent nodes repeatedly.
         */
        uint8_t out_array[MAX_SIMD_DEGREE_OR_2 * BLAKE3_OUT_LEN / 2];
        while (num_cvs > 2) {
                num_cvs = compress_parents_parallel(ops, cv_array, num_cvs, key,
                    flags, out_array);
                memcpy(cv_array, out_array, num_cvs * BLAKE3_OUT_LEN);
        }
        memcpy(out, cv_array, 2 * BLAKE3_OUT_LEN);
}

static void hasher_init_base(BLAKE3_CTX *ctx, const uint32_t key[8],
    uint8_t flags)
{
        memcpy(ctx->key, key, BLAKE3_KEY_LEN);
        chunk_state_init(&ctx->chunk, key, flags);
        ctx->cv_stack_len = 0;
        ctx->ops = blake3_get_ops();
}

/*
 * As described in hasher_push_cv() below, we do "lazy merging", delaying
 * merges until right before the next CV is about to be added. This is
 * different from the reference implementation. Another difference is that we
 * aren't always merging 1 chunk at a time. Instead, each CV might represent
 * any power-of-two number of chunks, as long as the smaller-above-larger
 * stack order is maintained. Instead of the "count the trailing 0-bits"
 * algorithm described in the spec, we use a "count the total number of
 * 1-bits" variant that doesn't require us to retain the subtree size of the
 * CV on top of the stack. The principle is the same: each CV that should
 * remain in the stack is represented by a 1-bit in the total number of chunks
 * (or bytes) so far.
 */
static void hasher_merge_cv_stack(BLAKE3_CTX *ctx, uint64_t total_len)
{
        size_t post_merge_stack_len = (size_t)popcnt(total_len);
        while (ctx->cv_stack_len > post_merge_stack_len) {
                uint8_t *parent_node =
                    &ctx->cv_stack[(ctx->cv_stack_len - 2) * BLAKE3_OUT_LEN];
                output_t output =
                    parent_output(parent_node, ctx->key, ctx->chunk.flags);
                output_chaining_value(ctx->ops, &output, parent_node);
                ctx->cv_stack_len -= 1;
        }
}

/*
 * In reference_impl.rs, we merge the new CV with existing CVs from the stack
 * before pushing it. We can do that because we know more input is coming, so
 * we know none of the merges are root.
 *
 * This setting is different. We want to feed as much input as possible to
 * compress_subtree_wide(), without setting aside anything for the chunk_state.
 * If the user gives us 64 KiB, we want to parallelize over all 64 KiB at once
 * as a single subtree, if at all possible.
 *
 * This leads to two problems:
 * 1) This 64 KiB input might be the only call that ever gets made to update.
 *    In this case, the root node of the 64 KiB subtree would be the root node
 *    of the whole tree, and it would need to be ROOT finalized. We can't
 *    compress it until we know.
 * 2) This 64 KiB input might complete a larger tree, whose root node is
 *    similarly going to be the the root of the whole tree. For example, maybe
 *    we have 196 KiB (that is, 128 + 64) hashed so far. We can't compress the
 *    node at the root of the 256 KiB subtree until we know how to finalize it.
 *
 * The second problem is solved with "lazy merging". That is, when we're about
 * to add a CV to the stack, we don't merge it with anything first, as the
 * reference impl does. Instead we do merges using the *previous* CV that was
 * added, which is sitting on top of the stack, and we put the new CV
 * (unmerged) on top of the stack afterwards. This guarantees that we never
 * merge the root node until finalize().
 *
 * Solving the first problem requires an additional tool,
 * compress_subtree_to_parent_node(). That function always returns the top
 * *two* chaining values of the subtree it's compressing. We then do lazy
 * merging with each of them separately, so that the second CV will always
 * remain unmerged. (That also helps us support extendable output when we're
 * hashing an input all-at-once.)
 */
static void hasher_push_cv(BLAKE3_CTX *ctx, uint8_t new_cv[BLAKE3_OUT_LEN],
    uint64_t chunk_counter)
{
        hasher_merge_cv_stack(ctx, chunk_counter);
        memcpy(&ctx->cv_stack[ctx->cv_stack_len * BLAKE3_OUT_LEN], new_cv,
            BLAKE3_OUT_LEN);
        ctx->cv_stack_len += 1;
}

void
Blake3_Init(BLAKE3_CTX *ctx)
{
        hasher_init_base(ctx, BLAKE3_IV, 0);
}

void
Blake3_InitKeyed(BLAKE3_CTX *ctx, const uint8_t key[BLAKE3_KEY_LEN])
{
        uint32_t key_words[8];
        load_key_words(key, key_words);
        hasher_init_base(ctx, key_words, KEYED_HASH);
}

static void
Blake3_Update2(BLAKE3_CTX *ctx, const void *input, size_t input_len)
{
        /*
         * Explicitly checking for zero avoids causing UB by passing a null
         * pointer to memcpy. This comes up in practice with things like:
         *   std::vector<uint8_t> v;
         *   blake3_hasher_update(&hasher, v.data(), v.size());
         */
        if (input_len == 0) {
                return;
        }

        const uint8_t *input_bytes = (const uint8_t *)input;

        /*
         * If we have some partial chunk bytes in the internal chunk_state, we
         * need to finish that chunk first.
         */
        if (chunk_state_len(&ctx->chunk) > 0) {
                size_t take = BLAKE3_CHUNK_LEN - chunk_state_len(&ctx->chunk);
                if (take > input_len) {
                        take = input_len;
                }
                chunk_state_update(ctx->ops, &ctx->chunk, input_bytes, take);
                input_bytes += take;
                input_len -= take;
                /*
                 * If we've filled the current chunk and there's more coming,
                 * finalize this chunk and proceed. In this case we know it's
                 * not the root.
                 */
                if (input_len > 0) {
                        output_t output = chunk_state_output(&ctx->chunk);
                        uint8_t chunk_cv[32];
                        output_chaining_value(ctx->ops, &output, chunk_cv);
                        hasher_push_cv(ctx, chunk_cv, ctx->chunk.chunk_counter);
                        chunk_state_reset(&ctx->chunk, ctx->key,
                            ctx->chunk.chunk_counter + 1);
                } else {
                        return;
                }
        }

        /*
         * Now the chunk_state is clear, and we have more input. If there's
         * more than a single chunk (so, definitely not the root chunk), hash
         * the largest whole subtree we can, with the full benefits of SIMD
         * (and maybe in the future, multi-threading) parallelism. Two
         * restrictions:
         * - The subtree has to be a power-of-2 number of chunks. Only
         *   subtrees along the right edge can be incomplete, and we don't know
         *   where the right edge is going to be until we get to finalize().
         * - The subtree must evenly divide the total number of chunks up
         *   until this point (if total is not 0). If the current incomplete
         *   subtree is only waiting for 1 more chunk, we can't hash a subtree
         *   of 4 chunks. We have to complete the current subtree first.
         * Because we might need to break up the input to form powers of 2, or
         * to evenly divide what we already have, this part runs in a loop.
         */
        while (input_len > BLAKE3_CHUNK_LEN) {
                size_t subtree_len = round_down_to_power_of_2(input_len);
                uint64_t count_so_far =
                    ctx->chunk.chunk_counter * BLAKE3_CHUNK_LEN;
                /*
                 * Shrink the subtree_len until it evenly divides the count so
                 * far. We know that subtree_len itself is a power of 2, so we
                 * can use a bitmasking trick instead of an actual remainder
                 * operation. (Note that if the caller consistently passes
                 * power-of-2 inputs of the same size, as is hopefully
                 * typical, this loop condition will always fail, and
                 * subtree_len will always be the full length of the input.)
                 *
                 * An aside: We don't have to shrink subtree_len quite this
                 * much. For example, if count_so_far is 1, we could pass 2
                 * chunks to compress_subtree_to_parent_node. Since we'll get
                 * 2 CVs back, we'll still get the right answer in the end,
                 * and we might get to use 2-way SIMD parallelism. The problem
                 * with this optimization, is that it gets us stuck always
                 * hashing 2 chunks. The total number of chunks will remain
                 * odd, and we'll never graduate to higher degrees of
                 * parallelism. See
                 * https://github.com/BLAKE3-team/BLAKE3/issues/69.
                 */
                while ((((uint64_t)(subtree_len - 1)) & count_so_far) != 0) {
                        subtree_len /= 2;
                }
                /*
                 * The shrunken subtree_len might now be 1 chunk long. If so,
                 * hash that one chunk by itself. Otherwise, compress the
                 * subtree into a pair of CVs.
                 */
                uint64_t subtree_chunks = subtree_len / BLAKE3_CHUNK_LEN;
                if (subtree_len <= BLAKE3_CHUNK_LEN) {
                        blake3_chunk_state_t chunk_state;
                        chunk_state_init(&chunk_state, ctx->key,
                            ctx->chunk.flags);
                        chunk_state.chunk_counter = ctx->chunk.chunk_counter;
                        chunk_state_update(ctx->ops, &chunk_state, input_bytes,
                            subtree_len);
                        output_t output = chunk_state_output(&chunk_state);
                        uint8_t cv[BLAKE3_OUT_LEN];
                        output_chaining_value(ctx->ops, &output, cv);
                        hasher_push_cv(ctx, cv, chunk_state.chunk_counter);
                } else {
                        /*
                         * This is the high-performance happy path, though
                         * getting here depends on the caller giving us a long
                         * enough input.
                         */
                        uint8_t cv_pair[2 * BLAKE3_OUT_LEN];
                        compress_subtree_to_parent_node(ctx->ops, input_bytes,
                            subtree_len, ctx->key, ctx-> chunk.chunk_counter,
                            ctx->chunk.flags, cv_pair);
                        hasher_push_cv(ctx, cv_pair, ctx->chunk.chunk_counter);
                        hasher_push_cv(ctx, &cv_pair[BLAKE3_OUT_LEN],
                            ctx->chunk.chunk_counter + (subtree_chunks / 2));
                }
                ctx->chunk.chunk_counter += subtree_chunks;
                input_bytes += subtree_len;
                input_len -= subtree_len;
        }

        /*
         * If there's any remaining input less than a full chunk, add it to
         * the chunk state. In that case, also do a final merge loop to make
         * sure the subtree stack doesn't contain any unmerged pairs. The
         * remaining input means we know these merges are non-root. This merge
         * loop isn't strictly necessary here, because hasher_push_chunk_cv
         * already does its own merge loop, but it simplifies
         * blake3_hasher_finalize below.
         */
        if (input_len > 0) {
                chunk_state_update(ctx->ops, &ctx->chunk, input_bytes,
                    input_len);
                hasher_merge_cv_stack(ctx, ctx->chunk.chunk_counter);
        }
}

void
Blake3_Update(BLAKE3_CTX *ctx, const void *input, size_t todo)
{
        size_t done = 0;
        const uint8_t *data = input;
        const size_t block_max = 1024 * 64;

        /* max feed buffer to leave the stack size small */
        while (todo != 0) {
                size_t block = (todo >= block_max) ? block_max : todo;
                Blake3_Update2(ctx, data + done, block);
                done += block;
                todo -= block;
        }
}

void
Blake3_Final(const BLAKE3_CTX *ctx, uint8_t *out)
{
        Blake3_FinalSeek(ctx, 0, out, BLAKE3_OUT_LEN);
}

void
Blake3_FinalSeek(const BLAKE3_CTX *ctx, uint64_t seek, uint8_t *out,
    size_t out_len)
{
        /*
         * Explicitly checking for zero avoids causing UB by passing a null
         * pointer to memcpy. This comes up in practice with things like:
         *   std::vector<uint8_t> v;
         *   blake3_hasher_finalize(&hasher, v.data(), v.size());
         */
        if (out_len == 0) {
                return;
        }
        /* If the subtree stack is empty, then the current chunk is the root. */
        if (ctx->cv_stack_len == 0) {
                output_t output = chunk_state_output(&ctx->chunk);
                output_root_bytes(ctx->ops, &output, seek, out, out_len);
                return;
        }
        /*
         * If there are any bytes in the chunk state, finalize that chunk and
         * do a roll-up merge between that chunk hash and every subtree in the
         * stack. In this case, the extra merge loop at the end of
         * blake3_hasher_update guarantees that none of the subtrees in the
         * stack need to be merged with each other first. Otherwise, if there
         * are no bytes in the chunk state, then the top of the stack is a
         * chunk hash, and we start the merge from that.
         */
        output_t output;
        size_t cvs_remaining;
        if (chunk_state_len(&ctx->chunk) > 0) {
                cvs_remaining = ctx->cv_stack_len;
                output = chunk_state_output(&ctx->chunk);
        } else {
                /* There are always at least 2 CVs in the stack in this case. */
                cvs_remaining = ctx->cv_stack_len - 2;
                output = parent_output(&ctx->cv_stack[cvs_remaining * 32],
                    ctx->key, ctx->chunk.flags);
        }
        while (cvs_remaining > 0) {
                cvs_remaining -= 1;
                uint8_t parent_block[BLAKE3_BLOCK_LEN];
                memcpy(parent_block, &ctx->cv_stack[cvs_remaining * 32], 32);
                output_chaining_value(ctx->ops, &output, &parent_block[32]);
                output = parent_output(parent_block, ctx->key,
                    ctx->chunk.flags);
        }
        output_root_bytes(ctx->ops, &output, seek, out, out_len);
}