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.

//! HyperLogLog Array4 mode - 4-bit packed representation with exception handling
//!
//! Array4 stores HLL register values using 4 bits per slot (2 slots per byte).
//! When values exceed 4 bits after cur_min offset, they're stored in an auxiliary hash map.

use crate::codec::SketchBytes;
use crate::codec::SketchSlice;
use crate::codec::assert::insufficient_data;
use crate::codec::family::Family;
use crate::common::NumStdDev;
use crate::error::Error;
use crate::hll::Coupon;
use crate::hll::aux_map::AuxMap;
use crate::hll::estimator::EstimateState;
use crate::hll::estimator::Estimator;
use crate::hll::serialization::COMPACT_FLAG_MASK;
use crate::hll::serialization::COUPON_SIZE_BYTES;
use crate::hll::serialization::CUR_MODE_HLL;
use crate::hll::serialization::HLL_PREAMBLE_SIZE;
use crate::hll::serialization::HLL_PREINTS;
use crate::hll::serialization::OUT_OF_ORDER_FLAG_MASK;
use crate::hll::serialization::SERIAL_VERSION;
use crate::hll::serialization::TGT_HLL4;
use crate::hll::serialization::encode_mode_byte;

const AUX_TOKEN: u8 = 15;

#[derive(Clone, Copy)]
pub enum AuxFormat {
    Compact,
    Updatable { lg_arr: u8 },
}

impl AuxFormat {
    pub fn from_header(compact: bool, lg_arr: u8) -> Self {
        if compact {
            Self::Compact
        } else {
            Self::Updatable { lg_arr }
        }
    }
}

/// Core Array4 data structure - stores 4-bit values efficiently
#[derive(Debug, Clone, PartialEq)]
pub struct Array4 {
    lg_config_k: u8,
    /// Packed 4-bit values: 2 values per byte
    /// Even slots use low nibble, odd slots use high nibble
    bytes: Box<[u8]>,
    /// Current minimum value offset (optimization to delay aux map creation)
    cur_min: u8,
    /// Count of slots at exactly cur_min (when 0, increment cur_min)
    num_at_cur_min: u32,
    /// Exception table for values >= 15 after cur_min offset
    aux_map: Option<AuxMap>,
    estimator: Estimator,
}

impl Array4 {
    pub fn new(lg_config_k: u8) -> Self {
        let num_bytes = 1 << (lg_config_k - 1);
        let num_at_cur_min = 1 << lg_config_k;
        Self {
            lg_config_k,
            bytes: vec![0u8; num_bytes].into_boxed_slice(),
            cur_min: 0,
            num_at_cur_min,
            aux_map: None,
            estimator: Estimator::new(lg_config_k),
        }
    }

    /// Get raw 4-bit value from slot (not adjusted for cur_min)
    #[inline]
    fn get_raw(&self, slot: u32) -> u8 {
        debug_assert!(slot >> 1 < self.bytes.len() as u32);

        let byte = self.bytes[(slot >> 1) as usize];
        if slot & 1 == 0 {
            byte & 15 // low nibble for even slots
        } else {
            byte >> 4 // high nibble for odd slots
        }
    }

    /// Get the actual value at a slot (adjusted for cur_min and aux_map)
    ///
    /// Returns the true register value:
    /// * If raw < 15: value = cur_min + raw
    /// * If raw == 15 (AUX_TOKEN): value is in aux_map
    pub fn get(&self, slot: u32) -> u8 {
        let raw = self.get_raw(slot);

        if raw < AUX_TOKEN {
            self.cur_min + raw
        } else {
            // Value is in aux_map
            self.aux_map
                .as_ref()
                .and_then(|map| map.get(slot))
                .unwrap_or(self.cur_min) // Fallback (shouldn't happen)
        }
    }

    /// Get the number of registers (K = 2^lg_config_k)
    pub fn num_registers(&self) -> usize {
        1 << self.lg_config_k
    }

    /// Returns the estimate state independently from register-derived cached values.
    pub fn estimate_state(&self) -> EstimateState {
        self.estimator.estimate_state()
    }

    /// Set raw 4-bit value in slot
    #[inline]
    fn put_raw(&mut self, slot: u32, value: u8) {
        debug_assert!(value <= AUX_TOKEN);
        debug_assert!(slot >> 1 < self.bytes.len() as u32);

        let byte_idx = (slot >> 1) as usize;
        let old_byte = self.bytes[byte_idx];
        self.bytes[byte_idx] = if slot & 1 == 0 {
            (old_byte & 0xF0) | (value & 0x0F) // set low nibble
        } else {
            (old_byte & 0x0F) | (value << 4) // set high nibble
        };
    }

    pub fn update(&mut self, coupon: Coupon) {
        let mask = (1 << self.lg_config_k) - 1;
        let slot = coupon.slot() & mask;
        let new_value = coupon.value();

        // Quick rejection: if new value <= cur_min, no update needed
        if new_value <= self.cur_min {
            return;
        }

        let raw_stored = self.get_raw(slot);
        let lower_bound = raw_stored + self.cur_min;

        if new_value <= lower_bound {
            return;
        }

        // Get actual old value (might be in aux map)
        let old_value = if raw_stored < AUX_TOKEN {
            lower_bound
        } else {
            self.aux_map
                .as_ref()
                .expect("aux_map should be initialized since stored value is AUX_TOKEN")
                .get(slot)
                .expect("slot should be in aux_map since associated value is AUX_TOKEN")
        };

        if new_value <= old_value {
            return;
        }

        // Update HIP and KxQ registers via estimator
        self.estimator
            .update(self.lg_config_k, old_value, new_value);

        let shifted_new = new_value - self.cur_min;

        // Four cases based on old/new exception status
        match (raw_stored, shifted_new) {
            // Case 1: Both old and new are exceptions
            (AUX_TOKEN, shifted) if shifted >= AUX_TOKEN => {
                self.aux_map
                    .as_mut()
                    .expect("aux_map should be initialized since stored value is AUX_TOKEN")
                    .replace(slot, new_value);
            }
            // Case 2: Old is exception, new is not (impossible without cur_min change)
            (AUX_TOKEN, _) => {
                unreachable!("AUX_TOKEN present with non-exception new value");
            }
            // Case 3: Old not exception, new is exception
            (_, shifted) if shifted >= AUX_TOKEN => {
                self.put_raw(slot, AUX_TOKEN);
                let aux = self
                    .aux_map
                    .get_or_insert_with(|| AuxMap::new(self.lg_config_k));
                aux.insert(slot, new_value);
            }
            // Case 4: Neither is exception
            _ => {
                self.put_raw(slot, shifted_new);
            }
        }

        // Handle cur_min adjustment
        if old_value == self.cur_min {
            self.num_at_cur_min -= 1;
            while self.num_at_cur_min == 0 {
                self.shift_to_bigger_cur_min();
            }
        }
    }

    /// Increment cur_min and adjust all values
    ///
    /// This is called when no slots remain at cur_min value.
    /// All stored values are decremented by 1, and exceptions
    /// that fall back into the 4-bit range are moved from aux map.
    fn shift_to_bigger_cur_min(&mut self) {
        let new_cur_min = self.cur_min + 1;
        let k = 1 << self.lg_config_k;
        let mut num_at_new = 0;

        // Decrement all stored values in the main array
        for slot in 0..k {
            let raw = self.get_raw(slot);
            debug_assert_ne!(raw, 0, "value cannot be 0 when shifting cur_min");
            if raw < AUX_TOKEN {
                let decremented = raw - 1;
                self.put_raw(slot, decremented);
                if decremented == 0 {
                    num_at_new += 1;
                }
            }
        }

        // Rebuild aux map: some exceptions may no longer be exceptions
        if let Some(old_aux) = self.aux_map.take() {
            let mut new_aux = None;

            for (slot, old_actual_val) in old_aux.into_iter() {
                debug_assert_eq!(
                    self.get_raw(slot),
                    AUX_TOKEN,
                    "AuxMap contains slot != AUX_TOKEN"
                );

                let new_shifted = old_actual_val - new_cur_min;

                if new_shifted < AUX_TOKEN {
                    self.put_raw(slot, new_shifted);
                } else {
                    // Still an exception
                    let aux = new_aux.get_or_insert_with(|| AuxMap::new(self.lg_config_k));
                    aux.insert(slot, old_actual_val);
                }
            }
            self.aux_map = new_aux;
        }

        self.cur_min = new_cur_min;
        self.num_at_cur_min = num_at_new;
    }

    /// Returns the current cardinality estimate.
    pub fn estimate(&self) -> f64 {
        // Array4 tracks cur_min and num_at_cur_min dynamically
        self.estimator
            .estimate(self.lg_config_k, self.cur_min, self.num_at_cur_min)
    }

    /// Get upper bound for cardinality estimate
    pub fn upper_bound(&self, num_std_dev: NumStdDev) -> f64 {
        self.estimator.upper_bound(
            self.lg_config_k,
            self.cur_min,
            self.num_at_cur_min,
            num_std_dev,
        )
    }

    /// Get lower bound for cardinality estimate
    pub fn lower_bound(&self, num_std_dev: NumStdDev) -> f64 {
        self.estimator.lower_bound(
            self.lg_config_k,
            self.cur_min,
            self.num_at_cur_min,
            num_std_dev,
        )
    }

    /// Restores estimate state after copying or transforming the same logical sketch.
    pub fn restore_estimate_state(&mut self, state: EstimateState) {
        self.estimator.restore_estimate_state(state);
    }

    /// Check if the sketch is empty (all slots are zero)
    pub fn is_empty(&self) -> bool {
        self.num_at_cur_min == (1 << self.lg_config_k) && self.cur_min == 0
    }

    /// Deserialize Array4 from HLL mode bytes
    ///
    /// Expects full HLL preamble (40 bytes) followed by packed 4-bit data and optional aux map.
    pub fn deserialize(
        mut cursor: SketchSlice,
        cur_min: u8,
        lg_config_k: u8,
        aux_format: AuxFormat,
        ooo: bool,
    ) -> Result<Self, Error> {
        let k = 1usize << lg_config_k;
        let num_bytes = 1usize << (lg_config_k - 1); // k/2 bytes for 4-bit packing

        // Read estimator values from preamble
        let hip_accum = cursor
            .read_f64_le()
            .map_err(insufficient_data("hip_accum"))?;
        let kxq0 = cursor.read_f64_le().map_err(insufficient_data("kxq0"))?;
        let kxq1 = cursor.read_f64_le().map_err(insufficient_data("kxq1"))?;

        // Read num_at_cur_min and aux_count
        let num_at_cur_min = cursor
            .read_u32_le()
            .map_err(insufficient_data("num_at_cur_min"))?;
        let aux_count = cursor
            .read_u32_le()
            .map_err(insufficient_data("aux_count"))?;
        if num_at_cur_min as usize > k || aux_count as usize > k {
            return Err(Error::deserial(
                "HLL4 register or auxiliary count exceeds k",
            ));
        }
        let (aux_slots, compact) = match aux_format {
            AuxFormat::Compact => (aux_count as usize, true),
            AuxFormat::Updatable { lg_arr } => {
                let slots = if aux_count == 0 {
                    0
                } else {
                    1usize
                        .checked_shl(u32::from(lg_arr))
                        .ok_or_else(|| Error::deserial(format!("invalid HLL4 lg_arr: {lg_arr}")))?
                };
                (slots, false)
            }
        };
        let required_bytes = aux_slots
            .checked_mul(COUPON_SIZE_BYTES)
            .and_then(|aux_bytes| num_bytes.checked_add(aux_bytes))
            .ok_or_else(|| Error::deserial("HLL4 payload length overflows"))?;
        if required_bytes > cursor.remaining().len() {
            return Err(Error::insufficient_data(format!(
                "HLL4 payload requires {required_bytes} bytes, got {}",
                cursor.remaining().len()
            )));
        }

        // Read packed 4-bit byte array
        let mut data = vec![0u8; num_bytes];
        cursor
            .read_exact(&mut data)
            .map_err(insufficient_data("data"))?;

        // Read aux map if present
        let mut aux_map = None;
        if aux_count > 0 {
            let mut aux = AuxMap::new(lg_config_k);
            let mut decoded_count = 0;
            for i in 0..aux_slots {
                let coupon = cursor.read_u32_le().map_err(|_| {
                    Error::insufficient_data(format!(
                        "expected {aux_slots} HLL4 auxiliary slots, failed at index {i}",
                    ))
                })?;
                let coupon = Coupon(coupon);
                if coupon.is_empty() && !compact {
                    continue;
                }
                let slot = coupon.slot() & ((1 << lg_config_k) - 1);
                let value = coupon.value();
                if coupon.is_empty() || aux.get(slot).is_some() {
                    return Err(Error::deserial(
                        "HLL4 auxiliary entries must be non-empty and unique",
                    ));
                }
                aux.insert(slot, value);
                decoded_count += 1;
            }
            if decoded_count != aux_count as usize {
                return Err(Error::deserial(format!(
                    "HLL4 auxiliary count is {aux_count}, decoded {decoded_count}"
                )));
            }
            aux_map = Some(aux);
        }

        let estimator = Estimator::from_serialized(hip_accum, kxq0, kxq1, ooo);

        Ok(Self {
            lg_config_k,
            bytes: data.into_boxed_slice(),
            cur_min,
            num_at_cur_min,
            aux_map,
            estimator,
        })
    }

    /// Serialize Array4 to bytes
    ///
    /// Produces full HLL preamble (40 bytes) followed by packed 4-bit data and optional aux map.
    pub fn serialize(&self, lg_config_k: u8) -> Vec<u8> {
        let num_bytes = 1 << (lg_config_k - 1); // k/2 bytes for 4-bit packing

        // Collect aux map entries if present
        let aux_entries: Vec<(u32, u8)> = if let Some(aux) = &self.aux_map {
            aux.iter().collect()
        } else {
            vec![]
        };

        let aux_count = aux_entries.len() as u32;
        let total_size = HLL_PREAMBLE_SIZE + num_bytes + (aux_count as usize * COUPON_SIZE_BYTES);
        let mut bytes = SketchBytes::with_capacity(total_size);

        // Write standard header
        bytes.write_u8(HLL_PREINTS);
        bytes.write_u8(SERIAL_VERSION);
        bytes.write_u8(Family::HLL.id);
        bytes.write_u8(lg_config_k);
        bytes.write_u8(0); // unused for HLL mode

        // Write flags.
        // COMPACT_FLAG_MASK is always set: aux map entries are written as a compact sequential
        // list of populated entries only.
        let mut flags = COMPACT_FLAG_MASK;
        if self.estimator.uses_composite_estimate() {
            flags |= OUT_OF_ORDER_FLAG_MASK;
        }
        bytes.write_u8(flags);

        // Write cur_min
        bytes.write_u8(self.cur_min);

        // Mode byte: HLL mode with HLL4 type
        bytes.write_u8(encode_mode_byte(CUR_MODE_HLL, TGT_HLL4));

        // Write estimator values
        bytes.write_f64_le(self.estimator.hip_accum());
        bytes.write_f64_le(self.estimator.kxq0());
        bytes.write_f64_le(self.estimator.kxq1());

        // Write num_at_cur_min
        bytes.write_u32_le(self.num_at_cur_min);

        // Write aux_count
        bytes.write_u32_le(aux_count);

        // Write packed 4-bit byte array
        bytes.write(&self.bytes);

        // Write aux map entries if present
        for (slot, value) in aux_entries.iter().copied() {
            bytes.write_u32_le(Coupon::pack(slot, value).raw());
        }

        bytes.into_bytes()
    }

    /// Returns the estimated size of the heap allocations in bytes
    pub fn estimated_size(&self) -> usize {
        self.bytes.len()
            + self
                .aux_map
                .as_ref()
                .map(|a| a.estimated_size())
                .unwrap_or(0)
    }
}

#[cfg(test)]
mod tests {
    use googletest::assert_that;
    use googletest::prelude::gt;
    use googletest::prelude::is_finite;
    use googletest::prelude::lt;
    use googletest::prelude::none;

    use super::*;
    use crate::hll::Coupon;

    #[test]
    fn test_get_set_raw() {
        let mut data = Array4::new(4); // 16 buckets

        // Test even slot (low nibble)
        data.put_raw(0, 5);
        assert_eq!(data.get_raw(0), 5);

        // Test odd slot (high nibble)
        data.put_raw(1, 7);
        assert_eq!(data.get_raw(1), 7);

        // Both values should be stored in the same byte
        assert_eq!(data.bytes[0], 0x75); // 0111_0101 = 7 << 4 | 5

        // Test multiple slots
        data.put_raw(2, 15);
        data.put_raw(3, 3);
        assert_eq!(data.get_raw(2), 15);
        assert_eq!(data.get_raw(3), 3);
    }

    #[test]
    fn test_hip_estimator_basic() {
        let mut arr = Array4::new(10); // 1024 buckets

        // Initially estimate should be 0
        assert_eq!(arr.estimate(), 0.0);

        // Add some unique values to different slots
        for i in 0..10_000u32 {
            arr.update(Coupon::from_value(i));
        }

        // Estimate should be positive and roughly in the ballpark
        // (not exact, but should be non-zero and not NaN/Inf)
        let estimate = arr.estimate();

        assert_that!(estimate, gt(0.0));
        assert_that!(estimate, is_finite());
        assert_that!(estimate, lt(100_000.0));

        // Rough sanity check: with 100 updates to different slots,
        // estimate should be in a reasonable range (very loose bounds)
        assert_that!(estimate, gt(1_000.0));
        assert_that!(estimate, lt(100_000.0));
    }

    #[test]
    fn test_kxq_register_split() {
        let mut arr = Array4::new(8); // 256 buckets

        // Test that values < 32 and >= 32 are handled correctly
        arr.update(Coupon::pack(0, 10)); // value < 32, goes to kxq0
        arr.update(Coupon::pack(1, 40)); // value >= 32, goes to kxq1

        // Verify registers were updated (not exact values, just check they changed)
        // kxq0 should have decreased (we removed a 0 and added a 10)
        // Initial kxq0 = 256 (all zeros = 1.0 each)
        assert_that!(arr.estimator.kxq0(), lt(256.0));

        // kxq1 should have a small positive value (from 1/2^40)
        assert_that!(arr.estimator.kxq1(), gt(0.0));
        assert_that!(arr.estimator.kxq1(), lt(0.001));
    }

    #[test]
    fn test_shift_cur_min_rebuilds_aux_entry() {
        let lg_config_k = 4;
        let num_slots = 1_u32 << lg_config_k;
        let mut arr = Array4::new(lg_config_k);

        arr.update(Coupon::pack(0, 15));
        assert_eq!(arr.get_raw(0), AUX_TOKEN);
        assert_eq!(arr.aux_map.as_ref().and_then(|aux| aux.get(0)), Some(15));

        for slot in 1..num_slots {
            arr.update(Coupon::pack(slot, 1));
        }

        assert_eq!(arr.cur_min, 1);
        assert_eq!(arr.num_at_cur_min, num_slots - 1);
        assert_eq!(arr.get_raw(0), 14);
        assert_eq!(arr.get(0), 15);
        assert_that!(arr.aux_map, none());

        for slot in 1..num_slots {
            assert_eq!(arr.get(slot), 1);
        }
    }
}