<commit_before>use wasm_bindgen::prelude::*;

use anyhow::Error;
use lancelot::{
    analysis::dis::{self, zydis},
    arch::Arch,
    aspace::AddressSpace,
    loader::pe::{PEError, PE as lPE},
    module::ModuleError,
    pagemap::PageMapError,
    util::UtilError,
};

#[cfg(feature = "wee_alloc")]
#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

#[wasm_bindgen]
pub fn version() -> String {
    String::from("v0.7.0")
}

#[wasm_bindgen(start)]
pub fn main_js() -> Result<(), JsValue> {
    #[cfg(debug_assertions)]
    console_error_panic_hook::set_once();

    Ok(())
}

// for better function signatures
// no functional difference.
type JsError = JsValue;

fn to_value_error(e: anyhow::Error) -> JsError {
    js_sys::Error::new(&format!("{}", e)).into()
}

fn to_js_err(e: Error) -> JsError {
    match e.downcast_ref::<PEError>() {
        Some(PEError::FormatNotSupported(_)) => return to_value_error(e),
        Some(PEError::MalformedPEFile(_)) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<ModuleError>() {
        Some(ModuleError::InvalidAddress(_)) => return to_value_error(e),
        None => (),
    };

    match e.downcast_ref::<UtilError>() {
        Some(UtilError::FileAccess) => return to_value_error(e),
        Some(UtilError::FileFormat) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<PageMapError>() {
        Some(PageMapError::NotMapped) => return to_value_error(e),
        None => (),
    };

    to_value_error(e)
}

#[wasm_bindgen]
pub fn from_bytes(buf: Vec<u8>) -> Result<PE, JsError> {
    let pe = lPE::from_bytes(&buf).map_err(to_js_err)?;
    let dec = dis::get_disassembler(&pe.module).map_err(to_js_err)?;
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("valid formatter options");
    Ok(PE {
        inner: pe,
        decoder: dec,
        formatter,
    })
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub struct AddressRange {
    pub start: u64,
    pub end:   u64,
}

#[wasm_bindgen]
impl AddressRange {
    #[wasm_bindgen(getter)]
    pub fn size(&self) -> u64 {
        self.end - self.start
    }
}

#[wasm_bindgen]
pub struct Section {
    pub physical_range: AddressRange,
    pub virtual_range:  AddressRange,
    // String are not clone, so we have to manually implement getters
    permissions:        String,
    name:               String,
}

#[wasm_bindgen]
impl Section {
    #[wasm_bindgen(getter)]
    pub fn permissions(&self) -> String {
        self.permissions.to_string()
    }

    #[wasm_bindgen(getter)]
    pub fn name(&self) -> String {
        self.name.clone()
    }
}

impl From<&lancelot::module::Section> for Section {
    fn from(other: &lancelot::module::Section) -> Section {
        Section {
            physical_range: AddressRange {
                start: other.physical_range.start,
                end:   other.physical_range.end,
            },
            virtual_range:  AddressRange {
                start: other.virtual_range.start,
                end:   other.virtual_range.end,
            },
            permissions:    other.permissions.to_string(),
            name:           other.name.to_string(),
        }
    }
}

#[wasm_bindgen]
pub struct Instruction {
    pub address: u64,
    pub size:    u8,
    bytes:       Vec<u8>,
    string:      String,
}

#[wasm_bindgen]
impl Instruction {
    #[wasm_bindgen(getter)]
    pub fn bytes(&self) -> Vec<u8> {
        self.bytes.clone()
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[wasm_bindgen]
pub struct PE {
    inner:     lPE,
    decoder:   zydis::Decoder,
    formatter: zydis::Formatter,
}

#[wasm_bindgen]
impl PE {
    #[wasm_bindgen(getter)]
    pub fn arch(&self) -> JsValue {
        match self.inner.module.arch {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            Arch::X32 => JsValue::from(wasm_bindgen::intern("x32")),
            Arch::X64 => JsValue::from(wasm_bindgen::intern("x64")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn base_address(&self) -> u64 {
        self.inner.module.address_space.base_address
    }

    // not cheap, so its a function, not getter
    #[wasm_bindgen]
    pub fn functions(&self) -> Result<Vec<u64>, JsError> {
        Ok(lancelot::analysis::pe::find_functions(&self.inner)
            .map_err(to_js_err)?
            .into_iter()
            .filter(|f| matches!(f, lancelot::analysis::pe::Function::Local(_)))
            .map(|f| match f {
                lancelot::analysis::pe::Function::Local(va) => va,
                _ => unreachable!(),
            })
            .collect())
    }

    // Vec<T> cannot be serialized by wasm-bindgen
    // so we have to manually convert to Vec<JsValue>
    // and then annotate the TS type.
    //
    // NB: JS now owns the objects, must explicitly drop.
    #[wasm_bindgen(getter, typescript_type = "Array<Section>")]
    pub fn sections(&self) -> Vec<JsValue> {
        self.inner
            .module
            .sections
            .iter()
            .map(Section::from)
            .map(JsValue::from)
            .collect()
    }

    #[wasm_bindgen]
    pub fn read_bytes(&self, va: u64, size: usize) -> Result<Vec<u8>, JsError> {
        self.inner.module.address_space.read_bytes(va, size).map_err(to_js_err)
    }

    #[wasm_bindgen]
    pub fn read_insn(&self, va: u64) -> Result<Instruction, JsError> {
        let mut insn_buf = [0u8; 16];
        self.inner
            .module
            .address_space
            .read_into(va, &mut insn_buf)
            .map_err(to_js_err)?;

        if let Ok(Some(insn)) = self.decoder.decode(&insn_buf) {
            let mut out_buf = [0u8; 200];
            let mut out_buf = zydis::OutputBuffer::new(&mut out_buf[..]);
            self.formatter
                .format_instruction(&insn, &mut out_buf, Some(va), None)
                .map_err(|e| -> JsError {
                    js_sys::Error::new(&format!("failed to format instruction: {}", e)).into()
                })?;

            Ok(Instruction {
                address: va,
                size:    insn.length,
                bytes:   insn_buf[..insn.length as usize].to_vec(),
                string:  out_buf
                    .as_str()
                    .expect("disassembly should always be utf-8")
                    .to_string(),
            })
        } else {
            Err(js_sys::Error::new("invalid instruction").into())
        }
    }
}
<commit_msg>jslancelot: strings<commit_after>use wasm_bindgen::prelude::*;

use anyhow::Error;
use lancelot::{
    analysis::dis::{self, zydis},
    arch::Arch,
    aspace::AddressSpace,
    loader::pe::{PEError, PE as lPE},
    module::ModuleError,
    pagemap::PageMapError,
    util::UtilError,
};

#[cfg(feature = "wee_alloc")]
#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

#[wasm_bindgen]
pub fn version() -> String {
    String::from("v0.7.0")
}

#[wasm_bindgen(start)]
pub fn main_js() -> Result<(), JsValue> {
    #[cfg(debug_assertions)]
    console_error_panic_hook::set_once();

    Ok(())
}

// for better function signatures
// no functional difference.
type JsError = JsValue;

fn to_value_error(e: anyhow::Error) -> JsError {
    js_sys::Error::new(&format!("{}", e)).into()
}

fn to_js_err(e: Error) -> JsError {
    match e.downcast_ref::<PEError>() {
        Some(PEError::FormatNotSupported(_)) => return to_value_error(e),
        Some(PEError::MalformedPEFile(_)) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<ModuleError>() {
        Some(ModuleError::InvalidAddress(_)) => return to_value_error(e),
        None => (),
    };

    match e.downcast_ref::<UtilError>() {
        Some(UtilError::FileAccess) => return to_value_error(e),
        Some(UtilError::FileFormat) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<PageMapError>() {
        Some(PageMapError::NotMapped) => return to_value_error(e),
        None => (),
    };

    to_value_error(e)
}

#[wasm_bindgen]
pub fn from_bytes(buf: Vec<u8>) -> Result<PE, JsError> {
    let pe = lPE::from_bytes(&buf).map_err(to_js_err)?;
    let dec = dis::get_disassembler(&pe.module).map_err(to_js_err)?;
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("valid formatter options");
    Ok(PE {
        inner: pe,
        decoder: dec,
        formatter,
    })
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub struct AddressRange {
    pub start: u64,
    pub end:   u64,
}

#[wasm_bindgen]
impl AddressRange {
    #[wasm_bindgen(getter)]
    pub fn size(&self) -> u64 {
        self.end - self.start
    }
}

#[wasm_bindgen]
pub struct Section {
    pub physical_range: AddressRange,
    pub virtual_range:  AddressRange,
    // String are not clone, so we have to manually implement getters
    permissions:        String,
    name:               String,
}

#[wasm_bindgen]
impl Section {
    #[wasm_bindgen(getter)]
    pub fn permissions(&self) -> String {
        self.permissions.to_string()
    }

    #[wasm_bindgen(getter)]
    pub fn name(&self) -> String {
        self.name.clone()
    }
}

impl From<&lancelot::module::Section> for Section {
    fn from(other: &lancelot::module::Section) -> Section {
        Section {
            physical_range: AddressRange {
                start: other.physical_range.start,
                end:   other.physical_range.end,
            },
            virtual_range:  AddressRange {
                start: other.virtual_range.start,
                end:   other.virtual_range.end,
            },
            permissions:    other.permissions.to_string(),
            name:           other.name.to_string(),
        }
    }
}

#[wasm_bindgen]
pub struct Instruction {
    pub address: u64,
    pub size:    u8,
    bytes:       Vec<u8>,
    string:      String,
}

#[wasm_bindgen]
impl Instruction {
    #[wasm_bindgen(getter)]
    pub fn bytes(&self) -> Vec<u8> {
        self.bytes.clone()
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[allow(clippy::upper_case_acronyms)]
enum StringEncoding {
    ASCII,
    UTF16,
}

#[wasm_bindgen(js_name = "String")]
pub struct EncodedString {
    encoding:    StringEncoding,
    pub address: u64,
    pub size:    usize,
    string:      String,
}

#[wasm_bindgen(js_class = "String")]
impl EncodedString {
    #[wasm_bindgen(getter, typescript_type = "string")]
    pub fn encoding(&self) -> JsValue {
        match self.encoding {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            StringEncoding::ASCII => JsValue::from(wasm_bindgen::intern("ASCII")),
            StringEncoding::UTF16 => JsValue::from(wasm_bindgen::intern("UTF-16")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[wasm_bindgen]
pub struct PE {
    inner:     lPE,
    decoder:   zydis::Decoder,
    formatter: zydis::Formatter,
}

#[wasm_bindgen]
impl PE {
    #[wasm_bindgen(getter, typescript_type = "string")]
    pub fn arch(&self) -> JsValue {
        match self.inner.module.arch {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            Arch::X32 => JsValue::from(wasm_bindgen::intern("x32")),
            Arch::X64 => JsValue::from(wasm_bindgen::intern("x64")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn base_address(&self) -> u64 {
        self.inner.module.address_space.base_address
    }

    // not cheap, so its a function, not getter
    #[wasm_bindgen]
    pub fn functions(&self) -> Result<Vec<u64>, JsError> {
        Ok(lancelot::analysis::pe::find_functions(&self.inner)
            .map_err(to_js_err)?
            .into_iter()
            .filter(|f| matches!(f, lancelot::analysis::pe::Function::Local(_)))
            .map(|f| match f {
                lancelot::analysis::pe::Function::Local(va) => va,
                _ => unreachable!(),
            })
            .collect())
    }

    // Vec<T> cannot be serialized by wasm-bindgen
    // so we have to manually convert to Vec<JsValue>
    // and then annotate the TS type.
    //
    // NB: JS now owns the objects, must explicitly drop.
    #[wasm_bindgen(getter, typescript_type = "Array<Section>")]
    pub fn sections(&self) -> Vec<JsValue> {
        self.inner
            .module
            .sections
            .iter()
            .map(Section::from)
            .map(JsValue::from)
            .collect()
    }

    #[wasm_bindgen]
    pub fn read_bytes(&self, va: u64, size: usize) -> Result<Vec<u8>, JsError> {
        self.inner.module.address_space.read_bytes(va, size).map_err(to_js_err)
    }

    #[wasm_bindgen]
    pub fn read_insn(&self, va: u64) -> Result<Instruction, JsError> {
        let mut insn_buf = [0u8; 16];
        self.inner
            .module
            .address_space
            .read_into(va, &mut insn_buf)
            .map_err(to_js_err)?;

        if let Ok(Some(insn)) = self.decoder.decode(&insn_buf) {
            let mut out_buf = [0u8; 200];
            let mut out_buf = zydis::OutputBuffer::new(&mut out_buf[..]);
            self.formatter
                .format_instruction(&insn, &mut out_buf, Some(va), None)
                .map_err(|e| -> JsError {
                    js_sys::Error::new(&format!("failed to format instruction: {}", e)).into()
                })?;

            Ok(Instruction {
                address: va,
                size:    insn.length,
                bytes:   insn_buf[..insn.length as usize].to_vec(),
                string:  out_buf
                    .as_str()
                    .expect("disassembly should always be utf-8")
                    .to_string(),
            })
        } else {
            Err(js_sys::Error::new("invalid instruction").into())
        }
    }

    #[wasm_bindgen(typescript_type = "Array<EncodedString>")]
    pub fn strings(&self) -> Vec<JsValue> {
        let mut ret: Vec<EncodedString> = Default::default();

        for (range, s) in lancelot::util::find_ascii_strings(&self.inner.buf) {
            if let Ok(va) = self.inner.module.virtual_address(range.start as u64) {
                ret.push(EncodedString {
                    encoding: StringEncoding::ASCII,
                    address:  va,
                    size:     range.end - range.start,
                    string:   s,
                });
            }
        }

        for (range, s) in lancelot::util::find_unicode_strings(&self.inner.buf) {
            if let Ok(va) = self.inner.module.virtual_address(range.start as u64) {
                ret.push(EncodedString {
                    encoding: StringEncoding::UTF16,
                    address:  va,
                    size:     range.end - range.start,
                    string:   s,
                });
            }
        }

        ret.sort_by_key(|string| string.address);

        ret.into_iter().map(JsValue::from).collect()
    }
}
<|endoftext|>


<commit_before>#[allow(dead_code)]
fn guess_this () -> i32 {
    let one = 5;
    let two = 7;
    let three = 3;
    let result = (one + two) / three;
    return result;
}

mod tests {
    use super::*;

    pub fn test_complicated () {
        verify!(1, guess_this(), "Complicated example");
    }
}

pub fn exec () {
    tests::test_complicated();
}
<commit_msg>add another example<commit_after>fn guess_this () -> i32 {
    let one = 5;
    let two = 7;
    let three = 3;
    let result = (one + two) / three;
    return result;
}

fn simple () -> &'static str {
    let hello = "Hello World!";
    return hello;
}

mod tests {
    use super::*;

    pub fn test_simple () {
        verify!("Hello World!", simple(), "Simple example");
    }

    pub fn test_complicated () {
        verify!(1, guess_this(), "Complicated example");
    }
}

pub fn exec () {
    tests::test_simple();
    tests::test_complicated();
}
<|endoftext|>


<commit_before>use env::Env;
use gobjects::*;
use library;
use nameutil::*;
use super::*;
use super::type_kind::*;

pub struct Info {
    pub full_name: String,
    pub class_tid: library::TypeId,
    pub kind: type_kind::TypeKind,
    pub name: String,
    pub parents: Vec<general::StatusedTypeId>,
    pub has_children: bool,
    pub has_ignored_parents: bool,
    pub functions: Vec<functions::Info>,
    pub has_constructors: bool,
    pub has_methods: bool,
}

impl Info {
    //TODO: add test in tests/ for panic
    pub fn type_<'a>(&self, library: &'a library::Library) -> &'a library::Class {
        let type_ = library.type_(self.class_tid).as_class()
            .unwrap_or_else(|| panic!("{} is not a class.", self.full_name));
        type_
    }

    ///TODO: return iterator
    pub fn constructors(&self) -> Vec<&functions::Info> {
        self.functions.iter()
            .filter(|f| f.kind == library::FunctionKind::Constructor)
            .collect()
    }

    pub fn methods(&self) -> Vec<&functions::Info> {
        self.functions.iter()
            .filter(|f| f.kind == library::FunctionKind::Method)
            .collect()
    }
}

pub fn new(env: &Env, obj: &GObject) -> Info {
    let full_name = obj.name.clone();

    let class_tid = env.library.find_type_unwrapped(0, &full_name, "Class");

    let type_ = env.library.type_(class_tid);
    let kind = type_.to_type_kind(&env.library);
    assert_eq!(kind, TypeKind::Widget);

    let name = split_namespace_name(&full_name).1.into();

    let klass = type_.to_class();
    let (parents, has_ignored_parents) = parents::analyze(env, klass);

    let has_children = klass.has_children;

    let functions = functions::analyze(env, klass, class_tid);

    let has_constructors = functions.iter().find(|f| f.kind == library::FunctionKind::Constructor).is_some();
    let has_methods = functions.iter().find(|f| f.kind == library::FunctionKind::Method).is_some();

    Info {
        full_name: full_name,
        class_tid: class_tid,
        kind: kind,
        name: name,
        parents: parents,
        has_children: has_children,
        has_ignored_parents: has_ignored_parents,
        functions: functions,
        has_constructors: has_constructors,
        has_methods: has_methods,
    }
}
<commit_msg>Add check for widget's functions<commit_after>use env::Env;
use gobjects::*;
use library;
use nameutil::*;
use super::*;
use super::type_kind::*;

pub struct Info {
    pub full_name: String,
    pub class_tid: library::TypeId,
    pub kind: type_kind::TypeKind,
    pub name: String,
    pub parents: Vec<general::StatusedTypeId>,
    pub has_children: bool,
    pub has_ignored_parents: bool,
    pub functions: Vec<functions::Info>,
    pub has_constructors: bool,
    pub has_methods: bool,
}

impl Info {
    //TODO: add test in tests/ for panic
    pub fn type_<'a>(&self, library: &'a library::Library) -> &'a library::Class {
        let type_ = library.type_(self.class_tid).as_class()
            .unwrap_or_else(|| panic!("{} is not a class.", self.full_name));
        type_
    }

    ///TODO: return iterator
    pub fn constructors(&self) -> Vec<&functions::Info> {
        self.functions.iter()
            .filter(|f| f.kind == library::FunctionKind::Constructor)
            .collect()
    }

    pub fn methods(&self) -> Vec<&functions::Info> {
        self.functions.iter()
            .filter(|f| f.kind == library::FunctionKind::Method)
            .collect()
    }
}

pub fn new(env: &Env, obj: &GObject) -> Info {
    let full_name = obj.name.clone();

    let class_tid = env.library.find_type_unwrapped(0, &full_name, "Class");

    let type_ = env.library.type_(class_tid);
    let kind = type_.to_type_kind(&env.library);
    assert_eq!(kind, TypeKind::Widget);

    let name = split_namespace_name(&full_name).1.into();

    let klass = type_.to_class();
    let (parents, has_ignored_parents) = parents::analyze(env, klass);

    let has_children = klass.has_children;

    let functions = functions::analyze(env, klass, class_tid);

    let has_constructors = functions.iter().find(|f| f.kind == library::FunctionKind::Constructor).is_some();
    let has_methods = functions.iter().find(|f| f.kind == library::FunctionKind::Method).is_some();

    let has_functions = functions.iter().find(|f| f.kind == library::FunctionKind::Function).is_some();
    assert!(!has_functions, "Widget {} has functions, to do functions code generation", full_name);

    Info {
        full_name: full_name,
        class_tid: class_tid,
        kind: kind,
        name: name,
        parents: parents,
        has_children: has_children,
        has_ignored_parents: has_ignored_parents,
        functions: functions,
        has_constructors: has_constructors,
        has_methods: has_methods,
    }
}
<|endoftext|>


<commit_before>// Take a look at the license at the top of the repository in the LICENSE file.

use crate::sys::{
    ffi,
    utils::{self, CFReleaser},
};
use crate::{DiskExt, DiskType};

use core_foundation_sys::array::CFArrayCreate;
use core_foundation_sys::base::kCFAllocatorDefault;
use core_foundation_sys::dictionary::{CFDictionaryGetValueIfPresent, CFDictionaryRef};
use core_foundation_sys::number::{kCFBooleanTrue, CFBooleanRef, CFNumberGetValue};
use core_foundation_sys::string::{self as cfs, CFStringRef};

use libc::c_void;

use std::ffi::{CStr, OsStr, OsString};
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use std::ptr;

#[doc = include_str!("../../md_doc/disk.md")]
pub struct Disk {
    pub(crate) type_: DiskType,
    pub(crate) name: OsString,
    pub(crate) file_system: Vec<u8>,
    pub(crate) mount_point: PathBuf,
    volume_url: RetainedCFURL,
    pub(crate) total_space: u64,
    pub(crate) available_space: u64,
    pub(crate) is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        unsafe {
            let requested_properties = build_requested_properties(&[
                ffi::kCFURLVolumeAvailableCapacityKey,
                ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey,
            ]);
            match get_disk_properties(&self.volume_url, &requested_properties) {
                Some(disk_props) => {
                    self.available_space = get_available_volume_space(&disk_props);
                    true
                }
                None => false,
            }
        }
    }
}

pub(super) unsafe fn get_disks() -> Vec<Disk> {
    let raw_disks = {
        let count = libc::getfsstat(ptr::null_mut(), 0, libc::MNT_NOWAIT);
        if count < 1 {
            return Vec::new();
        }
        let bufsize = count * std::mem::size_of::<libc::statfs>() as libc::c_int;
        let mut disks = Vec::with_capacity(count as _);
        let count = libc::getfsstat(disks.as_mut_ptr(), bufsize, libc::MNT_NOWAIT);

        if count < 1 {
            return Vec::new();
        }

        disks.set_len(count as usize);

        disks
    };

    // Create a list of properties about the disk that we want to fetch.
    let requested_properties = build_requested_properties(&[
        ffi::kCFURLVolumeIsEjectableKey,
        ffi::kCFURLVolumeIsRemovableKey,
        ffi::kCFURLVolumeIsInternalKey,
        ffi::kCFURLVolumeTotalCapacityKey,
        ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey,
        ffi::kCFURLVolumeAvailableCapacityKey,
        ffi::kCFURLVolumeNameKey,
        ffi::kCFURLVolumeIsBrowsableKey,
        ffi::kCFURLVolumeIsLocalKey,
    ]);

    let mut disks = Vec::with_capacity(raw_disks.len());
    for c_disk in raw_disks {
        let volume_url = CFReleaser::new(
            core_foundation_sys::url::CFURLCreateFromFileSystemRepresentation(
                kCFAllocatorDefault,
                c_disk.f_mntonname.as_ptr() as *const _,
                c_disk.f_mntonname.len() as _,
                false as _,
            ),
        )
        .unwrap();

        let prop_dict = match get_disk_properties(&volume_url, &requested_properties) {
            Some(props) => props,
            None => continue,
        };

        let browsable = get_bool_value(
            prop_dict.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsBrowsableKey),
        )
        .unwrap_or_default();

        // Do not return invisible "disks". Most of the time, these are APFS snapshots, hidden
        // system volumes, etc. Browsable is defined to be visible in the system's UI like Finder,
        // disk utility, system information, etc.
        //
        // To avoid seemingly duplicating many disks and creating an inaccurate view of the system's resources,
        // these are skipped entirely.
        if !browsable {
            continue;
        }

        let local_only = get_bool_value(
            prop_dict.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsLocalKey),
        )
        .unwrap_or(true);

        // Skip any drive that is not locally attached to the system.
        //
        // This includes items like SMB mounts, and matches the other platform's behavior.
        if !local_only {
            continue;
        }

        let mount_point = PathBuf::from(OsStr::from_bytes(
            CStr::from_ptr(c_disk.f_mntonname.as_ptr()).to_bytes(),
        ));

        disks.extend(new_disk(mount_point, volume_url, c_disk, &prop_dict))
    }

    disks
}

type RetainedCFArray = CFReleaser<core_foundation_sys::array::__CFArray>;
type RetainedCFDictionary = CFReleaser<core_foundation_sys::dictionary::__CFDictionary>;
type RetainedCFURL = CFReleaser<core_foundation_sys::url::__CFURL>;

unsafe fn build_requested_properties(properties: &[CFStringRef]) -> RetainedCFArray {
    CFReleaser::new(CFArrayCreate(
        ptr::null_mut(),
        properties.as_ptr() as *const *const c_void,
        properties.len() as _,
        &core_foundation_sys::array::kCFTypeArrayCallBacks,
    ))
    .unwrap()
}

fn get_disk_properties(
    volume_url: &RetainedCFURL,
    requested_properties: &RetainedCFArray,
) -> Option<RetainedCFDictionary> {
    CFReleaser::new(unsafe {
        ffi::CFURLCopyResourcePropertiesForKeys(
            volume_url.inner(),
            requested_properties.inner(),
            ptr::null_mut(),
        )
    })
}

fn get_available_volume_space(disk_props: &RetainedCFDictionary) -> u64 {
    // We prefer `AvailableCapacityForImportantUsage` over `AvailableCapacity` because
    // it takes more of the system's properties into account, like the trash, system-managed caches,
    // etc. It generally also returns higher values too, because of the above, so it's a more accurate
    // representation of what the system _could_ still use.
    unsafe {
        get_int_value(
            disk_props.inner(),
            DictKey::Extern(ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey),
        )
        .filter(|bytes| *bytes != 0)
        .or_else(|| {
            get_int_value(
                disk_props.inner(),
                DictKey::Extern(ffi::kCFURLVolumeAvailableCapacityKey),
            )
        })
    }
    .unwrap_or_default() as u64
}

pub(super) enum DictKey {
    Extern(CFStringRef),
    #[cfg(target_os = "macos")]
    Defined(&'static str),
}

unsafe fn get_dict_value<T, F: FnOnce(*const c_void) -> Option<T>>(
    dict: CFDictionaryRef,
    key: DictKey,
    callback: F,
) -> Option<T> {
    #[cfg(target_os = "macos")]
    let _defined;
    let key = match key {
        DictKey::Extern(val) => val,
        #[cfg(target_os = "macos")]
        DictKey::Defined(val) => {
            _defined = CFReleaser::new(cfs::CFStringCreateWithBytesNoCopy(
                kCFAllocatorDefault,
                val.as_ptr(),
                val.len() as _,
                cfs::kCFStringEncodingUTF8,
                false as _,
                core_foundation_sys::base::kCFAllocatorNull,
            ))
            .unwrap();

            _defined.inner()
        }
    };

    let mut value = std::ptr::null();
    if CFDictionaryGetValueIfPresent(dict, key.cast(), &mut value) != 0 {
        callback(value)
    } else {
        None
    }
}

pub(super) unsafe fn get_str_value(dict: CFDictionaryRef, key: DictKey) -> Option<String> {
    get_dict_value(dict, key, |v| {
        let v = v as cfs::CFStringRef;

        let len_utf16 = cfs::CFStringGetLength(v) as usize;
        let len_bytes = len_utf16 * 2; // Two bytes per UTF-16 codepoint.

        let v_ptr = cfs::CFStringGetCStringPtr(v, cfs::kCFStringEncodingUTF8);
        if v_ptr.is_null() {
            // Fallback on CFStringGetString to read the underlying bytes from the CFString.
            let mut buf = vec![0; len_bytes];
            let success = cfs::CFStringGetCString(
                v,
                buf.as_mut_ptr(),
                len_bytes as _,
                cfs::kCFStringEncodingUTF8,
            );

            if success != 0 {
                utils::vec_to_rust(buf)
            } else {
                None
            }
        } else {
            utils::cstr_to_rust_with_size(v_ptr, Some(len_bytes))
        }
    })
}

unsafe fn get_bool_value(dict: CFDictionaryRef, key: DictKey) -> Option<bool> {
    get_dict_value(dict, key, |v| Some(v as CFBooleanRef == kCFBooleanTrue))
}

unsafe fn get_int_value(dict: CFDictionaryRef, key: DictKey) -> Option<i64> {
    get_dict_value(dict, key, |v| {
        let mut val: i64 = 0;
        if CFNumberGetValue(
            v.cast(),
            core_foundation_sys::number::kCFNumberSInt64Type,
            &mut val as *mut i64 as *mut c_void,
        ) {
            Some(val)
        } else {
            None
        }
    })
}

unsafe fn new_disk(
    mount_point: PathBuf,
    volume_url: RetainedCFURL,
    c_disk: libc::statfs,
    disk_props: &RetainedCFDictionary,
) -> Option<Disk> {
    // IOKit is not available on any but the most recent (16+) iOS and iPadOS versions.
    // Due to this, we can't query the medium type. All iOS devices use flash-based storage
    // so we just assume the disk type is an SSD until Rust has a way to conditionally link to
    // IOKit in more recent deployment versions.
    #[cfg(target_os = "macos")]
    let type_ = crate::sys::inner::disk::get_disk_type(&c_disk);
    #[cfg(not(target_os = "macos"))]
    let type_ = DiskType::SSD;

    // Note: Since we requested these properties from the system, we don't expect
    // these property retrievals to fail.

    let name = get_str_value(
        disk_props.inner(),
        DictKey::Extern(ffi::kCFURLVolumeNameKey),
    )
    .map(OsString::from)?;

    let is_removable = {
        let ejectable = get_bool_value(
            disk_props.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsEjectableKey),
        )
        .unwrap_or_default();

        let removable = get_bool_value(
            disk_props.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsRemovableKey),
        )
        .unwrap_or_default();

        let is_removable = ejectable || removable;

        if is_removable {
            is_removable
        } else {
            // If neither `ejectable` or `removable` return `true`, fallback to checking
            // if the disk is attached to the internal system.
            let internal = get_bool_value(
                disk_props.inner(),
                DictKey::Extern(ffi::kCFURLVolumeIsInternalKey),
            )
            .unwrap_or_default();

            !internal
        }
    };

    let total_space = get_int_value(
        disk_props.inner(),
        DictKey::Extern(ffi::kCFURLVolumeTotalCapacityKey),
    )? as u64;

    let available_space = get_available_volume_space(disk_props);

    let file_system = IntoIterator::into_iter(c_disk.f_fstypename)
        .filter_map(|b| if b != 0 { Some(b as u8) } else { None })
        .collect();

    Some(Disk {
        type_,
        name,
        file_system,
        mount_point,
        volume_url,
        total_space,
        available_space,
        is_removable,
    })
}
<commit_msg>Add implementation comments<commit_after>// Take a look at the license at the top of the repository in the LICENSE file.

use crate::sys::{
    ffi,
    utils::{self, CFReleaser},
};
use crate::{DiskExt, DiskType};

use core_foundation_sys::array::CFArrayCreate;
use core_foundation_sys::base::kCFAllocatorDefault;
use core_foundation_sys::dictionary::{CFDictionaryGetValueIfPresent, CFDictionaryRef};
use core_foundation_sys::number::{kCFBooleanTrue, CFBooleanRef, CFNumberGetValue};
use core_foundation_sys::string::{self as cfs, CFStringRef};

use libc::c_void;

use std::ffi::{CStr, OsStr, OsString};
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use std::ptr;

#[doc = include_str!("../../md_doc/disk.md")]
pub struct Disk {
    pub(crate) type_: DiskType,
    pub(crate) name: OsString,
    pub(crate) file_system: Vec<u8>,
    pub(crate) mount_point: PathBuf,
    volume_url: RetainedCFURL,
    pub(crate) total_space: u64,
    pub(crate) available_space: u64,
    pub(crate) is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        unsafe {
            let requested_properties = build_requested_properties(&[
                ffi::kCFURLVolumeAvailableCapacityKey,
                ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey,
            ]);
            match get_disk_properties(&self.volume_url, &requested_properties) {
                Some(disk_props) => {
                    self.available_space = get_available_volume_space(&disk_props);
                    true
                }
                None => false,
            }
        }
    }
}

pub(super) unsafe fn get_disks() -> Vec<Disk> {
    let raw_disks = {
        let count = libc::getfsstat(ptr::null_mut(), 0, libc::MNT_NOWAIT);
        if count < 1 {
            return Vec::new();
        }
        let bufsize = count * std::mem::size_of::<libc::statfs>() as libc::c_int;
        let mut disks = Vec::with_capacity(count as _);
        let count = libc::getfsstat(disks.as_mut_ptr(), bufsize, libc::MNT_NOWAIT);

        if count < 1 {
            return Vec::new();
        }

        disks.set_len(count as usize);

        disks
    };

    // Create a list of properties about the disk that we want to fetch.
    let requested_properties = build_requested_properties(&[
        ffi::kCFURLVolumeIsEjectableKey,
        ffi::kCFURLVolumeIsRemovableKey,
        ffi::kCFURLVolumeIsInternalKey,
        ffi::kCFURLVolumeTotalCapacityKey,
        ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey,
        ffi::kCFURLVolumeAvailableCapacityKey,
        ffi::kCFURLVolumeNameKey,
        ffi::kCFURLVolumeIsBrowsableKey,
        ffi::kCFURLVolumeIsLocalKey,
    ]);

    let mut disks = Vec::with_capacity(raw_disks.len());
    for c_disk in raw_disks {
        let volume_url = CFReleaser::new(
            core_foundation_sys::url::CFURLCreateFromFileSystemRepresentation(
                kCFAllocatorDefault,
                c_disk.f_mntonname.as_ptr() as *const _,
                c_disk.f_mntonname.len() as _,
                false as _,
            ),
        )
        .unwrap();

        let prop_dict = match get_disk_properties(&volume_url, &requested_properties) {
            Some(props) => props,
            None => continue,
        };

        // Future note: There is a difference between `kCFURLVolumeIsBrowsableKey` and the
        // `kCFURLEnumeratorSkipInvisibles` option of `CFURLEnumeratorOptions`. Specifically,
        // the first one considers the writable `Data`(`/System/Volumes/Data`) partition to be
        // browsable, while it is classified as "invisible" by CoreFoundation's volume emumerator.
        let browsable = get_bool_value(
            prop_dict.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsBrowsableKey),
        )
        .unwrap_or_default();

        // Do not return invisible "disks". Most of the time, these are APFS snapshots, hidden
        // system volumes, etc. Browsable is defined to be visible in the system's UI like Finder,
        // disk utility, system information, etc.
        //
        // To avoid seemingly duplicating many disks and creating an inaccurate view of the system's resources,
        // these are skipped entirely.
        if !browsable {
            continue;
        }

        let local_only = get_bool_value(
            prop_dict.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsLocalKey),
        )
        .unwrap_or(true);

        // Skip any drive that is not locally attached to the system.
        //
        // This includes items like SMB mounts, and matches the other platform's behavior.
        if !local_only {
            continue;
        }

        let mount_point = PathBuf::from(OsStr::from_bytes(
            CStr::from_ptr(c_disk.f_mntonname.as_ptr()).to_bytes(),
        ));

        disks.extend(new_disk(mount_point, volume_url, c_disk, &prop_dict))
    }

    disks
}

type RetainedCFArray = CFReleaser<core_foundation_sys::array::__CFArray>;
type RetainedCFDictionary = CFReleaser<core_foundation_sys::dictionary::__CFDictionary>;
type RetainedCFURL = CFReleaser<core_foundation_sys::url::__CFURL>;

unsafe fn build_requested_properties(properties: &[CFStringRef]) -> RetainedCFArray {
    CFReleaser::new(CFArrayCreate(
        ptr::null_mut(),
        properties.as_ptr() as *const *const c_void,
        properties.len() as _,
        &core_foundation_sys::array::kCFTypeArrayCallBacks,
    ))
    .unwrap()
}

fn get_disk_properties(
    volume_url: &RetainedCFURL,
    requested_properties: &RetainedCFArray,
) -> Option<RetainedCFDictionary> {
    CFReleaser::new(unsafe {
        ffi::CFURLCopyResourcePropertiesForKeys(
            volume_url.inner(),
            requested_properties.inner(),
            ptr::null_mut(),
        )
    })
}

fn get_available_volume_space(disk_props: &RetainedCFDictionary) -> u64 {
    // We prefer `AvailableCapacityForImportantUsage` over `AvailableCapacity` because
    // it takes more of the system's properties into account, like the trash, system-managed caches,
    // etc. It generally also returns higher values too, because of the above, so it's a more accurate
    // representation of what the system _could_ still use.
    unsafe {
        get_int_value(
            disk_props.inner(),
            DictKey::Extern(ffi::kCFURLVolumeAvailableCapacityForImportantUsageKey),
        )
        .filter(|bytes| *bytes != 0)
        .or_else(|| {
            get_int_value(
                disk_props.inner(),
                DictKey::Extern(ffi::kCFURLVolumeAvailableCapacityKey),
            )
        })
    }
    .unwrap_or_default() as u64
}

pub(super) enum DictKey {
    Extern(CFStringRef),
    #[cfg(target_os = "macos")]
    Defined(&'static str),
}

unsafe fn get_dict_value<T, F: FnOnce(*const c_void) -> Option<T>>(
    dict: CFDictionaryRef,
    key: DictKey,
    callback: F,
) -> Option<T> {
    #[cfg(target_os = "macos")]
    let _defined;
    let key = match key {
        DictKey::Extern(val) => val,
        #[cfg(target_os = "macos")]
        DictKey::Defined(val) => {
            _defined = CFReleaser::new(cfs::CFStringCreateWithBytesNoCopy(
                kCFAllocatorDefault,
                val.as_ptr(),
                val.len() as _,
                cfs::kCFStringEncodingUTF8,
                false as _,
                core_foundation_sys::base::kCFAllocatorNull,
            ))
            .unwrap();

            _defined.inner()
        }
    };

    let mut value = std::ptr::null();
    if CFDictionaryGetValueIfPresent(dict, key.cast(), &mut value) != 0 {
        callback(value)
    } else {
        None
    }
}

pub(super) unsafe fn get_str_value(dict: CFDictionaryRef, key: DictKey) -> Option<String> {
    get_dict_value(dict, key, |v| {
        let v = v as cfs::CFStringRef;

        let len_utf16 = cfs::CFStringGetLength(v) as usize;
        let len_bytes = len_utf16 * 2; // Two bytes per UTF-16 codepoint.

        let v_ptr = cfs::CFStringGetCStringPtr(v, cfs::kCFStringEncodingUTF8);
        if v_ptr.is_null() {
            // Fallback on CFStringGetString to read the underlying bytes from the CFString.
            let mut buf = vec![0; len_bytes];
            let success = cfs::CFStringGetCString(
                v,
                buf.as_mut_ptr(),
                len_bytes as _,
                cfs::kCFStringEncodingUTF8,
            );

            if success != 0 {
                utils::vec_to_rust(buf)
            } else {
                None
            }
        } else {
            utils::cstr_to_rust_with_size(v_ptr, Some(len_bytes))
        }
    })
}

unsafe fn get_bool_value(dict: CFDictionaryRef, key: DictKey) -> Option<bool> {
    get_dict_value(dict, key, |v| Some(v as CFBooleanRef == kCFBooleanTrue))
}

unsafe fn get_int_value(dict: CFDictionaryRef, key: DictKey) -> Option<i64> {
    get_dict_value(dict, key, |v| {
        let mut val: i64 = 0;
        if CFNumberGetValue(
            v.cast(),
            core_foundation_sys::number::kCFNumberSInt64Type,
            &mut val as *mut i64 as *mut c_void,
        ) {
            Some(val)
        } else {
            None
        }
    })
}

unsafe fn new_disk(
    mount_point: PathBuf,
    volume_url: RetainedCFURL,
    c_disk: libc::statfs,
    disk_props: &RetainedCFDictionary,
) -> Option<Disk> {
    // IOKit is not available on any but the most recent (16+) iOS and iPadOS versions.
    // Due to this, we can't query the medium type. All iOS devices use flash-based storage
    // so we just assume the disk type is an SSD until Rust has a way to conditionally link to
    // IOKit in more recent deployment versions.
    #[cfg(target_os = "macos")]
    let type_ = crate::sys::inner::disk::get_disk_type(&c_disk);
    #[cfg(not(target_os = "macos"))]
    let type_ = DiskType::SSD;

    // Note: Since we requested these properties from the system, we don't expect
    // these property retrievals to fail.

    let name = get_str_value(
        disk_props.inner(),
        DictKey::Extern(ffi::kCFURLVolumeNameKey),
    )
    .map(OsString::from)?;

    let is_removable = {
        let ejectable = get_bool_value(
            disk_props.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsEjectableKey),
        )
        .unwrap_or_default();

        let removable = get_bool_value(
            disk_props.inner(),
            DictKey::Extern(ffi::kCFURLVolumeIsRemovableKey),
        )
        .unwrap_or_default();

        let is_removable = ejectable || removable;

        if is_removable {
            is_removable
        } else {
            // If neither `ejectable` or `removable` return `true`, fallback to checking
            // if the disk is attached to the internal system.
            let internal = get_bool_value(
                disk_props.inner(),
                DictKey::Extern(ffi::kCFURLVolumeIsInternalKey),
            )
            .unwrap_or_default();

            !internal
        }
    };

    let total_space = get_int_value(
        disk_props.inner(),
        DictKey::Extern(ffi::kCFURLVolumeTotalCapacityKey),
    )? as u64;

    let available_space = get_available_volume_space(disk_props);

    let file_system = IntoIterator::into_iter(c_disk.f_fstypename)
        .filter_map(|b| if b != 0 { Some(b as u8) } else { None })
        .collect();

    Some(Disk {
        type_,
        name,
        file_system,
        mount_point,
        volume_url,
        total_space,
        available_space,
        is_removable,
    })
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate orbclient;
extern crate orbimage;

use std::env;

use orbclient::{Color, EventOption, Renderer, Window, WindowFlag};
use orbimage::Image;

#[derive(Clone, Copy)]
enum BackgroundMode {
    /// Do not resize the image, just center it
    Center,
    /// Resize the image to the display size
    Fill,
    /// Resize the image - keeping its aspect ratio, and fit it to the display with blank space
    Scale,
    /// Resize the image - keeping its aspect ratio, and crop to remove all blank space
    Zoom,
}

impl BackgroundMode {
    fn from_str(string: &str) -> BackgroundMode {
        match string {
            "fill" => BackgroundMode::Fill,
            "scale" => BackgroundMode::Scale,
            "zoom" => BackgroundMode::Zoom,
            _ => BackgroundMode::Center
        }
    }
}

fn find_scale(image: &Image, mode: BackgroundMode, display_width: u32, display_height: u32) -> (u32, u32) {
    match mode {
        BackgroundMode::Center => {
            (image.width(), image.height())
        },
        BackgroundMode::Fill => {
            (display_width, display_height)
        },
        BackgroundMode::Scale => {
            let d_w = display_width as f64;
            let d_h = display_height as f64;
            let i_w = image.width() as f64;
            let i_h = image.height() as f64;

            let scale = if d_w / d_h > i_w / i_h {
                d_h / i_h
            } else {
                d_w / i_w
            };

            ((i_w * scale) as u32, (i_h * scale) as u32)
        },
        BackgroundMode::Zoom => {
            let d_w = display_width as f64;
            let d_h = display_height as f64;
            let i_w = image.width() as f64;
            let i_h = image.height() as f64;

            let scale = if d_w / d_h < i_w / i_h {
                d_h / i_h
            } else {
                d_w / i_w
            };

            ((i_w * scale) as u32, (i_h * scale) as u32)
        }
    }
}

fn main() {
    let mut args = env::args().skip(1);

    let path = match args.next() {
        Some(arg) => arg,
        None => "/ui/background.png".to_string(),
    };

    let mode = BackgroundMode::from_str(&args.next().unwrap_or(String::new()));

    match Image::from_path(&path) {
        Ok(image) => {
            let (display_width, display_height) = orbclient::get_display_size().expect("background: failed to get display size");

            let mut window = Window::new_flags(
                0, 0, display_width, display_height, &format!("{} - Background", path),
                &[WindowFlag::Back, WindowFlag::Borderless, WindowFlag::Unclosable]
            ).unwrap();

            let mut scaled_image = image.clone();
            let mut resize = Some((display_width, display_height));
            loop {
                if let Some((w, h)) = resize.take() {
                    let (width, height) = find_scale(&image, mode, w, h);

                    if width == scaled_image.width() && height == scaled_image.height() {
                        // Do not resize scaled image
                    } else if width == image.width() && height == image.height() {
                        scaled_image = image.clone();
                    } else {
                        scaled_image = image.resize(width, height, orbimage::ResizeType::Lanczos3).unwrap();
                    }

                    window.set(Color::rgb(0, 0, 0));

                    let x = (window.width() as i32 - scaled_image.width() as i32)/2;
                    let y = (window.height() as i32 - scaled_image.height() as i32)/2;
                    scaled_image.draw(&mut window, x, y);

                    window.sync();
                }

                for event in window.events() {
                    match event.to_option() {
                        EventOption::Resize(resize_event) => {
                            resize = Some((resize_event.width, resize_event.height));
                        },
                        EventOption::Screen(screen_event) => {
                            window.set_size(screen_event.width, screen_event.height);
                            resize = Some((screen_event.width, screen_event.height));
                        },
                        EventOption::Quit(_) => return,
                        _ => ()
                    }
                }
            }
        },
        Err(err) => {
            println!("background: error loading {}: {}", path, err);
        }
    }
}
<commit_msg>Crop background images<commit_after>#![deny(warnings)]

extern crate orbclient;
extern crate orbimage;

use std::env;

use orbclient::{Color, EventOption, Renderer, Window, WindowFlag};
use orbimage::Image;

#[derive(Clone, Copy, Debug)]
enum BackgroundMode {
    /// Do not resize the image, just center it
    Center,
    /// Resize the image to the display size
    Fill,
    /// Resize the image - keeping its aspect ratio, and fit it to the display with blank space
    Scale,
    /// Resize the image - keeping its aspect ratio, and crop to remove all blank space
    Zoom,
}

impl BackgroundMode {
    fn from_str(string: &str) -> BackgroundMode {
        match string {
            "fill" => BackgroundMode::Fill,
            "scale" => BackgroundMode::Scale,
            "zoom" => BackgroundMode::Zoom,
            _ => BackgroundMode::Center
        }
    }
}

fn find_scale(image: &Image, mode: BackgroundMode, display_width: u32, display_height: u32) -> (u32, u32) {
    match mode {
        BackgroundMode::Center => {
            (image.width(), image.height())
        },
        BackgroundMode::Fill => {
            (display_width, display_height)
        },
        BackgroundMode::Scale => {
            let d_w = display_width as f64;
            let d_h = display_height as f64;
            let i_w = image.width() as f64;
            let i_h = image.height() as f64;

            let scale = if d_w / d_h > i_w / i_h {
                d_h / i_h
            } else {
                d_w / i_w
            };

            ((i_w * scale) as u32, (i_h * scale) as u32)
        },
        BackgroundMode::Zoom => {
            let d_w = display_width as f64;
            let d_h = display_height as f64;
            let i_w = image.width() as f64;
            let i_h = image.height() as f64;

            let scale = if d_w / d_h < i_w / i_h {
                d_h / i_h
            } else {
                d_w / i_w
            };

            ((i_w * scale) as u32, (i_h * scale) as u32)
        }
    }
}

fn main() {
    let mut args = env::args().skip(1);

    let path = match args.next() {
        Some(arg) => arg,
        None => "/ui/background.png".to_string(),
    };

    let mode = BackgroundMode::from_str(&args.next().unwrap_or(String::new()));

    match Image::from_path(&path) {
        Ok(image) => {
            let (display_width, display_height) = orbclient::get_display_size().expect("background: failed to get display size");

            let mut window = Window::new_flags(
                0, 0, display_width, display_height, &format!("{} - Background", path),
                &[WindowFlag::Back, WindowFlag::Borderless, WindowFlag::Unclosable]
            ).unwrap();

            let mut scaled_image = image.clone();
            let mut resize = Some((display_width, display_height));
            loop {
                if let Some((w, h)) = resize.take() {
                    let (width, height) = find_scale(&image, mode, w, h);

                    if width == scaled_image.width() && height == scaled_image.height() {
                        // Do not resize scaled image
                    } else if width == image.width() && height == image.height() {
                        scaled_image = image.clone();
                    } else {
                        scaled_image = image.resize(width, height, orbimage::ResizeType::Lanczos3).unwrap();
                    }

                    let (crop_x, crop_w) = if width > w  {
                        ((width - w)/2, w)
                    } else {
                        (0, width)
                    };

                    let (crop_y, crop_h) = if height > h {
                        ((height - h)/2, h)
                    } else {
                        (0, height)
                    };

                    window.set(Color::rgb(0, 0, 0));

                    let x = (w as i32 - crop_w as i32)/2;
                    let y = (h as i32 - crop_h as i32)/2;
                    scaled_image.roi(
                        crop_x, crop_y,
                        crop_w, crop_h,
                    ).draw(
                        &mut window,
                        x, y
                    );

                    window.sync();
                }

                for event in window.events() {
                    match event.to_option() {
                        EventOption::Resize(resize_event) => {
                            resize = Some((resize_event.width, resize_event.height));
                        },
                        EventOption::Screen(screen_event) => {
                            window.set_size(screen_event.width, screen_event.height);
                            resize = Some((screen_event.width, screen_event.height));
                        },
                        EventOption::Quit(_) => return,
                        _ => ()
                    }
                }
            }
        },
        Err(err) => {
            println!("background: error loading {}: {}", path, err);
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    fs,
    path::{Path, PathBuf},
    sync::{mpsc, Arc},
    thread,
    time::Duration,
};
use crossbeam_channel::{Receiver, Sender, SendError};
use relative_path::RelativePathBuf;
use thread_worker::WorkerHandle;
use walkdir::WalkDir;
use parking_lot::Mutex;
use notify::{DebouncedEvent, RecommendedWatcher, RecursiveMode, Watcher as _Watcher};

use crate::{RootConfig, Roots, VfsRoot};

pub(crate) enum Task {
    AddRoot {
        root: VfsRoot,
        config: Arc<RootConfig>,
    },
}

#[derive(Debug)]
pub enum TaskResult {
    BulkLoadRoot {
        root: VfsRoot,
        files: Vec<(RelativePathBuf, String)>,
    },
    AddSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    ChangeSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    RemoveSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
    },
}

#[derive(Debug)]
enum ChangeKind {
    Create,
    Write,
    Remove,
}

const WATCHER_DELAY: Duration = Duration::from_millis(250);

pub(crate) struct Worker {
    worker: thread_worker::Worker<Task, TaskResult>,
    worker_handle: WorkerHandle,
}

impl Worker {
    pub(crate) fn start(roots: Arc<Roots>) -> Worker {
        let (worker, worker_handle) =
            thread_worker::spawn("vfs", 128, move |input_receiver, output_sender| {
                let (notify_sender, notify_receiver) = mpsc::channel();
                let watcher = notify::watcher(notify_sender, WATCHER_DELAY)
                    .map_err(|e| log::error!("failed to spawn notify {}", e))
                    .ok();
                let ctx = WatcherCtx {
                    roots,
                    watcher: Arc::new(Mutex::new(watcher)),
                    sender: output_sender,
                };
                let thread = thread::spawn({
                    let ctx = ctx.clone();
                    move || {
                        let _ = notify_receiver
                            .into_iter()
                            // forward relevant events only
                            .try_for_each(|change| ctx.handle_debounced_event(change));
                    }
                });
                let res1 = input_receiver.into_iter().try_for_each(|t| match t {
                    Task::AddRoot { root, config } => watch_root(&ctx, root, Arc::clone(&config)),
                });
                drop(ctx.watcher.lock().take());
                drop(ctx);
                let res2 = thread.join();
                match &res2 {
                    Ok(()) => log::info!("... Watcher terminated with ok"),
                    Err(_) => log::error!("... Watcher terminated with err"),
                }
                res1.unwrap();
                res2.unwrap();
            });
        Worker {
            worker,
            worker_handle,
        }
    }

    pub(crate) fn sender(&self) -> &Sender<Task> {
        &self.worker.inp
    }

    pub(crate) fn receiver(&self) -> &Receiver<TaskResult> {
        &self.worker.out
    }

    pub(crate) fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }
}

fn watch_root(
    woker: &WatcherCtx,
    root: VfsRoot,
    config: Arc<RootConfig>,
) -> Result<(), SendError<TaskResult>> {
    let mut guard = woker.watcher.lock();
    log::debug!("loading {} ...", config.root.as_path().display());
    let files = watch_recursive(guard.as_mut(), config.root.as_path(), &*config)
        .into_iter()
        .filter_map(|path| {
            let abs_path = path.to_path(&config.root);
            let text = fs::read_to_string(abs_path)
                .map_err(|e| log::warn!("watcher error: {}", e))
                .ok()?;
            Some((path, text))
        })
        .collect();
    woker
        .sender
        .send(TaskResult::BulkLoadRoot { root, files })?;
    log::debug!("... loaded {}", config.root.as_path().display());
    Ok(())
}

#[derive(Clone)]
struct WatcherCtx {
    roots: Arc<Roots>,
    watcher: Arc<Mutex<Option<RecommendedWatcher>>>,
    sender: Sender<TaskResult>,
}

impl WatcherCtx {
    fn handle_debounced_event(&self, ev: DebouncedEvent) -> Result<(), SendError<TaskResult>> {
        match ev {
            DebouncedEvent::NoticeWrite(_)
            | DebouncedEvent::NoticeRemove(_)
            | DebouncedEvent::Chmod(_) => {
                // ignore
            }
            DebouncedEvent::Rescan => {
                // TODO rescan all roots
            }
            DebouncedEvent::Create(path) => {
                self.handle_change(path, ChangeKind::Create)?;
            }
            DebouncedEvent::Write(path) => {
                self.handle_change(path, ChangeKind::Write)?;
            }
            DebouncedEvent::Remove(path) => {
                self.handle_change(path, ChangeKind::Remove)?;
            }
            DebouncedEvent::Rename(src, dst) => {
                self.handle_change(src, ChangeKind::Remove)?;
                self.handle_change(dst, ChangeKind::Create)?;
            }
            DebouncedEvent::Error(err, path) => {
                // TODO should we reload the file contents?
                log::warn!("watcher error \"{}\", {:?}", err, path);
            }
        }
        Ok(())
    }

    fn handle_change(&self, path: PathBuf, kind: ChangeKind) -> Result<(), SendError<TaskResult>> {
        let (root, rel_path) = match self.roots.find(&path) {
            None => return Ok(()),
            Some(it) => it,
        };
        let config = &self.roots[root];
        match kind {
            ChangeKind::Create => {
                let mut paths = Vec::new();
                if path.is_dir() {
                    let mut guard = self.watcher.lock();
                    paths.extend(watch_recursive(guard.as_mut(), &path, &config));
                } else {
                    paths.push(rel_path);
                }
                paths
                    .into_iter()
                    .filter_map(|rel_path| {
                        let abs_path = rel_path.to_path(&config.root);
                        let text = fs::read_to_string(&abs_path)
                            .map_err(|e| log::warn!("watcher failed {}", e))
                            .ok()?;
                        Some((rel_path, text))
                    })
                    .try_for_each(|(path, text)| {
                        self.sender
                            .send(TaskResult::AddSingleFile { root, path, text })
                    })?
            }
            ChangeKind::Write => match fs::read_to_string(&path) {
                Err(e) => log::warn!("watcher failed {}", e),
                Ok(text) => self.sender.send(TaskResult::ChangeSingleFile {
                    root,
                    path: rel_path,
                    text,
                })?,
            },
            ChangeKind::Remove => self.sender.send(TaskResult::RemoveSingleFile {
                root,
                path: rel_path,
            })?,
        }
        Ok(())
    }
}

fn watch_recursive(
    mut watcher: Option<&mut RecommendedWatcher>,
    dir: &Path,
    config: &RootConfig,
) -> Vec<RelativePathBuf> {
    let mut files = Vec::new();
    for entry in WalkDir::new(dir)
        .into_iter()
        .filter_entry(|it| config.contains(it.path()).is_some())
        .filter_map(|it| it.map_err(|e| log::warn!("watcher error: {}", e)).ok())
    {
        if entry.file_type().is_dir() {
            if let Some(watcher) = &mut watcher {
                watch_one(watcher, entry.path());
            }
        } else {
            let path = config.contains(entry.path()).unwrap();
            files.push(path.to_owned());
        }
    }
    files
}

fn watch_one(watcher: &mut RecommendedWatcher, dir: &Path) {
    match watcher.watch(dir, RecursiveMode::NonRecursive) {
        Ok(()) => log::debug!("watching \"{}\"", dir.display()),
        Err(e) => log::warn!("could not watch \"{}\": {}", dir.display(), e),
    }
}
<commit_msg>consolidate error handling<commit_after>use std::{
    fs,
    path::{Path, PathBuf},
    sync::{mpsc, Arc},
    thread,
    time::Duration,
};
use crossbeam_channel::{Receiver, Sender, SendError};
use relative_path::RelativePathBuf;
use thread_worker::WorkerHandle;
use walkdir::WalkDir;
use parking_lot::Mutex;
use notify::{DebouncedEvent, RecommendedWatcher, RecursiveMode, Watcher as _Watcher};

use crate::{RootConfig, Roots, VfsRoot};

pub(crate) enum Task {
    AddRoot {
        root: VfsRoot,
        config: Arc<RootConfig>,
    },
}

#[derive(Debug)]
pub enum TaskResult {
    BulkLoadRoot {
        root: VfsRoot,
        files: Vec<(RelativePathBuf, String)>,
    },
    AddSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    ChangeSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    RemoveSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
    },
}

#[derive(Debug)]
enum ChangeKind {
    Create,
    Write,
    Remove,
}

const WATCHER_DELAY: Duration = Duration::from_millis(250);

pub(crate) struct Worker {
    worker: thread_worker::Worker<Task, TaskResult>,
    worker_handle: WorkerHandle,
}

impl Worker {
    pub(crate) fn start(roots: Arc<Roots>) -> Worker {
        let (worker, worker_handle) =
            thread_worker::spawn("vfs", 128, move |input_receiver, output_sender| {
                let (notify_sender, notify_receiver) = mpsc::channel();
                let watcher = notify::watcher(notify_sender, WATCHER_DELAY)
                    .map_err(|e| log::error!("failed to spawn notify {}", e))
                    .ok();
                let ctx = WatcherCtx {
                    roots,
                    watcher: Arc::new(Mutex::new(watcher)),
                    sender: output_sender,
                };
                let thread = thread::spawn({
                    let ctx = ctx.clone();
                    move || {
                        let _ = notify_receiver
                            .into_iter()
                            // forward relevant events only
                            .try_for_each(|change| ctx.handle_debounced_event(change));
                    }
                });
                let res1 = input_receiver.into_iter().try_for_each(|t| match t {
                    Task::AddRoot { root, config } => watch_root(&ctx, root, Arc::clone(&config)),
                });
                drop(ctx.watcher.lock().take());
                drop(ctx);
                let res2 = thread.join();
                match &res2 {
                    Ok(()) => log::info!("... Watcher terminated with ok"),
                    Err(_) => log::error!("... Watcher terminated with err"),
                }
                res1.unwrap();
                res2.unwrap();
            });
        Worker {
            worker,
            worker_handle,
        }
    }

    pub(crate) fn sender(&self) -> &Sender<Task> {
        &self.worker.inp
    }

    pub(crate) fn receiver(&self) -> &Receiver<TaskResult> {
        &self.worker.out
    }

    pub(crate) fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }
}

fn watch_root(
    woker: &WatcherCtx,
    root: VfsRoot,
    config: Arc<RootConfig>,
) -> Result<(), SendError<TaskResult>> {
    let mut guard = woker.watcher.lock();
    log::debug!("loading {} ...", config.root.as_path().display());
    let files = watch_recursive(guard.as_mut(), config.root.as_path(), &*config)
        .into_iter()
        .filter_map(|path| {
            let abs_path = path.to_path(&config.root);
            let text = read_to_string(&abs_path)?;
            Some((path, text))
        })
        .collect();
    woker
        .sender
        .send(TaskResult::BulkLoadRoot { root, files })?;
    log::debug!("... loaded {}", config.root.as_path().display());
    Ok(())
}

#[derive(Clone)]
struct WatcherCtx {
    roots: Arc<Roots>,
    watcher: Arc<Mutex<Option<RecommendedWatcher>>>,
    sender: Sender<TaskResult>,
}

impl WatcherCtx {
    fn handle_debounced_event(&self, ev: DebouncedEvent) -> Result<(), SendError<TaskResult>> {
        match ev {
            DebouncedEvent::NoticeWrite(_)
            | DebouncedEvent::NoticeRemove(_)
            | DebouncedEvent::Chmod(_) => {
                // ignore
            }
            DebouncedEvent::Rescan => {
                // TODO rescan all roots
            }
            DebouncedEvent::Create(path) => {
                self.handle_change(path, ChangeKind::Create)?;
            }
            DebouncedEvent::Write(path) => {
                self.handle_change(path, ChangeKind::Write)?;
            }
            DebouncedEvent::Remove(path) => {
                self.handle_change(path, ChangeKind::Remove)?;
            }
            DebouncedEvent::Rename(src, dst) => {
                self.handle_change(src, ChangeKind::Remove)?;
                self.handle_change(dst, ChangeKind::Create)?;
            }
            DebouncedEvent::Error(err, path) => {
                // TODO should we reload the file contents?
                log::warn!("watcher error \"{}\", {:?}", err, path);
            }
        }
        Ok(())
    }

    fn handle_change(&self, path: PathBuf, kind: ChangeKind) -> Result<(), SendError<TaskResult>> {
        let (root, rel_path) = match self.roots.find(&path) {
            None => return Ok(()),
            Some(it) => it,
        };
        let config = &self.roots[root];
        match kind {
            ChangeKind::Create => {
                let mut paths = Vec::new();
                if path.is_dir() {
                    let mut guard = self.watcher.lock();
                    paths.extend(watch_recursive(guard.as_mut(), &path, &config));
                } else {
                    paths.push(rel_path);
                }
                paths
                    .into_iter()
                    .filter_map(|rel_path| {
                        let abs_path = rel_path.to_path(&config.root);
                        let text = read_to_string(&abs_path)?;
                        Some((rel_path, text))
                    })
                    .try_for_each(|(path, text)| {
                        self.sender
                            .send(TaskResult::AddSingleFile { root, path, text })
                    })?
            }
            ChangeKind::Write => {
                if let Some(text) = read_to_string(&path) {
                    self.sender.send(TaskResult::ChangeSingleFile {
                        root,
                        path: rel_path,
                        text,
                    })?;
                }
            }
            ChangeKind::Remove => self.sender.send(TaskResult::RemoveSingleFile {
                root,
                path: rel_path,
            })?,
        }
        Ok(())
    }
}

fn watch_recursive(
    mut watcher: Option<&mut RecommendedWatcher>,
    dir: &Path,
    config: &RootConfig,
) -> Vec<RelativePathBuf> {
    let mut files = Vec::new();
    for entry in WalkDir::new(dir)
        .into_iter()
        .filter_entry(|it| config.contains(it.path()).is_some())
        .filter_map(|it| it.map_err(|e| log::warn!("watcher error: {}", e)).ok())
    {
        if entry.file_type().is_dir() {
            if let Some(watcher) = &mut watcher {
                watch_one(watcher, entry.path());
            }
        } else {
            let path = config.contains(entry.path()).unwrap();
            files.push(path.to_owned());
        }
    }
    files
}

fn watch_one(watcher: &mut RecommendedWatcher, dir: &Path) {
    match watcher.watch(dir, RecursiveMode::NonRecursive) {
        Ok(()) => log::debug!("watching \"{}\"", dir.display()),
        Err(e) => log::warn!("could not watch \"{}\": {}", dir.display(), e),
    }
}

fn read_to_string(path: &Path) -> Option<String> {
    fs::read_to_string(&path)
        .map_err(|e| log::warn!("failed to read file {}", e))
        .ok()
}
<|endoftext|>


<commit_before>use super::*;

use std::{collections::HashSet, ops::Bound};

#[derive(Debug, Clone)]
pub(crate) struct View<'a> {
    pub(crate) pid: PageId,
    pub(crate) lo: &'a IVec,
    pub(crate) hi: &'a IVec,
    pub(crate) is_index: bool,
    pub(crate) next: Option<PageId>,
    pub(crate) ptr: TreePtr<'a>,
    frags: Vec<&'a Frag>,
    base_offset: usize,
    pub(crate) base_data: &'a Data,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
    min_children: usize,
    max_children: usize,
    size: u64,
}

impl<'a> View<'a> {
    pub(crate) fn new(
        pid: PageId,
        ptr: TreePtr<'a>,
        frags: Vec<&'a Frag>,
    ) -> View<'a> {
        let mut merging = false;
        let mut merging_child = None;
        let mut merge_confirmed = false;
        let mut min_children: isize = 0;
        let mut max_children: isize = 0;
        let mut size = 0;

        for (offset, frag) in frags.iter().enumerate() {
            size += frag.size_in_bytes();
            match frag {
                Frag::Base(node) => {
                    // NB if we re-add ParentSplit, we must
                    // handle the hi & next members differently
                    // here

                    min_children += node.data.len() as isize;
                    max_children += node.data.len() as isize;

                    return View {
                        hi: &node.hi,
                        lo: &node.lo,
                        is_index: node.data.index_ref().is_some(),
                        next: node.next,
                        base_offset: offset,
                        base_data: &node.data,
                        pid,
                        ptr,
                        frags,
                        merging,
                        merging_child,
                        min_children: std::cmp::max(0, min_children) as usize,
                        max_children: std::cmp::max(0, max_children) as usize,
                        size,
                    };
                }
                Frag::ParentMergeIntention(pid) => {
                    if merge_confirmed {
                        merge_confirmed = false;
                    } else {
                        assert!(merging_child.is_none());
                        merging_child = Some(*pid);
                    }
                }
                Frag::ParentMergeConfirm => {
                    assert!(!merge_confirmed);
                    merge_confirmed = true;
                }
                Frag::ChildMergeCap => {
                    assert!(!merging);
                    assert_eq!(offset, 0, "frags: {:?}", frags);
                    merging = true;
                }
                Frag::Del(..) => {
                    min_children -= 1;
                    max_children -= 1;
                }
                Frag::Set(..) | Frag::Merge(..) => max_children += 1,
            }
        }

        panic!("view was never initialized with a base")
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Excluded(bound) if self.hi >= bound => true,
            Bound::Included(bound) if self.hi > bound => true,
            _ => self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(
        &self,
        bound: &Bound<IVec>,
        is_forward: bool,
    ) -> bool {
        match bound {
            Bound::Excluded(bound)
                if self.lo < bound || (is_forward && bound == self.lo) =>
            {
                true
            }
            Bound::Included(bound) if self.lo <= bound => true,
            Bound::Unbounded if !is_forward => self.hi.is_empty(),
            _ => self.lo.is_empty(),
        }
    }

    fn keys(&self) -> Vec<&IVec> {
        let mut keys: HashSet<&IVec> = self
            .base_data
            .leaf_ref()
            .unwrap()
            .iter()
            .map(|(k, _v)| k)
            .collect();

        for offset in (0..self.base_offset).rev() {
            match self.frags[offset] {
                Frag::Set(k, _) => {
                    keys.insert(k);
                }
                Frag::Del(k) => {
                    keys.remove(k);
                }
                Frag::Merge(k, _) => {
                    keys.insert(k);
                }
                Frag::Base(_) => {
                    panic!(
                        "somehow hit 2 base nodes while \
                         searching for a successor"
                    );
                }
                Frag::ChildMergeCap => {
                    assert_eq!(offset, 0);
                }
                Frag::ParentMergeIntention(_) | Frag::ParentMergeConfirm => {
                    panic!(
                        "somehow hit parent merge \
                         frags while searching for a \
                         successor"
                    )
                }
            }
        }

        let mut keys: Vec<_> = keys.into_iter().collect();
        keys.sort_unstable_by(|a, b| prefix_cmp(&a, &b));

        keys
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(self.lo, self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
        };

        let keys = self.keys();

        let successor_keys = keys.iter().filter(|k| {
            let ord = prefix_cmp(k, &predecessor_key);

            ord == std::cmp::Ordering::Greater
                || ord == std::cmp::Ordering::Equal
        });

        for encoded_key in successor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    prefix_encode(self.lo, &[255; 1024 * 1024])
                } else {
                    prefix_encode(self.lo, self.hi)
                }
            }
            Bound::Included(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
        };

        let keys = self.keys();

        let predecessor_keys = keys.iter().filter(|k| {
            let ord = prefix_cmp(k, &successor_key);

            ord == std::cmp::Ordering::Less || ord == std::cmp::Ordering::Equal
        });

        for encoded_key in predecessor_keys.rev() {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn leaf_value_for_key(
        &self,
        key: &[u8],
        config: &Config,
    ) -> Option<IVec> {
        assert!(!self.is_index);

        let mut merge_base = None;
        let mut merges = vec![];

        for frag in self.frags[..=self.base_offset].iter() {
            match frag {
                Frag::Set(k, val) if self.key_eq(k, key) => {
                    if merges.is_empty() {
                        return Some(val.clone());
                    } else {
                        merge_base = Some(val);
                        break;
                    }
                }
                Frag::Del(k) if self.key_eq(k, key) => {
                    // we should ignore "earlier" Frag's for
                    // this key, but still need to handle
                    // merges we encountered "after" this
                    // deletion
                    break;
                }
                Frag::Merge(k, val) if self.key_eq(k, key) => merges.push(val),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.leaf_ref().expect("last_node should be a leaf");
                    let search = items
                        .binary_search_by(|&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key, &self.lo)
                        })
                        .ok();

                    let val = search.map(|idx| &items[idx].1);
                    if merges.is_empty() {
                        return val.cloned();
                    } else {
                        merge_base = val;
                    }
                }
                _ => {}
            }
        }

        if merges.is_empty() {
            None
        } else {
            let merge_fn_ptr = config
                .merge_operator
                .expect("must have a merge operator set");

            unsafe {
                let merge_fn: MergeOperator = std::mem::transmute(merge_fn_ptr);

                let mut ret = merge_fn(
                    key,
                    merge_base.map(|iv| &**iv),
                    &merges.pop().unwrap(),
                );
                                       ;
                for merge in merges.into_iter().rev() {
                    if let Some(v) = ret {
                        ret = merge_fn(key, Some(&*v), merge);
                    } else {
                        ret = merge_fn(key, None, merge);
                    }
                }

                ret.map(IVec::from)
            }
        }
    }

    #[inline]
    fn key_eq(&self, encoded: &[u8], not_encoded: &[u8]) -> bool {
        prefix_cmp_encoded(encoded, not_encoded, self.lo)
            == std::cmp::Ordering::Equal
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> (usize, PageId) {
        assert!(self.is_index);
        let removed = self.removed_children();

        let items = self
            .base_data
            .index_ref()
            .expect("last_node should be a leaf");

        let search = binary_search_lub(items, |&(ref k, ref _v)| {
            prefix_cmp_encoded(k, key, &self.lo)
        });

        // This might be none if ord is Less and we're
        // searching for the empty key
        let mut index = search.expect("failed to traverse index");

        while removed.contains(&items[index].1) {
            index = index
                .checked_sub(1)
                .expect("leftmost child should never have been merged");
        }

        return (index, items[index].1);
    }

    pub(crate) fn removed_children(&self) -> Vec<PageId> {
        self.frags[..=self.base_offset]
            .iter()
            .filter_map(|frag| match frag {
                Frag::ParentMergeIntention(child_pid) => Some(*child_pid),
                Frag::Base(node) => node.merging_child,
                _ => None,
            })
            .collect()
    }

    pub(crate) fn should_split(&self, max_sz: u64) -> bool {
        let size_checks = self.min_children > 4 && self.size > max_sz;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn should_merge(&self, min_sz: u64) -> bool {
        let size_checks = self.max_children < 2 && self.size < min_sz;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn compact(&self, config: &Config) -> Node {
        let mut lhs = self.frags[self.base_offset].unwrap_base().clone();
        for offset in (0..self.base_offset).rev() {
            let frag = self.frags[offset];
            lhs.apply(frag, config.merge_operator);
        }

        assert_eq!(lhs.merging, self.merging);
        assert_eq!(lhs.merging_child, self.merging_child);
        assert_eq!(lhs.hi, self.hi);
        assert_eq!(lhs.lo, self.lo);

        lhs
    }

    pub(crate) fn split(&self, config: &Config) -> (Node, Node) {
        let mut lhs = self.compact(config);
        let rhs = lhs.split();

        lhs.data.drop_gte(&rhs.lo, &lhs.lo);
        lhs.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        lhs.next = None;

        (lhs, rhs)
    }
}
<commit_msg>Properly track merging in the tree view<commit_after>use super::*;

use std::{collections::HashSet, ops::Bound};

#[derive(Debug, Clone)]
pub(crate) struct View<'a> {
    pub(crate) pid: PageId,
    pub(crate) lo: &'a IVec,
    pub(crate) hi: &'a IVec,
    pub(crate) is_index: bool,
    pub(crate) next: Option<PageId>,
    pub(crate) ptr: TreePtr<'a>,
    frags: Vec<&'a Frag>,
    base_offset: usize,
    pub(crate) base_data: &'a Data,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
    min_children: usize,
    max_children: usize,
    size: u64,
}

impl<'a> View<'a> {
    pub(crate) fn new(
        pid: PageId,
        ptr: TreePtr<'a>,
        frags: Vec<&'a Frag>,
    ) -> View<'a> {
        let mut merging = false;
        let mut merging_child = None;
        let mut last_merge_confirmed = false;
        let mut min_children: isize = 0;
        let mut max_children: isize = 0;
        let mut size = 0;

        for (offset, frag) in frags.iter().enumerate() {
            size += frag.size_in_bytes();
            match frag {
                Frag::Base(node) => {
                    // NB if we re-add ParentSplit, we must
                    // handle the hi & next members differently
                    // here

                    min_children += node.data.len() as isize;
                    max_children += node.data.len() as isize;

                    if !last_merge_confirmed {
                        merging_child = merging_child.or(node.merging_child);
                    }

                    merging |= node.merging;

                    return View {
                        hi: &node.hi,
                        lo: &node.lo,
                        is_index: node.data.index_ref().is_some(),
                        next: node.next,
                        base_offset: offset,
                        base_data: &node.data,
                        pid,
                        ptr,
                        frags,
                        merging,
                        merging_child,
                        min_children: std::cmp::max(0, min_children) as usize,
                        max_children: std::cmp::max(0, max_children) as usize,
                        size,
                    };
                }
                Frag::ParentMergeIntention(pid) => {
                    if !last_merge_confirmed {
                        merging_child = merging_child.or(Some(*pid));
                    }
                }
                Frag::ParentMergeConfirm => {
                    if merging_child.is_none() {
                        last_merge_confirmed = true;
                    }
                }
                Frag::ChildMergeCap => {
                    assert!(!merging);
                    assert_eq!(offset, 0, "frags: {:?}", frags);
                    merging = true;
                }
                Frag::Del(..) => {
                    min_children -= 1;
                    max_children -= 1;
                }
                Frag::Set(..) | Frag::Merge(..) => max_children += 1,
            }
        }

        panic!("view was never initialized with a base")
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Excluded(bound) if self.hi >= bound => true,
            Bound::Included(bound) if self.hi > bound => true,
            _ => self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(
        &self,
        bound: &Bound<IVec>,
        is_forward: bool,
    ) -> bool {
        match bound {
            Bound::Excluded(bound)
                if self.lo < bound || (is_forward && bound == self.lo) =>
            {
                true
            }
            Bound::Included(bound) if self.lo <= bound => true,
            Bound::Unbounded if !is_forward => self.hi.is_empty(),
            _ => self.lo.is_empty(),
        }
    }

    fn keys(&self) -> Vec<&IVec> {
        let mut keys: HashSet<&IVec> = self
            .base_data
            .leaf_ref()
            .unwrap()
            .iter()
            .map(|(k, _v)| k)
            .collect();

        for offset in (0..self.base_offset).rev() {
            match self.frags[offset] {
                Frag::Set(k, _) => {
                    keys.insert(k);
                }
                Frag::Del(k) => {
                    keys.remove(k);
                }
                Frag::Merge(k, _) => {
                    keys.insert(k);
                }
                Frag::Base(_) => {
                    panic!(
                        "somehow hit 2 base nodes while \
                         searching for a successor"
                    );
                }
                Frag::ChildMergeCap => {
                    assert_eq!(offset, 0);
                }
                Frag::ParentMergeIntention(_) | Frag::ParentMergeConfirm => {
                    panic!(
                        "somehow hit parent merge \
                         frags while searching for a \
                         successor"
                    )
                }
            }
        }

        let mut keys: Vec<_> = keys.into_iter().collect();
        keys.sort_unstable_by(|a, b| prefix_cmp(&a, &b));

        keys
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(self.lo, self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
        };

        let keys = self.keys();

        let successor_keys = keys.iter().filter(|k| {
            let ord = prefix_cmp(k, &predecessor_key);

            ord == std::cmp::Ordering::Greater
                || ord == std::cmp::Ordering::Equal
        });

        for encoded_key in successor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    prefix_encode(self.lo, &[255; 1024 * 1024])
                } else {
                    prefix_encode(self.lo, self.hi)
                }
            }
            Bound::Included(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
        };

        let keys = self.keys();

        let predecessor_keys = keys.iter().filter(|k| {
            let ord = prefix_cmp(k, &successor_key);

            ord == std::cmp::Ordering::Less || ord == std::cmp::Ordering::Equal
        });

        for encoded_key in predecessor_keys.rev() {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn leaf_value_for_key(
        &self,
        key: &[u8],
        config: &Config,
    ) -> Option<IVec> {
        assert!(!self.is_index);

        let mut merge_base = None;
        let mut merges = vec![];

        for frag in self.frags[..=self.base_offset].iter() {
            match frag {
                Frag::Set(k, val) if self.key_eq(k, key) => {
                    if merges.is_empty() {
                        return Some(val.clone());
                    } else {
                        merge_base = Some(val);
                        break;
                    }
                }
                Frag::Del(k) if self.key_eq(k, key) => {
                    // we should ignore "earlier" Frag's for
                    // this key, but still need to handle
                    // merges we encountered "after" this
                    // deletion
                    break;
                }
                Frag::Merge(k, val) if self.key_eq(k, key) => merges.push(val),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.leaf_ref().expect("last_node should be a leaf");
                    let search = items
                        .binary_search_by(|&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key, &self.lo)
                        })
                        .ok();

                    let val = search.map(|idx| &items[idx].1);
                    if merges.is_empty() {
                        return val.cloned();
                    } else {
                        merge_base = val;
                    }
                }
                _ => {}
            }
        }

        if merges.is_empty() {
            None
        } else {
            let merge_fn_ptr = config
                .merge_operator
                .expect("must have a merge operator set");

            unsafe {
                let merge_fn: MergeOperator = std::mem::transmute(merge_fn_ptr);

                let mut ret = merge_fn(
                    key,
                    merge_base.map(|iv| &**iv),
                    &merges.pop().unwrap(),
                );
                                       ;
                for merge in merges.into_iter().rev() {
                    if let Some(v) = ret {
                        ret = merge_fn(key, Some(&*v), merge);
                    } else {
                        ret = merge_fn(key, None, merge);
                    }
                }

                ret.map(IVec::from)
            }
        }
    }

    #[inline]
    fn key_eq(&self, encoded: &[u8], not_encoded: &[u8]) -> bool {
        prefix_cmp_encoded(encoded, not_encoded, self.lo)
            == std::cmp::Ordering::Equal
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> (usize, PageId) {
        assert!(self.is_index);
        let removed = self.removed_children();

        let items = self
            .base_data
            .index_ref()
            .expect("last_node should be a leaf");

        let search = binary_search_lub(items, |&(ref k, ref _v)| {
            prefix_cmp_encoded(k, key, &self.lo)
        });

        // This might be none if ord is Less and we're
        // searching for the empty key
        let mut index = search.expect("failed to traverse index");

        while removed.contains(&items[index].1) {
            index = index
                .checked_sub(1)
                .expect("leftmost child should never have been merged");
        }

        return (index, items[index].1);
    }

    pub(crate) fn removed_children(&self) -> Vec<PageId> {
        self.frags[..=self.base_offset]
            .iter()
            .filter_map(|frag| match frag {
                Frag::ParentMergeIntention(child_pid) => Some(*child_pid),
                Frag::Base(node) => node.merging_child,
                _ => None,
            })
            .collect()
    }

    pub(crate) fn should_split(&self, max_sz: u64) -> bool {
        let size_checks = self.min_children > 4 && self.size > max_sz;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn should_merge(&self, min_sz: u64) -> bool {
        let size_checks = self.max_children < 2 && self.size < min_sz;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn compact(&self, config: &Config) -> Node {
        let mut lhs = self.frags[self.base_offset].unwrap_base().clone();
        for offset in (0..self.base_offset).rev() {
            let frag = self.frags[offset];
            lhs.apply(frag, config.merge_operator);
        }

        assert_eq!(lhs.merging, self.merging);
        assert_eq!(lhs.merging_child, self.merging_child);
        assert_eq!(lhs.hi, self.hi);
        assert_eq!(lhs.lo, self.lo);

        lhs
    }

    pub(crate) fn split(&self, config: &Config) -> (Node, Node) {
        let mut lhs = self.compact(config);
        let rhs = lhs.split();

        lhs.data.drop_gte(&rhs.lo, &lhs.lo);
        lhs.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        lhs.next = None;

        (lhs, rhs)
    }
}
<|endoftext|>