<commit_before>use std::marker::PhantomData;

use ra_db::{LocationIntener, FileId};
use ra_syntax::{TreeArc, SyntaxNode, SourceFile, AstNode, ast};
use ra_arena::{Arena, RawId, impl_arena_id};

use crate::{
    HirDatabase, Def, Struct, Enum, EnumVariant, Crate,
    Module, Trait, Type, Static, Const,
};

#[derive(Debug, Default)]
pub struct HirInterner {
    defs: LocationIntener<DefLoc, DefId>,
    macros: LocationIntener<MacroCallLoc, MacroCallId>,
    fns: LocationIntener<FunctionLoc, FunctionId>,
}

impl HirInterner {
    pub fn len(&self) -> usize {
        self.defs.len() + self.macros.len()
    }
}

/// hir makes heavy use of ids: integer (u32) handlers to various things. You
/// can think of id as a pointer (but without a lifetime) or a file descriptor
/// (but for hir objects).
///
/// This module defines a bunch of ids we are using. The most important ones are
/// probably `HirFileId` and `DefId`.

/// Input to the analyzer is a set of files, where each file is indentified by
/// `FileId` and contains source code. However, another source of source code in
/// Rust are macros: each macro can be thought of as producing a "temporary
/// file". To assign an id to such a file, we use the id of the macro call that
/// produced the file. So, a `HirFileId` is either a `FileId` (source code
/// written by user), or a `MacroCallId` (source code produced by macro).
///
/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file containin
/// the call plus the offset of the macro call in the file. Note that this is a
/// recursive definition! However, the size_of of `HirFileId` is finite
/// (because everything bottoms out at the real `FileId`) and small
/// (`MacroCallId` uses the location interner).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HirFileId(HirFileIdRepr);

impl HirFileId {
    /// For macro-expansion files, returns the file original source file the
    /// expansionoriginated from.
    pub fn original_file(self, db: &impl HirDatabase) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(macro_call_id) => {
                let loc = macro_call_id.loc(db);
                loc.source_item_id.file_id.original_file(db)
            }
        }
    }

    pub(crate) fn as_original_file(self) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(_r) => panic!("macro generated file: {:?}", self),
        }
    }

    pub(crate) fn as_macro_call_id(self) -> Option<MacroCallId> {
        match self.0 {
            HirFileIdRepr::Macro(it) => Some(it),
            _ => None,
        }
    }

    pub(crate) fn hir_source_file(
        db: &impl HirDatabase,
        file_id: HirFileId,
    ) -> TreeArc<SourceFile> {
        match file_id.0 {
            HirFileIdRepr::File(file_id) => db.source_file(file_id),
            HirFileIdRepr::Macro(m) => {
                if let Some(exp) = db.expand_macro_invocation(m) {
                    return exp.file();
                }
                // returning an empty string looks fishy...
                SourceFile::parse("")
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum HirFileIdRepr {
    File(FileId),
    Macro(MacroCallId),
}

impl From<FileId> for HirFileId {
    fn from(file_id: FileId) -> HirFileId {
        HirFileId(HirFileIdRepr::File(file_id))
    }
}

impl From<MacroCallId> for HirFileId {
    fn from(macro_call_id: MacroCallId) -> HirFileId {
        HirFileId(HirFileIdRepr::Macro(macro_call_id))
    }
}

/// `MacroCallId` identifies a particular macro invocation, like
/// `println!("Hello, {}", world)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroCallId(RawId);
impl_arena_id!(MacroCallId);

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MacroCallLoc {
    pub(crate) module: Module,
    pub(crate) source_item_id: SourceItemId,
}

impl MacroCallId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> MacroCallLoc {
        db.as_ref().macros.id2loc(self)
    }
}

impl MacroCallLoc {
    #[allow(unused)]
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> MacroCallId {
        db.as_ref().macros.loc2id(&self)
    }
}

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct ItemLoc<N: AstNode> {
    pub(crate) module: Module,
    raw: SourceItemId,
    _ty: PhantomData<N>,
}

impl<N: AstNode> ItemLoc<N> {
    pub(crate) fn from_ast(
        db: &impl HirDatabase,
        module: Module,
        file_id: HirFileId,
        ast: &N,
    ) -> ItemLoc<N> {
        let items = db.file_items(file_id);
        let raw = SourceItemId {
            file_id,
            item_id: Some(items.id_of(file_id, ast.syntax())),
        };
        ItemLoc {
            module,
            raw,
            _ty: PhantomData,
        }
    }

    pub(crate) fn source(&self, db: &impl HirDatabase) -> (HirFileId, TreeArc<N>) {
        let syntax = db.file_item(self.raw);
        let ast = N::cast(&syntax)
            .unwrap_or_else(|| panic!("invalid ItemLoc: {:?}", self.raw))
            .to_owned();
        (self.raw.file_id, ast)
    }
}

impl<N: AstNode> Clone for ItemLoc<N> {
    fn clone(&self) -> ItemLoc<N> {
        ItemLoc {
            module: self.module,
            raw: self.raw,
            _ty: PhantomData,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct FunctionId(RawId);
impl_arena_id!(FunctionId);

pub(crate) type FunctionLoc = ItemLoc<ast::FnDef>;

impl FunctionId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> FunctionLoc {
        db.as_ref().fns.id2loc(self)
    }
}

impl FunctionLoc {
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> FunctionId {
        db.as_ref().fns.loc2id(&self)
    }
}

/// Def's are a core concept of hir. A `Def` is an Item (function, module, etc)
/// in a specific module.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DefId(RawId);
impl_arena_id!(DefId);

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct DefLoc {
    pub(crate) kind: DefKind,
    pub(crate) module: Module,
    pub(crate) source_item_id: SourceItemId,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum DefKind {
    Struct,
    Enum,
    EnumVariant,
    Const,
    Static,
    Trait,
    Type,
    Item,

    /// The constructor of a struct. E.g. if we have `struct Foo(usize)`, the
    /// name `Foo` needs to resolve to different types depending on whether we
    /// are in the types or values namespace: As a type, `Foo` of course refers
    /// to the struct `Foo`; as a value, `Foo` is a callable type with signature
    /// `(usize) -> Foo`. The cleanest approach to handle this seems to be to
    /// have different defs in the two namespaces.
    ///
    /// rustc does the same; note that it even creates a struct constructor if
    /// the struct isn't a tuple struct (see `CtorKind::Fictive` in rustc).
    StructCtor,
}

impl DefId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> DefLoc {
        db.as_ref().defs.id2loc(self)
    }

    pub fn resolve(self, db: &impl HirDatabase) -> Def {
        let loc = self.loc(db);
        match loc.kind {
            DefKind::Struct => {
                let struct_def = Struct::new(self);
                Def::Struct(struct_def)
            }
            DefKind::Enum => Def::Enum(Enum::new(self)),
            DefKind::EnumVariant => Def::EnumVariant(EnumVariant::new(self)),
            DefKind::Const => {
                let def = Const::new(self);
                Def::Const(def)
            }
            DefKind::Static => {
                let def = Static::new(self);
                Def::Static(def)
            }
            DefKind::Trait => {
                let def = Trait::new(self);
                Def::Trait(def)
            }
            DefKind::Type => {
                let def = Type::new(self);
                Def::Type(def)
            }

            DefKind::StructCtor => Def::Item,
            DefKind::Item => Def::Item,
        }
    }

    pub(crate) fn source(self, db: &impl HirDatabase) -> (HirFileId, TreeArc<SyntaxNode>) {
        let loc = self.loc(db);
        let syntax = db.file_item(loc.source_item_id);
        (loc.source_item_id.file_id, syntax)
    }

    /// For a module, returns that module; for any other def, returns the containing module.
    pub fn module(self, db: &impl HirDatabase) -> Module {
        self.loc(db).module
    }

    /// Returns the containing crate.
    pub fn krate(&self, db: &impl HirDatabase) -> Option<Crate> {
        self.module(db).krate(db)
    }
}

impl DefLoc {
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> DefId {
        db.as_ref().defs.loc2id(&self)
    }
}

/// Identifier of item within a specific file. This is stable over reparses, so
/// it's OK to use it as a salsa key/value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SourceFileItemId(RawId);
impl_arena_id!(SourceFileItemId);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceItemId {
    pub(crate) file_id: HirFileId,
    /// None for the whole file.
    pub(crate) item_id: Option<SourceFileItemId>,
}

/// Maps items' `SyntaxNode`s to `SourceFileItemId`s and back.
#[derive(Debug, PartialEq, Eq)]
pub struct SourceFileItems {
    file_id: HirFileId,
    arena: Arena<SourceFileItemId, TreeArc<SyntaxNode>>,
}

impl SourceFileItems {
    pub(crate) fn new(file_id: HirFileId, source_file: &SourceFile) -> SourceFileItems {
        let mut res = SourceFileItems {
            file_id,
            arena: Arena::default(),
        };
        res.init(source_file);
        res
    }

    fn init(&mut self, source_file: &SourceFile) {
        // By walking the tree in bread-first order we make sure that parents
        // get lower ids then children. That is, addding a new child does not
        // change parent's id. This means that, say, adding a new function to a
        // trait does not chage ids of top-level items, which helps caching.
        bfs(source_file.syntax(), |it| {
            if let Some(enum_variant) = ast::EnumVariant::cast(it) {
                self.alloc(enum_variant.syntax().to_owned());
            } else if let Some(module_item) = ast::ModuleItem::cast(it) {
                self.alloc(module_item.syntax().to_owned());
            } else if let Some(macro_call) = ast::MacroCall::cast(it) {
                self.alloc(macro_call.syntax().to_owned());
            }
        })
    }

    fn alloc(&mut self, item: TreeArc<SyntaxNode>) -> SourceFileItemId {
        self.arena.alloc(item)
    }
    pub(crate) fn id_of(&self, file_id: HirFileId, item: &SyntaxNode) -> SourceFileItemId {
        assert_eq!(
            self.file_id, file_id,
            "SourceFileItems: wrong file, expected {:?}, got {:?}",
            self.file_id, file_id
        );
        self.id_of_unchecked(item)
    }
    pub(crate) fn id_of_unchecked(&self, item: &SyntaxNode) -> SourceFileItemId {
        if let Some((id, _)) = self.arena.iter().find(|(_id, i)| *i == item) {
            return id;
        }
        // This should not happen. Let's try to give a sensible diagnostics.
        if let Some((id, i)) = self.arena.iter().find(|(_id, i)| i.range() == item.range()) {
            // FIXME(#288): whyyy are we getting here?
            log::error!(
                "unequal syntax nodes with the same range:\n{:?}\n{:?}",
                item,
                i
            );
            return id;
        }
        panic!(
            "Can't find {:?} in SourceFileItems:\n{:?}",
            item,
            self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
        );
    }
    pub fn id_of_source_file(&self) -> SourceFileItemId {
        let (id, _syntax) = self.arena.iter().next().unwrap();
        id
    }
}

impl std::ops::Index<SourceFileItemId> for SourceFileItems {
    type Output = SyntaxNode;
    fn index(&self, idx: SourceFileItemId) -> &SyntaxNode {
        &self.arena[idx]
    }
}

/// Walks the subtree in bfs order, calling `f` for each node.
fn bfs(node: &SyntaxNode, mut f: impl FnMut(&SyntaxNode)) {
    let mut curr_layer = vec![node];
    let mut next_layer = vec![];
    while !curr_layer.is_empty() {
        curr_layer.drain(..).for_each(|node| {
            next_layer.extend(node.children());
            f(node);
        });
        std::mem::swap(&mut curr_layer, &mut next_layer);
    }
}
<commit_msg>add StructId<commit_after>use std::marker::PhantomData;

use ra_db::{LocationIntener, FileId};
use ra_syntax::{TreeArc, SyntaxNode, SourceFile, AstNode, ast};
use ra_arena::{Arena, RawId, impl_arena_id};

use crate::{
    HirDatabase, Def, Struct, Enum, EnumVariant, Crate,
    Module, Trait, Type, Static, Const,
};

#[derive(Debug, Default)]
pub struct HirInterner {
    defs: LocationIntener<DefLoc, DefId>,
    macros: LocationIntener<MacroCallLoc, MacroCallId>,
    fns: LocationIntener<FunctionLoc, FunctionId>,
    structs: LocationIntener<StructLoc, StructId>,
}

impl HirInterner {
    pub fn len(&self) -> usize {
        self.defs.len() + self.macros.len()
    }
}

/// hir makes heavy use of ids: integer (u32) handlers to various things. You
/// can think of id as a pointer (but without a lifetime) or a file descriptor
/// (but for hir objects).
///
/// This module defines a bunch of ids we are using. The most important ones are
/// probably `HirFileId` and `DefId`.

/// Input to the analyzer is a set of files, where each file is indentified by
/// `FileId` and contains source code. However, another source of source code in
/// Rust are macros: each macro can be thought of as producing a "temporary
/// file". To assign an id to such a file, we use the id of the macro call that
/// produced the file. So, a `HirFileId` is either a `FileId` (source code
/// written by user), or a `MacroCallId` (source code produced by macro).
///
/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file containin
/// the call plus the offset of the macro call in the file. Note that this is a
/// recursive definition! However, the size_of of `HirFileId` is finite
/// (because everything bottoms out at the real `FileId`) and small
/// (`MacroCallId` uses the location interner).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HirFileId(HirFileIdRepr);

impl HirFileId {
    /// For macro-expansion files, returns the file original source file the
    /// expansionoriginated from.
    pub fn original_file(self, db: &impl HirDatabase) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(macro_call_id) => {
                let loc = macro_call_id.loc(db);
                loc.source_item_id.file_id.original_file(db)
            }
        }
    }

    pub(crate) fn as_original_file(self) -> FileId {
        match self.0 {
            HirFileIdRepr::File(file_id) => file_id,
            HirFileIdRepr::Macro(_r) => panic!("macro generated file: {:?}", self),
        }
    }

    pub(crate) fn as_macro_call_id(self) -> Option<MacroCallId> {
        match self.0 {
            HirFileIdRepr::Macro(it) => Some(it),
            _ => None,
        }
    }

    pub(crate) fn hir_source_file(
        db: &impl HirDatabase,
        file_id: HirFileId,
    ) -> TreeArc<SourceFile> {
        match file_id.0 {
            HirFileIdRepr::File(file_id) => db.source_file(file_id),
            HirFileIdRepr::Macro(m) => {
                if let Some(exp) = db.expand_macro_invocation(m) {
                    return exp.file();
                }
                // returning an empty string looks fishy...
                SourceFile::parse("")
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum HirFileIdRepr {
    File(FileId),
    Macro(MacroCallId),
}

impl From<FileId> for HirFileId {
    fn from(file_id: FileId) -> HirFileId {
        HirFileId(HirFileIdRepr::File(file_id))
    }
}

impl From<MacroCallId> for HirFileId {
    fn from(macro_call_id: MacroCallId) -> HirFileId {
        HirFileId(HirFileIdRepr::Macro(macro_call_id))
    }
}

/// `MacroCallId` identifies a particular macro invocation, like
/// `println!("Hello, {}", world)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroCallId(RawId);
impl_arena_id!(MacroCallId);

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MacroCallLoc {
    pub(crate) module: Module,
    pub(crate) source_item_id: SourceItemId,
}

impl MacroCallId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> MacroCallLoc {
        db.as_ref().macros.id2loc(self)
    }
}

impl MacroCallLoc {
    #[allow(unused)]
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> MacroCallId {
        db.as_ref().macros.loc2id(&self)
    }
}

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct ItemLoc<N: AstNode> {
    pub(crate) module: Module,
    raw: SourceItemId,
    _ty: PhantomData<N>,
}

impl<N: AstNode> ItemLoc<N> {
    pub(crate) fn from_ast(
        db: &impl HirDatabase,
        module: Module,
        file_id: HirFileId,
        ast: &N,
    ) -> ItemLoc<N> {
        let items = db.file_items(file_id);
        let raw = SourceItemId {
            file_id,
            item_id: Some(items.id_of(file_id, ast.syntax())),
        };
        ItemLoc {
            module,
            raw,
            _ty: PhantomData,
        }
    }

    pub(crate) fn source(&self, db: &impl HirDatabase) -> (HirFileId, TreeArc<N>) {
        let syntax = db.file_item(self.raw);
        let ast = N::cast(&syntax)
            .unwrap_or_else(|| panic!("invalid ItemLoc: {:?}", self.raw))
            .to_owned();
        (self.raw.file_id, ast)
    }
}

impl<N: AstNode> Clone for ItemLoc<N> {
    fn clone(&self) -> ItemLoc<N> {
        ItemLoc {
            module: self.module,
            raw: self.raw,
            _ty: PhantomData,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct FunctionId(RawId);
impl_arena_id!(FunctionId);

pub(crate) type FunctionLoc = ItemLoc<ast::FnDef>;

impl FunctionId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> FunctionLoc {
        db.as_ref().fns.id2loc(self)
    }
}

impl FunctionLoc {
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> FunctionId {
        db.as_ref().fns.loc2id(&self)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct StructId(RawId);
impl_arena_id!(StructId);

pub(crate) type StructLoc = ItemLoc<ast::StructDef>;

impl StructId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> StructLoc {
        db.as_ref().structs.id2loc(self)
    }
}

impl StructLoc {
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> StructId {
        db.as_ref().structs.loc2id(&self)
    }
}

/// Def's are a core concept of hir. A `Def` is an Item (function, module, etc)
/// in a specific module.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DefId(RawId);
impl_arena_id!(DefId);

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct DefLoc {
    pub(crate) kind: DefKind,
    pub(crate) module: Module,
    pub(crate) source_item_id: SourceItemId,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum DefKind {
    Struct,
    Enum,
    EnumVariant,
    Const,
    Static,
    Trait,
    Type,
    Item,

    /// The constructor of a struct. E.g. if we have `struct Foo(usize)`, the
    /// name `Foo` needs to resolve to different types depending on whether we
    /// are in the types or values namespace: As a type, `Foo` of course refers
    /// to the struct `Foo`; as a value, `Foo` is a callable type with signature
    /// `(usize) -> Foo`. The cleanest approach to handle this seems to be to
    /// have different defs in the two namespaces.
    ///
    /// rustc does the same; note that it even creates a struct constructor if
    /// the struct isn't a tuple struct (see `CtorKind::Fictive` in rustc).
    StructCtor,
}

impl DefId {
    pub(crate) fn loc(self, db: &impl AsRef<HirInterner>) -> DefLoc {
        db.as_ref().defs.id2loc(self)
    }

    pub fn resolve(self, db: &impl HirDatabase) -> Def {
        let loc = self.loc(db);
        match loc.kind {
            DefKind::Struct => {
                let struct_def = Struct::new(self);
                Def::Struct(struct_def)
            }
            DefKind::Enum => Def::Enum(Enum::new(self)),
            DefKind::EnumVariant => Def::EnumVariant(EnumVariant::new(self)),
            DefKind::Const => {
                let def = Const::new(self);
                Def::Const(def)
            }
            DefKind::Static => {
                let def = Static::new(self);
                Def::Static(def)
            }
            DefKind::Trait => {
                let def = Trait::new(self);
                Def::Trait(def)
            }
            DefKind::Type => {
                let def = Type::new(self);
                Def::Type(def)
            }

            DefKind::StructCtor => Def::Item,
            DefKind::Item => Def::Item,
        }
    }

    pub(crate) fn source(self, db: &impl HirDatabase) -> (HirFileId, TreeArc<SyntaxNode>) {
        let loc = self.loc(db);
        let syntax = db.file_item(loc.source_item_id);
        (loc.source_item_id.file_id, syntax)
    }

    /// For a module, returns that module; for any other def, returns the containing module.
    pub fn module(self, db: &impl HirDatabase) -> Module {
        self.loc(db).module
    }

    /// Returns the containing crate.
    pub fn krate(&self, db: &impl HirDatabase) -> Option<Crate> {
        self.module(db).krate(db)
    }
}

impl DefLoc {
    pub(crate) fn id(&self, db: &impl AsRef<HirInterner>) -> DefId {
        db.as_ref().defs.loc2id(&self)
    }
}

/// Identifier of item within a specific file. This is stable over reparses, so
/// it's OK to use it as a salsa key/value.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SourceFileItemId(RawId);
impl_arena_id!(SourceFileItemId);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceItemId {
    pub(crate) file_id: HirFileId,
    /// None for the whole file.
    pub(crate) item_id: Option<SourceFileItemId>,
}

/// Maps items' `SyntaxNode`s to `SourceFileItemId`s and back.
#[derive(Debug, PartialEq, Eq)]
pub struct SourceFileItems {
    file_id: HirFileId,
    arena: Arena<SourceFileItemId, TreeArc<SyntaxNode>>,
}

impl SourceFileItems {
    pub(crate) fn new(file_id: HirFileId, source_file: &SourceFile) -> SourceFileItems {
        let mut res = SourceFileItems {
            file_id,
            arena: Arena::default(),
        };
        res.init(source_file);
        res
    }

    fn init(&mut self, source_file: &SourceFile) {
        // By walking the tree in bread-first order we make sure that parents
        // get lower ids then children. That is, addding a new child does not
        // change parent's id. This means that, say, adding a new function to a
        // trait does not chage ids of top-level items, which helps caching.
        bfs(source_file.syntax(), |it| {
            if let Some(enum_variant) = ast::EnumVariant::cast(it) {
                self.alloc(enum_variant.syntax().to_owned());
            } else if let Some(module_item) = ast::ModuleItem::cast(it) {
                self.alloc(module_item.syntax().to_owned());
            } else if let Some(macro_call) = ast::MacroCall::cast(it) {
                self.alloc(macro_call.syntax().to_owned());
            }
        })
    }

    fn alloc(&mut self, item: TreeArc<SyntaxNode>) -> SourceFileItemId {
        self.arena.alloc(item)
    }
    pub(crate) fn id_of(&self, file_id: HirFileId, item: &SyntaxNode) -> SourceFileItemId {
        assert_eq!(
            self.file_id, file_id,
            "SourceFileItems: wrong file, expected {:?}, got {:?}",
            self.file_id, file_id
        );
        self.id_of_unchecked(item)
    }
    pub(crate) fn id_of_unchecked(&self, item: &SyntaxNode) -> SourceFileItemId {
        if let Some((id, _)) = self.arena.iter().find(|(_id, i)| *i == item) {
            return id;
        }
        // This should not happen. Let's try to give a sensible diagnostics.
        if let Some((id, i)) = self.arena.iter().find(|(_id, i)| i.range() == item.range()) {
            // FIXME(#288): whyyy are we getting here?
            log::error!(
                "unequal syntax nodes with the same range:\n{:?}\n{:?}",
                item,
                i
            );
            return id;
        }
        panic!(
            "Can't find {:?} in SourceFileItems:\n{:?}",
            item,
            self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
        );
    }
    pub fn id_of_source_file(&self) -> SourceFileItemId {
        let (id, _syntax) = self.arena.iter().next().unwrap();
        id
    }
}

impl std::ops::Index<SourceFileItemId> for SourceFileItems {
    type Output = SyntaxNode;
    fn index(&self, idx: SourceFileItemId) -> &SyntaxNode {
        &self.arena[idx]
    }
}

/// Walks the subtree in bfs order, calling `f` for each node.
fn bfs(node: &SyntaxNode, mut f: impl FnMut(&SyntaxNode)) {
    let mut curr_layer = vec![node];
    let mut next_layer = vec![];
    while !curr_layer.is_empty() {
        curr_layer.drain(..).for_each(|node| {
            next_layer.extend(node.children());
            f(node);
        });
        std::mem::swap(&mut curr_layer, &mut next_layer);
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::path::Path;
use std::path::PathBuf;
use std::fs::File;
use std::io::Write;

extern crate tools;
extern crate ipdl_parser;
extern crate getopts;

use getopts::Options;

use tools::file_format::analysis::{read_analysis, read_target, WithLocation, AnalysisTarget, AnalysisKind};

use ipdl_parser::parser;
use ipdl_parser::ast;

type TargetAnalysis = Vec<WithLocation<Vec<AnalysisTarget>>>;

fn get_options_parser() -> Options {
    let mut opts = Options::new();
    opts.optmulti("I", "include",
                  "Additional directory to search for included protocol specifications",
                  "DIR");
    opts.reqopt("d", "outheaders-dir",
                "Directory into which C++ headers analysis data is location.",
                "HDR_DIR");
    opts.reqopt("b", "base-input-prefix",
                "Base directory where IPDL input files are found.",
                "BASE_DIR");
    opts.reqopt("a", "analysis-prefix",
                "Base directory where analysis output files are found.",
                "ANALYSIS_DIR");
    opts
}

fn header_file_name(outheaders_dir: &str, ns: &ast::Namespace, parent_or_child: &str) -> String {
    format!("{}/{}/{}{}.h",
            outheaders_dir,
            ns.namespaces.clone().join("/"),
            ns.name.id,
            parent_or_child)
}

fn mangle_simple(s: &str) -> String {
    format!("{}{}", s.len(), s)
}

fn mangle_nested_name(ns: &[String], protocol: &str, name: &str) -> String {
    format!("_ZN{}{}{}E",
            ns.iter().map(|id| mangle_simple(&id)).collect::<Vec<_>>().join(""),
            mangle_simple(protocol),
            mangle_simple(name))
}

fn find_analysis<'a>(analysis: &'a TargetAnalysis, mangled: &str) -> &'a AnalysisTarget
{
    for datum in analysis {
        for piece in &datum.data {
            if piece.kind == AnalysisKind::Decl && piece.sym.contains(mangled) {
                return &piece;
            }
        }
    }

    panic!("No analysis target found for {}", mangled);
}

fn output_data(outputf: &mut File, locstr: &str, datum: &AnalysisTarget) {
    write!(outputf, r#"{{"loc": "{}", "target": 1, "kind": "idl", "pretty": "{}", "sym": "{}"}}"#,
           locstr, datum.pretty, datum.sym).unwrap();
    write!(outputf, "\n").unwrap();
    write!(outputf, r#"{{"loc": "{}", "source": 1, "pretty": "{}", "sym": "{}"}}"#,
           locstr, datum.pretty, datum.sym).unwrap();
    write!(outputf, "\n").unwrap();
}

fn output_send_recv(outputf: &mut File,
                    locstr: &str,
                    protocol: &ast::Namespace,
                    message: &ast::MessageDecl,
                    is_ctor: bool,
                    send_side: &str, send_analysis: &TargetAnalysis,
                    recv_side: &str, recv_analysis: &TargetAnalysis)
{
    let send_prefix = if message.send_semantics == ast::SendSemantics::Intr { "Call" } else { "Send" };
    let recv_prefix = if message.send_semantics == ast::SendSemantics::Intr { "Answer" } else { "Recv" };

    let ctor_suffix = if is_ctor { "Constructor" } else { "" };

    let mangled = mangle_nested_name(&protocol.namespaces,
                                     &format!("{}{}", protocol.name.id, send_side),
                                     &format!("{}{}{}", send_prefix, message.name.id, ctor_suffix));
    let send_datum = find_analysis(send_analysis, &mangled);
    output_data(outputf, &locstr, &send_datum);

    let mangled = mangle_nested_name(&protocol.namespaces,
                                     &format!("{}{}", protocol.name.id, recv_side),
                                     &format!("{}{}{}", recv_prefix, message.name.id, ctor_suffix));
    let recv_datum = find_analysis(recv_analysis, &mangled);
    output_data(outputf, &locstr, &recv_datum);
}

fn main() {
    let args : Vec<String> = env::args().collect();

    let opts = get_options_parser();

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m },
        Err(f) => { panic!(f.to_string()) },
    };

    let mut include_dirs = Vec::new();
    for i in matches.opt_strs("I") {
        include_dirs.push(PathBuf::from(i))
    }

    let outheaders_dir = matches.opt_str("d").unwrap();
    let base_dir = matches.opt_str("b").unwrap();
    let analysis_dir = matches.opt_str("a").unwrap();
    let base_path = Path::new(&base_dir);
    let analysis_path = Path::new(&analysis_dir);

    let mut file_names = Vec::new();
    for f in matches.free {
        file_names.push(PathBuf::from(f));
    }

    let maybe_tus = parser::parse(&include_dirs, file_names);

    if maybe_tus.is_none() {
        println!("Specification could not be parsed.");
        return;
    }

    let tus = maybe_tus.unwrap();

    for (_, tu) in tus {
        println!("Analyzing {:?}", tu.file_name);

        let path = tu.file_name.as_path();
        let relative = path.strip_prefix(base_path).unwrap();
        let absolute = analysis_path.join(relative);
        let mut outputf = File::create(absolute).unwrap();

        if let Some((ns, protocol)) = tu.protocol {
            let parent_fname = header_file_name(&outheaders_dir, &ns, "Parent");
            let parent_analysis = read_analysis(&parent_fname, &read_target);
            let child_fname = header_file_name(&outheaders_dir, &ns, "Child");
            let child_analysis = read_analysis(&child_fname, &read_target);

            let is_toplevel = protocol.managers.len() == 0;

            for message in protocol.messages {
                let loc = &message.name.loc;
                let locstr = format!("{}:{}-{}", loc.lineno, loc.colno, loc.colno + message.name.id.len());

                if is_toplevel && message.name.id == "__delete__" {
                    continue;
                }

                let is_ctor = protocol.manages.iter().any(|e| e.id == message.name.id);

                if message.direction == ast::Direction::ToChild || message.direction == ast::Direction::ToParentOrChild {
                    output_send_recv(&mut outputf, &locstr, &ns, &message, is_ctor,
                                     "Parent", &parent_analysis, "Child", &child_analysis);
                }

                if message.direction == ast::Direction::ToParent || message.direction == ast::Direction::ToParentOrChild {
                    output_send_recv(&mut outputf, &locstr, &ns, &message, is_ctor,
                                     "Child", &child_analysis, "Parent", &parent_analysis);
                }
            }
        }
    }
}
<commit_msg>Make failures in find_analysis non-fatal.<commit_after>use std::env;
use std::path::Path;
use std::path::PathBuf;
use std::fs::File;
use std::io::Write;

extern crate tools;
extern crate ipdl_parser;
extern crate getopts;

use getopts::Options;

use tools::file_format::analysis::{read_analysis, read_target, WithLocation, AnalysisTarget, AnalysisKind};

use ipdl_parser::parser;
use ipdl_parser::ast;

type TargetAnalysis = Vec<WithLocation<Vec<AnalysisTarget>>>;

fn get_options_parser() -> Options {
    let mut opts = Options::new();
    opts.optmulti("I", "include",
                  "Additional directory to search for included protocol specifications",
                  "DIR");
    opts.reqopt("d", "outheaders-dir",
                "Directory into which C++ headers analysis data is location.",
                "HDR_DIR");
    opts.reqopt("b", "base-input-prefix",
                "Base directory where IPDL input files are found.",
                "BASE_DIR");
    opts.reqopt("a", "analysis-prefix",
                "Base directory where analysis output files are found.",
                "ANALYSIS_DIR");
    opts
}

fn header_file_name(outheaders_dir: &str, ns: &ast::Namespace, parent_or_child: &str) -> String {
    format!("{}/{}/{}{}.h",
            outheaders_dir,
            ns.namespaces.clone().join("/"),
            ns.name.id,
            parent_or_child)
}

fn mangle_simple(s: &str) -> String {
    format!("{}{}", s.len(), s)
}

fn mangle_nested_name(ns: &[String], protocol: &str, name: &str) -> String {
    format!("_ZN{}{}{}E",
            ns.iter().map(|id| mangle_simple(&id)).collect::<Vec<_>>().join(""),
            mangle_simple(protocol),
            mangle_simple(name))
}

fn find_analysis<'a>(analysis: &'a TargetAnalysis, mangled: &str) -> Option<&'a AnalysisTarget>
{
    for datum in analysis {
        for piece in &datum.data {
            if piece.kind == AnalysisKind::Decl && piece.sym.contains(mangled) {
                return Some(&piece);
            }
        }
    }

    println!("No analysis target found for {}", mangled);
    return None
}

fn output_data(outputf: &mut File, locstr: &str, datum: &AnalysisTarget) {
    write!(outputf, r#"{{"loc": "{}", "target": 1, "kind": "idl", "pretty": "{}", "sym": "{}"}}"#,
           locstr, datum.pretty, datum.sym).unwrap();
    write!(outputf, "\n").unwrap();
    write!(outputf, r#"{{"loc": "{}", "source": 1, "pretty": "{}", "sym": "{}"}}"#,
           locstr, datum.pretty, datum.sym).unwrap();
    write!(outputf, "\n").unwrap();
}

fn output_send_recv(outputf: &mut File,
                    locstr: &str,
                    protocol: &ast::Namespace,
                    message: &ast::MessageDecl,
                    is_ctor: bool,
                    send_side: &str, send_analysis: &TargetAnalysis,
                    recv_side: &str, recv_analysis: &TargetAnalysis)
{
    let send_prefix = if message.send_semantics == ast::SendSemantics::Intr { "Call" } else { "Send" };
    let recv_prefix = if message.send_semantics == ast::SendSemantics::Intr { "Answer" } else { "Recv" };

    let ctor_suffix = if is_ctor { "Constructor" } else { "" };

    let mangled = mangle_nested_name(&protocol.namespaces,
                                     &format!("{}{}", protocol.name.id, send_side),
                                     &format!("{}{}{}", send_prefix, message.name.id, ctor_suffix));
    if let Some(send_datum) = find_analysis(send_analysis, &mangled) {
        output_data(outputf, &locstr, &send_datum);
    }

    let mangled = mangle_nested_name(&protocol.namespaces,
                                     &format!("{}{}", protocol.name.id, recv_side),
                                     &format!("{}{}{}", recv_prefix, message.name.id, ctor_suffix));
    if let Some(recv_datum) = find_analysis(recv_analysis, &mangled) {
        output_data(outputf, &locstr, &recv_datum);
    }
}

fn main() {
    let args : Vec<String> = env::args().collect();

    let opts = get_options_parser();

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m },
        Err(f) => { panic!(f.to_string()) },
    };

    let mut include_dirs = Vec::new();
    for i in matches.opt_strs("I") {
        include_dirs.push(PathBuf::from(i))
    }

    let outheaders_dir = matches.opt_str("d").unwrap();
    let base_dir = matches.opt_str("b").unwrap();
    let analysis_dir = matches.opt_str("a").unwrap();
    let base_path = Path::new(&base_dir);
    let analysis_path = Path::new(&analysis_dir);

    let mut file_names = Vec::new();
    for f in matches.free {
        file_names.push(PathBuf::from(f));
    }

    let maybe_tus = parser::parse(&include_dirs, file_names);

    if maybe_tus.is_none() {
        println!("Specification could not be parsed.");
        return;
    }

    let tus = maybe_tus.unwrap();

    for (_, tu) in tus {
        println!("Analyzing {:?}", tu.file_name);

        let path = tu.file_name.as_path();
        let relative = path.strip_prefix(base_path).unwrap();
        let absolute = analysis_path.join(relative);
        let mut outputf = File::create(absolute).unwrap();

        if let Some((ns, protocol)) = tu.protocol {
            let parent_fname = header_file_name(&outheaders_dir, &ns, "Parent");
            let parent_analysis = read_analysis(&parent_fname, &read_target);
            let child_fname = header_file_name(&outheaders_dir, &ns, "Child");
            let child_analysis = read_analysis(&child_fname, &read_target);

            let is_toplevel = protocol.managers.len() == 0;

            for message in protocol.messages {
                let loc = &message.name.loc;
                let locstr = format!("{}:{}-{}", loc.lineno, loc.colno, loc.colno + message.name.id.len());

                if is_toplevel && message.name.id == "__delete__" {
                    continue;
                }

                let is_ctor = protocol.manages.iter().any(|e| e.id == message.name.id);

                if message.direction == ast::Direction::ToChild || message.direction == ast::Direction::ToParentOrChild {
                    output_send_recv(&mut outputf, &locstr, &ns, &message, is_ctor,
                                     "Parent", &parent_analysis, "Child", &child_analysis);
                }

                if message.direction == ast::Direction::ToParent || message.direction == ast::Direction::ToParentOrChild {
                    output_send_recv(&mut outputf, &locstr, &ns, &message, is_ctor,
                                     "Child", &child_analysis, "Parent", &parent_analysis);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use cairo;
use cairo_sys;
use glib::translate::*;
use libc;

use regex::{Captures, Regex};
use std::borrow::Cow;
use std::cell::RefCell;

use attributes::Attribute;
use drawing_ctx::DrawingCtx;
use error::RenderingError;
use node::*;
use property_bag::PropertyBag;

pub struct NodeLink {
    link: RefCell<Option<String>>,
}

impl NodeLink {
    pub fn new() -> NodeLink {
        NodeLink {
            link: RefCell::new(None),
        }
    }
}

impl NodeTrait for NodeLink {
    fn set_atts(&self, _: &RsvgNode, pbag: &PropertyBag<'_>) -> NodeResult {
        for (attr, value) in pbag.iter() {
            match attr {
                Attribute::XlinkHref => *self.link.borrow_mut() = Some(value.to_owned()),

                _ => (),
            }
        }

        Ok(())
    }

    fn draw(
        &self,
        node: &RsvgNode,
        cascaded: &CascadedValues<'_>,
        draw_ctx: &mut DrawingCtx,
        clipping: bool,
    ) -> Result<(), RenderingError> {
        let link = self.link.borrow();

        let cascaded = CascadedValues::new(cascaded, node);
        let values = cascaded.get();

        draw_ctx.with_discrete_layer(node, values, clipping, &mut |dc| {
            if link.is_some() && link.as_ref().unwrap() != "" {
                const CAIRO_TAG_LINK: &str = "Link";

                let attributes = link.as_ref().map(|i| format!("uri='{}'", escape_value(i)));

                let cr = dc.get_cairo_context();

                cr.tag_begin(CAIRO_TAG_LINK, attributes.as_ref().map(|i| i.as_str()));

                let res = node.draw_children(&cascaded, dc, clipping);

                cr.tag_end(CAIRO_TAG_LINK);

                res
            } else {
                node.draw_children(&cascaded, dc, clipping)
            }
        })
    }
}

/// escape quotes and backslashes with backslash
fn escape_value(value: &str) -> Cow<'_, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"['\\]").unwrap();
    }

    REGEX.replace_all(value, |caps: &Captures<'_>| {
        match caps.get(0).unwrap().as_str() {
            "'" => "\\'".to_owned(),
            "\\" => "\\\\".to_owned(),
            _ => unreachable!(),
        }
    })
}

extern "C" {
    fn cairo_tag_begin(
        cr: *mut cairo_sys::cairo_t,
        tag_name: *const libc::c_char,
        attibutes: *const libc::c_char,
    );
    fn cairo_tag_end(cr: *mut cairo_sys::cairo_t, tag_name: *const libc::c_char);
}

/// Bindings that aren't supported by `cairo-rs` for now
trait CairoTagging {
    fn tag_begin(&self, tag_name: &str, attributes: Option<&str>);
    fn tag_end(&self, tag_name: &str);
}

impl CairoTagging for cairo::Context {
    fn tag_begin(&self, tag_name: &str, attributes: Option<&str>) {
        unsafe {
            cairo_tag_begin(
                self.to_glib_none().0,
                tag_name.to_glib_none().0,
                attributes.to_glib_none().0,
            );
        }
    }

    fn tag_end(&self, tag_name: &str) {
        unsafe {
            cairo_tag_end(self.to_glib_none().0, tag_name.to_glib_none().0);
        }
    }
}
<commit_msg>use tag_begin and tag_end of cairo v1.16, fix #435<commit_after>use cairo;
use cairo_sys;
use glib::translate::*;
use libc;

use regex::{Captures, Regex};
use std::borrow::Cow;
use std::cell::RefCell;

use attributes::Attribute;
use drawing_ctx::DrawingCtx;
use error::RenderingError;
use node::*;
use property_bag::PropertyBag;

pub struct NodeLink {
    link: RefCell<Option<String>>,
}

impl NodeLink {
    pub fn new() -> NodeLink {
        NodeLink {
            link: RefCell::new(None),
        }
    }
}

impl NodeTrait for NodeLink {
    fn set_atts(&self, _: &RsvgNode, pbag: &PropertyBag<'_>) -> NodeResult {
        for (attr, value) in pbag.iter() {
            match attr {
                Attribute::XlinkHref => *self.link.borrow_mut() = Some(value.to_owned()),

                _ => (),
            }
        }

        Ok(())
    }

    fn draw(
        &self,
        node: &RsvgNode,
        cascaded: &CascadedValues<'_>,
        draw_ctx: &mut DrawingCtx,
        clipping: bool,
    ) -> Result<(), RenderingError> {
        let link = self.link.borrow();

        let cascaded = CascadedValues::new(cascaded, node);
        let values = cascaded.get();

        draw_ctx.with_discrete_layer(node, values, clipping, &mut |dc| match link.as_ref() {
            Some(l) if !l.is_empty() => {
                const CAIRO_TAG_LINK: &str = "Link";

                let attributes = format!("uri='{}'", escape_value(l));

                let cr = dc.get_cairo_context();

                cr.tag_begin(CAIRO_TAG_LINK, &attributes);

                let res = node.draw_children(&cascaded, dc, clipping);

                cr.tag_end(CAIRO_TAG_LINK);

                res
            }
            _ => node.draw_children(&cascaded, dc, clipping),
        })
    }
}

/// escape quotes and backslashes with backslash
fn escape_value(value: &str) -> Cow<'_, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"['\\]").unwrap();
    }

    REGEX.replace_all(value, |caps: &Captures<'_>| {
        match caps.get(0).unwrap().as_str() {
            "'" => "\\'".to_owned(),
            "\\" => "\\\\".to_owned(),
            _ => unreachable!(),
        }
    })
}

#[cfg(not(feature = "v1_16"))]
extern "C" {
    fn cairo_tag_begin(
        cr: *mut cairo_sys::cairo_t,
        tag_name: *const libc::c_char,
        attibutes: *const libc::c_char,
    );
    fn cairo_tag_end(cr: *mut cairo_sys::cairo_t, tag_name: *const libc::c_char);
}

/// Bindings that aren't supported by `cairo-rs` for now
#[cfg(not(feature = "v1_16"))]
trait CairoTagging {
    fn tag_begin(&self, tag_name: &str, attributes: &str);
    fn tag_end(&self, tag_name: &str);
}

#[cfg(not(feature = "v1_16"))]
impl CairoTagging for cairo::Context {
    fn tag_begin(&self, tag_name: &str, attributes: &str) {
        unsafe {
            cairo_tag_begin(
                self.to_glib_none().0,
                tag_name.to_glib_none().0,
                attributes.to_glib_none().0,
            );
        }
    }

    fn tag_end(&self, tag_name: &str) {
        unsafe {
            cairo_tag_end(self.to_glib_none().0, tag_name.to_glib_none().0);
        }
    }
}
<|endoftext|>


<commit_before>use mio::*;
use mio::net::*;
use mio::unix::UnixReady;
use network::buffer_queue::BufferQueue;
use network::pool::Checkout;
use network::{ClientResult,Readiness};
use network::protocol::ProtocolResult;
use openssl::ssl::{self,HandshakeError,MidHandshakeSslStream,Ssl,SslStream};

pub enum TlsState {
  Initial,
  Handshake,
  Established,
  Error(HandshakeError<TcpStream>),
}

pub struct TlsHandshake {
  pub readiness:       Readiness,
  pub front:           Option<TcpStream>,
  pub ssl:             Option<Ssl>,
  pub stream:          Option<SslStream<TcpStream>>,
  mid:                 Option<MidHandshakeSslStream<TcpStream>>,
  state:               TlsState,
}

impl TlsHandshake {
  pub fn new(ssl:Ssl, sock: TcpStream) -> TlsHandshake {
    TlsHandshake {
      front:          Some(sock),
      ssl:            Some(ssl),
      mid:            None,
      stream:         None,
      state:          TlsState::Initial,
      readiness:      Readiness {
                        interest: UnixReady::from(Ready::readable()) | UnixReady::hup() | UnixReady::error(),
                        event:    UnixReady::from(Ready::empty()),
      },
    }
  }

  pub fn readable(&mut self) -> (ProtocolResult,ClientResult) {
    match self.state {
      TlsState::Error(_) => return (ProtocolResult::Continue, ClientResult::CloseClient),
      TlsState::Initial => {
        let ssl     = self.ssl.take().expect("TlsHandshake should have a Ssl backend");
        let sock    = self.front.take().expect("TlsHandshake should have a front socket");
        match ssl.accept(sock) {
          Ok(stream) => {
            self.stream = Some(stream);
            self.state = TlsState::Established;
            return (ProtocolResult::Upgrade, ClientResult::Continue);
          },
          Err(HandshakeError::SetupFailure(e)) => {
            error!("accept: handshake setup failed: {:?}", e);
            self.state = TlsState::Error(HandshakeError::SetupFailure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::Failure(e)) => {
            {
              if let Some(error_stack) = e.error().ssl_error() {
                let errors = error_stack.errors();
                if errors.len() == 2 && errors[0].code() == 0x1412E0E2 && errors[1].code() == 0x1408A0E3 {
                  incr!("openssl_sni_error");
                } else {
                  error!("accept: handshake failed: {:?}", e);
                }
              } else {
                error!("accept: handshake failed: {:?}", e);
              }
            }
            self.state = TlsState::Error(HandshakeError::Failure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::WouldBlock(mid)) => {
            self.state = TlsState::Handshake;
            self.mid = Some(mid);
            self.readiness.event.remove(Ready::readable());
            return (ProtocolResult::Continue, ClientResult::Continue);
          }
        }
      },
      TlsState::Handshake => {
        let mid = self.mid.take().expect("TlsHandshake should have a MidHandshakeSslStream backend");
        match mid.handshake() {
          Ok(stream) => {
            self.stream = Some(stream);
            self.state = TlsState::Established;
            return (ProtocolResult::Upgrade, ClientResult::Continue);
          },
          Err(HandshakeError::SetupFailure(e)) => {
            debug!("mid handshake setup failed: {:?}", e);
            self.state = TlsState::Error(HandshakeError::SetupFailure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::Failure(e)) => {
            debug!("mid handshake failed: {:?}", e);
            self.state = TlsState::Error(HandshakeError::Failure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::WouldBlock(new_mid)) => {
            self.state = TlsState::Handshake;
            self.mid = Some(new_mid);
            self.readiness.event.remove(Ready::readable());
            return (ProtocolResult::Continue, ClientResult::Continue);
          }
        }
      },
      TlsState::Established => {
        return (ProtocolResult::Upgrade, ClientResult::Continue);
      }
    }

  }

  pub fn socket(&self) -> Option<&TcpStream> {
    match self.state {
      TlsState::Initial => self.front.as_ref(),
      TlsState::Handshake => self.mid.as_ref().map(|mid| mid.get_ref()),
      TlsState::Established => self.stream.as_ref().map(|stream| stream.get_ref()),
      TlsState::Error(ref error) => {
        match error {
          &HandshakeError::SetupFailure(_) => {
            self.front.as_ref().or_else(|| self.mid.as_ref().map(|mid| mid.get_ref()))
          },
          &HandshakeError::Failure(ref mid) | &HandshakeError::WouldBlock(ref mid) => Some(mid.get_ref())
        }
      }
    }
  }
}

<commit_msg>log more openssl errors<commit_after>use mio::*;
use mio::net::*;
use mio::unix::UnixReady;
use network::buffer_queue::BufferQueue;
use network::pool::Checkout;
use network::{ClientResult,Readiness};
use network::protocol::ProtocolResult;
use openssl::ssl::{self,HandshakeError,MidHandshakeSslStream,Ssl,SslStream};

pub enum TlsState {
  Initial,
  Handshake,
  Established,
  Error(HandshakeError<TcpStream>),
}

pub struct TlsHandshake {
  pub readiness:       Readiness,
  pub front:           Option<TcpStream>,
  pub ssl:             Option<Ssl>,
  pub stream:          Option<SslStream<TcpStream>>,
  mid:                 Option<MidHandshakeSslStream<TcpStream>>,
  state:               TlsState,
}

impl TlsHandshake {
  pub fn new(ssl:Ssl, sock: TcpStream) -> TlsHandshake {
    TlsHandshake {
      front:          Some(sock),
      ssl:            Some(ssl),
      mid:            None,
      stream:         None,
      state:          TlsState::Initial,
      readiness:      Readiness {
                        interest: UnixReady::from(Ready::readable()) | UnixReady::hup() | UnixReady::error(),
                        event:    UnixReady::from(Ready::empty()),
      },
    }
  }

  pub fn readable(&mut self) -> (ProtocolResult,ClientResult) {
    match self.state {
      TlsState::Error(_) => return (ProtocolResult::Continue, ClientResult::CloseClient),
      TlsState::Initial => {
        let ssl     = self.ssl.take().expect("TlsHandshake should have a Ssl backend");
        let sock    = self.front.take().expect("TlsHandshake should have a front socket");
        match ssl.accept(sock) {
          Ok(stream) => {
            self.stream = Some(stream);
            self.state = TlsState::Established;
            return (ProtocolResult::Upgrade, ClientResult::Continue);
          },
          Err(HandshakeError::SetupFailure(e)) => {
            error!("accept: handshake setup failed: {:?}", e);
            self.state = TlsState::Error(HandshakeError::SetupFailure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::Failure(e)) => {
            {
              if let Some(error_stack) = e.error().ssl_error() {
                let errors = error_stack.errors();
                if errors.len() == 1 {
                  if errors[0].code() == 0x140A1175 {
                    incr!("openssl_inappropriate_fallback_error");
                  } else if errors[0].code() == 0x1408A10B {
                    incr!("openssl_wrong_version_number_error");
                  } else if errors[0].code() == 0x140760FC {
                    incr!("openssl_unknown_protocol_error");
                  } else if errors[0].code() == 0x1407609C {
                    //someone tried to connect in plain HTTP to a TLS server
                    incr!("openssl_http_request_error");
                  } else {
                    error!("accept: handshake failed: {:?}", e);
                  }
                } else if errors.len() == 2 && errors[0].code() == 0x1412E0E2 && errors[1].code() == 0x1408A0E3 {
                  incr!("openssl_sni_error");
                } else {
                  error!("accept: handshake failed: {:?}", e);
                }
              } else {
                error!("accept: handshake failed: {:?}", e);
              }
            }
            self.state = TlsState::Error(HandshakeError::Failure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::WouldBlock(mid)) => {
            self.state = TlsState::Handshake;
            self.mid = Some(mid);
            self.readiness.event.remove(Ready::readable());
            return (ProtocolResult::Continue, ClientResult::Continue);
          }
        }
      },
      TlsState::Handshake => {
        let mid = self.mid.take().expect("TlsHandshake should have a MidHandshakeSslStream backend");
        match mid.handshake() {
          Ok(stream) => {
            self.stream = Some(stream);
            self.state = TlsState::Established;
            return (ProtocolResult::Upgrade, ClientResult::Continue);
          },
          Err(HandshakeError::SetupFailure(e)) => {
            debug!("mid handshake setup failed: {:?}", e);
            self.state = TlsState::Error(HandshakeError::SetupFailure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::Failure(e)) => {
            debug!("mid handshake failed: {:?}", e);
            self.state = TlsState::Error(HandshakeError::Failure(e));
            return (ProtocolResult::Continue, ClientResult::CloseClient);
          },
          Err(HandshakeError::WouldBlock(new_mid)) => {
            self.state = TlsState::Handshake;
            self.mid = Some(new_mid);
            self.readiness.event.remove(Ready::readable());
            return (ProtocolResult::Continue, ClientResult::Continue);
          }
        }
      },
      TlsState::Established => {
        return (ProtocolResult::Upgrade, ClientResult::Continue);
      }
    }

  }

  pub fn socket(&self) -> Option<&TcpStream> {
    match self.state {
      TlsState::Initial => self.front.as_ref(),
      TlsState::Handshake => self.mid.as_ref().map(|mid| mid.get_ref()),
      TlsState::Established => self.stream.as_ref().map(|stream| stream.get_ref()),
      TlsState::Error(ref error) => {
        match error {
          &HandshakeError::SetupFailure(_) => {
            self.front.as_ref().or_else(|| self.mid.as_ref().map(|mid| mid.get_ref()))
          },
          &HandshakeError::Failure(ref mid) | &HandshakeError::WouldBlock(ref mid) => Some(mid.get_ref())
        }
      }
    }
  }
}

<|endoftext|>


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

use crate::*;

use super::{
    arr_to_u32, gc_blobs, pwrite_all, raw_segment_iter_from, u32_to_arr,
    u64_to_arr, BasedBuf, DiskPtr, LogIter, LogKind, LogOffset, Lsn,
    MessageKind,
};

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
#[derive(Debug, Default)]
#[cfg_attr(test, derive(Clone, PartialEq))]
pub struct Snapshot {
    /// The last read message lsn
    pub stable_lsn: Lsn,
    /// The last read message lid
    pub tip_lid: Option<LogOffset>,
    /// the mapping from pages to (lsn, lid)
    pub pt: Vec<PageState>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum PageState {
    Present { base: (Lsn, DiskPtr, u64), frags: Vec<(Lsn, DiskPtr, u64)> },
    Free(Lsn, DiskPtr),
    Uninitialized,
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr, u64)) {
        match *self {
            PageState::Present { base, ref mut frags } => {
                if frags.last().map_or(base.0, |f| f.0) < item.0 {
                    frags.push(item)
                } else {
                    debug!(
                        "skipping merging item {:?} into \
                        existing PageState::Present({:?})",
                        item, frags
                    );
                }
            }
            _ => panic!("pushed frags to {:?}", self),
        }
    }

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

impl Snapshot {
    fn apply(
        &mut self,
        log_kind: LogKind,
        pid: PageId,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        sz: u64,
    ) {
        // 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 _measure = Measure::new(&M.snapshot_apply);

        if self.pt.len() <= usize::try_from(pid).unwrap() {
            self.pt.resize(
                usize::try_from(pid + 1).unwrap(),
                PageState::Uninitialized,
            );
        }

        match log_kind {
            LogKind::Replace => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn,
                );

                self.pt[usize::try_from(pid).unwrap()] = PageState::Present {
                    base: (lsn, disk_ptr, sz),
                    frags: vec![],
                };
            }
            LogKind::Link => {
                // 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 @ PageState::Present { .. }) =
                    self.pt.get_mut(usize::try_from(pid).unwrap())
                {
                    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;
                    }

                    lids.push((lsn, disk_ptr, sz));
                } else {
                    trace!(
                        "skipping dangling append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );
                }
            }
            LogKind::Free => {
                trace!("free of pid {} at ptr {} lsn {}", pid, disk_ptr, lsn);
                self.pt[usize::try_from(pid).unwrap()] =
                    PageState::Free(lsn, disk_ptr);
            }
            LogKind::Corrupted | LogKind::Skip => panic!(
                "unexppected messagekind in snapshot application: {:?}",
                log_kind
            ),
        }
    }
}

fn advance_snapshot(
    mut iter: LogIter,
    mut snapshot: Snapshot,
    config: &RunningConfig,
) -> Result<Snapshot> {
    let _measure = Measure::new(&M.advance_snapshot);

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

    let old_stable_lsn = snapshot.stable_lsn;

    while let Some((log_kind, pid, lsn, ptr, sz)) = iter.next() {
        trace!(
            "in advance_snapshot looking at item with pid {} lsn {} ptr {}",
            pid,
            lsn,
            ptr
        );

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

        snapshot.apply(log_kind, pid, lsn, ptr, sz);
    }

    // `snapshot.tip_lid` can be set based on 4 possibilities for the tip of the
    // log:
    // 1. an empty DB - tip set to None, causing a fresh segment to be
    //    allocated on initialization
    // 2. the recovered tip is at the end of a
    //    segment with less space left than would fit MAX_MSG_HEADER_LEN -
    //    tip set to None, causing a fresh segment to be allocated on
    //    initialization, as in #1 above
    // 3. the recovered tip is in the middle of a segment - both set to the end
    //    of the last valid message, causing the system to be initialized to
    //    that point without allocating a new segment
    // 4. the recovered tip is at the beginning of a new segment, but without
    //    any valid messages in it yet. treat as #3 above, but also take care
    //    in te SA initialization to properly initialize any segment tracking
    //    state despite not having any pages currently residing there.
    //
    // The information we have at this point is iter.cur_lsn and the
    // iter.segment_base.offset. Situation 1 & 2 happen if iter.cur_lsn %
    // segment_size == 0, and result in None for the recovered log tip.
    // Situations 3 & 4 otherwise. We don't need to differentiate between them
    // here.

    let segment_progress: Lsn = iter.cur_lsn % (config.segment_size as Lsn);
    snapshot.stable_lsn = iter.cur_lsn;

    if segment_progress == 0 || iter.segment_base.is_none() {
        // either situation 1 or situation 2
        snapshot.tip_lid = None;
    } else if let Some(BasedBuf { offset, .. }) = iter.segment_base {
        // either situation 3 or situation 4. we need to zero the

        snapshot.tip_lid = Some(offset + segment_progress as LogOffset);

        // zero the tail of the segment after the recovered tip
        let shred_len = config.segment_size - segment_progress as usize;
        let shred_zone = vec![MessageKind::Corrupted.into(); shred_len];
        debug!(
            "zeroing the end of the recovered segment between {} and {}",
            snapshot.tip_lid.unwrap(),
            snapshot.tip_lid.unwrap() + shred_len as LogOffset
        );
        pwrite_all(&config.file, &shred_zone, snapshot.tip_lid.unwrap())?;
        config.file.sync_all()?;
    } else {
        unreachable!()
    }

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

    if snapshot.stable_lsn != old_stable_lsn {
        write_snapshot(config, &snapshot)?;
    }

    for to_zero in iter.segments.values().copied() {
        debug!("zeroing torn segment at lid {}", to_zero);

        // NB we intentionally corrupt this header to prevent any segment
        // from being allocated which would duplicate its LSN, messing
        // up recovery in the future.
        io_fail!(config, "segment initial free zero");
        pwrite_all(
            &config.file,
            &*vec![MessageKind::Corrupted.into(); SEG_HEADER_LEN],
            to_zero,
        )?;
        if !config.temporary {
            config.file.sync_all()?;
        }
    }

    // remove all blob files larger than our stable offset
    gc_blobs(config, iter.cur_lsn)?;

    #[cfg(feature = "event_log")]
    config.event_log.recovered_lsn(iter.cur_lsn);

    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(config: &RunningConfig) -> Result<Snapshot> {
    let last_snap = read_snapshot(config)?.unwrap_or_else(Snapshot::default);

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

    let res = advance_snapshot(log_iter, last_snap, config)?;

    Ok(res)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(config: &RunningConfig) -> std::io::Result<Option<Snapshot>> {
    let mut f = loop {
        let mut candidates = config.get_snapshot_files()?;
        if candidates.is_empty() {
            debug!("no previous snapshot found");
            return Ok(None);
        }

        candidates.sort();

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

        match std::fs::OpenOptions::new().read(true).open(&path) {
            Ok(f) => break f,
            Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                // this can happen if there's a race
                continue;
            }
            Err(other) => return Err(other),
        }
    };

    let mut buf = vec![];
    let _read = f.read_to_end(&mut buf)?;
    let len = buf.len();
    if len <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut len_expected_bytes = [0; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    let _ = buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        use std::convert::TryInto;

        let len_expected: u64 =
            u64::from_le_bytes(len_expected_bytes.as_ref().try_into().unwrap());

        decompress(&*buf, usize::try_from(len_expected).unwrap()).unwrap()
    } else {
        buf
    };

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

    Ok(Snapshot::deserialize(&mut bytes.as_slice()).ok())
}

fn write_snapshot(config: &RunningConfig, snapshot: &Snapshot) -> Result<()> {
    trace!("writing snapshot {:?}", snapshot);

    let raw_bytes = snapshot.serialize();
    let decompressed_len = raw_bytes.len();

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

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

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

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

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

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

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

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    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
    io_fail!(config, "snap write");
    f.write_all(&*bytes)?;
    io_fail!(config, "snap write len");
    f.write_all(&len_bytes)?;
    io_fail!(config, "snap write crc");
    f.write_all(&crc32)?;
    io_fail!(config, "snap write post");
    f.sync_all()?;

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

    io_fail!(config, "snap write mv");
    std::fs::rename(&path_1, &path_2)?;
    io_fail!(config, "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);

            io_fail!(config, "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>Bump snapshot stable_lsn to the end of the segment to enforce monotonic LSN progress, corresponding to runtime stabilization<commit_after>#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use crate::*;

use super::{
    arr_to_u32, gc_blobs, pwrite_all, raw_segment_iter_from, u32_to_arr,
    u64_to_arr, BasedBuf, DiskPtr, LogIter, LogKind, LogOffset, Lsn,
    MessageKind,
};

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
#[derive(Debug, Default)]
#[cfg_attr(test, derive(Clone, PartialEq))]
pub struct Snapshot {
    /// The last read message lsn
    pub stable_lsn: Lsn,
    /// The last read message lid
    pub tip_lid: Option<LogOffset>,
    /// the mapping from pages to (lsn, lid)
    pub pt: Vec<PageState>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum PageState {
    Present { base: (Lsn, DiskPtr, u64), frags: Vec<(Lsn, DiskPtr, u64)> },
    Free(Lsn, DiskPtr),
    Uninitialized,
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr, u64)) {
        match *self {
            PageState::Present { base, ref mut frags } => {
                if frags.last().map_or(base.0, |f| f.0) < item.0 {
                    frags.push(item)
                } else {
                    debug!(
                        "skipping merging item {:?} into \
                        existing PageState::Present({:?})",
                        item, frags
                    );
                }
            }
            _ => panic!("pushed frags to {:?}", self),
        }
    }

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

impl Snapshot {
    fn apply(
        &mut self,
        log_kind: LogKind,
        pid: PageId,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        sz: u64,
    ) {
        // 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 _measure = Measure::new(&M.snapshot_apply);

        if self.pt.len() <= usize::try_from(pid).unwrap() {
            self.pt.resize(
                usize::try_from(pid + 1).unwrap(),
                PageState::Uninitialized,
            );
        }

        match log_kind {
            LogKind::Replace => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn,
                );

                self.pt[usize::try_from(pid).unwrap()] = PageState::Present {
                    base: (lsn, disk_ptr, sz),
                    frags: vec![],
                };
            }
            LogKind::Link => {
                // 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 @ PageState::Present { .. }) =
                    self.pt.get_mut(usize::try_from(pid).unwrap())
                {
                    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;
                    }

                    lids.push((lsn, disk_ptr, sz));
                } else {
                    trace!(
                        "skipping dangling append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );
                }
            }
            LogKind::Free => {
                trace!("free of pid {} at ptr {} lsn {}", pid, disk_ptr, lsn);
                self.pt[usize::try_from(pid).unwrap()] =
                    PageState::Free(lsn, disk_ptr);
            }
            LogKind::Corrupted | LogKind::Skip => panic!(
                "unexppected messagekind in snapshot application: {:?}",
                log_kind
            ),
        }
    }
}

fn advance_snapshot(
    mut iter: LogIter,
    mut snapshot: Snapshot,
    config: &RunningConfig,
) -> Result<Snapshot> {
    let _measure = Measure::new(&M.advance_snapshot);

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

    let old_stable_lsn = snapshot.stable_lsn;

    while let Some((log_kind, pid, lsn, ptr, sz)) = iter.next() {
        trace!(
            "in advance_snapshot looking at item with pid {} lsn {} ptr {}",
            pid,
            lsn,
            ptr
        );

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

        snapshot.apply(log_kind, pid, lsn, ptr, sz);
    }

    // `snapshot.tip_lid` can be set based on 4 possibilities for the tip of the
    // log:
    // 1. an empty DB - tip set to None, causing a fresh segment to be
    //    allocated on initialization
    // 2. the recovered tip is at the end of a
    //    segment with less space left than would fit MAX_MSG_HEADER_LEN -
    //    tip set to None, causing a fresh segment to be allocated on
    //    initialization, as in #1 above
    // 3. the recovered tip is in the middle of a segment - both set to the end
    //    of the last valid message, causing the system to be initialized to
    //    that point without allocating a new segment
    // 4. the recovered tip is at the beginning of a new segment, but without
    //    any valid messages in it yet. treat as #3 above, but also take care
    //    in te SA initialization to properly initialize any segment tracking
    //    state despite not having any pages currently residing there.
    //
    // The information we have at this point is iter.cur_lsn and the
    // iter.segment_base.offset. Situation 1 & 2 happen if iter.cur_lsn %
    // segment_size == 0, and result in None for the recovered log tip.
    // Situations 3 & 4 otherwise. We don't need to differentiate between them
    // here.

    let segment_progress: Lsn = iter.cur_lsn % (config.segment_size as Lsn);
    snapshot.stable_lsn = if segment_progress + MAX_MSG_HEADER_LEN as Lsn
        >= config.segment_size as Lsn
    {
        config.normalize(iter.cur_lsn) + config.segment_size as Lsn
    } else {
        iter.cur_lsn
    };

    if segment_progress == 0 || iter.segment_base.is_none() {
        // either situation 1 or situation 2
        snapshot.tip_lid = None;
    } else if let Some(BasedBuf { offset, .. }) = iter.segment_base {
        // either situation 3 or situation 4. we need to zero the

        snapshot.tip_lid = Some(offset + segment_progress as LogOffset);

        // zero the tail of the segment after the recovered tip
        let shred_len = config.segment_size - segment_progress as usize;
        let shred_zone = vec![MessageKind::Corrupted.into(); shred_len];
        debug!(
            "zeroing the end of the recovered segment between {} and {}",
            snapshot.tip_lid.unwrap(),
            snapshot.tip_lid.unwrap() + shred_len as LogOffset
        );
        pwrite_all(&config.file, &shred_zone, snapshot.tip_lid.unwrap())?;
        config.file.sync_all()?;
    } else {
        unreachable!()
    }

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

    if snapshot.stable_lsn != old_stable_lsn {
        write_snapshot(config, &snapshot)?;
    }

    for to_zero in iter.segments.values().copied() {
        debug!("zeroing torn segment at lid {}", to_zero);

        // NB we intentionally corrupt this header to prevent any segment
        // from being allocated which would duplicate its LSN, messing
        // up recovery in the future.
        io_fail!(config, "segment initial free zero");
        pwrite_all(
            &config.file,
            &*vec![MessageKind::Corrupted.into(); SEG_HEADER_LEN],
            to_zero,
        )?;
        if !config.temporary {
            config.file.sync_all()?;
        }
    }

    // remove all blob files larger than our stable offset
    gc_blobs(config, snapshot.stable_lsn)?;

    #[cfg(feature = "event_log")]
    config.event_log.recovered_lsn(snapshot.stable_lsn);

    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(config: &RunningConfig) -> Result<Snapshot> {
    let last_snap = read_snapshot(config)?.unwrap_or_else(Snapshot::default);

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

    let res = advance_snapshot(log_iter, last_snap, config)?;

    Ok(res)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(config: &RunningConfig) -> std::io::Result<Option<Snapshot>> {
    let mut f = loop {
        let mut candidates = config.get_snapshot_files()?;
        if candidates.is_empty() {
            debug!("no previous snapshot found");
            return Ok(None);
        }

        candidates.sort();

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

        match std::fs::OpenOptions::new().read(true).open(&path) {
            Ok(f) => break f,
            Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                // this can happen if there's a race
                continue;
            }
            Err(other) => return Err(other),
        }
    };

    let mut buf = vec![];
    let _read = f.read_to_end(&mut buf)?;
    let len = buf.len();
    if len <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut len_expected_bytes = [0; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    let _ = buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        use std::convert::TryInto;

        let len_expected: u64 =
            u64::from_le_bytes(len_expected_bytes.as_ref().try_into().unwrap());

        decompress(&*buf, usize::try_from(len_expected).unwrap()).unwrap()
    } else {
        buf
    };

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

    Ok(Snapshot::deserialize(&mut bytes.as_slice()).ok())
}

fn write_snapshot(config: &RunningConfig, snapshot: &Snapshot) -> Result<()> {
    trace!("writing snapshot {:?}", snapshot);

    let raw_bytes = snapshot.serialize();
    let decompressed_len = raw_bytes.len();

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

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

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

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

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

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

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

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    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
    io_fail!(config, "snap write");
    f.write_all(&*bytes)?;
    io_fail!(config, "snap write len");
    f.write_all(&len_bytes)?;
    io_fail!(config, "snap write crc");
    f.write_all(&crc32)?;
    io_fail!(config, "snap write post");
    f.sync_all()?;

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

    io_fail!(config, "snap write mv");
    std::fs::rename(&path_1, &path_2)?;
    io_fail!(config, "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);

            io_fail!(config, "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|>