<commit_before>//! This pass eliminates casting of arrays into slices when their length
//! is taken using `.len()` method. Handy to preserve information in MIR for const prop

use crate::transform::MirPass;
use rustc_data_structures::fx::FxIndexMap;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, TyCtxt};

const MAX_NUM_BLOCKS: usize = 800;
const MAX_NUM_LOCALS: usize = 3000;

pub struct NormalizeArrayLen;

impl<'tcx> MirPass<'tcx> for NormalizeArrayLen {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        // if tcx.sess.mir_opt_level() < 3 {
        //     return;
        // }

        // early returns for edge cases of highly unrolled functions
        if body.basic_blocks().len() > MAX_NUM_BLOCKS {
            return;
        }
        if body.local_decls().len() > MAX_NUM_LOCALS {
            return;
        }
        normalize_array_len_calls(tcx, body)
    }
}

pub fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
    let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();

    // do a preliminary analysis to see if we ever have locals of type `[T;N]` or `&[T;N]`
    let mut interesting_locals = BitSet::new_empty(local_decls.len());
    for (local, decl) in local_decls.iter_enumerated() {
        match decl.ty.kind() {
            ty::Array(..) => {
                interesting_locals.insert(local);
            }
            ty::Ref(.., ty, Mutability::Not) => match ty.kind() {
                ty::Array(..) => {
                    interesting_locals.insert(local);
                }
                _ => {}
            },
            _ => {}
        }
    }
    if interesting_locals.is_empty() {
        // we have found nothing to analyze
        return;
    }
    let num_intesting_locals = interesting_locals.count();
    let mut state = FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
    let mut patches_scratchpad =
        FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
    let mut replacements_scratchpad =
        FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
    for block in basic_blocks {
        // make length calls for arrays [T; N] not to decay into length calls for &[T]
        // that forbids constant propagation
        normalize_array_len_call(
            tcx,
            block,
            local_decls,
            &interesting_locals,
            &mut state,
            &mut patches_scratchpad,
            &mut replacements_scratchpad,
        );
        state.clear();
        patches_scratchpad.clear();
        replacements_scratchpad.clear();
    }
}

struct Patcher<'a, 'tcx> {
    tcx: TyCtxt<'tcx>,
    patches_scratchpad: &'a FxIndexMap<usize, usize>,
    replacements_scratchpad: &'a mut FxIndexMap<usize, Local>,
    local_decls: &'a mut IndexVec<Local, LocalDecl<'tcx>>,
    statement_idx: usize,
}

impl<'a, 'tcx> Patcher<'a, 'tcx> {
    fn patch_expand_statement(
        &mut self,
        statement: &mut Statement<'tcx>,
    ) -> Option<std::vec::IntoIter<Statement<'tcx>>> {
        let idx = self.statement_idx;
        if let Some(len_statemnt_idx) = self.patches_scratchpad.get(&idx).copied() {
            let mut statements = Vec::with_capacity(2);

            // we are at statement that performs a cast. The only sound way is
            // to create another local that performs a similar copy without a cast and then
            // use this copy in the Len operation

            match &statement.kind {
                StatementKind::Assign(box (
                    ..,
                    Rvalue::Cast(
                        CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
                        operand,
                        _,
                    ),
                )) => {
                    match operand {
                        Operand::Copy(place) | Operand::Move(place) => {
                            // create new local
                            let ty = operand.ty(self.local_decls, self.tcx);
                            let local_decl =
                                LocalDecl::with_source_info(ty, statement.source_info.clone());
                            let local = self.local_decls.push(local_decl);
                            // make it live
                            let mut make_live_statement = statement.clone();
                            make_live_statement.kind = StatementKind::StorageLive(local);
                            statements.push(make_live_statement);
                            // copy into it

                            let operand = Operand::Copy(*place);
                            let mut make_copy_statement = statement.clone();
                            let assign_to = Place::from(local);
                            let rvalue = Rvalue::Use(operand);
                            make_copy_statement.kind =
                                StatementKind::Assign(box (assign_to, rvalue));
                            statements.push(make_copy_statement);

                            // to reorder we have to copy and make NOP
                            statements.push(statement.clone());
                            statement.make_nop();

                            self.replacements_scratchpad.insert(len_statemnt_idx, local);
                        }
                        _ => {
                            unreachable!("it's a bug in the implementation")
                        }
                    }
                }
                _ => {
                    unreachable!("it's a bug in the implementation")
                }
            }

            self.statement_idx += 1;

            Some(statements.into_iter())
        } else if let Some(local) = self.replacements_scratchpad.get(&idx).copied() {
            let mut statements = Vec::with_capacity(2);

            match &statement.kind {
                StatementKind::Assign(box (into, Rvalue::Len(place))) => {
                    let add_deref = if let Some(..) = place.as_local() {
                        false
                    } else if let Some(..) = place.local_or_deref_local() {
                        true
                    } else {
                        unreachable!("it's a bug in the implementation")
                    };
                    // replace len statement
                    let mut len_statement = statement.clone();
                    let mut place = Place::from(local);
                    if add_deref {
                        place = self.tcx.mk_place_deref(place);
                    }
                    len_statement.kind = StatementKind::Assign(box (*into, Rvalue::Len(place)));
                    statements.push(len_statement);

                    // make temporary dead
                    let mut make_dead_statement = statement.clone();
                    make_dead_statement.kind = StatementKind::StorageDead(local);
                    statements.push(make_dead_statement);

                    // make original statement NOP
                    statement.make_nop();
                }
                _ => {
                    unreachable!("it's a bug in the implementation")
                }
            }

            self.statement_idx += 1;

            Some(statements.into_iter())
        } else {
            self.statement_idx += 1;
            None
        }
    }
}

fn normalize_array_len_call<'tcx>(
    tcx: TyCtxt<'tcx>,
    block: &mut BasicBlockData<'tcx>,
    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
    interesting_locals: &BitSet<Local>,
    state: &mut FxIndexMap<Local, usize>,
    patches_scratchpad: &mut FxIndexMap<usize, usize>,
    replacements_scratchpad: &mut FxIndexMap<usize, Local>,
) {
    for (statement_idx, statement) in block.statements.iter_mut().enumerate() {
        match &mut statement.kind {
            StatementKind::Assign(box (place, rvalue)) => {
                match rvalue {
                    Rvalue::Cast(
                        CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
                        operand,
                        cast_ty,
                    ) => {
                        let local = if let Some(local) = place.as_local() { local } else { return };
                        match operand {
                            Operand::Copy(place) | Operand::Move(place) => {
                                let operand_local =
                                    if let Some(local) = place.local_or_deref_local() {
                                        local
                                    } else {
                                        return;
                                    };
                                if !interesting_locals.contains(operand_local) {
                                    return;
                                }
                                let operand_ty = local_decls[operand_local].ty;
                                match (operand_ty.kind(), cast_ty.kind()) {
                                    (ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
                                        if of_ty_src == of_ty_dst {
                                            // this is a cast from [T; N] into [T], so we are good
                                            state.insert(local, statement_idx);
                                        }
                                    }
                                    // current way of patching doesn't allow to work with `mut`
                                    (
                                        ty::Ref(
                                            ty::RegionKind::ReErased,
                                            operand_ty,
                                            Mutability::Not,
                                        ),
                                        ty::Ref(ty::RegionKind::ReErased, cast_ty, Mutability::Not),
                                    ) => {
                                        match (operand_ty.kind(), cast_ty.kind()) {
                                            // current way of patching doesn't allow to work with `mut`
                                            (ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
                                                if of_ty_src == of_ty_dst {
                                                    // this is a cast from [T; N] into [T], so we are good
                                                    state.insert(local, statement_idx);
                                                }
                                            }
                                            _ => {}
                                        }
                                    }
                                    _ => {}
                                }
                            }
                            _ => {}
                        }
                    }
                    Rvalue::Len(place) => {
                        let local = if let Some(local) = place.local_or_deref_local() {
                            local
                        } else {
                            return;
                        };
                        if let Some(cast_statement_idx) = state.get(&local).copied() {
                            patches_scratchpad.insert(cast_statement_idx, statement_idx);
                        }
                    }
                    _ => {
                        // invalidate
                        state.remove(&place.local);
                    }
                }
            }
            _ => {}
        }
    }

    let mut patcher = Patcher {
        tcx,
        patches_scratchpad: &*patches_scratchpad,
        replacements_scratchpad,
        local_decls,
        statement_idx: 0,
    };

    block.expand_statements(|st| patcher.patch_expand_statement(st));
}
<commit_msg>fix import<commit_after>//! This pass eliminates casting of arrays into slices when their length
//! is taken using `.len()` method. Handy to preserve information in MIR for const prop

use crate::MirPass;
use rustc_data_structures::fx::FxIndexMap;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, TyCtxt};

const MAX_NUM_BLOCKS: usize = 800;
const MAX_NUM_LOCALS: usize = 3000;

pub struct NormalizeArrayLen;

impl<'tcx> MirPass<'tcx> for NormalizeArrayLen {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        // if tcx.sess.mir_opt_level() < 3 {
        //     return;
        // }

        // early returns for edge cases of highly unrolled functions
        if body.basic_blocks().len() > MAX_NUM_BLOCKS {
            return;
        }
        if body.local_decls().len() > MAX_NUM_LOCALS {
            return;
        }
        normalize_array_len_calls(tcx, body)
    }
}

pub fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
    let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();

    // do a preliminary analysis to see if we ever have locals of type `[T;N]` or `&[T;N]`
    let mut interesting_locals = BitSet::new_empty(local_decls.len());
    for (local, decl) in local_decls.iter_enumerated() {
        match decl.ty.kind() {
            ty::Array(..) => {
                interesting_locals.insert(local);
            }
            ty::Ref(.., ty, Mutability::Not) => match ty.kind() {
                ty::Array(..) => {
                    interesting_locals.insert(local);
                }
                _ => {}
            },
            _ => {}
        }
    }
    if interesting_locals.is_empty() {
        // we have found nothing to analyze
        return;
    }
    let num_intesting_locals = interesting_locals.count();
    let mut state = FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
    let mut patches_scratchpad =
        FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
    let mut replacements_scratchpad =
        FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
    for block in basic_blocks {
        // make length calls for arrays [T; N] not to decay into length calls for &[T]
        // that forbids constant propagation
        normalize_array_len_call(
            tcx,
            block,
            local_decls,
            &interesting_locals,
            &mut state,
            &mut patches_scratchpad,
            &mut replacements_scratchpad,
        );
        state.clear();
        patches_scratchpad.clear();
        replacements_scratchpad.clear();
    }
}

struct Patcher<'a, 'tcx> {
    tcx: TyCtxt<'tcx>,
    patches_scratchpad: &'a FxIndexMap<usize, usize>,
    replacements_scratchpad: &'a mut FxIndexMap<usize, Local>,
    local_decls: &'a mut IndexVec<Local, LocalDecl<'tcx>>,
    statement_idx: usize,
}

impl<'a, 'tcx> Patcher<'a, 'tcx> {
    fn patch_expand_statement(
        &mut self,
        statement: &mut Statement<'tcx>,
    ) -> Option<std::vec::IntoIter<Statement<'tcx>>> {
        let idx = self.statement_idx;
        if let Some(len_statemnt_idx) = self.patches_scratchpad.get(&idx).copied() {
            let mut statements = Vec::with_capacity(2);

            // we are at statement that performs a cast. The only sound way is
            // to create another local that performs a similar copy without a cast and then
            // use this copy in the Len operation

            match &statement.kind {
                StatementKind::Assign(box (
                    ..,
                    Rvalue::Cast(
                        CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
                        operand,
                        _,
                    ),
                )) => {
                    match operand {
                        Operand::Copy(place) | Operand::Move(place) => {
                            // create new local
                            let ty = operand.ty(self.local_decls, self.tcx);
                            let local_decl =
                                LocalDecl::with_source_info(ty, statement.source_info.clone());
                            let local = self.local_decls.push(local_decl);
                            // make it live
                            let mut make_live_statement = statement.clone();
                            make_live_statement.kind = StatementKind::StorageLive(local);
                            statements.push(make_live_statement);
                            // copy into it

                            let operand = Operand::Copy(*place);
                            let mut make_copy_statement = statement.clone();
                            let assign_to = Place::from(local);
                            let rvalue = Rvalue::Use(operand);
                            make_copy_statement.kind =
                                StatementKind::Assign(box (assign_to, rvalue));
                            statements.push(make_copy_statement);

                            // to reorder we have to copy and make NOP
                            statements.push(statement.clone());
                            statement.make_nop();

                            self.replacements_scratchpad.insert(len_statemnt_idx, local);
                        }
                        _ => {
                            unreachable!("it's a bug in the implementation")
                        }
                    }
                }
                _ => {
                    unreachable!("it's a bug in the implementation")
                }
            }

            self.statement_idx += 1;

            Some(statements.into_iter())
        } else if let Some(local) = self.replacements_scratchpad.get(&idx).copied() {
            let mut statements = Vec::with_capacity(2);

            match &statement.kind {
                StatementKind::Assign(box (into, Rvalue::Len(place))) => {
                    let add_deref = if let Some(..) = place.as_local() {
                        false
                    } else if let Some(..) = place.local_or_deref_local() {
                        true
                    } else {
                        unreachable!("it's a bug in the implementation")
                    };
                    // replace len statement
                    let mut len_statement = statement.clone();
                    let mut place = Place::from(local);
                    if add_deref {
                        place = self.tcx.mk_place_deref(place);
                    }
                    len_statement.kind = StatementKind::Assign(box (*into, Rvalue::Len(place)));
                    statements.push(len_statement);

                    // make temporary dead
                    let mut make_dead_statement = statement.clone();
                    make_dead_statement.kind = StatementKind::StorageDead(local);
                    statements.push(make_dead_statement);

                    // make original statement NOP
                    statement.make_nop();
                }
                _ => {
                    unreachable!("it's a bug in the implementation")
                }
            }

            self.statement_idx += 1;

            Some(statements.into_iter())
        } else {
            self.statement_idx += 1;
            None
        }
    }
}

fn normalize_array_len_call<'tcx>(
    tcx: TyCtxt<'tcx>,
    block: &mut BasicBlockData<'tcx>,
    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
    interesting_locals: &BitSet<Local>,
    state: &mut FxIndexMap<Local, usize>,
    patches_scratchpad: &mut FxIndexMap<usize, usize>,
    replacements_scratchpad: &mut FxIndexMap<usize, Local>,
) {
    for (statement_idx, statement) in block.statements.iter_mut().enumerate() {
        match &mut statement.kind {
            StatementKind::Assign(box (place, rvalue)) => {
                match rvalue {
                    Rvalue::Cast(
                        CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
                        operand,
                        cast_ty,
                    ) => {
                        let local = if let Some(local) = place.as_local() { local } else { return };
                        match operand {
                            Operand::Copy(place) | Operand::Move(place) => {
                                let operand_local =
                                    if let Some(local) = place.local_or_deref_local() {
                                        local
                                    } else {
                                        return;
                                    };
                                if !interesting_locals.contains(operand_local) {
                                    return;
                                }
                                let operand_ty = local_decls[operand_local].ty;
                                match (operand_ty.kind(), cast_ty.kind()) {
                                    (ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
                                        if of_ty_src == of_ty_dst {
                                            // this is a cast from [T; N] into [T], so we are good
                                            state.insert(local, statement_idx);
                                        }
                                    }
                                    // current way of patching doesn't allow to work with `mut`
                                    (
                                        ty::Ref(
                                            ty::RegionKind::ReErased,
                                            operand_ty,
                                            Mutability::Not,
                                        ),
                                        ty::Ref(ty::RegionKind::ReErased, cast_ty, Mutability::Not),
                                    ) => {
                                        match (operand_ty.kind(), cast_ty.kind()) {
                                            // current way of patching doesn't allow to work with `mut`
                                            (ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
                                                if of_ty_src == of_ty_dst {
                                                    // this is a cast from [T; N] into [T], so we are good
                                                    state.insert(local, statement_idx);
                                                }
                                            }
                                            _ => {}
                                        }
                                    }
                                    _ => {}
                                }
                            }
                            _ => {}
                        }
                    }
                    Rvalue::Len(place) => {
                        let local = if let Some(local) = place.local_or_deref_local() {
                            local
                        } else {
                            return;
                        };
                        if let Some(cast_statement_idx) = state.get(&local).copied() {
                            patches_scratchpad.insert(cast_statement_idx, statement_idx);
                        }
                    }
                    _ => {
                        // invalidate
                        state.remove(&place.local);
                    }
                }
            }
            _ => {}
        }
    }

    let mut patcher = Patcher {
        tcx,
        patches_scratchpad: &*patches_scratchpad,
        replacements_scratchpad,
        local_decls,
        statement_idx: 0,
    };

    block.expand_statements(|st| patcher.patch_expand_statement(st));
}
<|endoftext|>


<commit_before>use clear::Clear;
use std::collections::hash_map;
use std::collections::hash_map::DefaultHasher;
use std::collections::HashMap;
use std::default::Default;
use std::hash::Hash;
use std::hash::Hasher;
use std::slice;
use wire_format;
use zigzag::encode_zig_zag_32;
use zigzag::encode_zig_zag_64;

#[derive(Debug)]
pub enum UnknownValue {
    Fixed32(u32),
    Fixed64(u64),
    Varint(u64),
    LengthDelimited(Vec<u8>),
}

impl UnknownValue {
    pub fn wire_type(&self) -> wire_format::WireType {
        self.get_ref().wire_type()
    }

    pub fn get_ref<'s>(&'s self) -> UnknownValueRef<'s> {
        match *self {
            UnknownValue::Fixed32(fixed32) => UnknownValueRef::Fixed32(fixed32),
            UnknownValue::Fixed64(fixed64) => UnknownValueRef::Fixed64(fixed64),
            UnknownValue::Varint(varint) => UnknownValueRef::Varint(varint),
            UnknownValue::LengthDelimited(ref bytes) => UnknownValueRef::LengthDelimited(&bytes),
        }
    }

    /// Construct unknown value from `sint32` value.
    pub fn sint32(i: i32) -> UnknownValue {
        UnknownValue::Varint(encode_zig_zag_32(i) as u64)
    }

    /// Construct unknown value from `sint64` value.
    pub fn sint64(i: i64) -> UnknownValue {
        UnknownValue::Varint(encode_zig_zag_64(i))
    }
}

pub enum UnknownValueRef<'o> {
    Fixed32(u32),
    Fixed64(u64),
    Varint(u64),
    LengthDelimited(&'o [u8]),
}

impl<'o> UnknownValueRef<'o> {
    pub fn wire_type(&self) -> wire_format::WireType {
        match *self {
            UnknownValueRef::Fixed32(_) => wire_format::WireTypeFixed32,
            UnknownValueRef::Fixed64(_) => wire_format::WireTypeFixed64,
            UnknownValueRef::Varint(_) => wire_format::WireTypeVarint,
            UnknownValueRef::LengthDelimited(_) => wire_format::WireTypeLengthDelimited,
        }
    }
}

#[derive(Clone, PartialEq, Eq, Debug, Default, Hash)]
pub struct UnknownValues {
    pub fixed32: Vec<u32>,
    pub fixed64: Vec<u64>,
    pub varint: Vec<u64>,
    pub length_delimited: Vec<Vec<u8>>,
}

impl UnknownValues {
    pub fn add_value(&mut self, value: UnknownValue) {
        match value {
            UnknownValue::Fixed64(fixed64) => self.fixed64.push(fixed64),
            UnknownValue::Fixed32(fixed32) => self.fixed32.push(fixed32),
            UnknownValue::Varint(varint) => self.varint.push(varint),
            UnknownValue::LengthDelimited(length_delimited) => {
                self.length_delimited.push(length_delimited)
            }
        };
    }

    pub fn iter<'s>(&'s self) -> UnknownValuesIter<'s> {
        UnknownValuesIter {
            fixed32: self.fixed32.iter(),
            fixed64: self.fixed64.iter(),
            varint: self.varint.iter(),
            length_delimited: self.length_delimited.iter(),
        }
    }
}

impl<'a> IntoIterator for &'a UnknownValues {
    type Item = UnknownValueRef<'a>;
    type IntoIter = UnknownValuesIter<'a>;

    fn into_iter(self) -> UnknownValuesIter<'a> {
        self.iter()
    }
}

pub struct UnknownValuesIter<'o> {
    fixed32: slice::Iter<'o, u32>,
    fixed64: slice::Iter<'o, u64>,
    varint: slice::Iter<'o, u64>,
    length_delimited: slice::Iter<'o, Vec<u8>>,
}

impl<'o> Iterator for UnknownValuesIter<'o> {
    type Item = UnknownValueRef<'o>;

    fn next(&mut self) -> Option<UnknownValueRef<'o>> {
        let fixed32 = self.fixed32.next();
        if fixed32.is_some() {
            return Some(UnknownValueRef::Fixed32(*fixed32.unwrap()));
        }
        let fixed64 = self.fixed64.next();
        if fixed64.is_some() {
            return Some(UnknownValueRef::Fixed64(*fixed64.unwrap()));
        }
        let varint = self.varint.next();
        if varint.is_some() {
            return Some(UnknownValueRef::Varint(*varint.unwrap()));
        }
        let length_delimited = self.length_delimited.next();
        if length_delimited.is_some() {
            return Some(UnknownValueRef::LengthDelimited(&length_delimited.unwrap()));
        }
        None
    }
}

#[derive(Clone, PartialEq, Eq, Debug, Default)]
pub struct UnknownFields {
    // option is needed, because HashMap constructor performs allocation,
    // and very expensive
    pub fields: Option<Box<HashMap<u32, UnknownValues>>>,
}

/// Very simple hash implementation of `Hash` for `UnknownFields`.
/// Since map is unordered, we cannot put entry hashes into hasher,
/// instead we summing hashes of entries.
impl Hash for UnknownFields {
    fn hash<H: Hasher>(&self, state: &mut H) {
        if let Some(ref map) = self.fields {
            if !map.is_empty() {
                let mut hash: u64 = 0;
                for (k, v) in &**map {
                    let mut entry_hasher = DefaultHasher::new();
                    Hash::hash(&(k, v), &mut entry_hasher);
                    hash.wrapping_add(entry_hasher.finish());
                }
                Hash::hash(&hash, state);
            }
        }
    }
}

impl UnknownFields {
    pub fn new() -> UnknownFields {
        Default::default()
    }

    fn init_map(&mut self) {
        if self.fields.is_none() {
            self.fields = Some(Default::default());
        }
    }

    fn find_field<'a>(&'a mut self, number: &'a u32) -> &'a mut UnknownValues {
        self.init_map();

        match self.fields.as_mut().unwrap().entry(*number) {
            hash_map::Entry::Occupied(e) => e.into_mut(),
            hash_map::Entry::Vacant(e) => e.insert(Default::default()),
        }
    }

    pub fn add_fixed32(&mut self, number: u32, fixed32: u32) {
        self.find_field(&number).fixed32.push(fixed32);
    }

    pub fn add_fixed64(&mut self, number: u32, fixed64: u64) {
        self.find_field(&number).fixed64.push(fixed64);
    }

    pub fn add_varint(&mut self, number: u32, varint: u64) {
        self.find_field(&number).varint.push(varint);
    }

    pub fn add_length_delimited(&mut self, number: u32, length_delimited: Vec<u8>) {
        self.find_field(&number)
            .length_delimited
            .push(length_delimited);
    }

    pub fn add_value(&mut self, number: u32, value: UnknownValue) {
        self.find_field(&number).add_value(value);
    }

    pub fn iter<'s>(&'s self) -> UnknownFieldsIter<'s> {
        UnknownFieldsIter {
            entries: self.fields.as_ref().map(|m| m.iter()),
        }
    }

    pub fn get(&self, field_number: u32) -> Option<&UnknownValues> {
        match self.fields {
            Some(ref map) => map.get(&field_number),
            None => None,
        }
    }
}

impl Clear for UnknownFields {
    fn clear(&mut self) {
        if let Some(ref mut fields) = self.fields {
            fields.clear();
        }
    }
}

impl<'a> IntoIterator for &'a UnknownFields {
    type Item = (u32, &'a UnknownValues);
    type IntoIter = UnknownFieldsIter<'a>;

    fn into_iter(self) -> UnknownFieldsIter<'a> {
        self.iter()
    }
}

pub struct UnknownFieldsIter<'s> {
    entries: Option<hash_map::Iter<'s, u32, UnknownValues>>,
}

impl<'s> Iterator for UnknownFieldsIter<'s> {
    type Item = (u32, &'s UnknownValues);

    fn next(&mut self) -> Option<(u32, &'s UnknownValues)> {
        match self.entries {
            Some(ref mut entries) => entries.next().map(|(&number, values)| (number, values)),
            None => None,
        }
    }
}

#[cfg(test)]
mod test {
    use super::UnknownFields;
    use std::collections::hash_map::DefaultHasher;
    use std::hash::Hash;
    use std::hash::Hasher;

    #[test]
    fn unknown_fields_hash() {
        let mut unknown_fields_1 = UnknownFields::new();
        let mut unknown_fields_2 = UnknownFields::new();

        // Check field order is not important

        unknown_fields_1.add_fixed32(10, 222);
        unknown_fields_1.add_fixed32(10, 223);
        unknown_fields_1.add_fixed64(14, 224);

        unknown_fields_2.add_fixed32(10, 222);
        unknown_fields_2.add_fixed64(14, 224);
        unknown_fields_2.add_fixed32(10, 223);

        fn hash(unknown_fields: &UnknownFields) -> u64 {
            let mut hasher = DefaultHasher::new();
            Hash::hash(unknown_fields, &mut hasher);
            hasher.finish()
        }

        assert_eq!(hash(&unknown_fields_1), hash(&unknown_fields_2));
    }
}
<commit_msg>Fix UnknownFields hash<commit_after>use clear::Clear;
use std::collections::hash_map;
use std::collections::hash_map::DefaultHasher;
use std::collections::HashMap;
use std::default::Default;
use std::hash::Hash;
use std::hash::Hasher;
use std::slice;
use wire_format;
use zigzag::encode_zig_zag_32;
use zigzag::encode_zig_zag_64;

#[derive(Debug)]
pub enum UnknownValue {
    Fixed32(u32),
    Fixed64(u64),
    Varint(u64),
    LengthDelimited(Vec<u8>),
}

impl UnknownValue {
    pub fn wire_type(&self) -> wire_format::WireType {
        self.get_ref().wire_type()
    }

    pub fn get_ref<'s>(&'s self) -> UnknownValueRef<'s> {
        match *self {
            UnknownValue::Fixed32(fixed32) => UnknownValueRef::Fixed32(fixed32),
            UnknownValue::Fixed64(fixed64) => UnknownValueRef::Fixed64(fixed64),
            UnknownValue::Varint(varint) => UnknownValueRef::Varint(varint),
            UnknownValue::LengthDelimited(ref bytes) => UnknownValueRef::LengthDelimited(&bytes),
        }
    }

    /// Construct unknown value from `sint32` value.
    pub fn sint32(i: i32) -> UnknownValue {
        UnknownValue::Varint(encode_zig_zag_32(i) as u64)
    }

    /// Construct unknown value from `sint64` value.
    pub fn sint64(i: i64) -> UnknownValue {
        UnknownValue::Varint(encode_zig_zag_64(i))
    }
}

pub enum UnknownValueRef<'o> {
    Fixed32(u32),
    Fixed64(u64),
    Varint(u64),
    LengthDelimited(&'o [u8]),
}

impl<'o> UnknownValueRef<'o> {
    pub fn wire_type(&self) -> wire_format::WireType {
        match *self {
            UnknownValueRef::Fixed32(_) => wire_format::WireTypeFixed32,
            UnknownValueRef::Fixed64(_) => wire_format::WireTypeFixed64,
            UnknownValueRef::Varint(_) => wire_format::WireTypeVarint,
            UnknownValueRef::LengthDelimited(_) => wire_format::WireTypeLengthDelimited,
        }
    }
}

#[derive(Clone, PartialEq, Eq, Debug, Default, Hash)]
pub struct UnknownValues {
    pub fixed32: Vec<u32>,
    pub fixed64: Vec<u64>,
    pub varint: Vec<u64>,
    pub length_delimited: Vec<Vec<u8>>,
}

impl UnknownValues {
    pub fn add_value(&mut self, value: UnknownValue) {
        match value {
            UnknownValue::Fixed64(fixed64) => self.fixed64.push(fixed64),
            UnknownValue::Fixed32(fixed32) => self.fixed32.push(fixed32),
            UnknownValue::Varint(varint) => self.varint.push(varint),
            UnknownValue::LengthDelimited(length_delimited) => {
                self.length_delimited.push(length_delimited)
            }
        };
    }

    pub fn iter<'s>(&'s self) -> UnknownValuesIter<'s> {
        UnknownValuesIter {
            fixed32: self.fixed32.iter(),
            fixed64: self.fixed64.iter(),
            varint: self.varint.iter(),
            length_delimited: self.length_delimited.iter(),
        }
    }
}

impl<'a> IntoIterator for &'a UnknownValues {
    type Item = UnknownValueRef<'a>;
    type IntoIter = UnknownValuesIter<'a>;

    fn into_iter(self) -> UnknownValuesIter<'a> {
        self.iter()
    }
}

pub struct UnknownValuesIter<'o> {
    fixed32: slice::Iter<'o, u32>,
    fixed64: slice::Iter<'o, u64>,
    varint: slice::Iter<'o, u64>,
    length_delimited: slice::Iter<'o, Vec<u8>>,
}

impl<'o> Iterator for UnknownValuesIter<'o> {
    type Item = UnknownValueRef<'o>;

    fn next(&mut self) -> Option<UnknownValueRef<'o>> {
        let fixed32 = self.fixed32.next();
        if fixed32.is_some() {
            return Some(UnknownValueRef::Fixed32(*fixed32.unwrap()));
        }
        let fixed64 = self.fixed64.next();
        if fixed64.is_some() {
            return Some(UnknownValueRef::Fixed64(*fixed64.unwrap()));
        }
        let varint = self.varint.next();
        if varint.is_some() {
            return Some(UnknownValueRef::Varint(*varint.unwrap()));
        }
        let length_delimited = self.length_delimited.next();
        if length_delimited.is_some() {
            return Some(UnknownValueRef::LengthDelimited(&length_delimited.unwrap()));
        }
        None
    }
}

#[derive(Clone, PartialEq, Eq, Debug, Default)]
pub struct UnknownFields {
    // option is needed, because HashMap constructor performs allocation,
    // and very expensive
    pub fields: Option<Box<HashMap<u32, UnknownValues>>>,
}

/// Very simple hash implementation of `Hash` for `UnknownFields`.
/// Since map is unordered, we cannot put entry hashes into hasher,
/// instead we summing hashes of entries.
impl Hash for UnknownFields {
    fn hash<H: Hasher>(&self, state: &mut H) {
        if let Some(ref map) = self.fields {
            if !map.is_empty() {
                let mut hash: u64 = 0;
                for (k, v) in &**map {
                    let mut entry_hasher = DefaultHasher::new();
                    Hash::hash(&(k, v), &mut entry_hasher);
                    hash = hash.wrapping_add(entry_hasher.finish());
                }
                Hash::hash(&hash, state);
            }
        }
    }
}

impl UnknownFields {
    pub fn new() -> UnknownFields {
        Default::default()
    }

    fn init_map(&mut self) {
        if self.fields.is_none() {
            self.fields = Some(Default::default());
        }
    }

    fn find_field<'a>(&'a mut self, number: &'a u32) -> &'a mut UnknownValues {
        self.init_map();

        match self.fields.as_mut().unwrap().entry(*number) {
            hash_map::Entry::Occupied(e) => e.into_mut(),
            hash_map::Entry::Vacant(e) => e.insert(Default::default()),
        }
    }

    pub fn add_fixed32(&mut self, number: u32, fixed32: u32) {
        self.find_field(&number).fixed32.push(fixed32);
    }

    pub fn add_fixed64(&mut self, number: u32, fixed64: u64) {
        self.find_field(&number).fixed64.push(fixed64);
    }

    pub fn add_varint(&mut self, number: u32, varint: u64) {
        self.find_field(&number).varint.push(varint);
    }

    pub fn add_length_delimited(&mut self, number: u32, length_delimited: Vec<u8>) {
        self.find_field(&number)
            .length_delimited
            .push(length_delimited);
    }

    pub fn add_value(&mut self, number: u32, value: UnknownValue) {
        self.find_field(&number).add_value(value);
    }

    pub fn iter<'s>(&'s self) -> UnknownFieldsIter<'s> {
        UnknownFieldsIter {
            entries: self.fields.as_ref().map(|m| m.iter()),
        }
    }

    pub fn get(&self, field_number: u32) -> Option<&UnknownValues> {
        match self.fields {
            Some(ref map) => map.get(&field_number),
            None => None,
        }
    }
}

impl Clear for UnknownFields {
    fn clear(&mut self) {
        if let Some(ref mut fields) = self.fields {
            fields.clear();
        }
    }
}

impl<'a> IntoIterator for &'a UnknownFields {
    type Item = (u32, &'a UnknownValues);
    type IntoIter = UnknownFieldsIter<'a>;

    fn into_iter(self) -> UnknownFieldsIter<'a> {
        self.iter()
    }
}

pub struct UnknownFieldsIter<'s> {
    entries: Option<hash_map::Iter<'s, u32, UnknownValues>>,
}

impl<'s> Iterator for UnknownFieldsIter<'s> {
    type Item = (u32, &'s UnknownValues);

    fn next(&mut self) -> Option<(u32, &'s UnknownValues)> {
        match self.entries {
            Some(ref mut entries) => entries.next().map(|(&number, values)| (number, values)),
            None => None,
        }
    }
}

#[cfg(test)]
mod test {
    use super::UnknownFields;
    use std::collections::hash_map::DefaultHasher;
    use std::hash::Hash;
    use std::hash::Hasher;

    #[test]
    fn unknown_fields_hash() {
        let mut unknown_fields_1 = UnknownFields::new();
        let mut unknown_fields_2 = UnknownFields::new();

        // Check field order is not important

        unknown_fields_1.add_fixed32(10, 222);
        unknown_fields_1.add_fixed32(10, 223);
        unknown_fields_1.add_fixed64(14, 224);

        unknown_fields_2.add_fixed32(10, 222);
        unknown_fields_2.add_fixed64(14, 224);
        unknown_fields_2.add_fixed32(10, 223);

        fn hash(unknown_fields: &UnknownFields) -> u64 {
            let mut hasher = DefaultHasher::new();
            Hash::hash(unknown_fields, &mut hasher);
            hasher.finish()
        }

        assert_eq!(hash(&unknown_fields_1), hash(&unknown_fields_2));
    }
}
<|endoftext|>


<commit_before>use super::node::write_node;
use super::node::NodeWriteGuard;
use super::search::MutSearchResult;
use super::BPlusTree;
use super::LevelTree;
use super::NodeCellRef;
use super::*;
use itertools::Itertools;
use std::cmp::min;
use std::fmt::Debug;
use std::sync::atomic::Ordering;

pub const LEVEL_PAGE_DIFF_MULTIPLIER: usize = 8;
pub const LEVEL_TREE_DEPTH: u32 = 2;

pub const LEVEL_M: usize = 8;
pub const LEVEL_0: usize = LEVEL_M * LEVEL_M; // Smaller can be faster but more fragmented
pub const LEVEL_1: usize = LEVEL_0 * LEVEL_PAGE_DIFF_MULTIPLIER;
pub const LEVEL_2: usize = LEVEL_1 * LEVEL_PAGE_DIFF_MULTIPLIER;

pub const NUM_LEVELS: usize = 3;

enum NodeSelection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    WholePage(Vec<NodeWriteGuard<KS, PS>>, NodeWriteGuard<KS, PS>),
    PartialPage(Vec<NodeWriteGuard<KS, PS>>, NodeWriteGuard<KS, PS>),
}

// Both src and dest tree MUST be immutable
fn merge_prune<KS, PS>(
    level: usize,
    node: &NodeCellRef,
    src_tree: &BPlusTree<KS, PS>,
    dest_tree: &dyn LevelTree,
    boundary: &EntryKey,
    prune: bool,
    lsm: usize,
) -> (usize, Option<NodeWriteGuard<KS, PS>>)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_first::<KS, PS>(node);
    match search {
        MutSearchResult::External => {
            if let NodeSelection::WholePage(nodes, mut new_first) =
                select_nodes_in_boundary::<KS, PS>(node, boundary, level, lsm)
            {
                let new_first_ref = new_first.node_ref().clone();
                let head_id = nodes[0].ext_id();
                new_first.extnode_mut(src_tree).id = head_id;
                new_first.left_ref_mut().map(|lr| *lr = Default::default());
                drop(new_first);
                debug_assert!(
                    !nodes.is_empty(),
                    "Cannot find any keys in range in external with boundary {:?}",
                    boundary
                );
                let mut num_keys_merged = 0;
                debug!("Selected {} pages to merge", nodes.len());
                for mut node in nodes {
                    let node_id = node.ext_id();
                    let node_len = node.len();
                    let deleted_keys = node
                        .keys()
                        .iter()
                        .filter(|k| src_tree.deletion.contains(k))
                        .cloned()
                        .collect_vec();
                    let merging_keys = node.extnode_mut_no_persist().keys.as_slice()[..node_len]
                        .iter_mut()
                        .filter(|k| !src_tree.deletion.contains(k))
                        .map(|k| mem::take(k))
                        .collect_vec();
                    let num_merging_keys = merging_keys.len();
                    num_keys_merged += num_merging_keys;
                    debug!(
                        "Collected {} keys, merging to destination tree",
                        num_merging_keys
                    );
                    dest_tree.merge_with_keys(box merging_keys);
                    if prune {
                        debug!("Merge completed, cleanup");
                        clear_node(node, &new_first_ref);
                        if node_id != head_id {
                            external::make_deleted::<KS, PS>(&node_id);
                        }
                        src_tree.len.fetch_sub(num_merging_keys, Ordering::Relaxed);
                    }
                }
                debug!("Merge prune completed in external level");
                return (num_keys_merged, None);
            } else {
                unreachable!();
            }
        }
        MutSearchResult::Internal(sub_node) => {
            enum RightCheck<KS, PS>
            where
                KS: Slice<EntryKey> + Debug + 'static,
                PS: Slice<NodeCellRef> + 'static,
            {
                SinglePtr,
                LevelTerminal(NodeWriteGuard<KS, PS>),
                Normal,
            }
            let check_right = |node: NodeWriteGuard<KS, PS>| {
                if node.right_ref().unwrap().is_default() {
                    debug_assert_eq!(node.right_bound(), &EntryKey::max());
                    if node.len() == 0 {
                        RightCheck::SinglePtr
                    } else {
                        RightCheck::LevelTerminal(node)
                    }
                } else {
                    debug_assert_ne!(node.right_bound(), &EntryKey::max());
                    debug_assert!(node.len() > 0);
                    RightCheck::Normal
                }
            };
            let (num_keys_merged, sub_level_new_root) = merge_prune(
                level + 1,
                &sub_node,
                src_tree,
                dest_tree,
                boundary,
                prune,
                lsm,
            );
            let mut new_root = None;
            if prune {
                let right_node =
                    match select_nodes_in_boundary::<KS, PS>(node, boundary, level, lsm) {
                        NodeSelection::WholePage(nodes, right_node) => {
                            let right_ref = right_node.node_ref().clone();
                            let right_checks = check_right(right_node);
                            clear_nodes(nodes, &right_ref);
                            right_checks
                        }
                        NodeSelection::PartialPage(nodes, mut terminal_node) => {
                            clear_nodes(nodes, terminal_node.node_ref());
                            let search_pos = terminal_node.keys().binary_search(boundary);
                            let pos = match search_pos {
                                Ok(n) => n + 1,
                                Err(n) => n,
                            };
                            let mut new_keys = KS::init();
                            let mut new_ptrs = PS::init();
                            let mut num_keys = 0;
                            for (i, k) in terminal_node.keys()[pos..].iter().enumerate() {
                                new_keys.as_slice()[i] = k.clone();
                                num_keys += 1;
                            }
                            for (i, p) in terminal_node.ptrs()[pos..].iter().enumerate() {
                                new_ptrs.as_slice()[i] = p.clone();
                            }
                            {
                                let innode = terminal_node.innode_mut();
                                innode.keys = new_keys;
                                innode.ptrs = new_ptrs;
                                innode.len = num_keys;
                            }
                            check_right(terminal_node)
                        }
                    };
                new_root = {
                    if let Some(sub_level_new_root) = sub_level_new_root {
                        if cfg!(debug_assertions) {
                            match right_node {
                                RightCheck::SinglePtr => {}
                                _ => panic!("Unexpected node status"),
                            }
                        }
                        Some(sub_level_new_root)
                    } else {
                        match right_node {
                            RightCheck::Normal => None,
                            RightCheck::LevelTerminal(node) => Some(node),
                            RightCheck::SinglePtr => unreachable!(),
                        }
                    }
                };
            }
            return (num_keys_merged, new_root);
        }
    }
}

fn clear_node<KS, PS>(mut node: NodeWriteGuard<KS, PS>, right_ref: &NodeCellRef)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    node.make_empty_node(false);
    *node.right_ref_mut().unwrap() = right_ref.clone();
    node.left_ref_mut().map(|r| *r = right_ref.clone());
}

fn clear_nodes<KS, PS>(nodes: Vec<NodeWriteGuard<KS, PS>>, right_ref: &NodeCellRef)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    for node in nodes {
        clear_node(node, right_ref);
    }
}

fn select_nodes_in_boundary<KS, PS>(
    first_node: &NodeCellRef,
    right_boundary: &EntryKey,
    level: usize,
    lsm: usize,
) -> NodeSelection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_node = write_node::<KS, PS>(first_node);
    if first_node.right_bound() > right_boundary {
        return NodeSelection::PartialPage(vec![], first_node);
    }
    let mut next_node = write_node(first_node.right_ref().unwrap());
    let mut collected = vec![first_node];
    loop {
        let first_key = next_node.first_key();
        let node_right = next_node.right_bound();
        debug_assert!(
            first_key < right_boundary,
            "First key {:?} out of bound {:?}, collected {}. Level {}, LSM {}",
            first_key,
            right_boundary,
            collected.len(),
            level,
            lsm
        );
        if node_right <= right_boundary {
            let next_right = write_node(next_node.right_ref().unwrap());
            collected.push(next_node);
            next_node = next_right;
            if next_node.first_key() >= right_boundary {
                trace!("Collected {} nodes at level {}", collected.len(), level);
                return NodeSelection::WholePage(collected, next_node);
            }
        } else {
            return NodeSelection::PartialPage(collected, next_node);
        }
    }
}

fn select_boundary<KS, PS>(node: &NodeCellRef) -> EntryKey
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let res = read_node(node, |node: &NodeReadHandler<KS, PS>| {
        let node_keys = node.keys();
        let node_len = node_keys.len();
        if node.len() < 1 {
            return Err(node.ptrs()[0].clone());
        }
        // Pick half of the keys in the root
        // Genreally, higher level sub tree in LSM tree will select more keys to merged
        // into next level
        let mid_idx = min(node_len / 2, 8);
        // Return the mid key as the boundary for selection (cut)
        Ok(node_keys[mid_idx].clone())
    });
    match res {
        Ok(key) => key.clone(),
        Err(r) => select_boundary::<KS, PS>(&r),
    }
}

pub async fn level_merge<KS, PS>(
    level: usize,
    src_tree: &BPlusTree<KS, PS>,
    dest_tree: &dyn LevelTree,
    prune: bool,
) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    debug!("Merging LSM tree level {}", level);
    let root = src_tree.get_root();
    let key_boundary = select_boundary::<KS, PS>(&root);
    merge_with_boundary(level, src_tree, dest_tree, &key_boundary, prune).await
}

pub async fn merge_with_boundary<KS, PS>(
    level: usize,
    src_tree: &BPlusTree<KS, PS>,
    dest_tree: &dyn LevelTree,
    key_boundary: &EntryKey,
    prune: bool,
) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    debug!(
        "Level merge level {} with boundary {:?}",
        level, key_boundary
    );
    debug_assert!(verification::tree_has_no_empty_node(&src_tree));
    debug_assert!(verification::is_tree_in_order(&src_tree, level));
    let (num_keys, new_root) = merge_prune(
        0,
        &src_tree.get_root(),
        src_tree,
        dest_tree,
        &key_boundary,
        prune,
        level,
    );
    if let Some(new_root) = new_root {
        let new_root_ref = new_root.node_ref().clone();
        debug!("Level merge update source root {:?}", &new_root_ref);
        *src_tree.root.write() = new_root_ref;
    }
    debug_assert!(verification::tree_has_no_empty_node(&src_tree));
    debug_assert!(verification::is_tree_in_order(&src_tree, level));
    debug!("Merge and pruned level {}, waiting for storage", level);
    storage::wait_until_updated().await;
    debug!("MERGE LEVEL {} COMPLETED", level);
    return num_keys;
}
<commit_msg>adjust tree scale<commit_after>use super::node::write_node;
use super::node::NodeWriteGuard;
use super::search::MutSearchResult;
use super::BPlusTree;
use super::LevelTree;
use super::NodeCellRef;
use super::*;
use itertools::Itertools;
use std::cmp::min;
use std::fmt::Debug;
use std::sync::atomic::Ordering;

pub const LEVEL_PAGE_DIFF_MULTIPLIER: usize = 4;
pub const LEVEL_TREE_DEPTH: u32 = 2;

pub const LEVEL_M: usize = 4;
pub const LEVEL_0: usize = LEVEL_M * LEVEL_M; // Smaller can be faster but more fragmented
pub const LEVEL_1: usize = LEVEL_0 * LEVEL_PAGE_DIFF_MULTIPLIER;
pub const LEVEL_2: usize = LEVEL_1 * LEVEL_PAGE_DIFF_MULTIPLIER;

pub const NUM_LEVELS: usize = 3;

enum NodeSelection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    WholePage(Vec<NodeWriteGuard<KS, PS>>, NodeWriteGuard<KS, PS>),
    PartialPage(Vec<NodeWriteGuard<KS, PS>>, NodeWriteGuard<KS, PS>),
}

// Both src and dest tree MUST be immutable
fn merge_prune<KS, PS>(
    level: usize,
    node: &NodeCellRef,
    src_tree: &BPlusTree<KS, PS>,
    dest_tree: &dyn LevelTree,
    boundary: &EntryKey,
    prune: bool,
    lsm: usize,
) -> (usize, Option<NodeWriteGuard<KS, PS>>)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_first::<KS, PS>(node);
    match search {
        MutSearchResult::External => {
            if let NodeSelection::WholePage(nodes, mut new_first) =
                select_nodes_in_boundary::<KS, PS>(node, boundary, level, lsm)
            {
                let new_first_ref = new_first.node_ref().clone();
                let head_id = nodes[0].ext_id();
                new_first.extnode_mut(src_tree).id = head_id;
                new_first.left_ref_mut().map(|lr| *lr = Default::default());
                drop(new_first);
                debug_assert!(
                    !nodes.is_empty(),
                    "Cannot find any keys in range in external with boundary {:?}",
                    boundary
                );
                let mut num_keys_merged = 0;
                debug!("Selected {} pages to merge", nodes.len());
                for mut node in nodes {
                    let node_id = node.ext_id();
                    let node_len = node.len();
                    let deleted_keys = node
                        .keys()
                        .iter()
                        .filter(|k| src_tree.deletion.contains(k))
                        .cloned()
                        .collect_vec();
                    let merging_keys = node.extnode_mut_no_persist().keys.as_slice()[..node_len]
                        .iter_mut()
                        .filter(|k| !src_tree.deletion.contains(k))
                        .map(|k| mem::take(k))
                        .collect_vec();
                    let num_merging_keys = merging_keys.len();
                    num_keys_merged += num_merging_keys;
                    debug!(
                        "Collected {} keys, merging to destination tree",
                        num_merging_keys
                    );
                    dest_tree.merge_with_keys(box merging_keys);
                    if prune {
                        debug!("Merge completed, cleanup");
                        clear_node(node, &new_first_ref);
                        if node_id != head_id {
                            external::make_deleted::<KS, PS>(&node_id);
                        }
                        src_tree.len.fetch_sub(num_merging_keys, Ordering::Relaxed);
                    }
                }
                debug!("Merge prune completed in external level");
                return (num_keys_merged, None);
            } else {
                unreachable!();
            }
        }
        MutSearchResult::Internal(sub_node) => {
            enum RightCheck<KS, PS>
            where
                KS: Slice<EntryKey> + Debug + 'static,
                PS: Slice<NodeCellRef> + 'static,
            {
                SinglePtr,
                LevelTerminal(NodeWriteGuard<KS, PS>),
                Normal,
            }
            let check_right = |node: NodeWriteGuard<KS, PS>| {
                if node.right_ref().unwrap().is_default() {
                    debug_assert_eq!(node.right_bound(), &EntryKey::max());
                    if node.len() == 0 {
                        RightCheck::SinglePtr
                    } else {
                        RightCheck::LevelTerminal(node)
                    }
                } else {
                    debug_assert_ne!(node.right_bound(), &EntryKey::max());
                    debug_assert!(node.len() > 0);
                    RightCheck::Normal
                }
            };
            let (num_keys_merged, sub_level_new_root) = merge_prune(
                level + 1,
                &sub_node,
                src_tree,
                dest_tree,
                boundary,
                prune,
                lsm,
            );
            let mut new_root = None;
            if prune {
                let right_node =
                    match select_nodes_in_boundary::<KS, PS>(node, boundary, level, lsm) {
                        NodeSelection::WholePage(nodes, right_node) => {
                            let right_ref = right_node.node_ref().clone();
                            let right_checks = check_right(right_node);
                            clear_nodes(nodes, &right_ref);
                            right_checks
                        }
                        NodeSelection::PartialPage(nodes, mut terminal_node) => {
                            clear_nodes(nodes, terminal_node.node_ref());
                            let search_pos = terminal_node.keys().binary_search(boundary);
                            let pos = match search_pos {
                                Ok(n) => n + 1,
                                Err(n) => n,
                            };
                            let mut new_keys = KS::init();
                            let mut new_ptrs = PS::init();
                            let mut num_keys = 0;
                            for (i, k) in terminal_node.keys()[pos..].iter().enumerate() {
                                new_keys.as_slice()[i] = k.clone();
                                num_keys += 1;
                            }
                            for (i, p) in terminal_node.ptrs()[pos..].iter().enumerate() {
                                new_ptrs.as_slice()[i] = p.clone();
                            }
                            {
                                let innode = terminal_node.innode_mut();
                                innode.keys = new_keys;
                                innode.ptrs = new_ptrs;
                                innode.len = num_keys;
                            }
                            check_right(terminal_node)
                        }
                    };
                new_root = {
                    if let Some(sub_level_new_root) = sub_level_new_root {
                        if cfg!(debug_assertions) {
                            match right_node {
                                RightCheck::SinglePtr => {}
                                _ => panic!("Unexpected node status"),
                            }
                        }
                        Some(sub_level_new_root)
                    } else {
                        match right_node {
                            RightCheck::Normal => None,
                            RightCheck::LevelTerminal(node) => Some(node),
                            RightCheck::SinglePtr => unreachable!(),
                        }
                    }
                };
            }
            return (num_keys_merged, new_root);
        }
    }
}

fn clear_node<KS, PS>(mut node: NodeWriteGuard<KS, PS>, right_ref: &NodeCellRef)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    node.make_empty_node(false);
    *node.right_ref_mut().unwrap() = right_ref.clone();
    node.left_ref_mut().map(|r| *r = right_ref.clone());
}

fn clear_nodes<KS, PS>(nodes: Vec<NodeWriteGuard<KS, PS>>, right_ref: &NodeCellRef)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    for node in nodes {
        clear_node(node, right_ref);
    }
}

fn select_nodes_in_boundary<KS, PS>(
    first_node: &NodeCellRef,
    right_boundary: &EntryKey,
    level: usize,
    lsm: usize,
) -> NodeSelection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_node = write_node::<KS, PS>(first_node);
    if first_node.right_bound() > right_boundary {
        return NodeSelection::PartialPage(vec![], first_node);
    }
    let mut next_node = write_node(first_node.right_ref().unwrap());
    let mut collected = vec![first_node];
    loop {
        let first_key = next_node.first_key();
        let node_right = next_node.right_bound();
        debug_assert!(
            first_key < right_boundary,
            "First key {:?} out of bound {:?}, collected {}. Level {}, LSM {}",
            first_key,
            right_boundary,
            collected.len(),
            level,
            lsm
        );
        if node_right <= right_boundary {
            let next_right = write_node(next_node.right_ref().unwrap());
            collected.push(next_node);
            next_node = next_right;
            if next_node.first_key() >= right_boundary {
                trace!("Collected {} nodes at level {}", collected.len(), level);
                return NodeSelection::WholePage(collected, next_node);
            }
        } else {
            return NodeSelection::PartialPage(collected, next_node);
        }
    }
}

fn select_boundary<KS, PS>(node: &NodeCellRef) -> EntryKey
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let res = read_node(node, |node: &NodeReadHandler<KS, PS>| {
        let node_keys = node.keys();
        let node_len = node_keys.len();
        if node.len() < 1 {
            return Err(node.ptrs()[0].clone());
        }
        // Pick half of the keys in the root
        // Genreally, higher level sub tree in LSM tree will select more keys to merged
        // into next level
        let mid_idx = min(node_len / 2, 8);
        // Return the mid key as the boundary for selection (cut)
        Ok(node_keys[mid_idx].clone())
    });
    match res {
        Ok(key) => key.clone(),
        Err(r) => select_boundary::<KS, PS>(&r),
    }
}

pub async fn level_merge<KS, PS>(
    level: usize,
    src_tree: &BPlusTree<KS, PS>,
    dest_tree: &dyn LevelTree,
    prune: bool,
) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    debug!("Merging LSM tree level {}", level);
    let root = src_tree.get_root();
    let key_boundary = select_boundary::<KS, PS>(&root);
    merge_with_boundary(level, src_tree, dest_tree, &key_boundary, prune).await
}

pub async fn merge_with_boundary<KS, PS>(
    level: usize,
    src_tree: &BPlusTree<KS, PS>,
    dest_tree: &dyn LevelTree,
    key_boundary: &EntryKey,
    prune: bool,
) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    debug!(
        "Level merge level {} with boundary {:?}",
        level, key_boundary
    );
    debug_assert!(verification::tree_has_no_empty_node(&src_tree));
    debug_assert!(verification::is_tree_in_order(&src_tree, level));
    let (num_keys, new_root) = merge_prune(
        0,
        &src_tree.get_root(),
        src_tree,
        dest_tree,
        &key_boundary,
        prune,
        level,
    );
    if let Some(new_root) = new_root {
        let new_root_ref = new_root.node_ref().clone();
        debug!("Level merge update source root {:?}", &new_root_ref);
        *src_tree.root.write() = new_root_ref;
    }
    debug_assert!(verification::tree_has_no_empty_node(&src_tree));
    debug_assert!(verification::is_tree_in_order(&src_tree, level));
    debug!("Merge and pruned level {}, waiting for storage", level);
    storage::wait_until_updated().await;
    debug!("MERGE LEVEL {} COMPLETED", level);
    return num_keys;
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

pub use self::RelationDir::*;
use self::TypeVariableValue::*;
use self::UndoEntry::*;
use hir::def_id::{DefId};
use syntax::util::small_vector::SmallVector;
use syntax::ast;
use syntax_pos::Span;
use ty::{self, Ty};

use std::cmp::min;
use std::marker::PhantomData;
use std::mem;
use std::u32;
use rustc_data_structures::snapshot_vec as sv;
use rustc_data_structures::unify as ut;

pub struct TypeVariableTable<'tcx> {
    values: sv::SnapshotVec<Delegate<'tcx>>,
    eq_relations: ut::UnificationTable<ty::TyVid>,
}

/// Reasons to create a type inference variable
pub enum TypeVariableOrigin {
    MiscVariable(Span),
    NormalizeProjectionType(Span),
    TypeInference(Span),
    TypeParameterDefinition(Span, ast::Name),
    TransformedUpvar(Span),
    SubstitutionPlaceholder(Span),
    AutoDeref(Span),
    AdjustmentType(Span),
    DivergingStmt(Span),
    DivergingBlockExpr(Span),
    LatticeVariable(Span),
}

struct TypeVariableData<'tcx> {
    value: TypeVariableValue<'tcx>,
    origin: TypeVariableOrigin,
    diverging: bool
}

enum TypeVariableValue<'tcx> {
    Known(Ty<'tcx>),
    Bounded {
        relations: Vec<Relation>,
        default: Option<Default<'tcx>>
    }
}

// We will use this to store the required information to recapitulate what happened when
// an error occurs.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Default<'tcx> {
    pub ty: Ty<'tcx>,
    /// The span where the default was incurred
    pub origin_span: Span,
    /// The definition that the default originates from
    pub def_id: DefId
}

pub struct Snapshot {
    snapshot: sv::Snapshot,
    eq_snapshot: ut::Snapshot<ty::TyVid>,
}

enum UndoEntry<'tcx> {
    // The type of the var was specified.
    SpecifyVar(ty::TyVid, Vec<Relation>, Option<Default<'tcx>>),
    Relate(ty::TyVid, ty::TyVid),
    RelateRange(ty::TyVid, usize),
}

struct Delegate<'tcx>(PhantomData<&'tcx ()>);

type Relation = (RelationDir, ty::TyVid);

#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub enum RelationDir {
    SubtypeOf, SupertypeOf, EqTo, BiTo
}

impl RelationDir {
    fn opposite(self) -> RelationDir {
        match self {
            SubtypeOf => SupertypeOf,
            SupertypeOf => SubtypeOf,
            EqTo => EqTo,
            BiTo => BiTo,
        }
    }
}

impl<'tcx> TypeVariableTable<'tcx> {
    pub fn new() -> TypeVariableTable<'tcx> {
        TypeVariableTable {
            values: sv::SnapshotVec::new(),
            eq_relations: ut::UnificationTable::new(),
        }
    }

    fn relations<'a>(&'a mut self, a: ty::TyVid) -> &'a mut Vec<Relation> {
        relations(self.values.get_mut(a.index as usize))
    }

    pub fn default(&self, vid: ty::TyVid) -> Option<Default<'tcx>> {
        match &self.values.get(vid.index as usize).value {
            &Known(_) => None,
            &Bounded { ref default, .. } => default.clone()
        }
    }

    pub fn var_diverges<'a>(&'a self, vid: ty::TyVid) -> bool {
        self.values.get(vid.index as usize).diverging
    }

    pub fn var_origin(&self, vid: ty::TyVid) -> &TypeVariableOrigin {
        &self.values.get(vid.index as usize).origin
    }

    /// Records that `a <: b`, `a :> b`, or `a == b`, depending on `dir`.
    ///
    /// Precondition: neither `a` nor `b` are known.
    pub fn relate_vars(&mut self, a: ty::TyVid, dir: RelationDir, b: ty::TyVid) {
        let a = self.root_var(a);
        let b = self.root_var(b);
        if a != b {
            if dir == EqTo {
                // a and b must be equal which we mark in the unification table
                let root = self.eq_relations.union(a, b);
                // In addition to being equal, all relations from the variable which is no longer
                // the root must be added to the root so they are not forgotten as the other
                // variable should no longer be referenced (other than to get the root)
                let other = if a == root { b } else { a };
                let count = {
                    let (relations, root_relations) = if other.index < root.index {
                        let (pre, post) = self.values.split_at_mut(root.index as usize);
                        (relations(&mut pre[other.index as usize]), relations(&mut post[0]))
                    } else {
                        let (pre, post) = self.values.split_at_mut(other.index as usize);
                        (relations(&mut post[0]), relations(&mut pre[root.index as usize]))
                    };
                    root_relations.extend_from_slice(relations);
                    relations.len()
                };
                self.values.record(RelateRange(root, count));
            } else {
                self.relations(a).push((dir, b));
                self.relations(b).push((dir.opposite(), a));
                self.values.record(Relate(a, b));
            }
        }
    }

    /// Instantiates `vid` with the type `ty` and then pushes an entry onto `stack` for each of the
    /// relations of `vid` to other variables. The relations will have the form `(ty, dir, vid1)`
    /// where `vid1` is some other variable id.
    ///
    /// Precondition: `vid` must be a root in the unification table
    pub fn instantiate_and_push(
        &mut self,
        vid: ty::TyVid,
        ty: Ty<'tcx>,
        stack: &mut SmallVector<(Ty<'tcx>, RelationDir, ty::TyVid)>)
    {
        debug_assert!(self.root_var(vid) == vid);
        let old_value = {
            let value_ptr = &mut self.values.get_mut(vid.index as usize).value;
            mem::replace(value_ptr, Known(ty))
        };

        let (relations, default) = match old_value {
            Bounded { relations, default } => (relations, default),
            Known(_) => bug!("Asked to instantiate variable that is \
                              already instantiated")
        };

        for &(dir, vid) in &relations {
            stack.push((ty, dir, vid));
        }

        self.values.record(SpecifyVar(vid, relations, default));
    }

    pub fn new_var(&mut self,
                   diverging: bool,
                   origin: TypeVariableOrigin,
                   default: Option<Default<'tcx>>,) -> ty::TyVid {
        self.eq_relations.new_key(());
        let index = self.values.push(TypeVariableData {
            value: Bounded { relations: vec![], default: default },
            origin: origin,
            diverging: diverging
        });
        let v = ty::TyVid { index: index as u32 };
        debug!("new_var() -> {:?}", v);
        v
    }

    pub fn num_vars(&self) -> usize {
        self.values.len()
    }

    pub fn root_var(&mut self, vid: ty::TyVid) -> ty::TyVid {
        self.eq_relations.find(vid)
    }

    pub fn probe(&mut self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
        let vid = self.root_var(vid);
        self.probe_root(vid)
    }

    /// Retrieves the type of `vid` given that it is currently a root in the unification table
    pub fn probe_root(&mut self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
        debug_assert!(self.root_var(vid) == vid);
        match self.values.get(vid.index as usize).value {
            Bounded { .. } => None,
            Known(t) => Some(t)
        }
    }

    pub fn replace_if_possible(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
        match t.sty {
            ty::TyInfer(ty::TyVar(v)) => {
                match self.probe(v) {
                    None => t,
                    Some(u) => u
                }
            }
            _ => t,
        }
    }

    pub fn snapshot(&mut self) -> Snapshot {
        Snapshot {
            snapshot: self.values.start_snapshot(),
            eq_snapshot: self.eq_relations.snapshot(),
        }
    }

    pub fn rollback_to(&mut self, s: Snapshot) {
        debug!("rollback_to{:?}", {
            for action in self.values.actions_since_snapshot(&s.snapshot) {
                match *action {
                    sv::UndoLog::NewElem(index) => {
                        debug!("inference variable _#{}t popped", index)
                    }
                    _ => { }
                }
            }
        });

        self.values.rollback_to(s.snapshot);
        self.eq_relations.rollback_to(s.eq_snapshot);
    }

    pub fn commit(&mut self, s: Snapshot) {
        self.values.commit(s.snapshot);
        self.eq_relations.commit(s.eq_snapshot);
    }

    pub fn types_escaping_snapshot(&mut self, s: &Snapshot) -> Vec<Ty<'tcx>> {
        /*!
         * Find the set of type variables that existed *before* `s`
         * but which have only been unified since `s` started, and
         * return the types with which they were unified. So if we had
         * a type variable `V0`, then we started the snapshot, then we
         * created a type variable `V1`, unifed `V0` with `T0`, and
         * unified `V1` with `T1`, this function would return `{T0}`.
         */

        let mut new_elem_threshold = u32::MAX;
        let mut escaping_types = Vec::new();
        let actions_since_snapshot = self.values.actions_since_snapshot(&s.snapshot);
        debug!("actions_since_snapshot.len() = {}", actions_since_snapshot.len());
        for action in actions_since_snapshot {
            match *action {
                sv::UndoLog::NewElem(index) => {
                    // if any new variables were created during the
                    // snapshot, remember the lower index (which will
                    // always be the first one we see). Note that this
                    // action must precede those variables being
                    // specified.
                    new_elem_threshold = min(new_elem_threshold, index as u32);
                    debug!("NewElem({}) new_elem_threshold={}", index, new_elem_threshold);
                }

                sv::UndoLog::Other(SpecifyVar(vid, ..)) => {
                    if vid.index < new_elem_threshold {
                        // quick check to see if this variable was
                        // created since the snapshot started or not.
                        let escaping_type = match self.values.get(vid.index as usize).value {
                            Bounded { .. } => bug!(),
                            Known(ty) => ty,
                        };
                        escaping_types.push(escaping_type);
                    }
                    debug!("SpecifyVar({:?}) new_elem_threshold={}", vid, new_elem_threshold);
                }

                _ => { }
            }
        }

        escaping_types
    }

    pub fn unsolved_variables(&mut self) -> Vec<ty::TyVid> {
        (0..self.values.len())
            .filter_map(|i| {
                let vid = ty::TyVid { index: i as u32 };
                if self.probe(vid).is_some() {
                    None
                } else {
                    Some(vid)
                }
            })
            .collect()
    }
}

impl<'tcx> sv::SnapshotVecDelegate for Delegate<'tcx> {
    type Value = TypeVariableData<'tcx>;
    type Undo = UndoEntry<'tcx>;

    fn reverse(values: &mut Vec<TypeVariableData<'tcx>>, action: UndoEntry<'tcx>) {
        match action {
            SpecifyVar(vid, relations, default) => {
                values[vid.index as usize].value = Bounded {
                    relations: relations,
                    default: default
                };
            }

            Relate(a, b) => {
                relations(&mut (*values)[a.index as usize]).pop();
                relations(&mut (*values)[b.index as usize]).pop();
            }

            RelateRange(i, n) => {
                let relations = relations(&mut (*values)[i.index as usize]);
                for _ in 0..n {
                    relations.pop();
                }
            }
        }
    }
}

fn relations<'a>(v: &'a mut TypeVariableData) -> &'a mut Vec<Relation> {
    match v.value {
        Known(_) => bug!("var_sub_var: variable is known"),
        Bounded { ref mut relations, .. } => relations
    }
}
<commit_msg>add some debug logs to type_variable.rs<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

pub use self::RelationDir::*;
use self::TypeVariableValue::*;
use self::UndoEntry::*;
use hir::def_id::{DefId};
use syntax::util::small_vector::SmallVector;
use syntax::ast;
use syntax_pos::Span;
use ty::{self, Ty};

use std::cmp::min;
use std::marker::PhantomData;
use std::mem;
use std::u32;
use rustc_data_structures::snapshot_vec as sv;
use rustc_data_structures::unify as ut;

pub struct TypeVariableTable<'tcx> {
    values: sv::SnapshotVec<Delegate<'tcx>>,
    eq_relations: ut::UnificationTable<ty::TyVid>,
}

/// Reasons to create a type inference variable
#[derive(Debug)]
pub enum TypeVariableOrigin {
    MiscVariable(Span),
    NormalizeProjectionType(Span),
    TypeInference(Span),
    TypeParameterDefinition(Span, ast::Name),
    TransformedUpvar(Span),
    SubstitutionPlaceholder(Span),
    AutoDeref(Span),
    AdjustmentType(Span),
    DivergingStmt(Span),
    DivergingBlockExpr(Span),
    LatticeVariable(Span),
}

struct TypeVariableData<'tcx> {
    value: TypeVariableValue<'tcx>,
    origin: TypeVariableOrigin,
    diverging: bool
}

enum TypeVariableValue<'tcx> {
    Known(Ty<'tcx>),
    Bounded {
        relations: Vec<Relation>,
        default: Option<Default<'tcx>>
    }
}

// We will use this to store the required information to recapitulate what happened when
// an error occurs.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Default<'tcx> {
    pub ty: Ty<'tcx>,
    /// The span where the default was incurred
    pub origin_span: Span,
    /// The definition that the default originates from
    pub def_id: DefId
}

pub struct Snapshot {
    snapshot: sv::Snapshot,
    eq_snapshot: ut::Snapshot<ty::TyVid>,
}

enum UndoEntry<'tcx> {
    // The type of the var was specified.
    SpecifyVar(ty::TyVid, Vec<Relation>, Option<Default<'tcx>>),
    Relate(ty::TyVid, ty::TyVid),
    RelateRange(ty::TyVid, usize),
}

struct Delegate<'tcx>(PhantomData<&'tcx ()>);

type Relation = (RelationDir, ty::TyVid);

#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub enum RelationDir {
    SubtypeOf, SupertypeOf, EqTo, BiTo
}

impl RelationDir {
    fn opposite(self) -> RelationDir {
        match self {
            SubtypeOf => SupertypeOf,
            SupertypeOf => SubtypeOf,
            EqTo => EqTo,
            BiTo => BiTo,
        }
    }
}

impl<'tcx> TypeVariableTable<'tcx> {
    pub fn new() -> TypeVariableTable<'tcx> {
        TypeVariableTable {
            values: sv::SnapshotVec::new(),
            eq_relations: ut::UnificationTable::new(),
        }
    }

    fn relations<'a>(&'a mut self, a: ty::TyVid) -> &'a mut Vec<Relation> {
        relations(self.values.get_mut(a.index as usize))
    }

    pub fn default(&self, vid: ty::TyVid) -> Option<Default<'tcx>> {
        match &self.values.get(vid.index as usize).value {
            &Known(_) => None,
            &Bounded { ref default, .. } => default.clone()
        }
    }

    pub fn var_diverges<'a>(&'a self, vid: ty::TyVid) -> bool {
        self.values.get(vid.index as usize).diverging
    }

    pub fn var_origin(&self, vid: ty::TyVid) -> &TypeVariableOrigin {
        &self.values.get(vid.index as usize).origin
    }

    /// Records that `a <: b`, `a :> b`, or `a == b`, depending on `dir`.
    ///
    /// Precondition: neither `a` nor `b` are known.
    pub fn relate_vars(&mut self, a: ty::TyVid, dir: RelationDir, b: ty::TyVid) {
        let a = self.root_var(a);
        let b = self.root_var(b);
        if a != b {
            if dir == EqTo {
                // a and b must be equal which we mark in the unification table
                let root = self.eq_relations.union(a, b);
                // In addition to being equal, all relations from the variable which is no longer
                // the root must be added to the root so they are not forgotten as the other
                // variable should no longer be referenced (other than to get the root)
                let other = if a == root { b } else { a };
                let count = {
                    let (relations, root_relations) = if other.index < root.index {
                        let (pre, post) = self.values.split_at_mut(root.index as usize);
                        (relations(&mut pre[other.index as usize]), relations(&mut post[0]))
                    } else {
                        let (pre, post) = self.values.split_at_mut(other.index as usize);
                        (relations(&mut post[0]), relations(&mut pre[root.index as usize]))
                    };
                    root_relations.extend_from_slice(relations);
                    relations.len()
                };
                self.values.record(RelateRange(root, count));
            } else {
                self.relations(a).push((dir, b));
                self.relations(b).push((dir.opposite(), a));
                self.values.record(Relate(a, b));
            }
        }
    }

    /// Instantiates `vid` with the type `ty` and then pushes an entry onto `stack` for each of the
    /// relations of `vid` to other variables. The relations will have the form `(ty, dir, vid1)`
    /// where `vid1` is some other variable id.
    ///
    /// Precondition: `vid` must be a root in the unification table
    pub fn instantiate_and_push(
        &mut self,
        vid: ty::TyVid,
        ty: Ty<'tcx>,
        stack: &mut SmallVector<(Ty<'tcx>, RelationDir, ty::TyVid)>)
    {
        debug_assert!(self.root_var(vid) == vid);
        let old_value = {
            let value_ptr = &mut self.values.get_mut(vid.index as usize).value;
            mem::replace(value_ptr, Known(ty))
        };

        let (relations, default) = match old_value {
            Bounded { relations, default } => (relations, default),
            Known(_) => bug!("Asked to instantiate variable that is \
                              already instantiated")
        };

        for &(dir, vid) in &relations {
            stack.push((ty, dir, vid));
        }

        self.values.record(SpecifyVar(vid, relations, default));
    }

    pub fn new_var(&mut self,
                   diverging: bool,
                   origin: TypeVariableOrigin,
                   default: Option<Default<'tcx>>,) -> ty::TyVid {
        debug!("new_var(diverging={:?}, origin={:?})", diverging, origin);
        self.eq_relations.new_key(());
        let index = self.values.push(TypeVariableData {
            value: Bounded { relations: vec![], default: default },
            origin: origin,
            diverging: diverging
        });
        let v = ty::TyVid { index: index as u32 };
        debug!("new_var: diverging={:?} index={:?}", diverging, v);
        v
    }

    pub fn num_vars(&self) -> usize {
        self.values.len()
    }

    pub fn root_var(&mut self, vid: ty::TyVid) -> ty::TyVid {
        self.eq_relations.find(vid)
    }

    pub fn probe(&mut self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
        let vid = self.root_var(vid);
        self.probe_root(vid)
    }

    /// Retrieves the type of `vid` given that it is currently a root in the unification table
    pub fn probe_root(&mut self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
        debug_assert!(self.root_var(vid) == vid);
        match self.values.get(vid.index as usize).value {
            Bounded { .. } => None,
            Known(t) => Some(t)
        }
    }

    pub fn replace_if_possible(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
        match t.sty {
            ty::TyInfer(ty::TyVar(v)) => {
                match self.probe(v) {
                    None => t,
                    Some(u) => u
                }
            }
            _ => t,
        }
    }

    pub fn snapshot(&mut self) -> Snapshot {
        Snapshot {
            snapshot: self.values.start_snapshot(),
            eq_snapshot: self.eq_relations.snapshot(),
        }
    }

    pub fn rollback_to(&mut self, s: Snapshot) {
        debug!("rollback_to{:?}", {
            for action in self.values.actions_since_snapshot(&s.snapshot) {
                match *action {
                    sv::UndoLog::NewElem(index) => {
                        debug!("inference variable _#{}t popped", index)
                    }
                    _ => { }
                }
            }
        });

        self.values.rollback_to(s.snapshot);
        self.eq_relations.rollback_to(s.eq_snapshot);
    }

    pub fn commit(&mut self, s: Snapshot) {
        self.values.commit(s.snapshot);
        self.eq_relations.commit(s.eq_snapshot);
    }

    pub fn types_escaping_snapshot(&mut self, s: &Snapshot) -> Vec<Ty<'tcx>> {
        /*!
         * Find the set of type variables that existed *before* `s`
         * but which have only been unified since `s` started, and
         * return the types with which they were unified. So if we had
         * a type variable `V0`, then we started the snapshot, then we
         * created a type variable `V1`, unifed `V0` with `T0`, and
         * unified `V1` with `T1`, this function would return `{T0}`.
         */

        let mut new_elem_threshold = u32::MAX;
        let mut escaping_types = Vec::new();
        let actions_since_snapshot = self.values.actions_since_snapshot(&s.snapshot);
        debug!("actions_since_snapshot.len() = {}", actions_since_snapshot.len());
        for action in actions_since_snapshot {
            match *action {
                sv::UndoLog::NewElem(index) => {
                    // if any new variables were created during the
                    // snapshot, remember the lower index (which will
                    // always be the first one we see). Note that this
                    // action must precede those variables being
                    // specified.
                    new_elem_threshold = min(new_elem_threshold, index as u32);
                    debug!("NewElem({}) new_elem_threshold={}", index, new_elem_threshold);
                }

                sv::UndoLog::Other(SpecifyVar(vid, ..)) => {
                    if vid.index < new_elem_threshold {
                        // quick check to see if this variable was
                        // created since the snapshot started or not.
                        let escaping_type = match self.values.get(vid.index as usize).value {
                            Bounded { .. } => bug!(),
                            Known(ty) => ty,
                        };
                        escaping_types.push(escaping_type);
                    }
                    debug!("SpecifyVar({:?}) new_elem_threshold={}", vid, new_elem_threshold);
                }

                _ => { }
            }
        }

        escaping_types
    }

    pub fn unsolved_variables(&mut self) -> Vec<ty::TyVid> {
        (0..self.values.len())
            .filter_map(|i| {
                let vid = ty::TyVid { index: i as u32 };
                if self.probe(vid).is_some() {
                    None
                } else {
                    Some(vid)
                }
            })
            .collect()
    }
}

impl<'tcx> sv::SnapshotVecDelegate for Delegate<'tcx> {
    type Value = TypeVariableData<'tcx>;
    type Undo = UndoEntry<'tcx>;

    fn reverse(values: &mut Vec<TypeVariableData<'tcx>>, action: UndoEntry<'tcx>) {
        match action {
            SpecifyVar(vid, relations, default) => {
                values[vid.index as usize].value = Bounded {
                    relations: relations,
                    default: default
                };
            }

            Relate(a, b) => {
                relations(&mut (*values)[a.index as usize]).pop();
                relations(&mut (*values)[b.index as usize]).pop();
            }

            RelateRange(i, n) => {
                let relations = relations(&mut (*values)[i.index as usize]);
                for _ in 0..n {
                    relations.pop();
                }
            }
        }
    }
}

fn relations<'a>(v: &'a mut TypeVariableData) -> &'a mut Vec<Relation> {
    match v.value {
        Known(_) => bug!("var_sub_var: variable is known"),
        Bounded { ref mut relations, .. } => relations
    }
}
<|endoftext|>