datasketches 0.5.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Hash set for storing unique coupons with linear probing
//!
//! Uses open addressing with a custom stride function to handle collisions.
//! Provides better performance than List when many coupons are stored.

use crate::codec::SketchBytes;
use crate::codec::SketchSlice;
use crate::codec::assert::insufficient_data;
use crate::codec::family::Family;
use crate::error::Error;
use crate::hll::Coupon;
use crate::hll::HllType;
use crate::hll::KEY_MASK_26;
use crate::hll::container::Container;
use crate::hll::serialization::COMPACT_FLAG_MASK;
use crate::hll::serialization::CUR_MODE_SET;
use crate::hll::serialization::HASH_SET_PREINTS;
use crate::hll::serialization::SERIAL_VERSION;
use crate::hll::serialization::SET_PREAMBLE_SIZE;
use crate::hll::serialization::encode_mode_byte;

/// Hash set for efficient coupon storage with collision handling
#[derive(Debug, Clone, PartialEq)]
pub struct HashSet {
    container: Container,
}

impl Default for HashSet {
    fn default() -> Self {
        const LG_INIT_SET_SIZE: usize = 5;
        Self::new(LG_INIT_SET_SIZE)
    }
}

impl HashSet {
    pub fn new(lg_size: usize) -> Self {
        Self {
            container: Container::new(lg_size),
        }
    }

    /// Insert coupon into hash set, ignoring duplicates
    pub fn update(&mut self, coupon: Coupon) {
        let mask = (1 << self.container.lg_size()) - 1;

        // Initial probe position from low bits of coupon
        let mut probe = coupon.raw() & mask;
        let starting_position = probe;

        loop {
            let slot = &mut self.container.coupons[probe as usize];
            if slot.is_empty() {
                // Found empty slot, insert new coupon
                *slot = coupon;
                self.container.len += 1;
                break;
            } else if *slot == coupon {
                // Duplicate found, nothing to do
                break;
            }

            // Collision: compute stride and probe next position
            // Stride is always odd to ensure all slots are visited
            let stride = ((coupon.raw() & KEY_MASK_26) >> self.container.lg_size()) | 1;
            probe = (probe + stride) & mask;
            if probe == starting_position {
                // Invariant: the caller (HllSketch) is responsible for
                // growing / upgrading the HashSet when it's full
                unreachable!("HashSet full; no empty slots");
            }
        }
    }

    pub fn container(&self) -> &Container {
        &self.container
    }

    /// Deserialize a HashSet from bytes
    pub fn deserialize(
        mut cursor: SketchSlice,
        lg_arr: usize,
        compact: bool,
    ) -> Result<Self, Error> {
        // Read coupon count from bytes 8-11
        let coupon_count = cursor
            .read_u32_le()
            .map_err(insufficient_data("coupon_count"))?;
        let coupon_count = coupon_count as usize;
        let array_size = 1usize << lg_arr;
        if coupon_count >= array_size {
            return Err(Error::deserial(format!(
                "SET mode coupon count {coupon_count} must be below capacity {array_size}"
            )));
        }
        let read_count = if compact { coupon_count } else { array_size };
        let required_bytes = read_count * size_of::<u32>();
        if required_bytes > cursor.remaining().len() {
            return Err(Error::insufficient_data(format!(
                "SET mode coupons require {required_bytes} bytes, got {}",
                cursor.remaining().len()
            )));
        }

        if compact {
            // Compact mode: only couponCount coupons are stored
            // Create a new hash set and insert coupons one by one
            let mut hash_set = HashSet::new(lg_arr);
            for i in 0..coupon_count {
                let coupon = cursor.read_u32_le().map_err(|_| {
                    Error::insufficient_data(format!(
                        "expected {coupon_count} coupons, failed at index {i}"
                    ))
                })?;
                hash_set.update(Coupon(coupon));
            }
            if hash_set.container.len() != coupon_count {
                return Err(Error::deserial("SET mode contains duplicate coupons"));
            }
            Ok(hash_set)
        } else {
            // Non-compact mode: full hash table with empty slots
            // Read entire hash table including empty slots
            let mut coupons = vec![Coupon::EMPTY; array_size];
            for (i, coupon) in coupons.iter_mut().enumerate() {
                let raw = cursor.read_u32_le().map_err(|_| {
                    Error::insufficient_data(format!(
                        "expected {array_size} coupons, failed at index {i}"
                    ))
                })?;
                *coupon = Coupon(raw);
            }
            if coupons.iter().filter(|coupon| !coupon.is_empty()).count() != coupon_count {
                return Err(Error::deserial(
                    "SET mode coupon count does not match occupied slots",
                ));
            }

            Ok(Self {
                container: Container::from_coupons(
                    lg_arr,
                    coupons.into_boxed_slice(),
                    coupon_count,
                ),
            })
        }
    }

    /// Serialize a HashSet to bytes
    pub fn serialize(&self, lg_config_k: u8, hll_type: HllType) -> Vec<u8> {
        let compact = true; // Always use compact format
        let coupon_count = self.container.len();
        let lg_arr = self.container.lg_size();

        // Compute size
        let array_size = if compact { coupon_count } else { 1 << lg_arr };
        let total_size = SET_PREAMBLE_SIZE + (array_size * 4);

        let mut bytes = SketchBytes::with_capacity(total_size);

        // Write preamble
        bytes.write_u8(HASH_SET_PREINTS);
        bytes.write_u8(SERIAL_VERSION);
        bytes.write_u8(Family::HLL.id);
        bytes.write_u8(lg_config_k);
        bytes.write_u8(lg_arr as u8);

        // Write flags
        let mut flags = 0u8;
        if compact {
            flags |= COMPACT_FLAG_MASK;
        }
        bytes.write_u8(flags);

        // Write unused byte
        bytes.write_u8(0);

        // Write mode byte: SET mode with target HLL type
        bytes.write_u8(encode_mode_byte(CUR_MODE_SET, hll_type as u8));

        // Write coupon count
        bytes.write_u32_le(coupon_count as u32);

        // Write coupons
        if compact {
            for coupon in self.container.iter() {
                bytes.write_u32_le(coupon.raw());
            }
        } else {
            // Non-compact mode: write entire hash table
            for coupon in self.container.coupons.iter().copied() {
                bytes.write_u32_le(coupon.raw());
            }
        }

        bytes.into_bytes()
    }
}