<commit_before>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{ConsoleLogger, FeatureFlag, FeatureFlags, NamedItem, SourceLocationKey};
use fixture_tests::Fixture;
use graphql_ir::{
    build_ir_with_extra_features, BuilderOptions, FragmentDefinition, FragmentVariablesSemantic,
    OperationDefinition, Program, RelayMode,
};
use graphql_syntax::parse_executable;
use graphql_test_helpers::diagnostics_to_sorted_string;
use graphql_text_printer::print_full_operation;

use intern::string_key::Intern;
use relay_codegen::{
    build_request_params, print_fragment, print_operation, print_request, JsModuleFormat,
};
use relay_compiler::{validate, ConfigFileProject, ProjectConfig};
use relay_test_schema::{get_test_schema, get_test_schema_with_extensions};
use relay_transforms::{apply_transforms, DIRECTIVE_SPLIT_OPERATION};
use std::{array, sync::Arc};

pub fn transform_fixture(fixture: &Fixture<'_>) -> Result<String, String> {
    let source_location = SourceLocationKey::standalone(fixture.file_name);

    if fixture.content.contains("%TODO%") {
        if fixture.content.contains("expected-to-throw") {
            return Err("TODO".to_string());
        }
        return Ok("TODO".to_string());
    }

    let feature_flags = FeatureFlags {
        enable_flight_transform: true,
        hash_supported_argument: FeatureFlag::Limited {
            allowlist: array::IntoIter::new(["UserNameRenderer".intern()]).collect(),
        },
        no_inline: FeatureFlag::Enabled,
        enable_relay_resolver_transform: true,
        enable_3d_branch_arg_generation: true,
        actor_change_support: FeatureFlag::Enabled,
        text_artifacts: FeatureFlag::Disabled,
        enable_client_edges: FeatureFlag::Enabled,
        enable_provided_variables: FeatureFlag::Enabled,
        skip_printing_nulls: FeatureFlag::Disabled,
    };

    let default_project_config = ProjectConfig {
        name: "test".intern(),
        feature_flags: Arc::new(feature_flags),
        js_module_format: JsModuleFormat::Haste,
        ..Default::default()
    };
    // Adding %project_config section on top of the fixture will allow
    // us to validate output changes with different configurations
    let parts: Vec<_> = fixture.content.split("%project_config%").collect();
    let (project_config, other_parts) = match parts.as_slice() {
        [fixture_content, project_config_str] => (
            {
                let config_file_project: ConfigFileProject =
                    serde_json::from_str(project_config_str).unwrap();
                ProjectConfig {
                    schema_config: config_file_project.schema_config,
                    typegen_config: config_file_project.typegen_config,
                    feature_flags: config_file_project
                        .feature_flags
                        .map_or(default_project_config.feature_flags, |flags| {
                            Arc::new(flags)
                        }),
                    js_module_format: config_file_project.js_module_format,
                    ..default_project_config
                }
            },
            fixture_content.split("%extensions%").collect::<Vec<&str>>(),
        ),
        [fixture_content] => (
            default_project_config,
            fixture_content.split("%extensions%").collect::<Vec<&str>>(),
        ),
        _ => panic!("Invalid fixture input {}", fixture.content),
    };

    let (base, schema) = match other_parts.as_slice() {
        [base, extensions] => (base, get_test_schema_with_extensions(extensions)),
        [base] => (base, get_test_schema()),
        _ => panic!("Invalid fixture input {}", fixture.content),
    };

    let ast = parse_executable(base, source_location)
        .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;
    let ir_result = build_ir_with_extra_features(
        &schema,
        &ast.definitions,
        &BuilderOptions {
            allow_undefined_fragment_spreads: false,
            fragment_variables_semantic: FragmentVariablesSemantic::PassedValue,
            relay_mode: Some(RelayMode {
                enable_provided_variables: &FeatureFlag::Enabled,
            }),
            default_anonymous_operation_name: None,
        },
    );
    let ir = ir_result
        .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;
    let program = Program::from_definitions(Arc::clone(&schema), ir);

    validate(&program, &project_config, &None)
        .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;

    // TODO pass base fragment names
    let programs = apply_transforms(
        &project_config,
        Arc::new(program),
        Default::default(),
        Arc::new(ConsoleLogger),
        None,
        None,
    )
    .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;

    let mut operations: Vec<&std::sync::Arc<OperationDefinition>> =
        programs.normalization.operations().collect();
    operations.sort_by_key(|operation| operation.name.item);
    let result = operations
        .into_iter()
        .map(|operation| {
            if operation
                .directives
                .named(*DIRECTIVE_SPLIT_OPERATION)
                .is_some()
            {
                let mut import_statements = Default::default();
                let operation =
                    print_operation(&schema, operation, &project_config, &mut import_statements);
                format!("{}{}", import_statements, operation)
            } else {
                let name = operation.name.item;
                let print_operation_node = programs
                    .operation_text
                    .operation(name)
                    .expect("a query text operation should be generated for this operation");
                let text = print_full_operation(&programs.operation_text, print_operation_node);

                let reader_operation = programs
                    .reader
                    .operation(name)
                    .expect("a reader fragment should be generated for this operation");
                let operation_fragment = FragmentDefinition {
                    name: reader_operation.name,
                    variable_definitions: reader_operation.variable_definitions.clone(),
                    selections: reader_operation.selections.clone(),
                    used_global_variables: Default::default(),
                    directives: reader_operation.directives.clone(),
                    type_condition: reader_operation.type_,
                };
                let mut import_statements = Default::default();
                let request_parameters = build_request_params(operation);
                let request = print_request(
                    &schema,
                    operation,
                    &operation_fragment,
                    request_parameters,
                    &project_config,
                    &mut import_statements,
                );
                format!("{}{}\n\nQUERY:\n\n{}", import_statements, request, text)
            }
        })
        .chain({
            let mut fragments: Vec<&std::sync::Arc<FragmentDefinition>> =
                programs.reader.fragments().collect();
            fragments.sort_by_key(|fragment| fragment.name.item);
            fragments.into_iter().map(|fragment| {
                let mut import_statements = Default::default();
                let fragment =
                    print_fragment(&schema, fragment, &project_config, &mut import_statements);
                format!("{}{}", import_statements, fragment)
            })
        })
        .collect::<Vec<_>>();
    Ok(result.join("\n\n"))
}
<commit_msg>Unblock local development in rust?<commit_after>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{ConsoleLogger, FeatureFlag, FeatureFlags, NamedItem, SourceLocationKey};
use fixture_tests::Fixture;
use graphql_ir::{
    build_ir_with_extra_features, BuilderOptions, FragmentDefinition, FragmentVariablesSemantic,
    OperationDefinition, Program, RelayMode,
};
use graphql_syntax::parse_executable;
use graphql_test_helpers::diagnostics_to_sorted_string;
use graphql_text_printer::print_full_operation;

use intern::string_key::Intern;
use relay_codegen::{
    build_request_params, print_fragment, print_operation, print_request, JsModuleFormat,
};
use relay_compiler::{validate, ConfigFileProject, ProjectConfig};
use relay_test_schema::{get_test_schema, get_test_schema_with_extensions};
use relay_transforms::{apply_transforms, DIRECTIVE_SPLIT_OPERATION};
use std::sync::Arc;

pub fn transform_fixture(fixture: &Fixture<'_>) -> Result<String, String> {
    let source_location = SourceLocationKey::standalone(fixture.file_name);

    if fixture.content.contains("%TODO%") {
        if fixture.content.contains("expected-to-throw") {
            return Err("TODO".to_string());
        }
        return Ok("TODO".to_string());
    }
    let hash_supported_argument_allowlist = vec!["UserNameRenderer".intern()];

    let feature_flags = FeatureFlags {
        enable_flight_transform: true,
        hash_supported_argument: FeatureFlag::Limited {
            allowlist: hash_supported_argument_allowlist.into_iter().collect(),
        },
        no_inline: FeatureFlag::Enabled,
        enable_relay_resolver_transform: true,
        enable_3d_branch_arg_generation: true,
        actor_change_support: FeatureFlag::Enabled,
        text_artifacts: FeatureFlag::Disabled,
        enable_client_edges: FeatureFlag::Enabled,
        enable_provided_variables: FeatureFlag::Enabled,
        skip_printing_nulls: FeatureFlag::Disabled,
    };

    let default_project_config = ProjectConfig {
        name: "test".intern(),
        feature_flags: Arc::new(feature_flags),
        js_module_format: JsModuleFormat::Haste,
        ..Default::default()
    };
    // Adding %project_config section on top of the fixture will allow
    // us to validate output changes with different configurations
    let parts: Vec<_> = fixture.content.split("%project_config%").collect();
    let (project_config, other_parts) = match parts.as_slice() {
        [fixture_content, project_config_str] => (
            {
                let config_file_project: ConfigFileProject =
                    serde_json::from_str(project_config_str).unwrap();
                ProjectConfig {
                    schema_config: config_file_project.schema_config,
                    typegen_config: config_file_project.typegen_config,
                    feature_flags: config_file_project
                        .feature_flags
                        .map_or(default_project_config.feature_flags, |flags| {
                            Arc::new(flags)
                        }),
                    js_module_format: config_file_project.js_module_format,
                    ..default_project_config
                }
            },
            fixture_content.split("%extensions%").collect::<Vec<&str>>(),
        ),
        [fixture_content] => (
            default_project_config,
            fixture_content.split("%extensions%").collect::<Vec<&str>>(),
        ),
        _ => panic!("Invalid fixture input {}", fixture.content),
    };

    let (base, schema) = match other_parts.as_slice() {
        [base, extensions] => (base, get_test_schema_with_extensions(extensions)),
        [base] => (base, get_test_schema()),
        _ => panic!("Invalid fixture input {}", fixture.content),
    };

    let ast = parse_executable(base, source_location)
        .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;
    let ir_result = build_ir_with_extra_features(
        &schema,
        &ast.definitions,
        &BuilderOptions {
            allow_undefined_fragment_spreads: false,
            fragment_variables_semantic: FragmentVariablesSemantic::PassedValue,
            relay_mode: Some(RelayMode {
                enable_provided_variables: &FeatureFlag::Enabled,
            }),
            default_anonymous_operation_name: None,
        },
    );
    let ir = ir_result
        .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;
    let program = Program::from_definitions(Arc::clone(&schema), ir);

    validate(&program, &project_config, &None)
        .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;

    // TODO pass base fragment names
    let programs = apply_transforms(
        &project_config,
        Arc::new(program),
        Default::default(),
        Arc::new(ConsoleLogger),
        None,
        None,
    )
    .map_err(|diagnostics| diagnostics_to_sorted_string(fixture.content, &diagnostics))?;

    let mut operations: Vec<&std::sync::Arc<OperationDefinition>> =
        programs.normalization.operations().collect();
    operations.sort_by_key(|operation| operation.name.item);
    let result = operations
        .into_iter()
        .map(|operation| {
            if operation
                .directives
                .named(*DIRECTIVE_SPLIT_OPERATION)
                .is_some()
            {
                let mut import_statements = Default::default();
                let operation =
                    print_operation(&schema, operation, &project_config, &mut import_statements);
                format!("{}{}", import_statements, operation)
            } else {
                let name = operation.name.item;
                let print_operation_node = programs
                    .operation_text
                    .operation(name)
                    .expect("a query text operation should be generated for this operation");
                let text = print_full_operation(&programs.operation_text, print_operation_node);

                let reader_operation = programs
                    .reader
                    .operation(name)
                    .expect("a reader fragment should be generated for this operation");
                let operation_fragment = FragmentDefinition {
                    name: reader_operation.name,
                    variable_definitions: reader_operation.variable_definitions.clone(),
                    selections: reader_operation.selections.clone(),
                    used_global_variables: Default::default(),
                    directives: reader_operation.directives.clone(),
                    type_condition: reader_operation.type_,
                };
                let mut import_statements = Default::default();
                let request_parameters = build_request_params(operation);
                let request = print_request(
                    &schema,
                    operation,
                    &operation_fragment,
                    request_parameters,
                    &project_config,
                    &mut import_statements,
                );
                format!("{}{}\n\nQUERY:\n\n{}", import_statements, request, text)
            }
        })
        .chain({
            let mut fragments: Vec<&std::sync::Arc<FragmentDefinition>> =
                programs.reader.fragments().collect();
            fragments.sort_by_key(|fragment| fragment.name.item);
            fragments.into_iter().map(|fragment| {
                let mut import_statements = Default::default();
                let fragment =
                    print_fragment(&schema, fragment, &project_config, &mut import_statements);
                format!("{}{}", import_statements, fragment)
            })
        })
        .collect::<Vec<_>>();
    Ok(result.join("\n\n"))
}
<|endoftext|>


<commit_before>#![no_std]

#![feature(asm)]

#[allow(unstable)]
extern crate core;

use console::Console;
use console::Color;

// Must be re-exported to be called from assembly
pub use interrupt::int_handler;

mod console;
mod asm;
mod interrupt;

static mut console: console::Console = console::Console {position: 0};

#[no_mangle]
#[no_stack_check]
pub fn main() {
    unsafe {
        console.clear_screen(Color::LightRed);
        console.print_string("Hello world");
        interrupt::register_handler(on_keyboard_interrupt);
        console.print_string("End world");
    }
}

#[no_stack_check]
fn on_keyboard_interrupt() {
    let scancode = asm::inb(0x60);
    // Keydown
    if (scancode & 0b10000000 == 0) {
        let c = scancode_to_char(scancode);
        unsafe {
            console.print_char(c);
        }
    }
}

fn scancode_to_char(scancode: u8) -> char {
    let mut charmap:[char;256] = ['?';256];
    return charmap[scancode as usize];
}

<commit_msg>Add scancodes for 0-9<commit_after>#![no_std]

#![feature(asm)]

#[allow(unstable)]
extern crate core;

use console::Console;
use console::Color;

// Must be re-exported to be called from assembly
pub use interrupt::int_handler;

mod console;
mod asm;
mod interrupt;

static mut console: console::Console = console::Console {position: 0};

#[no_mangle]
#[no_stack_check]
pub fn main() {
    unsafe {
        console.clear_screen(Color::LightRed);
        console.print_string("Hello world");
        interrupt::register_handler(on_keyboard_interrupt);
        console.print_string("End world");
    }
}

#[no_stack_check]
fn on_keyboard_interrupt() {
    let scancode = asm::inb(0x60);
    // Keydown
    if (scancode & 0b10000000 == 0) {
        let c = scancode_to_char(scancode);
        unsafe {
            console.print_char(c);
        }
    }
}

fn scancode_to_char(scancode: u8) -> char {
    let mut charmap:[char;256] = ['?';256];
    charmap[0x2] = '1';
    charmap[0x3] = '2';
    charmap[0x4] = '3';
    charmap[0x5] = '4';
    charmap[0x6] = '5';
    charmap[0x7] = '6';
    charmap[0x8] = '7';
    charmap[0x9] = '8';
    charmap[0xA] = '9';
    charmap[0xB] = '0';

    return charmap[scancode as usize];
}

<|endoftext|>


<commit_before>use std::{
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpStream as StdTcpStream},
    ops::{Deref, DerefMut},
    os::unix::io::{AsRawFd, FromRawFd, IntoRawFd},
    pin::Pin,
    sync::atomic::{AtomicBool, Ordering},
    task::{self, Poll},
};

use cfg_if::cfg_if;
use futures::ready;
use log::error;
use pin_project::pin_project;
use socket2::SockAddr;
use tokio::{
    io::{AsyncRead, AsyncWrite, Interest, ReadBuf},
    net::{TcpSocket, TcpStream as TokioTcpStream, UdpSocket},
};

use crate::net::{
    sys::{set_common_sockopt_after_connect, set_common_sockopt_for_connect},
    AddrFamily,
    ConnectOpts,
};

enum TcpStreamState {
    Connected,
    FastOpenConnect(SocketAddr),
    FastOpenWrite,
}

/// A `TcpStream` that supports TFO (TCP Fast Open)
#[pin_project(project = TcpStreamProj)]
pub struct TcpStream {
    #[pin]
    inner: TokioTcpStream,
    state: TcpStreamState,
}

impl TcpStream {
    pub async fn connect(addr: SocketAddr, opts: &ConnectOpts) -> io::Result<TcpStream> {
        let socket = match addr {
            SocketAddr::V4(..) => TcpSocket::new_v4()?,
            SocketAddr::V6(..) => TcpSocket::new_v6()?,
        };

        // Any traffic to localhost should not be protected
        // This is a workaround for VPNService
        #[cfg(target_os = "android")]
        if !addr.ip().is_loopback() {
            use std::time::Duration;
            use tokio::time;

            if let Some(ref path) = opts.vpn_protect_path {
                // RPC calls to `VpnService.protect()`
                // Timeout in 3 seconds like shadowsocks-libev
                match time::timeout(Duration::from_secs(3), vpn_protect(path, socket.as_raw_fd())).await {
                    Ok(Ok(..)) => {}
                    Ok(Err(err)) => return Err(err),
                    Err(..) => return Err(io::Error::new(ErrorKind::TimedOut, "protect() timeout")),
                }
            }
        }

        // Set SO_MARK for mark-based routing on Linux (since 2.6.25)
        // NOTE: This will require CAP_NET_ADMIN capability (root in most cases)
        if let Some(mark) = opts.fwmark {
            let ret = unsafe {
                libc::setsockopt(
                    socket.as_raw_fd(),
                    libc::SOL_SOCKET,
                    libc::SO_MARK,
                    &mark as *const _ as *const _,
                    mem::size_of_val(&mark) as libc::socklen_t,
                )
            };
            if ret != 0 {
                let err = io::Error::last_os_error();
                error!("set SO_MARK error: {}", err);
                return Err(err);
            }
        }

        // Set SO_BINDTODEVICE for binding to a specific interface
        if let Some(ref iface) = opts.bind_interface {
            set_bindtodevice(&socket, iface)?;
        }

        set_common_sockopt_for_connect(addr, &socket, opts)?;

        if !opts.tcp.fastopen {
            // If TFO is not enabled, it just works like a normal TcpStream
            let stream = socket.connect(addr).await?;
            set_common_sockopt_after_connect(&stream, opts)?;

            return Ok(TcpStream {
                inner: stream,
                state: TcpStreamState::Connected,
            });
        }

        let mut connected = false;

        // TFO in Linux was supported since 3.7
        //
        // But TCP_FASTOPEN_CONNECT was supported since 4.1, so we have to be compatible with it
        static SUPPORT_TCP_FASTOPEN_CONNECT: AtomicBool = AtomicBool::new(true);
        if SUPPORT_TCP_FASTOPEN_CONNECT.load(Ordering::Relaxed) {
            unsafe {
                let enable: libc::c_int = 1;

                let ret = libc::setsockopt(
                    socket.as_raw_fd(),
                    libc::IPPROTO_TCP,
                    libc::TCP_FASTOPEN_CONNECT,
                    &enable as *const _ as *const libc::c_void,
                    mem::size_of_val(&enable) as libc::socklen_t,
                );

                if ret != 0 {
                    let err = io::Error::last_os_error();
                    if let Some(libc::ENOPROTOOPT) = err.raw_os_error() {
                        // `TCP_FASTOPEN_CONNECT` is not supported, maybe kernel version < 4.11
                        // Fallback to `sendto` with `MSG_FASTOPEN` (Supported after 3.7)
                        SUPPORT_TCP_FASTOPEN_CONNECT.store(false, Ordering::Relaxed);
                    } else {
                        error!("set TCP_FASTOPEN_CONNECT error: {}", err);
                        return Err(err);
                    }
                } else {
                    connected = true;
                }
            }
        }

        let stream = if connected {
            // call connect() if TCP_FASTOPEN_CONNECT is set
            socket.connect(addr).await?
        } else {
            // call sendto() with MSG_FASTOPEN in poll_read
            TokioTcpStream::from_std(unsafe { StdTcpStream::from_raw_fd(socket.into_raw_fd()) })?
        };

        set_common_sockopt_after_connect(&stream, opts)?;

        Ok(TcpStream {
            inner: stream,
            state: if connected {
                TcpStreamState::FastOpenWrite
            } else {
                TcpStreamState::FastOpenConnect(addr)
            },
        })
    }
}

impl Deref for TcpStream {
    type Target = TokioTcpStream;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl DerefMut for TcpStream {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl AsyncRead for TcpStream {
    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        self.project().inner.poll_read(cx, buf)
    }
}

impl AsyncWrite for TcpStream {
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        loop {
            let TcpStreamProj { inner, state } = self.as_mut().project();

            match *state {
                TcpStreamState::Connected => return inner.poll_write(cx, buf),

                TcpStreamState::FastOpenConnect(addr) => {
                    // Fallback mode. Must be kernal < 4.11
                    //
                    // Uses sendto as BSD-like systems

                    let saddr = SockAddr::from(addr);

                    let stream = inner.get_mut();

                    // Ensure socket is writable
                    ready!(stream.poll_write_ready(cx))?;

                    let mut connecting = false;
                    let send_result = stream.try_io(Interest::WRITABLE, || {
                        unsafe {
                            let ret = libc::sendto(
                                stream.as_raw_fd(),
                                buf.as_ptr() as *const libc::c_void,
                                buf.len(),
                                libc::MSG_FASTOPEN,
                                saddr.as_ptr(),
                                saddr.len(),
                            );

                            if ret >= 0 {
                                Ok(ret as usize)
                            } else {
                                // Error occurs
                                let err = io::Error::last_os_error();

                                // EINPROGRESS
                                if let Some(libc::EINPROGRESS) = err.raw_os_error() {
                                    // For non-blocking socket, it returns the number of bytes queued (and transmitted in the SYN-data packet) if cookie is available.
                                    // If cookie is not available, it transmits a data-less SYN packet with Fast Open cookie request option and returns -EINPROGRESS like connect().
                                    //
                                    // So in this state. We have to loop again to call `poll_write` for sending the first packet.
                                    connecting = true;

                                    // Let `try_io` clears the write readiness.
                                    Err(ErrorKind::WouldBlock.into())
                                } else {
                                    // Other errors, including EAGAIN, EWOULDBLOCK
                                    Err(err)
                                }
                            }
                        }
                    });

                    match send_result {
                        Ok(n) => {
                            // Connected successfully with fast open
                            *state = TcpStreamState::Connected;
                            return Ok(n).into();
                        }
                        Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                            if connecting {
                                // Connecting with normal TCP handshakes, write the first packet after connected
                                *state = TcpStreamState::Connected;
                            }
                        }
                        Err(err) => return Err(err).into(),
                    }
                }

                TcpStreamState::FastOpenWrite => {
                    // First `write` after `TCP_FASTOPEN_CONNECT`
                    // Kernel >= 4.11

                    let stream = inner.get_mut();

                    // Ensure socket is writable
                    ready!(stream.poll_write_ready(cx))?;

                    let mut connecting = false;
                    let send_result = stream.try_io(Interest::WRITABLE, || {
                        unsafe {
                            let ret = libc::send(stream.as_raw_fd(), buf.as_ptr() as *const libc::c_void, buf.len(), 0);

                            if ret >= 0 {
                                Ok(ret as usize)
                            } else {
                                let err = io::Error::last_os_error();
                                // EINPROGRESS
                                if let Some(libc::EINPROGRESS) = err.raw_os_error() {
                                    // For non-blocking socket, it returns the number of bytes queued (and transmitted in the SYN-data packet) if cookie is available.
                                    // If cookie is not available, it transmits a data-less SYN packet with Fast Open cookie request option and returns -EINPROGRESS like connect().
                                    //
                                    // So in this state. We have to loop again to call `poll_write` for sending the first packet.
                                    connecting = true;

                                    // Let `poll_write_io` clears the write readiness.
                                    Err(ErrorKind::WouldBlock.into())
                                } else {
                                    // Other errors, including EAGAIN, EWOULDBLOCK
                                    Err(err)
                                }
                            }
                        }
                    });

                    match send_result {
                        Ok(n) => {
                            // Connected successfully with fast open
                            *state = TcpStreamState::Connected;
                            return Ok(n).into();
                        }
                        Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                            if connecting {
                                // Connecting with normal TCP handshakes, write the first packet after connected
                                *state = TcpStreamState::Connected;
                            }
                        }
                        Err(err) => return Err(err).into(),
                    }
                }
            }
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().inner.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().inner.poll_shutdown(cx)
    }
}

/// Enable `TCP_FASTOPEN`
///
/// `TCP_FASTOPEN` was supported since Linux 3.7
pub fn set_tcp_fastopen<S: AsRawFd>(socket: &S) -> io::Result<()> {
    let queue: libc::c_int = 5;

    unsafe {
        let ret = libc::setsockopt(
            socket.as_raw_fd(),
            libc::IPPROTO_TCP,
            libc::TCP_FASTOPEN,
            &queue as *const _ as *const libc::c_void,
            mem::size_of_val(&queue) as libc::socklen_t,
        );

        if ret != 0 {
            let err = io::Error::last_os_error();
            error!("set TCP_FASTOPEN error: {}", err);
            return Err(err);
        }
    }

    Ok(())
}

/// Create a `UdpSocket` for connecting to `addr`
pub async fn create_outbound_udp_socket(af: AddrFamily, config: &ConnectOpts) -> io::Result<UdpSocket> {
    let bind_addr = match (af, config.bind_local_addr) {
        (AddrFamily::Ipv4, Some(IpAddr::V4(ip))) => SocketAddr::new(ip.into(), 0),
        (AddrFamily::Ipv6, Some(IpAddr::V6(ip))) => SocketAddr::new(ip.into(), 0),
        (AddrFamily::Ipv4, ..) => SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0),
        (AddrFamily::Ipv6, ..) => SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 0),
    };

    let socket = UdpSocket::bind(bind_addr).await?;

    // Any traffic except localhost should be protected
    // This is a workaround for VPNService
    #[cfg(target_os = "android")]
    {
        use std::time::Duration;
        use tokio::time;

        if let Some(ref path) = config.vpn_protect_path {
            // RPC calls to `VpnService.protect()`
            // Timeout in 3 seconds like shadowsocks-libev
            match time::timeout(Duration::from_secs(3), vpn_protect(path, socket.as_raw_fd())).await {
                Ok(Ok(..)) => {}
                Ok(Err(err)) => return Err(err),
                Err(..) => return Err(io::Error::new(ErrorKind::TimedOut, "protect() timeout")),
            }
        }
    }

    // Set SO_MARK for mark-based routing on Linux (since 2.6.25)
    // NOTE: This will require CAP_NET_ADMIN capability (root in most cases)
    if let Some(mark) = config.fwmark {
        let ret = unsafe {
            libc::setsockopt(
                socket.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_MARK,
                &mark as *const _ as *const _,
                mem::size_of_val(&mark) as libc::socklen_t,
            )
        };
        if ret != 0 {
            let err = io::Error::last_os_error();
            error!("set SO_MARK error: {}", err);
            return Err(err);
        }
    }

    // Set SO_BINDTODEVICE for binding to a specific interface
    if let Some(ref iface) = config.bind_interface {
        set_bindtodevice(&socket, iface)?;
    }

    Ok(socket)
}

fn set_bindtodevice<S: AsRawFd>(socket: &S, iface: &str) -> io::Result<()> {
    let iface_bytes = iface.as_bytes();

    unsafe {
        let ret = libc::setsockopt(
            socket.as_raw_fd(),
            libc::SOL_SOCKET,
            libc::SO_BINDTODEVICE,
            iface_bytes.as_ptr() as *const _ as *const libc::c_void,
            iface_bytes.len() as libc::socklen_t,
        );

        if ret != 0 {
            let err = io::Error::last_os_error();
            error!("set SO_BINDTODEVICE error: {}", err);
            return Err(err);
        }
    }

    Ok(())
}

cfg_if! {
    if #[cfg(target_os = "android")] {
        use std::path::Path;
        use std::os::unix::io::RawFd;

        use super::uds::UnixStream;

        /// This is a RPC for Android to `protect()` socket for connecting to remote servers
        ///
        /// https://developer.android.com/reference/android/net/VpnService#protect(java.net.Socket)
        ///
        /// More detail could be found in [shadowsocks-android](https://github.com/shadowsocks/shadowsocks-android) project.
        async fn vpn_protect<P: AsRef<Path>>(protect_path: P, fd: RawFd) -> io::Result<()> {
            use tokio::io::AsyncReadExt;

            let mut stream = UnixStream::connect(protect_path).await?;

            // send fds
            let dummy: [u8; 1] = [1];
            let fds: [RawFd; 1] = [fd];
            stream.send_with_fd(&dummy, &fds).await?;

            // receive the return value
            let mut response = [0; 1];
            stream.read_exact(&mut response).await?;

            if response[0] == 0xFF {
                return Err(io::Error::new(ErrorKind::Other, "protect() failed"));
            }

            Ok(())
        }
    }
}
<commit_msg>set TFO queue length to 1024 in Linux for matching backlogs<commit_after>use std::{
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpStream as StdTcpStream},
    ops::{Deref, DerefMut},
    os::unix::io::{AsRawFd, FromRawFd, IntoRawFd},
    pin::Pin,
    sync::atomic::{AtomicBool, Ordering},
    task::{self, Poll},
};

use cfg_if::cfg_if;
use futures::ready;
use log::error;
use pin_project::pin_project;
use socket2::SockAddr;
use tokio::{
    io::{AsyncRead, AsyncWrite, Interest, ReadBuf},
    net::{TcpSocket, TcpStream as TokioTcpStream, UdpSocket},
};

use crate::net::{
    sys::{set_common_sockopt_after_connect, set_common_sockopt_for_connect},
    AddrFamily,
    ConnectOpts,
};

enum TcpStreamState {
    Connected,
    FastOpenConnect(SocketAddr),
    FastOpenWrite,
}

/// A `TcpStream` that supports TFO (TCP Fast Open)
#[pin_project(project = TcpStreamProj)]
pub struct TcpStream {
    #[pin]
    inner: TokioTcpStream,
    state: TcpStreamState,
}

impl TcpStream {
    pub async fn connect(addr: SocketAddr, opts: &ConnectOpts) -> io::Result<TcpStream> {
        let socket = match addr {
            SocketAddr::V4(..) => TcpSocket::new_v4()?,
            SocketAddr::V6(..) => TcpSocket::new_v6()?,
        };

        // Any traffic to localhost should not be protected
        // This is a workaround for VPNService
        #[cfg(target_os = "android")]
        if !addr.ip().is_loopback() {
            use std::time::Duration;
            use tokio::time;

            if let Some(ref path) = opts.vpn_protect_path {
                // RPC calls to `VpnService.protect()`
                // Timeout in 3 seconds like shadowsocks-libev
                match time::timeout(Duration::from_secs(3), vpn_protect(path, socket.as_raw_fd())).await {
                    Ok(Ok(..)) => {}
                    Ok(Err(err)) => return Err(err),
                    Err(..) => return Err(io::Error::new(ErrorKind::TimedOut, "protect() timeout")),
                }
            }
        }

        // Set SO_MARK for mark-based routing on Linux (since 2.6.25)
        // NOTE: This will require CAP_NET_ADMIN capability (root in most cases)
        if let Some(mark) = opts.fwmark {
            let ret = unsafe {
                libc::setsockopt(
                    socket.as_raw_fd(),
                    libc::SOL_SOCKET,
                    libc::SO_MARK,
                    &mark as *const _ as *const _,
                    mem::size_of_val(&mark) as libc::socklen_t,
                )
            };
            if ret != 0 {
                let err = io::Error::last_os_error();
                error!("set SO_MARK error: {}", err);
                return Err(err);
            }
        }

        // Set SO_BINDTODEVICE for binding to a specific interface
        if let Some(ref iface) = opts.bind_interface {
            set_bindtodevice(&socket, iface)?;
        }

        set_common_sockopt_for_connect(addr, &socket, opts)?;

        if !opts.tcp.fastopen {
            // If TFO is not enabled, it just works like a normal TcpStream
            let stream = socket.connect(addr).await?;
            set_common_sockopt_after_connect(&stream, opts)?;

            return Ok(TcpStream {
                inner: stream,
                state: TcpStreamState::Connected,
            });
        }

        let mut connected = false;

        // TFO in Linux was supported since 3.7
        //
        // But TCP_FASTOPEN_CONNECT was supported since 4.1, so we have to be compatible with it
        static SUPPORT_TCP_FASTOPEN_CONNECT: AtomicBool = AtomicBool::new(true);
        if SUPPORT_TCP_FASTOPEN_CONNECT.load(Ordering::Relaxed) {
            unsafe {
                let enable: libc::c_int = 1;

                let ret = libc::setsockopt(
                    socket.as_raw_fd(),
                    libc::IPPROTO_TCP,
                    libc::TCP_FASTOPEN_CONNECT,
                    &enable as *const _ as *const libc::c_void,
                    mem::size_of_val(&enable) as libc::socklen_t,
                );

                if ret != 0 {
                    let err = io::Error::last_os_error();
                    if let Some(libc::ENOPROTOOPT) = err.raw_os_error() {
                        // `TCP_FASTOPEN_CONNECT` is not supported, maybe kernel version < 4.11
                        // Fallback to `sendto` with `MSG_FASTOPEN` (Supported after 3.7)
                        SUPPORT_TCP_FASTOPEN_CONNECT.store(false, Ordering::Relaxed);
                    } else {
                        error!("set TCP_FASTOPEN_CONNECT error: {}", err);
                        return Err(err);
                    }
                } else {
                    connected = true;
                }
            }
        }

        let stream = if connected {
            // call connect() if TCP_FASTOPEN_CONNECT is set
            socket.connect(addr).await?
        } else {
            // call sendto() with MSG_FASTOPEN in poll_read
            TokioTcpStream::from_std(unsafe { StdTcpStream::from_raw_fd(socket.into_raw_fd()) })?
        };

        set_common_sockopt_after_connect(&stream, opts)?;

        Ok(TcpStream {
            inner: stream,
            state: if connected {
                TcpStreamState::FastOpenWrite
            } else {
                TcpStreamState::FastOpenConnect(addr)
            },
        })
    }
}

impl Deref for TcpStream {
    type Target = TokioTcpStream;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl DerefMut for TcpStream {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl AsyncRead for TcpStream {
    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        self.project().inner.poll_read(cx, buf)
    }
}

impl AsyncWrite for TcpStream {
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        loop {
            let TcpStreamProj { inner, state } = self.as_mut().project();

            match *state {
                TcpStreamState::Connected => return inner.poll_write(cx, buf),

                TcpStreamState::FastOpenConnect(addr) => {
                    // Fallback mode. Must be kernal < 4.11
                    //
                    // Uses sendto as BSD-like systems

                    let saddr = SockAddr::from(addr);

                    let stream = inner.get_mut();

                    // Ensure socket is writable
                    ready!(stream.poll_write_ready(cx))?;

                    let mut connecting = false;
                    let send_result = stream.try_io(Interest::WRITABLE, || {
                        unsafe {
                            let ret = libc::sendto(
                                stream.as_raw_fd(),
                                buf.as_ptr() as *const libc::c_void,
                                buf.len(),
                                libc::MSG_FASTOPEN,
                                saddr.as_ptr(),
                                saddr.len(),
                            );

                            if ret >= 0 {
                                Ok(ret as usize)
                            } else {
                                // Error occurs
                                let err = io::Error::last_os_error();

                                // EINPROGRESS
                                if let Some(libc::EINPROGRESS) = err.raw_os_error() {
                                    // For non-blocking socket, it returns the number of bytes queued (and transmitted in the SYN-data packet) if cookie is available.
                                    // If cookie is not available, it transmits a data-less SYN packet with Fast Open cookie request option and returns -EINPROGRESS like connect().
                                    //
                                    // So in this state. We have to loop again to call `poll_write` for sending the first packet.
                                    connecting = true;

                                    // Let `try_io` clears the write readiness.
                                    Err(ErrorKind::WouldBlock.into())
                                } else {
                                    // Other errors, including EAGAIN, EWOULDBLOCK
                                    Err(err)
                                }
                            }
                        }
                    });

                    match send_result {
                        Ok(n) => {
                            // Connected successfully with fast open
                            *state = TcpStreamState::Connected;
                            return Ok(n).into();
                        }
                        Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                            if connecting {
                                // Connecting with normal TCP handshakes, write the first packet after connected
                                *state = TcpStreamState::Connected;
                            }
                        }
                        Err(err) => return Err(err).into(),
                    }
                }

                TcpStreamState::FastOpenWrite => {
                    // First `write` after `TCP_FASTOPEN_CONNECT`
                    // Kernel >= 4.11

                    let stream = inner.get_mut();

                    // Ensure socket is writable
                    ready!(stream.poll_write_ready(cx))?;

                    let mut connecting = false;
                    let send_result = stream.try_io(Interest::WRITABLE, || {
                        unsafe {
                            let ret = libc::send(stream.as_raw_fd(), buf.as_ptr() as *const libc::c_void, buf.len(), 0);

                            if ret >= 0 {
                                Ok(ret as usize)
                            } else {
                                let err = io::Error::last_os_error();
                                // EINPROGRESS
                                if let Some(libc::EINPROGRESS) = err.raw_os_error() {
                                    // For non-blocking socket, it returns the number of bytes queued (and transmitted in the SYN-data packet) if cookie is available.
                                    // If cookie is not available, it transmits a data-less SYN packet with Fast Open cookie request option and returns -EINPROGRESS like connect().
                                    //
                                    // So in this state. We have to loop again to call `poll_write` for sending the first packet.
                                    connecting = true;

                                    // Let `poll_write_io` clears the write readiness.
                                    Err(ErrorKind::WouldBlock.into())
                                } else {
                                    // Other errors, including EAGAIN, EWOULDBLOCK
                                    Err(err)
                                }
                            }
                        }
                    });

                    match send_result {
                        Ok(n) => {
                            // Connected successfully with fast open
                            *state = TcpStreamState::Connected;
                            return Ok(n).into();
                        }
                        Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                            if connecting {
                                // Connecting with normal TCP handshakes, write the first packet after connected
                                *state = TcpStreamState::Connected;
                            }
                        }
                        Err(err) => return Err(err).into(),
                    }
                }
            }
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().inner.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().inner.poll_shutdown(cx)
    }
}

/// Enable `TCP_FASTOPEN`
///
/// `TCP_FASTOPEN` was supported since Linux 3.7
pub fn set_tcp_fastopen<S: AsRawFd>(socket: &S) -> io::Result<()> {
    // https://lwn.net/Articles/508865/
    //
    // The option value, qlen, specifies this server's limit on the size of the queue of TFO requests that have
    // not yet completed the three-way handshake (see the remarks on prevention of resource-exhaustion attacks above).
    //
    // It was recommended to be `5` in this document.
    //
    // But since mio's TcpListener sets backlogs to 1024, it would be nice to have 1024 slots for handshaking TFO requests.
    let queue: libc::c_int = 1024;

    unsafe {
        let ret = libc::setsockopt(
            socket.as_raw_fd(),
            libc::IPPROTO_TCP,
            libc::TCP_FASTOPEN,
            &queue as *const _ as *const libc::c_void,
            mem::size_of_val(&queue) as libc::socklen_t,
        );

        if ret != 0 {
            let err = io::Error::last_os_error();
            error!("set TCP_FASTOPEN error: {}", err);
            return Err(err);
        }
    }

    Ok(())
}

/// Create a `UdpSocket` for connecting to `addr`
pub async fn create_outbound_udp_socket(af: AddrFamily, config: &ConnectOpts) -> io::Result<UdpSocket> {
    let bind_addr = match (af, config.bind_local_addr) {
        (AddrFamily::Ipv4, Some(IpAddr::V4(ip))) => SocketAddr::new(ip.into(), 0),
        (AddrFamily::Ipv6, Some(IpAddr::V6(ip))) => SocketAddr::new(ip.into(), 0),
        (AddrFamily::Ipv4, ..) => SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0),
        (AddrFamily::Ipv6, ..) => SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 0),
    };

    let socket = UdpSocket::bind(bind_addr).await?;

    // Any traffic except localhost should be protected
    // This is a workaround for VPNService
    #[cfg(target_os = "android")]
    {
        use std::time::Duration;
        use tokio::time;

        if let Some(ref path) = config.vpn_protect_path {
            // RPC calls to `VpnService.protect()`
            // Timeout in 3 seconds like shadowsocks-libev
            match time::timeout(Duration::from_secs(3), vpn_protect(path, socket.as_raw_fd())).await {
                Ok(Ok(..)) => {}
                Ok(Err(err)) => return Err(err),
                Err(..) => return Err(io::Error::new(ErrorKind::TimedOut, "protect() timeout")),
            }
        }
    }

    // Set SO_MARK for mark-based routing on Linux (since 2.6.25)
    // NOTE: This will require CAP_NET_ADMIN capability (root in most cases)
    if let Some(mark) = config.fwmark {
        let ret = unsafe {
            libc::setsockopt(
                socket.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_MARK,
                &mark as *const _ as *const _,
                mem::size_of_val(&mark) as libc::socklen_t,
            )
        };
        if ret != 0 {
            let err = io::Error::last_os_error();
            error!("set SO_MARK error: {}", err);
            return Err(err);
        }
    }

    // Set SO_BINDTODEVICE for binding to a specific interface
    if let Some(ref iface) = config.bind_interface {
        set_bindtodevice(&socket, iface)?;
    }

    Ok(socket)
}

fn set_bindtodevice<S: AsRawFd>(socket: &S, iface: &str) -> io::Result<()> {
    let iface_bytes = iface.as_bytes();

    unsafe {
        let ret = libc::setsockopt(
            socket.as_raw_fd(),
            libc::SOL_SOCKET,
            libc::SO_BINDTODEVICE,
            iface_bytes.as_ptr() as *const _ as *const libc::c_void,
            iface_bytes.len() as libc::socklen_t,
        );

        if ret != 0 {
            let err = io::Error::last_os_error();
            error!("set SO_BINDTODEVICE error: {}", err);
            return Err(err);
        }
    }

    Ok(())
}

cfg_if! {
    if #[cfg(target_os = "android")] {
        use std::path::Path;
        use std::os::unix::io::RawFd;

        use super::uds::UnixStream;

        /// This is a RPC for Android to `protect()` socket for connecting to remote servers
        ///
        /// https://developer.android.com/reference/android/net/VpnService#protect(java.net.Socket)
        ///
        /// More detail could be found in [shadowsocks-android](https://github.com/shadowsocks/shadowsocks-android) project.
        async fn vpn_protect<P: AsRef<Path>>(protect_path: P, fd: RawFd) -> io::Result<()> {
            use tokio::io::AsyncReadExt;

            let mut stream = UnixStream::connect(protect_path).await?;

            // send fds
            let dummy: [u8; 1] = [1];
            let fds: [RawFd; 1] = [fd];
            stream.send_with_fd(&dummy, &fds).await?;

            // receive the return value
            let mut response = [0; 1];
            stream.read_exact(&mut response).await?;

            if response[0] == 0xFF {
                return Err(io::Error::new(ErrorKind::Other, "protect() failed"));
            }

            Ok(())
        }
    }
}
<|endoftext|>


<commit_before>use arch::cpu;
use arch::pci;
use process;
use spin::{StaticMutex,STATIC_MUTEX_INIT};
use syscall::{ErrNum,Result};

static MUTEX: StaticMutex = STATIC_MUTEX_INIT;

const VBE_DISPI_IOPORT_INDEX: u16 = 0x1ce;
const VBE_DISPI_IOPORT_DATA: u16 = 0x1cf;
const VBE_DISPI_INDEX_XRES: u16 = 1;
const VBE_DISPI_INDEX_YRES: u16 = 2;
const VBE_DISPI_INDEX_BPP: u16 = 3;
const VBE_DISPI_INDEX_ENABLE: u16 = 4;
const VBE_DISPI_DISABLED: u16 = 0;
const VBE_DISPI_ENABLED: u16 = 1;
const VBE_DISPI_LFB_ENABLED: u16 = 64;

unsafe fn vbe_write(index: u16, value: u16) {
   cpu::outw(VBE_DISPI_IOPORT_INDEX, index);
   cpu::outw(VBE_DISPI_IOPORT_DATA, value);
}

pub fn init(xres: u16, yres: u16, bpp: u8) -> Result<&'static mut [u8]> {
    let _d = lock!(MUTEX);
    unsafe {
        let bus_slot =
            pci::find(0x1234, 0x1111)                   // QEmu
                .or_else(|| pci::find(0x80EE, 0xBEEF))  // VirtualBox
                .ok_or(ErrNum::NotSupported)?;

        let bar0 = pci::inl(bus_slot, 0, 0x10);
        let bar1 = pci::inl(bus_slot, 0, 0x14);
        let base_address = (bar1 as u64) << 32 | (bar0 as u64) & 0xffff_fff0;
        vbe_write(VBE_DISPI_INDEX_ENABLE, VBE_DISPI_DISABLED);
        vbe_write(VBE_DISPI_INDEX_XRES, xres);
        vbe_write(VBE_DISPI_INDEX_YRES, yres);
        vbe_write(VBE_DISPI_INDEX_BPP, bpp as u16);
        vbe_write(VBE_DISPI_INDEX_ENABLE, VBE_DISPI_ENABLED | VBE_DISPI_LFB_ENABLED);

        let len = (xres as usize * yres as usize * bpp as usize) / 8;
        process::map_phys(base_address as usize, len, true, true)
    }
}
<commit_msg>Clear the screen<commit_after>use arch::cpu;
use arch::pci;
use process;
use spin::{StaticMutex,STATIC_MUTEX_INIT};
use syscall::{ErrNum,Result};

static MUTEX: StaticMutex = STATIC_MUTEX_INIT;

const VBE_DISPI_IOPORT_INDEX: u16 = 0x1ce;
const VBE_DISPI_IOPORT_DATA: u16 = 0x1cf;
const VBE_DISPI_INDEX_XRES: u16 = 1;
const VBE_DISPI_INDEX_YRES: u16 = 2;
const VBE_DISPI_INDEX_BPP: u16 = 3;
const VBE_DISPI_INDEX_ENABLE: u16 = 4;
const VBE_DISPI_DISABLED: u16 = 0;
const VBE_DISPI_ENABLED: u16 = 1;
const VBE_DISPI_LFB_ENABLED: u16 = 64;

unsafe fn vbe_write(index: u16, value: u16) {
   cpu::outw(VBE_DISPI_IOPORT_INDEX, index);
   cpu::outw(VBE_DISPI_IOPORT_DATA, value);
}

pub fn init(xres: u16, yres: u16, bpp: u8) -> Result<&'static mut [u8]> {
    let _d = lock!(MUTEX);
    let len = (xres as usize * yres as usize * bpp as usize) / 8;

    let slice =
        unsafe {
            let bus_slot =
                pci::find(0x1234, 0x1111)                   // QEmu
                    .or_else(|| pci::find(0x80EE, 0xBEEF))  // VirtualBox
                    .ok_or(ErrNum::NotSupported)?;

            let bar0 = pci::inl(bus_slot, 0, 0x10);
            let bar1 = pci::inl(bus_slot, 0, 0x14);
            let base_address = (bar1 as u64) << 32 | (bar0 as u64) & 0xffff_fff0;
            vbe_write(VBE_DISPI_INDEX_ENABLE, VBE_DISPI_DISABLED);
            vbe_write(VBE_DISPI_INDEX_XRES, xres);
            vbe_write(VBE_DISPI_INDEX_YRES, yres);
            vbe_write(VBE_DISPI_INDEX_BPP, bpp as u16);
            vbe_write(VBE_DISPI_INDEX_ENABLE, VBE_DISPI_ENABLED | VBE_DISPI_LFB_ENABLED);
            process::map_phys(base_address as usize, len, true, true)?
        };

    for b in slice.iter_mut() {
        *b = 0;
    }

    Ok(slice)
}
<|endoftext|>


<commit_before>// Copyright 2019-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
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Fuzzing library for the persistent store.
//!
//! The overall design principles are (in order of precedence):
//! - Determinism: fuzzing is a function from seeds (byte slices) to sequences of store
//!   manipulations (things like creating a store, applying operations, interrupting operations,
//!   interrupting reboots, checking invariant, etc). We can replay this function on the same input
//!   to get the same sequence of manipulations (for the same fuzzing and store code).
//! - Coverage: fuzzing tries to coverage as much different behaviors as possible for small seeds.
//!   Ideally, each seed bit would control a branch decision in the tree of execution paths.
//! - Surjectivity: all sequences of manipulations are reachable by fuzzing for some seed. The only
//!   situation where coverage takes precedence over surjectivity is for the value of insert updates
//!   where a pseudo-random generator is used to avoid wasting entropy.

// TODO(ia0): Remove when used.
#![allow(dead_code)]

/// Bit-level entropy source based on a byte slice shared reference.
///
/// The entropy has the following properties (in order of precedence):
/// - It always returns a result.
/// - It is deterministic: for a given slice and a given sequence of operations, the same results
///   are returned. This permits to replay and debug fuzzing artifacts.
/// - It uses the slice as a bit stream. In particular, it doesn't do big number arithmetic. This
///   permits to have a simple implementation.
/// - It doesn't waste information: for a given operation, the minimum integer number of bits is
///   used to produce the result. As a consequence fractional bits can be wasted at each operation.
/// - It uses the information uniformly: each bit is used exactly once, except when only a fraction
///   of it is used. In particular, a bit is not used more than once. A consequence of each bit
///   being used essentially once, is that the results are mostly uniformly distributed.
///
/// # Invariant
///
/// - The bit is a valid position in the slice, or one past: `bit <= 8 * data.len()`.
struct Entropy<'a> {
    /// The byte slice shared reference providing the entropy.
    data: &'a [u8],

    /// The bit position in the byte slice of the next entropy bit.
    bit: usize,
}

impl Entropy<'_> {
    /// Creates a bit-level entropy given a byte slice.
    fn new(data: &[u8]) -> Entropy {
        let bit = 0;
        Entropy { data, bit }
    }

    /// Consumes the remaining entropy.
    fn consume_all(&mut self) {
        self.bit = 8 * self.data.len();
    }

    /// Returns whether there is entropy remaining.
    fn is_empty(&self) -> bool {
        assert!(self.bit <= 8 * self.data.len());
        self.bit == 8 * self.data.len()
    }

    /// Reads a bit.
    fn read_bit(&mut self) -> bool {
        if self.is_empty() {
            return false;
        }
        let b = self.bit;
        self.bit += 1;
        self.data[b / 8] & 1 << (b % 8) != 0
    }

    /// Reads a number with a given bit-width.
    ///
    /// # Preconditions
    ///
    /// - The number should fit in the return type: `n <= 8 * size_of::<usize>()`.
    fn read_bits(&mut self, n: usize) -> usize {
        assert!(n <= 8 * std::mem::size_of::<usize>());
        let mut r = 0;
        for i in 0..n {
            r |= (self.read_bit() as usize) << i;
        }
        r
    }

    /// Reads a byte.
    fn read_byte(&mut self) -> u8 {
        self.read_bits(8) as u8
    }

    /// Reads a slice.
    fn read_slice(&mut self, length: usize) -> Vec<u8> {
        let mut result = Vec::with_capacity(length);
        for _ in 0..length {
            result.push(self.read_byte());
        }
        result
    }

    /// Reads a number between `min` and `max` (inclusive bounds).
    ///
    /// The distribution is uniform if the range width is a power of two. Otherwise, the minimum
    /// amount of entropy is used (the next power of two) and the distribution is the closest to
    /// uniform for that entropy.
    ///
    /// # Preconditions
    ///
    /// - The bounds should be correctly ordered: `min <= max`.
    /// - The upper-bound should not be too large: `max < usize::max_value()`.
    fn read_range(&mut self, min: usize, max: usize) -> usize {
        assert!(min <= max && max < usize::max_value());
        let count = max - min + 1;
        let delta = self.read_bits(num_bits(count - 1)) % count;
        min + delta
    }
}

/// Returns the number of bits necessary to represent a number.
fn num_bits(x: usize) -> usize {
    8 * std::mem::size_of::<usize>() - x.leading_zeros() as usize
}

#[test]
fn num_bits_ok() {
    assert_eq!(num_bits(0), 0);
    assert_eq!(num_bits(1), 1);
    assert_eq!(num_bits(2), 2);
    assert_eq!(num_bits(3), 2);
    assert_eq!(num_bits(4), 3);
    assert_eq!(num_bits(7), 3);
    assert_eq!(num_bits(8), 4);
    assert_eq!(num_bits(15), 4);
    assert_eq!(num_bits(16), 5);
    assert_eq!(
        num_bits(usize::max_value()),
        8 * std::mem::size_of::<usize>()
    );
}

#[test]
fn read_bit_ok() {
    let mut entropy = Entropy::new(&[0b10110010]);
    assert!(!entropy.read_bit());
    assert!(entropy.read_bit());
    assert!(!entropy.read_bit());
    assert!(!entropy.read_bit());
    assert!(entropy.read_bit());
    assert!(entropy.read_bit());
    assert!(!entropy.read_bit());
    assert!(entropy.read_bit());
}

#[test]
fn read_bits_ok() {
    let mut entropy = Entropy::new(&[0x83, 0x92]);
    assert_eq!(entropy.read_bits(4), 0x3);
    assert_eq!(entropy.read_bits(8), 0x28);
    assert_eq!(entropy.read_bits(2), 0b01);
    assert_eq!(entropy.read_bits(2), 0b10);
}

#[test]
fn read_range_ok() {
    let mut entropy = Entropy::new(&[0b00101011]);
    assert_eq!(entropy.read_range(0, 7), 0b011);
    assert_eq!(entropy.read_range(1, 8), 1 + 0b101);
    assert_eq!(entropy.read_range(4, 6), 4 + 0b00);
    let mut entropy = Entropy::new(&[0b00101011]);
    assert_eq!(entropy.read_range(0, 8), 0b1011 % 9);
    assert_eq!(entropy.read_range(3, 15), 3 + 0b0010);
    let mut entropy = Entropy::new(&[0x12, 0x34, 0x56, 0x78]);
    assert_eq!(entropy.read_range(0, usize::max_value() - 1), 0x78563412);
}
<commit_msg>Improve documentation<commit_after>// Copyright 2019-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
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Fuzzing library for the persistent store.
//!
//! The overall design principles are (in order of precedence):
//! - Determinism: fuzzing is a function from seeds (byte slices) to sequences of store
//!   manipulations (things like creating a store, applying operations, interrupting operations,
//!   interrupting reboots, checking invariant, etc). We can replay this function on the same input
//!   to get the same sequence of manipulations (for the same fuzzing and store code).
//! - Coverage: fuzzing tries to coverage as much different behaviors as possible for small seeds.
//!   Ideally, each seed bit would control a branch decision in the tree of execution paths.
//! - Surjectivity: all sequences of manipulations are reachable by fuzzing for some seed. The only
//!   situation where coverage takes precedence over surjectivity is for the value of insert updates
//!   where a pseudo-random generator is used to avoid wasting entropy.

// TODO(ia0): Remove when used.
#![allow(dead_code)]

/// Bit-level entropy source based on a byte slice shared reference.
///
/// This is used to convert the byte slice provided by the fuzzer into the entropy used by the
/// fuzzing code to generate a sequence of store manipulations, among other things. Entropy
/// operations use the shortest necessary sequence of bits from the byte slice, such that fuzzer
/// mutations of the byte slice have local impact or cascading effects towards future operations
/// only.
///
/// The entropy has the following properties (in order of precedence):
/// - It always returns a result.
/// - It is deterministic: for a given slice and a given sequence of operations, the same results
///   are returned. This permits to replay and debug fuzzing artifacts.
/// - It uses the slice as a bit stream. In particular, it doesn't do big number arithmetic. This
///   permits to have a simple implementation.
/// - It doesn't waste information: for a given operation, the minimum integer number of bits is
///   used to produce the result. As a consequence fractional bits can be wasted at each operation.
/// - It uses the information uniformly: each bit is used exactly once, except when only a fraction
///   of it is used. In particular, a bit is not used more than once. A consequence of each bit
///   being used essentially once, is that the results are mostly uniformly distributed.
///
/// # Invariant
///
/// - The bit is a valid position in the slice, or one past: `bit <= 8 * data.len()`.
struct Entropy<'a> {
    /// The byte slice shared reference providing the entropy.
    data: &'a [u8],

    /// The bit position in the byte slice of the next entropy bit.
    bit: usize,
}

impl Entropy<'_> {
    /// Creates a bit-level entropy given a byte slice.
    fn new(data: &[u8]) -> Entropy {
        let bit = 0;
        Entropy { data, bit }
    }

    /// Consumes the remaining entropy.
    fn consume_all(&mut self) {
        self.bit = 8 * self.data.len();
    }

    /// Returns whether there is entropy remaining.
    fn is_empty(&self) -> bool {
        assert!(self.bit <= 8 * self.data.len());
        self.bit == 8 * self.data.len()
    }

    /// Reads a bit.
    fn read_bit(&mut self) -> bool {
        if self.is_empty() {
            return false;
        }
        let b = self.bit;
        self.bit += 1;
        self.data[b / 8] & 1 << (b % 8) != 0
    }

    /// Reads a number with a given bit-width.
    ///
    /// # Preconditions
    ///
    /// - The number should fit in the return type: `n <= 8 * size_of::<usize>()`.
    fn read_bits(&mut self, n: usize) -> usize {
        assert!(n <= 8 * std::mem::size_of::<usize>());
        let mut r = 0;
        for i in 0..n {
            r |= (self.read_bit() as usize) << i;
        }
        r
    }

    /// Reads a byte.
    fn read_byte(&mut self) -> u8 {
        self.read_bits(8) as u8
    }

    /// Reads a slice.
    fn read_slice(&mut self, length: usize) -> Vec<u8> {
        let mut result = Vec::with_capacity(length);
        for _ in 0..length {
            result.push(self.read_byte());
        }
        result
    }

    /// Reads a number between `min` and `max` (inclusive bounds).
    ///
    /// The distribution is uniform if the range width is a power of two. Otherwise, the minimum
    /// amount of entropy is used (the next power of two) and the distribution is the closest to
    /// uniform for that entropy.
    ///
    /// # Preconditions
    ///
    /// - The bounds should be correctly ordered: `min <= max`.
    /// - The upper-bound should not be too large: `max < usize::max_value()`.
    fn read_range(&mut self, min: usize, max: usize) -> usize {
        assert!(min <= max && max < usize::max_value());
        let count = max - min + 1;
        let delta = self.read_bits(num_bits(count - 1)) % count;
        min + delta
    }
}

/// Returns the number of bits necessary to represent a number.
fn num_bits(x: usize) -> usize {
    8 * std::mem::size_of::<usize>() - x.leading_zeros() as usize
}

#[test]
fn num_bits_ok() {
    assert_eq!(num_bits(0), 0);
    assert_eq!(num_bits(1), 1);
    assert_eq!(num_bits(2), 2);
    assert_eq!(num_bits(3), 2);
    assert_eq!(num_bits(4), 3);
    assert_eq!(num_bits(7), 3);
    assert_eq!(num_bits(8), 4);
    assert_eq!(num_bits(15), 4);
    assert_eq!(num_bits(16), 5);
    assert_eq!(
        num_bits(usize::max_value()),
        8 * std::mem::size_of::<usize>()
    );
}

#[test]
fn read_bit_ok() {
    let mut entropy = Entropy::new(&[0b10110010]);
    assert!(!entropy.read_bit());
    assert!(entropy.read_bit());
    assert!(!entropy.read_bit());
    assert!(!entropy.read_bit());
    assert!(entropy.read_bit());
    assert!(entropy.read_bit());
    assert!(!entropy.read_bit());
    assert!(entropy.read_bit());
}

#[test]
fn read_bits_ok() {
    let mut entropy = Entropy::new(&[0x83, 0x92]);
    assert_eq!(entropy.read_bits(4), 0x3);
    assert_eq!(entropy.read_bits(8), 0x28);
    assert_eq!(entropy.read_bits(2), 0b01);
    assert_eq!(entropy.read_bits(2), 0b10);
}

#[test]
fn read_range_ok() {
    let mut entropy = Entropy::new(&[0b00101011]);
    assert_eq!(entropy.read_range(0, 7), 0b011);
    assert_eq!(entropy.read_range(1, 8), 1 + 0b101);
    assert_eq!(entropy.read_range(4, 6), 4 + 0b00);
    let mut entropy = Entropy::new(&[0b00101011]);
    assert_eq!(entropy.read_range(0, 8), 0b1011 % 9);
    assert_eq!(entropy.read_range(3, 15), 3 + 0b0010);
    let mut entropy = Entropy::new(&[0x12, 0x34, 0x56, 0x78]);
    assert_eq!(entropy.read_range(0, usize::max_value() - 1), 0x78563412);
}
<|endoftext|>


<commit_before>use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen};
use crate::ops::Try;

/// An iterator with a `peek()` that returns an optional reference to the next
/// element.
///
/// This `struct` is created by the [`peekable`] method on [`Iterator`]. See its
/// documentation for more.
///
/// [`peekable`]: Iterator::peekable
/// [`Iterator`]: trait.Iterator.html
#[derive(Clone, Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Peekable<I: Iterator> {
    iter: I,
    /// Remember a peeked value, even if it was None.
    peeked: Option<Option<I::Item>>,
}

impl<I: Iterator> Peekable<I> {
    pub(in crate::iter) fn new(iter: I) -> Peekable<I> {
        Peekable { iter, peeked: None }
    }
}

// Peekable must remember if a None has been seen in the `.peek()` method.
// It ensures that `.peek(); .peek();` or `.peek(); .next();` only advances the
// underlying iterator at most once. This does not by itself make the iterator
// fused.
#[stable(feature = "rust1", since = "1.0.0")]
impl<I: Iterator> Iterator for Peekable<I> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<I::Item> {
        match self.peeked.take() {
            Some(v) => v,
            None => self.iter.next(),
        }
    }

    #[inline]
    #[rustc_inherit_overflow_checks]
    fn count(mut self) -> usize {
        match self.peeked.take() {
            Some(None) => 0,
            Some(Some(_)) => 1 + self.iter.count(),
            None => self.iter.count(),
        }
    }

    #[inline]
    fn nth(&mut self, n: usize) -> Option<I::Item> {
        match self.peeked.take() {
            Some(None) => None,
            Some(v @ Some(_)) if n == 0 => v,
            Some(Some(_)) => self.iter.nth(n - 1),
            None => self.iter.nth(n),
        }
    }

    #[inline]
    fn last(mut self) -> Option<I::Item> {
        let peek_opt = match self.peeked.take() {
            Some(None) => return None,
            Some(v) => v,
            None => None,
        };
        self.iter.last().or(peek_opt)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let peek_len = match self.peeked {
            Some(None) => return (0, Some(0)),
            Some(Some(_)) => 1,
            None => 0,
        };
        let (lo, hi) = self.iter.size_hint();
        let lo = lo.saturating_add(peek_len);
        let hi = match hi {
            Some(x) => x.checked_add(peek_len),
            None => None,
        };
        (lo, hi)
    }

    #[inline]
    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> R,
        R: Try<Ok = B>,
    {
        let acc = match self.peeked.take() {
            Some(None) => return try { init },
            Some(Some(v)) => f(init, v)?,
            None => init,
        };
        self.iter.try_fold(acc, f)
    }

    #[inline]
    fn fold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
    where
        Fold: FnMut(Acc, Self::Item) -> Acc,
    {
        let acc = match self.peeked {
            Some(None) => return init,
            Some(Some(v)) => fold(init, v),
            None => init,
        };
        self.iter.fold(acc, fold)
    }
}

#[stable(feature = "double_ended_peek_iterator", since = "1.38.0")]
impl<I> DoubleEndedIterator for Peekable<I>
where
    I: DoubleEndedIterator,
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        match self.peeked.as_mut() {
            Some(v @ Some(_)) => self.iter.next_back().or_else(|| v.take()),
            Some(None) => None,
            None => self.iter.next_back(),
        }
    }

    #[inline]
    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> R,
        R: Try<Ok = B>,
    {
        match self.peeked.take() {
            Some(None) => try { init },
            Some(Some(v)) => match self.iter.try_rfold(init, &mut f).into_result() {
                Ok(acc) => f(acc, v),
                Err(e) => {
                    self.peeked = Some(Some(v));
                    Try::from_error(e)
                }
            },
            None => self.iter.try_rfold(init, f),
        }
    }

    #[inline]
    fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
    where
        Fold: FnMut(Acc, Self::Item) -> Acc,
    {
        match self.peeked {
            Some(None) => init,
            Some(Some(v)) => {
                let acc = self.iter.rfold(init, &mut fold);
                fold(acc, v)
            }
            None => self.iter.rfold(init, fold),
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<I: ExactSizeIterator> ExactSizeIterator for Peekable<I> {}

#[stable(feature = "fused", since = "1.26.0")]
impl<I: FusedIterator> FusedIterator for Peekable<I> {}

impl<I: Iterator> Peekable<I> {
    /// Returns a reference to the next() value without advancing the iterator.
    ///
    /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
    /// But if the iteration is over, `None` is returned.
    ///
    /// [`next`]: Iterator::next
    ///
    /// Because `peek()` returns a reference, and many iterators iterate over
    /// references, there can be a possibly confusing situation where the
    /// return value is a double reference. You can see this effect in the
    /// examples below.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let xs = [1, 2, 3];
    ///
    /// let mut iter = xs.iter().peekable();
    ///
    /// // peek() lets us see into the future
    /// assert_eq!(iter.peek(), Some(&&1));
    /// assert_eq!(iter.next(), Some(&1));
    ///
    /// assert_eq!(iter.next(), Some(&2));
    ///
    /// // The iterator does not advance even if we `peek` multiple times
    /// assert_eq!(iter.peek(), Some(&&3));
    /// assert_eq!(iter.peek(), Some(&&3));
    ///
    /// assert_eq!(iter.next(), Some(&3));
    ///
    /// // After the iterator is finished, so is `peek()`
    /// assert_eq!(iter.peek(), None);
    /// assert_eq!(iter.next(), None);
    /// ```
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn peek(&mut self) -> Option<&I::Item> {
        let iter = &mut self.iter;
        self.peeked.get_or_insert_with(|| iter.next()).as_ref()
    }

    /// Returns a mutable reference to the next() value without advancing the iterator.
    ///
    /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
    /// But if the iteration is over, `None` is returned.
    ///
    /// Because `peek_mut()` returns a reference, and many iterators iterate over
    /// references, there can be a possibly confusing situation where the
    /// return value is a double reference. You can see this effect in the examples
    /// below.
    ///
    /// [`next`]: Iterator::next
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let mut iter = [1, 2, 3].iter().peekable();
    ///
    /// // Like with `peek()`, we can see into the future without advancing the iterator.
    /// assert_eq!(iter.peek_mut(), Some(&mut &1));
    /// assert_eq!(iter.peek_mut(), Some(&mut &1));
    /// assert_eq!(iter.next(), Some(&1));
    ///
    /// // Peek into the iterator and set the value behind the mutable reference.
    /// if let Some(p) = iter.peek_mut() {
    ///     assert_eq!(*p, &2);
    ///     *p = &5;
    /// }
    ///
    /// // The value we put in reappears as the iterator continues.
    /// assert_eq!(iter.collect::<Vec<_>>(), vec![&5, &3]);
    /// ```
    #[inline]
    #[stable(feature = "peekable_peek_mut", since = "1.52.0")]
    pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
        let iter = &mut self.iter;
        self.peeked.get_or_insert_with(|| iter.next()).as_mut()
    }

    /// Consume and return the next value of this iterator if a condition is true.
    ///
    /// If `func` returns `true` for the next value of this iterator, consume and return it.
    /// Otherwise, return `None`.
    ///
    /// # Examples
    /// Consume a number if it's equal to 0.
    /// ```
    /// let mut iter = (0..5).peekable();
    /// // The first item of the iterator is 0; consume it.
    /// assert_eq!(iter.next_if(|&x| x == 0), Some(0));
    /// // The next item returned is now 1, so `consume` will return `false`.
    /// assert_eq!(iter.next_if(|&x| x == 0), None);
    /// // `next_if` saves the value of the next item if it was not equal to `expected`.
    /// assert_eq!(iter.next(), Some(1));
    /// ```
    ///
    /// Consume any number less than 10.
    /// ```
    /// let mut iter = (1..20).peekable();
    /// // Consume all numbers less than 10
    /// while iter.next_if(|&x| x < 10).is_some() {}
    /// // The next value returned will be 10
    /// assert_eq!(iter.next(), Some(10));
    /// ```
    #[stable(feature = "peekable_next_if", since = "1.51.0")]
    pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> {
        match self.next() {
            Some(matched) if func(&matched) => Some(matched),
            other => {
                // Since we called `self.next()`, we consumed `self.peeked`.
                assert!(self.peeked.is_none());
                self.peeked = Some(other);
                None
            }
        }
    }

    /// Consume and return the next item if it is equal to `expected`.
    ///
    /// # Example
    /// Consume a number if it's equal to 0.
    /// ```
    /// let mut iter = (0..5).peekable();
    /// // The first item of the iterator is 0; consume it.
    /// assert_eq!(iter.next_if_eq(&0), Some(0));
    /// // The next item returned is now 1, so `consume` will return `false`.
    /// assert_eq!(iter.next_if_eq(&0), None);
    /// // `next_if_eq` saves the value of the next item if it was not equal to `expected`.
    /// assert_eq!(iter.next(), Some(1));
    /// ```
    #[stable(feature = "peekable_next_if", since = "1.51.0")]
    pub fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item>
    where
        T: ?Sized,
        I::Item: PartialEq<T>,
    {
        self.next_if(|next| next == expected)
    }
}

#[unstable(feature = "trusted_len", issue = "37572")]
unsafe impl<I> TrustedLen for Peekable<I> where I: TrustedLen {}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<S: Iterator, I: Iterator> SourceIter for Peekable<I>
where
    I: SourceIter<Source = S>,
{
    type Source = S;

    #[inline]
    unsafe fn as_inner(&mut self) -> &mut S {
        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
        unsafe { SourceIter::as_inner(&mut self.iter) }
    }
}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<I: InPlaceIterable> InPlaceIterable for Peekable<I> {}
<commit_msg>Update library/core/src/iter/adapters/peekable.rs<commit_after>use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen};
use crate::ops::Try;

/// An iterator with a `peek()` that returns an optional reference to the next
/// element.
///
/// This `struct` is created by the [`peekable`] method on [`Iterator`]. See its
/// documentation for more.
///
/// [`peekable`]: Iterator::peekable
/// [`Iterator`]: trait.Iterator.html
#[derive(Clone, Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Peekable<I: Iterator> {
    iter: I,
    /// Remember a peeked value, even if it was None.
    peeked: Option<Option<I::Item>>,
}

impl<I: Iterator> Peekable<I> {
    pub(in crate::iter) fn new(iter: I) -> Peekable<I> {
        Peekable { iter, peeked: None }
    }
}

// Peekable must remember if a None has been seen in the `.peek()` method.
// It ensures that `.peek(); .peek();` or `.peek(); .next();` only advances the
// underlying iterator at most once. This does not by itself make the iterator
// fused.
#[stable(feature = "rust1", since = "1.0.0")]
impl<I: Iterator> Iterator for Peekable<I> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<I::Item> {
        match self.peeked.take() {
            Some(v) => v,
            None => self.iter.next(),
        }
    }

    #[inline]
    #[rustc_inherit_overflow_checks]
    fn count(mut self) -> usize {
        match self.peeked.take() {
            Some(None) => 0,
            Some(Some(_)) => 1 + self.iter.count(),
            None => self.iter.count(),
        }
    }

    #[inline]
    fn nth(&mut self, n: usize) -> Option<I::Item> {
        match self.peeked.take() {
            Some(None) => None,
            Some(v @ Some(_)) if n == 0 => v,
            Some(Some(_)) => self.iter.nth(n - 1),
            None => self.iter.nth(n),
        }
    }

    #[inline]
    fn last(mut self) -> Option<I::Item> {
        let peek_opt = match self.peeked.take() {
            Some(None) => return None,
            Some(v) => v,
            None => None,
        };
        self.iter.last().or(peek_opt)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let peek_len = match self.peeked {
            Some(None) => return (0, Some(0)),
            Some(Some(_)) => 1,
            None => 0,
        };
        let (lo, hi) = self.iter.size_hint();
        let lo = lo.saturating_add(peek_len);
        let hi = match hi {
            Some(x) => x.checked_add(peek_len),
            None => None,
        };
        (lo, hi)
    }

    #[inline]
    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> R,
        R: Try<Ok = B>,
    {
        let acc = match self.peeked.take() {
            Some(None) => return try { init },
            Some(Some(v)) => f(init, v)?,
            None => init,
        };
        self.iter.try_fold(acc, f)
    }

    #[inline]
    fn fold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
    where
        Fold: FnMut(Acc, Self::Item) -> Acc,
    {
        let acc = match self.peeked {
            Some(None) => return init,
            Some(Some(v)) => fold(init, v),
            None => init,
        };
        self.iter.fold(acc, fold)
    }
}

#[stable(feature = "double_ended_peek_iterator", since = "1.38.0")]
impl<I> DoubleEndedIterator for Peekable<I>
where
    I: DoubleEndedIterator,
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        match self.peeked.as_mut() {
            Some(v @ Some(_)) => self.iter.next_back().or_else(|| v.take()),
            Some(None) => None,
            None => self.iter.next_back(),
        }
    }

    #[inline]
    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> R,
        R: Try<Ok = B>,
    {
        match self.peeked.take() {
            Some(None) => try { init },
            Some(Some(v)) => match self.iter.try_rfold(init, &mut f).into_result() {
                Ok(acc) => f(acc, v),
                Err(e) => {
                    self.peeked = Some(Some(v));
                    Try::from_error(e)
                }
            },
            None => self.iter.try_rfold(init, f),
        }
    }

    #[inline]
    fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
    where
        Fold: FnMut(Acc, Self::Item) -> Acc,
    {
        match self.peeked {
            Some(None) => init,
            Some(Some(v)) => {
                let acc = self.iter.rfold(init, &mut fold);
                fold(acc, v)
            }
            None => self.iter.rfold(init, fold),
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<I: ExactSizeIterator> ExactSizeIterator for Peekable<I> {}

#[stable(feature = "fused", since = "1.26.0")]
impl<I: FusedIterator> FusedIterator for Peekable<I> {}

impl<I: Iterator> Peekable<I> {
    /// Returns a reference to the next() value without advancing the iterator.
    ///
    /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
    /// But if the iteration is over, `None` is returned.
    ///
    /// [`next`]: Iterator::next
    ///
    /// Because `peek()` returns a reference, and many iterators iterate over
    /// references, there can be a possibly confusing situation where the
    /// return value is a double reference. You can see this effect in the
    /// examples below.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let xs = [1, 2, 3];
    ///
    /// let mut iter = xs.iter().peekable();
    ///
    /// // peek() lets us see into the future
    /// assert_eq!(iter.peek(), Some(&&1));
    /// assert_eq!(iter.next(), Some(&1));
    ///
    /// assert_eq!(iter.next(), Some(&2));
    ///
    /// // The iterator does not advance even if we `peek` multiple times
    /// assert_eq!(iter.peek(), Some(&&3));
    /// assert_eq!(iter.peek(), Some(&&3));
    ///
    /// assert_eq!(iter.next(), Some(&3));
    ///
    /// // After the iterator is finished, so is `peek()`
    /// assert_eq!(iter.peek(), None);
    /// assert_eq!(iter.next(), None);
    /// ```
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn peek(&mut self) -> Option<&I::Item> {
        let iter = &mut self.iter;
        self.peeked.get_or_insert_with(|| iter.next()).as_ref()
    }

    /// Returns a mutable reference to the next() value without advancing the iterator.
    ///
    /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
    /// But if the iteration is over, `None` is returned.
    ///
    /// Because `peek_mut()` returns a reference, and many iterators iterate over
    /// references, there can be a possibly confusing situation where the
    /// return value is a double reference. You can see this effect in the examples
    /// below.
    ///
    /// [`next`]: Iterator::next
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let mut iter = [1, 2, 3].iter().peekable();
    ///
    /// // Like with `peek()`, we can see into the future without advancing the iterator.
    /// assert_eq!(iter.peek_mut(), Some(&mut &1));
    /// assert_eq!(iter.peek_mut(), Some(&mut &1));
    /// assert_eq!(iter.next(), Some(&1));
    ///
    /// // Peek into the iterator and set the value behind the mutable reference.
    /// if let Some(p) = iter.peek_mut() {
    ///     assert_eq!(*p, &2);
    ///     *p = &5;
    /// }
    ///
    /// // The value we put in reappears as the iterator continues.
    /// assert_eq!(iter.collect::<Vec<_>>(), vec![&5, &3]);
    /// ```
    #[inline]
    #[stable(feature = "peekable_peek_mut", since = "1.53.0")]
    pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
        let iter = &mut self.iter;
        self.peeked.get_or_insert_with(|| iter.next()).as_mut()
    }

    /// Consume and return the next value of this iterator if a condition is true.
    ///
    /// If `func` returns `true` for the next value of this iterator, consume and return it.
    /// Otherwise, return `None`.
    ///
    /// # Examples
    /// Consume a number if it's equal to 0.
    /// ```
    /// let mut iter = (0..5).peekable();
    /// // The first item of the iterator is 0; consume it.
    /// assert_eq!(iter.next_if(|&x| x == 0), Some(0));
    /// // The next item returned is now 1, so `consume` will return `false`.
    /// assert_eq!(iter.next_if(|&x| x == 0), None);
    /// // `next_if` saves the value of the next item if it was not equal to `expected`.
    /// assert_eq!(iter.next(), Some(1));
    /// ```
    ///
    /// Consume any number less than 10.
    /// ```
    /// let mut iter = (1..20).peekable();
    /// // Consume all numbers less than 10
    /// while iter.next_if(|&x| x < 10).is_some() {}
    /// // The next value returned will be 10
    /// assert_eq!(iter.next(), Some(10));
    /// ```
    #[stable(feature = "peekable_next_if", since = "1.51.0")]
    pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> {
        match self.next() {
            Some(matched) if func(&matched) => Some(matched),
            other => {
                // Since we called `self.next()`, we consumed `self.peeked`.
                assert!(self.peeked.is_none());
                self.peeked = Some(other);
                None
            }
        }
    }

    /// Consume and return the next item if it is equal to `expected`.
    ///
    /// # Example
    /// Consume a number if it's equal to 0.
    /// ```
    /// let mut iter = (0..5).peekable();
    /// // The first item of the iterator is 0; consume it.
    /// assert_eq!(iter.next_if_eq(&0), Some(0));
    /// // The next item returned is now 1, so `consume` will return `false`.
    /// assert_eq!(iter.next_if_eq(&0), None);
    /// // `next_if_eq` saves the value of the next item if it was not equal to `expected`.
    /// assert_eq!(iter.next(), Some(1));
    /// ```
    #[stable(feature = "peekable_next_if", since = "1.51.0")]
    pub fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item>
    where
        T: ?Sized,
        I::Item: PartialEq<T>,
    {
        self.next_if(|next| next == expected)
    }
}

#[unstable(feature = "trusted_len", issue = "37572")]
unsafe impl<I> TrustedLen for Peekable<I> where I: TrustedLen {}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<S: Iterator, I: Iterator> SourceIter for Peekable<I>
where
    I: SourceIter<Source = S>,
{
    type Source = S;

    #[inline]
    unsafe fn as_inner(&mut self) -> &mut S {
        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
        unsafe { SourceIter::as_inner(&mut self.iter) }
    }
}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<I: InPlaceIterable> InPlaceIterable for Peekable<I> {}
<|endoftext|>


<commit_before>fn main() {
    let _ = main as usize as *const fn();
}

// END RUST SOURCE
// START rustc.main.ConstProp.before.mir
//  bb0: {
//      ...
//      _3 = const main as fn() (Pointer(ReifyFnPointer));
//      _2 = move _3 as usize (Misc);
//      ...
//      _1 = move _2 as *const fn() (Misc);
//      ...
//  }
// END rustc.main.ConstProp.before.mir
// START rustc.main.ConstProp.after.mir
//  bb0: {
//      ...
//      _3 = const Scalar(AllocId(0).0x0) : fn();
//      _2 = move _3 as usize (Misc);
//      ...
//      _1 = move _2 as *const fn() (Misc);
//      ...
//  }
// END rustc.main.ConstProp.after.mir
<commit_msg>Fix reify_fn_ptr test as we now pretty print const function pointers.<commit_after>fn main() {
    let _ = main as usize as *const fn();
}

// END RUST SOURCE
// START rustc.main.ConstProp.before.mir
//  bb0: {
//      ...
//      _3 = const main as fn() (Pointer(ReifyFnPointer));
//      _2 = move _3 as usize (Misc);
//      ...
//      _1 = move _2 as *const fn() (Misc);
//      ...
//  }
// END rustc.main.ConstProp.before.mir
// START rustc.main.ConstProp.after.mir
//  bb0: {
//      ...
//      _3 = const main;
//      _2 = move _3 as usize (Misc);
//      ...
//      _1 = move _2 as *const fn() (Misc);
//      ...
//  }
// END rustc.main.ConstProp.after.mir
<|endoftext|>