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

use proc_macro2::TokenStream;
use quote::quote;
use syn::{
    parse::Parser,
    parse_quote,
    punctuated::Punctuated,
    token::{Comma, Unsafe},
    Attribute, FnArg, ForeignItem, Ident, ImplItem, Item, ReturnType,
};

use super::{
    unqualify::{unqualify_params, unqualify_ret_type},
    RsCodegenResult, Use,
};
use crate::{
    conversion::{
        analysis::fun::{
            ArgumentAnalysis, FnAnalysis, FnKind, MethodKind, RustRenameStrategy, UnsafetyNeeded,
        },
        api::ImplBlockDetails,
        codegen_rs::lifetime::add_lifetime_to_all_params,
    },
    types::{Namespace, QualifiedName},
};
use crate::{
    conversion::{api::FuncToConvert, codegen_rs::lifetime::add_explicit_lifetime_if_necessary},
    types::make_ident,
};

impl UnsafetyNeeded {
    fn bridge_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::None => None,
            _ => Some(parse_quote! { unsafe }),
        }
    }

    fn wrapper_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::Always => Some(parse_quote! { unsafe }),
            _ => None,
        }
    }
}

pub(super) fn gen_function(
    ns: &Namespace,
    fun: FuncToConvert,
    analysis: FnAnalysis,
    cpp_call_name: String,
) -> RsCodegenResult {
    if !analysis.generate_code {
        return RsCodegenResult::default();
    }
    let cxxbridge_name = analysis.cxxbridge_name;
    let rust_name = analysis.rust_name;
    let ret_type = analysis.ret_type;
    let param_details = analysis.param_details;
    let wrapper_function_needed = analysis.cpp_wrapper.is_some();
    let params = analysis.params;
    let vis = analysis.vis;
    let kind = analysis.kind;
    let doc_attr = fun.doc_attr;

    let mut cpp_name_attr = Vec::new();
    let mut impl_entry = None;
    let rust_name_attr: Vec<_> = match &analysis.rust_rename_strategy {
        RustRenameStrategy::RenameUsingRustAttr => Attribute::parse_outer
            .parse2(quote!(
                #[rust_name = #rust_name]
            ))
            .unwrap(),
        _ => Vec::new(),
    };
    let mut materialization = match kind {
        FnKind::Method(..) => None,
        FnKind::Function => match analysis.rust_rename_strategy {
            RustRenameStrategy::RenameInOutputMod(alias) => {
                Some(Use::UsedFromCxxBridgeWithAlias(alias))
            }
            _ => Some(Use::UsedFromCxxBridge),
        },
    };
    let any_param_needs_rust_conversion = param_details
        .iter()
        .any(|pd| pd.conversion.rust_work_needed());
    let rust_wrapper_needed = any_param_needs_rust_conversion
        || (cxxbridge_name != rust_name && matches!(kind, FnKind::Method(..)));
    let wrapper_unsafety = analysis.requires_unsafe.wrapper_token();
    if rust_wrapper_needed {
        match kind {
            FnKind::Method(ref type_name, MethodKind::Constructor) => {
                // Constructor.
                impl_entry = Some(generate_constructor_impl(
                    &param_details,
                    type_name,
                    &cxxbridge_name,
                    &rust_name,
                    &wrapper_unsafety,
                    &doc_attr,
                ));
            }
            FnKind::Method(ref type_name, ref method_kind) => {
                // Method, or static method.
                impl_entry = Some(generate_method_impl(
                    &param_details,
                    matches!(
                        method_kind,
                        MethodKind::MakeUnique | MethodKind::Constructor
                    ),
                    type_name,
                    &cxxbridge_name,
                    &rust_name,
                    &ret_type,
                    &wrapper_unsafety,
                    &doc_attr,
                ));
            }
            _ => {
                // Generate plain old function
                materialization = Some(Use::Custom(generate_function_impl(
                    &param_details,
                    &rust_name,
                    &ret_type,
                    &wrapper_unsafety,
                    &doc_attr,
                )));
            }
        }
    }
    if cxxbridge_name != cpp_call_name && !wrapper_function_needed {
        cpp_name_attr = Attribute::parse_outer
            .parse2(quote!(
                #[cxx_name = #cpp_call_name]
            ))
            .unwrap();
    }
    // In very rare occasions, we might need to give an explicit lifetime.
    let (lifetime_tokens, params, ret_type) =
        add_explicit_lifetime_if_necessary(&param_details, params, &ret_type);

    // Finally - namespace support. All the Types in everything
    // above this point are fully qualified. We need to unqualify them.
    // We need to do that _after_ the above wrapper_function_needed
    // work, because it relies upon spotting fully qualified names like
    // std::unique_ptr. However, after it's done its job, all such
    // well-known types should be unqualified already (e.g. just UniquePtr)
    // and the following code will act to unqualify only those types
    // which the user has declared.
    let params = unqualify_params(params);
    let ret_type = unqualify_ret_type(ret_type.into_owned());
    // And we need to make an attribute for the namespace that the function
    // itself is in.
    let namespace_attr = if ns.is_empty() || wrapper_function_needed {
        Vec::new()
    } else {
        let namespace_string = ns.to_string();
        Attribute::parse_outer
            .parse2(quote!(
                #[namespace = #namespace_string]
            ))
            .unwrap()
    };
    // At last, actually generate the cxx::bridge entry.
    let bridge_unsafety = analysis.requires_unsafe.bridge_token();
    let extern_c_mod_item = ForeignItem::Fn(parse_quote!(
        #(#namespace_attr)*
        #(#rust_name_attr)*
        #(#cpp_name_attr)*
        #doc_attr
        #vis #bridge_unsafety fn #cxxbridge_name #lifetime_tokens ( #params ) #ret_type;
    ));
    RsCodegenResult {
        extern_c_mod_items: vec![extern_c_mod_item],
        bridge_items: Vec::new(),
        global_items: Vec::new(),
        bindgen_mod_items: Vec::new(),
        impl_entry,
        materializations: materialization.into_iter().collect(),
        extern_rust_mod_items: Vec::new(),
    }
}

fn generate_arg_lists(
    param_details: &[ArgumentAnalysis],
    avoid_self: bool,
) -> (Punctuated<FnArg, Comma>, Vec<TokenStream>) {
    let mut wrapper_params: Punctuated<FnArg, Comma> = Punctuated::new();
    let mut arg_list = Vec::new();

    for pd in param_details {
        let type_name = pd.conversion.rust_wrapper_unconverted_type();
        let wrapper_arg_name = if pd.self_type.is_some() && !avoid_self {
            parse_quote!(self)
        } else {
            pd.name.clone()
        };
        wrapper_params.push(parse_quote!(
            #wrapper_arg_name: #type_name
        ));
        arg_list.push(pd.conversion.rust_conversion(wrapper_arg_name));
    }
    (wrapper_params, arg_list)
}

/// Generate an 'impl Type { methods-go-here }' item
#[allow(clippy::too_many_arguments)] // it's true, but probably best for now
fn generate_method_impl(
    param_details: &[ArgumentAnalysis],
    avoid_self: bool,
    impl_block_type_name: &QualifiedName,
    cxxbridge_name: &Ident,
    rust_name: &str,
    ret_type: &ReturnType,
    unsafety: &Option<Unsafe>,
    doc_attr: &Option<Attribute>,
) -> Box<ImplBlockDetails> {
    let (wrapper_params, arg_list) = generate_arg_lists(param_details, avoid_self);
    let (lifetime_tokens, wrapper_params, ret_type) =
        add_explicit_lifetime_if_necessary(param_details, wrapper_params, ret_type);
    let rust_name = make_ident(&rust_name);
    Box::new(ImplBlockDetails {
        item: ImplItem::Method(parse_quote! {
            #doc_attr
            pub #unsafety fn #rust_name #lifetime_tokens ( #wrapper_params ) #ret_type {
                cxxbridge::#cxxbridge_name ( #(#arg_list),* )
            }
        }),
        ty: impl_block_type_name.get_final_ident(),
    })
}

/// Generate a 'impl Type { methods-go-here }' item which is a constructor
/// for use with moveit traits.
fn generate_constructor_impl(
    param_details: &[ArgumentAnalysis],
    impl_block_type_name: &QualifiedName,
    cxxbridge_name: &Ident,
    rust_name: &str,
    unsafety: &Option<Unsafe>,
    doc_attr: &Option<Attribute>,
) -> Box<ImplBlockDetails> {
    let (wrapper_params, arg_list) = generate_arg_lists(param_details, true);
    let mut wrapper_params: Punctuated<FnArg, Comma> = wrapper_params.into_iter().skip(1).collect();
    let ptr_arg_name = &arg_list[0];
    let rust_name = make_ident(&rust_name);
    let any_references = param_details.iter().any(|pd| pd.was_reference);
    let (lifetime_param, lifetime_addition) = if any_references {
        add_lifetime_to_all_params(&mut wrapper_params);
        (quote! { <'a> }, quote! { + 'a })
    } else {
        (quote! {}, quote! {})
    };
    Box::new(ImplBlockDetails {
        item: ImplItem::Method(parse_quote! {
            #doc_attr
            pub #unsafety fn #rust_name #lifetime_param ( #wrapper_params ) -> impl autocxx::moveit::new::New<Output=Self> #lifetime_addition {
                unsafe {
                    autocxx::moveit::new::by_raw(move |#ptr_arg_name| {
                        let #ptr_arg_name = #ptr_arg_name.get_unchecked_mut().as_mut_ptr();
                        cxxbridge::#cxxbridge_name(#(#arg_list),* )
                    })
                }
            }
        }),
        ty: impl_block_type_name.get_final_ident(),
    })
}

/// Generate a function call wrapper
fn generate_function_impl(
    param_details: &[ArgumentAnalysis],
    rust_name: &str,
    ret_type: &ReturnType,
    unsafety: &Option<Unsafe>,
    doc_attr: &Option<Attribute>,
) -> Box<Item> {
    let (wrapper_params, arg_list) = generate_arg_lists(param_details, false);
    let rust_name = make_ident(&rust_name);
    Box::new(Item::Fn(parse_quote! {
        #doc_attr
        pub #unsafety fn #rust_name ( #wrapper_params ) #ret_type {
            cxxbridge::#rust_name ( #(#arg_list),* )
        }
    }))
}
<commit_msg>Fix unnecessary 'unsafe' warning.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::TokenStream;
use quote::quote;
use syn::{
    parse::Parser,
    parse_quote,
    punctuated::Punctuated,
    token::{Comma, Unsafe},
    Attribute, FnArg, ForeignItem, Ident, ImplItem, Item, ReturnType,
};

use super::{
    unqualify::{unqualify_params, unqualify_ret_type},
    RsCodegenResult, Use,
};
use crate::{
    conversion::{
        analysis::fun::{
            ArgumentAnalysis, FnAnalysis, FnKind, MethodKind, RustRenameStrategy, UnsafetyNeeded,
        },
        api::ImplBlockDetails,
        codegen_rs::lifetime::add_lifetime_to_all_params,
    },
    types::{Namespace, QualifiedName},
};
use crate::{
    conversion::{api::FuncToConvert, codegen_rs::lifetime::add_explicit_lifetime_if_necessary},
    types::make_ident,
};

impl UnsafetyNeeded {
    fn bridge_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::None => None,
            _ => Some(parse_quote! { unsafe }),
        }
    }

    fn wrapper_token(&self) -> Option<Unsafe> {
        match self {
            UnsafetyNeeded::Always => Some(parse_quote! { unsafe }),
            _ => None,
        }
    }
}

pub(super) fn gen_function(
    ns: &Namespace,
    fun: FuncToConvert,
    analysis: FnAnalysis,
    cpp_call_name: String,
) -> RsCodegenResult {
    if !analysis.generate_code {
        return RsCodegenResult::default();
    }
    let cxxbridge_name = analysis.cxxbridge_name;
    let rust_name = analysis.rust_name;
    let ret_type = analysis.ret_type;
    let param_details = analysis.param_details;
    let wrapper_function_needed = analysis.cpp_wrapper.is_some();
    let params = analysis.params;
    let vis = analysis.vis;
    let kind = analysis.kind;
    let doc_attr = fun.doc_attr;

    let mut cpp_name_attr = Vec::new();
    let mut impl_entry = None;
    let rust_name_attr: Vec<_> = match &analysis.rust_rename_strategy {
        RustRenameStrategy::RenameUsingRustAttr => Attribute::parse_outer
            .parse2(quote!(
                #[rust_name = #rust_name]
            ))
            .unwrap(),
        _ => Vec::new(),
    };
    let mut materialization = match kind {
        FnKind::Method(..) => None,
        FnKind::Function => match analysis.rust_rename_strategy {
            RustRenameStrategy::RenameInOutputMod(alias) => {
                Some(Use::UsedFromCxxBridgeWithAlias(alias))
            }
            _ => Some(Use::UsedFromCxxBridge),
        },
    };
    let any_param_needs_rust_conversion = param_details
        .iter()
        .any(|pd| pd.conversion.rust_work_needed());
    let rust_wrapper_needed = any_param_needs_rust_conversion
        || (cxxbridge_name != rust_name && matches!(kind, FnKind::Method(..)));
    let wrapper_unsafety = analysis.requires_unsafe.wrapper_token();
    if rust_wrapper_needed {
        match kind {
            FnKind::Method(ref type_name, MethodKind::Constructor) => {
                // Constructor.
                impl_entry = Some(generate_constructor_impl(
                    &param_details,
                    type_name,
                    &cxxbridge_name,
                    &rust_name,
                    &wrapper_unsafety,
                    &doc_attr,
                ));
            }
            FnKind::Method(ref type_name, ref method_kind) => {
                // Method, or static method.
                impl_entry = Some(generate_method_impl(
                    &param_details,
                    matches!(
                        method_kind,
                        MethodKind::MakeUnique | MethodKind::Constructor
                    ),
                    type_name,
                    &cxxbridge_name,
                    &rust_name,
                    &ret_type,
                    &wrapper_unsafety,
                    &doc_attr,
                ));
            }
            _ => {
                // Generate plain old function
                materialization = Some(Use::Custom(generate_function_impl(
                    &param_details,
                    &rust_name,
                    &ret_type,
                    &wrapper_unsafety,
                    &doc_attr,
                )));
            }
        }
    }
    if cxxbridge_name != cpp_call_name && !wrapper_function_needed {
        cpp_name_attr = Attribute::parse_outer
            .parse2(quote!(
                #[cxx_name = #cpp_call_name]
            ))
            .unwrap();
    }
    // In very rare occasions, we might need to give an explicit lifetime.
    let (lifetime_tokens, params, ret_type) =
        add_explicit_lifetime_if_necessary(&param_details, params, &ret_type);

    // Finally - namespace support. All the Types in everything
    // above this point are fully qualified. We need to unqualify them.
    // We need to do that _after_ the above wrapper_function_needed
    // work, because it relies upon spotting fully qualified names like
    // std::unique_ptr. However, after it's done its job, all such
    // well-known types should be unqualified already (e.g. just UniquePtr)
    // and the following code will act to unqualify only those types
    // which the user has declared.
    let params = unqualify_params(params);
    let ret_type = unqualify_ret_type(ret_type.into_owned());
    // And we need to make an attribute for the namespace that the function
    // itself is in.
    let namespace_attr = if ns.is_empty() || wrapper_function_needed {
        Vec::new()
    } else {
        let namespace_string = ns.to_string();
        Attribute::parse_outer
            .parse2(quote!(
                #[namespace = #namespace_string]
            ))
            .unwrap()
    };
    // At last, actually generate the cxx::bridge entry.
    let bridge_unsafety = analysis.requires_unsafe.bridge_token();
    let extern_c_mod_item = ForeignItem::Fn(parse_quote!(
        #(#namespace_attr)*
        #(#rust_name_attr)*
        #(#cpp_name_attr)*
        #doc_attr
        #vis #bridge_unsafety fn #cxxbridge_name #lifetime_tokens ( #params ) #ret_type;
    ));
    RsCodegenResult {
        extern_c_mod_items: vec![extern_c_mod_item],
        bridge_items: Vec::new(),
        global_items: Vec::new(),
        bindgen_mod_items: Vec::new(),
        impl_entry,
        materializations: materialization.into_iter().collect(),
        extern_rust_mod_items: Vec::new(),
    }
}

fn generate_arg_lists(
    param_details: &[ArgumentAnalysis],
    avoid_self: bool,
) -> (Punctuated<FnArg, Comma>, Vec<TokenStream>) {
    let mut wrapper_params: Punctuated<FnArg, Comma> = Punctuated::new();
    let mut arg_list = Vec::new();

    for pd in param_details {
        let type_name = pd.conversion.rust_wrapper_unconverted_type();
        let wrapper_arg_name = if pd.self_type.is_some() && !avoid_self {
            parse_quote!(self)
        } else {
            pd.name.clone()
        };
        wrapper_params.push(parse_quote!(
            #wrapper_arg_name: #type_name
        ));
        arg_list.push(pd.conversion.rust_conversion(wrapper_arg_name));
    }
    (wrapper_params, arg_list)
}

/// Generate an 'impl Type { methods-go-here }' item
#[allow(clippy::too_many_arguments)] // it's true, but probably best for now
fn generate_method_impl(
    param_details: &[ArgumentAnalysis],
    avoid_self: bool,
    impl_block_type_name: &QualifiedName,
    cxxbridge_name: &Ident,
    rust_name: &str,
    ret_type: &ReturnType,
    unsafety: &Option<Unsafe>,
    doc_attr: &Option<Attribute>,
) -> Box<ImplBlockDetails> {
    let (wrapper_params, arg_list) = generate_arg_lists(param_details, avoid_self);
    let (lifetime_tokens, wrapper_params, ret_type) =
        add_explicit_lifetime_if_necessary(param_details, wrapper_params, ret_type);
    let rust_name = make_ident(&rust_name);
    Box::new(ImplBlockDetails {
        item: ImplItem::Method(parse_quote! {
            #doc_attr
            pub #unsafety fn #rust_name #lifetime_tokens ( #wrapper_params ) #ret_type {
                cxxbridge::#cxxbridge_name ( #(#arg_list),* )
            }
        }),
        ty: impl_block_type_name.get_final_ident(),
    })
}

/// Generate a 'impl Type { methods-go-here }' item which is a constructor
/// for use with moveit traits.
fn generate_constructor_impl(
    param_details: &[ArgumentAnalysis],
    impl_block_type_name: &QualifiedName,
    cxxbridge_name: &Ident,
    rust_name: &str,
    unsafety: &Option<Unsafe>,
    doc_attr: &Option<Attribute>,
) -> Box<ImplBlockDetails> {
    let (wrapper_params, arg_list) = generate_arg_lists(param_details, true);
    let mut wrapper_params: Punctuated<FnArg, Comma> = wrapper_params.into_iter().skip(1).collect();
    let ptr_arg_name = &arg_list[0];
    let rust_name = make_ident(&rust_name);
    let any_references = param_details.iter().any(|pd| pd.was_reference);
    let (lifetime_param, lifetime_addition) = if any_references {
        add_lifetime_to_all_params(&mut wrapper_params);
        (quote! { <'a> }, quote! { + 'a })
    } else {
        (quote! {}, quote! {})
    };
    let body = quote! {
        autocxx::moveit::new::by_raw(move |#ptr_arg_name| {
            let #ptr_arg_name = #ptr_arg_name.get_unchecked_mut().as_mut_ptr();
            cxxbridge::#cxxbridge_name(#(#arg_list),* )
        })
    };
    let body = if unsafety.is_some() {
        // No need for `unsafe` inside the function
        body
    } else {
        quote! {
            unsafe { #body }
        }
    };
    Box::new(ImplBlockDetails {
        item: ImplItem::Method(parse_quote! {
            #doc_attr
            pub #unsafety fn #rust_name #lifetime_param ( #wrapper_params ) -> impl autocxx::moveit::new::New<Output=Self> #lifetime_addition {
                #body
            }
        }),
        ty: impl_block_type_name.get_final_ident(),
    })
}

/// Generate a function call wrapper
fn generate_function_impl(
    param_details: &[ArgumentAnalysis],
    rust_name: &str,
    ret_type: &ReturnType,
    unsafety: &Option<Unsafe>,
    doc_attr: &Option<Attribute>,
) -> Box<Item> {
    let (wrapper_params, arg_list) = generate_arg_lists(param_details, false);
    let rust_name = make_ident(&rust_name);
    Box::new(Item::Fn(parse_quote! {
        #doc_attr
        pub #unsafety fn #rust_name ( #wrapper_params ) #ret_type {
            cxxbridge::#rust_name ( #(#arg_list),* )
        }
    }))
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use address::Address;
use instruction::Instruction;
use petgraph::EdgeDirection;

/// When is this edge taken? Conditionally or unconditionally?
#[derive(Clone,Copy,Debug,Eq,PartialEq)]
pub enum EdgeType {
    /// This edge is taken when a conditional is true.
    ConditionalTrue,
    /// This edge is taken when a conditional is false.
    ConditionalFalse,
    /// This is edge is always taken.
    Unconditional,
}

/// A [basic block] is a sequence of instructions with no inward-bound
/// branches except to the entry point and no outward-bound branches
/// except at the exit.
///
/// [basic block]: https://en.wikipedia.org/wiki/Basic_block
#[derive(Debug)]
pub struct BasicBlock<'f> {
    /// The name of the basic block. Not all blocks have meaningful names.
    pub name: Option<String>,
    /// The address of the first instruction in the basic block.
    pub address: Address,
    /// The instructions within this basic block.
    pub instructions: Vec<&'f Box<Instruction>>,
}

impl<'f> BasicBlock<'f> {
    /// Construct a new `BasicBlock`.
    pub fn new(name: Option<&str>, address: Address) -> Self {
        BasicBlock {
            name: name.map(|s| s.to_owned()),
            address: address,
            instructions: vec![],
        }
    }
}

/// Information about an edge between 2 [basic blocks].
///
/// This represents a branch, jump or other form of control flow
/// transfer within the control flow graph.
///
/// [basic blocks]: struct.BasicBlock.html
#[derive(Debug)]
pub struct BasicBlockEdge {
    /// Is this edge taken [conditionally or unconditionally]?
    ///
    /// [conditionally or unconditionally]: enum.EdgeType.html
    pub edge_type: EdgeType,
    /// Is this an inbound or outbound edge?
    pub direction: EdgeDirection,
}
<commit_msg>Rename conditional edge types.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use address::Address;
use instruction::Instruction;
use petgraph::EdgeDirection;

/// When is this edge taken? Conditionally or unconditionally?
#[derive(Clone,Copy,Debug,Eq,PartialEq)]
pub enum EdgeType {
    /// This edge is when a conditional is taken.
    ConditionalTaken,
    /// This edge is when a conditional falls through to the next instruction.
    ConditionalFallthrough,
    /// This is edge is always taken.
    Unconditional,
}

/// A [basic block] is a sequence of instructions with no inward-bound
/// branches except to the entry point and no outward-bound branches
/// except at the exit.
///
/// [basic block]: https://en.wikipedia.org/wiki/Basic_block
#[derive(Debug)]
pub struct BasicBlock<'f> {
    /// The name of the basic block. Not all blocks have meaningful names.
    pub name: Option<String>,
    /// The address of the first instruction in the basic block.
    pub address: Address,
    /// The instructions within this basic block.
    pub instructions: Vec<&'f Box<Instruction>>,
}

impl<'f> BasicBlock<'f> {
    /// Construct a new `BasicBlock`.
    pub fn new(name: Option<&str>, address: Address) -> Self {
        BasicBlock {
            name: name.map(|s| s.to_owned()),
            address: address,
            instructions: vec![],
        }
    }
}

/// Information about an edge between 2 [basic blocks].
///
/// This represents a branch, jump or other form of control flow
/// transfer within the control flow graph.
///
/// [basic blocks]: struct.BasicBlock.html
#[derive(Debug)]
pub struct BasicBlockEdge {
    /// Is this edge taken [conditionally or unconditionally]?
    ///
    /// [conditionally or unconditionally]: enum.EdgeType.html
    pub edge_type: EdgeType,
    /// Is this an inbound or outbound edge?
    pub direction: EdgeDirection,
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86_64` execution contexts.
//!
//! This is inteded to be general-purpose and composable, so that the same
//! code can be reused for interrupts and for multithreading.

use core::mem;
use core::fmt;

/// Registers pushed to the stack when handling an interrupt or context switch.
#[repr(C, packed)]
#[derive(Copy, Clone)]
pub struct Registers { pub rsi: u64
                     , pub rdi: u64
                     , pub r11: u64
                     , pub r10: u64
                     , pub r9:  u64
                     , pub r8:  u64
                     , pub rdx: u64
                     , pub rcx: u64
                     , pub rax: u64
                     }

 impl Registers {
     /// Transform this struct into an array of `u64`s
     /// (if you would ever want to do this)
     pub unsafe fn to_array(&self) -> [u64; 9] {
        //  [ self.rsi, self.rdi, self.r11
        //  , self.r10, self.r9, self.r8
        //  , self.rdx, self.rcx, self.rax
        //  ]
        // using transmute is probably faster and we're already unsafe...
        mem::transmute(*self)
     }

     /// Create a new empty set of Registers
     pub const fn empty() -> Self {
         Registers { rsi: 0, rdi: 0, r11: 0
                   , r10: 0, r9:  0, r8:  0
                   , rdx: 0, rcx: 0, rax: 0
                   }
     }

     /// Push the caller-saved registers to the stack
     /// (such as when handling a context switch or interrupt).
     ///
     /// THIS FUNCTION IS NAKED. DO NOT CALL IT NORMALLY.
     #[naked]
     #[inline(always)]
     pub unsafe fn push() {
         asm!( "push rax
                push rcx
                push rdx
                push r8
                push r9
                push r10
                push r11
                push rdi
                push rsi"
             :::: "intel"
                , "volatile");
     }

     /// Push the caller-saved registers off the stack
     /// (such as when handling a context switch or interrupt).
     ///
     /// THIS FUNCTION IS NAKED. DO NOT CALL IT NORMALLY.
     #[naked]
     #[inline(always)]
     pub unsafe fn pop() {
         asm!( "pop rsi
                pop rdi
                pop r11
                pop r10
                pop r9
                pop r8
                pop rdx
                pop rcx
                pop rax"
             :::: "intel"
                , "volatile");
     }
 }

 impl fmt::Debug for Registers {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!( f
               , "    RSI: {:#018x} RDI: {:#018x} R11: {:#018x}\n    \
                      R10: {:#018x} R9:  {:#018x} R8:  {:#018x}\n    \
                      RDX: {:#018x} RCX: {:#018x} RAX: {:#018x}"
              , self.rsi, self.rdi, self.r11
              , self.r10, self.r9,  self.r8
              , self.rdx, self.rcx, self.rax)
     }
 }

/// Thread execution context
#[repr(C, packed)]
pub struct Context { /// Value of the stack pointer (`rsp`) register
                     pub rsp: *mut u8
                   , /// Value of the caller-saved registers
                     pub registers: Registers
                   , /// Value of the instruction pointer (`rip`) register
                     pub rip: *mut u8
                 //, pub stack: [u8] // TODO: should be box
                   }

impl Context {
    pub fn empty() -> Self {
        unsafe {
            Context { rsp: mem::transmute(0u64)
                    , registers: Registers::empty()
                    , rip: mem::transmute(0u64)
                  //, stack: [0u8; 8]
                    }
        }
    }
}
<commit_msg>[x86_64] start on fixing stack frame alignment in interrupts<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86_64` execution contexts.
//!
//! This is inteded to be general-purpose and composable, so that the same
//! code can be reused for interrupts and for multithreading.

use core::mem;
use core::fmt;
use super::flags::Flags;
use super::segment;

/// Registers pushed to the stack when handling an interrupt or context switch.
#[repr(C, packed)]
#[derive(Copy, Clone)]
pub struct Registers { pub rsi: u64
                     , pub rdi: u64
                     , pub r11: u64
                     , pub r10: u64
                     , pub r9:  u64
                     , pub r8:  u64
                     , pub rdx: u64
                     , pub rcx: u64
                     , pub rax: u64
                     }

 impl Registers {
     /// Transform this struct into an array of `u64`s
     /// (if you would ever want to do this)
     pub unsafe fn to_array(&self) -> [u64; 9] {
        //  [ self.rsi, self.rdi, self.r11
        //  , self.r10, self.r9, self.r8
        //  , self.rdx, self.rcx, self.rax
        //  ]
        // using transmute is probably faster and we're already unsafe...
        mem::transmute(*self)
     }

     /// Create a new empty set of Registers
     pub const fn empty() -> Self {
         Registers { rsi: 0, rdi: 0, r11: 0
                   , r10: 0, r9:  0, r8:  0
                   , rdx: 0, rcx: 0, rax: 0
                   }
     }

     /// Push the caller-saved registers to the stack
     /// (such as when handling a context switch or interrupt).
     ///
     /// THIS FUNCTION IS NAKED. DO NOT CALL IT NORMALLY.
     #[naked]
     #[inline(always)]
     pub unsafe fn push() {
         asm!( "push rax
                push rcx
                push rdx
                push r8
                push r9
                push r10
                push r11
                push rdi
                push rsi"
             :::: "intel"
                , "volatile");
     }

     /// Push the caller-saved registers off the stack
     /// (such as when handling a context switch or interrupt).
     ///
     /// THIS FUNCTION IS NAKED. DO NOT CALL IT NORMALLY.
     #[naked]
     #[inline(always)]
     pub unsafe fn pop() {
         asm!( "pop rsi
                pop rdi
                pop r11
                pop r10
                pop r9
                pop r8
                pop rdx
                pop rcx
                pop rax"
             :::: "intel"
                , "volatile");
     }
 }

 impl fmt::Debug for Registers {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!( f
               , "    RSI: {:#018x} RDI: {:#018x} R11: {:#018x}\n    \
                      R10: {:#018x} R9:  {:#018x} R8:  {:#018x}\n    \
                      RDX: {:#018x} RCX: {:#018x} RAX: {:#018x}"
              , self.rsi, self.rdi, self.r11
              , self.r10, self.r9,  self.r8
              , self.rdx, self.rcx, self.rax)
     }
}

#[derive(Debug)]
#[repr(C)]
pub struct InterruptFrame {
    //  this is the actual value of the interrupt stack frame context,
    //  not the old one (which is wrong). note that the old one seems to cause
    //  stack misalignment.
    //          -- eliza, october 4th, 2016
    /// Value of the instruction pointer (`$rip`) register
    pub rip: u64
  , __pad_1: u32
  , __pad_2: u16
  , /// Value of the code segment (`$cs`) register
    pub cs: segment::Selector
  , /// Value of the CPU flags (`$rflags`) register
    pub rflags: Flags
  , /// Value of the stack pointer (`$rsp`) register
    //  TODO: should this actually be a pointer?
    pub rsp: u64
  , __pad_3: u32
  , __pad_4: u16
  , /// Value of the stack segment (`$ss`) register
    pub ss: segment::Selector
}

/// Thread execution context
#[repr(C, packed)]
pub struct Context { /// Value of the stack pointer (`rsp`) register
                     pub rsp: *mut u8
                   , /// Value of the caller-saved registers
                     pub registers: Registers
                   , /// Value of the instruction pointer (`rip`) register
                     pub rip: *mut u8
                 //, pub stack: [u8] // TODO: should be box
                   }

impl Context {
    pub fn empty() -> Self {
        unsafe {
            Context { rsp: mem::transmute(0u64)
                    , registers: Registers::empty()
                    , rip: mem::transmute(0u64)
                  //, stack: [0u8; 8]
                    }
        }
    }
}
<|endoftext|>


<commit_before>use core::convert::{Into, TryInto};
use core::fmt;
use core::ops::{Add, AddAssign, Sub, SubAssign};

use bit_field::BitField;
use usize_conversions::{usize_from, FromUsize};
use ux::*;

/// A canonical 64-bit virtual memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 48 lower bits of a virtual address can be used. The top 16 bits need
/// to be copies of bit 47, i.e. the most significant bit. Addresses that fulfil this criterium
/// are called “canonical”. This type guarantees that it always represents a canonical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct VirtAddr(u64);

/// A 64-bit physical memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 52 lower bits of a physical address can be used. The top 12 bits need
/// to be zero. This type guarantees that it always represents a valid physical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct PhysAddr(u64);

/// A passed `u64` was not a valid virtual address.
///
/// This means that bits 48 to 64 are not
/// a valid sign extension and are not null either. So automatic sign extension would have
/// overwritten possibly meaningful bits. This likely indicates a bug, for example an invalid
/// address calculation.
#[derive(Debug)]
pub struct VirtAddrNotValid(u64);

impl VirtAddr {
    /// Creates a new canonical virtual address.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical. Panics
    /// if the bits in the range 48 to 64 contain data (i.e. are not null and no sign extension).
    pub fn new(addr: u64) -> VirtAddr {
        Self::try_new(addr).expect(
            "address passed to VirtAddr::new must not contain any data \
             in bits 48 to 64",
        )
    }

    /// Tries to create a new canonical virtual address.
    ///
    /// This function tries to performs sign
    /// extension of bit 47 to make the address canonical. It succeeds if bits 48 to 64 are
    /// either a correct sign extension (i.e. copies of bit 47) or all null. Else, an error
    /// is returned.
    pub fn try_new(addr: u64) -> Result<VirtAddr, VirtAddrNotValid> {
        match addr.get_bits(47..64) {
            0 | 0x1ffff => Ok(VirtAddr(addr)),      // address is canonical
            1 => Ok(VirtAddr::new_unchecked(addr)), // address needs sign extension
            other => Err(VirtAddrNotValid(other)),
        }
    }

    /// Creates a new canonical virtual address without checks.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical, so
    /// bits 48 to 64 are overwritten. If you want to check that these bits contain no data,
    /// use `new` or `try_new`.
    pub fn new_unchecked(mut addr: u64) -> VirtAddr {
        if addr.get_bit(47) {
            addr.set_bits(48..64, 0xffff);
        } else {
            addr.set_bits(48..64, 0);
        }
        VirtAddr(addr)
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    pub fn as_ptr<T>(self) -> *const T {
        usize_from(self.as_u64()) as *const T
    }

    pub fn as_mut_ptr<T>(self) -> *mut T {
        self.as_ptr::<T>() as *mut T
    }

    /// Aligns the virtual address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        VirtAddr(align_up(self.0, align.into()))
    }

    /// Aligns the virtual address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        VirtAddr(align_down(self.0, align.into()))
    }

    /// Returns the 12-bit page offset of this virtual address.
    pub fn page_offset(&self) -> u12 {
        u12::new((self.0 & 0xfff).try_into().unwrap())
    }

    /// Returns the 9-bit level 1 page table index.
    pub fn p1_index(&self) -> u9 {
        u9::new(((self.0 >> 12) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 2 page table index.
    pub fn p2_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 3 page table index.
    pub fn p3_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 4 page table index.
    pub fn p4_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9 >> 9) & 0o777).try_into().unwrap())
    }
}

impl fmt::Debug for VirtAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "VirtAddr({:#x})", self.0)
    }
}

impl Add<u64> for VirtAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for VirtAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Add<usize> for VirtAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn add(self, rhs: usize) -> Self::Output {
        self + u64::from_usize(rhs)
    }
}

impl AddAssign<usize> for VirtAddr
where
    u64: FromUsize,
{
    fn add_assign(&mut self, rhs: usize) {
        self.add_assign(u64::from_usize(rhs))
    }
}

impl Sub<u64> for VirtAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0.checked_sub(rhs).unwrap())
    }
}

impl SubAssign<u64> for VirtAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<usize> for VirtAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn sub(self, rhs: usize) -> Self::Output {
        self - u64::from_usize(rhs)
    }
}

impl SubAssign<usize> for VirtAddr
where
    u64: FromUsize,
{
    fn sub_assign(&mut self, rhs: usize) {
        self.sub_assign(u64::from_usize(rhs))
    }
}

impl Sub<VirtAddr> for VirtAddr {
    type Output = u64;
    fn sub(self, rhs: VirtAddr) -> Self::Output {
        self.as_u64().checked_sub(rhs.as_u64()).unwrap()
    }
}

/// A passed `u64` was not a valid physical address.
///
/// This means that bits 52 to 64 are not were not all null.
#[derive(Debug)]
pub struct PhysAddrNotValid(u64);

impl PhysAddr {
    /// Creates a new physical address.
    ///
    /// Panics if a bit in the range 52 to 64 is set.
    pub fn new(addr: u64) -> PhysAddr {
        assert_eq!(
            addr.get_bits(52..64),
            0,
            "physical addresses must not have any bits in the range 52 to 64 set"
        );
        PhysAddr(addr)
    }

    /// Tries to create a new physical address.
    ///
    /// Fails if any bits in the range 52 to 64 are set.
    pub fn try_new(addr: u64) -> Result<PhysAddr, PhysAddrNotValid> {
        match addr.get_bits(52..64) {
            0 => Ok(PhysAddr(addr)), // address is valid
            other => Err(PhysAddrNotValid(other)),
        }
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    /// Convenience method for checking if a physical address is null.
    pub fn is_null(&self) -> bool {
        self.0 == 0
    }

    /// Aligns the physical address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        PhysAddr(align_up(self.0, align.into()))
    }

    /// Aligns the physical address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        PhysAddr(align_down(self.0, align.into()))
    }
}

impl fmt::Debug for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "PhysAddr({:#x})", self.0)
    }
}

impl fmt::Binary for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::LowerHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::Octal for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::UpperHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Add<u64> for PhysAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for PhysAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Add<usize> for PhysAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn add(self, rhs: usize) -> Self::Output {
        self + u64::from_usize(rhs)
    }
}

impl AddAssign<usize> for PhysAddr
where
    u64: FromUsize,
{
    fn add_assign(&mut self, rhs: usize) {
        self.add_assign(u64::from_usize(rhs))
    }
}

impl Sub<u64> for PhysAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0.checked_sub(rhs).unwrap())
    }
}

impl SubAssign<u64> for PhysAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<usize> for PhysAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn sub(self, rhs: usize) -> Self::Output {
        self - u64::from_usize(rhs)
    }
}

impl SubAssign<usize> for PhysAddr
where
    u64: FromUsize,
{
    fn sub_assign(&mut self, rhs: usize) {
        self.sub_assign(u64::from_usize(rhs))
    }
}

impl Sub<PhysAddr> for PhysAddr {
    type Output = u64;
    fn sub(self, rhs: PhysAddr) -> Self::Output {
        self.as_u64().checked_sub(rhs.as_u64()).unwrap()
    }
}

/// Align address downwards.
///
/// Returns the greatest x with alignment `align` so that x <= addr. The alignment must be
///  a power of 2.
pub fn align_down(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    addr & !(align - 1)
}

/// Align address upwards.
///
/// Returns the smallest x with alignment `align` so that x >= addr. The alignment must be
/// a power of 2.
pub fn align_up(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    let align_mask = align - 1;
    if addr & align_mask == 0 {
        addr // already aligned
    } else {
        (addr | align_mask) + 1
    }
}

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

    #[test]
    pub fn test_align_up() {
        // align 1
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(1234, 1), 1234);
        assert_eq!(align_up(0xffffffffffffffff, 1), 0xffffffffffffffff);
        // align 2
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(1233, 2), 1234);
        assert_eq!(align_up(0xfffffffffffffffe, 2), 0xfffffffffffffffe);
        // address 0
        assert_eq!(align_up(0, 128), 0);
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(0, 0x8000000000000000), 0);
    }
}
<commit_msg>Add {Virt, Phys}Addr::is_aligned method<commit_after>use core::convert::{Into, TryInto};
use core::fmt;
use core::ops::{Add, AddAssign, Sub, SubAssign};

use bit_field::BitField;
use usize_conversions::{usize_from, FromUsize};
use ux::*;

/// A canonical 64-bit virtual memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 48 lower bits of a virtual address can be used. The top 16 bits need
/// to be copies of bit 47, i.e. the most significant bit. Addresses that fulfil this criterium
/// are called “canonical”. This type guarantees that it always represents a canonical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct VirtAddr(u64);

/// A 64-bit physical memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 52 lower bits of a physical address can be used. The top 12 bits need
/// to be zero. This type guarantees that it always represents a valid physical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct PhysAddr(u64);

/// A passed `u64` was not a valid virtual address.
///
/// This means that bits 48 to 64 are not
/// a valid sign extension and are not null either. So automatic sign extension would have
/// overwritten possibly meaningful bits. This likely indicates a bug, for example an invalid
/// address calculation.
#[derive(Debug)]
pub struct VirtAddrNotValid(u64);

impl VirtAddr {
    /// Creates a new canonical virtual address.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical. Panics
    /// if the bits in the range 48 to 64 contain data (i.e. are not null and no sign extension).
    pub fn new(addr: u64) -> VirtAddr {
        Self::try_new(addr).expect(
            "address passed to VirtAddr::new must not contain any data \
             in bits 48 to 64",
        )
    }

    /// Tries to create a new canonical virtual address.
    ///
    /// This function tries to performs sign
    /// extension of bit 47 to make the address canonical. It succeeds if bits 48 to 64 are
    /// either a correct sign extension (i.e. copies of bit 47) or all null. Else, an error
    /// is returned.
    pub fn try_new(addr: u64) -> Result<VirtAddr, VirtAddrNotValid> {
        match addr.get_bits(47..64) {
            0 | 0x1ffff => Ok(VirtAddr(addr)),      // address is canonical
            1 => Ok(VirtAddr::new_unchecked(addr)), // address needs sign extension
            other => Err(VirtAddrNotValid(other)),
        }
    }

    /// Creates a new canonical virtual address without checks.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical, so
    /// bits 48 to 64 are overwritten. If you want to check that these bits contain no data,
    /// use `new` or `try_new`.
    pub fn new_unchecked(mut addr: u64) -> VirtAddr {
        if addr.get_bit(47) {
            addr.set_bits(48..64, 0xffff);
        } else {
            addr.set_bits(48..64, 0);
        }
        VirtAddr(addr)
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    pub fn as_ptr<T>(self) -> *const T {
        usize_from(self.as_u64()) as *const T
    }

    pub fn as_mut_ptr<T>(self) -> *mut T {
        self.as_ptr::<T>() as *mut T
    }

    /// Aligns the virtual address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        VirtAddr(align_up(self.0, align.into()))
    }

    /// Aligns the virtual address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        VirtAddr(align_down(self.0, align.into()))
    }

    /// Checks whether the virtual address has the demanded alignment.
    pub fn is_aligned<U>(self, align: U) -> bool
    where
        U: Into<u64>,
    {
        self.align_down(align) == self
    }

    /// Returns the 12-bit page offset of this virtual address.
    pub fn page_offset(&self) -> u12 {
        u12::new((self.0 & 0xfff).try_into().unwrap())
    }

    /// Returns the 9-bit level 1 page table index.
    pub fn p1_index(&self) -> u9 {
        u9::new(((self.0 >> 12) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 2 page table index.
    pub fn p2_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 3 page table index.
    pub fn p3_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 4 page table index.
    pub fn p4_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9 >> 9) & 0o777).try_into().unwrap())
    }
}

impl fmt::Debug for VirtAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "VirtAddr({:#x})", self.0)
    }
}

impl Add<u64> for VirtAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for VirtAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Add<usize> for VirtAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn add(self, rhs: usize) -> Self::Output {
        self + u64::from_usize(rhs)
    }
}

impl AddAssign<usize> for VirtAddr
where
    u64: FromUsize,
{
    fn add_assign(&mut self, rhs: usize) {
        self.add_assign(u64::from_usize(rhs))
    }
}

impl Sub<u64> for VirtAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0.checked_sub(rhs).unwrap())
    }
}

impl SubAssign<u64> for VirtAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<usize> for VirtAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn sub(self, rhs: usize) -> Self::Output {
        self - u64::from_usize(rhs)
    }
}

impl SubAssign<usize> for VirtAddr
where
    u64: FromUsize,
{
    fn sub_assign(&mut self, rhs: usize) {
        self.sub_assign(u64::from_usize(rhs))
    }
}

impl Sub<VirtAddr> for VirtAddr {
    type Output = u64;
    fn sub(self, rhs: VirtAddr) -> Self::Output {
        self.as_u64().checked_sub(rhs.as_u64()).unwrap()
    }
}

/// A passed `u64` was not a valid physical address.
///
/// This means that bits 52 to 64 are not were not all null.
#[derive(Debug)]
pub struct PhysAddrNotValid(u64);

impl PhysAddr {
    /// Creates a new physical address.
    ///
    /// Panics if a bit in the range 52 to 64 is set.
    pub fn new(addr: u64) -> PhysAddr {
        assert_eq!(
            addr.get_bits(52..64),
            0,
            "physical addresses must not have any bits in the range 52 to 64 set"
        );
        PhysAddr(addr)
    }

    /// Tries to create a new physical address.
    ///
    /// Fails if any bits in the range 52 to 64 are set.
    pub fn try_new(addr: u64) -> Result<PhysAddr, PhysAddrNotValid> {
        match addr.get_bits(52..64) {
            0 => Ok(PhysAddr(addr)), // address is valid
            other => Err(PhysAddrNotValid(other)),
        }
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    /// Convenience method for checking if a physical address is null.
    pub fn is_null(&self) -> bool {
        self.0 == 0
    }

    /// Aligns the physical address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        PhysAddr(align_up(self.0, align.into()))
    }

    /// Aligns the physical address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
    where
        U: Into<u64>,
    {
        PhysAddr(align_down(self.0, align.into()))
    }

    /// Checks whether the physical address has the demanded alignment.
    pub fn is_aligned<U>(self, align: U) -> bool
    where
        U: Into<u64>,
    {
        self.align_down(align) == self
    }
}

impl fmt::Debug for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "PhysAddr({:#x})", self.0)
    }
}

impl fmt::Binary for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::LowerHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::Octal for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::UpperHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Add<u64> for PhysAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for PhysAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Add<usize> for PhysAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn add(self, rhs: usize) -> Self::Output {
        self + u64::from_usize(rhs)
    }
}

impl AddAssign<usize> for PhysAddr
where
    u64: FromUsize,
{
    fn add_assign(&mut self, rhs: usize) {
        self.add_assign(u64::from_usize(rhs))
    }
}

impl Sub<u64> for PhysAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0.checked_sub(rhs).unwrap())
    }
}

impl SubAssign<u64> for PhysAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<usize> for PhysAddr
where
    u64: FromUsize,
{
    type Output = Self;
    fn sub(self, rhs: usize) -> Self::Output {
        self - u64::from_usize(rhs)
    }
}

impl SubAssign<usize> for PhysAddr
where
    u64: FromUsize,
{
    fn sub_assign(&mut self, rhs: usize) {
        self.sub_assign(u64::from_usize(rhs))
    }
}

impl Sub<PhysAddr> for PhysAddr {
    type Output = u64;
    fn sub(self, rhs: PhysAddr) -> Self::Output {
        self.as_u64().checked_sub(rhs.as_u64()).unwrap()
    }
}

/// Align address downwards.
///
/// Returns the greatest x with alignment `align` so that x <= addr. The alignment must be
///  a power of 2.
pub fn align_down(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    addr & !(align - 1)
}

/// Align address upwards.
///
/// Returns the smallest x with alignment `align` so that x >= addr. The alignment must be
/// a power of 2.
pub fn align_up(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    let align_mask = align - 1;
    if addr & align_mask == 0 {
        addr // already aligned
    } else {
        (addr | align_mask) + 1
    }
}

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

    #[test]
    pub fn test_align_up() {
        // align 1
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(1234, 1), 1234);
        assert_eq!(align_up(0xffffffffffffffff, 1), 0xffffffffffffffff);
        // align 2
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(1233, 2), 1234);
        assert_eq!(align_up(0xfffffffffffffffe, 2), 0xfffffffffffffffe);
        // address 0
        assert_eq!(align_up(0, 128), 0);
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(0, 0x8000000000000000), 0);
    }
}
<|endoftext|>


<commit_before>// Parse ANSI attr code
use std::borrow::Cow;
use std::default::Default;
use std::mem;

use tuikit::prelude::*;
use vte::Perform;

/// An ANSI Parser, will parse one line at a time.
///
/// It will cache the latest attribute used, that means if an attribute affect multiple
/// lines, the parser will recognize it.
pub struct ANSIParser {
    partial_str: String,
    last_attr: Attr,

    stripped: String,
    fragments: Vec<(Attr, (u32, u32))>,
}

impl Default for ANSIParser {
    fn default() -> Self {
        ANSIParser {
            partial_str: String::new(),
            last_attr: Attr::default(),

            stripped: String::new(),
            fragments: Vec::new(),
        }
    }
}

impl Perform for ANSIParser {
    fn print(&mut self, ch: char) {
        self.partial_str.push(ch);
    }

    fn execute(&mut self, byte: u8) {
        match byte {
            // \b to delete character back
            0x08 => {
                self.partial_str.pop();
            }
            // put back \0 \r \n \t
            0x00 | 0x0d | 0x0A | 0x09 => self.partial_str.push(byte as char),
            // ignore all others
            _ => trace!("AnsiParser:execute ignored {:?}", byte),
        }
    }

    fn hook(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool) {
        trace!("AnsiParser:hook ignored {:?}", params);
    }

    fn put(&mut self, byte: u8) {
        trace!("AnsiParser:put ignored {:?}", byte);
    }

    fn unhook(&mut self) {
        trace!("AnsiParser:unhook ignored");
    }

    fn osc_dispatch(&mut self, params: &[&[u8]]) {
        trace!("AnsiParser:osc ignored {:?}", params);
    }

    fn csi_dispatch(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool, mode: char) {
        // https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_(Select_Graphic_Rendition)_parameters
        // Only care about graphic modes, ignore all others

        if mode != 'm' {
            trace!("ignore: params: {:?}, mode: {:?}", params, mode);
            return;
        }

        // \[[m => means reset
        let mut attr = if params.is_empty() {
            Attr::default()
        } else {
            self.last_attr
        };

        let mut iter = params.iter();
        while let Some(&code) = iter.next() {
            match code {
                0 => attr = Attr::default(),
                1 => attr.effect |= Effect::BOLD,
                2 => attr.effect |= !Effect::BOLD,
                4 => attr.effect |= Effect::UNDERLINE,
                5 => attr.effect |= Effect::BLINK,
                7 => attr.effect |= Effect::REVERSE,
                num if num >= 30 && num <= 37 => {
                    attr.fg = Color::AnsiValue((num - 30) as u8);
                }
                38 => match iter.next() {
                    Some(2) => {
                        // ESC[ 38;2;<r>;<g>;<b> m Select RGB foreground color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.fg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 38;5;<n> m Select foreground color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.fg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("error on parsing CSI {:?} m", params);
                    }
                },
                39 => attr.fg = Color::Default,
                num if num >= 40 && num <= 47 => {
                    attr.bg = Color::AnsiValue((num - 40) as u8);
                }
                48 => match iter.next() {
                    Some(2) => {
                        // ESC[ 48;2;<r>;<g>;<b> m Select RGB background color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.bg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 48;5;<n> m Select background color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.bg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("ignore CSI {:?} m", params);
                    }
                },
                49 => attr.bg = Color::Default,
                _ => {
                    trace!("ignore CSI {:?} m", params);
                }
            }
        }

        self.attr_change(attr);
    }

    fn esc_dispatch(&mut self, _params: &[i64], _intermediates: &[u8], _ignore: bool, _byte: u8) {
        // ESC characters are replaced with \[
        self.partial_str.push('"');
        self.partial_str.push('[');
    }
}

impl ANSIParser {
    /// save the partial_str into fragments with current attr
    fn save_str(&mut self) {
        if self.partial_str.is_empty() {
            return;
        }

        let string = mem::replace(&mut self.partial_str, String::new());
        self.fragments.push((
            self.last_attr,
            (self.stripped.len() as u32, (self.stripped.len() + string.len()) as u32),
        ));
        self.stripped.push_str(&string);
    }

    // accept a new attr
    fn attr_change(&mut self, new_attr: Attr) {
        if new_attr == self.last_attr {
            return;
        }

        self.save_str();
        self.last_attr = new_attr;
    }

    pub fn parse_ansi(&mut self, text: &str) -> AnsiString<'static> {
        let mut statemachine = vte::Parser::new();

        for byte in text.as_bytes() {
            statemachine.advance(self, *byte);
        }
        self.save_str();

        let stripped = mem::replace(&mut self.stripped, String::new());
        let fragments = mem::replace(&mut self.fragments, Vec::new());
        AnsiString::new(stripped, fragments)
    }
}

#[derive(Clone, Debug)]
/// A String that contains ANSI state (e.g. colors)
///
/// It is internally represented as Vec<(attr, string)>
pub struct AnsiString<'a> {
    stripped: Cow<'a, str>,
    // attr: start, end
    fragments: Vec<(Attr, (u32, u32))>,
}

impl<'a> AnsiString<'a> {
    pub fn new_empty() -> Self {
        Self {
            stripped: Cow::Borrowed(""),
            fragments: Vec::new(),
        }
    }

    fn new_string(string: String) -> Self {
        let string_len = string.len() as u32;
        Self {
            stripped: Cow::Owned(string),
            fragments: vec![(Attr::default(), (0, string_len))],
        }
    }

    fn new_str(str_ref: &'a str) -> Self {
        let string_len = str_ref.len() as u32;
        Self {
            stripped: Cow::Borrowed(str_ref),
            fragments: vec![(Attr::default(), (0, string_len))],
        }
    }

    fn new(stripped: String, fragments: Vec<(Attr, (u32, u32))>) -> Self {
        Self {
            stripped: Cow::Owned(stripped),
            fragments,
        }
    }

    pub fn parse(raw: &'a str) -> AnsiString<'static> {
        ANSIParser::default().parse_ansi(raw)
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.fragments.is_empty()
    }

    #[inline]
    pub fn into_inner(self) -> Cow<'a, str> {
        self.stripped
    }

    pub fn iter(&self) -> AnsiStringIterator {
        AnsiStringIterator::new(&self.stripped, &self.fragments)
    }

    pub fn has_attrs(&self) -> bool {
        // more than 1 fragments or is not default attr
        self.fragments.len() > 1 || (!self.fragments.is_empty() && self.fragments[0].0 != Attr::default())
    }

    #[inline]
    pub fn stripped(&self) -> &str {
        &self.stripped
    }
}

impl<'a> From<&'a str> for AnsiString<'a> {
    fn from(s: &'a str) -> AnsiString<'a> {
        AnsiString::new_str(s)
    }
}

impl From<String> for AnsiString<'static> {
    fn from(s: String) -> Self {
        AnsiString::new_string(s)
    }
}

/// An iterator over all the (char, attr) characters.
pub struct AnsiStringIterator<'a> {
    fragments: &'a [(Attr, (u32, u32))],
    fragment_idx: usize,
    chars_iter: std::str::CharIndices<'a>,
}

impl<'a> AnsiStringIterator<'a> {
    pub fn new(stripped: &'a str, fragments: &'a [(Attr, (u32, u32))]) -> Self {
        Self {
            fragments,
            fragment_idx: 0,
            chars_iter: stripped.char_indices(),
        }
    }
}

impl<'a> Iterator for AnsiStringIterator<'a> {
    type Item = (char, Attr);

    fn next(&mut self) -> Option<Self::Item> {
        match self.chars_iter.next() {
            Some((byte_idx, char)) => {
                // update fragment_idx
                loop {
                    if self.fragment_idx >= self.fragments.len() {
                        return None;
                    }

                    let (_attr, (_start, end)) = self.fragments[self.fragment_idx];
                    if byte_idx < (end as usize) {
                        break;
                    } else {
                        self.fragment_idx += 1;
                    }
                }

                let (attr, _) = self.fragments[self.fragment_idx];
                Some((char, attr))
            }
            None => None,
        }
    }
}

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

    #[test]
    fn test_ansi_iterator() {
        let input = "\x1B[48;2;5;10;15m\x1B[38;2;70;130;180mhi\x1B[0m";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::Rgb(70, 130, 180),
            bg: Color::Rgb(5, 10, 15),
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
    }

    #[test]
    fn test_normal_string() {
        let input = "ab";
        let ansistring = ANSIParser::default().parse_ansi(input);

        assert_eq!(false, ansistring.has_attrs());

        let mut it = ansistring.iter();
        assert_eq!(Some(('a', Attr::default())), it.next());
        assert_eq!(Some(('b', Attr::default())), it.next());
        assert_eq!(None, it.next());

        assert_eq!("ab", ansistring.into_inner())
    }

    #[test]
    fn test_multiple_attributes() {
        let input = "\x1B[1;31mhi";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::AnsiValue(1),
            effect: Effect::BOLD,
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
    }

    #[test]
    fn test_reset() {
        let input = "\x1B[35mA\x1B[mB";
        let ansistring = ANSIParser::default().parse_ansi(input);
        assert_eq!(ansistring.fragments.len(), 2);
    }
}
<commit_msg>[ansi] store attributes as None if no ansi is detected<commit_after>// Parse ANSI attr code
use std::borrow::Cow;
use std::default::Default;
use std::mem;

use tuikit::prelude::*;
use vte::Perform;

/// An ANSI Parser, will parse one line at a time.
///
/// It will cache the latest attribute used, that means if an attribute affect multiple
/// lines, the parser will recognize it.
pub struct ANSIParser {
    partial_str: String,
    last_attr: Attr,

    stripped: String,
    fragments: Vec<(Attr, (u32, u32))>,
}

impl Default for ANSIParser {
    fn default() -> Self {
        ANSIParser {
            partial_str: String::new(),
            last_attr: Attr::default(),

            stripped: String::new(),
            fragments: Vec::new(),
        }
    }
}

impl Perform for ANSIParser {
    fn print(&mut self, ch: char) {
        self.partial_str.push(ch);
    }

    fn execute(&mut self, byte: u8) {
        match byte {
            // \b to delete character back
            0x08 => {
                self.partial_str.pop();
            }
            // put back \0 \r \n \t
            0x00 | 0x0d | 0x0A | 0x09 => self.partial_str.push(byte as char),
            // ignore all others
            _ => trace!("AnsiParser:execute ignored {:?}", byte),
        }
    }

    fn hook(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool) {
        trace!("AnsiParser:hook ignored {:?}", params);
    }

    fn put(&mut self, byte: u8) {
        trace!("AnsiParser:put ignored {:?}", byte);
    }

    fn unhook(&mut self) {
        trace!("AnsiParser:unhook ignored");
    }

    fn osc_dispatch(&mut self, params: &[&[u8]]) {
        trace!("AnsiParser:osc ignored {:?}", params);
    }

    fn csi_dispatch(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool, mode: char) {
        // https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_(Select_Graphic_Rendition)_parameters
        // Only care about graphic modes, ignore all others

        if mode != 'm' {
            trace!("ignore: params: {:?}, mode: {:?}", params, mode);
            return;
        }

        // \[[m => means reset
        let mut attr = if params.is_empty() {
            Attr::default()
        } else {
            self.last_attr
        };

        let mut iter = params.iter();
        while let Some(&code) = iter.next() {
            match code {
                0 => attr = Attr::default(),
                1 => attr.effect |= Effect::BOLD,
                2 => attr.effect |= !Effect::BOLD,
                4 => attr.effect |= Effect::UNDERLINE,
                5 => attr.effect |= Effect::BLINK,
                7 => attr.effect |= Effect::REVERSE,
                num if num >= 30 && num <= 37 => {
                    attr.fg = Color::AnsiValue((num - 30) as u8);
                }
                38 => match iter.next() {
                    Some(2) => {
                        // ESC[ 38;2;<r>;<g>;<b> m Select RGB foreground color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.fg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 38;5;<n> m Select foreground color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.fg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("error on parsing CSI {:?} m", params);
                    }
                },
                39 => attr.fg = Color::Default,
                num if num >= 40 && num <= 47 => {
                    attr.bg = Color::AnsiValue((num - 40) as u8);
                }
                48 => match iter.next() {
                    Some(2) => {
                        // ESC[ 48;2;<r>;<g>;<b> m Select RGB background color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.bg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 48;5;<n> m Select background color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.bg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("ignore CSI {:?} m", params);
                    }
                },
                49 => attr.bg = Color::Default,
                _ => {
                    trace!("ignore CSI {:?} m", params);
                }
            }
        }

        self.attr_change(attr);
    }

    fn esc_dispatch(&mut self, _params: &[i64], _intermediates: &[u8], _ignore: bool, _byte: u8) {
        // ESC characters are replaced with \[
        self.partial_str.push('"');
        self.partial_str.push('[');
    }
}

impl ANSIParser {
    /// save the partial_str into fragments with current attr
    fn save_str(&mut self) {
        if self.partial_str.is_empty() {
            return;
        }

        let string = mem::replace(&mut self.partial_str, String::new());
        self.fragments.push((
            self.last_attr,
            (self.stripped.len() as u32, (self.stripped.len() + string.len()) as u32),
        ));
        self.stripped.push_str(&string);
    }

    // accept a new attr
    fn attr_change(&mut self, new_attr: Attr) {
        if new_attr == self.last_attr {
            return;
        }

        self.save_str();
        self.last_attr = new_attr;
    }

    pub fn parse_ansi(&mut self, text: &str) -> AnsiString<'static> {
        let mut statemachine = vte::Parser::new();

        for byte in text.as_bytes() {
            statemachine.advance(self, *byte);
        }
        self.save_str();

        let stripped = mem::replace(&mut self.stripped, String::new());
        let fragments = mem::replace(&mut self.fragments, Vec::new());
        AnsiString::new(stripped, fragments)
    }
}

#[derive(Clone, Debug)]
/// A String that contains ANSI state (e.g. colors)
///
/// It is internally represented as Vec<(attr, string)>
pub struct AnsiString<'a> {
    stripped: Cow<'a, str>,
    // attr: start, end
    fragments: Option<Vec<(Attr, (u32, u32))>>,
}

impl<'a> AnsiString<'a> {
    pub fn new_empty() -> Self {
        Self {
            stripped: Cow::Borrowed(""),
            fragments: None,
        }
    }

    fn new_string(string: String) -> Self {
        Self {
            stripped: Cow::Owned(string),
            fragments: None,
        }
    }

    fn new_str(str_ref: &'a str) -> Self {
        Self {
            stripped: Cow::Borrowed(str_ref),
            fragments: None,
        }
    }

    fn new(stripped: String, fragments: Vec<(Attr, (u32, u32))>) -> Self {
        let fragments_empty = fragments.is_empty() || (fragments.len() == 1 && fragments[0].0 == Attr::default());
        Self {
            stripped: Cow::Owned(stripped),
            fragments: if fragments_empty { None } else { Some(fragments) },
        }
    }

    pub fn parse(raw: &'a str) -> AnsiString<'static> {
        ANSIParser::default().parse_ansi(raw)
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.stripped.is_empty()
    }

    #[inline]
    pub fn into_inner(self) -> Cow<'a, str> {
        self.stripped
    }

    pub fn iter(&'a self) -> Box<dyn Iterator<Item = (char, Attr)> + 'a> {
        if self.fragments.is_none() {
            return Box::new(self.stripped.chars().map(|c| (c, Attr::default())));
        }

        Box::new(AnsiStringIterator::new(
            &self.stripped,
            self.fragments.as_ref().unwrap(),
        ))
    }

    pub fn has_attrs(&self) -> bool {
        self.fragments.is_some()
    }

    #[inline]
    pub fn stripped(&self) -> &str {
        &self.stripped
    }
}

impl<'a> From<&'a str> for AnsiString<'a> {
    fn from(s: &'a str) -> AnsiString<'a> {
        AnsiString::new_str(s)
    }
}

impl From<String> for AnsiString<'static> {
    fn from(s: String) -> Self {
        AnsiString::new_string(s)
    }
}

/// An iterator over all the (char, attr) characters.
pub struct AnsiStringIterator<'a> {
    fragments: &'a [(Attr, (u32, u32))],
    fragment_idx: usize,
    chars_iter: std::str::CharIndices<'a>,
}

impl<'a> AnsiStringIterator<'a> {
    pub fn new(stripped: &'a str, fragments: &'a [(Attr, (u32, u32))]) -> Self {
        Self {
            fragments,
            fragment_idx: 0,
            chars_iter: stripped.char_indices(),
        }
    }
}

impl<'a> Iterator for AnsiStringIterator<'a> {
    type Item = (char, Attr);

    fn next(&mut self) -> Option<Self::Item> {
        match self.chars_iter.next() {
            Some((byte_idx, char)) => {
                // update fragment_idx
                loop {
                    if self.fragment_idx >= self.fragments.len() {
                        return None;
                    }

                    let (_attr, (_start, end)) = self.fragments[self.fragment_idx];
                    if byte_idx < (end as usize) {
                        break;
                    } else {
                        self.fragment_idx += 1;
                    }
                }

                let (attr, _) = self.fragments[self.fragment_idx];
                Some((char, attr))
            }
            None => None,
        }
    }
}

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

    #[test]
    fn test_ansi_iterator() {
        let input = "\x1B[48;2;5;10;15m\x1B[38;2;70;130;180mhi\x1B[0m";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::Rgb(70, 130, 180),
            bg: Color::Rgb(5, 10, 15),
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
        assert_eq!(ansistring.stripped(), "hi");
    }

    #[test]
    fn test_normal_string() {
        let input = "ab";
        let ansistring = ANSIParser::default().parse_ansi(input);

        assert_eq!(false, ansistring.has_attrs());

        let mut it = ansistring.iter();
        assert_eq!(Some(('a', Attr::default())), it.next());
        assert_eq!(Some(('b', Attr::default())), it.next());
        assert_eq!(None, it.next());

        assert_eq!(ansistring.stripped(), "ab");
    }

    #[test]
    fn test_multiple_attributes() {
        let input = "\x1B[1;31mhi";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::AnsiValue(1),
            effect: Effect::BOLD,
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
        assert_eq!(ansistring.stripped(), "hi");
    }

    #[test]
    fn test_reset() {
        let input = "\x1B[35mA\x1B[mB";
        let ansistring = ANSIParser::default().parse_ansi(input);
        assert_eq!(ansistring.fragments.as_ref().map(|x| x.len()).unwrap(), 2);
        assert_eq!(ansistring.stripped(), "AB");
    }
}
<|endoftext|>