fn invariants_hold_on_many_insertions()

in backends/v3/src/radix.rs [928:1012]

• 68 lines of code
• 7 McCabe index (conditional complexity)

    fn invariants_hold_on_many_insertions(remove_all: bool) {
        // Small vocabulary sizes lead to violations more quickly due to
        // prefix sharing, etc.
        const VOCAB_SIZE: u32 = 2;
        const DATA_LEN: usize = 1_000;

        const MAX_PREFILL_LEN: usize = 8;
        const MAX_DECODE_LEN: usize = 8;

        let vocab_range = Uniform::new(0, VOCAB_SIZE);
        let data_range = Uniform::new(0, DATA_LEN);
        let prefill_len_range = Uniform::new(0, MAX_PREFILL_LEN);
        let decode_len_range = Uniform::new(0, MAX_DECODE_LEN);

        let mut rng = SmallRng::seed_from_u64(64);
        let data = (0..DATA_LEN)
            .map(|_| vocab_range.sample(&mut rng))
            .collect::<Vec<_>>();
        let mut allocator = RadixAllocator::new(1, 100, None);

        let mut allocations = Vec::new();

        for i in 0..100_000 {
            // Allocate until all blocks are used.
            'allocation: loop {
                // Use offset 0 half of the times for prefix sharing.
                let prefill_offset = data_range.sample(&mut rng);
                let prefill_len = prefill_len_range.sample(&mut rng);
                let decode_len = decode_len_range.sample(&mut rng);

                let prefill =
                    data[prefill_offset..data.len().min(prefill_offset + prefill_len)].to_vec();

                let allocation = match allocator
                    .allocate((prefill.len() + decode_len) as u32, Some(Arc::new(prefill)))
                {
                    Some(allocation) => allocation,
                    None => break 'allocation,
                };
                let non_prefix_blocks = allocation.blocks[allocation.prefix_len as usize..]
                    .iter()
                    .copied()
                    .collect::<FxHashSet<_>>();
                let blockset = allocation.blocks.iter().copied().collect::<FxHashSet<_>>();

                // No duplicate blocks in an allocation.
                assert_eq!(
                    allocation.blocks.len(),
                    blockset.len(),
                    "Duplicate blocks in allocation"
                );

                allocations.push(AllocationWithInfo {
                    allocation,
                    blockset,
                    non_prefix_blocks,
                });
            }

            // Check invariants. Skip first iteration, since there is no prefix sharing yet.
            if i > 1 {
                check_allocation_invariants(&allocations);
            }

            // Remove 20% of the allocations, randomly.
            if remove_all {
                allocations.into_iter().for_each(|allocation| {
                    allocator.free(
                        allocation.allocation.blocks.clone(),
                        allocation.allocation.allocation_id,
                    )
                });
                allocations = Vec::new();
            } else {
                allocations.shuffle(&mut rng);
                let remove_index = (allocations.len() as f64 * 0.8) as usize;
                for allocation in allocations.drain(remove_index..) {
                    allocator.free(
                        allocation.allocation.blocks.clone(),
                        allocation.allocation.allocation_id,
                    );
                }
            }
        }
    }