datasketches 0.6.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.

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::estimator::EstimateState;
use crate::hll::estimator::Estimator;
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_HLL6;
use crate::hll::serialization::encode_mode_byte;

const VAL_MASK_6: u16 = 0x3F;

/// Registers packed into consecutive six-bit fields.
#[derive(Debug, Clone, PartialEq)]
pub struct Array6 {
    lg_config_k: u8,
    bytes: Box<[u8]>,
    num_zeros: u32,
    estimator: Estimator,
}

impl Array6 {
    pub fn new(lg_config_k: u8) -> Self {
        let k = 1 << lg_config_k;
        let num_bytes = num_bytes_for_k(k);

        Self {
            lg_config_k,
            bytes: vec![0u8; num_bytes].into_boxed_slice(),
            num_zeros: k,
            estimator: Estimator::new(lg_config_k),
        }
    }

    #[inline]
    fn get_raw(&self, slot: u32) -> u8 {
        let start_bit = slot * 6;
        let byte_idx = (start_bit >> 3) as usize;
        let shift = (start_bit & 7) as u8;

        let two_bytes = u16::from_le_bytes([self.bytes[byte_idx], self.bytes[byte_idx + 1]]);

        ((two_bytes >> shift) & VAL_MASK_6) as u8
    }

    #[inline]
    pub fn get(&self, slot: u32) -> u8 {
        self.get_raw(slot)
    }

    pub fn num_registers(&self) -> usize {
        1 << self.lg_config_k
    }

    pub fn estimate_state(&self) -> EstimateState {
        self.estimator.estimate_state()
    }

    #[inline]
    fn put_raw(&mut self, slot: u32, value: u8) {
        debug_assert!(value <= 63, "6-bit value must be 0-63");

        let start_bit = slot * 6;
        let byte_idx = (start_bit >> 3) as usize;
        let shift = (start_bit & 0x7) as u8;

        let mut two_bytes = u16::from_le_bytes([self.bytes[byte_idx], self.bytes[byte_idx + 1]]);

        two_bytes &= !(VAL_MASK_6 << shift);

        two_bytes |= ((value as u16) & VAL_MASK_6) << shift;

        let bytes_out = two_bytes.to_le_bytes();
        self.bytes[byte_idx] = bytes_out[0];
        self.bytes[byte_idx + 1] = bytes_out[1];
    }

    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();

        let old_value = self.get_raw(slot);

        if new_value > old_value {
            self.estimator
                .update(self.lg_config_k, old_value, new_value);

            self.put_raw(slot, new_value);

            if old_value == 0 {
                self.num_zeros -= 1;
            }
        }
    }

    pub fn estimate(&self) -> f64 {
        // Array6 doesn't use cur_min (always 0), so num_at_cur_min = num_zeros
        self.estimator.estimate(self.lg_config_k, 0, self.num_zeros)
    }

    pub fn composite_estimate(&self) -> f64 {
        self.estimator
            .composite_estimate(self.lg_config_k, 0, self.num_zeros)
    }

    pub fn upper_bound(&self, num_std_dev: NumStdDev) -> f64 {
        self.estimator
            .upper_bound(self.lg_config_k, 0, self.num_zeros, num_std_dev)
    }

    pub fn lower_bound(&self, num_std_dev: NumStdDev) -> f64 {
        self.estimator
            .lower_bound(self.lg_config_k, 0, self.num_zeros, num_std_dev)
    }

    pub fn restore_estimate_state(&mut self, state: EstimateState) {
        self.estimator.restore_estimate_state(state);
    }

    pub fn is_empty(&self) -> bool {
        self.num_zeros == (1 << self.lg_config_k)
    }

    pub fn deserialize_registers(
        mut cursor: SketchSlice,
        lg_config_k: u8,
        ooo: bool,
    ) -> Result<Self, Error> {
        let k = 1 << lg_config_k;
        let num_bytes = num_bytes_for_k(k);

        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 (for Array6, this is num_zeros since cur_min=0)
        let num_zeros = cursor
            .read_u32_le()
            .map_err(insufficient_data("num_zeros"))?;
        let aux_count = cursor
            .read_u32_le()
            .map_err(insufficient_data("aux_count"))?;
        if num_zeros > k || aux_count != 0 {
            return Err(Error::deserial(
                "HLL6 zero count must not exceed k and auxiliary count must be zero",
            ));
        }
        let available_bytes = cursor.remaining().len();
        if available_bytes < num_bytes {
            return Err(Error::insufficient_data_of(
                "HLL6 payload",
                format_args!("expected {num_bytes} bytes, got {available_bytes}"),
            ));
        }

        let mut data = vec![0u8; num_bytes];
        cursor
            .read_exact(&mut data)
            .map_err(insufficient_data("data"))?;

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

        Ok(Self {
            lg_config_k,
            bytes: data.into_boxed_slice(),
            num_zeros,
            estimator,
        })
    }

    pub fn serialize(&self, lg_config_k: u8) -> Vec<u8> {
        let k = 1 << lg_config_k;
        let num_bytes = num_bytes_for_k(k);
        let total_size = HLL_PREAMBLE_SIZE + num_bytes;
        let mut bytes = SketchBytes::with_capacity(total_size);

        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

        let mut flags = 0u8;
        if self.estimator.uses_composite_estimate() {
            flags |= OUT_OF_ORDER_FLAG_MASK;
        }
        bytes.write_u8(flags);

        // cur_min is always 0 for Array6
        bytes.write_u8(0);

        bytes.write_u8(encode_mode_byte(CUR_MODE_HLL, TGT_HLL6));

        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 (num_zeros for Array6)
        bytes.write_u32_le(self.num_zeros);

        // Write aux_count (always 0 for Array6)
        bytes.write_u32_le(0);

        bytes.write(&self.bytes);

        bytes.into_bytes()
    }

    /// Returns an estimate of memory usage in bytes.
    pub fn estimated_size(&self) -> usize {
        self.bytes.len()
    }
}

fn num_bytes_for_k(k: u32) -> usize {
    // The final register's 16-bit read/write window needs one padding byte.
    (((k * 3) >> 2) + 1) as usize
}

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

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

    #[test]
    fn test_num_bytes_calculation() {
        assert_eq!(num_bytes_for_k(16), (16 * 3 / 4) + 1);

        assert_eq!(num_bytes_for_k(1024), (1024 * 3 / 4) + 1);
    }

    #[test]
    fn test_get_set_raw() {
        let mut arr = Array6::new(4); // 16 slots

        arr.put_raw(0, 0);
        arr.put_raw(1, 1);
        arr.put_raw(2, 31);
        arr.put_raw(3, 63); // Max 6-bit value

        assert_eq!(arr.get_raw(0), 0);
        assert_eq!(arr.get_raw(1), 1);
        assert_eq!(arr.get_raw(2), 31);
        assert_eq!(arr.get_raw(3), 63);

        arr.put_raw(5, 42);
        assert_eq!(arr.get_raw(5), 42);
        assert_eq!(arr.get_raw(3), 63); // Earlier value unchanged

        // Test all slots to ensure no cross-slot corruption
        for slot in 0..16 {
            arr.put_raw(slot, (slot % 64) as u8);
        }
        for slot in 0..16 {
            assert_eq!(arr.get_raw(slot), (slot % 64) as u8);
        }
    }

    #[test]
    fn test_boundary_crossing() {
        let mut arr = Array6::new(8); // 256 slots

        // Slot 1: starts at bit 6 (crosses byte 0/1 boundary)
        arr.put_raw(1, 0b111111);
        assert_eq!(arr.get_raw(1), 63);

        // Slot 2: starts at bit 12 (in byte 1)
        arr.put_raw(2, 0b101010);
        assert_eq!(arr.get_raw(2), 42);

        // Slot 3: starts at bit 18 (crosses byte 2/3 boundary)
        arr.put_raw(3, 0b110011);
        assert_eq!(arr.get_raw(3), 51);

        assert_eq!(arr.get_raw(1), 63);
        assert_eq!(arr.get_raw(2), 42);
        assert_eq!(arr.get_raw(3), 51);
    }

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

        assert_eq!(arr.estimate(), 0.0);

        for i in 0..10_000u32 {
            arr.update(Coupon::from_value(i));
        }

        let estimate = arr.estimate();

        assert_that!(estimate, gt(0.0));
        assert_that!(estimate, is_finite());

        assert_that!(estimate, gt(1_000.0));
        assert_that!(estimate, lt(100_000.0));
    }

    #[test]
    fn test_full_range() {
        let mut arr = Array6::new(6); // 64 slots

        for val in 0..64u8 {
            arr.put_raw(val as u32, val);
        }

        for val in 0..64u8 {
            assert_eq!(arr.get_raw(val as u32), val);
        }
    }

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

        arr.update(Coupon::pack(0, 10)); // value < 32, goes to kxq0
        arr.update(Coupon::pack(1, 40)); // value >= 32, goes to kxq1

        // 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));
    }
}