<commit_before>// Copyright 2015 The Gfx-rs Developers.
//
// 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.

use syntax::{abi, ast, codemap, ext};
use syntax::ext::base::ItemDecorator;
use syntax::ext::build::AstBuilder;
use syntax::owned_slice::OwnedSlice;
use syntax::parse::token;
use syntax::ptr::P;

#[derive(Copy, PartialEq, Debug)]
enum Param {
    Uniform,
    Block,
    Texture,
    Special,
}

#[derive(Copy, Debug)]
enum ParamError {
    DeprecatedTexture,
}

/// Classify variable types (`i32`, `TextureParam`, etc) into the `Param`
fn classify(node: &ast::Ty_) -> Result<Param, ParamError> {
    match *node {
        ast::TyPath(_,ref path) => match path.segments.last() {
            Some(segment) => match segment.identifier.name.as_str() {
                "RawBufferHandle" => Ok(Param::Block),
                "TextureParam"    => Ok(Param::Texture),
                "TextureHandle"   => Err(ParamError::DeprecatedTexture),
                "PhantomData"     => Ok(Param::Special),
                _ => Ok(Param::Uniform),
            },
            None => Ok(Param::Uniform),
        },
        _ => Ok(Param::Uniform),
    }
}

/// Generates the the method body for `gfx::shade::ShaderParam::create_link`
fn method_create(cx: &mut ext::base::ExtCtxt,
                 span: codemap::Span,
                 definition: &ast::StructDef,
                 input: ast::Ident,
                 link_ident: ast::Ident,
                 path_root: ast::Ident)
                 -> P<ast::Block> {
    let init_expr = cx.expr_struct_ident(
        span, link_ident,
        definition.fields.iter().scan((), |_, field| {
            field.node.ident().map(|name|
                cx.field_imm(field.span, name, cx.expr_none(field.span))
            )
        }).collect()
    );
    let class_info: Vec<(Param, P<ast::Expr>)> = definition.fields.iter().scan((), |_, field|
        match (field.node.ident(), classify(&field.node.ty.node)) {
            (None, _) => {
                cx.span_err(field.span, "Named fields are required for `ShaderParam`");
                None
            },
            (Some(fname), Ok(c)) => {
                let name = match super::find_name(cx, field.span, &field.node.attrs) {
                    Some(name) => name,
                    None => token::get_ident(fname),
                };
                Some((c, cx.expr_str(field.span, name)))
            },
            (Some(fname), Err(e)) => {
                cx.span_err(field.span, &format!(
                    "Unrecognized parameter ({:?}) type {:?}",
                    fname.as_str(), e
                ));
                None
            },
        }
    ).collect();
    let gen_arms = |ptype: Param, var: ast::Ident| -> Vec<ast::Arm> {
        class_info.iter().zip(definition.fields.iter())
                  .filter(|&(&(ref class, _), _)| *class == ptype)
                  .scan((), |_, (&(_, ref name_expr), field)|
            field.node.ident().map(|name| quote_arm!(cx,
                $name_expr => {out.$name = Some(i as $path_root::gfx::shade::$var)}
            ))
        ).collect()
    };
    let uniform_arms = gen_arms(Param::Uniform, cx.ident_of("VarUniform"));
    let block_arms = gen_arms(Param::Block, cx.ident_of("VarBlock"));
    let texture_arms = gen_arms(Param::Texture, cx.ident_of("VarTexture"));
    let expr = quote_expr!(cx, {
        let mut out = $init_expr;
        for (i, u) in $input.uniforms.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &u.name[..] {
                $uniform_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingUniform(u.name.clone())),
            }
        }
        for (i, b) in $input.blocks.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &b.name[..] {
                $block_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingBlock(b.name.clone())),
            }
        }
        for (i, t) in $input.textures.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &t.name[..] {
                $texture_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingTexture(t.name.clone())),
            }
        }
        Ok(out)
    });
    cx.block_expr(expr)
}

/// Generates the the method body for `gfx::shade::ShaderParam::fill_params`
fn method_fill(cx: &mut ext::base::ExtCtxt,
               span: codemap::Span,
               definition: &ast::StructDef,
               path_root: ast::Ident)
               -> P<ast::Block> {
    let max_num = cx.expr_usize(span, definition.fields.len());
    let mut calls = vec![
        quote_stmt!(cx, use self::$path_root::gfx::shade::ToUniform;),
        quote_stmt!(cx, out.uniforms.reserve($max_num);),
        quote_stmt!(cx, out.blocks.reserve($max_num);),
        quote_stmt!(cx, out.textures.reserve($max_num);),
    ];
    calls.extend(definition.fields.iter().scan((), |_, field| {
        let name = match field.node.ident() {
            Some(n) => n,
            None => {
                cx.span_err(span, "Named fields are required for `ShaderParam`");
                return None
            }
        };
        classify(&field.node.ty.node).ok().map(|param| match param {
            Param::Uniform => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    out.uniforms[id as usize] = Some(self.$name.to_uniform());
                })
            ),
            Param::Block   => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    out.blocks[id as usize] = Some(self.$name.clone());
                })
            ),
            Param::Texture => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    out.textures[id as usize] = Some(self.$name.clone());
                })
            ),
            Param::Special => quote_stmt!(cx, ()),
        })
    }));
    cx.block_all(span, calls, None)
}

/// A helper function that translates variable type (`i32`, `TextureHandle`, etc)
/// into the corresponding shader var id type (`VarUniform`, `VarBlock`, or `VarTexture`)
fn node_to_var_type(cx: &mut ext::base::ExtCtxt,
                    span: codemap::Span, node: &ast::Ty_,
                    path_root: ast::Ident) -> P<ast::Ty> {
    let id = cx.ident_of(match classify(node) {
        Ok(Param::Uniform) => "VarUniform",
        Ok(Param::Block)   => "VarBlock",
        Ok(Param::Texture) => "VarTexture",
        Ok(Param::Special) => return quote_ty!(cx, Option<()>),
        Err(ParamError::DeprecatedTexture) => {
            cx.span_err(span, "Use gfx::shade::TextureParam for texture vars instead of gfx::shade::TextureHandle");
            ""
        },
    });
    quote_ty!(cx, Option<$path_root::gfx::shade::$id>)
}

fn impl_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
             name: &str, type_ident: ast::Ident) -> ast::ImplItem {
    ast::ImplItem {
        id: ast::DUMMY_NODE_ID,
        ident: cx.ident_of(name),
        vis: ast::Visibility::Inherited,
        attrs: Vec::new(),
        node: ast::TypeImplItem(cx.ty_ident(span, type_ident)),
        span: span
    }
}

fn impl_method(cx: &mut ext::base::ExtCtxt, span: codemap::Span, name: &str,
               with_self: bool, declaration: P<ast::FnDecl>, body: P<ast::Block>)
               -> ast::ImplItem {
    ast::ImplItem {
        id: ast::DUMMY_NODE_ID,
        ident: cx.ident_of(name),
        vis: ast::Visibility::Inherited,
        attrs: Vec::new(),
        node: ast::MethodImplItem(
            ast::MethodSig {
                unsafety: ast::Unsafety::Normal,
                abi: abi::Abi::Rust,
                decl: declaration,
                generics: ast::Generics {
                    lifetimes: Vec::new(),
                    ty_params: OwnedSlice::empty(),
                    where_clause: ast::WhereClause {
                        id: ast::DUMMY_NODE_ID,
                        predicates: Vec::new()
                    }
                },
                explicit_self: codemap::Spanned {
                    node: if with_self {
                        ast::SelfRegion(None, ast::MutImmutable, cx.ident_of("self"))
                    } else {
                        ast::SelfStatic
                    },
                    span: span
                }
            },
            body
        ),
        span: span
    }
}

#[derive(Copy)]
pub struct ShaderParam;
impl ItemDecorator for ShaderParam {
    /// Decorator for `shader_param` attribute
    fn expand(&self, context: &mut ext::base::ExtCtxt, span: codemap::Span,
              meta_item: &ast::MetaItem, item: &ast::Item,
              push: &mut FnMut(P<ast::Item>)) {
        // Insert the `gfx` reexport module
        let extern_hack = context.ident_of(super::EXTERN_CRATE_HACK);
        let path_root = super::extern_crate_hack(context, span, |i| (*push)(i));

        // constructing the Link struct
        let (base_def, generics, link_def) = match item.node {
            ast::ItemStruct(ref definition, ref generics) => {
                (definition, generics.clone(), ast::StructDef {
                    fields: definition.fields.iter()
                        .map(|f| codemap::Spanned {
                            node: ast::StructField_ {
                                kind: f.node.kind,
                                id: f.node.id,
                                ty: node_to_var_type(context, f.span, &f.node.ty.node, path_root),
                                attrs: Vec::new(),
                            },
                            span: f.span,
                        }).collect(),
                    ctor_id: None,
                })
            },
            _ => {
                context.span_err(span, "Only free-standing named structs allowed to derive ShaderParam");
                return;
            }
        };

        // derive and push
        let link_name = format!("_{}Link", item.ident.as_str());
        let link_ident = context.ident_of(&link_name);
        let link_item = context.item_struct(span, link_ident, link_def)
                               .map(|mut item| {
            item.attrs.push(context.attribute(span,
                context.meta_list(span, token::InternedString::new("derive"), vec![
                        context.meta_word(span, token::InternedString::new("Copy")),
                        context.meta_word(span, token::InternedString::new("Debug")),
                ])
            ));
            item
        });
        (*push)(link_item);

        // process meta parameters (expecting none)
        match meta_item.node {
            ast::MetaWord(_) => (), //expected
            _ => {
                context.span_err(meta_item.span, "#[shader_param] needs no param");
            }
        }

        // find the generic `Resources` bound
        let resource_ident = match generics.ty_params.iter().find(|typ|
            typ.bounds.iter().find(|b| match **b {
                ast::TraitTyParamBound(ref poly_trait, _) =>
                    poly_trait.trait_ref.path.segments.last().unwrap().identifier.as_str() == "Resources",
                ast::RegionTyParamBound(_) => false,
            }).is_some()
        ){
            Some(typ) => typ.ident,
            None => {
                context.span_err(meta_item.span, "#[shader_param] unable to find generic `gfx::Resources` bound");
                context.ident_of("R")
            }
        };

        // construct `create_link()`
        let lifetimes = generics.lifetimes.iter().map(|ld| ld.lifetime).collect();
        let generic_parameters = generics.ty_params.iter().map(|ty|
            context.ty_ident(span, ty.ident)
        ).collect();
        let struct_ty = context.ty_path(context.path_all(
            span, false,
            vec![item.ident],
            lifetimes,
            generic_parameters,
            Vec::new(),
        ));
        let create_param = context.ident_of("params");
        let body_create = method_create(context, span, base_def, create_param, link_ident, path_root);
        let decl_create = context.fn_decl(
            vec![
                ast::Arg {
                    ty: quote_ty!(context, Option<&$struct_ty>),
                    pat: context.pat_wild(span),
                    id: ast::DUMMY_NODE_ID,
                },
                ast::Arg {
                    ty: quote_ty!(context, &$extern_hack::gfx::ProgramInfo),
                    pat: context.pat_ident(span, create_param),
                    id: ast::DUMMY_NODE_ID,
                },
            ],
            context.ty_path(context.path_all(
                span, false, vec![context.ident_of("Result")],
                Vec::new(), vec![
                    context.ty_ident(span, link_ident),
                    quote_ty!(context, $extern_hack::gfx::shade::ParameterError),
                ], Vec::new()
            )),
        );

        // construct `fill_params()`
        let body_fill = method_fill(context, span, base_def, path_root);
        let decl_fill = context.fn_decl(
             vec![
                ast::Arg::new_self(span, ast::MutImmutable, context.ident_of("self")),
                ast::Arg {
                    ty: quote_ty!(context, &$link_ident),
                    pat: context.pat_ident(span, context.ident_of("link")),
                    id: ast::DUMMY_NODE_ID,
                },
                ast::Arg {
                    ty: quote_ty!(context, &mut $extern_hack::gfx::ParamStorage<$resource_ident>),
                    pat: context.pat_ident(span, context.ident_of("out")),
                    id: ast::DUMMY_NODE_ID,
                },
            ],
            quote_ty!(context, ())
        );

        // construct implementations for types and methods
        let impls = vec![
            P(impl_type(context, span, "Resources", resource_ident)),
            P(impl_type(context, span, "Link", link_ident)),
            P(impl_method(context, span, "create_link", false, decl_create, body_create)),
            P(impl_method(context, span, "fill_params", true, decl_fill, body_fill))
        ];

        // final implementation item
        let item = context.item(span, item.ident, Vec::new(), ast::Item_::ItemImpl(
                ast::Unsafety::Normal,
                ast::ImplPolarity::Positive,
                generics,
                Some(context.trait_ref(context.path(span, vec![
                    extern_hack,
                    context.ident_of("gfx"),
                    context.ident_of("shade"),
                    context.ident_of("ShaderParam"),
                ]))),
                struct_ty,
                impls
        ));
        (*push)(super::fixup_extern_crate_paths(item, path_root));
    }
}
<commit_msg>Update for latest Rust<commit_after>// Copyright 2015 The Gfx-rs Developers.
//
// 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.

use syntax::{abi, ast, codemap, ext};
use syntax::ext::base::ItemDecorator;
use syntax::ext::build::AstBuilder;
use syntax::owned_slice::OwnedSlice;
use syntax::parse::token;
use syntax::ptr::P;

#[derive(Copy, PartialEq, Debug)]
enum Param {
    Uniform,
    Block,
    Texture,
    Special,
}

#[derive(Copy, Debug)]
enum ParamError {
    DeprecatedTexture,
}

/// Classify variable types (`i32`, `TextureParam`, etc) into the `Param`
fn classify(node: &ast::Ty_) -> Result<Param, ParamError> {
    match *node {
        ast::TyPath(_,ref path) => match path.segments.last() {
            Some(segment) => match segment.identifier.name.as_str() {
                "RawBufferHandle" => Ok(Param::Block),
                "TextureParam"    => Ok(Param::Texture),
                "TextureHandle"   => Err(ParamError::DeprecatedTexture),
                "PhantomData"     => Ok(Param::Special),
                _ => Ok(Param::Uniform),
            },
            None => Ok(Param::Uniform),
        },
        _ => Ok(Param::Uniform),
    }
}

/// Generates the the method body for `gfx::shade::ShaderParam::create_link`
fn method_create(cx: &mut ext::base::ExtCtxt,
                 span: codemap::Span,
                 definition: &ast::StructDef,
                 input: ast::Ident,
                 link_ident: ast::Ident,
                 path_root: ast::Ident)
                 -> P<ast::Block> {
    let init_expr = cx.expr_struct_ident(
        span, link_ident,
        definition.fields.iter().scan((), |_, field| {
            field.node.ident().map(|name|
                cx.field_imm(field.span, name, cx.expr_none(field.span))
            )
        }).collect()
    );
    let class_info: Vec<(Param, P<ast::Expr>)> = definition.fields.iter().scan((), |_, field|
        match (field.node.ident(), classify(&field.node.ty.node)) {
            (None, _) => {
                cx.span_err(field.span, "Named fields are required for `ShaderParam`");
                None
            },
            (Some(fname), Ok(c)) => {
                let name = match super::find_name(cx, field.span, &field.node.attrs) {
                    Some(name) => name,
                    None => token::get_ident(fname),
                };
                Some((c, cx.expr_str(field.span, name)))
            },
            (Some(fname), Err(e)) => {
                cx.span_err(field.span, &format!(
                    "Unrecognized parameter ({:?}) type {:?}",
                    fname.as_str(), e
                ));
                None
            },
        }
    ).collect();
    let gen_arms = |ptype: Param, var: ast::Ident| -> Vec<ast::Arm> {
        class_info.iter().zip(definition.fields.iter())
                  .filter(|&(&(ref class, _), _)| *class == ptype)
                  .scan((), |_, (&(_, ref name_expr), field)|
            field.node.ident().map(|name| quote_arm!(cx,
                $name_expr => {out.$name = Some(i as $path_root::gfx::shade::$var)}
            ))
        ).collect()
    };
    let uniform_arms = gen_arms(Param::Uniform, cx.ident_of("VarUniform"));
    let block_arms = gen_arms(Param::Block, cx.ident_of("VarBlock"));
    let texture_arms = gen_arms(Param::Texture, cx.ident_of("VarTexture"));
    let expr = quote_expr!(cx, {
        let mut out = $init_expr;
        for (i, u) in $input.uniforms.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &u.name[..] {
                $uniform_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingUniform(u.name.clone())),
            }
        }
        for (i, b) in $input.blocks.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &b.name[..] {
                $block_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingBlock(b.name.clone())),
            }
        }
        for (i, t) in $input.textures.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &t.name[..] {
                $texture_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingTexture(t.name.clone())),
            }
        }
        Ok(out)
    });
    cx.block_expr(expr)
}

/// Generates the the method body for `gfx::shade::ShaderParam::fill_params`
fn method_fill(cx: &mut ext::base::ExtCtxt,
               span: codemap::Span,
               definition: &ast::StructDef,
               path_root: ast::Ident)
               -> P<ast::Block> {
    let max_num = cx.expr_usize(span, definition.fields.len());
    let mut calls = vec![
        quote_stmt!(cx, use self::$path_root::gfx::shade::ToUniform;),
        quote_stmt!(cx, out.uniforms.reserve($max_num);),
        quote_stmt!(cx, out.blocks.reserve($max_num);),
        quote_stmt!(cx, out.textures.reserve($max_num);),
    ];
    calls.extend(definition.fields.iter().scan((), |_, field| {
        let name = match field.node.ident() {
            Some(n) => n,
            None => {
                cx.span_err(span, "Named fields are required for `ShaderParam`");
                return None
            }
        };
        classify(&field.node.ty.node).ok().map(|param| match param {
            Param::Uniform => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    out.uniforms[id as usize] = Some(self.$name.to_uniform());
                })
            ),
            Param::Block   => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    out.blocks[id as usize] = Some(self.$name.clone());
                })
            ),
            Param::Texture => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    out.textures[id as usize] = Some(self.$name.clone());
                })
            ),
            Param::Special => quote_stmt!(cx, ()),
        })
    }));
    let calls = calls.map_in_place(Option::unwrap);
    cx.block_all(span, calls, None)
}

/// A helper function that translates variable type (`i32`, `TextureHandle`, etc)
/// into the corresponding shader var id type (`VarUniform`, `VarBlock`, or `VarTexture`)
fn node_to_var_type(cx: &mut ext::base::ExtCtxt,
                    span: codemap::Span, node: &ast::Ty_,
                    path_root: ast::Ident) -> P<ast::Ty> {
    let id = cx.ident_of(match classify(node) {
        Ok(Param::Uniform) => "VarUniform",
        Ok(Param::Block)   => "VarBlock",
        Ok(Param::Texture) => "VarTexture",
        Ok(Param::Special) => return quote_ty!(cx, Option<()>),
        Err(ParamError::DeprecatedTexture) => {
            cx.span_err(span, "Use gfx::shade::TextureParam for texture vars instead of gfx::shade::TextureHandle");
            ""
        },
    });
    quote_ty!(cx, Option<$path_root::gfx::shade::$id>)
}

fn impl_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
             name: &str, type_ident: ast::Ident) -> ast::ImplItem {
    ast::ImplItem {
        id: ast::DUMMY_NODE_ID,
        ident: cx.ident_of(name),
        vis: ast::Visibility::Inherited,
        attrs: Vec::new(),
        node: ast::TypeImplItem(cx.ty_ident(span, type_ident)),
        span: span
    }
}

fn impl_method(cx: &mut ext::base::ExtCtxt, span: codemap::Span, name: &str,
               with_self: bool, declaration: P<ast::FnDecl>, body: P<ast::Block>)
               -> ast::ImplItem {
    ast::ImplItem {
        id: ast::DUMMY_NODE_ID,
        ident: cx.ident_of(name),
        vis: ast::Visibility::Inherited,
        attrs: Vec::new(),
        node: ast::MethodImplItem(
            ast::MethodSig {
                unsafety: ast::Unsafety::Normal,
                abi: abi::Abi::Rust,
                decl: declaration,
                generics: ast::Generics {
                    lifetimes: Vec::new(),
                    ty_params: OwnedSlice::empty(),
                    where_clause: ast::WhereClause {
                        id: ast::DUMMY_NODE_ID,
                        predicates: Vec::new()
                    }
                },
                explicit_self: codemap::Spanned {
                    node: if with_self {
                        ast::SelfRegion(None, ast::MutImmutable, cx.ident_of("self"))
                    } else {
                        ast::SelfStatic
                    },
                    span: span
                }
            },
            body
        ),
        span: span
    }
}

#[derive(Copy)]
pub struct ShaderParam;
impl ItemDecorator for ShaderParam {
    /// Decorator for `shader_param` attribute
    fn expand(&self, context: &mut ext::base::ExtCtxt, span: codemap::Span,
              meta_item: &ast::MetaItem, item: &ast::Item,
              push: &mut FnMut(P<ast::Item>)) {
        // Insert the `gfx` reexport module
        let extern_hack = context.ident_of(super::EXTERN_CRATE_HACK);
        let path_root = super::extern_crate_hack(context, span, |i| (*push)(i));

        // constructing the Link struct
        let (base_def, generics, link_def) = match item.node {
            ast::ItemStruct(ref definition, ref generics) => {
                (definition, generics.clone(), ast::StructDef {
                    fields: definition.fields.iter()
                        .map(|f| codemap::Spanned {
                            node: ast::StructField_ {
                                kind: f.node.kind,
                                id: f.node.id,
                                ty: node_to_var_type(context, f.span, &f.node.ty.node, path_root),
                                attrs: Vec::new(),
                            },
                            span: f.span,
                        }).collect(),
                    ctor_id: None,
                })
            },
            _ => {
                context.span_err(span, "Only free-standing named structs allowed to derive ShaderParam");
                return;
            }
        };

        // derive and push
        let link_name = format!("_{}Link", item.ident.as_str());
        let link_ident = context.ident_of(&link_name);
        let link_item = context.item_struct(span, link_ident, link_def)
                               .map(|mut item| {
            item.attrs.push(context.attribute(span,
                context.meta_list(span, token::InternedString::new("derive"), vec![
                        context.meta_word(span, token::InternedString::new("Copy")),
                        context.meta_word(span, token::InternedString::new("Debug")),
                ])
            ));
            item
        });
        (*push)(link_item);

        // process meta parameters (expecting none)
        match meta_item.node {
            ast::MetaWord(_) => (), //expected
            _ => {
                context.span_err(meta_item.span, "#[shader_param] needs no param");
            }
        }

        // find the generic `Resources` bound
        let resource_ident = match generics.ty_params.iter().find(|typ|
            typ.bounds.iter().find(|b| match **b {
                ast::TraitTyParamBound(ref poly_trait, _) =>
                    poly_trait.trait_ref.path.segments.last().unwrap().identifier.as_str() == "Resources",
                ast::RegionTyParamBound(_) => false,
            }).is_some()
        ){
            Some(typ) => typ.ident,
            None => {
                context.span_err(meta_item.span, "#[shader_param] unable to find generic `gfx::Resources` bound");
                context.ident_of("R")
            }
        };

        // construct `create_link()`
        let lifetimes = generics.lifetimes.iter().map(|ld| ld.lifetime).collect();
        let generic_parameters = generics.ty_params.iter().map(|ty|
            context.ty_ident(span, ty.ident)
        ).collect();
        let struct_ty = context.ty_path(context.path_all(
            span, false,
            vec![item.ident],
            lifetimes,
            generic_parameters,
            Vec::new(),
        ));
        let create_param = context.ident_of("params");
        let body_create = method_create(context, span, base_def, create_param, link_ident, path_root);
        let decl_create = context.fn_decl(
            vec![
                ast::Arg {
                    ty: quote_ty!(context, Option<&$struct_ty>),
                    pat: context.pat_wild(span),
                    id: ast::DUMMY_NODE_ID,
                },
                ast::Arg {
                    ty: quote_ty!(context, &$extern_hack::gfx::ProgramInfo),
                    pat: context.pat_ident(span, create_param),
                    id: ast::DUMMY_NODE_ID,
                },
            ],
            context.ty_path(context.path_all(
                span, false, vec![context.ident_of("Result")],
                Vec::new(), vec![
                    context.ty_ident(span, link_ident),
                    quote_ty!(context, $extern_hack::gfx::shade::ParameterError),
                ], Vec::new()
            )),
        );

        // construct `fill_params()`
        let body_fill = method_fill(context, span, base_def, path_root);
        let decl_fill = context.fn_decl(
             vec![
                ast::Arg::new_self(span, ast::MutImmutable, context.ident_of("self")),
                ast::Arg {
                    ty: quote_ty!(context, &$link_ident),
                    pat: context.pat_ident(span, context.ident_of("link")),
                    id: ast::DUMMY_NODE_ID,
                },
                ast::Arg {
                    ty: quote_ty!(context, &mut $extern_hack::gfx::ParamStorage<$resource_ident>),
                    pat: context.pat_ident(span, context.ident_of("out")),
                    id: ast::DUMMY_NODE_ID,
                },
            ],
            quote_ty!(context, ())
        );

        // construct implementations for types and methods
        let impls = vec![
            P(impl_type(context, span, "Resources", resource_ident)),
            P(impl_type(context, span, "Link", link_ident)),
            P(impl_method(context, span, "create_link", false, decl_create, body_create)),
            P(impl_method(context, span, "fill_params", true, decl_fill, body_fill))
        ];

        // final implementation item
        let item = context.item(span, item.ident, Vec::new(), ast::Item_::ItemImpl(
                ast::Unsafety::Normal,
                ast::ImplPolarity::Positive,
                generics,
                Some(context.trait_ref(context.path(span, vec![
                    extern_hack,
                    context.ident_of("gfx"),
                    context.ident_of("shade"),
                    context.ident_of("ShaderParam"),
                ]))),
                struct_ty,
                impls
        ));
        (*push)(super::fixup_extern_crate_paths(item, path_root));
    }
}
<|endoftext|>


<commit_before>//! Matrix room membership.

use std::convert::TryInto;
use std::error::Error;

use diesel::{
    Connection,
    ExpressionMethods,
    ExecuteDsl,
    LoadDsl,
    FilterDsl,
    FindDsl,
    SaveChangesDsl,
    SelectDsl,
    insert,
    update,
};
use diesel::associations::Identifiable;
use diesel::expression::dsl::*;
use diesel::pg::PgConnection;
use diesel::pg::data_types::PgTimestamp;
use diesel::result::Error as DieselError;
use ruma_events::EventType;
use ruma_events::room::join_rules::JoinRule;
use ruma_events::room::member::{
    MemberEvent,
    MembershipState,
    MemberEventContent,
    MemberEventExtraContent
};
use ruma_identifiers::{EventId, RoomId, UserId};
use serde_json::{Error as SerdeJsonError, Value, from_value};

use error::ApiError;
use event::{NewEvent, Event};
use profile::Profile;
use room::Room;
use schema::{events, room_memberships};

/// Room membership update or create data.
#[derive(Debug, Clone)]
pub struct RoomMembershipOptions {
    /// The room's ID.
    pub room_id: RoomId,
    /// The user's ID.
    pub user_id: UserId,
    /// The ID of the user who created the membership.
    pub sender: UserId,
    /// The current membership state.
    pub membership: String,
}

/// A new Matrix room membership, not yet saved.
#[derive(Debug, Clone)]
#[insertable_into(room_memberships)]
pub struct NewRoomMembership {
    /// The eventID.
    pub event_id: EventId,
    /// The room's ID.
    pub room_id: RoomId,
    /// The user's ID.
    pub user_id: UserId,
    /// The ID of the user who created the membership.
    pub sender: UserId,
    /// The current membership state.
    pub membership: String,
}

/// A Matrix room membership.
#[derive(Debug, Clone, Queryable)]
#[changeset_for(room_memberships)]
pub struct RoomMembership {
    /// The eventID.
    pub event_id: EventId,
    /// The room's ID.
    pub room_id: RoomId,
    /// The user's ID.
    pub user_id: UserId,
    /// The ID of the user who created the membership.
    pub sender: UserId,
    /// The current membership state.
    pub membership: String,
    /// The time the room was created.
    pub created_at: PgTimestamp,
}

impl Identifiable for RoomMembership {
    type Id = EventId;
    type Table = room_memberships::table;

    fn id(&self) -> &Self::Id {
        &self.event_id
    }

    fn table() -> Self::Table {
        room_memberships::table
    }
}

impl RoomMembership {
    /// Creates a new `RoomMembership` in the database.
    pub fn create(connection: &PgConnection, homeserver_domain: &str, options: RoomMembershipOptions)
    -> Result<RoomMembership, ApiError> {
        connection.transaction::<RoomMembership, ApiError, _>(|| {
            let join_rules_event = Event::find_room_join_rules_by_room_id(
                &connection,
                options.clone().room_id
            )?;

            let room = Room::find(connection, &options.room_id)?;

            // Only the creator of the room can join an invite-only room,
            // without an invite.
            if options.sender != room.user_id {
                if options.membership == "join" && join_rules_event.content.join_rule == JoinRule::Invite {
                    return Err(ApiError::unauthorized(Some("You are not invited to this room")));
                }

                let power_levels = room.current_power_levels(connection)?;
                let user_power_level = power_levels
                    .users
                    .get(&options.sender)
                    .unwrap_or(&power_levels.users_default);

                if options.membership == "invite" {
                    if power_levels.invite > *user_power_level {
                        return Err(
                            ApiError::unauthorized(Some("Insufficient power level to invite"))
                        );
                    }
                }
            }

            let profile = Profile::find_by_uid(connection, options.user_id.clone())?;

            let new_member_event = RoomMembership::create_new_room_member_event(
                homeserver_domain,
                &options,
                profile
            )?;

            let new_room_membership = NewRoomMembership {
                event_id: new_member_event.id.clone(),
                room_id: options.room_id.clone(),
                user_id: options.user_id.clone(),
                sender: options.sender.clone(),
                membership: options.membership.clone(),
            };

            insert(&new_member_event)
                .into(events::table)
                .execute(connection)
                .map_err(ApiError::from)?;

            let room_membership: RoomMembership = insert(&new_room_membership)
                                                    .into(room_memberships::table)
                                                    .get_result(connection)
                                                    .map_err(ApiError::from)?;

            Ok(room_membership)
        }).map_err(ApiError::from)
    }

    /// Return `RoomMembership` for given `RoomId` and `UserId`.
    pub fn find(connection: &PgConnection, room_id: &RoomId, user_id: &UserId)
    -> Result<Option<RoomMembership>, ApiError> {
        let membership = room_memberships::table
            .filter(room_memberships::room_id.eq(room_id))
            .filter(room_memberships::user_id.eq(user_id))
            .first(connection);

        match membership {
            Ok(membership) => Ok(Some(membership)),
            Err(DieselError::NotFound) => Ok(None),
            Err(err) => Err(ApiError::from(err)),
        }
    }

    /// Return `RoomMembership`'s for given `UserId`.
    pub fn find_by_uid(connection: &PgConnection, user_id: UserId) -> Result<Vec<RoomMembership>, ApiError> {
        let room_memberships: Vec<RoomMembership> = room_memberships::table
            .filter(room_memberships::user_id.eq(user_id))
            .get_results(connection)
            .map_err(|err| match err {
                DieselError::NotFound => ApiError::not_found(Some(err.description())),
                _ => ApiError::from(err),
            })?;

        Ok(room_memberships)
    }

    /// Update an existing `RoomMembership` entry or insert a new one.
    pub fn upsert(connection: &PgConnection, domain: &str, options: RoomMembershipOptions)
    -> Result<RoomMembership, ApiError> {
        let room_membership = RoomMembership::find(
            connection,
            &options.room_id,
            &options.user_id
        )?;

        match room_membership {
            Some(mut entry) => entry.update(connection, domain, options),
            None => RoomMembership::create(connection, domain, options)
        }
    }

    /// Update a `RoomMembership` entry using new `RoomMembershipOptions`.
    ///
    /// After the update a new `MemberEvent` is created.
    pub fn update(&mut self, connection: &PgConnection, homeserver_domain: &str, options: RoomMembershipOptions)
    -> Result<RoomMembership, ApiError> {
        let profile = Profile::find_by_uid(connection, options.user_id.clone())?;

        let event = RoomMembership::create_new_room_member_event(
            &homeserver_domain,
            &options,
            profile,
        )?;

        self.membership = options.membership.clone();
        self.sender = options.sender.clone();

        connection.transaction::<RoomMembership, ApiError, _>(|| {
            insert(&event)
                .into(events::table)
                .execute(connection)
                .map_err(ApiError::from)?;

            self.save_changes::<RoomMembership>(connection)
                .map_err(ApiError::from)?;

            // Use the new `EventId` as primary key.
            update(room_memberships::table.find(self.event_id.clone()))
                .set(room_memberships::event_id.eq(event.id.clone()))
                .get_result(connection)
                .map_err(ApiError::from)
        }).map_err(ApiError::from)
    }

    /// Create a new `MemberEvent`.
    pub fn create_new_room_member_event(
        homeserver_domain: &str,
        options: &RoomMembershipOptions,
        profile: Option<Profile>
    ) -> Result<NewEvent, ApiError> {
        let event_id = EventId::new(&homeserver_domain).map_err(ApiError::from)?;
        let membership_string = Value::String(options.membership.clone());
        let membership: MembershipState = from_value(membership_string)?;

        let (avatar_url, displayname) = match profile {
            Some(profile) => (profile.avatar_url, profile.displayname),
            None => (None, None),
        };

        let new_member_event: NewEvent = MemberEvent {
            content: MemberEventContent {
                avatar_url: avatar_url,
                displayname: displayname,
                membership: membership,
                third_party_invite: (),
            },
            event_id: event_id.clone(),
            event_type: EventType::RoomMember,
            extra_content: MemberEventExtraContent { invite_room_state: None },
            prev_content: None,
            room_id: options.room_id.clone(),
            state_key: format!("@{}:{}", options.user_id.clone(), &homeserver_domain),
            unsigned: None,
            user_id: options.user_id.clone(),
        }.try_into()?;

        Ok(new_member_event)
    }

    /// Return member event's for given `room_id`.
    pub fn get_events_by_room(connection: &PgConnection, room_id: RoomId) -> Result<Vec<MemberEvent>, ApiError> {
        let event_ids = room_memberships::table
            .filter(room_memberships::room_id.eq(room_id))
            .select(room_memberships::event_id);

        let events: Vec<Event> = events::table
            .filter(events::id.eq(any(event_ids)))
            .get_results(connection)
            .map_err(|err| match err {
                DieselError::NotFound => ApiError::not_found(None),
                _ => ApiError::from(err),
            })?;

        let member_events: Result<Vec<MemberEvent>, SerdeJsonError> = events.into_iter()
                                                                        .map(TryInto::try_into)
                                                                        .collect();
        member_events.map_err(ApiError::from)
    }
}
<commit_msg>Remove unnecessary transaction wrapping<commit_after>//! Matrix room membership.

use std::convert::TryInto;
use std::error::Error;

use diesel::{
    Connection,
    ExpressionMethods,
    ExecuteDsl,
    LoadDsl,
    FilterDsl,
    FindDsl,
    SaveChangesDsl,
    SelectDsl,
    insert,
    update,
};
use diesel::associations::Identifiable;
use diesel::expression::dsl::*;
use diesel::pg::PgConnection;
use diesel::pg::data_types::PgTimestamp;
use diesel::result::Error as DieselError;
use ruma_events::EventType;
use ruma_events::room::join_rules::JoinRule;
use ruma_events::room::member::{
    MemberEvent,
    MembershipState,
    MemberEventContent,
    MemberEventExtraContent
};
use ruma_identifiers::{EventId, RoomId, UserId};
use serde_json::{Error as SerdeJsonError, Value, from_value};

use error::ApiError;
use event::{NewEvent, Event};
use profile::Profile;
use room::Room;
use schema::{events, room_memberships};

/// Room membership update or create data.
#[derive(Debug, Clone)]
pub struct RoomMembershipOptions {
    /// The room's ID.
    pub room_id: RoomId,
    /// The user's ID.
    pub user_id: UserId,
    /// The ID of the user who created the membership.
    pub sender: UserId,
    /// The current membership state.
    pub membership: String,
}

/// A new Matrix room membership, not yet saved.
#[derive(Debug, Clone)]
#[insertable_into(room_memberships)]
pub struct NewRoomMembership {
    /// The eventID.
    pub event_id: EventId,
    /// The room's ID.
    pub room_id: RoomId,
    /// The user's ID.
    pub user_id: UserId,
    /// The ID of the user who created the membership.
    pub sender: UserId,
    /// The current membership state.
    pub membership: String,
}

/// A Matrix room membership.
#[derive(Debug, Clone, Queryable)]
#[changeset_for(room_memberships)]
pub struct RoomMembership {
    /// The eventID.
    pub event_id: EventId,
    /// The room's ID.
    pub room_id: RoomId,
    /// The user's ID.
    pub user_id: UserId,
    /// The ID of the user who created the membership.
    pub sender: UserId,
    /// The current membership state.
    pub membership: String,
    /// The time the room was created.
    pub created_at: PgTimestamp,
}

impl Identifiable for RoomMembership {
    type Id = EventId;
    type Table = room_memberships::table;

    fn id(&self) -> &Self::Id {
        &self.event_id
    }

    fn table() -> Self::Table {
        room_memberships::table
    }
}

impl RoomMembership {
    /// Creates a new `RoomMembership` in the database.
    pub fn create(connection: &PgConnection, homeserver_domain: &str, options: RoomMembershipOptions)
    -> Result<RoomMembership, ApiError> {
        let join_rules_event = Event::find_room_join_rules_by_room_id(
            &connection,
            options.clone().room_id,
        )?;

        let room = Room::find(connection, &options.room_id)?;

        // Only the creator of the room can join an invite-only room, without an invite.
        if options.sender != room.user_id {
            if options.membership == "join" && join_rules_event.content.join_rule == JoinRule::Invite {
                return Err(ApiError::unauthorized(Some("You are not invited to this room")));
            }

            let power_levels = room.current_power_levels(connection)?;
            let user_power_level = power_levels
                .users
                .get(&options.sender)
                .unwrap_or(&power_levels.users_default);

            if options.membership == "invite" {
                if power_levels.invite > *user_power_level {
                    return Err(ApiError::unauthorized(Some("Insufficient power level to invite")));
                }
            }
        }

        let profile = Profile::find_by_uid(connection, options.user_id.clone())?;

        let new_member_event = RoomMembership::create_new_room_member_event(
            homeserver_domain,
            &options,
            profile,
        )?;

        let new_room_membership = NewRoomMembership {
            event_id: new_member_event.id.clone(),
            room_id: options.room_id.clone(),
            user_id: options.user_id.clone(),
            sender: options.sender.clone(),
            membership: options.membership.clone(),
        };

        connection.transaction::<RoomMembership, ApiError, _>(|| {
            insert(&new_member_event)
                .into(events::table)
                .execute(connection)
                .map_err(ApiError::from)?;

            let room_membership: RoomMembership = insert(&new_room_membership)
                                                    .into(room_memberships::table)
                                                    .get_result(connection)
                                                    .map_err(ApiError::from)?;
            Ok(room_membership)
        }).map_err(ApiError::from)
    }

    /// Return `RoomMembership` for given `RoomId` and `UserId`.
    pub fn find(connection: &PgConnection, room_id: &RoomId, user_id: &UserId)
    -> Result<Option<RoomMembership>, ApiError> {
        let membership = room_memberships::table
            .filter(room_memberships::room_id.eq(room_id))
            .filter(room_memberships::user_id.eq(user_id))
            .first(connection);

        match membership {
            Ok(membership) => Ok(Some(membership)),
            Err(DieselError::NotFound) => Ok(None),
            Err(err) => Err(ApiError::from(err)),
        }
    }

    /// Return `RoomMembership`'s for given `UserId`.
    pub fn find_by_uid(connection: &PgConnection, user_id: UserId) -> Result<Vec<RoomMembership>, ApiError> {
        let room_memberships: Vec<RoomMembership> = room_memberships::table
            .filter(room_memberships::user_id.eq(user_id))
            .get_results(connection)
            .map_err(|err| match err {
                DieselError::NotFound => ApiError::not_found(Some(err.description())),
                _ => ApiError::from(err),
            })?;

        Ok(room_memberships)
    }

    /// Update an existing `RoomMembership` entry or insert a new one.
    pub fn upsert(connection: &PgConnection, domain: &str, options: RoomMembershipOptions)
    -> Result<RoomMembership, ApiError> {
        let room_membership = RoomMembership::find(
            connection,
            &options.room_id,
            &options.user_id
        )?;

        match room_membership {
            Some(mut entry) => entry.update(connection, domain, options),
            None => RoomMembership::create(connection, domain, options)
        }
    }

    /// Update a `RoomMembership` entry using new `RoomMembershipOptions`.
    ///
    /// After the update a new `MemberEvent` is created.
    pub fn update(&mut self, connection: &PgConnection, homeserver_domain: &str, options: RoomMembershipOptions)
    -> Result<RoomMembership, ApiError> {
        let profile = Profile::find_by_uid(connection, options.user_id.clone())?;

        let event = RoomMembership::create_new_room_member_event(
            &homeserver_domain,
            &options,
            profile,
        )?;

        self.membership = options.membership.clone();
        self.sender = options.sender.clone();

        connection.transaction::<RoomMembership, ApiError, _>(|| {
            insert(&event)
                .into(events::table)
                .execute(connection)
                .map_err(ApiError::from)?;

            self.save_changes::<RoomMembership>(connection)
                .map_err(ApiError::from)?;

            // Use the new `EventId` as primary key.
            update(room_memberships::table.find(self.event_id.clone()))
                .set(room_memberships::event_id.eq(event.id.clone()))
                .get_result(connection)
                .map_err(ApiError::from)
        }).map_err(ApiError::from)
    }

    /// Create a new `MemberEvent`.
    pub fn create_new_room_member_event(
        homeserver_domain: &str,
        options: &RoomMembershipOptions,
        profile: Option<Profile>
    ) -> Result<NewEvent, ApiError> {
        let event_id = EventId::new(&homeserver_domain).map_err(ApiError::from)?;
        let membership_string = Value::String(options.membership.clone());
        let membership: MembershipState = from_value(membership_string)?;

        let (avatar_url, displayname) = match profile {
            Some(profile) => (profile.avatar_url, profile.displayname),
            None => (None, None),
        };

        let new_member_event: NewEvent = MemberEvent {
            content: MemberEventContent {
                avatar_url: avatar_url,
                displayname: displayname,
                membership: membership,
                third_party_invite: (),
            },
            event_id: event_id.clone(),
            event_type: EventType::RoomMember,
            extra_content: MemberEventExtraContent { invite_room_state: None },
            prev_content: None,
            room_id: options.room_id.clone(),
            state_key: format!("@{}:{}", options.user_id.clone(), &homeserver_domain),
            unsigned: None,
            user_id: options.user_id.clone(),
        }.try_into()?;

        Ok(new_member_event)
    }

    /// Return member event's for given `room_id`.
    pub fn get_events_by_room(connection: &PgConnection, room_id: RoomId) -> Result<Vec<MemberEvent>, ApiError> {
        let event_ids = room_memberships::table
            .filter(room_memberships::room_id.eq(room_id))
            .select(room_memberships::event_id);

        let events: Vec<Event> = events::table
            .filter(events::id.eq(any(event_ids)))
            .get_results(connection)
            .map_err(|err| match err {
                DieselError::NotFound => ApiError::not_found(None),
                _ => ApiError::from(err),
            })?;

        let member_events: Result<Vec<MemberEvent>, SerdeJsonError> = events.into_iter()
                                                                        .map(TryInto::try_into)
                                                                        .collect();
        member_events.map_err(ApiError::from)
    }
}
<|endoftext|>


<commit_before>use super::Compressor;
use super::{footer::DocStoreFooter, index::SkipIndex};
use crate::directory::{FileSlice, OwnedBytes};
use crate::schema::Document;
use crate::space_usage::StoreSpaceUsage;
use crate::store::index::Checkpoint;
use crate::DocId;
use crate::{
    common::{BinarySerializable, HasLen, VInt},
    error::DataCorruption,
    fastfield::DeleteBitSet,
};
use lru::LruCache;
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

const LRU_CACHE_CAPACITY: usize = 100;

type Block = OwnedBytes;

type BlockCache = Arc<Mutex<LruCache<usize, Block>>>;

/// Reads document off tantivy's [`Store`](./index.html)
pub struct StoreReader {
    compressor: Compressor,
    data: FileSlice,
    cache: BlockCache,
    cache_hits: Arc<AtomicUsize>,
    cache_misses: Arc<AtomicUsize>,
    skip_index: Arc<SkipIndex>,
    space_usage: StoreSpaceUsage,
}

impl StoreReader {
    /// Opens a store reader
    pub fn open(store_file: FileSlice) -> io::Result<StoreReader> {
        let (footer, data_and_offset) = DocStoreFooter::extract_footer(store_file)?;

        let (data_file, offset_index_file) = data_and_offset.split(footer.offset as usize);
        let index_data = offset_index_file.read_bytes()?;
        let space_usage = StoreSpaceUsage::new(data_file.len(), offset_index_file.len());
        let skip_index = SkipIndex::open(index_data);
        Ok(StoreReader {
            compressor: footer.compressor,
            data: data_file,
            cache: Arc::new(Mutex::new(LruCache::new(LRU_CACHE_CAPACITY))),
            cache_hits: Default::default(),
            cache_misses: Default::default(),
            skip_index: Arc::new(skip_index),
            space_usage,
        })
    }

    pub(crate) fn block_checkpoints(&self) -> impl Iterator<Item = Checkpoint> + '_ {
        self.skip_index.checkpoints()
    }

    pub(crate) fn compressor(&self) -> Compressor {
        self.compressor
    }

    fn block_checkpoint(&self, doc_id: DocId) -> Option<Checkpoint> {
        self.skip_index.seek(doc_id)
    }

    pub(crate) fn block_data(&self) -> io::Result<OwnedBytes> {
        self.data.read_bytes()
    }

    fn compressed_block(&self, checkpoint: &Checkpoint) -> io::Result<OwnedBytes> {
        self.data.slice(checkpoint.byte_range.clone()).read_bytes()
    }

    fn read_block(&self, checkpoint: &Checkpoint) -> io::Result<Block> {
        if let Some(block) = self.cache.lock().unwrap().get(&checkpoint.byte_range.start) {
            self.cache_hits.fetch_add(1, Ordering::SeqCst);
            return Ok(block.clone());
        }

        self.cache_misses.fetch_add(1, Ordering::SeqCst);

        let compressed_block = self.compressed_block(checkpoint)?;
        let mut decompressed_block = vec![];
        self.compressor
            .decompress(compressed_block.as_slice(), &mut decompressed_block)?;

        let block = OwnedBytes::new(decompressed_block);
        self.cache
            .lock()
            .unwrap()
            .put(checkpoint.byte_range.start, block.clone());

        Ok(block)
    }

    /// Reads a given document.
    ///
    /// Calling `.get(doc)` is relatively costly as it requires
    /// decompressing a compressed block. The store utilizes a LRU cache,
    /// so accessing docs from the same compressed block should be faster.
    /// For that reason a store reader should be kept and reused.
    ///
    /// It should not be called to score documents
    /// for instance.
    pub fn get(&self, doc_id: DocId) -> crate::Result<Document> {
        let mut doc_bytes = self.get_document_bytes(doc_id)?;
        Ok(Document::deserialize(&mut doc_bytes)?)
    }

    /// Reads raw bytes of a given document. Returns `RawDocument`, which contains the block of a document and its start and end
    /// position within the block.
    ///
    /// Calling `.get(doc)` is relatively costly as it requires
    /// decompressing a compressed block. The store utilizes a LRU cache,
    /// so accessing docs from the same compressed block should be faster.
    /// For that reason a store reader should be kept and reused.
    ///
    pub fn get_document_bytes(&self, doc_id: DocId) -> crate::Result<OwnedBytes> {
        let checkpoint = self.block_checkpoint(doc_id).ok_or_else(|| {
            crate::TantivyError::InvalidArgument(format!("Failed to lookup Doc #{}.", doc_id))
        })?;
        let block = self.read_block(&checkpoint)?;
        let mut cursor = &block[..];
        let cursor_len_before = cursor.len();
        for _ in checkpoint.doc_range.start..doc_id {
            let doc_length = VInt::deserialize(&mut cursor)?.val() as usize;
            cursor = &cursor[doc_length..];
        }

        let doc_length = VInt::deserialize(&mut cursor)?.val() as usize;
        let start_pos = cursor_len_before - cursor.len();
        let end_pos = cursor_len_before - cursor.len() + doc_length;
        Ok(block.slice(start_pos..end_pos))
    }

    /// Iterator over all Documents in their order as they are stored in the doc store.
    /// Use this, if you want to extract all Documents from the doc store.
    /// The delete_bitset has to be forwarded from the `SegmentReader` or the results maybe wrong.
    pub fn iter<'a: 'b, 'b>(
        &'b self,
        delete_bitset: Option<&'a DeleteBitSet>,
    ) -> impl Iterator<Item = crate::Result<Document>> + 'b {
        self.iter_raw(delete_bitset).map(|doc_bytes_res| {
            let mut doc_bytes = doc_bytes_res?;
            Ok(Document::deserialize(&mut doc_bytes)?)
        })
    }

    /// Iterator over all RawDocuments in their order as they are stored in the doc store.
    /// Use this, if you want to extract all Documents from the doc store.
    /// The delete_bitset has to be forwarded from the `SegmentReader` or the results maybe wrong.
    pub(crate) fn iter_raw<'a: 'b, 'b>(
        &'b self,
        delete_bitset: Option<&'a DeleteBitSet>,
    ) -> impl Iterator<Item = crate::Result<OwnedBytes>> + 'b {
        let last_docid = self
            .block_checkpoints()
            .last()
            .map(|checkpoint| checkpoint.doc_range.end)
            .unwrap_or(0);
        let mut checkpoint_block_iter = self.block_checkpoints();
        let mut curr_checkpoint = checkpoint_block_iter.next();
        let mut curr_block = curr_checkpoint
            .as_ref()
            .map(|checkpoint| self.read_block(&checkpoint).map_err(|e| e.kind())); // map error in order to enable cloning
        let mut block_start_pos = 0;
        let mut num_skipped = 0;
        (0..last_docid)
            .filter_map(move |doc_id| {
                // filter_map is only used to resolve lifetime issues between the two closures on
                // the outer variables
                let alive = delete_bitset.map_or(true, |bitset| bitset.is_alive(doc_id));
                if !alive {
                    // we keep the number of skipped documents to move forward in the map block
                    num_skipped += 1;
                }
                // check move to next checkpoint
                let mut reset_block_pos = false;
                if doc_id >= curr_checkpoint.as_ref().unwrap().doc_range.end {
                    curr_checkpoint = checkpoint_block_iter.next();
                    curr_block = curr_checkpoint
                        .as_ref()
                        .map(|checkpoint| self.read_block(&checkpoint).map_err(|e| e.kind()));
                    reset_block_pos = true;
                    num_skipped = 0;
                }

                if alive {
                    let ret = Some((curr_block.clone(), num_skipped, reset_block_pos));
                    // the map block will move over the num_skipped, so we reset to 0
                    num_skipped = 0;
                    ret
                } else {
                    None
                }
            })
            .map(move |(block, num_skipped, reset_block_pos)| {
                let block = block
                    .ok_or_else(|| {
                        DataCorruption::comment_only(
                            "the current checkpoint in the doc store iterator is none, this should never happen",
                        )
                    })?
                    .map_err(|error_kind| {
                        std::io::Error::new(error_kind, "error when reading block in doc store")
                    })?;
                // this flag is set, when filter_map moved to the next block
                if reset_block_pos {
                    block_start_pos = 0;
                }
                let mut cursor = &block[block_start_pos..];
                let mut pos = 0;
                // move forward 1 doc + num_skipped in block and return length of current doc
                let doc_length = loop {
                    let doc_length = VInt::deserialize(&mut cursor)?.val() as usize;
                    let num_bytes_read = block[block_start_pos..].len() - cursor.len();
                    block_start_pos += num_bytes_read;

                    pos += 1;
                    if pos == num_skipped + 1 {
                        break doc_length;
                    } else {
                        block_start_pos += doc_length;
                        cursor = &block[block_start_pos..];
                    }
                };
                let end_pos = block_start_pos + doc_length;
                let doc_bytes = block.slice(block_start_pos..end_pos);
                block_start_pos = end_pos;
                Ok(doc_bytes)
            })
    }

    /// Summarize total space usage of this store reader.
    pub fn space_usage(&self) -> StoreSpaceUsage {
        self.space_usage.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::Document;
    use crate::schema::Field;
    use crate::{directory::RamDirectory, store::tests::write_lorem_ipsum_store, Directory};
    use std::path::Path;

    fn get_text_field<'a>(doc: &'a Document, field: &'a Field) -> Option<&'a str> {
        doc.get_first(*field).and_then(|f| f.text())
    }

    #[test]
    fn test_store_lru_cache() -> crate::Result<()> {
        let directory = RamDirectory::create();
        let path = Path::new("store");
        let writer = directory.open_write(path)?;
        let schema = write_lorem_ipsum_store(writer, 500, Compressor::default());
        let title = schema.get_field("title").unwrap();
        let store_file = directory.open_read(path)?;
        let store = StoreReader::open(store_file)?;

        assert_eq!(store.cache.lock().unwrap().len(), 0);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 0);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 0);

        let doc = store.get(0)?;
        assert_eq!(get_text_field(&doc, &title), Some("Doc 0"));

        assert_eq!(store.cache.lock().unwrap().len(), 1);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 0);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 1);
        assert_eq!(
            store
                .cache
                .lock()
                .unwrap()
                .peek_lru()
                .map(|(&k, _)| k as usize),
            Some(0)
        );

        let doc = store.get(499)?;
        assert_eq!(get_text_field(&doc, &title), Some("Doc 499"));

        assert_eq!(store.cache.lock().unwrap().len(), 2);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 0);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 2);

        assert_eq!(
            store
                .cache
                .lock()
                .unwrap()
                .peek_lru()
                .map(|(&k, _)| k as usize),
            Some(0)
        );

        let doc = store.get(0)?;
        assert_eq!(get_text_field(&doc, &title), Some("Doc 0"));

        assert_eq!(store.cache.lock().unwrap().len(), 2);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 1);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 2);
        assert_eq!(
            store
                .cache
                .lock()
                .unwrap()
                .peek_lru()
                .map(|(&k, _)| k as usize),
            Some(9210)
        );

        Ok(())
    }
}
<commit_msg>fix store reader iterator<commit_after>use super::Compressor;
use super::{footer::DocStoreFooter, index::SkipIndex};
use crate::directory::{FileSlice, OwnedBytes};
use crate::schema::Document;
use crate::space_usage::StoreSpaceUsage;
use crate::store::index::Checkpoint;
use crate::DocId;
use crate::{
    common::{BinarySerializable, HasLen, VInt},
    error::DataCorruption,
    fastfield::DeleteBitSet,
};
use lru::LruCache;
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

const LRU_CACHE_CAPACITY: usize = 100;

type Block = OwnedBytes;

type BlockCache = Arc<Mutex<LruCache<usize, Block>>>;

/// Reads document off tantivy's [`Store`](./index.html)
pub struct StoreReader {
    compressor: Compressor,
    data: FileSlice,
    cache: BlockCache,
    cache_hits: Arc<AtomicUsize>,
    cache_misses: Arc<AtomicUsize>,
    skip_index: Arc<SkipIndex>,
    space_usage: StoreSpaceUsage,
}

impl StoreReader {
    /// Opens a store reader
    pub fn open(store_file: FileSlice) -> io::Result<StoreReader> {
        let (footer, data_and_offset) = DocStoreFooter::extract_footer(store_file)?;

        let (data_file, offset_index_file) = data_and_offset.split(footer.offset as usize);
        let index_data = offset_index_file.read_bytes()?;
        let space_usage = StoreSpaceUsage::new(data_file.len(), offset_index_file.len());
        let skip_index = SkipIndex::open(index_data);
        Ok(StoreReader {
            compressor: footer.compressor,
            data: data_file,
            cache: Arc::new(Mutex::new(LruCache::new(LRU_CACHE_CAPACITY))),
            cache_hits: Default::default(),
            cache_misses: Default::default(),
            skip_index: Arc::new(skip_index),
            space_usage,
        })
    }

    pub(crate) fn block_checkpoints(&self) -> impl Iterator<Item = Checkpoint> + '_ {
        self.skip_index.checkpoints()
    }

    pub(crate) fn compressor(&self) -> Compressor {
        self.compressor
    }

    fn block_checkpoint(&self, doc_id: DocId) -> Option<Checkpoint> {
        self.skip_index.seek(doc_id)
    }

    pub(crate) fn block_data(&self) -> io::Result<OwnedBytes> {
        self.data.read_bytes()
    }

    fn compressed_block(&self, checkpoint: &Checkpoint) -> io::Result<OwnedBytes> {
        self.data.slice(checkpoint.byte_range.clone()).read_bytes()
    }

    fn read_block(&self, checkpoint: &Checkpoint) -> io::Result<Block> {
        if let Some(block) = self.cache.lock().unwrap().get(&checkpoint.byte_range.start) {
            self.cache_hits.fetch_add(1, Ordering::SeqCst);
            return Ok(block.clone());
        }

        self.cache_misses.fetch_add(1, Ordering::SeqCst);

        let compressed_block = self.compressed_block(checkpoint)?;
        let mut decompressed_block = vec![];
        self.compressor
            .decompress(compressed_block.as_slice(), &mut decompressed_block)?;

        let block = OwnedBytes::new(decompressed_block);
        self.cache
            .lock()
            .unwrap()
            .put(checkpoint.byte_range.start, block.clone());

        Ok(block)
    }

    /// Reads a given document.
    ///
    /// Calling `.get(doc)` is relatively costly as it requires
    /// decompressing a compressed block. The store utilizes a LRU cache,
    /// so accessing docs from the same compressed block should be faster.
    /// For that reason a store reader should be kept and reused.
    ///
    /// It should not be called to score documents
    /// for instance.
    pub fn get(&self, doc_id: DocId) -> crate::Result<Document> {
        let mut doc_bytes = self.get_document_bytes(doc_id)?;
        Ok(Document::deserialize(&mut doc_bytes)?)
    }

    /// Reads raw bytes of a given document. Returns `RawDocument`, which contains the block of a document and its start and end
    /// position within the block.
    ///
    /// Calling `.get(doc)` is relatively costly as it requires
    /// decompressing a compressed block. The store utilizes a LRU cache,
    /// so accessing docs from the same compressed block should be faster.
    /// For that reason a store reader should be kept and reused.
    ///
    pub fn get_document_bytes(&self, doc_id: DocId) -> crate::Result<OwnedBytes> {
        let checkpoint = self.block_checkpoint(doc_id).ok_or_else(|| {
            crate::TantivyError::InvalidArgument(format!("Failed to lookup Doc #{}.", doc_id))
        })?;
        let block = self.read_block(&checkpoint)?;
        let mut cursor = &block[..];
        let cursor_len_before = cursor.len();
        for _ in checkpoint.doc_range.start..doc_id {
            let doc_length = VInt::deserialize(&mut cursor)?.val() as usize;
            cursor = &cursor[doc_length..];
        }

        let doc_length = VInt::deserialize(&mut cursor)?.val() as usize;
        let start_pos = cursor_len_before - cursor.len();
        let end_pos = cursor_len_before - cursor.len() + doc_length;
        Ok(block.slice(start_pos..end_pos))
    }

    /// Iterator over all Documents in their order as they are stored in the doc store.
    /// Use this, if you want to extract all Documents from the doc store.
    /// The delete_bitset has to be forwarded from the `SegmentReader` or the results maybe wrong.
    pub fn iter<'a: 'b, 'b>(
        &'b self,
        delete_bitset: Option<&'a DeleteBitSet>,
    ) -> impl Iterator<Item = crate::Result<Document>> + 'b {
        self.iter_raw(delete_bitset).map(|doc_bytes_res| {
            let mut doc_bytes = doc_bytes_res?;
            Ok(Document::deserialize(&mut doc_bytes)?)
        })
    }

    /// Iterator over all RawDocuments in their order as they are stored in the doc store.
    /// Use this, if you want to extract all Documents from the doc store.
    /// The delete_bitset has to be forwarded from the `SegmentReader` or the results maybe wrong.
    pub(crate) fn iter_raw<'a: 'b, 'b>(
        &'b self,
        delete_bitset: Option<&'a DeleteBitSet>,
    ) -> impl Iterator<Item = crate::Result<OwnedBytes>> + 'b {
        let last_docid = self
            .block_checkpoints()
            .last()
            .map(|checkpoint| checkpoint.doc_range.end)
            .unwrap_or(0);
        let mut checkpoint_block_iter = self.block_checkpoints();
        let mut curr_checkpoint = checkpoint_block_iter.next();
        let mut curr_block = curr_checkpoint
            .as_ref()
            .map(|checkpoint| self.read_block(&checkpoint).map_err(|e| e.kind())); // map error in order to enable cloning
        let mut block_start_pos = 0;
        let mut num_skipped = 0;
        let mut reset_block_pos = false;
        (0..last_docid)
            .filter_map(move |doc_id| {
                // filter_map is only used to resolve lifetime issues between the two closures on
                // the outer variables
                let alive = delete_bitset.map_or(true, |bitset| bitset.is_alive(doc_id));
                if !alive {
                    // we keep the number of skipped documents to move forward in the map block
                    num_skipped += 1;
                }

                // check move to next checkpoint
                if doc_id >= curr_checkpoint.as_ref().unwrap().doc_range.end {
                    curr_checkpoint = checkpoint_block_iter.next();
                    curr_block = curr_checkpoint
                        .as_ref()
                        .map(|checkpoint| self.read_block(&checkpoint).map_err(|e| e.kind()));
                    reset_block_pos = true;
                    num_skipped = 0;
                }

                if alive {
                    let ret = Some((curr_block.clone(), num_skipped, reset_block_pos));
                    // the map block will move over the num_skipped, so we reset to 0
                    num_skipped = 0;
                    reset_block_pos = false;
                    ret
                } else {
                    None
                }
            })
            .map(move |(block, num_skipped, reset_block_pos)| {
                let block = block
                    .ok_or_else(|| {
                        DataCorruption::comment_only(
                            "the current checkpoint in the doc store iterator is none, this should never happen",
                        )
                    })?
                    .map_err(|error_kind| {
                        std::io::Error::new(error_kind, "error when reading block in doc store")
                    })?;
                // this flag is set, when filter_map moved to the next block
                if reset_block_pos {
                    block_start_pos = 0;
                }
                let mut cursor = &block[block_start_pos..];
                let mut pos = 0;
                // move forward 1 doc + num_skipped in block and return length of current doc
                let doc_length = loop {
                    let doc_length = VInt::deserialize(&mut cursor)?.val() as usize;
                    let num_bytes_read = block[block_start_pos..].len() - cursor.len();
                    block_start_pos += num_bytes_read;

                    pos += 1;
                    if pos == num_skipped + 1 {
                        break doc_length;
                    } else {
                        block_start_pos += doc_length;
                        if block_start_pos > block[0..].len() {
                            println!("Hmm");
                        }

                        cursor = &block[block_start_pos..];
                    }
                };
                let end_pos = block_start_pos + doc_length;
                let doc_bytes = block.slice(block_start_pos..end_pos);
                block_start_pos = end_pos;
                Ok(doc_bytes)
            })
    }

    /// Summarize total space usage of this store reader.
    pub fn space_usage(&self) -> StoreSpaceUsage {
        self.space_usage.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::Document;
    use crate::schema::Field;
    use crate::{directory::RamDirectory, store::tests::write_lorem_ipsum_store, Directory};
    use std::path::Path;

    fn get_text_field<'a>(doc: &'a Document, field: &'a Field) -> Option<&'a str> {
        doc.get_first(*field).and_then(|f| f.text())
    }

    #[test]
    fn test_store_lru_cache() -> crate::Result<()> {
        let directory = RamDirectory::create();
        let path = Path::new("store");
        let writer = directory.open_write(path)?;
        let schema = write_lorem_ipsum_store(writer, 500, Compressor::default());
        let title = schema.get_field("title").unwrap();
        let store_file = directory.open_read(path)?;
        let store = StoreReader::open(store_file)?;

        assert_eq!(store.cache.lock().unwrap().len(), 0);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 0);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 0);

        let doc = store.get(0)?;
        assert_eq!(get_text_field(&doc, &title), Some("Doc 0"));

        assert_eq!(store.cache.lock().unwrap().len(), 1);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 0);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 1);
        assert_eq!(
            store
                .cache
                .lock()
                .unwrap()
                .peek_lru()
                .map(|(&k, _)| k as usize),
            Some(0)
        );

        let doc = store.get(499)?;
        assert_eq!(get_text_field(&doc, &title), Some("Doc 499"));

        assert_eq!(store.cache.lock().unwrap().len(), 2);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 0);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 2);

        assert_eq!(
            store
                .cache
                .lock()
                .unwrap()
                .peek_lru()
                .map(|(&k, _)| k as usize),
            Some(0)
        );

        let doc = store.get(0)?;
        assert_eq!(get_text_field(&doc, &title), Some("Doc 0"));

        assert_eq!(store.cache.lock().unwrap().len(), 2);
        assert_eq!(store.cache_hits.load(Ordering::SeqCst), 1);
        assert_eq!(store.cache_misses.load(Ordering::SeqCst), 2);
        assert_eq!(
            store
                .cache
                .lock()
                .unwrap()
                .peek_lru()
                .map(|(&k, _)| k as usize),
            Some(9210)
        );

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