Repository: pola-rs/polars
File: crates/polars-expr/src/hot_groups/fixed_index_table.rs (Rust)

Code before:
    /// This is useful for permanent entries which are handled externally.
    /// Returns the key index this would have taken up.
    pub fn push_unmapped_key(&mut self, key: K) -> IdxSize {
        let idx = self.keys.len();
        self.keys.push(key);
        idx as IdxSize
    }

    /// Tries to insert a key with a given hash.
    ///
    /// Returns Some((index, evict_old)) if successful, None otherwise.
    pub fn insert_key<Q, E, I, V>(
        &mut self,
        hash: u64,
        key: Q,
        mut eq: E,
        mut insert: I,
        mut evict_insert: V,
    ) -> Option<EvictIdx>
    where
        E: FnMut(&Q, &K) -> bool,
        I: FnMut(Q) -> K,
        V: FnMut(Q, &mut K),
    {
        let tag = hash as u32;
        let h1 = if (self.shift as u32) < u64::BITS {
            (hash >> self.shift) as usize
        } else {
            0
        };
        let h2 = if (self.shift as u32) < u64::BITS {
            (hash.wrapping_mul(H2_MULT) >> self.shift) as usize
        } else {
            0
        };

        unsafe {
            // We only want a single branch for the hot hit/miss check. This is
            // why we check both slots at once.
            let s1 = self.slots.get_unchecked(h1);
            let s2 = self.slots.get_unchecked(h2);
            let s1_delta = s1.tag ^ tag;
            let s2_delta = s2.tag ^ tag;
            // This check can have false positives (the binary AND of the deltas
            // happens to be zero by accident), but this is very unlikely (~1/10k)
            // and harmless if it does. False negatives are impossible. If this
            // branch succeeds we almost surely have a hit, if it fails
            // we're certain we have a miss.
            if s1_delta & s2_delta == 0 {
                // We want to branchlessly select the most likely candidate
                // first to ensure no further branch mispredicts in the vast

Review comment (performance):
Please revert this, there should not be branches/checks here. Just assert the shift fits during construction of the table (and also an assert that the table size must be a power of two). Supporting a hot table size of 1 is not worth lowering the performance for actual hot table sizes we'd use.

Code after:
    ///
    /// This is useful for permanent entries which are handled externally.
    /// Returns the key index this would have taken up.
    pub fn push_unmapped_key(&mut self, key: K) -> IdxSize {
        let idx = self.keys.len();
        self.keys.push(key);
        idx as IdxSize
    }

    /// Tries to insert a key with a given hash.
    ///
    /// Returns Some((index, evict_old)) if successful, None otherwise.
    pub fn insert_key<Q, E, I, V>(
        &mut self,
        hash: u64,
        key: Q,
        mut eq: E,
        mut insert: I,
        mut evict_insert: V,
    ) -> Option<EvictIdx>
    where
        E: FnMut(&Q, &K) -> bool,
        I: FnMut(Q) -> K,
        V: FnMut(Q, &mut K),
    {
        let tag = hash as u32;
        let h1 = (hash >> self.shift) as usize;
        let h2 = (hash.wrapping_mul(H2_MULT) >> self.shift) as usize;

        unsafe {
            // We only want a single branch for the hot hit/miss check. This is
            // why we check both slots at once.
            let s1 = self.slots.get_unchecked(h1);
            let s2 = self.slots.get_unchecked(h2);
            let s1_delta = s1.tag ^ tag;
            let s2_delta = s2.tag ^ tag;
            // This check can have false positives (the binary AND of the deltas
            // happens to be zero by accident), but this is very unlikely (~1/10k)
            // and harmless if it does. False negatives are impossible. If this
            // branch succeeds we almost surely have a hit, if it fails
            // we're certain we have a miss.
            if s1_delta & s2_delta == 0 {
                // We want to branchlessly select the most likely candidate
                // first to ensure no further branch mispredicts in the vast
                // majority of cases.
                let ha = select_unpredictable(s1_delta == 0, h1, h2);
                let sa = self.slots.get_unchecked_mut(ha);
                if let Some(sak) = self.keys.get(sa.key_index as usize) {
                    if eq(&key, sak) {
                        sa.last_access_tag = tag;
                        return Some(EvictIdx::new(sa.key_index, false));