<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{Error as EngineError, IncludeCppEngine, RebuildDependencyRecorder};
use proc_macro2::TokenStream;
use quote::ToTokens;
use std::{fmt::Display, io::Read, path::PathBuf};
use std::{panic::UnwindSafe, path::Path, rc::Rc};
use syn::Item;

/// Errors which may occur when parsing a Rust source file to discover
/// and interpret include_cxx macros.
#[derive(Debug)]
pub enum ParseError {
    /// Unable to open the source file
    FileOpen(std::io::Error),
    /// The .rs file couldn't be read.
    FileRead(std::io::Error),
    /// The .rs file couldn't be parsed.
    Syntax(syn::Error),
    /// The include CPP macro could not be expanded into
    /// Rust bindings to C++, because of some problem during the conversion
    /// process. This could be anything from a C++ parsing error to some
    /// C++ feature that autocxx can't yet handle and isn't able to skip
    /// over. It could also cover errors in your syntax of the `include_cpp`
    /// macro or the directives inside.
    AutocxxCodegenError(EngineError),
}

impl Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ParseError::FileOpen(err) => write!(f, "Unable to open file: {}", err)?,
            ParseError::FileRead(err) => write!(f, "Unable to read file: {}", err)?,
            ParseError::Syntax(err) => write!(f, "Syntax error parsing Rust file: {}", err)?,
            ParseError::AutocxxCodegenError(err) => {
                write!(f, "Unable to parse include_cpp! macro: {}", err)?
            }
        }
        Ok(())
    }
}

/// Parse a Rust file, and spot any include_cpp macros within it.
pub fn parse_file<P1: AsRef<Path>>(rs_file: P1) -> Result<ParsedFile, ParseError> {
    let mut source = String::new();
    let mut file = std::fs::File::open(rs_file).map_err(ParseError::FileOpen)?;
    file.read_to_string(&mut source)
        .map_err(ParseError::FileRead)?;
    proc_macro2::fallback::force();
    let source = syn::parse_file(&source).map_err(ParseError::Syntax)?;
    parse_file_contents(source)
}

fn parse_file_contents(source: syn::File) -> Result<ParsedFile, ParseError> {
    let mut results = Vec::new();
    for item in source.items {
        if let Item::Macro(ref mac) = item {
            if mac.mac.path.is_ident("include_cpp") {
                let include_cpp = crate::IncludeCppEngine::new_from_syn(mac.mac.clone())
                    .map_err(ParseError::AutocxxCodegenError)?;
                results.push(Segment::Autocxx(include_cpp));
                continue;
            }
        }
        results.push(Segment::Other(item));
    }
    Ok(ParsedFile(results))
}

/// A Rust file parsed by autocxx. May contain zero or more autocxx 'engines',
/// i.e. the `IncludeCpp` class, corresponding to zero or more include_cpp
/// macros within this file. Also contains `syn::Item` structures for all
/// the rest of the Rust code, such that it can be reconstituted if necessary.
pub struct ParsedFile(Vec<Segment>);

#[allow(clippy::large_enum_variant)]
enum Segment {
    Autocxx(IncludeCppEngine),
    Other(Item),
}

impl ParsedFile {
    /// Get all the autocxxes in this parsed file.
    pub fn get_autocxxes(&self) -> Vec<&IncludeCppEngine> {
        self.0
            .iter()
            .filter_map(|s| match s {
                Segment::Autocxx(includecpp) => Some(includecpp),
                Segment::Other(_) => None,
            })
            .collect()
    }

    pub fn get_autocxxes_mut(&mut self) -> Vec<&mut IncludeCppEngine> {
        self.0
            .iter_mut()
            .filter_map(|s| match s {
                Segment::Autocxx(includecpp) => Some(includecpp),
                Segment::Other(_) => None,
            })
            .collect()
    }

    pub fn resolve_all(
        &mut self,
        autocxx_inc: Vec<PathBuf>,
        extra_clang_args: &[&str],
        dep_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
    ) -> Result<(), ParseError> {
        let inner_dep_recorder: Option<Rc<dyn RebuildDependencyRecorder>> =
            dep_recorder.map(Rc::from);
        for include_cpp in self.get_autocxxes_mut() {
            #[allow(clippy::manual_map)] // because of dyn shenanigans
            let dep_recorder: Option<Box<dyn RebuildDependencyRecorder>> = match &inner_dep_recorder
            {
                None => None,
                Some(inner_dep_recorder) => Some(Box::new(CompositeDepRecorder::new(
                    inner_dep_recorder.clone(),
                ))),
            };
            include_cpp
                .generate(autocxx_inc.clone(), extra_clang_args, dep_recorder)
                .map_err(ParseError::AutocxxCodegenError)?
        }
        Ok(())
    }
}

impl ToTokens for ParsedFile {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        for seg in &self.0 {
            match seg {
                Segment::Other(item) => item.to_tokens(tokens),
                Segment::Autocxx(autocxx) => {
                    let these_tokens = autocxx.generate_rs();
                    tokens.extend(these_tokens);
                }
            }
        }
    }
}

/// Shenanigans required to share the same RebuildDependencyRecorder
/// with all of the include_cpp instances in this one file.
#[derive(Debug, Clone)]
struct CompositeDepRecorder(Rc<dyn RebuildDependencyRecorder>);

impl CompositeDepRecorder {
    fn new(inner: Rc<dyn RebuildDependencyRecorder>) -> Self {
        CompositeDepRecorder(inner)
    }
}

impl UnwindSafe for CompositeDepRecorder {}

impl RebuildDependencyRecorder for CompositeDepRecorder {
    fn record_header_file_dependency(&self, filename: &str) {
        self.0.record_header_file_dependency(filename);
    }
}
<commit_msg>Avoid unnecessary collection.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{Error as EngineError, IncludeCppEngine, RebuildDependencyRecorder};
use proc_macro2::TokenStream;
use quote::ToTokens;
use std::{fmt::Display, io::Read, path::PathBuf};
use std::{panic::UnwindSafe, path::Path, rc::Rc};
use syn::Item;

/// Errors which may occur when parsing a Rust source file to discover
/// and interpret include_cxx macros.
#[derive(Debug)]
pub enum ParseError {
    /// Unable to open the source file
    FileOpen(std::io::Error),
    /// The .rs file couldn't be read.
    FileRead(std::io::Error),
    /// The .rs file couldn't be parsed.
    Syntax(syn::Error),
    /// The include CPP macro could not be expanded into
    /// Rust bindings to C++, because of some problem during the conversion
    /// process. This could be anything from a C++ parsing error to some
    /// C++ feature that autocxx can't yet handle and isn't able to skip
    /// over. It could also cover errors in your syntax of the `include_cpp`
    /// macro or the directives inside.
    AutocxxCodegenError(EngineError),
}

impl Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ParseError::FileOpen(err) => write!(f, "Unable to open file: {}", err)?,
            ParseError::FileRead(err) => write!(f, "Unable to read file: {}", err)?,
            ParseError::Syntax(err) => write!(f, "Syntax error parsing Rust file: {}", err)?,
            ParseError::AutocxxCodegenError(err) => {
                write!(f, "Unable to parse include_cpp! macro: {}", err)?
            }
        }
        Ok(())
    }
}

/// Parse a Rust file, and spot any include_cpp macros within it.
pub fn parse_file<P1: AsRef<Path>>(rs_file: P1) -> Result<ParsedFile, ParseError> {
    let mut source = String::new();
    let mut file = std::fs::File::open(rs_file).map_err(ParseError::FileOpen)?;
    file.read_to_string(&mut source)
        .map_err(ParseError::FileRead)?;
    proc_macro2::fallback::force();
    let source = syn::parse_file(&source).map_err(ParseError::Syntax)?;
    parse_file_contents(source)
}

fn parse_file_contents(source: syn::File) -> Result<ParsedFile, ParseError> {
    let mut results = Vec::new();
    for item in source.items {
        if let Item::Macro(ref mac) = item {
            if mac.mac.path.is_ident("include_cpp") {
                let include_cpp = crate::IncludeCppEngine::new_from_syn(mac.mac.clone())
                    .map_err(ParseError::AutocxxCodegenError)?;
                results.push(Segment::Autocxx(include_cpp));
                continue;
            }
        }
        results.push(Segment::Other(item));
    }
    Ok(ParsedFile(results))
}

/// A Rust file parsed by autocxx. May contain zero or more autocxx 'engines',
/// i.e. the `IncludeCpp` class, corresponding to zero or more include_cpp
/// macros within this file. Also contains `syn::Item` structures for all
/// the rest of the Rust code, such that it can be reconstituted if necessary.
pub struct ParsedFile(Vec<Segment>);

#[allow(clippy::large_enum_variant)]
enum Segment {
    Autocxx(IncludeCppEngine),
    Other(Item),
}

impl ParsedFile {
    /// Get all the autocxxes in this parsed file.
    pub fn get_autocxxes(&self) -> impl Iterator<Item = &IncludeCppEngine> {
        self.0.iter().filter_map(|s| match s {
            Segment::Autocxx(includecpp) => Some(includecpp),
            Segment::Other(_) => None,
        })
    }

    pub fn get_autocxxes_mut(&mut self) -> impl Iterator<Item = &mut IncludeCppEngine> {
        self.0.iter_mut().filter_map(|s| match s {
            Segment::Autocxx(includecpp) => Some(includecpp),
            Segment::Other(_) => None,
        })
    }

    pub fn resolve_all(
        &mut self,
        autocxx_inc: Vec<PathBuf>,
        extra_clang_args: &[&str],
        dep_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
    ) -> Result<(), ParseError> {
        let inner_dep_recorder: Option<Rc<dyn RebuildDependencyRecorder>> =
            dep_recorder.map(Rc::from);
        for include_cpp in self.get_autocxxes_mut() {
            #[allow(clippy::manual_map)] // because of dyn shenanigans
            let dep_recorder: Option<Box<dyn RebuildDependencyRecorder>> = match &inner_dep_recorder
            {
                None => None,
                Some(inner_dep_recorder) => Some(Box::new(CompositeDepRecorder::new(
                    inner_dep_recorder.clone(),
                ))),
            };
            include_cpp
                .generate(autocxx_inc.clone(), extra_clang_args, dep_recorder)
                .map_err(ParseError::AutocxxCodegenError)?
        }
        Ok(())
    }
}

impl ToTokens for ParsedFile {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        for seg in &self.0 {
            match seg {
                Segment::Other(item) => item.to_tokens(tokens),
                Segment::Autocxx(autocxx) => {
                    let these_tokens = autocxx.generate_rs();
                    tokens.extend(these_tokens);
                }
            }
        }
    }
}

/// Shenanigans required to share the same RebuildDependencyRecorder
/// with all of the include_cpp instances in this one file.
#[derive(Debug, Clone)]
struct CompositeDepRecorder(Rc<dyn RebuildDependencyRecorder>);

impl CompositeDepRecorder {
    fn new(inner: Rc<dyn RebuildDependencyRecorder>) -> Self {
        CompositeDepRecorder(inner)
    }
}

impl UnwindSafe for CompositeDepRecorder {}

impl RebuildDependencyRecorder for CompositeDepRecorder {
    fn record_header_file_dependency(&self, filename: &str) {
        self.0.record_header_file_dependency(filename);
    }
}
<|endoftext|>


<commit_before>use Thingie as Base;
use nalgebra::{Pnt2, Vec2, Cross, Dot};
use super::{Intersect, Intersection};

pub struct Circle<S: Base> {
    pub pos: Pnt2<S>,
    pub radius: S
}

///Lines go from a to b
pub struct Line<S: Base> {
    pub a: Pnt2<S>,
    pub b: Pnt2<S>
}

impl<S> Line<S> where S: Base {
    pub fn new(a: Pnt2<S>, b: Pnt2<S>) -> Line<S> {
        Line {
            a: a,
            b: b
        }
    }
}

impl<S> Intersect<Line<S>, S> for Line<S> where S: Base {
    fn intersection(&self, other: &Line<S>) -> Intersection<S> {
        let r = Vec2::new(self.b.x - self.a.x, self.b.y - self.a.y);
        let s = Vec2::new(other.b.x - other.a.x, other.b.y - other.a.y);

        //cross returns a Vec1... wtf?!
        let unum = (other.a - self.a).cross(&r).x;
        let denom = r.cross(&s).x;

        if unum == S::zero() && denom == S::zero() {
            let t0 = (other.a - self.a).dot(&r) / r.dot(&r);
            let t1 = (other.a + s - self.a).dot(&r) / r.dot(&r);
            //do they overlap
            // s and r go in opposite directions so it's t1 -> t0

            if s.dot(&r) < S::zero() && t1 <= S::one() && S::zero() <= t0 {
                //TODO: fix
                return Intersection::Overlap(other.a.clone(), self.b.clone())
            } else if t0 <= S::one() && S::zero() <= t1 {
                //TODO: fix
                return Intersection::Overlap(other.a.clone(), self.b.clone())
            }

            return Intersection::Outside
        }

        if denom == S::zero() {
            //lines are parallel
            return Intersection::Parallel
        }

        let u = unum / denom;
        let t = (other.a - self.a).cross(&s).x / denom;

        if t >= S::zero() && t <= S::one() && u >= S::zero() && u <= S::one() {
            return Intersection::Intersects(self.a + (r * t), None)
        }

        return Intersection::Outside
    }

    fn intersects(&self, other: &Line<S>) -> bool {
        match self.intersection(other) {
            Intersection::Overlap(_, _) => true,
            Intersection::Intersects(_, _) => true,
            _ => false
        }
    }

    fn contains(&self, other: &Line<S>) -> bool {
        false
    }
}

#[test]
fn line_intersects() {
    // no intersect
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(2.0, 8.0))
               .intersects(&Line::new(Pnt2::new(8.0, 0.0), Pnt2::new(0.0, 20.0))), false);
    // intersect
    assert_eq!(Line::new(Pnt2::new(0.0,10.0), Pnt2::new(2.0, 0.0))
               .intersects(&Line::new(Pnt2::new(10.0, 0.0), Pnt2::new(0.0, 5.0))), true);
    // parallel, vertical
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0))
               .intersects(&Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0))), false);
    // parallel, diagonal
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0))), false);
    //collinear, overlap
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0))), true);
    //collinear, no overlap
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0))), false);
}

#[test]
fn line_intersection() {
    let mut l1 = Line::new(Pnt2::new(-5.0, 5.0), Pnt2::new(-1.0, 5.0));
    let mut l2 = Line::new(Pnt2::new(0.0, -5.0), Pnt2::new(0.0, 5.0));

    // no intersection
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 5.0;
    // intersection
    assert_eq!(l1.intersection(&l2), Intersection::Intersects(Pnt2 {x: 0.0, y: 5.0}, None));

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0));
    // parallel, vertical
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0));

    // parallel, diagonal
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0));

    //collinear, overlap
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(l2.a.clone(), l1.b.clone()));
    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0));

    //collinear, no overlap
    assert_eq!(l1.intersection(&l2), Intersection::Outside);
}

<commit_msg>Fixed collinear case<commit_after>use Thingie as Base;
use nalgebra::{Pnt2, Vec2, Cross, Dot};
use super::{Intersect, Intersection};

pub struct Circle<S: Base> {
    pub pos: Pnt2<S>,
    pub radius: S
}

///Lines go from a to b
pub struct Line<S: Base> {
    pub a: Pnt2<S>,
    pub b: Pnt2<S>
}

impl<S> Line<S> where S: Base {
    pub fn new(a: Pnt2<S>, b: Pnt2<S>) -> Line<S> {
        Line {
            a: a,
            b: b
        }
    }
}

impl<S> Intersect<Line<S>, S> for Line<S> where S: Base {
    fn intersection(&self, other: &Line<S>) -> Intersection<S> {
        let r = Vec2::new(self.b.x - self.a.x, self.b.y - self.a.y);
        let s = Vec2::new(other.b.x - other.a.x, other.b.y - other.a.y);

        //cross returns a Vec1... wtf?!
        let unum = (other.a - self.a).cross(&r).x;
        let denom = r.cross(&s).x;

        if unum == S::zero() && denom == S::zero() {
            let t0 = (other.a - self.a).dot(&r) / r.dot(&r);
            let t1 = t0 + s.dot(&r) / r.dot(&r);
            //do they overlap
            // s and r go in opposite directions so s dot r < 0 it's t1 -> t0
            if s.dot(&r) < S::zero() && t1 <= S::one() && S::zero() <= t0 {
                return Intersection::Overlap(self.a + (r * t1), self.a + (r * t0))
            } else if t0 <= S::one() && S::zero() <= t1 {
                return Intersection::Overlap(self.a + (r * t0), self.a + (r * t1))
            }

            return Intersection::Outside
        }

        if denom == S::zero() {
            //lines are parallel
            return Intersection::Parallel
        }

        let u = unum / denom;
        let t = (other.a - self.a).cross(&s).x / denom;

        if t >= S::zero() && t <= S::one() && u >= S::zero() && u <= S::one() {
            return Intersection::Intersects(self.a + (r * t), None)
        }

        return Intersection::Outside
    }

    fn intersects(&self, other: &Line<S>) -> bool {
        match self.intersection(other) {
            Intersection::Overlap(_, _) => true,
            Intersection::Intersects(_, _) => true,
            _ => false
        }
    }

    fn contains(&self, other: &Line<S>) -> bool {
        false
    }
}

#[test]
fn line_intersects() {
    // no intersect
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(2.0, 8.0))
               .intersects(&Line::new(Pnt2::new(8.0, 0.0), Pnt2::new(0.0, 20.0))), false);
    // intersect
    assert_eq!(Line::new(Pnt2::new(0.0,10.0), Pnt2::new(2.0, 0.0))
               .intersects(&Line::new(Pnt2::new(10.0, 0.0), Pnt2::new(0.0, 5.0))), true);
    // parallel, vertical
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0))
               .intersects(&Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0))), false);
    // parallel, diagonal
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0))), false);
    //collinear, overlap
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0))), true);
    //collinear, no overlap
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0))), false);
}

#[test]
fn line_intersection() {
    let mut l1 = Line::new(Pnt2::new(-5.0, 5.0), Pnt2::new(-1.0, 5.0));
    let mut l2 = Line::new(Pnt2::new(0.0, -5.0), Pnt2::new(0.0, 5.0));

    // no intersection
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 5.0;
    // intersection
    assert_eq!(l1.intersection(&l2), Intersection::Intersects(Pnt2 {x: 0.0, y: 5.0}, None));

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0));
    // parallel, vertical
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0));

    // parallel, diagonal
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0));

    //collinear, overlap
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(l2.a.clone(), l2.b.clone()));
    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0));

    //collinear, no overlap
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 15.0;
    l1.b.y = 15.0;
    //it is right but there's stupid floating point mess
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(Pnt2 {x: 7.0, y: 7.0 }, Pnt2 {x: 10.000000000000002, y: 10.000000000000002}));
}

<|endoftext|>


<commit_before>use std::env;
use std::error::Error;
use std::io::Write;

use libs;
use shell;
use tools;
use tools::clog;

fn apply_underlined(result: &mut String) {
    result.push_str(libs::colored::UNDERLINED);
}

fn apply_user(result: &mut String) {
    let username = tools::get_user_name();
    result.push_str(&username);
}

fn apply_black(result: &mut String) {
    result.push_str(libs::colored::BLACK);
}

fn apply_black_b(result: &mut String) {
    result.push_str(libs::colored::BLACK_B);
}

fn apply_black_bg(result: &mut String) {
    result.push_str(libs::colored::BLACK_BG);
}

fn apply_blue(result: &mut String) {
    result.push_str(libs::colored::BLUE);
}

fn apply_blue_b(result: &mut String) {
    result.push_str(libs::colored::BLUE_B);
}

fn apply_blue_bg(result: &mut String) {
    result.push_str(libs::colored::BLUE_BG);
}

fn apply_bold(result: &mut String) {
    result.push_str(libs::colored::BOLD);
}

fn apply_green(result: &mut String) {
    result.push_str(libs::colored::GREEN);
}

fn apply_green_b(result: &mut String) {
    result.push_str(libs::colored::GREEN_B);
}

fn apply_green_bg(result: &mut String) {
    result.push_str(libs::colored::GREEN_BG);
}

fn apply_red(result: &mut String) {
    result.push_str(libs::colored::RED);
}

fn apply_red_b(result: &mut String) {
    result.push_str(libs::colored::RED_B);
}

fn apply_red_bg(result: &mut String) {
    result.push_str(libs::colored::RED_BG);
}

fn apply_white(result: &mut String) {
    result.push_str(libs::colored::WHITE);
}

fn apply_white_b(result: &mut String) {
    result.push_str(libs::colored::WHITE_B);
}

fn apply_white_bg(result: &mut String) {
    result.push_str(libs::colored::WHITE_BG);
}

fn apply_reset(result: &mut String) {
    result.push_str(libs::colored::RESET);
}

fn apply_color_status(sh: &shell::Shell, result: &mut String) {
    if sh.previous_status == 0 {
        result.push_str(libs::colored::GREEN_B);
    } else {
        result.push_str(libs::colored::RED_B);
    }
}

fn apply_cwd(result: &mut String) {
    let _current_dir;
    match env::current_dir() {
        Ok(x) => _current_dir = x,
        Err(e) => {
            println_stderr!("cicada: PROMPT: env current_dir error: {}", e.description());
            return;
        }
    }
    let current_dir;
    match _current_dir.to_str() {
        Some(x) => current_dir = x,
        None => {
            println_stderr!("cicada: PROMPT: to_str error");
            return;
        }
    }
    let _tokens: Vec<&str> = current_dir.split('/').collect();

    let last;
    match _tokens.last() {
        Some(x) => last = x,
        None => {
            log!("cicada: PROMPT: token last error");
            return;
        }
    }

    let home = tools::get_user_home();
    let pwd = if last.is_empty() {
        "/"
    } else if current_dir == home {
        "~"
    } else {
        last
    };
    result.push_str(pwd);
}

fn apply_hostname(result: &mut String) {
    let hostname = tools::get_hostname();
    result.push_str(&hostname);
}

fn apply_newline(result: &mut String) {
    result.push('\n');
}

pub fn apply_pyenv(result: &mut String) {
    if let Ok(x) = env::var("VIRTUAL_ENV") {
        if !x.is_empty() {
            let _tokens: Vec<&str> = x.split('/').collect();
            let env_name;
            match _tokens.last() {
                Some(x) => env_name = x,
                None => {
                    log!("prompt token last error");
                    return;
                }
            }
            result.push('(');
            apply_blue(result);
            result.push_str(env_name);
            apply_reset(result);
            result.push(')');
        }
    }
}

fn apply_others(result: &mut String, token: &str) {
    log!("unknown prompt item: {:?}", token);
    let s = format!("<{}>", token);
    result.push_str(&s);
}

pub fn apply_preset_item(sh: &shell::Shell, result: &mut String, token: &str) {
    match token.to_ascii_lowercase().as_ref() {
        "black" => apply_black(result),
        "black_b" => apply_black_b(result),
        "black_bg" => apply_black_bg(result),
        "blue" => apply_blue(result),
        "blue_b" => apply_blue_b(result),
        "blue_bg" => apply_blue_bg(result),
        "bold" => apply_bold(result),
        "color_status" => apply_color_status(sh, result),
        "cwd" => apply_cwd(result),
        "green" => apply_green(result),
        "green_b" => apply_green_b(result),
        "green_bg" => apply_green_bg(result),
        "hostname" => apply_hostname(result),
        "newline" => apply_newline(result),
        "red" => apply_red(result),
        "red_b" => apply_red_b(result),
        "red_bg" => apply_red_bg(result),
        "reset" => apply_reset(result),
        "underlined" => apply_underlined(result),
        "user" => apply_user(result),
        "white" => apply_white(result),
        "white_b" => apply_white_b(result),
        "white_bg" => apply_white_bg(result),
        _ => apply_others(result, token),
    }
}
<commit_msg>minor updates on pyenv part in prompt<commit_after>use std::env;
use std::error::Error;
use std::io::Write;

use libs;
use shell;
use tools;
use tools::clog;

fn apply_underlined(prompt: &mut String) {
    prompt.push_str(libs::colored::UNDERLINED);
}

fn apply_user(prompt: &mut String) {
    let username = tools::get_user_name();
    prompt.push_str(&username);
}

fn apply_black(prompt: &mut String) {
    prompt.push_str(libs::colored::BLACK);
}

fn apply_black_b(prompt: &mut String) {
    prompt.push_str(libs::colored::BLACK_B);
}

fn apply_black_bg(prompt: &mut String) {
    prompt.push_str(libs::colored::BLACK_BG);
}

fn apply_blue(prompt: &mut String) {
    prompt.push_str(libs::colored::BLUE);
}

fn apply_blue_b(prompt: &mut String) {
    prompt.push_str(libs::colored::BLUE_B);
}

fn apply_blue_bg(prompt: &mut String) {
    prompt.push_str(libs::colored::BLUE_BG);
}

fn apply_bold(prompt: &mut String) {
    prompt.push_str(libs::colored::BOLD);
}

fn apply_green(prompt: &mut String) {
    prompt.push_str(libs::colored::GREEN);
}

fn apply_green_b(prompt: &mut String) {
    prompt.push_str(libs::colored::GREEN_B);
}

fn apply_green_bg(prompt: &mut String) {
    prompt.push_str(libs::colored::GREEN_BG);
}

fn apply_red(prompt: &mut String) {
    prompt.push_str(libs::colored::RED);
}

fn apply_red_b(prompt: &mut String) {
    prompt.push_str(libs::colored::RED_B);
}

fn apply_red_bg(prompt: &mut String) {
    prompt.push_str(libs::colored::RED_BG);
}

fn apply_white(prompt: &mut String) {
    prompt.push_str(libs::colored::WHITE);
}

fn apply_white_b(prompt: &mut String) {
    prompt.push_str(libs::colored::WHITE_B);
}

fn apply_white_bg(prompt: &mut String) {
    prompt.push_str(libs::colored::WHITE_BG);
}

fn apply_reset(prompt: &mut String) {
    prompt.push_str(libs::colored::RESET);
}

fn apply_color_status(sh: &shell::Shell, prompt: &mut String) {
    if sh.previous_status == 0 {
        prompt.push_str(libs::colored::GREEN_B);
    } else {
        prompt.push_str(libs::colored::RED_B);
    }
}

fn apply_cwd(prompt: &mut String) {
    let _current_dir;
    match env::current_dir() {
        Ok(x) => _current_dir = x,
        Err(e) => {
            println_stderr!("cicada: PROMPT: env current_dir error: {}", e.description());
            return;
        }
    }
    let current_dir;
    match _current_dir.to_str() {
        Some(x) => current_dir = x,
        None => {
            println_stderr!("cicada: PROMPT: to_str error");
            return;
        }
    }
    let _tokens: Vec<&str> = current_dir.split('/').collect();

    let last;
    match _tokens.last() {
        Some(x) => last = x,
        None => {
            log!("cicada: PROMPT: token last error");
            return;
        }
    }

    let home = tools::get_user_home();
    let pwd = if last.is_empty() {
        "/"
    } else if current_dir == home {
        "~"
    } else {
        last
    };
    prompt.push_str(pwd);
}

fn apply_hostname(prompt: &mut String) {
    let hostname = tools::get_hostname();
    prompt.push_str(&hostname);
}

fn apply_newline(prompt: &mut String) {
    prompt.push('\n');
}

pub fn apply_pyenv(prompt: &mut String) {
    if let Ok(x) = env::var("VIRTUAL_ENV") {
        if !x.is_empty() {
            let _tokens: Vec<&str> = x.split('/').collect();
            let env_name;
            match _tokens.last() {
                Some(x) => env_name = x,
                None => {
                    log!("prompt token last error");
                    return;
                }
            }
            apply_blue_b(prompt);
            prompt.push('(');
            prompt.push_str(env_name);
            prompt.push(')');
            apply_reset(prompt);
        }
    }
}

fn apply_others(prompt: &mut String, token: &str) {
    log!("unknown prompt item: {:?}", token);
    let s = format!("<{}>", token);
    prompt.push_str(&s);
}

pub fn apply_preset_item(sh: &shell::Shell, prompt: &mut String, token: &str) {
    match token.to_ascii_lowercase().as_ref() {
        "black" => apply_black(prompt),
        "black_b" => apply_black_b(prompt),
        "black_bg" => apply_black_bg(prompt),
        "blue" => apply_blue(prompt),
        "blue_b" => apply_blue_b(prompt),
        "blue_bg" => apply_blue_bg(prompt),
        "bold" => apply_bold(prompt),
        "color_status" => apply_color_status(sh, prompt),
        "cwd" => apply_cwd(prompt),
        "green" => apply_green(prompt),
        "green_b" => apply_green_b(prompt),
        "green_bg" => apply_green_bg(prompt),
        "hostname" => apply_hostname(prompt),
        "newline" => apply_newline(prompt),
        "red" => apply_red(prompt),
        "red_b" => apply_red_b(prompt),
        "red_bg" => apply_red_bg(prompt),
        "reset" => apply_reset(prompt),
        "underlined" => apply_underlined(prompt),
        "user" => apply_user(prompt),
        "white" => apply_white(prompt),
        "white_b" => apply_white_b(prompt),
        "white_bg" => apply_white_bg(prompt),
        _ => apply_others(prompt, token),
    }
}
<|endoftext|>


<commit_before>use crate::common::CompositeFile;
use crate::common::HasLen;
use crate::core::InvertedIndexReader;
use crate::core::Segment;
use crate::core::SegmentComponent;
use crate::core::SegmentId;
use crate::directory::ReadOnlySource;
use crate::fastfield::DeleteBitSet;
use crate::fastfield::FacetReader;
use crate::fastfield::FastFieldReaders;
use crate::fieldnorm::FieldNormReader;
use crate::schema::Field;
use crate::schema::FieldType;
use crate::schema::Schema;
use crate::space_usage::SegmentSpaceUsage;
use crate::store::StoreReader;
use crate::termdict::TermDictionary;
use crate::DocId;
use fail::fail_point;
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use std::sync::RwLock;

/// Entry point to access all of the datastructures of the `Segment`
///
/// - term dictionary
/// - postings
/// - store
/// - fast field readers
/// - field norm reader
///
/// The segment reader has a very low memory footprint,
/// as close to all of the memory data is mmapped.
///
///
/// TODO fix not decoding docfreq
#[derive(Clone)]
pub struct SegmentReader {
    inv_idx_reader_cache: Arc<RwLock<HashMap<Field, Arc<InvertedIndexReader>>>>,

    segment_id: SegmentId,
    max_doc: DocId,
    num_docs: DocId,

    termdict_composite: CompositeFile,
    postings_composite: CompositeFile,
    positions_composite: CompositeFile,
    positions_idx_composite: CompositeFile,
    fast_fields_readers: Arc<FastFieldReaders>,
    fieldnorms_composite: CompositeFile,

    store_source: ReadOnlySource,
    delete_bitset_opt: Option<DeleteBitSet>,
    schema: Schema,
}

impl SegmentReader {
    /// Returns the highest document id ever attributed in
    /// this segment + 1.
    /// Today, `tantivy` does not handle deletes, so it happens
    /// to also be the number of documents in the index.
    pub fn max_doc(&self) -> DocId {
        self.max_doc
    }

    /// Returns the number of documents.
    /// Deleted documents are not counted.
    ///
    /// Today, `tantivy` does not handle deletes so max doc and
    /// num_docs are the same.
    pub fn num_docs(&self) -> DocId {
        self.num_docs
    }

    /// Returns the schema of the index this segment belongs to.
    pub fn schema(&self) -> &Schema {
        &self.schema
    }

    /// Return the number of documents that have been
    /// deleted in the segment.
    pub fn num_deleted_docs(&self) -> DocId {
        self.delete_bitset()
            .map(|delete_set| delete_set.len() as DocId)
            .unwrap_or(0u32)
    }

    /// Returns true iff some of the documents of the segment have been deleted.
    pub fn has_deletes(&self) -> bool {
        self.delete_bitset().is_some()
    }

    /// Accessor to a segment's fast field reader given a field.
    ///
    /// Returns the u64 fast value reader if the field
    /// is a u64 field indexed as "fast".
    ///
    /// Return a FastFieldNotAvailableError if the field is not
    /// declared as a fast field in the schema.
    ///
    /// # Panics
    /// May panic if the index is corrupted.
    pub fn fast_fields(&self) -> &FastFieldReaders {
        &self.fast_fields_readers
    }

    /// Accessor to the `FacetReader` associated to a given `Field`.
    pub fn facet_reader(&self, field: Field) -> Option<FacetReader> {
        let field_entry = self.schema.get_field_entry(field);
        if field_entry.field_type() != &FieldType::HierarchicalFacet {
            return None;
        }
        let term_ords_reader = self.fast_fields().u64s(field)?;
        let termdict_source = self.termdict_composite.open_read(field)?;
        let termdict = TermDictionary::from_source(&termdict_source);
        let facet_reader = FacetReader::new(term_ords_reader, termdict);
        Some(facet_reader)
    }

    /// Accessor to the segment's `Field norms`'s reader.
    ///
    /// Field norms are the length (in tokens) of the fields.
    /// It is used in the computation of the [TfIdf](https://fulmicoton.gitbooks.io/tantivy-doc/content/tfidf.html).
    ///
    /// They are simply stored as a fast field, serialized in
    /// the `.fieldnorm` file of the segment.
    pub fn get_fieldnorms_reader(&self, field: Field) -> FieldNormReader {
        if let Some(fieldnorm_source) = self.fieldnorms_composite.open_read(field) {
            FieldNormReader::open(fieldnorm_source)
        } else {
            let field_name = self.schema.get_field_name(field);
            let err_msg = format!(
                "Field norm not found for field {:?}. Was it market as indexed during indexing.",
                field_name
            );
            panic!(err_msg);
        }
    }

    /// Accessor to the segment's `StoreReader`.
    pub fn get_store_reader(&self) -> StoreReader {
        StoreReader::from_source(self.store_source.clone())
    }

    /// Open a new segment for reading.
    pub fn open(segment: &Segment) -> crate::Result<SegmentReader> {
        let termdict_source = segment.open_read(SegmentComponent::TERMS)?;
        let termdict_composite = CompositeFile::open(&termdict_source)?;

        let store_source = segment.open_read(SegmentComponent::STORE)?;

        fail_point!("SegmentReader::open#middle");

        let postings_source = segment.open_read(SegmentComponent::POSTINGS)?;
        let postings_composite = CompositeFile::open(&postings_source)?;

        let positions_composite = {
            if let Ok(source) = segment.open_read(SegmentComponent::POSITIONS) {
                CompositeFile::open(&source)?
            } else {
                CompositeFile::empty()
            }
        };

        let positions_idx_composite = {
            if let Ok(source) = segment.open_read(SegmentComponent::POSITIONSSKIP) {
                CompositeFile::open(&source)?
            } else {
                CompositeFile::empty()
            }
        };

        let schema = segment.schema();

        let fast_fields_data = segment.open_read(SegmentComponent::FASTFIELDS)?;
        let fast_fields_composite = CompositeFile::open(&fast_fields_data)?;
        let fast_field_readers =
            Arc::new(FastFieldReaders::load_all(&schema, &fast_fields_composite)?);

        let fieldnorms_data = segment.open_read(SegmentComponent::FIELDNORMS)?;
        let fieldnorms_composite = CompositeFile::open(&fieldnorms_data)?;

        let delete_bitset_opt = if segment.meta().has_deletes() {
            let delete_data = segment.open_read(SegmentComponent::DELETE)?;
            Some(DeleteBitSet::open(delete_data))
        } else {
            None
        };

        Ok(SegmentReader {
            inv_idx_reader_cache: Arc::new(RwLock::new(HashMap::new())),
            max_doc: segment.meta().max_doc(),
            num_docs: segment.meta().num_docs(),
            termdict_composite,
            postings_composite,
            fast_fields_readers: fast_field_readers,
            fieldnorms_composite,
            segment_id: segment.id(),
            store_source,
            delete_bitset_opt,
            positions_composite,
            positions_idx_composite,
            schema,
        })
    }

    /// Returns a field reader associated to the field given in argument.
    /// If the field was not present in the index during indexing time,
    /// the InvertedIndexReader is empty.
    ///
    /// The field reader is in charge of iterating through the
    /// term dictionary associated to a specific field,
    /// and opening the posting list associated to any term.
    pub fn inverted_index(&self, field: Field) -> Arc<InvertedIndexReader> {
        if let Some(inv_idx_reader) = self
            .inv_idx_reader_cache
            .read()
            .expect("Lock poisoned. This should never happen")
            .get(&field)
        {
            return Arc::clone(inv_idx_reader);
        }
        let field_entry = self.schema.get_field_entry(field);
        let field_type = field_entry.field_type();
        let record_option_opt = field_type.get_index_record_option();

        if record_option_opt.is_none() {
            panic!("Field {:?} does not seem indexed.", field_entry.name());
        }

        let record_option = record_option_opt.unwrap();

        let postings_source_opt = self.postings_composite.open_read(field);

        if postings_source_opt.is_none() {
            // no documents in the segment contained this field.
            // As a result, no data is associated to the inverted index.
            //
            // Returns an empty inverted index.
            return Arc::new(InvertedIndexReader::empty(field_type));
        }

        let postings_source = postings_source_opt.unwrap();

        let termdict_source = self.termdict_composite.open_read(field).expect(
            "Failed to open field term dictionary in composite file. Is the field indexed?",
        );

        let positions_source = self
            .positions_composite
            .open_read(field)
            .expect("Index corrupted. Failed to open field positions in composite file.");

        let positions_idx_source = self
            .positions_idx_composite
            .open_read(field)
            .expect("Index corrupted. Failed to open field positions in composite file.");

        let inv_idx_reader = Arc::new(InvertedIndexReader::new(
            TermDictionary::from_source(&termdict_source),
            postings_source,
            positions_source,
            positions_idx_source,
            record_option,
        ));

        // by releasing the lock in between, we may end up opening the inverting index
        // twice, but this is fine.
        self.inv_idx_reader_cache
            .write()
            .expect("Field reader cache lock poisoned. This should never happen.")
            .insert(field, Arc::clone(&inv_idx_reader));

        inv_idx_reader
    }

    /// Returns the segment id
    pub fn segment_id(&self) -> SegmentId {
        self.segment_id
    }

    /// Returns the bitset representing
    /// the documents that have been deleted.
    pub fn delete_bitset(&self) -> Option<&DeleteBitSet> {
        self.delete_bitset_opt.as_ref()
    }

    /// Returns true iff the `doc` is marked
    /// as deleted.
    pub fn is_deleted(&self, doc: DocId) -> bool {
        self.delete_bitset()
            .map(|delete_set| delete_set.is_deleted(doc))
            .unwrap_or(false)
    }

    /// Returns an iterator that will iterate over the alive document ids
    pub fn doc_ids_alive(&self) -> SegmentReaderAliveDocsIterator<'_> {
        SegmentReaderAliveDocsIterator::new(&self)
    }

    /// Summarize total space usage of this segment.
    pub fn space_usage(&self) -> SegmentSpaceUsage {
        SegmentSpaceUsage::new(
            self.num_docs(),
            self.termdict_composite.space_usage(),
            self.postings_composite.space_usage(),
            self.positions_composite.space_usage(),
            self.positions_idx_composite.space_usage(),
            self.fast_fields_readers.space_usage(),
            self.fieldnorms_composite.space_usage(),
            self.get_store_reader().space_usage(),
            self.delete_bitset_opt
                .as_ref()
                .map(DeleteBitSet::space_usage)
                .unwrap_or(0),
        )
    }
}

impl fmt::Debug for SegmentReader {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SegmentReader({:?})", self.segment_id)
    }
}

/// Implements the iterator trait to allow easy iteration
/// over non-deleted ("alive") DocIds in a SegmentReader
pub struct SegmentReaderAliveDocsIterator<'a> {
    reader: &'a SegmentReader,
    max_doc: DocId,
    current: DocId,
}

impl<'a> SegmentReaderAliveDocsIterator<'a> {
    pub fn new(reader: &'a SegmentReader) -> SegmentReaderAliveDocsIterator<'a> {
        SegmentReaderAliveDocsIterator {
            reader,
            max_doc: reader.max_doc(),
            current: 0,
        }
    }
}

impl<'a> Iterator for SegmentReaderAliveDocsIterator<'a> {
    type Item = DocId;

    fn next(&mut self) -> Option<Self::Item> {
        // TODO: Use TinySet (like in BitSetDocSet) to speed this process up
        if self.current >= self.max_doc {
            return None;
        }

        // find the next alive doc id
        while self.reader.is_deleted(self.current) {
            self.current += 1;

            if self.current >= self.max_doc {
                return None;
            }
        }

        // capture the current alive DocId
        let result = Some(self.current);

        // move down the chain
        self.current += 1;

        result
    }
}

#[cfg(test)]
mod test {
    use crate::core::Index;
    use crate::schema::{Schema, Term, STORED, TEXT};
    use crate::DocId;

    #[test]
    fn test_alive_docs_iterator() {
        let mut schema_builder = Schema::builder();
        schema_builder.add_text_field("name", TEXT | STORED);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema.clone());
        let name = schema.get_field("name").unwrap();

        {
            let mut index_writer = index.writer_with_num_threads(1, 3_000_000).unwrap();
            index_writer.add_document(doc!(name => "tantivy"));
            index_writer.add_document(doc!(name => "horse"));
            index_writer.add_document(doc!(name => "jockey"));
            index_writer.add_document(doc!(name => "cap"));

            // we should now have one segment with two docs
            index_writer.commit().unwrap();
        }

        {
            let mut index_writer2 = index.writer(50_000_000).unwrap();
            index_writer2.delete_term(Term::from_field_text(name, "horse"));
            index_writer2.delete_term(Term::from_field_text(name, "cap"));

            // ok, now we should have a deleted doc
            index_writer2.commit().unwrap();
        }
        let searcher = index.reader().unwrap().searcher();
        let docs: Vec<DocId> = searcher.segment_reader(0).doc_ids_alive().collect();
        assert_eq!(vec![0u32, 2u32], docs);
    }
}
<commit_msg>Alive doc iterator. (#837)<commit_after>use crate::common::CompositeFile;
use crate::common::HasLen;
use crate::core::InvertedIndexReader;
use crate::core::Segment;
use crate::core::SegmentComponent;
use crate::core::SegmentId;
use crate::directory::ReadOnlySource;
use crate::fastfield::DeleteBitSet;
use crate::fastfield::FacetReader;
use crate::fastfield::FastFieldReaders;
use crate::fieldnorm::FieldNormReader;
use crate::schema::Field;
use crate::schema::FieldType;
use crate::schema::Schema;
use crate::space_usage::SegmentSpaceUsage;
use crate::store::StoreReader;
use crate::termdict::TermDictionary;
use crate::DocId;
use fail::fail_point;
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use std::sync::RwLock;

/// Entry point to access all of the datastructures of the `Segment`
///
/// - term dictionary
/// - postings
/// - store
/// - fast field readers
/// - field norm reader
///
/// The segment reader has a very low memory footprint,
/// as close to all of the memory data is mmapped.
///
///
/// TODO fix not decoding docfreq
#[derive(Clone)]
pub struct SegmentReader {
    inv_idx_reader_cache: Arc<RwLock<HashMap<Field, Arc<InvertedIndexReader>>>>,

    segment_id: SegmentId,
    max_doc: DocId,
    num_docs: DocId,

    termdict_composite: CompositeFile,
    postings_composite: CompositeFile,
    positions_composite: CompositeFile,
    positions_idx_composite: CompositeFile,
    fast_fields_readers: Arc<FastFieldReaders>,
    fieldnorms_composite: CompositeFile,

    store_source: ReadOnlySource,
    delete_bitset_opt: Option<DeleteBitSet>,
    schema: Schema,
}

impl SegmentReader {
    /// Returns the highest document id ever attributed in
    /// this segment + 1.
    /// Today, `tantivy` does not handle deletes, so it happens
    /// to also be the number of documents in the index.
    pub fn max_doc(&self) -> DocId {
        self.max_doc
    }

    /// Returns the number of documents.
    /// Deleted documents are not counted.
    ///
    /// Today, `tantivy` does not handle deletes so max doc and
    /// num_docs are the same.
    pub fn num_docs(&self) -> DocId {
        self.num_docs
    }

    /// Returns the schema of the index this segment belongs to.
    pub fn schema(&self) -> &Schema {
        &self.schema
    }

    /// Return the number of documents that have been
    /// deleted in the segment.
    pub fn num_deleted_docs(&self) -> DocId {
        self.delete_bitset()
            .map(|delete_set| delete_set.len() as DocId)
            .unwrap_or(0u32)
    }

    /// Returns true iff some of the documents of the segment have been deleted.
    pub fn has_deletes(&self) -> bool {
        self.delete_bitset().is_some()
    }

    /// Accessor to a segment's fast field reader given a field.
    ///
    /// Returns the u64 fast value reader if the field
    /// is a u64 field indexed as "fast".
    ///
    /// Return a FastFieldNotAvailableError if the field is not
    /// declared as a fast field in the schema.
    ///
    /// # Panics
    /// May panic if the index is corrupted.
    pub fn fast_fields(&self) -> &FastFieldReaders {
        &self.fast_fields_readers
    }

    /// Accessor to the `FacetReader` associated to a given `Field`.
    pub fn facet_reader(&self, field: Field) -> Option<FacetReader> {
        let field_entry = self.schema.get_field_entry(field);
        if field_entry.field_type() != &FieldType::HierarchicalFacet {
            return None;
        }
        let term_ords_reader = self.fast_fields().u64s(field)?;
        let termdict_source = self.termdict_composite.open_read(field)?;
        let termdict = TermDictionary::from_source(&termdict_source);
        let facet_reader = FacetReader::new(term_ords_reader, termdict);
        Some(facet_reader)
    }

    /// Accessor to the segment's `Field norms`'s reader.
    ///
    /// Field norms are the length (in tokens) of the fields.
    /// It is used in the computation of the [TfIdf](https://fulmicoton.gitbooks.io/tantivy-doc/content/tfidf.html).
    ///
    /// They are simply stored as a fast field, serialized in
    /// the `.fieldnorm` file of the segment.
    pub fn get_fieldnorms_reader(&self, field: Field) -> FieldNormReader {
        if let Some(fieldnorm_source) = self.fieldnorms_composite.open_read(field) {
            FieldNormReader::open(fieldnorm_source)
        } else {
            let field_name = self.schema.get_field_name(field);
            let err_msg = format!(
                "Field norm not found for field {:?}. Was it market as indexed during indexing.",
                field_name
            );
            panic!(err_msg);
        }
    }

    /// Accessor to the segment's `StoreReader`.
    pub fn get_store_reader(&self) -> StoreReader {
        StoreReader::from_source(self.store_source.clone())
    }

    /// Open a new segment for reading.
    pub fn open(segment: &Segment) -> crate::Result<SegmentReader> {
        let termdict_source = segment.open_read(SegmentComponent::TERMS)?;
        let termdict_composite = CompositeFile::open(&termdict_source)?;

        let store_source = segment.open_read(SegmentComponent::STORE)?;

        fail_point!("SegmentReader::open#middle");

        let postings_source = segment.open_read(SegmentComponent::POSTINGS)?;
        let postings_composite = CompositeFile::open(&postings_source)?;

        let positions_composite = {
            if let Ok(source) = segment.open_read(SegmentComponent::POSITIONS) {
                CompositeFile::open(&source)?
            } else {
                CompositeFile::empty()
            }
        };

        let positions_idx_composite = {
            if let Ok(source) = segment.open_read(SegmentComponent::POSITIONSSKIP) {
                CompositeFile::open(&source)?
            } else {
                CompositeFile::empty()
            }
        };

        let schema = segment.schema();

        let fast_fields_data = segment.open_read(SegmentComponent::FASTFIELDS)?;
        let fast_fields_composite = CompositeFile::open(&fast_fields_data)?;
        let fast_field_readers =
            Arc::new(FastFieldReaders::load_all(&schema, &fast_fields_composite)?);

        let fieldnorms_data = segment.open_read(SegmentComponent::FIELDNORMS)?;
        let fieldnorms_composite = CompositeFile::open(&fieldnorms_data)?;

        let delete_bitset_opt = if segment.meta().has_deletes() {
            let delete_data = segment.open_read(SegmentComponent::DELETE)?;
            Some(DeleteBitSet::open(delete_data))
        } else {
            None
        };

        Ok(SegmentReader {
            inv_idx_reader_cache: Arc::new(RwLock::new(HashMap::new())),
            max_doc: segment.meta().max_doc(),
            num_docs: segment.meta().num_docs(),
            termdict_composite,
            postings_composite,
            fast_fields_readers: fast_field_readers,
            fieldnorms_composite,
            segment_id: segment.id(),
            store_source,
            delete_bitset_opt,
            positions_composite,
            positions_idx_composite,
            schema,
        })
    }

    /// Returns a field reader associated to the field given in argument.
    /// If the field was not present in the index during indexing time,
    /// the InvertedIndexReader is empty.
    ///
    /// The field reader is in charge of iterating through the
    /// term dictionary associated to a specific field,
    /// and opening the posting list associated to any term.
    pub fn inverted_index(&self, field: Field) -> Arc<InvertedIndexReader> {
        if let Some(inv_idx_reader) = self
            .inv_idx_reader_cache
            .read()
            .expect("Lock poisoned. This should never happen")
            .get(&field)
        {
            return Arc::clone(inv_idx_reader);
        }
        let field_entry = self.schema.get_field_entry(field);
        let field_type = field_entry.field_type();
        let record_option_opt = field_type.get_index_record_option();

        if record_option_opt.is_none() {
            panic!("Field {:?} does not seem indexed.", field_entry.name());
        }

        let record_option = record_option_opt.unwrap();

        let postings_source_opt = self.postings_composite.open_read(field);

        if postings_source_opt.is_none() {
            // no documents in the segment contained this field.
            // As a result, no data is associated to the inverted index.
            //
            // Returns an empty inverted index.
            return Arc::new(InvertedIndexReader::empty(field_type));
        }

        let postings_source = postings_source_opt.unwrap();

        let termdict_source = self.termdict_composite.open_read(field).expect(
            "Failed to open field term dictionary in composite file. Is the field indexed?",
        );

        let positions_source = self
            .positions_composite
            .open_read(field)
            .expect("Index corrupted. Failed to open field positions in composite file.");

        let positions_idx_source = self
            .positions_idx_composite
            .open_read(field)
            .expect("Index corrupted. Failed to open field positions in composite file.");

        let inv_idx_reader = Arc::new(InvertedIndexReader::new(
            TermDictionary::from_source(&termdict_source),
            postings_source,
            positions_source,
            positions_idx_source,
            record_option,
        ));

        // by releasing the lock in between, we may end up opening the inverting index
        // twice, but this is fine.
        self.inv_idx_reader_cache
            .write()
            .expect("Field reader cache lock poisoned. This should never happen.")
            .insert(field, Arc::clone(&inv_idx_reader));

        inv_idx_reader
    }

    /// Returns the segment id
    pub fn segment_id(&self) -> SegmentId {
        self.segment_id
    }

    /// Returns the bitset representing
    /// the documents that have been deleted.
    pub fn delete_bitset(&self) -> Option<&DeleteBitSet> {
        self.delete_bitset_opt.as_ref()
    }

    /// Returns true iff the `doc` is marked
    /// as deleted.
    pub fn is_deleted(&self, doc: DocId) -> bool {
        self.delete_bitset()
            .map(|delete_set| delete_set.is_deleted(doc))
            .unwrap_or(false)
    }

    /// Returns an iterator that will iterate over the alive document ids
    pub fn doc_ids_alive<'a>(&'a self) -> impl Iterator<Item = DocId> + 'a {
        (0u32..self.max_doc).filter(move |doc| !self.is_deleted(*doc))
    }

    /// Summarize total space usage of this segment.
    pub fn space_usage(&self) -> SegmentSpaceUsage {
        SegmentSpaceUsage::new(
            self.num_docs(),
            self.termdict_composite.space_usage(),
            self.postings_composite.space_usage(),
            self.positions_composite.space_usage(),
            self.positions_idx_composite.space_usage(),
            self.fast_fields_readers.space_usage(),
            self.fieldnorms_composite.space_usage(),
            self.get_store_reader().space_usage(),
            self.delete_bitset_opt
                .as_ref()
                .map(DeleteBitSet::space_usage)
                .unwrap_or(0),
        )
    }
}

impl fmt::Debug for SegmentReader {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SegmentReader({:?})", self.segment_id)
    }
}

#[cfg(test)]
mod test {
    use crate::core::Index;
    use crate::schema::{Schema, Term, STORED, TEXT};
    use crate::DocId;

    #[test]
    fn test_alive_docs_iterator() {
        let mut schema_builder = Schema::builder();
        schema_builder.add_text_field("name", TEXT | STORED);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema.clone());
        let name = schema.get_field("name").unwrap();

        {
            let mut index_writer = index.writer_with_num_threads(1, 3_000_000).unwrap();
            index_writer.add_document(doc!(name => "tantivy"));
            index_writer.add_document(doc!(name => "horse"));
            index_writer.add_document(doc!(name => "jockey"));
            index_writer.add_document(doc!(name => "cap"));

            // we should now have one segment with two docs
            index_writer.commit().unwrap();
        }

        {
            let mut index_writer2 = index.writer(50_000_000).unwrap();
            index_writer2.delete_term(Term::from_field_text(name, "horse"));
            index_writer2.delete_term(Term::from_field_text(name, "cap"));

            // ok, now we should have a deleted doc
            index_writer2.commit().unwrap();
        }
        let searcher = index.reader().unwrap().searcher();
        let docs: Vec<DocId> = searcher.segment_reader(0).doc_ids_alive().collect();
        assert_eq!(vec![0u32, 2u32], docs);
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use crate::PythonLogLevel;

use colored::*;
use std::cell::RefCell;
use std::collections::HashMap;
use std::convert::{TryFrom, TryInto};
use std::fs::File;
use std::fs::OpenOptions;
use std::future::Future;
use std::io::{stderr, Stderr, Write};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};

use lazy_static::lazy_static;
use log::{debug, log, set_logger, set_max_level, LevelFilter, Log, Metadata, Record};
use parking_lot::Mutex;
use simplelog::{ConfigBuilder, LevelPadding, WriteLogger};
use tokio::task_local;
use uuid::Uuid;

const TIME_FORMAT_STR: &str = "%H:%M:%S";

pub type StdioHandler = Box<dyn Fn(&str) -> Result<(), ()> + Send>;

lazy_static! {
  pub static ref LOGGER: Logger = Logger::new();
}

pub struct Logger {
  pantsd_log: Mutex<MaybeWriteLogger<File>>,
  stderr_log: Mutex<MaybeWriteLogger<Stderr>>,
  use_color: AtomicBool,
  show_rust_3rdparty_logs: AtomicBool,
  stderr_handlers: Mutex<HashMap<Uuid, StdioHandler>>,
}

impl Logger {
  pub fn new() -> Logger {
    Logger {
      pantsd_log: Mutex::new(MaybeWriteLogger::empty()),
      stderr_log: Mutex::new(MaybeWriteLogger::empty()),
      show_rust_3rdparty_logs: AtomicBool::new(true),
      use_color: AtomicBool::new(true),
      stderr_handlers: Mutex::new(HashMap::new()),
    }
  }

  pub fn init(max_level: u64, show_rust_3rdparty_logs: bool, use_color: bool) {
    let max_python_level: Result<PythonLogLevel, _> = max_level.try_into();
    match max_python_level {
      Ok(python_level) => {
        let level: log::LevelFilter = python_level.into();
        set_max_level(level);
        LOGGER.use_color.store(use_color, Ordering::SeqCst);
        LOGGER
          .show_rust_3rdparty_logs
          .store(show_rust_3rdparty_logs, Ordering::SeqCst);
        if set_logger(&*LOGGER).is_err() {
          debug!("Logging already initialized.");
        }
      }
      Err(err) => panic!("Unrecognised log level from python: {}: {}", max_level, err),
    };
  }

  fn maybe_increase_global_verbosity(&self, new_level: log::LevelFilter) {
    if log::max_level() < new_level {
      set_max_level(new_level);
    }
  }

  pub fn set_stderr_logger(&self, python_level: u64) -> Result<(), String> {
    python_level
      .try_into()
      .map_err(|err| format!("{}", err))
      .map(|level: PythonLogLevel| {
        self.maybe_increase_global_verbosity(level.into());
        *self.stderr_log.lock() = MaybeWriteLogger::new(
          stderr(),
          level.into(),
          self.show_rust_3rdparty_logs.load(Ordering::SeqCst),
        )
      })
  }

  ///
  /// Set up a file logger which logs at python_level to log_file_path.
  /// Returns the file descriptor of the log file.
  ///
  #[cfg(unix)]
  pub fn set_pantsd_logger(
    &self,
    log_file_path: PathBuf,
    python_level: u64,
  ) -> Result<std::os::unix::io::RawFd, String> {
    use std::os::unix::io::AsRawFd;
    python_level
      .try_into()
      .map_err(|err| format!("{}", err))
      .and_then(|level: PythonLogLevel| {
        {
          // Maybe close open file by dropping the existing logger
          *self.pantsd_log.lock() = MaybeWriteLogger::empty();
        }
        OpenOptions::new()
          .create(true)
          .append(true)
          .open(log_file_path)
          .map(|file| {
            let fd = file.as_raw_fd();
            self.maybe_increase_global_verbosity(level.into());
            *self.pantsd_log.lock() = MaybeWriteLogger::new(
              file,
              level.into(),
              self.show_rust_3rdparty_logs.load(Ordering::SeqCst),
            );
            fd
          })
          .map_err(|err| format!("Error opening pantsd logfile: {}", err))
      })
  }

  /// log_from_python is only used in the Python FFI, which in turn is only called within the
  /// Python `NativeHandler` class. Every logging call from Python should get proxied through this
  /// function, which translates the log message into the Rust log paradigm provided by
  /// the `log` crate.
  pub fn log_from_python(message: &str, python_level: u64, target: &str) -> Result<(), String> {
    python_level
      .try_into()
      .map_err(|err| format!("{}", err))
      .map(|level: PythonLogLevel| {
        log!(target: target, level.into(), "{}", message);
      })
  }

  pub fn register_stderr_handler(&self, callback: StdioHandler) -> Uuid {
    let mut handlers = self.stderr_handlers.lock();
    let unique_id = Uuid::new_v4();
    handlers.insert(unique_id, callback);
    unique_id
  }

  pub fn deregister_stderr_handler(&self, unique_id: Uuid) {
    let mut handlers = self.stderr_handlers.lock();
    handlers.remove(&unique_id);
  }
}

impl Log for Logger {
  fn enabled(&self, _metadata: &Metadata) -> bool {
    // Individual log levels are handled by each sub-logger,
    // And a global filter is applied to set_max_level.
    // No need to filter here.
    true
  }

  fn log(&self, record: &Record) {
    use chrono::Timelike;
    use log::Level;
    let destination = get_destination();
    match destination {
      Destination::Stderr => {
        let cur_date = chrono::Local::now();
        let time_str = format!(
          "{}.{:02}",
          cur_date.format(TIME_FORMAT_STR),
          cur_date.time().nanosecond() / 10_000_000 // two decimal places of precision
        );

        let level = record.level();
        let use_color = self.use_color.load(Ordering::SeqCst);

        let level_marker = match level {
          _ if !use_color => format!("[{}]", level).normal().clear(),
          Level::Info => format!("[{}]", level).normal(),
          Level::Error => format!("[{}]", level).red(),
          Level::Warn => format!("[{}]", level).red(),
          Level::Debug => format!("[{}]", level).green(),
          Level::Trace => format!("[{}]", level).magenta(),
        };

        let log_string = format!("{} {} {}", time_str, level_marker, record.args());

        {
          // If there are no handlers, or sending to any of the handlers failed, send to stderr
          // directly.
          let handlers_map = self.stderr_handlers.lock();
          let mut any_handler_failed = false;
          for callback in handlers_map.values() {
            if callback(&log_string).is_err() {
              any_handler_failed = true;
            }
          }
          if handlers_map.len() == 0 || any_handler_failed {
            self.stderr_log.lock().log(record);
          }
        }
      }
      Destination::Pantsd => self.pantsd_log.lock().log(record),
    }
  }

  fn flush(&self) {
    self.stderr_log.lock().flush();
    self.pantsd_log.lock().flush();
  }
}

struct MaybeWriteLogger<W: Write + Send + 'static> {
  level: LevelFilter,
  show_rust_3rdparty_logs: bool,
  inner: Option<Box<WriteLogger<W>>>,
}

impl<W: Write + Send + 'static> MaybeWriteLogger<W> {
  pub fn empty() -> MaybeWriteLogger<W> {
    MaybeWriteLogger {
      level: LevelFilter::Off,
      show_rust_3rdparty_logs: true,
      inner: None,
    }
  }

  pub fn new(
    writable: W,
    level: LevelFilter,
    show_rust_3rdparty_logs: bool,
  ) -> MaybeWriteLogger<W> {
    // We initialize the inner WriteLogger with no filters so that we don't
    // have to create a new one every time we change the level of the outer
    // MaybeWriteLogger.

    let config = ConfigBuilder::new()
      .set_time_format_str(TIME_FORMAT_STR)
      .set_thread_level(LevelFilter::Off)
      .set_level_padding(LevelPadding::Off)
      .build();

    MaybeWriteLogger {
      level,
      show_rust_3rdparty_logs,
      inner: Some(WriteLogger::new(LevelFilter::max(), config, writable)),
    }
  }
}

impl<W: Write + Send + 'static> Log for MaybeWriteLogger<W> {
  fn enabled(&self, metadata: &Metadata) -> bool {
    metadata.level() <= self.level
  }

  fn log(&self, record: &Record) {
    if !self.enabled(record.metadata()) {
      return;
    }
    let mut should_log = self.show_rust_3rdparty_logs;
    if !self.show_rust_3rdparty_logs {
      if let Some(ref module_path) = record.module_path() {
        for pants_package in super::pants_packages::PANTS_PACKAGE_NAMES {
          if &module_path.split("::").next().unwrap() == pants_package {
            should_log = true;
            break;
          }
        }
      } else {
        should_log = true;
      }
    }
    if !should_log {
      return;
    }
    if let Some(ref logger) = self.inner {
      logger.log(record);
    }
  }

  fn flush(&self) {
    if let Some(ref logger) = self.inner {
      logger.flush();
    }
  }
}

///
/// Thread- or task-local context for where the Logger should send log statements.
///
/// We do this in a per-thread way because we find that Pants threads generally are either
/// daemon-specific or user-facing. We make sure that every time we spawn a thread on the Python
/// side, we set the thread-local information, and every time we submit a Future to a tokio Runtime
/// on the rust side, we set the task-local information.
///
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[repr(C)]
pub enum Destination {
  Pantsd,
  Stderr,
}

impl TryFrom<&str> for Destination {
  type Error = String;
  fn try_from(dest: &str) -> Result<Self, Self::Error> {
    match dest {
      "pantsd" => Ok(Destination::Pantsd),
      "stderr" => Ok(Destination::Stderr),
      other => Err(format!("Unknown log destination: {:?}", other)),
    }
  }
}

thread_local! {
  static THREAD_DESTINATION: RefCell<Destination> = RefCell::new(Destination::Stderr)
}

task_local! {
  static TASK_DESTINATION: Destination;
}

///
/// Set the current log destination for a Thread, but _not_ for a Task. Tasks must always be spawned
/// by callers using the `scope_task_destination` helper (generally via task_executor::Executor.)
///
pub fn set_thread_destination(destination: Destination) {
  THREAD_DESTINATION.with(|thread_destination| {
    *thread_destination.borrow_mut() = destination;
  })
}

///
/// Propagate the current log destination to a Future representing a newly spawned Task. Usage of
/// this method should mostly be contained to task_executor::Executor.
///
pub async fn scope_task_destination<F>(destination: Destination, f: F) -> F::Output
where
  F: Future,
{
  TASK_DESTINATION.scope(destination, f).await
}

///
/// Get the current log destination, from either a Task or a Thread.
///
/// TODO: Having this return an Option and tracking down all cases where it has defaulted would be
/// good.
///
pub fn get_destination() -> Destination {
  if let Ok(destination) = TASK_DESTINATION.try_with(|destination| *destination) {
    destination
  } else {
    THREAD_DESTINATION.with(|destination| *destination.borrow())
  }
}
<commit_msg> Print initial logs from pantsd using local time format<commit_after>// Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use crate::PythonLogLevel;

use colored::*;
use std::cell::RefCell;
use std::collections::HashMap;
use std::convert::{TryFrom, TryInto};
use std::fs::File;
use std::fs::OpenOptions;
use std::future::Future;
use std::io::{stderr, Stderr, Write};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};

use lazy_static::lazy_static;
use log::{debug, log, set_logger, set_max_level, LevelFilter, Log, Metadata, Record};
use parking_lot::Mutex;
use simplelog::{ConfigBuilder, LevelPadding, WriteLogger};
use tokio::task_local;
use uuid::Uuid;

const TIME_FORMAT_STR: &str = "%H:%M:%S";

pub type StdioHandler = Box<dyn Fn(&str) -> Result<(), ()> + Send>;

lazy_static! {
  pub static ref LOGGER: Logger = Logger::new();
}

pub struct Logger {
  pantsd_log: Mutex<MaybeWriteLogger<File>>,
  stderr_log: Mutex<MaybeWriteLogger<Stderr>>,
  use_color: AtomicBool,
  show_rust_3rdparty_logs: AtomicBool,
  stderr_handlers: Mutex<HashMap<Uuid, StdioHandler>>,
}

impl Logger {
  pub fn new() -> Logger {
    Logger {
      pantsd_log: Mutex::new(MaybeWriteLogger::empty()),
      stderr_log: Mutex::new(MaybeWriteLogger::empty()),
      show_rust_3rdparty_logs: AtomicBool::new(true),
      use_color: AtomicBool::new(true),
      stderr_handlers: Mutex::new(HashMap::new()),
    }
  }

  pub fn init(max_level: u64, show_rust_3rdparty_logs: bool, use_color: bool) {
    let max_python_level: Result<PythonLogLevel, _> = max_level.try_into();
    match max_python_level {
      Ok(python_level) => {
        let level: log::LevelFilter = python_level.into();
        set_max_level(level);
        LOGGER.use_color.store(use_color, Ordering::SeqCst);
        LOGGER
          .show_rust_3rdparty_logs
          .store(show_rust_3rdparty_logs, Ordering::SeqCst);
        if set_logger(&*LOGGER).is_err() {
          debug!("Logging already initialized.");
        }
      }
      Err(err) => panic!("Unrecognised log level from python: {}: {}", max_level, err),
    };
  }

  fn maybe_increase_global_verbosity(&self, new_level: log::LevelFilter) {
    if log::max_level() < new_level {
      set_max_level(new_level);
    }
  }

  pub fn set_stderr_logger(&self, python_level: u64) -> Result<(), String> {
    python_level
      .try_into()
      .map_err(|err| format!("{}", err))
      .map(|level: PythonLogLevel| {
        self.maybe_increase_global_verbosity(level.into());
        *self.stderr_log.lock() = MaybeWriteLogger::new(
          stderr(),
          level.into(),
          self.show_rust_3rdparty_logs.load(Ordering::SeqCst),
        )
      })
  }

  ///
  /// Set up a file logger which logs at python_level to log_file_path.
  /// Returns the file descriptor of the log file.
  ///
  #[cfg(unix)]
  pub fn set_pantsd_logger(
    &self,
    log_file_path: PathBuf,
    python_level: u64,
  ) -> Result<std::os::unix::io::RawFd, String> {
    use std::os::unix::io::AsRawFd;
    python_level
      .try_into()
      .map_err(|err| format!("{}", err))
      .and_then(|level: PythonLogLevel| {
        {
          // Maybe close open file by dropping the existing logger
          *self.pantsd_log.lock() = MaybeWriteLogger::empty();
        }
        OpenOptions::new()
          .create(true)
          .append(true)
          .open(log_file_path)
          .map(|file| {
            let fd = file.as_raw_fd();
            self.maybe_increase_global_verbosity(level.into());
            *self.pantsd_log.lock() = MaybeWriteLogger::new(
              file,
              level.into(),
              self.show_rust_3rdparty_logs.load(Ordering::SeqCst),
            );
            fd
          })
          .map_err(|err| format!("Error opening pantsd logfile: {}", err))
      })
  }

  /// log_from_python is only used in the Python FFI, which in turn is only called within the
  /// Python `NativeHandler` class. Every logging call from Python should get proxied through this
  /// function, which translates the log message into the Rust log paradigm provided by
  /// the `log` crate.
  pub fn log_from_python(message: &str, python_level: u64, target: &str) -> Result<(), String> {
    python_level
      .try_into()
      .map_err(|err| format!("{}", err))
      .map(|level: PythonLogLevel| {
        log!(target: target, level.into(), "{}", message);
      })
  }

  pub fn register_stderr_handler(&self, callback: StdioHandler) -> Uuid {
    let mut handlers = self.stderr_handlers.lock();
    let unique_id = Uuid::new_v4();
    handlers.insert(unique_id, callback);
    unique_id
  }

  pub fn deregister_stderr_handler(&self, unique_id: Uuid) {
    let mut handlers = self.stderr_handlers.lock();
    handlers.remove(&unique_id);
  }
}

impl Log for Logger {
  fn enabled(&self, _metadata: &Metadata) -> bool {
    // Individual log levels are handled by each sub-logger,
    // And a global filter is applied to set_max_level.
    // No need to filter here.
    true
  }

  fn log(&self, record: &Record) {
    use chrono::Timelike;
    use log::Level;
    let destination = get_destination();
    match destination {
      Destination::Stderr => {
        let cur_date = chrono::Local::now();
        let time_str = format!(
          "{}.{:02}",
          cur_date.format(TIME_FORMAT_STR),
          cur_date.time().nanosecond() / 10_000_000 // two decimal places of precision
        );

        let level = record.level();
        let use_color = self.use_color.load(Ordering::SeqCst);

        let level_marker = match level {
          _ if !use_color => format!("[{}]", level).normal().clear(),
          Level::Info => format!("[{}]", level).normal(),
          Level::Error => format!("[{}]", level).red(),
          Level::Warn => format!("[{}]", level).red(),
          Level::Debug => format!("[{}]", level).green(),
          Level::Trace => format!("[{}]", level).magenta(),
        };

        let log_string = format!("{} {} {}", time_str, level_marker, record.args());

        {
          // If there are no handlers, or sending to any of the handlers failed, send to stderr
          // directly.
          let handlers_map = self.stderr_handlers.lock();
          let mut any_handler_failed = false;
          for callback in handlers_map.values() {
            if callback(&log_string).is_err() {
              any_handler_failed = true;
            }
          }
          if handlers_map.len() == 0 || any_handler_failed {
            self.stderr_log.lock().log(record);
          }
        }
      }
      Destination::Pantsd => self.pantsd_log.lock().log(record),
    }
  }

  fn flush(&self) {
    self.stderr_log.lock().flush();
    self.pantsd_log.lock().flush();
  }
}

struct MaybeWriteLogger<W: Write + Send + 'static> {
  level: LevelFilter,
  show_rust_3rdparty_logs: bool,
  inner: Option<Box<WriteLogger<W>>>,
}

impl<W: Write + Send + 'static> MaybeWriteLogger<W> {
  pub fn empty() -> MaybeWriteLogger<W> {
    MaybeWriteLogger {
      level: LevelFilter::Off,
      show_rust_3rdparty_logs: true,
      inner: None,
    }
  }

  pub fn new(
    writable: W,
    level: LevelFilter,
    show_rust_3rdparty_logs: bool,
  ) -> MaybeWriteLogger<W> {
    // We initialize the inner WriteLogger with no filters so that we don't
    // have to create a new one every time we change the level of the outer
    // MaybeWriteLogger.

    let config = ConfigBuilder::new()
      .set_time_format_str(TIME_FORMAT_STR)
      .set_time_to_local(true)
      .set_thread_level(LevelFilter::Off)
      .set_level_padding(LevelPadding::Off)
      .build();

    MaybeWriteLogger {
      level,
      show_rust_3rdparty_logs,
      inner: Some(WriteLogger::new(LevelFilter::max(), config, writable)),
    }
  }
}

impl<W: Write + Send + 'static> Log for MaybeWriteLogger<W> {
  fn enabled(&self, metadata: &Metadata) -> bool {
    metadata.level() <= self.level
  }

  fn log(&self, record: &Record) {
    if !self.enabled(record.metadata()) {
      return;
    }
    let mut should_log = self.show_rust_3rdparty_logs;
    if !self.show_rust_3rdparty_logs {
      if let Some(ref module_path) = record.module_path() {
        for pants_package in super::pants_packages::PANTS_PACKAGE_NAMES {
          if &module_path.split("::").next().unwrap() == pants_package {
            should_log = true;
            break;
          }
        }
      } else {
        should_log = true;
      }
    }
    if !should_log {
      return;
    }
    if let Some(ref logger) = self.inner {
      logger.log(record);
    }
  }

  fn flush(&self) {
    if let Some(ref logger) = self.inner {
      logger.flush();
    }
  }
}

///
/// Thread- or task-local context for where the Logger should send log statements.
///
/// We do this in a per-thread way because we find that Pants threads generally are either
/// daemon-specific or user-facing. We make sure that every time we spawn a thread on the Python
/// side, we set the thread-local information, and every time we submit a Future to a tokio Runtime
/// on the rust side, we set the task-local information.
///
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[repr(C)]
pub enum Destination {
  Pantsd,
  Stderr,
}

impl TryFrom<&str> for Destination {
  type Error = String;
  fn try_from(dest: &str) -> Result<Self, Self::Error> {
    match dest {
      "pantsd" => Ok(Destination::Pantsd),
      "stderr" => Ok(Destination::Stderr),
      other => Err(format!("Unknown log destination: {:?}", other)),
    }
  }
}

thread_local! {
  static THREAD_DESTINATION: RefCell<Destination> = RefCell::new(Destination::Stderr)
}

task_local! {
  static TASK_DESTINATION: Destination;
}

///
/// Set the current log destination for a Thread, but _not_ for a Task. Tasks must always be spawned
/// by callers using the `scope_task_destination` helper (generally via task_executor::Executor.)
///
pub fn set_thread_destination(destination: Destination) {
  THREAD_DESTINATION.with(|thread_destination| {
    *thread_destination.borrow_mut() = destination;
  })
}

///
/// Propagate the current log destination to a Future representing a newly spawned Task. Usage of
/// this method should mostly be contained to task_executor::Executor.
///
pub async fn scope_task_destination<F>(destination: Destination, f: F) -> F::Output
where
  F: Future,
{
  TASK_DESTINATION.scope(destination, f).await
}

///
/// Get the current log destination, from either a Task or a Thread.
///
/// TODO: Having this return an Option and tracking down all cases where it has defaulted would be
/// good.
///
pub fn get_destination() -> Destination {
  if let Ok(destination) = TASK_DESTINATION.try_with(|destination| *destination) {
    destination
  } else {
    THREAD_DESTINATION.with(|destination| *destination.borrow())
  }
}
<|endoftext|>


<commit_before>extern crate serde_json;

pub mod hashutils;
pub mod tree;
pub mod proof;
pub mod merkletree;

use self::hashutils::*;
use self::tree::*;
use self::merkletree::*;

impl MerkleTree {
    fn count_bits(v: usize) -> usize {
        let mut ret = 0;
        let mut val = v;
        while val != 0 {
            val &= val - 1;
            ret += 1;
        }
        return ret;
    }

    pub fn append(&mut self, node: TreeLeafData) {
        if self.count == 0 {
            // empty tree
            self.root = Tree::new_leaf(node);
        }
        else if Self::count_bits(self.count) != 1 {
            // add to right subtree
            match self.root.clone() {
                Tree::Node { ref left, ref right, .. } => {
                    let mut iter = right.iter().map(|x| (*x).clone()).collect::<Vec<TreeLeafData>>();
                    iter.push(node);
                    let new_right = MerkleTree::from_vec(iter);
                    let combined_hash = DIGEST.hash_nodes(
                        left.hash(),
                        new_right.root_hash() as &Vec<u8>
                    );
                    self.root = Tree::Node {
                        left: (*left).clone(),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
        }
        else
        {
            // add tree layer
            let new_right = MerkleTree::from_vec(vec![ node ]);
            match self.root.clone() {
                Tree::Node { ref left, ref right, ref hash } => {
                    let combined_hash = DIGEST.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(self.root.clone()),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                Tree::Leaf { ref hash, ref value } => {
                    let combined_hash = DIGEST.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(Tree::new_leaf((*value).clone())),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
            self.height += 1;
        }
        self.count += 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use self::serde_json;

    #[test]
    fn test_merkletree_append() {
        let values = vec![
            "{\"data\":{\"alias\":\"Node1\",\"client_ip\":\"192.168.1.35\",\"client_port\":9702,\"node_ip\":\"192.168.1.35\",\"node_port\":9701,\"services\":[\"VALIDATOR\"]},\"dest\":\"Gw6pDLhcBcoQesN72qfotTgFa7cbuqZpkX3Xo6pLhPhv\",\"identifier\":\"FYmoFw55GeQH7SRFa37dkx1d2dZ3zUF8ckg7wmL7ofN4\",\"txnId\":\"fea82e10e894419fe2bea7d96296a6d46f50f93f9eeda954ec461b2ed2950b62\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node2\",\"client_ip\":\"192.168.1.35\",\"client_port\":9704,\"node_ip\":\"192.168.1.35\",\"node_port\":9703,\"services\":[\"VALIDATOR\"]},\"dest\":\"8ECVSk179mjsjKRLWiQtssMLgp6EPhWXtaYyStWPSGAb\",\"identifier\":\"8QhFxKxyaFsJy4CyxeYX34dFH8oWqyBv1P4HLQCsoeLy\",\"txnId\":\"1ac8aece2a18ced660fef8694b61aac3af08ba875ce3026a160acbc3a3af35fc\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node3\",\"client_ip\":\"192.168.1.35\",\"client_port\":9706,\"node_ip\":\"192.168.1.35\",\"node_port\":9705,\"services\":[\"VALIDATOR\"]},\"dest\":\"DKVxG2fXXTU8yT5N7hGEbXB3dfdAnYv1JczDUHpmDxya\",\"identifier\":\"2yAeV5ftuasWNgQwVYzeHeTuM7LwwNtPR3Zg9N4JiDgF\",\"txnId\":\"7e9f355dffa78ed24668f0e0e369fd8c224076571c51e2ea8be5f26479edebe4\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node4\",\"client_ip\":\"192.168.1.35\",\"client_port\":9708,\"node_ip\":\"192.168.1.35\",\"node_port\":9707,\"services\":[\"VALIDATOR\"]},\"dest\":\"4PS3EDQ3dW1tci1Bp6543CfuuebjFrg36kLAUcskGfaA\",\"identifier\":\"FTE95CVthRtrBnK2PYCBbC9LghTcGwi9Zfi1Gz2dnyNx\",\"txnId\":\"aa5e817d7cc626170eca175822029339a444eb0ee8f0bd20d3b0b76e566fb008\",\"type\":\"0\"}" ];
        let mut mt = MerkleTree::from_vec(vec![]);
        println!("root(0)={}", mt.root_hash_hex());
        let mut r = 1;
        for i in values {
            mt.append(String::from(i));
            println!("root({})={}", r, mt.root_hash_hex());
            r+=1;
        }
        assert_eq!(mt.root_hash_hex(), "1285070cf01debc1155cef8dfd5ba54c05abb919a4c08c8632b079fb1e1e5e7c");
    }

    #[test]
    fn test_valid_proof() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let tree      = MerkleTree::from_vec(values.clone());
        let root_hash = tree.root_hash();

        for value in values {
            let proof    = tree.gen_proof(value);
            let is_valid = proof.map(|p| p.validate(&root_hash)).unwrap_or(false);

            assert!(is_valid);
        }
    }

    #[test]
    fn test_merkletree_serialize() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let mt = MerkleTree::from_vec(values.clone());
        let serialized = serde_json::to_string(&mt).unwrap();
        println!("serialize mt: h={}, c={}, rhash={}", mt.height, mt.count, serialized);
        let newmt: MerkleTree = serde_json::from_str(serialized.as_str()).unwrap();
        println!("serialize newmt: h={}, c={}", newmt.height, newmt.count);

        let mut collected = Vec::new();
        for value in &newmt {
            collected.push(value);
        }
        let refs = values.iter().collect::<Vec<_>>();
        assert_eq!(refs, collected);

        assert_eq!(mt.root_hash(), newmt.root_hash());
    }
}
<commit_msg>implement tree node search<commit_after>extern crate serde_json;

pub mod hashutils;
pub mod tree;
pub mod proof;
pub mod merkletree;

use self::hashutils::*;
use self::tree::*;
use self::merkletree::*;

impl MerkleTree {
    fn count_bits(v: usize) -> usize {
        let mut ret = 0;
        let mut val = v;
        while val != 0 {
            val &= val - 1;
            ret += 1;
        }
        return ret;
    }

    pub fn hash_hex(rh: &Vec<u8>) -> String {
        let mut ret:String = String::with_capacity(DIGEST.output_len*2);
        for i in rh {
            ret.push_str(&format!("{:02x}", i));
        }
        return ret;
    }

    pub fn find_hash<'a>(from: &'a Tree, required_hash: &Vec<u8>) -> Option<&'a Tree> {
        match from {
            &Tree::Empty{ ref hash, .. } => {
                assert!(false);
                return None;
            },
            &Tree::Node{ ref left, ref right, ref hash, .. } => {
                if hash == required_hash {
                    return Some(from);
                } else {
                    let right = MerkleTree::find_hash(right, required_hash);
                    match right {
                        Some(r) => {
                            return Some(r);
                        },
                        None => {
                            let left = MerkleTree::find_hash(left, required_hash);
                            match left {
                                Some(r) => {
                                    return Some(r);
                                },
                                None => {
                                    return None;
                                }
                            }
                        }
                    }
                }
            },
            &Tree::Leaf{ ref hash, .. } => {
                if hash == required_hash {
                    return Some(from);
                } else {
                    return None;
                }
            }
        }
    }

    pub fn consistency_proof(new_hash: &Vec<u8>, proof: &Vec<Vec<u8>>) -> bool {
        return true;
    }

    pub fn append(&mut self, node: TreeLeafData) {
        if self.count == 0 {
            // empty tree
            self.root = Tree::new_leaf(node);
        }
        else if Self::count_bits(self.count) != 1 {
            // add to right subtree
            match self.root.clone() {
                Tree::Node { ref left, ref right, .. } => {
                    let mut iter = right.iter().map(|x| (*x).clone()).collect::<Vec<TreeLeafData>>();
                    iter.push(node);
                    let new_right = MerkleTree::from_vec(iter);
                    let combined_hash = DIGEST.hash_nodes(
                        left.hash(),
                        new_right.root_hash() as &Vec<u8>
                    );
                    self.root = Tree::Node {
                        left: (*left).clone(),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
        }
        else
        {
            // add tree layer
            let new_right = MerkleTree::from_vec(vec![ node ]);
            match self.root.clone() {
                Tree::Node { ref left, ref right, ref hash } => {
                    let combined_hash = DIGEST.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(self.root.clone()),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                Tree::Leaf { ref hash, ref value } => {
                    let combined_hash = DIGEST.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(Tree::new_leaf((*value).clone())),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
            self.height += 1;
        }
        self.count += 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use self::serde_json;

    #[test]
    fn test_merkletree_append() {
        let values = vec![
            "{\"data\":{\"alias\":\"Node1\",\"client_ip\":\"192.168.1.35\",\"client_port\":9702,\"node_ip\":\"192.168.1.35\",\"node_port\":9701,\"services\":[\"VALIDATOR\"]},\"dest\":\"Gw6pDLhcBcoQesN72qfotTgFa7cbuqZpkX3Xo6pLhPhv\",\"identifier\":\"FYmoFw55GeQH7SRFa37dkx1d2dZ3zUF8ckg7wmL7ofN4\",\"txnId\":\"fea82e10e894419fe2bea7d96296a6d46f50f93f9eeda954ec461b2ed2950b62\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node2\",\"client_ip\":\"192.168.1.35\",\"client_port\":9704,\"node_ip\":\"192.168.1.35\",\"node_port\":9703,\"services\":[\"VALIDATOR\"]},\"dest\":\"8ECVSk179mjsjKRLWiQtssMLgp6EPhWXtaYyStWPSGAb\",\"identifier\":\"8QhFxKxyaFsJy4CyxeYX34dFH8oWqyBv1P4HLQCsoeLy\",\"txnId\":\"1ac8aece2a18ced660fef8694b61aac3af08ba875ce3026a160acbc3a3af35fc\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node3\",\"client_ip\":\"192.168.1.35\",\"client_port\":9706,\"node_ip\":\"192.168.1.35\",\"node_port\":9705,\"services\":[\"VALIDATOR\"]},\"dest\":\"DKVxG2fXXTU8yT5N7hGEbXB3dfdAnYv1JczDUHpmDxya\",\"identifier\":\"2yAeV5ftuasWNgQwVYzeHeTuM7LwwNtPR3Zg9N4JiDgF\",\"txnId\":\"7e9f355dffa78ed24668f0e0e369fd8c224076571c51e2ea8be5f26479edebe4\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node4\",\"client_ip\":\"192.168.1.35\",\"client_port\":9708,\"node_ip\":\"192.168.1.35\",\"node_port\":9707,\"services\":[\"VALIDATOR\"]},\"dest\":\"4PS3EDQ3dW1tci1Bp6543CfuuebjFrg36kLAUcskGfaA\",\"identifier\":\"FTE95CVthRtrBnK2PYCBbC9LghTcGwi9Zfi1Gz2dnyNx\",\"txnId\":\"aa5e817d7cc626170eca175822029339a444eb0ee8f0bd20d3b0b76e566fb008\",\"type\":\"0\"}" ];
        let mut mt = MerkleTree::from_vec(vec![]);
        println!("root(0)={}", mt.root_hash_hex());
        let mut r = 1;
        for i in values {
            mt.append(String::from(i));
            println!("root({})={}", r, mt.root_hash_hex());
            r+=1;
        }
        assert_eq!(mt.root_hash_hex(), "1285070cf01debc1155cef8dfd5ba54c05abb919a4c08c8632b079fb1e1e5e7c");
    }

    #[test]
    fn test_merkletree_find_hash() {
        let values = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9" ];
        let mut mt = MerkleTree::from_vec(vec![]);
        println!("root(0)={}", mt.root_hash_hex());
        let mut r = 1;
        for i in values {
            mt.append(String::from(i));
            println!("root({})={}", r, mt.root_hash_hex());
            r+=1;
        }

        let mut rh: Vec<u8>;

        rh = vec![ 0xe8, 0xbc, 0xd9, 0x7e, 0x34, 0x96, 0x93, 0xdc,
                   0xfe, 0xc0, 0x54, 0xfe, 0x21, 0x9a, 0xb3, 0x57,
                   0xb7, 0x5d, 0x3c, 0x1c, 0xd9, 0xf8, 0xbe, 0x17,
                   0x67, 0xf6, 0x09, 0x0f, 0x9c, 0x86, 0xf9, 0xfd ];
        assert!(MerkleTree::find_hash(&mt.root, &rh) != None);

        rh = vec![ 0x22, 0x15, 0xe8, 0xac, 0x4e, 0x2b, 0x87, 0x1c,
                   0x2a, 0x48, 0x18, 0x9e, 0x79, 0x73, 0x8c, 0x95,
                   0x6c, 0x08, 0x1e, 0x23, 0xac, 0x2f, 0x24, 0x15,
                   0xbf, 0x77, 0xda, 0x19, 0x9d, 0xfd, 0x92, 0x0c ];
        assert!(MerkleTree::find_hash(&mt.root, &rh) != None);

        rh = vec![ 0x23, 0x15, 0xe8, 0xac, 0x4e, 0x2b, 0x87, 0x1c,
                   0x2a, 0x48, 0x18, 0x9e, 0x79, 0x73, 0x8c, 0x95,
                   0x6c, 0x08, 0x1e, 0x23, 0xac, 0x2f, 0x24, 0x15,
                   0xbf, 0x77, 0xda, 0x19, 0x9d, 0xfd, 0x92, 0x0d ];
        assert!(MerkleTree::find_hash(&mt.root, &rh) == None);
    }

    #[test]
    fn test_valid_proof() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let tree      = MerkleTree::from_vec(values.clone());
        let root_hash = tree.root_hash();

        for value in values {
            let proof    = tree.gen_proof(value);
            let is_valid = proof.map(|p| p.validate(&root_hash)).unwrap_or(false);

            assert!(is_valid);
        }
    }

    #[test]
    fn test_merkletree_serialize() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let mt = MerkleTree::from_vec(values.clone());
        let serialized = serde_json::to_string(&mt).unwrap();
        println!("serialize mt: h={}, c={}, rhash={}", mt.height, mt.count, serialized);
        let newmt: MerkleTree = serde_json::from_str(serialized.as_str()).unwrap();
        println!("serialize newmt: h={}, c={}", newmt.height, newmt.count);

        let mut collected = Vec::new();
        for value in &newmt {
            collected.push(value);
        }
        let refs = values.iter().collect::<Vec<_>>();
        assert_eq!(refs, collected);

        assert_eq!(mt.root_hash(), newmt.root_hash());
    }
}
<|endoftext|>