<commit_before>// Copyright 2017, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

extern crate winapi;

use std::ops::Deref;

#[cfg(feature = "use_glib")]
use glib::translate::*;
use ffi;
use ffi::enums::{Format, SurfaceType};
use surface::{Surface, SurfaceExt};

#[derive(Debug)]
pub struct Win32Surface(Surface);

impl Win32Surface {
    pub fn from(surface: Surface) -> Result<Win32Surface, Surface> {
        if surface.get_type() == SurfaceType::Win32 {
            Ok(Win32Surface(surface))
        } else {
            Err(surface)
        }
    }

    #[doc(hidden)]
    pub unsafe fn from_raw_full(ptr: *mut ffi::cairo_surface_t) -> Win32Surface {
        Self::from(Surface::from_raw_full(ptr)).unwrap()
    }

    pub fn create(hdc: winapi::HDC) -> Win32Surface {
        unsafe { Self::from_raw_full(ffi::cairo_win32_surface_create(hdc)) }
    }

    pub fn create_with_dib(format: Format, width: i32, height: i32) -> Win32Surface {
        unsafe { Self::from_raw_full(ffi::cairo_win32_surface_create_with_dib(format, width, height)) }
    }

    pub fn create_with_ddb(hdc: winapi::HDC,
                           format: Format,
                           width: i32,
                           height: i32)
                           -> Win32Surface {
        unsafe {
            Self::from_raw_full(ffi::cairo_win32_surface_create_with_ddb(hdc, format, width, height))
        }
    }

    pub fn printing_surface_create(hdc: winapi::HDC) -> Win32Surface {
        unsafe { Self::from_raw_full(ffi::cairo_win32_printing_surface_create(hdc)) }
    }
}

#[cfg(feature = "use_glib")]
impl<'a> ToGlibPtr<'a, *mut ffi::cairo_surface_t> for Win32Surface {
    type Storage = &'a Surface;

    #[inline]
    fn to_glib_none(&'a self) -> Stash<'a, *mut ffi::cairo_surface_t, Self> {
        let stash = self.0.to_glib_none();
        Stash(stash.0, stash.1)
    }
}

#[cfg(feature = "use_glib")]
impl FromGlibPtrNone<*mut ffi::cairo_surface_t> for Win32Surface {
    #[inline]
    unsafe fn from_glib_none(ptr: *mut ffi::cairo_surface_t) -> Win32Surface {
        Self::from(from_glib_none(ptr)).unwrap()
    }
}

#[cfg(feature = "use_glib")]
impl FromGlibPtrFull<*mut ffi::cairo_surface_t> for Win32Surface {
    #[inline]
    unsafe fn from_glib_full(ptr: *mut ffi::cairo_surface_t) -> Win32Surface {
        Self::from(from_glib_full(ptr)).unwrap()
    }
}

impl AsRef<Surface> for Win32Surface {
    fn as_ref(&self) -> &Surface {
        &self.0
    }
}

impl Deref for Win32Surface {
    type Target = Surface;

    fn deref(&self) -> &Surface {
        &self.0
    }
}

impl Clone for Win32Surface {
    fn clone(&self) -> Win32Surface {
        Win32Surface(self.0.clone())
    }
}
<commit_msg>Win32Surface: return Result from the creation functions<commit_after>// Copyright 2017, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

extern crate winapi;

use std::ops::Deref;

#[cfg(feature = "use_glib")]
use glib::translate::*;
use ffi;
use ffi::enums::{Format, SurfaceType};
use surface::{Surface, SurfaceExt};
use Status;

#[derive(Debug)]
pub struct Win32Surface(Surface);

impl Win32Surface {
    pub fn from(surface: Surface) -> Result<Win32Surface, Surface> {
        if surface.get_type() == SurfaceType::Win32 {
            Ok(Win32Surface(surface))
        } else {
            Err(surface)
        }
    }

    #[doc(hidden)]
    pub unsafe fn from_raw_full(ptr: *mut ffi::cairo_surface_t) -> Result<Win32Surface, Status> {
        let surface = Self::from(Surface::from_raw_full(ptr)).unwrap();
        let status = surface.status();
        match status {
            Status::Success => Ok(surface),
            _ => Err(status)
        }
    }

    pub fn create(hdc: winapi::HDC) -> Result<Win32Surface, Status> {
        unsafe { Self::from_raw_full(ffi::cairo_win32_surface_create(hdc)) }
    }

    pub fn create_with_dib(format: Format, width: i32, height: i32) -> Result<Win32Surface, Status> {
        unsafe { Self::from_raw_full(ffi::cairo_win32_surface_create_with_dib(format, width, height)) }
    }

    pub fn create_with_ddb(hdc: winapi::HDC,
                           format: Format,
                           width: i32,
                           height: i32)
                           -> Result<Win32Surface, Status> {
        unsafe {
            Self::from_raw_full(ffi::cairo_win32_surface_create_with_ddb(hdc, format, width, height))
        }
    }

    pub fn printing_surface_create(hdc: winapi::HDC) -> Result<Win32Surface, Status> {
        unsafe { Self::from_raw_full(ffi::cairo_win32_printing_surface_create(hdc)) }
    }
}

#[cfg(feature = "use_glib")]
impl<'a> ToGlibPtr<'a, *mut ffi::cairo_surface_t> for Win32Surface {
    type Storage = &'a Surface;

    #[inline]
    fn to_glib_none(&'a self) -> Stash<'a, *mut ffi::cairo_surface_t, Self> {
        let stash = self.0.to_glib_none();
        Stash(stash.0, stash.1)
    }
}

#[cfg(feature = "use_glib")]
impl FromGlibPtrNone<*mut ffi::cairo_surface_t> for Win32Surface {
    #[inline]
    unsafe fn from_glib_none(ptr: *mut ffi::cairo_surface_t) -> Win32Surface {
        Self::from(from_glib_none(ptr)).unwrap()
    }
}

#[cfg(feature = "use_glib")]
impl FromGlibPtrFull<*mut ffi::cairo_surface_t> for Win32Surface {
    #[inline]
    unsafe fn from_glib_full(ptr: *mut ffi::cairo_surface_t) -> Win32Surface {
        Self::from(from_glib_full(ptr)).unwrap()
    }
}

impl AsRef<Surface> for Win32Surface {
    fn as_ref(&self) -> &Surface {
        &self.0
    }
}

impl Deref for Win32Surface {
    type Target = Surface;

    fn deref(&self) -> &Surface {
        &self.0
    }
}

impl Clone for Win32Surface {
    fn clone(&self) -> Win32Surface {
        Win32Surface(self.0.clone())
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

#[derive(PartialEq)]
enum ScopeKind {
    Block,
    Function,
}

struct ScopeFrame<'ast> {
    kind: ScopeKind,
    vars: HashMap<&'ast str, &'ast str>
}

pub struct Scope<'ast> {
    frame: usize,
    frames: Vec<ScopeFrame<'ast>>
}

impl<'ast> Scope<'ast> {
    pub fn new() -> Self {
        let mut frames = Vec::with_capacity(8);

        frames.push(ScopeFrame {
            kind: ScopeKind::Function,
            vars: HashMap::new()
        });

        Scope {
            frame: 0,
            frames
        }
    }

    fn push(&mut self, kind: ScopeKind) {
        self.frame += 1;

        if let Some(frame) = self.frames.get_mut(self.frame) {
            frame.kind = kind;
            return;
        }

        self.frames.push(ScopeFrame {
            kind,
            vars: HashMap::new()
        });
    }

    pub fn block_frame(&mut self) {
        self.push(ScopeKind::Block)
    }

    pub fn function_frame(&mut self) {
        self.push(ScopeKind::Function)
    }

    pub fn pop(&mut self) {
        self.frames[self.frame].vars.clear();
        self.frame -= 1;
    }

    pub fn has_in_block(&self, var: &str) -> bool {
        self.frames[self.frame].vars.contains_key(var)
    }

    pub fn has_in_function(&self, var: &str) -> bool {
        for frame in self.frames[..self.frame + 1].iter().rev() {
            if frame.vars.contains_key(var) {
                return true;
            }

            if frame.kind == ScopeKind::Function {
                return false;
            }
        }

        unreachable!("Last ScopeFrame must be of Function kind")
    }

    pub fn set_in_block(&mut self, var: &'ast str) {
        self.frames[self.frame].vars.insert(var, var);
    }

    pub fn set_in_function(&mut self, var: &'ast str) {
        let frame = self.frames[..self.frame + 1]
                        .iter_mut()
                        .rev()
                        .skip_while(|frame| frame.kind != ScopeKind::Function)
                        .next()
                        .expect("Must have a Function kind scope");

        frame.vars.insert(var, var);
    }
}

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

    #[test]
    fn scope() {
        let mut scope = Scope::new();

        scope.set_in_block("foo");
        scope.set_in_function("bar");

        scope.function_frame();

        scope.set_in_block("fooz");

        scope.block_frame();

        scope.set_in_block("baz");
        scope.set_in_function("qux");

        assert_eq!(scope.has_in_function("fooz"), true);
        assert_eq!(scope.has_in_function("baz"), true);
        assert_eq!(scope.has_in_function("qux"), true);

        assert_eq!(scope.has_in_block("baz"), true);
        assert_eq!(scope.has_in_block("qux"), false);
        assert_eq!(scope.has_in_block("fooz"), false);

        assert_eq!(scope.has_in_function("foo"), false);
        assert_eq!(scope.has_in_function("bar"), false);

        scope.pop();

        assert_eq!(scope.has_in_block("baz"), false);

        assert_eq!(scope.has_in_block("qux"), true);
        assert_eq!(scope.has_in_block("fooz"), true);
        assert_eq!(scope.has_in_function("qux"), true);
        assert_eq!(scope.has_in_function("fooz"), true);

        scope.pop();

        assert_eq!(scope.has_in_block("foo"), true);
        assert_eq!(scope.has_in_block("bar"), true);
        assert_eq!(scope.has_in_function("foo"), true);
        assert_eq!(scope.has_in_function("bar"), true);
    }
}
<commit_msg>Ability to generate new unique variable names<commit_after>use std::collections::hash_map::{HashMap, Entry};
use arena::Arena;

#[derive(PartialEq)]
enum ScopeKind {
    Block,
    Function,
}

struct ScopeFrame<'ast> {
    kind: ScopeKind,
    vars: HashMap<&'ast str, &'ast str>,
}

pub struct Scope<'ast> {
    arena: &'ast Arena,
    frame: usize,
    frames: Vec<ScopeFrame<'ast>>,
}

impl<'ast> Scope<'ast> {
    pub fn new(arena: &'ast Arena) -> Self {
        let mut frames = Vec::with_capacity(8);

        frames.push(ScopeFrame {
            kind: ScopeKind::Function,
            vars: HashMap::new()
        });

        Scope {
            arena,
            frame: 0,
            frames
        }
    }

    fn push(&mut self, kind: ScopeKind) {
        self.frame += 1;

        if let Some(frame) = self.frames.get_mut(self.frame) {
            frame.kind = kind;
            return;
        }

        self.frames.push(ScopeFrame {
            kind,
            vars: HashMap::new()
        });
    }

    pub fn block_frame(&mut self) {
        self.push(ScopeKind::Block)
    }

    pub fn function_frame(&mut self) {
        self.push(ScopeKind::Function)
    }

    pub fn pop(&mut self) {
        self.frames[self.frame].vars.clear();
        self.frame -= 1;
    }

    pub fn has_in_block(&self, var: &str) -> bool {
        self.frames[self.frame].vars.contains_key(var)
    }

    pub fn has_in_function(&self, var: &str) -> bool {
        for frame in self.frames[..self.frame + 1].iter().rev() {
            if frame.vars.contains_key(var) {
                return true;
            }

            if frame.kind == ScopeKind::Function {
                return false;
            }
        }

        unreachable!("Last ScopeFrame must be of Function kind")
    }

    pub fn set_in_block(&mut self, var: &'ast str) {
        self.frames[self.frame].vars.insert(var, var);
    }

    pub fn set_in_function(&mut self, var: &'ast str) {
        let frame = self.frames[..self.frame + 1]
                        .iter_mut()
                        .rev()
                        .skip_while(|frame| frame.kind != ScopeKind::Function)
                        .next()
                        .expect("Must have a Function kind scope");

        frame.vars.insert(var, var);
    }

    pub fn set_unique_in_block(&mut self, var: &'ast str) -> &'ast str {
        let frame = &mut self.frames[self.frame];

        Scope::set_unique_in_frame(&self.arena, frame, var)
    }

    pub fn set_unique_in_function(&mut self, var: &'ast str) -> &'ast str {
        let frame = self.frames[..self.frame + 1]
                        .iter_mut()
                        .rev()
                        .skip_while(|frame| frame.kind != ScopeKind::Function)
                        .next()
                        .expect("Must have a Function kind scope");

        Scope::set_unique_in_frame(&self.arena, frame, var)
    }

    fn set_unique_in_frame(arena: &'ast Arena, frame: &mut ScopeFrame<'ast>, var: &'ast str) -> &'ast str {
        if let Entry::Vacant(vacant) = frame.vars.entry(var) {
            vacant.insert(var);
            return var;
        }

        let mut attempt = 1;

        loop {
            let altered = format!("{}${}", var, attempt);

            if !frame.vars.contains_key(altered.as_str()) {
                let var = arena.alloc_string(altered);
                frame.vars.insert(var, var);
                return var;
            }

            attempt += 1;
        }
    }
}

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

    #[test]
    fn scope() {
        let arena = Arena::new();
        let mut scope = Scope::new(&arena);

        scope.set_in_block("foo");
        scope.set_in_function("bar");

        scope.function_frame();

        scope.set_in_block("fooz");

        scope.block_frame();

        scope.set_in_block("baz");
        scope.set_in_function("qux");

        assert_eq!(scope.has_in_function("fooz"), true);
        assert_eq!(scope.has_in_function("baz"), true);
        assert_eq!(scope.has_in_function("qux"), true);

        assert_eq!(scope.has_in_block("baz"), true);
        assert_eq!(scope.has_in_block("qux"), false);
        assert_eq!(scope.has_in_block("fooz"), false);

        assert_eq!(scope.has_in_function("foo"), false);
        assert_eq!(scope.has_in_function("bar"), false);

        scope.pop();

        assert_eq!(scope.has_in_block("baz"), false);

        assert_eq!(scope.has_in_block("qux"), true);
        assert_eq!(scope.has_in_block("fooz"), true);
        assert_eq!(scope.has_in_function("qux"), true);
        assert_eq!(scope.has_in_function("fooz"), true);

        scope.pop();

        assert_eq!(scope.has_in_block("foo"), true);
        assert_eq!(scope.has_in_block("bar"), true);
        assert_eq!(scope.has_in_function("foo"), true);
        assert_eq!(scope.has_in_function("bar"), true);
    }

    #[test]
    fn set_unique_in_block() {
        let arena = Arena::new();
        let mut scope = Scope::new(&arena);

        assert_eq!(scope.set_unique_in_block("foo"), "foo");
        assert_eq!(scope.set_unique_in_block("foo"), "foo$1");
        assert_eq!(scope.set_unique_in_block("foo"), "foo$2");
        assert_eq!(scope.set_unique_in_block("foo"), "foo$3");

        assert_eq!(scope.set_unique_in_block("bar"), "bar");
        assert_eq!(scope.set_unique_in_block("bar"), "bar$1");
        assert_eq!(scope.set_unique_in_block("bar"), "bar$2");
        assert_eq!(scope.set_unique_in_block("bar"), "bar$3");

        assert_eq!(scope.has_in_block("foo"), true);
        assert_eq!(scope.has_in_block("foo$1"), true);
        assert_eq!(scope.has_in_block("foo$2"), true);
        assert_eq!(scope.has_in_block("foo$3"), true);
        assert_eq!(scope.has_in_block("foo$4"), false);
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::io::{Read, Seek, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
/// TODO consider splitting `Snapshot` into separate
/// snapshots for PC, SA, Materializer.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Snapshot<R> {
    /// the last lsn included in the `Snapshot`
    pub max_lsn: Lsn,
    /// the last lid included in the `Snapshot`
    pub last_lid: LogId,
    /// the highest allocated pid
    pub max_pid: PageId,
    /// the mapping from pages to (lsn, lid)
    pub pt: HashMap<PageId, PageState>,
    /// replaced pages per segment index
    pub replacements: HashMap<SegmentId, HashSet<(PageId, Lsn)>>,
    /// the free pids
    pub free: HashSet<PageId>,
    /// the `Materializer`-specific recovered state
    pub recovery: Option<R>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PageState {
    Present(Vec<(Lsn, DiskPtr)>),
    Allocated(Lsn, DiskPtr),
    Free(Lsn, DiskPtr),
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr)) {
        match self {
            &mut PageState::Present(ref mut items) => {
                items.push(item)
            }
            &mut PageState::Allocated(_, _) => {
                *self = PageState::Present(vec![item])
            }
            &mut PageState::Free(_, _) => {
                panic!("pushed items to a PageState::Free")
            }
        }
    }

    /// Iterate over the (lsn, lid) pairs that hold this page's state.
    pub fn iter(&self) -> Box<dyn Iterator<Item = (Lsn, DiskPtr)>> {
        match self {
            &PageState::Present(ref items) => {
                Box::new(items.clone().into_iter())
            }
            &PageState::Allocated(lsn, ptr)
            | &PageState::Free(lsn, ptr) => {
                Box::new(vec![(lsn, ptr)].into_iter())
            }
        }
    }

    fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl<R> Default for Snapshot<R> {
    fn default() -> Snapshot<R> {
        Snapshot {
            max_lsn: 0,
            last_lid: 0,
            max_pid: 0,
            pt: HashMap::new(),
            replacements: HashMap::new(),
            free: HashSet::new(),
            recovery: None,
        }
    }
}

impl<R> Snapshot<R> {
    fn apply<P>(
        &mut self,
        materializer: &dyn Materializer<PageFrag = P, Recovery = R>,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        bytes: &[u8],
        config: &Config,
    ) where
        P: 'static
            + Debug
            + Clone
            + Serialize
            + DeserializeOwned
            + Send
            + Sync,
        R: Debug + Clone + Serialize + DeserializeOwned + Send,
    {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!(
            "trying to deserialize buf for ptr {} lsn {}",
            disk_ptr,
            lsn
        );
        let deserialization = deserialize::<LoggedUpdate<P>>(&*bytes);

        if let Err(e) = deserialization {
            error!(
                "failed to deserialize buffer for item in log: lsn {} \
                    ptr {}: {:?}",
                lsn,
                disk_ptr,
                e
            );
            return;
        }

        let prepend = deserialization.unwrap();
        let pid = prepend.pid;

        if pid >= self.max_pid {
            self.max_pid = pid + 1;
        }

        let io_buf_size = config.io_buf_size;

        let replaced_at_idx =
            disk_ptr.lid() as SegmentId / io_buf_size;

        match prepend.update {
            Update::Append(partial_page) => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn
                    );

                    if lids.is_free() {
                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );
                        return;
                    }

                    if let Some(r) =
                        materializer.recover(&partial_page)
                    {
                        self.recovery = Some(r);
                    }

                    lids.push((lsn, disk_ptr));
                }
                self.free.remove(&pid);
            }
            Update::Compact(partial_page) => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                if let Some(r) = materializer.recover(&partial_page) {
                    self.recovery = Some(r);
                }

                self.replace_pid(
                    pid,
                    replaced_at_idx,
                    lsn,
                    io_buf_size,
                    config,
                );
                self.pt.insert(
                    pid,
                    PageState::Present(vec![(lsn, disk_ptr)]),
                );
                self.free.remove(&pid);
            }
            Update::Allocate => {
                trace!(
                    "allocate  of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                self.replace_pid(
                    pid,
                    replaced_at_idx,
                    lsn,
                    io_buf_size,
                    config,
                );
                self.pt
                    .insert(pid, PageState::Allocated(lsn, disk_ptr));
                self.free.remove(&pid);
            }
            Update::Free => {
                trace!(
                    "free of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                self.replace_pid(
                    pid,
                    replaced_at_idx,
                    lsn,
                    io_buf_size,
                    config,
                );
                self.pt.insert(pid, PageState::Free(lsn, disk_ptr));
                self.free.insert(pid);
            }
        }
    }

    fn replace_pid(
        &mut self,
        pid: PageId,
        replaced_at_idx: usize,
        replaced_at_lsn: Lsn,
        io_buf_size: usize,
        config: &Config,
    ) {
        match self.pt.remove(&pid) {
            Some(PageState::Present(coords)) => {
                for (_lsn, ptr) in &coords {
                    let idx = ptr.lid() as SegmentId / io_buf_size;
                    if replaced_at_idx == idx {
                        return;
                    }
                    let entry = self
                        .replacements
                        .entry(idx)
                        .or_insert(HashSet::new());
                    entry.insert((pid, replaced_at_lsn));
                }

                // re-run any blob removals in case
                // they were not completed.
                if coords.len() > 1 {
                    let blob_ptrs = coords
                        .iter()
                        .filter(|(_, ptr)| ptr.is_blob())
                        .map(|(_, ptr)| ptr.blob().1);

                    for blob_ptr in blob_ptrs {
                        trace!(
                            "removing blob while advancing \
                             snapshot: {}",
                            blob_ptr,
                        );

                        // we don't care if this actually works
                        // because it's possible that a previous
                        // snapshot has run over this log and
                        // removed the blob already.
                        let _ = remove_blob(blob_ptr, config);
                    }
                }
            }
            Some(PageState::Allocated(_lsn, ptr))
            | Some(PageState::Free(_lsn, ptr)) => {
                let idx = ptr.lid() as SegmentId / io_buf_size;
                if replaced_at_idx == idx {
                    return;
                }
                let entry = self
                    .replacements
                    .entry(idx)
                    .or_insert(HashSet::new());
                entry.insert((pid, replaced_at_lsn));
            }
            None => {
                // we just encountered a replace without it
                // existing in the pt. this means we've
                // relocated it to a later segment than
                // the one we're currently at during recovery.
                // so we can skip it here.
            }
        }
    }
}

pub(super) fn advance_snapshot<PM, P, R>(
    iter: LogIter,
    mut snapshot: Snapshot<R>,
    config: &Config,
) -> Result<Snapshot<R>, ()>
where
    PM: Materializer<Recovery = R, PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let start = clock();

    trace!("building on top of old snapshot: {:?}", snapshot);

    let materializer = PM::new(config.clone(), &snapshot.recovery);

    let io_buf_size = config.io_buf_size;

    for (lsn, ptr, bytes) in iter {
        let segment_idx = ptr.lid() as SegmentId / io_buf_size;

        trace!(
            "in advance_snapshot looking at item with lsn {} ptr {}",
            lsn,
            ptr
        );

        if lsn <= snapshot.max_lsn {
            // don't process already-processed Lsn's. max_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} max_lsn {}",
                lsn,
                ptr,
                snapshot.max_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.max_lsn);
        snapshot.max_lsn = lsn;
        snapshot.last_lid = ptr.lid();

        // invalidate any removed pids
        snapshot.replacements.remove(&segment_idx);

        if !PM::is_null() {
            snapshot.apply(&materializer, lsn, ptr, &*bytes, config);
        }
    }

    write_snapshot(config, &snapshot)?;

    trace!("generated new snapshot: {:?}", snapshot);

    M.advance_snapshot.measure(clock() - start);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default<PM, P, R>(
    config: &Config,
) -> Result<Snapshot<R>, ()>
where
    PM: Materializer<Recovery = R, PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let last_snap =
        read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter_from(last_snap.max_lsn, config)?;

    advance_snapshot::<PM, P, R>(log_iter, last_snap, config)
}

/// Read a `Snapshot` from disk.
fn read_snapshot<R>(
    config: &Config,
) -> std::io::Result<Option<Snapshot<R>>>
where
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let mut candidates = config.get_snapshot_files()?;
    if candidates.is_empty() {
        info!("no previous snapshot found");
        return Ok(None);
    }

    candidates.sort();

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path)?;
    if f.metadata()?.len() <= 16 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf)?;
    let len = buf.len();
    buf.split_off(len - 16);

    let mut len_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-16)).unwrap();
    f.read_exact(&mut len_expected_bytes).unwrap();

    let mut crc_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-8)).unwrap();
    f.read_exact(&mut crc_expected_bytes).unwrap();
    let crc_expected: u64 =
        unsafe { std::mem::transmute(crc_expected_bytes) };

    let crc_actual = crc64(&*buf);

    if crc_expected != crc_actual {
        error!("crc for snapshot file {:?} failed!", path);
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        let len_expected: u64 =
            unsafe { std::mem::transmute(len_expected_bytes) };
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(deserialize::<Snapshot<R>>(&*bytes).ok())
}

pub(crate) fn write_snapshot<R>(
    config: &Config,
    snapshot: &Snapshot<R>,
) -> Result<(), ()>
where
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let raw_bytes = serialize(&snapshot).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.zstd_compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc64: [u8; 8] =
        unsafe { std::mem::transmute(crc64(&*bytes)) };
    let len_bytes: [u8; 8] =
        unsafe { std::mem::transmute(decompressed_len as u64) };

    let path_1_suffix =
        format!("snap.{:016X}.in___motion", snapshot.max_lsn);

    let mut path_1 = config.snapshot_prefix();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.max_lsn);

    let mut path_2 = config.snapshot_prefix();
    path_2.push(path_2_suffix);

    let _res = std::fs::create_dir_all(path_1.parent().unwrap());
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    maybe_fail!("snap write");
    f.write_all(&*bytes)?;
    maybe_fail!("snap write len");
    f.write_all(&len_bytes)?;
    maybe_fail!("snap write crc");
    f.write_all(&crc64)?;
    f.sync_all()?;
    maybe_fail!("snap write post");

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    maybe_fail!("snap write mv");
    std::fs::rename(path_1, &path_2)?;
    maybe_fail!("snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            maybe_fail!("snap write rm old");

            if let Err(_e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
    Ok(())
}
<commit_msg>Bail out of recovery when we hit a corrupted blob<commit_after>use std::collections::{HashMap, HashSet};
use std::io::{Read, Seek, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
/// TODO consider splitting `Snapshot` into separate
/// snapshots for PC, SA, Materializer.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Snapshot<R> {
    /// the last lsn included in the `Snapshot`
    pub max_lsn: Lsn,
    /// the last lid included in the `Snapshot`
    pub last_lid: LogId,
    /// the highest allocated pid
    pub max_pid: PageId,
    /// the mapping from pages to (lsn, lid)
    pub pt: HashMap<PageId, PageState>,
    /// replaced pages per segment index
    pub replacements: HashMap<SegmentId, HashSet<(PageId, Lsn)>>,
    /// the free pids
    pub free: HashSet<PageId>,
    /// the `Materializer`-specific recovered state
    pub recovery: Option<R>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PageState {
    Present(Vec<(Lsn, DiskPtr)>),
    Allocated(Lsn, DiskPtr),
    Free(Lsn, DiskPtr),
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr)) {
        match self {
            &mut PageState::Present(ref mut items) => {
                items.push(item)
            }
            &mut PageState::Allocated(_, _) => {
                *self = PageState::Present(vec![item])
            }
            &mut PageState::Free(_, _) => {
                panic!("pushed items to a PageState::Free")
            }
        }
    }

    /// Iterate over the (lsn, lid) pairs that hold this page's state.
    pub fn iter(&self) -> Box<dyn Iterator<Item = (Lsn, DiskPtr)>> {
        match self {
            &PageState::Present(ref items) => {
                Box::new(items.clone().into_iter())
            }
            &PageState::Allocated(lsn, ptr)
            | &PageState::Free(lsn, ptr) => {
                Box::new(vec![(lsn, ptr)].into_iter())
            }
        }
    }

    fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl<R> Default for Snapshot<R> {
    fn default() -> Snapshot<R> {
        Snapshot {
            max_lsn: 0,
            last_lid: 0,
            max_pid: 0,
            pt: HashMap::new(),
            replacements: HashMap::new(),
            free: HashSet::new(),
            recovery: None,
        }
    }
}

impl<R> Snapshot<R> {
    fn apply<P>(
        &mut self,
        materializer: &dyn Materializer<PageFrag = P, Recovery = R>,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        bytes: &[u8],
        config: &Config,
    ) -> Result<(), ()>
    where
        P: 'static
            + Debug
            + Clone
            + Serialize
            + DeserializeOwned
            + Send
            + Sync,
        R: Debug + Clone + Serialize + DeserializeOwned + Send,
    {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!(
            "trying to deserialize buf for ptr {} lsn {}",
            disk_ptr,
            lsn
        );
        let deserialization = deserialize::<LoggedUpdate<P>>(&*bytes);

        if let Err(e) = deserialization {
            error!(
                "failed to deserialize buffer for item in log: lsn {} \
                    ptr {}: {:?}",
                lsn,
                disk_ptr,
                e
            );
            return Err(Error::Corruption { at: disk_ptr });
        }

        let prepend = deserialization.unwrap();
        let pid = prepend.pid;

        if pid >= self.max_pid {
            self.max_pid = pid + 1;
        }

        let io_buf_size = config.io_buf_size;

        let replaced_at_idx =
            disk_ptr.lid() as SegmentId / io_buf_size;

        match prepend.update {
            Update::Append(partial_page) => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return Ok(());
                    }

                    if let Some(r) =
                        materializer.recover(&partial_page)
                    {
                        self.recovery = Some(r);
                    }

                    lids.push((lsn, disk_ptr));
                }
                self.free.remove(&pid);
            }
            Update::Compact(partial_page) => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                if let Some(r) = materializer.recover(&partial_page) {
                    self.recovery = Some(r);
                }

                self.replace_pid(
                    pid,
                    replaced_at_idx,
                    lsn,
                    io_buf_size,
                    config,
                );
                self.pt.insert(
                    pid,
                    PageState::Present(vec![(lsn, disk_ptr)]),
                );
                self.free.remove(&pid);
            }
            Update::Allocate => {
                trace!(
                    "allocate  of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                self.replace_pid(
                    pid,
                    replaced_at_idx,
                    lsn,
                    io_buf_size,
                    config,
                );
                self.pt
                    .insert(pid, PageState::Allocated(lsn, disk_ptr));
                self.free.remove(&pid);
            }
            Update::Free => {
                trace!(
                    "free of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                self.replace_pid(
                    pid,
                    replaced_at_idx,
                    lsn,
                    io_buf_size,
                    config,
                );
                self.pt.insert(pid, PageState::Free(lsn, disk_ptr));
                self.free.insert(pid);
            }
        }

        Ok(())
    }

    fn replace_pid(
        &mut self,
        pid: PageId,
        replaced_at_idx: usize,
        replaced_at_lsn: Lsn,
        io_buf_size: usize,
        config: &Config,
    ) {
        match self.pt.remove(&pid) {
            Some(PageState::Present(coords)) => {
                for (_lsn, ptr) in &coords {
                    let idx = ptr.lid() as SegmentId / io_buf_size;
                    if replaced_at_idx == idx {
                        return;
                    }
                    let entry = self
                        .replacements
                        .entry(idx)
                        .or_insert(HashSet::new());
                    entry.insert((pid, replaced_at_lsn));
                }

                // re-run any blob removals in case
                // they were not completed.
                if coords.len() > 1 {
                    let blob_ptrs = coords
                        .iter()
                        .filter(|(_, ptr)| ptr.is_blob())
                        .map(|(_, ptr)| ptr.blob().1);

                    for blob_ptr in blob_ptrs {
                        trace!(
                            "removing blob while advancing \
                             snapshot: {}",
                            blob_ptr,
                        );

                        // we don't care if this actually works
                        // because it's possible that a previous
                        // snapshot has run over this log and
                        // removed the blob already.
                        let _ = remove_blob(blob_ptr, config);
                    }
                }
            }
            Some(PageState::Allocated(_lsn, ptr))
            | Some(PageState::Free(_lsn, ptr)) => {
                let idx = ptr.lid() as SegmentId / io_buf_size;
                if replaced_at_idx == idx {
                    return;
                }
                let entry = self
                    .replacements
                    .entry(idx)
                    .or_insert(HashSet::new());
                entry.insert((pid, replaced_at_lsn));
            }
            None => {
                // we just encountered a replace without it
                // existing in the pt. this means we've
                // relocated it to a later segment than
                // the one we're currently at during recovery.
                // so we can skip it here.
            }
        }
    }
}

pub(super) fn advance_snapshot<PM, P, R>(
    iter: LogIter,
    mut snapshot: Snapshot<R>,
    config: &Config,
) -> Result<Snapshot<R>, ()>
where
    PM: Materializer<Recovery = R, PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let start = clock();

    trace!("building on top of old snapshot: {:?}", snapshot);

    let materializer = PM::new(config.clone(), &snapshot.recovery);

    let io_buf_size = config.io_buf_size;

    for (lsn, ptr, bytes) in iter {
        let segment_idx = ptr.lid() as SegmentId / io_buf_size;

        trace!(
            "in advance_snapshot looking at item with lsn {} ptr {}",
            lsn,
            ptr
        );

        if lsn <= snapshot.max_lsn {
            // don't process already-processed Lsn's. max_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} max_lsn {}",
                lsn,
                ptr,
                snapshot.max_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.max_lsn);
        snapshot.max_lsn = lsn;
        snapshot.last_lid = ptr.lid();

        // invalidate any removed pids
        snapshot.replacements.remove(&segment_idx);

        if !PM::is_null() {
            if let Err(e) = snapshot.apply(
                &materializer,
                lsn,
                ptr,
                &*bytes,
                config,
            ) {
                error!(
                    "encountered error while reading log message: {}",
                    e
                );
                break;
            }
        }
    }

    write_snapshot(config, &snapshot)?;

    trace!("generated new snapshot: {:?}", snapshot);

    M.advance_snapshot.measure(clock() - start);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default<PM, P, R>(
    config: &Config,
) -> Result<Snapshot<R>, ()>
where
    PM: Materializer<Recovery = R, PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let last_snap =
        read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter_from(last_snap.max_lsn, config)?;

    advance_snapshot::<PM, P, R>(log_iter, last_snap, config)
}

/// Read a `Snapshot` from disk.
fn read_snapshot<R>(
    config: &Config,
) -> std::io::Result<Option<Snapshot<R>>>
where
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let mut candidates = config.get_snapshot_files()?;
    if candidates.is_empty() {
        info!("no previous snapshot found");
        return Ok(None);
    }

    candidates.sort();

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path)?;
    if f.metadata()?.len() <= 16 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf)?;
    let len = buf.len();
    buf.split_off(len - 16);

    let mut len_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-16)).unwrap();
    f.read_exact(&mut len_expected_bytes).unwrap();

    let mut crc_expected_bytes = [0u8; 8];
    f.seek(std::io::SeekFrom::End(-8)).unwrap();
    f.read_exact(&mut crc_expected_bytes).unwrap();
    let crc_expected: u64 =
        unsafe { std::mem::transmute(crc_expected_bytes) };

    let crc_actual = crc64(&*buf);

    if crc_expected != crc_actual {
        error!("crc for snapshot file {:?} failed!", path);
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        let len_expected: u64 =
            unsafe { std::mem::transmute(len_expected_bytes) };
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(deserialize::<Snapshot<R>>(&*bytes).ok())
}

pub(crate) fn write_snapshot<R>(
    config: &Config,
    snapshot: &Snapshot<R>,
) -> Result<(), ()>
where
    R: Debug + Clone + Serialize + DeserializeOwned + Send,
{
    let raw_bytes = serialize(&snapshot).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.zstd_compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc64: [u8; 8] =
        unsafe { std::mem::transmute(crc64(&*bytes)) };
    let len_bytes: [u8; 8] =
        unsafe { std::mem::transmute(decompressed_len as u64) };

    let path_1_suffix =
        format!("snap.{:016X}.in___motion", snapshot.max_lsn);

    let mut path_1 = config.snapshot_prefix();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.max_lsn);

    let mut path_2 = config.snapshot_prefix();
    path_2.push(path_2_suffix);

    let _res = std::fs::create_dir_all(path_1.parent().unwrap());
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    maybe_fail!("snap write");
    f.write_all(&*bytes)?;
    maybe_fail!("snap write len");
    f.write_all(&len_bytes)?;
    maybe_fail!("snap write crc");
    f.write_all(&crc64)?;
    f.sync_all()?;
    maybe_fail!("snap write post");

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    maybe_fail!("snap write mv");
    std::fs::rename(path_1, &path_2)?;
    maybe_fail!("snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            maybe_fail!("snap write rm old");

            if let Err(_e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>//! This module defines Concrete Syntax Tree (CST), used by rust-analyzer.
//!
//! The CST includes comments and whitespace, provides a single node type,
//! `SyntaxNode`, and a basic traversal API (parent, children, siblings).
//!
//! The *real* implementation is in the (language-agnostic) `rowan` crate, this
//! module just wraps its API.

use rowan::{GreenNodeBuilder, Language};

use crate::{Parse, SyntaxError, SyntaxKind, TextSize};

pub(crate) use rowan::{GreenNode, GreenToken, NodeOrToken};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum RustLanguage {}
impl Language for RustLanguage {
    type Kind = SyntaxKind;

    fn kind_from_raw(raw: rowan::SyntaxKind) -> SyntaxKind {
        SyntaxKind::from(raw.0)
    }

    fn kind_to_raw(kind: SyntaxKind) -> rowan::SyntaxKind {
        rowan::SyntaxKind(kind.into())
    }
}

pub type SyntaxNode = rowan::SyntaxNode<RustLanguage>;
pub type SyntaxToken = rowan::SyntaxToken<RustLanguage>;
pub type SyntaxElement = rowan::SyntaxElement<RustLanguage>;
pub type SyntaxNodeChildren = rowan::SyntaxNodeChildren<RustLanguage>;
pub type SyntaxElementChildren = rowan::SyntaxElementChildren<RustLanguage>;
pub type PreorderWithTokens = rowan::api::PreorderWithTokens<RustLanguage>;

#[derive(Default)]
pub struct SyntaxTreeBuilder {
    errors: Vec<SyntaxError>,
    inner: GreenNodeBuilder<'static>,
}

impl SyntaxTreeBuilder {
    pub(crate) fn finish_raw(self) -> (GreenNode, Vec<SyntaxError>) {
        let green = self.inner.finish();
        (green, self.errors)
    }

    pub fn finish(self) -> Parse<SyntaxNode> {
        let (green, errors) = self.finish_raw();
        // NOTE Why you're disabling assertions.
        // ref: https://github.com/rust-analyzer/rust-analyzer/pull/10357
        if cfg!(debug_assertions) && false {
            let node = SyntaxNode::new_root(green.clone());
            crate::validation::validate_block_structure(&node);
        }
        Parse::new(green, errors)
    }

    pub fn token(&mut self, kind: SyntaxKind, text: &str) {
        let kind = RustLanguage::kind_to_raw(kind);
        self.inner.token(kind, text);
    }

    pub fn start_node(&mut self, kind: SyntaxKind) {
        let kind = RustLanguage::kind_to_raw(kind);
        self.inner.start_node(kind);
    }

    pub fn finish_node(&mut self) {
        self.inner.finish_node();
    }

    pub fn error(&mut self, error: parser::ParseError, text_pos: TextSize) {
        self.errors.push(SyntaxError::new_at_offset(*error.0, text_pos));
    }
}
<commit_msg>remove comment<commit_after>//! This module defines Concrete Syntax Tree (CST), used by rust-analyzer.
//!
//! The CST includes comments and whitespace, provides a single node type,
//! `SyntaxNode`, and a basic traversal API (parent, children, siblings).
//!
//! The *real* implementation is in the (language-agnostic) `rowan` crate, this
//! module just wraps its API.

use rowan::{GreenNodeBuilder, Language};

use crate::{Parse, SyntaxError, SyntaxKind, TextSize};

pub(crate) use rowan::{GreenNode, GreenToken, NodeOrToken};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum RustLanguage {}
impl Language for RustLanguage {
    type Kind = SyntaxKind;

    fn kind_from_raw(raw: rowan::SyntaxKind) -> SyntaxKind {
        SyntaxKind::from(raw.0)
    }

    fn kind_to_raw(kind: SyntaxKind) -> rowan::SyntaxKind {
        rowan::SyntaxKind(kind.into())
    }
}

pub type SyntaxNode = rowan::SyntaxNode<RustLanguage>;
pub type SyntaxToken = rowan::SyntaxToken<RustLanguage>;
pub type SyntaxElement = rowan::SyntaxElement<RustLanguage>;
pub type SyntaxNodeChildren = rowan::SyntaxNodeChildren<RustLanguage>;
pub type SyntaxElementChildren = rowan::SyntaxElementChildren<RustLanguage>;
pub type PreorderWithTokens = rowan::api::PreorderWithTokens<RustLanguage>;

#[derive(Default)]
pub struct SyntaxTreeBuilder {
    errors: Vec<SyntaxError>,
    inner: GreenNodeBuilder<'static>,
}

impl SyntaxTreeBuilder {
    pub(crate) fn finish_raw(self) -> (GreenNode, Vec<SyntaxError>) {
        let green = self.inner.finish();
        (green, self.errors)
    }

    pub fn finish(self) -> Parse<SyntaxNode> {
        let (green, errors) = self.finish_raw();
        if cfg!(debug_assertions) && false {
            let node = SyntaxNode::new_root(green.clone());
            crate::validation::validate_block_structure(&node);
        }
        Parse::new(green, errors)
    }

    pub fn token(&mut self, kind: SyntaxKind, text: &str) {
        let kind = RustLanguage::kind_to_raw(kind);
        self.inner.token(kind, text);
    }

    pub fn start_node(&mut self, kind: SyntaxKind) {
        let kind = RustLanguage::kind_to_raw(kind);
        self.inner.start_node(kind);
    }

    pub fn finish_node(&mut self) {
        self.inner.finish_node();
    }

    pub fn error(&mut self, error: parser::ParseError, text_pos: TextSize) {
        self.errors.push(SyntaxError::new_at_offset(*error.0, text_pos));
    }
}
<|endoftext|>


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

//! This crate provides implementations of the streaming AEAD primitive.
//!
//! AEAD encryption assures the confidentiality and authenticity of the data.
//! This primitive is CPA secure.

#![deny(broken_intra_doc_links)]

use std::sync::Once;
use tink_core::registry::register_key_manager;

mod aes_ctr_hmac_key_manager;
pub use aes_ctr_hmac_key_manager::*;
mod aes_gcm_hkdf_key_manager;
pub use aes_gcm_hkdf_key_manager::*;
mod decrypt_reader;
use decrypt_reader::*;
mod streamingaead_factory;
pub use streamingaead_factory::*;
mod streamingaead_key_templates;
pub use streamingaead_key_templates::*;

pub mod subtle;

/// The [upstream Tink](https://github.com/google/tink) version that this Rust
/// port is based on.
pub const UPSTREAM_VERSION: &str = "1.6.0";

static INIT: Once = Once::new();

/// Initialize the `tink-streaming-aead` crate, registering its primitives so they are available via
/// tink-core.
pub fn init() {
    INIT.call_once(|| {
        register_key_manager(std::sync::Arc::new(AesCtrHmacKeyManager::default()))
            .expect("tink_streaming_aead::init() failed"); // safe: init
        register_key_manager(std::sync::Arc::new(AesGcmHkdfKeyManager::default()))
            .expect("tink_streaming_aead::init() failed"); // safe: init

        tink_core::registry::register_template_generator(
            "AES128_CTR_HMAC_SHA256_4KB",
            aes128_gcm_hkdf_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES128_CTR_HMAC_SHA256_1MB",
            aes128_gcm_hkdf_1mb_key_template,
        );

        tink_core::registry::register_template_generator(
            "AES256_CTR_HMAC_SHA256_4KB",
            aes256_gcm_hkdf_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_CTR_HMAC_SHA256_1MB",
            aes256_gcm_hkdf_1mb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES128_GCM_HKDF_4KB",
            aes128_ctr_hmac_sha256_segment_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES128_GCM_HKDF_1MB",
            aes128_ctr_hmac_sha256_segment_1mb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_GCM_HKDF_4KB",
            aes256_ctr_hmac_sha256_segment_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_GCM_HKDF_1MB",
            aes256_ctr_hmac_sha256_segment_1mb_key_template,
        );
    });
}
<commit_msg>tink-streaming-aead: fix template registration<commit_after>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

//! This crate provides implementations of the streaming AEAD primitive.
//!
//! AEAD encryption assures the confidentiality and authenticity of the data.
//! This primitive is CPA secure.

#![deny(broken_intra_doc_links)]

use std::sync::Once;
use tink_core::registry::register_key_manager;

mod aes_ctr_hmac_key_manager;
pub use aes_ctr_hmac_key_manager::*;
mod aes_gcm_hkdf_key_manager;
pub use aes_gcm_hkdf_key_manager::*;
mod decrypt_reader;
use decrypt_reader::*;
mod streamingaead_factory;
pub use streamingaead_factory::*;
mod streamingaead_key_templates;
pub use streamingaead_key_templates::*;

pub mod subtle;

/// The [upstream Tink](https://github.com/google/tink) version that this Rust
/// port is based on.
pub const UPSTREAM_VERSION: &str = "1.6.0";

static INIT: Once = Once::new();

/// Initialize the `tink-streaming-aead` crate, registering its primitives so they are available via
/// tink-core.
pub fn init() {
    INIT.call_once(|| {
        register_key_manager(std::sync::Arc::new(AesCtrHmacKeyManager::default()))
            .expect("tink_streaming_aead::init() failed"); // safe: init
        register_key_manager(std::sync::Arc::new(AesGcmHkdfKeyManager::default()))
            .expect("tink_streaming_aead::init() failed"); // safe: init

        tink_core::registry::register_template_generator(
            "AES128_CTR_HMAC_SHA256_4KB",
            aes128_ctr_hmac_sha256_segment_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES128_CTR_HMAC_SHA256_1MB",
            aes128_ctr_hmac_sha256_segment_1mb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_CTR_HMAC_SHA256_4KB",
            aes256_ctr_hmac_sha256_segment_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_CTR_HMAC_SHA256_1MB",
            aes256_ctr_hmac_sha256_segment_1mb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES128_GCM_HKDF_4KB",
            aes128_gcm_hkdf_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES128_GCM_HKDF_1MB",
            aes128_gcm_hkdf_1mb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_GCM_HKDF_4KB",
            aes256_gcm_hkdf_4kb_key_template,
        );
        tink_core::registry::register_template_generator(
            "AES256_GCM_HKDF_1MB",
            aes256_gcm_hkdf_1mb_key_template,
        );
    });
}
<|endoftext|>


<commit_before>use std::mem;

use ra_syntax::{
    TextUnit, TextRange, SyntaxNodeRef, File, AstNode, SyntaxKind,
    ast,
    algo::{
        find_covering_node, find_leaf_at_offset, LeafAtOffset,
    },
    text_utils::{intersect, contains_offset_nonstrict},
    SyntaxKind::*,
};

use {LocalEdit, EditBuilder, find_node_at_offset};

pub fn join_lines(file: &File, range: TextRange) -> LocalEdit {
    let range = if range.is_empty() {
        let syntax = file.syntax();
        let text = syntax.text().slice(range.start()..);
        let pos = match text.find('\n') {
            None => return LocalEdit {
                edit: EditBuilder::new().finish(),
                cursor_position: None
            },
            Some(pos) => pos
        };
        TextRange::offset_len(
            range.start() + pos,
            TextUnit::of_char('\n'),
        )
    } else {
        range
    };

    let node = find_covering_node(file.syntax(), range);
    let mut edit = EditBuilder::new();
    for node in node.descendants() {
        let text = match node.leaf_text() {
            Some(text) => text,
            None => continue,
        };
        let range = match intersect(range, node.range()) {
            Some(range) => range,
            None => continue,
        } - node.range().start();
        for (pos, _) in text[range].bytes().enumerate().filter(|&(_, b)| b == b'\n') {
            let pos: TextUnit = (pos as u32).into();
            let off = node.range().start() + range.start() + pos;
            if !edit.invalidates_offset(off) {
                remove_newline(&mut edit, node, text.as_str(), off);
            }
        }
    }

    LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    }
}

pub fn on_enter(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let comment = find_leaf_at_offset(file.syntax(), offset).left_biased().and_then(|it| ast::Comment::cast(it))?;

    if let ast::CommentFlavor::Multiline = comment.flavor() {
        return None;
    }

    let prefix = comment.prefix();
    if offset < comment.syntax().range().start() + TextUnit::of_str(prefix) + TextUnit::from(1) {
        return None;
    }

    let indent = node_indent(file, comment.syntax())?;
    let inserted = format!("\n{}{} ", indent, prefix);
    let cursor_position = offset + TextUnit::of_str(&inserted);
    let mut edit = EditBuilder::new();
    edit.insert(offset, inserted);
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: Some(cursor_position),
    })
}

fn node_indent<'a>(file: &'a File, node: SyntaxNodeRef) -> Option<&'a str> {
    let ws = match find_leaf_at_offset(file.syntax(), node.range().start()) {
        LeafAtOffset::Between(l, r) => {
            assert!(r == node);
            l
        }
        LeafAtOffset::Single(n) => {
            assert!(n == node);
            return Some("")
        }
        LeafAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.leaf_text().unwrap();
    let pos = text.as_str().rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(&text[pos..])
}

pub fn on_eq_typed(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
    if let_stmt.has_semi() {
        return None;
    }
    if let Some(expr) = let_stmt.initializer() {
        let expr_range = expr.syntax().range();
        if contains_offset_nonstrict(expr_range, offset) && offset != expr_range.start() {
            return None;
        }
        if file.syntax().text().slice(offset..expr_range.start()).contains('\n') {
            return None;
        }
    } else {
        return None;
    }
    let offset = let_stmt.syntax().range().end();
    let mut edit = EditBuilder::new();
    edit.insert(offset, ";".to_string());
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    })
}

fn remove_newline(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
    node_text: &str,
    offset: TextUnit,
) {
    if node.kind() == WHITESPACE && node_text.bytes().filter(|&b| b == b'\n').count() == 1 {
        // Special case that turns something like:
        //
        // ```
        // my_function({<|>
        //    <some-expr>
        // })
        // ```
        //
        // into `my_function(<some-expr>)`
        if join_single_expr_block(edit, node).is_some() {
            return
        }

        if let (Some(prev), Some(next)) = (node.prev_sibling(), node.next_sibling()) {
            let range = TextRange::from_to(prev.range().start(), node.range().end());
            if is_trailing_comma(prev.kind(), next.kind()) {
                // Removes: trailing comma, newline (incl. surrounding whitespace)
                edit.delete(range);
            } else if no_space_required(prev.kind(), next.kind()) {
                // Removes: newline (incl. surrounding whitespace)
                edit.delete(node.range());
            } else if prev.kind() == COMMA && next.kind() == R_CURLY {
                // Removes: comma, newline (incl. surrounding whitespace)
                // Adds: a single whitespace
                edit.replace(range, " ".to_string());
            } else if let (Some(_), Some(next)) = (ast::Comment::cast(prev), ast::Comment::cast(next)) {
                // Removes: newline (incl. surrounding whitespace), start of the next comment
                let comment_text = next.text();
                if let Some(newline_pos) = comment_text.find('\n') {
                    // Special case for multi-line c-like comments: join the comment content but
                    // keep the leading `/*`

                    let newline_offset = next.syntax().range().start()
                                        + TextUnit::from(newline_pos as u32)
                                        + TextUnit::of_char('\n');

                    edit.insert(newline_offset, "/*".to_string());
                    edit.delete(TextRange::from_to(
                        node.range().start(),
                        next.syntax().range().start() + TextUnit::of_str(next.prefix())
                    ));
                } else {
                    // Single-line comments
                    edit.delete(TextRange::from_to(
                        node.range().start(),
                        next.syntax().range().start() + TextUnit::of_str(next.prefix())
                    ));
                }
            } else {
                // Remove newline but add a computed amount of whitespace characters
                edit.replace(
                    node.range(),
                    compute_ws(prev, next).to_string(),
                );
            }

            return;
        }
    }

    // The node is either the first or the last in the file
    let suff = &node_text[TextRange::from_to(
        offset - node.range().start() + TextUnit::of_char('\n'),
        TextUnit::of_str(node_text),
    )];
    let spaces = suff.bytes().take_while(|&b| b == b' ').count();

    edit.replace(
        TextRange::offset_len(offset, ((spaces + 1) as u32).into()),
        " ".to_string(),
    );
}

fn is_trailing_comma(left: SyntaxKind, right: SyntaxKind) -> bool {
    match (left, right) {
       (COMMA, R_PAREN) | (COMMA, R_BRACK) => true,
       _ => false
    }
}

fn no_space_required(left: SyntaxKind, right: SyntaxKind) -> bool {
    match (left, right) {
       (_, DOT) => true,
        _ => false
    }
}

fn join_single_expr_block(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
) -> Option<()> {
    let block = ast::Block::cast(node.parent()?)?;
    let block_expr = ast::BlockExpr::cast(block.syntax().parent()?)?;
    let expr = single_expr(block)?;
    edit.replace(
        block_expr.syntax().range(),
        expr.syntax().text().to_string(),
    );
    Some(())
}

fn single_expr(block: ast::Block) -> Option<ast::Expr> {
    let mut res = None;
    for child in block.syntax().children() {
        if let Some(expr) = ast::Expr::cast(child) {
            if expr.syntax().text().contains('\n') {
                return None;
            }
            if mem::replace(&mut res, Some(expr)).is_some() {
                return None;
            }
        } else {
            match child.kind() {
                WHITESPACE | L_CURLY | R_CURLY => (),
                _ => return None,
            }
        }
    }
    res
}

fn compute_ws(left: SyntaxNodeRef, right: SyntaxNodeRef) -> &'static str {
    match left.kind() {
        L_PAREN | L_BRACK => return "",
        _ => (),
    }
    match right.kind() {
        R_PAREN | R_BRACK => return "",
        _ => (),
    }
    " "
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_utils::{check_action, extract_range, extract_offset, add_cursor};

    fn check_join_lines(before: &str, after: &str) {
        check_action(before, after, |file, offset| {
            let range = TextRange::offset_len(offset, 0.into());
            let res = join_lines(file, range);
            Some(res)
        })
    }

    #[test]
    fn test_join_lines_comma() {
        check_join_lines(r"
fn foo() {
    <|>foo(1,
    )
}
", r"
fn foo() {
    <|>foo(1)
}
");
    }

    #[test]
    fn test_join_lines_lambda_block() {
        check_join_lines(r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| {
        self.full_reparse(edit)
    })
}
", r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| self.full_reparse(edit))
}
");
    }

    #[test]
    fn test_join_lines_block() {
        check_join_lines(r"
fn foo() {
    foo(<|>{
        92
    })
}", r"
fn foo() {
    foo(<|>92)
}");
    }

    fn check_join_lines_sel(before: &str, after: &str) {
        let (sel, before) = extract_range(before);
        let file = File::parse(&before);
        let result = join_lines(&file, sel);
        let actual = result.edit.apply(&before);
        assert_eq_text!(after, &actual);
    }

    #[test]
    fn test_join_lines_selection_fn_args() {
        check_join_lines_sel(r"
fn foo() {
    <|>foo(1,
        2,
        3,
    <|>)
}
    ", r"
fn foo() {
    foo(1, 2, 3)
}
    ");
    }

    #[test]
    fn test_join_lines_selection_struct() {
        check_join_lines_sel(r"
struct Foo <|>{
    f: u32,
}<|>
    ", r"
struct Foo { f: u32 }
    ");
    }

    #[test]
    fn test_join_lines_selection_dot_chain() {
        check_join_lines_sel(r"
fn foo() {
    join(<|>type_params.type_params()
            .filter_map(|it| it.name())
            .map(|it| it.text())<|>)
}", r"
fn foo() {
    join(type_params.type_params().filter_map(|it| it.name()).map(|it| it.text()))
}");
    }

    #[test]
    fn test_join_lines_selection_lambda_block_body() {
        check_join_lines_sel(r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| <|>{
            world.analysis().matching_brace(&file, offset).unwrap_or(offset)
        }<|>)
        .collect();
}", r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| world.analysis().matching_brace(&file, offset).unwrap_or(offset))
        .collect();
}");
    }

    #[test]
    fn test_on_eq_typed() {
        fn do_check(before: &str, after: &str) {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_eq_typed(&file, offset).unwrap();
            let actual = result.edit.apply(&before);
            assert_eq_text!(after, &actual);
        }

        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        // }
        // ");
        do_check(r"
fn foo() {
    let foo =<|> 1 + 1
}
", r"
fn foo() {
    let foo = 1 + 1;
}
");
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        //     let bar = 1;
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        //     let bar = 1;
        // }
        // ");
    }

    #[test]
    fn test_on_enter() {
        fn apply_on_enter(before: &str) -> Option<String> {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_enter(&file, offset)?;
            let actual = result.edit.apply(&before);
            let actual = add_cursor(&actual, result.cursor_position.unwrap());
            Some(actual)
        }

        fn do_check(before: &str, after: &str) {
            let actual = apply_on_enter(before).unwrap();
            assert_eq_text!(after, &actual);
        }

        fn do_check_noop(text: &str) {
            assert!(apply_on_enter(text).is_none())
        }

        do_check(r"
/// Some docs<|>
fn foo() {
}
", r"
/// Some docs
/// <|>
fn foo() {
}
");
        do_check(r"
impl S {
    /// Some<|> docs.
    fn foo() {}
}
", r"
impl S {
    /// Some
    /// <|> docs.
    fn foo() {}
}
");
        do_check_noop(r"<|>//! docz");
    }
}
<commit_msg>Remove nesting<commit_after>use std::mem;

use ra_syntax::{
    TextUnit, TextRange, SyntaxNodeRef, File, AstNode, SyntaxKind,
    ast,
    algo::{
        find_covering_node, find_leaf_at_offset, LeafAtOffset,
    },
    text_utils::{intersect, contains_offset_nonstrict},
    SyntaxKind::*,
};

use {LocalEdit, EditBuilder, find_node_at_offset};

pub fn join_lines(file: &File, range: TextRange) -> LocalEdit {
    let range = if range.is_empty() {
        let syntax = file.syntax();
        let text = syntax.text().slice(range.start()..);
        let pos = match text.find('\n') {
            None => return LocalEdit {
                edit: EditBuilder::new().finish(),
                cursor_position: None
            },
            Some(pos) => pos
        };
        TextRange::offset_len(
            range.start() + pos,
            TextUnit::of_char('\n'),
        )
    } else {
        range
    };

    let node = find_covering_node(file.syntax(), range);
    let mut edit = EditBuilder::new();
    for node in node.descendants() {
        let text = match node.leaf_text() {
            Some(text) => text,
            None => continue,
        };
        let range = match intersect(range, node.range()) {
            Some(range) => range,
            None => continue,
        } - node.range().start();
        for (pos, _) in text[range].bytes().enumerate().filter(|&(_, b)| b == b'\n') {
            let pos: TextUnit = (pos as u32).into();
            let off = node.range().start() + range.start() + pos;
            if !edit.invalidates_offset(off) {
                remove_newline(&mut edit, node, text.as_str(), off);
            }
        }
    }

    LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    }
}

pub fn on_enter(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let comment = find_leaf_at_offset(file.syntax(), offset).left_biased().and_then(|it| ast::Comment::cast(it))?;

    if let ast::CommentFlavor::Multiline = comment.flavor() {
        return None;
    }

    let prefix = comment.prefix();
    if offset < comment.syntax().range().start() + TextUnit::of_str(prefix) + TextUnit::from(1) {
        return None;
    }

    let indent = node_indent(file, comment.syntax())?;
    let inserted = format!("\n{}{} ", indent, prefix);
    let cursor_position = offset + TextUnit::of_str(&inserted);
    let mut edit = EditBuilder::new();
    edit.insert(offset, inserted);
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: Some(cursor_position),
    })
}

fn node_indent<'a>(file: &'a File, node: SyntaxNodeRef) -> Option<&'a str> {
    let ws = match find_leaf_at_offset(file.syntax(), node.range().start()) {
        LeafAtOffset::Between(l, r) => {
            assert!(r == node);
            l
        }
        LeafAtOffset::Single(n) => {
            assert!(n == node);
            return Some("")
        }
        LeafAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.leaf_text().unwrap();
    let pos = text.as_str().rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(&text[pos..])
}

pub fn on_eq_typed(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
    if let_stmt.has_semi() {
        return None;
    }
    if let Some(expr) = let_stmt.initializer() {
        let expr_range = expr.syntax().range();
        if contains_offset_nonstrict(expr_range, offset) && offset != expr_range.start() {
            return None;
        }
        if file.syntax().text().slice(offset..expr_range.start()).contains('\n') {
            return None;
        }
    } else {
        return None;
    }
    let offset = let_stmt.syntax().range().end();
    let mut edit = EditBuilder::new();
    edit.insert(offset, ";".to_string());
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    })
}

fn remove_newline(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
    node_text: &str,
    offset: TextUnit,
) {
    if node.kind() != WHITESPACE || node_text.bytes().filter(|&b| b == b'\n').count() != 1 {
        // The node is either the first or the last in the file
        let suff = &node_text[TextRange::from_to(
            offset - node.range().start() + TextUnit::of_char('\n'),
            TextUnit::of_str(node_text),
        )];
        let spaces = suff.bytes().take_while(|&b| b == b' ').count();

        edit.replace(
            TextRange::offset_len(offset, ((spaces + 1) as u32).into()),
            " ".to_string(),
        );
        return;
    }

    // Special case that turns something like:
    //
    // ```
    // my_function({<|>
    //    <some-expr>
    // })
    // ```
    //
    // into `my_function(<some-expr>)`
    if join_single_expr_block(edit, node).is_some() {
        return
    }

    // The node is between two other nodes
    let prev = node.prev_sibling().unwrap();
    let next = node.next_sibling().unwrap();
    if is_trailing_comma(prev.kind(), next.kind()) {
        // Removes: trailing comma, newline (incl. surrounding whitespace)
        edit.delete(TextRange::from_to(prev.range().start(), node.range().end()));
    } else if prev.kind() == COMMA && next.kind() == R_CURLY {
        // Removes: comma, newline (incl. surrounding whitespace)
        // Adds: a single whitespace
        edit.replace(
            TextRange::from_to(prev.range().start(), node.range().end()),
            " ".to_string()
        );
    } else if let (Some(_), Some(next)) = (ast::Comment::cast(prev), ast::Comment::cast(next)) {
        // Removes: newline (incl. surrounding whitespace), start of the next comment
        let comment_text = next.text();
        if let Some(newline_pos) = comment_text.find('\n') {
            // Special case for multiline comments: join the comment content but
            // keep the leading `/*`

            let newline_offset = next.syntax().range().start()
                                + TextUnit::from(newline_pos as u32)
                                + TextUnit::of_char('\n');

            edit.insert(newline_offset, "/*".to_string());
            edit.delete(TextRange::from_to(
                node.range().start(),
                next.syntax().range().start() + TextUnit::of_str(next.prefix())
            ));
        } else {
            // Single-line comments
            edit.delete(TextRange::from_to(
                node.range().start(),
                next.syntax().range().start() + TextUnit::of_str(next.prefix())
            ));
        }
    } else {
        // Remove newline but add a computed amount of whitespace characters
        edit.replace(
            node.range(),
            compute_ws(prev, next).to_string(),
        );
    }
}

fn is_trailing_comma(left: SyntaxKind, right: SyntaxKind) -> bool {
    match (left, right) {
       (COMMA, R_PAREN) | (COMMA, R_BRACK) => true,
       _ => false
    }
}

fn join_single_expr_block(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
) -> Option<()> {
    let block = ast::Block::cast(node.parent()?)?;
    let block_expr = ast::BlockExpr::cast(block.syntax().parent()?)?;
    let expr = single_expr(block)?;
    edit.replace(
        block_expr.syntax().range(),
        expr.syntax().text().to_string(),
    );
    Some(())
}

fn single_expr(block: ast::Block) -> Option<ast::Expr> {
    let mut res = None;
    for child in block.syntax().children() {
        if let Some(expr) = ast::Expr::cast(child) {
            if expr.syntax().text().contains('\n') {
                return None;
            }
            if mem::replace(&mut res, Some(expr)).is_some() {
                return None;
            }
        } else {
            match child.kind() {
                WHITESPACE | L_CURLY | R_CURLY => (),
                _ => return None,
            }
        }
    }
    res
}

fn compute_ws(left: SyntaxNodeRef, right: SyntaxNodeRef) -> &'static str {
    match left.kind() {
        L_PAREN | L_BRACK => return "",
        _ => (),
    }
    match right.kind() {
        R_PAREN | R_BRACK => return "",
        DOT => return "",
        _ => (),
    }
    " "
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_utils::{check_action, extract_range, extract_offset, add_cursor};

    fn check_join_lines(before: &str, after: &str) {
        check_action(before, after, |file, offset| {
            let range = TextRange::offset_len(offset, 0.into());
            let res = join_lines(file, range);
            Some(res)
        })
    }

    #[test]
    fn test_join_lines_comma() {
        check_join_lines(r"
fn foo() {
    <|>foo(1,
    )
}
", r"
fn foo() {
    <|>foo(1)
}
");
    }

    #[test]
    fn test_join_lines_lambda_block() {
        check_join_lines(r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| {
        self.full_reparse(edit)
    })
}
", r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| self.full_reparse(edit))
}
");
    }

    #[test]
    fn test_join_lines_block() {
        check_join_lines(r"
fn foo() {
    foo(<|>{
        92
    })
}", r"
fn foo() {
    foo(<|>92)
}");
    }

    fn check_join_lines_sel(before: &str, after: &str) {
        let (sel, before) = extract_range(before);
        let file = File::parse(&before);
        let result = join_lines(&file, sel);
        let actual = result.edit.apply(&before);
        assert_eq_text!(after, &actual);
    }

    #[test]
    fn test_join_lines_selection_fn_args() {
        check_join_lines_sel(r"
fn foo() {
    <|>foo(1,
        2,
        3,
    <|>)
}
    ", r"
fn foo() {
    foo(1, 2, 3)
}
    ");
    }

    #[test]
    fn test_join_lines_selection_struct() {
        check_join_lines_sel(r"
struct Foo <|>{
    f: u32,
}<|>
    ", r"
struct Foo { f: u32 }
    ");
    }

    #[test]
    fn test_join_lines_selection_dot_chain() {
        check_join_lines_sel(r"
fn foo() {
    join(<|>type_params.type_params()
            .filter_map(|it| it.name())
            .map(|it| it.text())<|>)
}", r"
fn foo() {
    join(type_params.type_params().filter_map(|it| it.name()).map(|it| it.text()))
}");
    }

    #[test]
    fn test_join_lines_selection_lambda_block_body() {
        check_join_lines_sel(r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| <|>{
            world.analysis().matching_brace(&file, offset).unwrap_or(offset)
        }<|>)
        .collect();
}", r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| world.analysis().matching_brace(&file, offset).unwrap_or(offset))
        .collect();
}");
    }

    #[test]
    fn test_on_eq_typed() {
        fn do_check(before: &str, after: &str) {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_eq_typed(&file, offset).unwrap();
            let actual = result.edit.apply(&before);
            assert_eq_text!(after, &actual);
        }

        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        // }
        // ");
        do_check(r"
fn foo() {
    let foo =<|> 1 + 1
}
", r"
fn foo() {
    let foo = 1 + 1;
}
");
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        //     let bar = 1;
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        //     let bar = 1;
        // }
        // ");
    }

    #[test]
    fn test_on_enter() {
        fn apply_on_enter(before: &str) -> Option<String> {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_enter(&file, offset)?;
            let actual = result.edit.apply(&before);
            let actual = add_cursor(&actual, result.cursor_position.unwrap());
            Some(actual)
        }

        fn do_check(before: &str, after: &str) {
            let actual = apply_on_enter(before).unwrap();
            assert_eq_text!(after, &actual);
        }

        fn do_check_noop(text: &str) {
            assert!(apply_on_enter(text).is_none())
        }

        do_check(r"
/// Some docs<|>
fn foo() {
}
", r"
/// Some docs
/// <|>
fn foo() {
}
");
        do_check(r"
impl S {
    /// Some<|> docs.
    fn foo() {}
}
", r"
impl S {
    /// Some
    /// <|> docs.
    fn foo() {}
}
");
        do_check_noop(r"<|>//! docz");
    }
}
<|endoftext|>


<commit_before>// You can run this example from the root of the mio repo:
// cargo run --example tcp_server --features="os-poll net"
use mio::event::Event;
use mio::net::{TcpListener, TcpStream};
use mio::{Events, Interest, Poll, Registry, Token};
use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::str::from_utf8;

// Setup some tokens to allow us to identify which event is for which socket.
const SERVER: Token = Token(0);

// Some data we'll send over the connection.
const DATA: &[u8] = b"Hello world!\n";

fn main() -> io::Result<()> {
    env_logger::init();

    // Create a poll instance.
    let mut poll = Poll::new()?;
    // Create storage for events.
    let mut events = Events::with_capacity(128);

    // Setup the TCP server socket.
    let addr = "127.0.0.1:9000".parse().unwrap();
    let mut server = TcpListener::bind(addr)?;

    // Register the server with poll we can receive events for it.
    poll.registry()
        .register(&mut server, SERVER, Interest::READABLE)?;

    // Map of `Token` -> `TcpStream`.
    let mut connections = HashMap::new();
    // Unique token for each incoming connection.
    let mut unique_token = Token(SERVER.0 + 1);

    println!("You can connect to the server using `nc`:");
    println!(" $ nc 127.0.0.1 9000");
    println!("You'll see our welcome message and anything you type we'll be printed here.");

    loop {
        poll.poll(&mut events, None)?;

        for event in events.iter() {
            match event.token() {
                SERVER => loop {
                    // Received an event for the TCP server socket, which
                    // indicates we can accept an connection.
                    let (mut connection, address) = match server.accept() {
                        Ok((connection, address)) => (connection, address),
                        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                            // If we get a `WouldBlock` error we know our
                            // listener has no more incoming connections queued,
                            // so we can return to polling and wait for some
                            // more.
                            break;
                        }
                        Err(e) => {
                            // If it was any other kind of error, something went
                            // wrong and we terminate with an error.
                            return Err(e);
                        }
                    };

                    println!("Accepted connection from: {}", address);

                    let token = next(&mut unique_token);
                    poll.registry().register(
                        &mut connection,
                        token,
                        Interest::READABLE.add(Interest::WRITABLE),
                    )?;

                    connections.insert(token, connection);
                },
                token => {
                    // Maybe received an event for a TCP connection.
                    let done = if let Some(connection) = connections.get_mut(&token) {
                        handle_connection_event(poll.registry(), connection, event)?
                    } else {
                        // Sporadic events happen, we can safely ignore them.
                        false
                    };
                    if done {
                        connections.remove(&token);
                    }
                }
            }
        }
    }
}

fn next(current: &mut Token) -> Token {
    let next = current.0;
    current.0 += 1;
    Token(next)
}

/// Returns `true` if the connection is done.
fn handle_connection_event(
    registry: &Registry,
    connection: &mut TcpStream,
    event: &Event,
) -> io::Result<bool> {
    if event.is_writable() {
        // We can (maybe) write to the connection.
        match connection.write(DATA) {
            // We want to write the entire `DATA` buffer in a single go. If we
            // write less we'll return a short write error (same as
            // `io::Write::write_all` does).
            Ok(n) if n < DATA.len() => return Err(io::ErrorKind::WriteZero.into()),
            Ok(_) => {
                // After we've written something we'll reregister the connection
                // to only respond to readable events.
                registry.reregister(connection, event.token(), Interest::READABLE)?
            }
            // Would block "errors" are the OS's way of saying that the
            // connection is not actually ready to perform this I/O operation.
            Err(ref err) if would_block(err) => {}
            // Got interrupted (how rude!), we'll try again.
            Err(ref err) if interrupted(err) => {
                return handle_connection_event(registry, connection, event)
            }
            // Other errors we'll consider fatal.
            Err(err) => return Err(err),
        }
    }

    if event.is_readable() {
        let mut connection_closed = false;
        let mut received_data = vec![0; 4096];
        let mut bytes_read = 0;
        // We can (maybe) read from the connection.
        loop {
            match connection.read(&mut received_data[bytes_read..]) {
                Ok(0) => {
                    // Reading 0 bytes means the other side has closed the
                    // connection or is done writing, then so are we.
                    connection_closed = true;
                    break;
                }
                Ok(n) => {
                    bytes_read += n;
                    if bytes_read == received_data.len() {
                        received_data.resize(received_data.len() + 1024, 0);
                    }
                }
                // Would block "errors" are the OS's way of saying that the
                // connection is not actually ready to perform this I/O operation.
                Err(ref err) if would_block(err) => break,
                Err(ref err) if interrupted(err) => continue,
                // Other errors we'll consider fatal.
                Err(err) => return Err(err),
            }
        }

        if bytes_read != 0 {
            let received_data = &received_data[..bytes_read];
            if let Ok(str_buf) = from_utf8(received_data) {
                println!("Received data: {}", str_buf.trim_end());
            } else {
                println!("Received (none UTF-8) data: {:?}", received_data);
            }
        }

        if connection_closed {
            println!("Connection closed");
            return Ok(true);
        }
    }

    Ok(false)
}

fn would_block(err: &io::Error) -> bool {
    err.kind() == io::ErrorKind::WouldBlock
}

fn interrupted(err: &io::Error) -> bool {
    err.kind() == io::ErrorKind::Interrupted
}
<commit_msg>Small spelling correction in example<commit_after>// You can run this example from the root of the mio repo:
// cargo run --example tcp_server --features="os-poll net"
use mio::event::Event;
use mio::net::{TcpListener, TcpStream};
use mio::{Events, Interest, Poll, Registry, Token};
use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::str::from_utf8;

// Setup some tokens to allow us to identify which event is for which socket.
const SERVER: Token = Token(0);

// Some data we'll send over the connection.
const DATA: &[u8] = b"Hello world!\n";

fn main() -> io::Result<()> {
    env_logger::init();

    // Create a poll instance.
    let mut poll = Poll::new()?;
    // Create storage for events.
    let mut events = Events::with_capacity(128);

    // Setup the TCP server socket.
    let addr = "127.0.0.1:9000".parse().unwrap();
    let mut server = TcpListener::bind(addr)?;

    // Register the server with poll we can receive events for it.
    poll.registry()
        .register(&mut server, SERVER, Interest::READABLE)?;

    // Map of `Token` -> `TcpStream`.
    let mut connections = HashMap::new();
    // Unique token for each incoming connection.
    let mut unique_token = Token(SERVER.0 + 1);

    println!("You can connect to the server using `nc`:");
    println!(" $ nc 127.0.0.1 9000");
    println!("You'll see our welcome message and anything you type will be printed here.");

    loop {
        poll.poll(&mut events, None)?;

        for event in events.iter() {
            match event.token() {
                SERVER => loop {
                    // Received an event for the TCP server socket, which
                    // indicates we can accept an connection.
                    let (mut connection, address) = match server.accept() {
                        Ok((connection, address)) => (connection, address),
                        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                            // If we get a `WouldBlock` error we know our
                            // listener has no more incoming connections queued,
                            // so we can return to polling and wait for some
                            // more.
                            break;
                        }
                        Err(e) => {
                            // If it was any other kind of error, something went
                            // wrong and we terminate with an error.
                            return Err(e);
                        }
                    };

                    println!("Accepted connection from: {}", address);

                    let token = next(&mut unique_token);
                    poll.registry().register(
                        &mut connection,
                        token,
                        Interest::READABLE.add(Interest::WRITABLE),
                    )?;

                    connections.insert(token, connection);
                },
                token => {
                    // Maybe received an event for a TCP connection.
                    let done = if let Some(connection) = connections.get_mut(&token) {
                        handle_connection_event(poll.registry(), connection, event)?
                    } else {
                        // Sporadic events happen, we can safely ignore them.
                        false
                    };
                    if done {
                        connections.remove(&token);
                    }
                }
            }
        }
    }
}

fn next(current: &mut Token) -> Token {
    let next = current.0;
    current.0 += 1;
    Token(next)
}

/// Returns `true` if the connection is done.
fn handle_connection_event(
    registry: &Registry,
    connection: &mut TcpStream,
    event: &Event,
) -> io::Result<bool> {
    if event.is_writable() {
        // We can (maybe) write to the connection.
        match connection.write(DATA) {
            // We want to write the entire `DATA` buffer in a single go. If we
            // write less we'll return a short write error (same as
            // `io::Write::write_all` does).
            Ok(n) if n < DATA.len() => return Err(io::ErrorKind::WriteZero.into()),
            Ok(_) => {
                // After we've written something we'll reregister the connection
                // to only respond to readable events.
                registry.reregister(connection, event.token(), Interest::READABLE)?
            }
            // Would block "errors" are the OS's way of saying that the
            // connection is not actually ready to perform this I/O operation.
            Err(ref err) if would_block(err) => {}
            // Got interrupted (how rude!), we'll try again.
            Err(ref err) if interrupted(err) => {
                return handle_connection_event(registry, connection, event)
            }
            // Other errors we'll consider fatal.
            Err(err) => return Err(err),
        }
    }

    if event.is_readable() {
        let mut connection_closed = false;
        let mut received_data = vec![0; 4096];
        let mut bytes_read = 0;
        // We can (maybe) read from the connection.
        loop {
            match connection.read(&mut received_data[bytes_read..]) {
                Ok(0) => {
                    // Reading 0 bytes means the other side has closed the
                    // connection or is done writing, then so are we.
                    connection_closed = true;
                    break;
                }
                Ok(n) => {
                    bytes_read += n;
                    if bytes_read == received_data.len() {
                        received_data.resize(received_data.len() + 1024, 0);
                    }
                }
                // Would block "errors" are the OS's way of saying that the
                // connection is not actually ready to perform this I/O operation.
                Err(ref err) if would_block(err) => break,
                Err(ref err) if interrupted(err) => continue,
                // Other errors we'll consider fatal.
                Err(err) => return Err(err),
            }
        }

        if bytes_read != 0 {
            let received_data = &received_data[..bytes_read];
            if let Ok(str_buf) = from_utf8(received_data) {
                println!("Received data: {}", str_buf.trim_end());
            } else {
                println!("Received (none UTF-8) data: {:?}", received_data);
            }
        }

        if connection_closed {
            println!("Connection closed");
            return Ok(true);
        }
    }

    Ok(false)
}

fn would_block(err: &io::Error) -> bool {
    err.kind() == io::ErrorKind::WouldBlock
}

fn interrupted(err: &io::Error) -> bool {
    err.kind() == io::ErrorKind::Interrupted
}
<|endoftext|>