<commit_before>use crate::config;

use crate::infer_schema_internals::*;
use serde::de::{self, MapAccess, Visitor};
use serde::{Deserialize, Deserializer};
use serde_regex::Serde as RegexWrapper;
use std::error::Error;
use std::fmt::{self, Display, Formatter, Write};
use std::io::Write as IoWrite;

const SCHEMA_HEADER: &str = "// @generated automatically by Diesel CLI.\n";

type Regex = RegexWrapper<regex::Regex>;

pub enum Filtering {
    OnlyTables(Vec<TableName>),
    OnlyTableRegexes(Vec<Regex>),
    ExceptTables(Vec<TableName>),
    ExceptTableRegexes(Vec<Regex>),
    None,
}

impl Default for Filtering {
    fn default() -> Self {
        Filtering::None
    }
}

impl Filtering {
    pub fn should_ignore_table(&self, name: &TableName) -> bool {
        use self::Filtering::*;

        match *self {
            OnlyTables(ref names) => !names.contains(name),
            OnlyTableRegexes(ref regexes) => {
                !regexes.iter().any(|regex| regex.is_match(&name.sql_name))
            }
            ExceptTables(ref names) => names.contains(name),
            ExceptTableRegexes(ref regexes) => {
                regexes.iter().any(|regex| regex.is_match(&name.sql_name))
            }
            None => false,
        }
    }
}

pub fn run_print_schema<W: IoWrite>(
    database_url: &str,
    config: &config::PrintSchema,
    output: &mut W,
) -> Result<(), Box<dyn Error>> {
    let schema = output_schema(database_url, config)?;

    output.write_all(schema.as_bytes())?;
    Ok(())
}

pub fn output_schema(
    database_url: &str,
    config: &config::PrintSchema,
) -> Result<String, Box<dyn Error>> {
    let table_names = load_table_names(database_url, config.schema_name())?
        .into_iter()
        .filter(|t| !config.filter.should_ignore_table(t))
        .collect::<Vec<_>>();
    let foreign_keys = load_foreign_key_constraints(database_url, config.schema_name())?;
    let foreign_keys =
        remove_unsafe_foreign_keys_for_codegen(database_url, &foreign_keys, &table_names);
    let table_data = table_names
        .into_iter()
        .map(|t| load_table_data(database_url, t))
        .collect::<Result<_, Box<dyn Error>>>()?;
    let definitions = TableDefinitions {
        tables: table_data,
        fk_constraints: foreign_keys,
        include_docs: config.with_docs,
        import_types: config.import_types(),
    };

    let mut out = String::new();
    writeln!(out, "{}", SCHEMA_HEADER)?;

    if let Some(schema_name) = config.schema_name() {
        write!(out, "{}", ModuleDefinition(schema_name, definitions))?;
    } else {
        write!(out, "{}", definitions)?;
    }

    if let Some(ref patch_file) = config.patch_file {
        let patch = std::fs::read_to_string(patch_file)?;
        let patch = diffy::Patch::from_str(&patch)?;

        out = diffy::apply(&out, &patch)?;
    }

    Ok(out)
}

struct ModuleDefinition<'a>(&'a str, TableDefinitions<'a>);

impl<'a> Display for ModuleDefinition<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        {
            let mut out = PadAdapter::new(f);
            writeln!(out, "pub mod {} {{", self.0)?;
            write!(out, "{}", self.1)?;
        }
        writeln!(f, "}}")?;
        Ok(())
    }
}

struct TableDefinitions<'a> {
    tables: Vec<TableData>,
    fk_constraints: Vec<ForeignKeyConstraint>,
    include_docs: bool,
    import_types: Option<&'a [String]>,
}

impl<'a> Display for TableDefinitions<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let mut is_first = true;
        for table in &self.tables {
            if is_first {
                is_first = false;
            } else {
                writeln!(f)?;
            }
            writeln!(
                f,
                "{}",
                TableDefinition {
                    table,
                    include_docs: self.include_docs,
                    import_types: self.import_types,
                }
            )?;
        }

        if !self.fk_constraints.is_empty() {
            writeln!(f)?;
        }

        for foreign_key in &self.fk_constraints {
            writeln!(f, "{}", Joinable(foreign_key))?;
        }

        if self.tables.len() > 1 {
            write!(f, "\ndiesel::allow_tables_to_appear_in_same_query!(")?;
            {
                let mut out = PadAdapter::new(f);
                writeln!(out)?;
                for table in &self.tables {
                    if table.name.rust_name == table.name.sql_name {
                        writeln!(out, "{},", table.name.sql_name)?;
                    } else {
                        writeln!(out, "{},", table.name.rust_name)?;
                    }
                }
            }
            writeln!(f, ");")?;
        }

        Ok(())
    }
}

struct TableDefinition<'a> {
    table: &'a TableData,
    import_types: Option<&'a [String]>,
    include_docs: bool,
}

impl<'a> Display for TableDefinition<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "diesel::table! {{")?;
        {
            let mut out = PadAdapter::new(f);
            writeln!(out)?;

            if let Some(types) = self.import_types {
                for import in types {
                    writeln!(out, "use {};", import)?;
                }
                writeln!(out)?;
            }

            if self.include_docs {
                for d in self.table.docs.lines() {
                    writeln!(out, "///{}{}", if d.is_empty() { "" } else { " " }, d)?;
                }
            }

            if self.table.name.rust_name != self.table.name.sql_name {
                writeln!(
                    out,
                    r#"#[sql_name = "{}"]"#,
                    self.table.name.full_sql_name()
                )?;
            }

            write!(out, "{} (", self.table.name)?;

            for (i, pk) in self.table.primary_key.iter().enumerate() {
                if i != 0 {
                    write!(out, ", ")?;
                }
                write!(out, "{}", pk)?;
            }

            write!(
                out,
                ") {}",
                ColumnDefinitions {
                    columns: &self.table.column_data,
                    include_docs: self.include_docs,
                }
            )?;
        }
        write!(f, "}}")?;
        Ok(())
    }
}

struct ColumnDefinitions<'a> {
    columns: &'a [ColumnDefinition],
    include_docs: bool,
}

impl<'a> Display for ColumnDefinitions<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        {
            let mut out = PadAdapter::new(f);
            writeln!(out, "{{")?;
            for column in self.columns {
                if self.include_docs {
                    for d in column.docs.lines() {
                        writeln!(out, "///{}{}", if d.is_empty() { "" } else { " " }, d)?;
                    }
                }
                if column.rust_name == column.sql_name {
                    writeln!(out, "{} -> {},", column.sql_name, column.ty)?;
                } else {
                    writeln!(out, r#"#[sql_name = "{}"]"#, column.sql_name)?;
                    writeln!(out, "{} -> {},", column.rust_name, column.ty)?;
                }
            }
        }
        writeln!(f, "}}")?;
        Ok(())
    }
}

struct Joinable<'a>(&'a ForeignKeyConstraint);

impl<'a> Display for Joinable<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let child_table_name = &self.0.child_table.rust_name;

        let parent_table_name = &self.0.parent_table.rust_name;

        write!(
            f,
            "diesel::joinable!({} -> {} ({}));",
            child_table_name, parent_table_name, self.0.foreign_key_rust_name,
        )
    }
}

/// Lifted directly from libcore/fmt/builders.rs
struct PadAdapter<'a, 'b: 'a> {
    fmt: &'a mut Formatter<'b>,
    on_newline: bool,
}

impl<'a, 'b: 'a> PadAdapter<'a, 'b> {
    fn new(fmt: &'a mut Formatter<'b>) -> PadAdapter<'a, 'b> {
        PadAdapter {
            fmt,
            on_newline: false,
        }
    }
}

impl<'a, 'b: 'a> Write for PadAdapter<'a, 'b> {
    fn write_str(&mut self, mut s: &str) -> fmt::Result {
        while !s.is_empty() {
            let on_newline = self.on_newline;

            let split = match s.find('\n') {
                Some(pos) => {
                    self.on_newline = true;
                    pos + 1
                }
                None => {
                    self.on_newline = false;
                    s.len()
                }
            };

            let to_write = &s[..split];
            if on_newline && to_write != "\n" {
                self.fmt.write_str("    ")?;
            }
            self.fmt.write_str(to_write)?;

            s = &s[split..];
        }

        Ok(())
    }
}

impl<'de> Deserialize<'de> for Filtering {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct FilteringVisitor;

        impl<'de> Visitor<'de> for FilteringVisitor {
            type Value = Filtering;

            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
                f.write_str("either only_tables, only_table_regexes, except_tables, or except_table_regexes")
            }

            fn visit_map<V>(self, mut map: V) -> Result<Self::Value, V::Error>
            where
                V: MapAccess<'de>,
            {
                let mut only_tables = None::<Vec<TableName>>;
                let mut only_table_regexes = None::<Vec<Regex>>;
                let mut except_tables = None::<Vec<TableName>>;
                let mut except_table_regexes = None::<Vec<Regex>>;
                while let Some(key) = map.next_key()? {
                    match key {
                        "only_tables" => {
                            if only_tables.is_some() {
                                return Err(de::Error::duplicate_field("only_tables"));
                            }
                            only_tables = Some(map.next_value()?);
                        }
                        "only_table_regexes" => {
                            if only_table_regexes.is_some() {
                                return Err(de::Error::duplicate_field("only_table_regexes"));
                            }
                            only_table_regexes = Some(map.next_value()?);
                        }
                        "except_tables" => {
                            if except_tables.is_some() {
                                return Err(de::Error::duplicate_field("except_tables"));
                            }
                            except_tables = Some(map.next_value()?);
                        }
                        "except_table_regexes" => {
                            if except_table_regexes.is_some() {
                                return Err(de::Error::duplicate_field("except_table_regexes"));
                            }
                            except_table_regexes = Some(map.next_value()?);
                        }
                        _ => {
                            return Err(de::Error::unknown_field(
                                key,
                                &[
                                    "only_tables",
                                    "only_table_regexes",
                                    "except_tables",
                                    "except_table_regexes",
                                ],
                            ))
                        }
                    }
                }
                match (
                    only_tables,
                    only_table_regexes,
                    except_tables,
                    except_table_regexes,
                ) {
                    (Some(t), None, None, None) => Ok(Filtering::OnlyTables(t)),
                    (None, Some(t), None, None) => Ok(Filtering::OnlyTableRegexes(t)),
                    (None, None, Some(t), None) => Ok(Filtering::ExceptTables(t)),
                    (None, None, None, Some(t)) => Ok(Filtering::ExceptTableRegexes(t)),
                    (None, None, None, None) => Ok(Filtering::None),
                    _ => Err(de::Error::duplicate_field(
                        "only_tables, only_table_regexes, except_tables, except_table_regexes",
                    )),
                }
            }
        }

        deserializer.deserialize_map(FilteringVisitor)
    }
}
<commit_msg>use absolute path for backward compatibility<commit_after>use crate::config;

use crate::infer_schema_internals::*;
use serde::de::{self, MapAccess, Visitor};
use serde::{Deserialize, Deserializer};
use serde_regex::Serde as RegexWrapper;
use std::error::Error;
use std::fmt::{self, Display, Formatter, Write};
use std::io::Write as IoWrite;

const SCHEMA_HEADER: &str = "// @generated automatically by Diesel CLI.\n";

type Regex = RegexWrapper<::regex::Regex>;

pub enum Filtering {
    OnlyTables(Vec<TableName>),
    OnlyTableRegexes(Vec<Regex>),
    ExceptTables(Vec<TableName>),
    ExceptTableRegexes(Vec<Regex>),
    None,
}

impl Default for Filtering {
    fn default() -> Self {
        Filtering::None
    }
}

impl Filtering {
    pub fn should_ignore_table(&self, name: &TableName) -> bool {
        use self::Filtering::*;

        match *self {
            OnlyTables(ref names) => !names.contains(name),
            OnlyTableRegexes(ref regexes) => {
                !regexes.iter().any(|regex| regex.is_match(&name.sql_name))
            }
            ExceptTables(ref names) => names.contains(name),
            ExceptTableRegexes(ref regexes) => {
                regexes.iter().any(|regex| regex.is_match(&name.sql_name))
            }
            None => false,
        }
    }
}

pub fn run_print_schema<W: IoWrite>(
    database_url: &str,
    config: &config::PrintSchema,
    output: &mut W,
) -> Result<(), Box<dyn Error>> {
    let schema = output_schema(database_url, config)?;

    output.write_all(schema.as_bytes())?;
    Ok(())
}

pub fn output_schema(
    database_url: &str,
    config: &config::PrintSchema,
) -> Result<String, Box<dyn Error>> {
    let table_names = load_table_names(database_url, config.schema_name())?
        .into_iter()
        .filter(|t| !config.filter.should_ignore_table(t))
        .collect::<Vec<_>>();
    let foreign_keys = load_foreign_key_constraints(database_url, config.schema_name())?;
    let foreign_keys =
        remove_unsafe_foreign_keys_for_codegen(database_url, &foreign_keys, &table_names);
    let table_data = table_names
        .into_iter()
        .map(|t| load_table_data(database_url, t))
        .collect::<Result<_, Box<dyn Error>>>()?;
    let definitions = TableDefinitions {
        tables: table_data,
        fk_constraints: foreign_keys,
        include_docs: config.with_docs,
        import_types: config.import_types(),
    };

    let mut out = String::new();
    writeln!(out, "{}", SCHEMA_HEADER)?;

    if let Some(schema_name) = config.schema_name() {
        write!(out, "{}", ModuleDefinition(schema_name, definitions))?;
    } else {
        write!(out, "{}", definitions)?;
    }

    if let Some(ref patch_file) = config.patch_file {
        let patch = std::fs::read_to_string(patch_file)?;
        let patch = diffy::Patch::from_str(&patch)?;

        out = diffy::apply(&out, &patch)?;
    }

    Ok(out)
}

struct ModuleDefinition<'a>(&'a str, TableDefinitions<'a>);

impl<'a> Display for ModuleDefinition<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        {
            let mut out = PadAdapter::new(f);
            writeln!(out, "pub mod {} {{", self.0)?;
            write!(out, "{}", self.1)?;
        }
        writeln!(f, "}}")?;
        Ok(())
    }
}

struct TableDefinitions<'a> {
    tables: Vec<TableData>,
    fk_constraints: Vec<ForeignKeyConstraint>,
    include_docs: bool,
    import_types: Option<&'a [String]>,
}

impl<'a> Display for TableDefinitions<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let mut is_first = true;
        for table in &self.tables {
            if is_first {
                is_first = false;
            } else {
                writeln!(f)?;
            }
            writeln!(
                f,
                "{}",
                TableDefinition {
                    table,
                    include_docs: self.include_docs,
                    import_types: self.import_types,
                }
            )?;
        }

        if !self.fk_constraints.is_empty() {
            writeln!(f)?;
        }

        for foreign_key in &self.fk_constraints {
            writeln!(f, "{}", Joinable(foreign_key))?;
        }

        if self.tables.len() > 1 {
            write!(f, "\ndiesel::allow_tables_to_appear_in_same_query!(")?;
            {
                let mut out = PadAdapter::new(f);
                writeln!(out)?;
                for table in &self.tables {
                    if table.name.rust_name == table.name.sql_name {
                        writeln!(out, "{},", table.name.sql_name)?;
                    } else {
                        writeln!(out, "{},", table.name.rust_name)?;
                    }
                }
            }
            writeln!(f, ");")?;
        }

        Ok(())
    }
}

struct TableDefinition<'a> {
    table: &'a TableData,
    import_types: Option<&'a [String]>,
    include_docs: bool,
}

impl<'a> Display for TableDefinition<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "diesel::table! {{")?;
        {
            let mut out = PadAdapter::new(f);
            writeln!(out)?;

            if let Some(types) = self.import_types {
                for import in types {
                    writeln!(out, "use {};", import)?;
                }
                writeln!(out)?;
            }

            if self.include_docs {
                for d in self.table.docs.lines() {
                    writeln!(out, "///{}{}", if d.is_empty() { "" } else { " " }, d)?;
                }
            }

            if self.table.name.rust_name != self.table.name.sql_name {
                writeln!(
                    out,
                    r#"#[sql_name = "{}"]"#,
                    self.table.name.full_sql_name()
                )?;
            }

            write!(out, "{} (", self.table.name)?;

            for (i, pk) in self.table.primary_key.iter().enumerate() {
                if i != 0 {
                    write!(out, ", ")?;
                }
                write!(out, "{}", pk)?;
            }

            write!(
                out,
                ") {}",
                ColumnDefinitions {
                    columns: &self.table.column_data,
                    include_docs: self.include_docs,
                }
            )?;
        }
        write!(f, "}}")?;
        Ok(())
    }
}

struct ColumnDefinitions<'a> {
    columns: &'a [ColumnDefinition],
    include_docs: bool,
}

impl<'a> Display for ColumnDefinitions<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        {
            let mut out = PadAdapter::new(f);
            writeln!(out, "{{")?;
            for column in self.columns {
                if self.include_docs {
                    for d in column.docs.lines() {
                        writeln!(out, "///{}{}", if d.is_empty() { "" } else { " " }, d)?;
                    }
                }
                if column.rust_name == column.sql_name {
                    writeln!(out, "{} -> {},", column.sql_name, column.ty)?;
                } else {
                    writeln!(out, r#"#[sql_name = "{}"]"#, column.sql_name)?;
                    writeln!(out, "{} -> {},", column.rust_name, column.ty)?;
                }
            }
        }
        writeln!(f, "}}")?;
        Ok(())
    }
}

struct Joinable<'a>(&'a ForeignKeyConstraint);

impl<'a> Display for Joinable<'a> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let child_table_name = &self.0.child_table.rust_name;

        let parent_table_name = &self.0.parent_table.rust_name;

        write!(
            f,
            "diesel::joinable!({} -> {} ({}));",
            child_table_name, parent_table_name, self.0.foreign_key_rust_name,
        )
    }
}

/// Lifted directly from libcore/fmt/builders.rs
struct PadAdapter<'a, 'b: 'a> {
    fmt: &'a mut Formatter<'b>,
    on_newline: bool,
}

impl<'a, 'b: 'a> PadAdapter<'a, 'b> {
    fn new(fmt: &'a mut Formatter<'b>) -> PadAdapter<'a, 'b> {
        PadAdapter {
            fmt,
            on_newline: false,
        }
    }
}

impl<'a, 'b: 'a> Write for PadAdapter<'a, 'b> {
    fn write_str(&mut self, mut s: &str) -> fmt::Result {
        while !s.is_empty() {
            let on_newline = self.on_newline;

            let split = match s.find('\n') {
                Some(pos) => {
                    self.on_newline = true;
                    pos + 1
                }
                None => {
                    self.on_newline = false;
                    s.len()
                }
            };

            let to_write = &s[..split];
            if on_newline && to_write != "\n" {
                self.fmt.write_str("    ")?;
            }
            self.fmt.write_str(to_write)?;

            s = &s[split..];
        }

        Ok(())
    }
}

impl<'de> Deserialize<'de> for Filtering {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct FilteringVisitor;

        impl<'de> Visitor<'de> for FilteringVisitor {
            type Value = Filtering;

            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
                f.write_str("either only_tables, only_table_regexes, except_tables, or except_table_regexes")
            }

            fn visit_map<V>(self, mut map: V) -> Result<Self::Value, V::Error>
            where
                V: MapAccess<'de>,
            {
                let mut only_tables = None::<Vec<TableName>>;
                let mut only_table_regexes = None::<Vec<Regex>>;
                let mut except_tables = None::<Vec<TableName>>;
                let mut except_table_regexes = None::<Vec<Regex>>;
                while let Some(key) = map.next_key()? {
                    match key {
                        "only_tables" => {
                            if only_tables.is_some() {
                                return Err(de::Error::duplicate_field("only_tables"));
                            }
                            only_tables = Some(map.next_value()?);
                        }
                        "only_table_regexes" => {
                            if only_table_regexes.is_some() {
                                return Err(de::Error::duplicate_field("only_table_regexes"));
                            }
                            only_table_regexes = Some(map.next_value()?);
                        }
                        "except_tables" => {
                            if except_tables.is_some() {
                                return Err(de::Error::duplicate_field("except_tables"));
                            }
                            except_tables = Some(map.next_value()?);
                        }
                        "except_table_regexes" => {
                            if except_table_regexes.is_some() {
                                return Err(de::Error::duplicate_field("except_table_regexes"));
                            }
                            except_table_regexes = Some(map.next_value()?);
                        }
                        _ => {
                            return Err(de::Error::unknown_field(
                                key,
                                &[
                                    "only_tables",
                                    "only_table_regexes",
                                    "except_tables",
                                    "except_table_regexes",
                                ],
                            ))
                        }
                    }
                }
                match (
                    only_tables,
                    only_table_regexes,
                    except_tables,
                    except_table_regexes,
                ) {
                    (Some(t), None, None, None) => Ok(Filtering::OnlyTables(t)),
                    (None, Some(t), None, None) => Ok(Filtering::OnlyTableRegexes(t)),
                    (None, None, Some(t), None) => Ok(Filtering::ExceptTables(t)),
                    (None, None, None, Some(t)) => Ok(Filtering::ExceptTableRegexes(t)),
                    (None, None, None, None) => Ok(Filtering::None),
                    _ => Err(de::Error::duplicate_field(
                        "only_tables, only_table_regexes, except_tables, except_table_regexes",
                    )),
                }
            }
        }

        deserializer.deserialize_map(FilteringVisitor)
    }
}
<|endoftext|>


<commit_before>#![no_main]

use libfuzzer_sys::fuzz_target;
use serde_json::{from_slice, Value};

fuzz_target!(|data: &[u8]| {
    let _ = from_slice::<Value>(data);
});
<commit_msg>Ignore using destructuring assignment to wildcard pattern<commit_after>#![no_main]

use libfuzzer_sys::fuzz_target;
use serde_json::{from_slice, Value};

fuzz_target!(|data: &[u8]| {
    _ = from_slice::<Value>(data);
});
<|endoftext|>


<commit_before>use crate::io::AsyncRead;

use std::io;
use std::ops::DerefMut;
use std::pin::Pin;
use std::task::{Context, Poll};

/// Reads bytes asynchronously.
///
/// This trait inherits from [`std::io::BufRead`] and indicates that an I/O object is
/// **non-blocking**. All non-blocking I/O objects must return an error when
/// bytes are unavailable instead of blocking the current thread.
///
/// Utilities for working with `AsyncBufRead` values are provided by
/// [`AsyncBufReadExt`].
///
/// [`std::io::BufRead`]: std::io::BufRead
/// [`AsyncBufReadExt`]: crate::io::AsyncBufReadExt
pub trait AsyncBufRead: AsyncRead {
    /// Attempts to return the contents of the internal buffer, filling it with more data
    /// from the inner reader if it is empty.
    ///
    /// On success, returns `Poll::Ready(Ok(buf))`.
    ///
    /// If no data is available for reading, the method returns
    /// `Poll::Pending` and arranges for the current task (via
    /// `cx.waker().wake_by_ref()`) to receive a notification when the object becomes
    /// readable or is closed.
    ///
    /// This function is a lower-level call. It needs to be paired with the
    /// [`consume`] method to function properly. When calling this
    /// method, none of the contents will be "read" in the sense that later
    /// calling [`poll_read`] may return the same contents. As such, [`consume`] must
    /// be called with the number of bytes that are consumed from this buffer to
    /// ensure that the bytes are never returned twice.
    ///
    /// An empty buffer returned indicates that the stream has reached EOF.
    ///
    /// [`poll_read`]: AsyncRead::poll_read
    /// [`consume`]: AsyncBufRead::consume
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>>;

    /// Tells this buffer that `amt` bytes have been consumed from the buffer,
    /// so they should no longer be returned in calls to [`poll_read`].
    ///
    /// This function is a lower-level call. It needs to be paired with the
    /// [`poll_fill_buf`] method to function properly. This function does
    /// not perform any I/O, it simply informs this object that some amount of
    /// its buffer, returned from [`poll_fill_buf`], has been consumed and should
    /// no longer be returned. As such, this function may do odd things if
    /// [`poll_fill_buf`] isn't called before calling it.
    ///
    /// The `amt` must be `<=` the number of bytes in the buffer returned by
    /// [`poll_fill_buf`].
    ///
    /// [`poll_read`]: AsyncRead::poll_read
    /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf
    fn consume(self: Pin<&mut Self>, amt: usize);
}

macro_rules! deref_async_buf_read {
    () => {
        fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
            Pin::new(&mut **self.get_mut()).poll_fill_buf(cx)
        }

        fn consume(mut self: Pin<&mut Self>, amt: usize) {
            Pin::new(&mut **self).consume(amt)
        }
    };
}

impl<T: ?Sized + AsyncBufRead + Unpin> AsyncBufRead for Box<T> {
    deref_async_buf_read!();
}

impl<T: ?Sized + AsyncBufRead + Unpin> AsyncBufRead for &mut T {
    deref_async_buf_read!();
}

impl
 AsyncBufRead for Pin

where
    P: DerefMut + Unpin,
    P::Target: AsyncBufRead,
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        self.get_mut().as_mut().poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        self.get_mut().as_mut().consume(amt)
    }
}

impl AsyncBufRead for &[u8] {
    fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        Poll::Ready(Ok(*self))
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        *self = &self[amt..];
    }
}

impl<T: AsRef<[u8]> + Unpin> AsyncBufRead for io::Cursor<T> {
    fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        Poll::Ready(io::BufRead::fill_buf(self.get_mut()))
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        io::BufRead::consume(self.get_mut(), amt)
    }
}
<commit_msg>io: update AsyncBufRead documentation (#2564)<commit_after>use crate::io::AsyncRead;

use std::io;
use std::ops::DerefMut;
use std::pin::Pin;
use std::task::{Context, Poll};

/// Reads bytes asynchronously.
///
/// This trait is analogous to [`std::io::BufRead`], but integrates with
/// the asynchronous task system. In particular, the [`poll_fill_buf`] method,
/// unlike [`BufRead::fill_buf`], will automatically queue the current task for wakeup
/// and return if data is not yet available, rather than blocking the calling
/// thread.
///
/// Utilities for working with `AsyncBufRead` values are provided by
/// [`AsyncBufReadExt`].
///
/// [`std::io::BufRead`]: std::io::BufRead
/// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf
/// [`BufRead::fill_buf`]: std::io::BufRead::fill_buf
/// [`AsyncBufReadExt`]: crate::io::AsyncBufReadExt
pub trait AsyncBufRead: AsyncRead {
    /// Attempts to return the contents of the internal buffer, filling it with more data
    /// from the inner reader if it is empty.
    ///
    /// On success, returns `Poll::Ready(Ok(buf))`.
    ///
    /// If no data is available for reading, the method returns
    /// `Poll::Pending` and arranges for the current task (via
    /// `cx.waker().wake_by_ref()`) to receive a notification when the object becomes
    /// readable or is closed.
    ///
    /// This function is a lower-level call. It needs to be paired with the
    /// [`consume`] method to function properly. When calling this
    /// method, none of the contents will be "read" in the sense that later
    /// calling [`poll_read`] may return the same contents. As such, [`consume`] must
    /// be called with the number of bytes that are consumed from this buffer to
    /// ensure that the bytes are never returned twice.
    ///
    /// An empty buffer returned indicates that the stream has reached EOF.
    ///
    /// [`poll_read`]: AsyncRead::poll_read
    /// [`consume`]: AsyncBufRead::consume
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>>;

    /// Tells this buffer that `amt` bytes have been consumed from the buffer,
    /// so they should no longer be returned in calls to [`poll_read`].
    ///
    /// This function is a lower-level call. It needs to be paired with the
    /// [`poll_fill_buf`] method to function properly. This function does
    /// not perform any I/O, it simply informs this object that some amount of
    /// its buffer, returned from [`poll_fill_buf`], has been consumed and should
    /// no longer be returned. As such, this function may do odd things if
    /// [`poll_fill_buf`] isn't called before calling it.
    ///
    /// The `amt` must be `<=` the number of bytes in the buffer returned by
    /// [`poll_fill_buf`].
    ///
    /// [`poll_read`]: AsyncRead::poll_read
    /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf
    fn consume(self: Pin<&mut Self>, amt: usize);
}

macro_rules! deref_async_buf_read {
    () => {
        fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
            Pin::new(&mut **self.get_mut()).poll_fill_buf(cx)
        }

        fn consume(mut self: Pin<&mut Self>, amt: usize) {
            Pin::new(&mut **self).consume(amt)
        }
    };
}

impl<T: ?Sized + AsyncBufRead + Unpin> AsyncBufRead for Box<T> {
    deref_async_buf_read!();
}

impl<T: ?Sized + AsyncBufRead + Unpin> AsyncBufRead for &mut T {
    deref_async_buf_read!();
}

impl
 AsyncBufRead for Pin

where
    P: DerefMut + Unpin,
    P::Target: AsyncBufRead,
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        self.get_mut().as_mut().poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        self.get_mut().as_mut().consume(amt)
    }
}

impl AsyncBufRead for &[u8] {
    fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        Poll::Ready(Ok(*self))
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        *self = &self[amt..];
    }
}

impl<T: AsRef<[u8]> + Unpin> AsyncBufRead for io::Cursor<T> {
    fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        Poll::Ready(io::BufRead::fill_buf(self.get_mut()))
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        io::BufRead::consume(self.get_mut(), amt)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This file defines transaction store APIs that are related to committed signed transactions.

use crate::transaction_accumulator::TransactionAccumulatorSchema;
use crate::{
    errors::AptosDbError,
    schema::{
        transaction::TransactionSchema, transaction_by_account::TransactionByAccountSchema,
        transaction_by_hash::TransactionByHashSchema, write_set::WriteSetSchema,
    },
    transaction_info::TransactionInfoSchema,
};
use anyhow::{ensure, format_err, Result};
use aptos_crypto::{hash::CryptoHash, HashValue};
use aptos_types::{
    account_address::AccountAddress,
    proof::position::Position,
    transaction::{Transaction, Version},
    write_set::WriteSet,
};
use schemadb::{ReadOptions, SchemaBatch, SchemaIterator, DB};
use std::sync::Arc;

#[derive(Clone, Debug)]
pub struct TransactionStore {
    db: Arc<DB>,
}

impl TransactionStore {
    pub fn new(db: Arc<DB>) -> Self {
        Self { db }
    }

    /// Gets the version of a transaction by the sender `address` and `sequence_number`.
    pub fn get_account_transaction_version(
        &self,
        address: AccountAddress,
        sequence_number: u64,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        if let Some(version) = self
            .db
            .get::<TransactionByAccountSchema>(&(address, sequence_number))?
        {
            if version <= ledger_version {
                return Ok(Some(version));
            }
        }

        Ok(None)
    }

    /// Gets the version of a transaction by its hash.
    pub fn get_transaction_version_by_hash(
        &self,
        hash: &HashValue,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        Ok(match self.db.get::<TransactionByHashSchema>(hash)? {
            Some(version) if version <= ledger_version => Some(version),
            _ => None,
        })
    }

    /// Gets an iterator that yields `(sequence_number, version)` for each
    /// transaction sent by an account, with minimum sequence number greater
    /// `min_seq_num`, and returning at most `num_versions` results with
    /// `version <= ledger_version`.
    /// Guarantees that the returned sequence numbers are sequential, i.e.,
    /// `seq_num_{i} + 1 = seq_num_{i+1}`.
    pub fn get_account_transaction_version_iter(
        &self,
        address: AccountAddress,
        min_seq_num: u64,
        num_versions: u64,
        ledger_version: Version,
    ) -> Result<AccountTransactionVersionIter> {
        let mut iter = self
            .db
            .iter::<TransactionByAccountSchema>(ReadOptions::default())?;
        iter.seek(&(address, min_seq_num))?;
        Ok(AccountTransactionVersionIter {
            inner: iter,
            address,
            expected_next_seq_num: None,
            end_seq_num: min_seq_num
                .checked_add(num_versions)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
            prev_version: None,
            ledger_version,
        })
    }

    /// Get signed transaction given `version`
    pub fn get_transaction(&self, version: Version) -> Result<Transaction> {
        self.db
            .get::<TransactionSchema>(&version)?
            .ok_or_else(|| AptosDbError::NotFound(format!("Txn {}", version)).into())
    }

    /// Gets an iterator that yields `num_transactions` transactions starting from `start_version`.
    pub fn get_transaction_iter(
        &self,
        start_version: Version,
        num_transactions: usize,
    ) -> Result<TransactionIter> {
        let mut iter = self.db.iter::<TransactionSchema>(ReadOptions::default())?;
        iter.seek(&start_version)?;
        Ok(TransactionIter {
            inner: iter,
            expected_next_version: start_version,
            end_version: start_version
                .checked_add(num_transactions as u64)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
        })
    }

    /// Searches around the version to find the block's transactions
    pub fn get_block_boundaries(
        &self,
        version: Version,
        latest_ledger_version: Version,
    ) -> Result<(Version, Version)> {
        // Must be larger than a block size, otherwise a NotFound error will be raised wrongly.
        const MAX_VERSIONS_TO_SEARCH: usize = 1000 * 100;

        // Genesis is always 0,0
        if version == 0 {
            return Ok((0, 0));
        }

        // Linear searches via `DB::iter()` and `DB::rev_iter()` here, NOT expecting performance hit, due to the fact
        // that the iterator caches data block and that there are limited number of transactions in
        // each block.
        let mut iter = self.db.rev_iter::<TransactionSchema>(Default::default())?;
        iter.seek(&version)?;
        let mut start_version = None;
        for res in iter.take(MAX_VERSIONS_TO_SEARCH) {
            let (v, txn) = res?;
            // If we've found the beginning of the block, stop
            if matches!(
                txn,
                Transaction::GenesisTransaction(_) | Transaction::BlockMetadata(_)
            ) {
                start_version = Some(v);
                break;
            }
        }

        let mut iter = self.db.iter::<TransactionSchema>(Default::default())?;
        iter.seek(&version)?;
        let mut end_version = None;
        for res in iter.take(MAX_VERSIONS_TO_SEARCH) {
            let (v, txn) = res?;
            match txn {
                Transaction::BlockMetadata(_) => {
                    // If the current version is the beginning of the block, we need to find the
                    // rest of the block
                    if start_version != Some(v) {
                        // This block metadata is the next block, so you need the previous version
                        end_version = Some(v - 1);
                        break;
                    }
                }
                // Every txn up to the end or the block metadata should be included in the block
                _ => {
                    if v == latest_ledger_version {
                        end_version = Some(v);
                        break;
                    }
                }
            }
        }

        // If we've found both, we have the whole block
        match (start_version, end_version) {
            (Some(start), Some(end)) => Ok((start, end)),
            // If we didn't find the end, it's a single transaction block (reconfig)
            (Some(start), None) => Ok((start, start)),
            _ => Err(AptosDbError::NotFound(format!(
                "Block boundaries not found for version {}",
                version
            ))
            .into()),
        }
    }

    /// Save signed transaction at `version`
    pub fn put_transaction(
        &self,
        version: Version,
        transaction: &Transaction,
        batch: &mut SchemaBatch,
    ) -> Result<()> {
        if let Transaction::UserTransaction(txn) = transaction {
            batch.put::<TransactionByAccountSchema>(
                &(txn.sender(), txn.sequence_number()),
                &version,
            )?;
        }
        batch.put::<TransactionByHashSchema>(&transaction.hash(), &version)?;
        batch.put::<TransactionSchema>(&version, transaction)?;

        Ok(())
    }

    /// Get executed transaction vm output given `version`
    pub fn get_write_set(&self, version: Version) -> Result<WriteSet> {
        self.db.get::<WriteSetSchema>(&version)?.ok_or_else(|| {
            AptosDbError::NotFound(format!("WriteSet at version {}", version)).into()
        })
    }

    /// Get write sets in `[begin_version, end_version)` half-open range.
    ///
    /// N.b. an empty `Vec` is returned when `begin_version == end_version`
    pub fn get_write_sets(
        &self,
        begin_version: Version,
        end_version: Version,
    ) -> Result<Vec<WriteSet>> {
        if begin_version == end_version {
            return Ok(Vec::new());
        }
        ensure!(
            begin_version < end_version,
            "begin_version {} >= end_version {}",
            begin_version,
            end_version
        );

        let mut iter = self.db.iter::<WriteSetSchema>(Default::default())?;
        iter.seek(&begin_version)?;

        let mut ret = Vec::with_capacity((end_version - begin_version) as usize);
        for current_version in begin_version..end_version {
            let (version, write_set) = iter
                .next()
                .transpose()?
                .ok_or_else(|| format_err!("Write set missing for version {}", current_version))?;
            ensure!(
                version == current_version,
                "Write set missing for version {}, got version {}",
                current_version,
                version,
            );
            ret.push(write_set);
        }

        Ok(ret)
    }

    /// Save executed transaction vm output given `version`
    pub fn put_write_set(
        &self,
        version: Version,
        write_set: &WriteSet,
        batch: &mut SchemaBatch,
    ) -> Result<()> {
        batch.put::<WriteSetSchema>(&version, write_set)
    }

    /// Prune the transaction by hash store given a list of transaction
    pub fn prune_transaction_by_hash(
        &self,
        transactions: &[Transaction],
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for transaction in transactions {
            db_batch.delete::<TransactionByHashSchema>(&transaction.hash())?;
        }
        Ok(())
    }

    /// Prune the transaction by account store given a list of transaction
    pub fn prune_transaction_by_account(
        &self,
        transactions: &[Transaction],
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for transaction in transactions {
            if let Transaction::UserTransaction(txn) = transaction {
                db_batch
                    .delete::<TransactionByAccountSchema>(&(txn.sender(), txn.sequence_number()))?;
            }
        }
        Ok(())
    }

    /// Prune the transaction schema store between a range of version in [begin, end)
    pub fn prune_transaction_schema(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version in begin..end {
            db_batch.delete::<TransactionSchema>(&version)?;
        }
        Ok(())
    }

    /// Prune the transaction schema store between a range of version in [begin, end)
    pub fn prune_transaction_info_schema(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version in begin..end {
            db_batch.delete::<TransactionInfoSchema>(&version)?;
        }
        Ok(())
    }

    /// Prune the transaction accumulator between a range of version in [begin, end).
    ///
    /// To avoid always pruning a full left subtree, we uses the following algorithm.
    /// For each leaf with an odd leaf index.
    /// 1. From the bottom upwards, find the first ancestor that's a left child of its parent.
    /// (the position of which can be got by popping "1"s from the right of the leaf address).
    /// Note that this node DOES NOT become non-useful.
    /// 2. From the node found from the previous step, delete both its children non-useful, and go
    /// to the right child to repeat the process until we reach a leaf node.
    /// More details are in this issue https://github.com/aptos-labs/aptos-core/issues/1288.
    pub fn prune_transaction_accumulator(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version_to_delete in begin..end {
            // The even version will be pruned in the iteration of version + 1.
            if version_to_delete % 2 == 0 {
                continue;
            }

            let first_ancestor_that_is_a_left_child =
                self.find_first_ancestor_that_is_a_left_child(version_to_delete);

            // This assertion is true because we skip the leaf nodes with address which is a
            // a multiple of 2.
            assert!(!first_ancestor_that_is_a_left_child.is_leaf());

            let mut current = first_ancestor_that_is_a_left_child;
            while !current.is_leaf() {
                db_batch.delete::<TransactionAccumulatorSchema>(&current.left_child())?;
                db_batch.delete::<TransactionAccumulatorSchema>(&current.right_child())?;
                current = current.right_child();
            }
        }
        Ok(())
    }

    /// Finds the first ancestor that is a child of its parent.
    fn find_first_ancestor_that_is_a_left_child(&self, version: Version) -> Position {
        // We can get the first ancestor's position based on the two observations:
        // - floor(level position of a node / 2) = level position of its parent.
        // - if a node is a left child of its parent, its level position should be a multiple of 2.
        // - level position means the position counted from 0 of a single tree level. For example,
        //                a (level position = 0)
        //         /                                \
        //    b (level position = 0)      c(level position = 1)
        //
        // To find the first ancestor which is a left child of its parent, we can keep diving the
        // version by 2 (to find the ancestor) until we get a number which is a multiple of 2
        // (to make sure the ancestor is a left child of its parent). The number of time we
        // divide the version is the level of the ancestor. The remainder is the level position
        // of the ancestor.
        let first_ancestor_that_is_a_left_child_level = version.trailing_ones();
        let index_in_level = version >> first_ancestor_that_is_a_left_child_level;
        Position::from_level_and_pos(first_ancestor_that_is_a_left_child_level, index_in_level)
    }

    /// Prune the transaction schema store between a range of version in [begin, end)
    pub fn prune_write_set(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version in begin..end {
            db_batch.delete::<WriteSetSchema>(&version)?;
        }
        Ok(())
    }
}

pub struct TransactionIter<'a> {
    inner: SchemaIterator<'a, TransactionSchema>,
    expected_next_version: Version,
    end_version: Version,
}

impl<'a> TransactionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<Transaction>> {
        if self.expected_next_version >= self.end_version {
            return Ok(None);
        }

        let ret = match self.inner.next().transpose()? {
            Some((version, transaction)) => {
                ensure!(
                    version == self.expected_next_version,
                    "Transaction versions are not consecutive.",
                );
                self.expected_next_version += 1;
                Some(transaction)
            }
            None => None,
        };

        Ok(ret)
    }
}

impl<'a> Iterator for TransactionIter<'a> {
    type Item = Result<Transaction>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

pub struct AccountTransactionVersionIter<'a> {
    inner: SchemaIterator<'a, TransactionByAccountSchema>,
    address: AccountAddress,
    expected_next_seq_num: Option<u64>,
    end_seq_num: u64,
    prev_version: Option<Version>,
    ledger_version: Version,
}

impl<'a> AccountTransactionVersionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<(u64, Version)>> {
        Ok(match self.inner.next().transpose()? {
            Some(((address, seq_num), version)) => {
                // No more transactions sent by this account.
                if address != self.address {
                    return Ok(None);
                }
                if seq_num >= self.end_seq_num {
                    return Ok(None);
                }

                // Ensure seq_num_{i+1} == seq_num_{i} + 1
                if let Some(expected_seq_num) = self.expected_next_seq_num {
                    ensure!(
                        seq_num == expected_seq_num,
                        "DB corruption: account transactions sequence numbers are not contiguous: \
                     actual: {}, expected: {}",
                        seq_num,
                        expected_seq_num,
                    );
                };

                // Ensure version_{i+1} > version_{i}
                if let Some(prev_version) = self.prev_version {
                    ensure!(
                        prev_version < version,
                        "DB corruption: account transaction versions are not strictly increasing: \
                         previous version: {}, current version: {}",
                        prev_version,
                        version,
                    );
                }

                // No more transactions (in this view of the ledger).
                if version > self.ledger_version {
                    return Ok(None);
                }

                self.expected_next_seq_num = Some(seq_num + 1);
                self.prev_version = Some(version);
                Some((seq_num, version))
            }
            None => None,
        })
    }
}

impl<'a> Iterator for AccountTransactionVersionIter<'a> {
    type Item = Result<(u64, Version)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

#[cfg(test)]
mod test;
<commit_msg>[storage] trivial: remove get_block_boundaries<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This file defines transaction store APIs that are related to committed signed transactions.

use crate::transaction_accumulator::TransactionAccumulatorSchema;
use crate::{
    errors::AptosDbError,
    schema::{
        transaction::TransactionSchema, transaction_by_account::TransactionByAccountSchema,
        transaction_by_hash::TransactionByHashSchema, write_set::WriteSetSchema,
    },
    transaction_info::TransactionInfoSchema,
};
use anyhow::{ensure, format_err, Result};
use aptos_crypto::{hash::CryptoHash, HashValue};
use aptos_types::{
    account_address::AccountAddress,
    proof::position::Position,
    transaction::{Transaction, Version},
    write_set::WriteSet,
};
use schemadb::{ReadOptions, SchemaBatch, SchemaIterator, DB};
use std::sync::Arc;

#[derive(Clone, Debug)]
pub struct TransactionStore {
    db: Arc<DB>,
}

impl TransactionStore {
    pub fn new(db: Arc<DB>) -> Self {
        Self { db }
    }

    /// Gets the version of a transaction by the sender `address` and `sequence_number`.
    pub fn get_account_transaction_version(
        &self,
        address: AccountAddress,
        sequence_number: u64,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        if let Some(version) = self
            .db
            .get::<TransactionByAccountSchema>(&(address, sequence_number))?
        {
            if version <= ledger_version {
                return Ok(Some(version));
            }
        }

        Ok(None)
    }

    /// Gets the version of a transaction by its hash.
    pub fn get_transaction_version_by_hash(
        &self,
        hash: &HashValue,
        ledger_version: Version,
    ) -> Result<Option<Version>> {
        Ok(match self.db.get::<TransactionByHashSchema>(hash)? {
            Some(version) if version <= ledger_version => Some(version),
            _ => None,
        })
    }

    /// Gets an iterator that yields `(sequence_number, version)` for each
    /// transaction sent by an account, with minimum sequence number greater
    /// `min_seq_num`, and returning at most `num_versions` results with
    /// `version <= ledger_version`.
    /// Guarantees that the returned sequence numbers are sequential, i.e.,
    /// `seq_num_{i} + 1 = seq_num_{i+1}`.
    pub fn get_account_transaction_version_iter(
        &self,
        address: AccountAddress,
        min_seq_num: u64,
        num_versions: u64,
        ledger_version: Version,
    ) -> Result<AccountTransactionVersionIter> {
        let mut iter = self
            .db
            .iter::<TransactionByAccountSchema>(ReadOptions::default())?;
        iter.seek(&(address, min_seq_num))?;
        Ok(AccountTransactionVersionIter {
            inner: iter,
            address,
            expected_next_seq_num: None,
            end_seq_num: min_seq_num
                .checked_add(num_versions)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
            prev_version: None,
            ledger_version,
        })
    }

    /// Get signed transaction given `version`
    pub fn get_transaction(&self, version: Version) -> Result<Transaction> {
        self.db
            .get::<TransactionSchema>(&version)?
            .ok_or_else(|| AptosDbError::NotFound(format!("Txn {}", version)).into())
    }

    /// Gets an iterator that yields `num_transactions` transactions starting from `start_version`.
    pub fn get_transaction_iter(
        &self,
        start_version: Version,
        num_transactions: usize,
    ) -> Result<TransactionIter> {
        let mut iter = self.db.iter::<TransactionSchema>(ReadOptions::default())?;
        iter.seek(&start_version)?;
        Ok(TransactionIter {
            inner: iter,
            expected_next_version: start_version,
            end_version: start_version
                .checked_add(num_transactions as u64)
                .ok_or_else(|| format_err!("too many transactions requested"))?,
        })
    }

    /// Save signed transaction at `version`
    pub fn put_transaction(
        &self,
        version: Version,
        transaction: &Transaction,
        batch: &mut SchemaBatch,
    ) -> Result<()> {
        if let Transaction::UserTransaction(txn) = transaction {
            batch.put::<TransactionByAccountSchema>(
                &(txn.sender(), txn.sequence_number()),
                &version,
            )?;
        }
        batch.put::<TransactionByHashSchema>(&transaction.hash(), &version)?;
        batch.put::<TransactionSchema>(&version, transaction)?;

        Ok(())
    }

    /// Get executed transaction vm output given `version`
    pub fn get_write_set(&self, version: Version) -> Result<WriteSet> {
        self.db.get::<WriteSetSchema>(&version)?.ok_or_else(|| {
            AptosDbError::NotFound(format!("WriteSet at version {}", version)).into()
        })
    }

    /// Get write sets in `[begin_version, end_version)` half-open range.
    ///
    /// N.b. an empty `Vec` is returned when `begin_version == end_version`
    pub fn get_write_sets(
        &self,
        begin_version: Version,
        end_version: Version,
    ) -> Result<Vec<WriteSet>> {
        if begin_version == end_version {
            return Ok(Vec::new());
        }
        ensure!(
            begin_version < end_version,
            "begin_version {} >= end_version {}",
            begin_version,
            end_version
        );

        let mut iter = self.db.iter::<WriteSetSchema>(Default::default())?;
        iter.seek(&begin_version)?;

        let mut ret = Vec::with_capacity((end_version - begin_version) as usize);
        for current_version in begin_version..end_version {
            let (version, write_set) = iter
                .next()
                .transpose()?
                .ok_or_else(|| format_err!("Write set missing for version {}", current_version))?;
            ensure!(
                version == current_version,
                "Write set missing for version {}, got version {}",
                current_version,
                version,
            );
            ret.push(write_set);
        }

        Ok(ret)
    }

    /// Save executed transaction vm output given `version`
    pub fn put_write_set(
        &self,
        version: Version,
        write_set: &WriteSet,
        batch: &mut SchemaBatch,
    ) -> Result<()> {
        batch.put::<WriteSetSchema>(&version, write_set)
    }

    /// Prune the transaction by hash store given a list of transaction
    pub fn prune_transaction_by_hash(
        &self,
        transactions: &[Transaction],
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for transaction in transactions {
            db_batch.delete::<TransactionByHashSchema>(&transaction.hash())?;
        }
        Ok(())
    }

    /// Prune the transaction by account store given a list of transaction
    pub fn prune_transaction_by_account(
        &self,
        transactions: &[Transaction],
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for transaction in transactions {
            if let Transaction::UserTransaction(txn) = transaction {
                db_batch
                    .delete::<TransactionByAccountSchema>(&(txn.sender(), txn.sequence_number()))?;
            }
        }
        Ok(())
    }

    /// Prune the transaction schema store between a range of version in [begin, end)
    pub fn prune_transaction_schema(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version in begin..end {
            db_batch.delete::<TransactionSchema>(&version)?;
        }
        Ok(())
    }

    /// Prune the transaction schema store between a range of version in [begin, end)
    pub fn prune_transaction_info_schema(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version in begin..end {
            db_batch.delete::<TransactionInfoSchema>(&version)?;
        }
        Ok(())
    }

    /// Prune the transaction accumulator between a range of version in [begin, end).
    ///
    /// To avoid always pruning a full left subtree, we uses the following algorithm.
    /// For each leaf with an odd leaf index.
    /// 1. From the bottom upwards, find the first ancestor that's a left child of its parent.
    /// (the position of which can be got by popping "1"s from the right of the leaf address).
    /// Note that this node DOES NOT become non-useful.
    /// 2. From the node found from the previous step, delete both its children non-useful, and go
    /// to the right child to repeat the process until we reach a leaf node.
    /// More details are in this issue https://github.com/aptos-labs/aptos-core/issues/1288.
    pub fn prune_transaction_accumulator(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version_to_delete in begin..end {
            // The even version will be pruned in the iteration of version + 1.
            if version_to_delete % 2 == 0 {
                continue;
            }

            let first_ancestor_that_is_a_left_child =
                self.find_first_ancestor_that_is_a_left_child(version_to_delete);

            // This assertion is true because we skip the leaf nodes with address which is a
            // a multiple of 2.
            assert!(!first_ancestor_that_is_a_left_child.is_leaf());

            let mut current = first_ancestor_that_is_a_left_child;
            while !current.is_leaf() {
                db_batch.delete::<TransactionAccumulatorSchema>(&current.left_child())?;
                db_batch.delete::<TransactionAccumulatorSchema>(&current.right_child())?;
                current = current.right_child();
            }
        }
        Ok(())
    }

    /// Finds the first ancestor that is a child of its parent.
    fn find_first_ancestor_that_is_a_left_child(&self, version: Version) -> Position {
        // We can get the first ancestor's position based on the two observations:
        // - floor(level position of a node / 2) = level position of its parent.
        // - if a node is a left child of its parent, its level position should be a multiple of 2.
        // - level position means the position counted from 0 of a single tree level. For example,
        //                a (level position = 0)
        //         /                                \
        //    b (level position = 0)      c(level position = 1)
        //
        // To find the first ancestor which is a left child of its parent, we can keep diving the
        // version by 2 (to find the ancestor) until we get a number which is a multiple of 2
        // (to make sure the ancestor is a left child of its parent). The number of time we
        // divide the version is the level of the ancestor. The remainder is the level position
        // of the ancestor.
        let first_ancestor_that_is_a_left_child_level = version.trailing_ones();
        let index_in_level = version >> first_ancestor_that_is_a_left_child_level;
        Position::from_level_and_pos(first_ancestor_that_is_a_left_child_level, index_in_level)
    }

    /// Prune the transaction schema store between a range of version in [begin, end)
    pub fn prune_write_set(
        &self,
        begin: Version,
        end: Version,
        db_batch: &mut SchemaBatch,
    ) -> anyhow::Result<()> {
        for version in begin..end {
            db_batch.delete::<WriteSetSchema>(&version)?;
        }
        Ok(())
    }
}

pub struct TransactionIter<'a> {
    inner: SchemaIterator<'a, TransactionSchema>,
    expected_next_version: Version,
    end_version: Version,
}

impl<'a> TransactionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<Transaction>> {
        if self.expected_next_version >= self.end_version {
            return Ok(None);
        }

        let ret = match self.inner.next().transpose()? {
            Some((version, transaction)) => {
                ensure!(
                    version == self.expected_next_version,
                    "Transaction versions are not consecutive.",
                );
                self.expected_next_version += 1;
                Some(transaction)
            }
            None => None,
        };

        Ok(ret)
    }
}

impl<'a> Iterator for TransactionIter<'a> {
    type Item = Result<Transaction>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

pub struct AccountTransactionVersionIter<'a> {
    inner: SchemaIterator<'a, TransactionByAccountSchema>,
    address: AccountAddress,
    expected_next_seq_num: Option<u64>,
    end_seq_num: u64,
    prev_version: Option<Version>,
    ledger_version: Version,
}

impl<'a> AccountTransactionVersionIter<'a> {
    fn next_impl(&mut self) -> Result<Option<(u64, Version)>> {
        Ok(match self.inner.next().transpose()? {
            Some(((address, seq_num), version)) => {
                // No more transactions sent by this account.
                if address != self.address {
                    return Ok(None);
                }
                if seq_num >= self.end_seq_num {
                    return Ok(None);
                }

                // Ensure seq_num_{i+1} == seq_num_{i} + 1
                if let Some(expected_seq_num) = self.expected_next_seq_num {
                    ensure!(
                        seq_num == expected_seq_num,
                        "DB corruption: account transactions sequence numbers are not contiguous: \
                     actual: {}, expected: {}",
                        seq_num,
                        expected_seq_num,
                    );
                };

                // Ensure version_{i+1} > version_{i}
                if let Some(prev_version) = self.prev_version {
                    ensure!(
                        prev_version < version,
                        "DB corruption: account transaction versions are not strictly increasing: \
                         previous version: {}, current version: {}",
                        prev_version,
                        version,
                    );
                }

                // No more transactions (in this view of the ledger).
                if version > self.ledger_version {
                    return Ok(None);
                }

                self.expected_next_seq_num = Some(seq_num + 1);
                self.prev_version = Some(version);
                Some((seq_num, version))
            }
            None => None,
        })
    }
}

impl<'a> Iterator for AccountTransactionVersionIter<'a> {
    type Item = Result<(u64, Version)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_impl().transpose()
    }
}

#[cfg(test)]
mod test;
<|endoftext|>


<commit_before>use colored::*;
use regex::Regex;
use std::path::{Path, PathBuf};
use std::{env, process::Command};
use ui_test::{color_eyre::Result, Config, Mode, OutputConflictHandling};

fn miri_path() -> PathBuf {
    PathBuf::from(option_env!("MIRI").unwrap_or(env!("CARGO_BIN_EXE_miri")))
}

// Build the shared object file for testing external C function calls.
fn build_so_for_c_ffi_tests() -> PathBuf {
    let cc = option_env!("CC").unwrap_or("cc");
    // Target directory that we can write to.
    let so_target_dir = Path::new(&env::var_os("CARGO_TARGET_DIR").unwrap()).join("miri-extern-so");
    // Create the directory if it does not already exist.
    std::fs::create_dir_all(&so_target_dir)
        .expect("Failed to create directory for shared object file");
    let so_file_path = so_target_dir.join("libtestlib.so");
    let cc_output = Command::new(cc)
        .args([
            "-shared",
            "-o",
            so_file_path.to_str().unwrap(),
            "tests/extern-so/test.c",
            // Only add the functions specified in libcode.version to the shared object file.
            // This is to avoid automatically adding `malloc`, etc.
            // Source: https://anadoxin.org/blog/control-over-symbol-exports-in-gcc.html/
            "-fPIC",
            "-Wl,--version-script=tests/extern-so/libcode.version",
        ])
        .output()
        .expect("failed to generate shared object file for testing external C function calls");
    if !cc_output.status.success() {
        panic!("error in generating shared object file for testing external C function calls");
    }
    so_file_path
}

fn run_tests(
    mode: Mode,
    path: &str,
    target: Option<String>,
    with_dependencies: bool,
) -> Result<()> {
    let mut config = Config {
        target,
        stderr_filters: STDERR.clone(),
        stdout_filters: STDOUT.clone(),
        root_dir: PathBuf::from(path),
        mode,
        program: miri_path(),
        quiet: false,
        ..Config::default()
    };

    let in_rustc_test_suite = option_env!("RUSTC_STAGE").is_some();

    // Add some flags we always want.
    config.args.push("--edition".into());
    config.args.push("2018".into());
    if in_rustc_test_suite {
        // Less aggressive warnings to make the rustc toolstate management less painful.
        // (We often get warnings when e.g. a feature gets stabilized or some lint gets added/improved.)
        config.args.push("-Astable-features".into());
        config.args.push("-Aunused".into());
    } else {
        config.args.push("-Dwarnings".into());
        config.args.push("-Dunused".into());
    }
    if let Ok(extra_flags) = env::var("MIRIFLAGS") {
        for flag in extra_flags.split_whitespace() {
            config.args.push(flag.into());
        }
    }
    config.args.push("-Zui-testing".into());
    if let Some(target) = &config.target {
        config.args.push("--target".into());
        config.args.push(target.into());
    }

    // If we're on linux, and we're testing the extern-so functionality,
    // then build the shared object file for testing external C function calls
    // and push the relevant compiler flag.
    if cfg!(target_os = "linux") && path.starts_with("tests/extern-so/") {
        let so_file_path = build_so_for_c_ffi_tests();
        let mut flag = std::ffi::OsString::from("-Zmiri-extern-so-file=");
        flag.push(so_file_path.into_os_string());
        config.args.push(flag);
    }

    let skip_ui_checks = env::var_os("MIRI_SKIP_UI_CHECKS").is_some();

    config.output_conflict_handling = match (env::var_os("MIRI_BLESS").is_some(), skip_ui_checks) {
        (false, false) => OutputConflictHandling::Error,
        (true, false) => OutputConflictHandling::Bless,
        (false, true) => OutputConflictHandling::Ignore,
        (true, true) => panic!("cannot use MIRI_BLESS and MIRI_SKIP_UI_CHECKS at the same time"),
    };

    // Handle command-line arguments.
    config.path_filter.extend(std::env::args().skip(1).filter(|arg| {
        match &**arg {
            "--quiet" => {
                config.quiet = true;
                false
            }
            _ => true,
        }
    }));

    let use_std = env::var_os("MIRI_NO_STD").is_none();

    if with_dependencies && use_std {
        config.dependencies_crate_manifest_path =
            Some(Path::new("test_dependencies").join("Cargo.toml"));
        config.dependency_builder.args = vec![
            "run".into(),
            "--manifest-path".into(),
            "cargo-miri/Cargo.toml".into(),
            "--".into(),
            "miri".into(),
            "run".into(), // There is no `cargo miri build` so we just use `cargo miri run`.
        ];
    }
    ui_test::run_tests(config)
}

macro_rules! regexes {
    ($name:ident: $($regex:expr => $replacement:expr,)*) => {lazy_static::lazy_static! {
        static ref $name: Vec<(Regex, &'static str)> = vec![
            $((Regex::new($regex).unwrap(), $replacement),)*
        ];
    }};
}

regexes! {
    STDOUT:
    // Windows file paths
    r"\\"                           => "/",
}

regexes! {
    STDERR:
    // erase line and column info
    r"\.rs:[0-9]+:[0-9]+(: [0-9]+:[0-9]+)?" => ".rs:LL:CC",
    // erase alloc ids
    "alloc[0-9]+"                    => "ALLOC",
    // erase Stacked Borrows tags
    "<[0-9]+>"                       => "<TAG>",
    // erase whitespace that differs between platforms
    r" +at (.*\.rs)"                 => " at $1",
    // erase generics in backtraces
    "([0-9]+: .*)::<.*>"             => "$1",
    // erase addresses in backtraces
    "([0-9]+: ) +0x[0-9a-f]+ - (.*)" => "$1$2",
    // erase long hexadecimals
    r"0x[0-9a-fA-F]+[0-9a-fA-F]{2,2}" => "$$HEX",
    // erase specific alignments
    "alignment [0-9]+"               => "alignment ALIGN",
    // erase thread caller ids
    r"call [0-9]+"                  => "call ID",
    // erase platform module paths
    "sys::[a-z]+::"                  => "sys::PLATFORM::",
    // Windows file paths
    r"\\"                           => "/",
    // erase Rust stdlib path
    "[^ `]*/(rust[^/]*|checkout)/library/" => "RUSTLIB/",
    // erase platform file paths
    "sys/[a-z]+/"                    => "sys/PLATFORM/",
    // erase paths into the crate registry
    r"[^ ]*/\.?cargo/registry/.*/(.*\.rs)"  => "CARGO_REGISTRY/.../$1",
}

enum Dependencies {
    WithDependencies,
    WithoutDependencies,
}

use Dependencies::*;

fn ui(mode: Mode, path: &str, with_dependencies: Dependencies) -> Result<()> {
    let target = get_target();

    let msg = format!(
        "## Running ui tests in {path} against miri for {}",
        target.as_deref().unwrap_or("host")
    );
    eprintln!("{}", msg.green().bold());

    let with_dependencies = match with_dependencies {
        WithDependencies => true,
        WithoutDependencies => false,
    };
    run_tests(mode, path, target, with_dependencies)
}

fn get_target() -> Option<String> {
    env::var("MIRI_TEST_TARGET").ok()
}

fn main() -> Result<()> {
    ui_test::color_eyre::install()?;

    // Add a test env var to do environment communication tests.
    env::set_var("MIRI_ENV_VAR_TEST", "0");
    // Let the tests know where to store temp files (they might run for a different target, which can make this hard to find).
    env::set_var("MIRI_TEMP", env::temp_dir());

    ui(Mode::Pass, "tests/pass", WithoutDependencies)?;
    ui(Mode::Pass, "tests/pass-dep", WithDependencies)?;
    ui(Mode::Panic, "tests/panic", WithDependencies)?;
    ui(Mode::Fail { require_patterns: true }, "tests/fail", WithDependencies)?;
    if cfg!(target_os = "linux") {
        ui(Mode::Pass, "tests/extern-so/pass", WithoutDependencies)?;
        ui(Mode::Fail { require_patterns: true }, "tests/extern-so/fail", WithDependencies)?;
    }

    Ok(())
}
<commit_msg>run all extern-so tests consistently without dependencies<commit_after>use colored::*;
use regex::Regex;
use std::path::{Path, PathBuf};
use std::{env, process::Command};
use ui_test::{color_eyre::Result, Config, Mode, OutputConflictHandling};

fn miri_path() -> PathBuf {
    PathBuf::from(option_env!("MIRI").unwrap_or(env!("CARGO_BIN_EXE_miri")))
}

// Build the shared object file for testing external C function calls.
fn build_so_for_c_ffi_tests() -> PathBuf {
    let cc = option_env!("CC").unwrap_or("cc");
    // Target directory that we can write to.
    let so_target_dir = Path::new(&env::var_os("CARGO_TARGET_DIR").unwrap()).join("miri-extern-so");
    // Create the directory if it does not already exist.
    std::fs::create_dir_all(&so_target_dir)
        .expect("Failed to create directory for shared object file");
    let so_file_path = so_target_dir.join("libtestlib.so");
    let cc_output = Command::new(cc)
        .args([
            "-shared",
            "-o",
            so_file_path.to_str().unwrap(),
            "tests/extern-so/test.c",
            // Only add the functions specified in libcode.version to the shared object file.
            // This is to avoid automatically adding `malloc`, etc.
            // Source: https://anadoxin.org/blog/control-over-symbol-exports-in-gcc.html/
            "-fPIC",
            "-Wl,--version-script=tests/extern-so/libcode.version",
        ])
        .output()
        .expect("failed to generate shared object file for testing external C function calls");
    if !cc_output.status.success() {
        panic!("error in generating shared object file for testing external C function calls");
    }
    so_file_path
}

fn run_tests(
    mode: Mode,
    path: &str,
    target: Option<String>,
    with_dependencies: bool,
) -> Result<()> {
    let mut config = Config {
        target,
        stderr_filters: STDERR.clone(),
        stdout_filters: STDOUT.clone(),
        root_dir: PathBuf::from(path),
        mode,
        program: miri_path(),
        quiet: false,
        ..Config::default()
    };

    let in_rustc_test_suite = option_env!("RUSTC_STAGE").is_some();

    // Add some flags we always want.
    config.args.push("--edition".into());
    config.args.push("2018".into());
    if in_rustc_test_suite {
        // Less aggressive warnings to make the rustc toolstate management less painful.
        // (We often get warnings when e.g. a feature gets stabilized or some lint gets added/improved.)
        config.args.push("-Astable-features".into());
        config.args.push("-Aunused".into());
    } else {
        config.args.push("-Dwarnings".into());
        config.args.push("-Dunused".into());
    }
    if let Ok(extra_flags) = env::var("MIRIFLAGS") {
        for flag in extra_flags.split_whitespace() {
            config.args.push(flag.into());
        }
    }
    config.args.push("-Zui-testing".into());
    if let Some(target) = &config.target {
        config.args.push("--target".into());
        config.args.push(target.into());
    }

    // If we're on linux, and we're testing the extern-so functionality,
    // then build the shared object file for testing external C function calls
    // and push the relevant compiler flag.
    if cfg!(target_os = "linux") && path.starts_with("tests/extern-so/") {
        let so_file_path = build_so_for_c_ffi_tests();
        let mut flag = std::ffi::OsString::from("-Zmiri-extern-so-file=");
        flag.push(so_file_path.into_os_string());
        config.args.push(flag);
    }

    let skip_ui_checks = env::var_os("MIRI_SKIP_UI_CHECKS").is_some();

    config.output_conflict_handling = match (env::var_os("MIRI_BLESS").is_some(), skip_ui_checks) {
        (false, false) => OutputConflictHandling::Error,
        (true, false) => OutputConflictHandling::Bless,
        (false, true) => OutputConflictHandling::Ignore,
        (true, true) => panic!("cannot use MIRI_BLESS and MIRI_SKIP_UI_CHECKS at the same time"),
    };

    // Handle command-line arguments.
    config.path_filter.extend(std::env::args().skip(1).filter(|arg| {
        match &**arg {
            "--quiet" => {
                config.quiet = true;
                false
            }
            _ => true,
        }
    }));

    let use_std = env::var_os("MIRI_NO_STD").is_none();

    if with_dependencies && use_std {
        config.dependencies_crate_manifest_path =
            Some(Path::new("test_dependencies").join("Cargo.toml"));
        config.dependency_builder.args = vec![
            "run".into(),
            "--manifest-path".into(),
            "cargo-miri/Cargo.toml".into(),
            "--".into(),
            "miri".into(),
            "run".into(), // There is no `cargo miri build` so we just use `cargo miri run`.
        ];
    }
    ui_test::run_tests(config)
}

macro_rules! regexes {
    ($name:ident: $($regex:expr => $replacement:expr,)*) => {lazy_static::lazy_static! {
        static ref $name: Vec<(Regex, &'static str)> = vec![
            $((Regex::new($regex).unwrap(), $replacement),)*
        ];
    }};
}

regexes! {
    STDOUT:
    // Windows file paths
    r"\\"                           => "/",
}

regexes! {
    STDERR:
    // erase line and column info
    r"\.rs:[0-9]+:[0-9]+(: [0-9]+:[0-9]+)?" => ".rs:LL:CC",
    // erase alloc ids
    "alloc[0-9]+"                    => "ALLOC",
    // erase Stacked Borrows tags
    "<[0-9]+>"                       => "<TAG>",
    // erase whitespace that differs between platforms
    r" +at (.*\.rs)"                 => " at $1",
    // erase generics in backtraces
    "([0-9]+: .*)::<.*>"             => "$1",
    // erase addresses in backtraces
    "([0-9]+: ) +0x[0-9a-f]+ - (.*)" => "$1$2",
    // erase long hexadecimals
    r"0x[0-9a-fA-F]+[0-9a-fA-F]{2,2}" => "$$HEX",
    // erase specific alignments
    "alignment [0-9]+"               => "alignment ALIGN",
    // erase thread caller ids
    r"call [0-9]+"                  => "call ID",
    // erase platform module paths
    "sys::[a-z]+::"                  => "sys::PLATFORM::",
    // Windows file paths
    r"\\"                           => "/",
    // erase Rust stdlib path
    "[^ `]*/(rust[^/]*|checkout)/library/" => "RUSTLIB/",
    // erase platform file paths
    "sys/[a-z]+/"                    => "sys/PLATFORM/",
    // erase paths into the crate registry
    r"[^ ]*/\.?cargo/registry/.*/(.*\.rs)"  => "CARGO_REGISTRY/.../$1",
}

enum Dependencies {
    WithDependencies,
    WithoutDependencies,
}

use Dependencies::*;

fn ui(mode: Mode, path: &str, with_dependencies: Dependencies) -> Result<()> {
    let target = get_target();

    let msg = format!(
        "## Running ui tests in {path} against miri for {}",
        target.as_deref().unwrap_or("host")
    );
    eprintln!("{}", msg.green().bold());

    let with_dependencies = match with_dependencies {
        WithDependencies => true,
        WithoutDependencies => false,
    };
    run_tests(mode, path, target, with_dependencies)
}

fn get_target() -> Option<String> {
    env::var("MIRI_TEST_TARGET").ok()
}

fn main() -> Result<()> {
    ui_test::color_eyre::install()?;

    // Add a test env var to do environment communication tests.
    env::set_var("MIRI_ENV_VAR_TEST", "0");
    // Let the tests know where to store temp files (they might run for a different target, which can make this hard to find).
    env::set_var("MIRI_TEMP", env::temp_dir());

    ui(Mode::Pass, "tests/pass", WithoutDependencies)?;
    ui(Mode::Pass, "tests/pass-dep", WithDependencies)?;
    ui(Mode::Panic, "tests/panic", WithDependencies)?;
    ui(Mode::Fail { require_patterns: true }, "tests/fail", WithDependencies)?;
    if cfg!(target_os = "linux") {
        ui(Mode::Pass, "tests/extern-so/pass", WithoutDependencies)?;
        ui(Mode::Fail { require_patterns: true }, "tests/extern-so/fail", WithoutDependencies)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>extern crate thoughtfuck;

mod common;

use std::fs::File;
use std::io::Read;

use common::FauxStdout;
use thoughtfuck::vm::*;
use thoughtfuck::program::*;
use thoughtfuck::parse::parse;

fn read_file(filename: &str) -> String {
    let mut file = File::open(filename).unwrap();
    let mut source = String::new();
    file.read_to_string(&mut source).unwrap();
    source
}

fn test_output(input_filename: &str, expected_answer_filename: &str) {
    let source = read_file(input_filename);

    let commands = parse(&source);

    let mut program = Program::new();
    program.append(&commands);

    let mut stdout = FauxStdout::new();
    {
        let mut vm = VM::new(&mut stdout);
        program.execute(&mut vm);
    }
    assert_eq!(stdout.buffer, read_file(expected_answer_filename));
}

#[test]
fn hello_world () {
    const SOURCE: &str = "++++++++[>++++[>++>+++>+++>+<<<<-]>+>+>->>+[<]<-]>>.>---.+++++++..+++.>>.<-.<.+++.------.--------.>>+.>++.";
    let commands = parse(SOURCE);

    let mut program = Program::new();
    program.append(&commands);

    let mut stdout = FauxStdout::new();
    {
        let mut vm = VM::new(&mut stdout);
        program.execute(&mut vm);
    }
    assert_eq!(stdout.buffer, "Hello World!\n");
}
<commit_msg>Rewrite hello_world integration test to use input files<commit_after>extern crate thoughtfuck;

mod common;

use std::fs::File;
use std::io::Read;

use common::FauxStdout;
use thoughtfuck::vm::*;
use thoughtfuck::program::*;
use thoughtfuck::parse::parse;

fn read_file(filename: &str) -> String {
    let mut file = File::open(filename).unwrap();
    let mut source = String::new();
    file.read_to_string(&mut source).unwrap();
    source
}

fn test_output(input_filename: &str, expected_answer_filename: &str) {
    let source = read_file(input_filename);

    let commands = parse(&source);

    let mut program = Program::new();
    program.append(&commands);

    let mut stdout = FauxStdout::new();
    {
        let mut vm = VM::new(&mut stdout);
        program.execute(&mut vm);
    }
    assert_eq!(stdout.buffer, read_file(expected_answer_filename));
}

#[test]
fn hello_world () {
    test_output("./tests/hello.b", "./tests/hello.txt");
}
<|endoftext|>