<commit_before>//! A control flow graph represented as mappings of extended basic blocks to their predecessors
//! and successors.
//!
//! Successors are represented as extended basic blocks while predecessors are represented by basic
//! blocks. Basic blocks are denoted by tuples of EBB and branch/jump instructions. Each
//! predecessor tuple corresponds to the end of a basic block.
//!
//! ```c
//!     Ebb0:
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         brz vx, Ebb1 ; end of basic block
//!
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         jmp Ebb2     ; end of basic block
//! ```
//!
//! Here `Ebb1` and `Ebb2` would each have a single predecessor denoted as `(Ebb0, brz)`
//! and `(Ebb0, jmp Ebb2)` respectively.

use ir::{Function, Inst, Ebb};
use ir::instructions::BranchInfo;
use entity::EntityMap;
use std::mem;

/// A basic block denoted by its enclosing Ebb and last instruction.
pub type BasicBlock = (Ebb, Inst);

/// A container for the successors and predecessors of some Ebb.
#[derive(Debug, Clone, Default)]
pub struct CFGNode {
    /// EBBs that are the targets of branches and jumps in this EBB.
    pub successors: Vec<Ebb>,
    /// Basic blocks that can branch or jump to this EBB.
    pub predecessors: Vec<BasicBlock>,
}

/// The Control Flow Graph maintains a mapping of ebbs to their predecessors
/// and successors where predecessors are basic blocks and successors are
/// extended basic blocks.
#[derive(Debug)]
pub struct ControlFlowGraph {
    entry_block: Option<Ebb>,
    data: EntityMap<Ebb, CFGNode>,
    valid: bool,
}

impl ControlFlowGraph {
    /// Allocate a new blank control flow graph.
    pub fn new() -> ControlFlowGraph {
        ControlFlowGraph {
            entry_block: None,
            data: EntityMap::new(),
            valid: false,
        }
    }

    /// Allocate and compute the control flow graph for `func`.
    pub fn with_function(func: &Function) -> ControlFlowGraph {
        let mut cfg = ControlFlowGraph::new();
        cfg.compute(func);
        cfg
    }

    /// Compute the control flow graph of `func`.
    ///
    /// This will clear and overwrite any information already stored in this data structure.
    pub fn compute(&mut self, func: &Function) {
        self.entry_block = func.layout.entry_block();
        self.data.clear();
        self.data.resize(func.dfg.num_ebbs());

        for ebb in &func.layout {
            self.compute_ebb(func, ebb);
        }

        self.valid = true;
    }

    fn compute_ebb(&mut self, func: &Function, ebb: Ebb) {
        for inst in func.layout.ebb_insts(ebb) {
            match func.dfg[inst].analyze_branch(&func.dfg.value_lists) {
                BranchInfo::SingleDest(dest, _) => {
                    self.add_edge((ebb, inst), dest);
                }
                BranchInfo::Table(jt) => {
                    for (_, dest) in func.jump_tables[jt].entries() {
                        self.add_edge((ebb, inst), dest);
                    }
                }
                BranchInfo::NotABranch => {}
            }
        }
    }

    fn invalidate_ebb_successors(&mut self, ebb: Ebb) {
        // Temporarily take ownership because we need mutable access to self.data inside the loop.
        // Unfortunately borrowck cannot see that our mut accesses to predecessors don't alias
        // our iteration over successors.
        let mut successors = mem::replace(&mut self.data[ebb].successors, Vec::new());
        for suc in successors.iter().cloned() {
            self.data[suc].predecessors.retain(|&(e, _)| e != ebb);
        }
        successors.clear();
        self.data[ebb].successors = successors;
    }

    /// Recompute the control flow graph of `ebb`.
    ///
    /// This is for use after modifying instructions within a specific EBB. It recomputes all edges
    /// from `ebb` while leaving edges to `ebb` intact. Its functionality a subset of that of the
    /// more expensive `compute`, and should be used when we know we don't need to recompute the CFG
    /// from scratch, but rather that our changes have been restricted to specific EBBs.
    pub fn recompute_ebb(&mut self, func: &Function, ebb: Ebb) {
        debug_assert!(self.is_valid());
        self.invalidate_ebb_successors(ebb);
        self.compute_ebb(func, ebb);
    }

    fn add_edge(&mut self, from: BasicBlock, to: Ebb) {
        self.data[from.0].successors.push(to);
        self.data[to].predecessors.push(from);
    }

    /// Get the CFG predecessor basic blocks to `ebb`.
    pub fn get_predecessors(&self, ebb: Ebb) -> &[BasicBlock] {
        debug_assert!(self.is_valid());
        &self.data[ebb].predecessors
    }

    /// Get the CFG successors to `ebb`.
    pub fn get_successors(&self, ebb: Ebb) -> &[Ebb] {
        debug_assert!(self.is_valid());
        &self.data[ebb].successors
    }

    /// Check if the CFG is in a valid state.
    ///
    /// Note that this doesn't perform any kind of validity checks. It simply checks if the
    /// `compute()` method has been called since the last `clear()`. It does not check that the
    /// CFG is consistent with the function.
    pub fn is_valid(&self) -> bool {
        self.valid
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cursor::{Cursor, FuncCursor};
    use ir::{Function, InstBuilder, types};

    #[test]
    fn empty() {
        let func = Function::new();
        ControlFlowGraph::with_function(&func);
    }

    #[test]
    fn no_predecessors() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let cfg = ControlFlowGraph::with_function(&func);

        let mut fun_ebbs = func.layout.ebbs();
        for ebb in func.layout.ebbs() {
            assert_eq!(ebb, fun_ebbs.next().unwrap());
            assert_eq!(cfg.get_predecessors(ebb).len(), 0);
            assert_eq!(cfg.get_successors(ebb).len(), 0);
        }
    }

    #[test]
    fn branches_and_jumps() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let cond = func.dfg.append_ebb_param(ebb0, types::I32);
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();

        let br_ebb0_ebb2;
        let br_ebb1_ebb1;
        let jmp_ebb0_ebb1;
        let jmp_ebb1_ebb2;

        {
            let mut cur = FuncCursor::new(&mut func);

            cur.insert_ebb(ebb0);
            br_ebb0_ebb2 = cur.ins().brnz(cond, ebb2, &[]);
            jmp_ebb0_ebb1 = cur.ins().jump(ebb1, &[]);

            cur.insert_ebb(ebb1);
            br_ebb1_ebb1 = cur.ins().brnz(cond, ebb1, &[]);
            jmp_ebb1_ebb2 = cur.ins().jump(ebb2, &[]);

            cur.insert_ebb(ebb2);
        }

        let mut cfg = ControlFlowGraph::with_function(&func);

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 2);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, jmp_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 2);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), true);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }

        // Change some instructions and recompute ebb0
        func.dfg.replace(br_ebb0_ebb2).brnz(cond, ebb1, &[]);
        func.dfg.replace(jmp_ebb0_ebb1).return_(&[]);
        cfg.recompute_ebb(&mut func, ebb0);
        let br_ebb0_ebb1 = br_ebb0_ebb2;

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 1);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, br_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), false);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 1);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), false);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }
    }
}
<commit_msg>Remove the CFG entry_block member.<commit_after>//! A control flow graph represented as mappings of extended basic blocks to their predecessors
//! and successors.
//!
//! Successors are represented as extended basic blocks while predecessors are represented by basic
//! blocks. Basic blocks are denoted by tuples of EBB and branch/jump instructions. Each
//! predecessor tuple corresponds to the end of a basic block.
//!
//! ```c
//!     Ebb0:
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         brz vx, Ebb1 ; end of basic block
//!
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         jmp Ebb2     ; end of basic block
//! ```
//!
//! Here `Ebb1` and `Ebb2` would each have a single predecessor denoted as `(Ebb0, brz)`
//! and `(Ebb0, jmp Ebb2)` respectively.

use ir::{Function, Inst, Ebb};
use ir::instructions::BranchInfo;
use entity::EntityMap;
use std::mem;

/// A basic block denoted by its enclosing Ebb and last instruction.
pub type BasicBlock = (Ebb, Inst);

/// A container for the successors and predecessors of some Ebb.
#[derive(Debug, Clone, Default)]
pub struct CFGNode {
    /// EBBs that are the targets of branches and jumps in this EBB.
    pub successors: Vec<Ebb>,
    /// Basic blocks that can branch or jump to this EBB.
    pub predecessors: Vec<BasicBlock>,
}

/// The Control Flow Graph maintains a mapping of ebbs to their predecessors
/// and successors where predecessors are basic blocks and successors are
/// extended basic blocks.
#[derive(Debug)]
pub struct ControlFlowGraph {
    data: EntityMap<Ebb, CFGNode>,
    valid: bool,
}

impl ControlFlowGraph {
    /// Allocate a new blank control flow graph.
    pub fn new() -> ControlFlowGraph {
        ControlFlowGraph {
            data: EntityMap::new(),
            valid: false,
        }
    }

    /// Allocate and compute the control flow graph for `func`.
    pub fn with_function(func: &Function) -> ControlFlowGraph {
        let mut cfg = ControlFlowGraph::new();
        cfg.compute(func);
        cfg
    }

    /// Compute the control flow graph of `func`.
    ///
    /// This will clear and overwrite any information already stored in this data structure.
    pub fn compute(&mut self, func: &Function) {
        self.data.clear();
        self.data.resize(func.dfg.num_ebbs());

        for ebb in &func.layout {
            self.compute_ebb(func, ebb);
        }

        self.valid = true;
    }

    fn compute_ebb(&mut self, func: &Function, ebb: Ebb) {
        for inst in func.layout.ebb_insts(ebb) {
            match func.dfg[inst].analyze_branch(&func.dfg.value_lists) {
                BranchInfo::SingleDest(dest, _) => {
                    self.add_edge((ebb, inst), dest);
                }
                BranchInfo::Table(jt) => {
                    for (_, dest) in func.jump_tables[jt].entries() {
                        self.add_edge((ebb, inst), dest);
                    }
                }
                BranchInfo::NotABranch => {}
            }
        }
    }

    fn invalidate_ebb_successors(&mut self, ebb: Ebb) {
        // Temporarily take ownership because we need mutable access to self.data inside the loop.
        // Unfortunately borrowck cannot see that our mut accesses to predecessors don't alias
        // our iteration over successors.
        let mut successors = mem::replace(&mut self.data[ebb].successors, Vec::new());
        for suc in successors.iter().cloned() {
            self.data[suc].predecessors.retain(|&(e, _)| e != ebb);
        }
        successors.clear();
        self.data[ebb].successors = successors;
    }

    /// Recompute the control flow graph of `ebb`.
    ///
    /// This is for use after modifying instructions within a specific EBB. It recomputes all edges
    /// from `ebb` while leaving edges to `ebb` intact. Its functionality a subset of that of the
    /// more expensive `compute`, and should be used when we know we don't need to recompute the CFG
    /// from scratch, but rather that our changes have been restricted to specific EBBs.
    pub fn recompute_ebb(&mut self, func: &Function, ebb: Ebb) {
        debug_assert!(self.is_valid());
        self.invalidate_ebb_successors(ebb);
        self.compute_ebb(func, ebb);
    }

    fn add_edge(&mut self, from: BasicBlock, to: Ebb) {
        self.data[from.0].successors.push(to);
        self.data[to].predecessors.push(from);
    }

    /// Get the CFG predecessor basic blocks to `ebb`.
    pub fn get_predecessors(&self, ebb: Ebb) -> &[BasicBlock] {
        debug_assert!(self.is_valid());
        &self.data[ebb].predecessors
    }

    /// Get the CFG successors to `ebb`.
    pub fn get_successors(&self, ebb: Ebb) -> &[Ebb] {
        debug_assert!(self.is_valid());
        &self.data[ebb].successors
    }

    /// Check if the CFG is in a valid state.
    ///
    /// Note that this doesn't perform any kind of validity checks. It simply checks if the
    /// `compute()` method has been called since the last `clear()`. It does not check that the
    /// CFG is consistent with the function.
    pub fn is_valid(&self) -> bool {
        self.valid
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cursor::{Cursor, FuncCursor};
    use ir::{Function, InstBuilder, types};

    #[test]
    fn empty() {
        let func = Function::new();
        ControlFlowGraph::with_function(&func);
    }

    #[test]
    fn no_predecessors() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let cfg = ControlFlowGraph::with_function(&func);

        let mut fun_ebbs = func.layout.ebbs();
        for ebb in func.layout.ebbs() {
            assert_eq!(ebb, fun_ebbs.next().unwrap());
            assert_eq!(cfg.get_predecessors(ebb).len(), 0);
            assert_eq!(cfg.get_successors(ebb).len(), 0);
        }
    }

    #[test]
    fn branches_and_jumps() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let cond = func.dfg.append_ebb_param(ebb0, types::I32);
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();

        let br_ebb0_ebb2;
        let br_ebb1_ebb1;
        let jmp_ebb0_ebb1;
        let jmp_ebb1_ebb2;

        {
            let mut cur = FuncCursor::new(&mut func);

            cur.insert_ebb(ebb0);
            br_ebb0_ebb2 = cur.ins().brnz(cond, ebb2, &[]);
            jmp_ebb0_ebb1 = cur.ins().jump(ebb1, &[]);

            cur.insert_ebb(ebb1);
            br_ebb1_ebb1 = cur.ins().brnz(cond, ebb1, &[]);
            jmp_ebb1_ebb2 = cur.ins().jump(ebb2, &[]);

            cur.insert_ebb(ebb2);
        }

        let mut cfg = ControlFlowGraph::with_function(&func);

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 2);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, jmp_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 2);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), true);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }

        // Change some instructions and recompute ebb0
        func.dfg.replace(br_ebb0_ebb2).brnz(cond, ebb1, &[]);
        func.dfg.replace(jmp_ebb0_ebb1).return_(&[]);
        cfg.recompute_ebb(&mut func, ebb0);
        let br_ebb0_ebb1 = br_ebb0_ebb2;

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 1);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, br_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), false);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 1);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), false);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }
    }
}
<|endoftext|>


<commit_before>// A "shape" is a compact encoding of a type that is used by interpreted glue.
// This substitutes for the runtime tags used by e.g. MLs.

use lib::llvm::llvm;
use lib::llvm::{True, False, ModuleRef, TypeRef, ValueRef};
use driver::session;
use driver::session::session;
use trans::base;
use middle::trans::common::*;
use middle::trans::machine::*;
use back::abi;
use middle::ty;
use middle::ty::field;
use syntax::ast;
use syntax::ast_util::dummy_sp;
use syntax::util::interner;
use util::ppaux::ty_to_str;
use syntax::codemap::span;
use dvec::DVec;

use std::map::HashMap;
use option::is_some;

use ty_ctxt = middle::ty::ctxt;

type ctxt = {mut next_tag_id: u16, pad: u16, pad2: u32};

fn mk_global(ccx: @crate_ctxt, name: ~str, llval: ValueRef, internal: bool) ->
   ValueRef {
    let llglobal = do str::as_c_str(name) |buf| {
        lib::llvm::llvm::LLVMAddGlobal(ccx.llmod, val_ty(llval), buf)
    };
    lib::llvm::llvm::LLVMSetInitializer(llglobal, llval);
    lib::llvm::llvm::LLVMSetGlobalConstant(llglobal, True);

    if internal {
        lib::llvm::SetLinkage(llglobal, lib::llvm::InternalLinkage);
    }

    return llglobal;
}

fn mk_ctxt(llmod: ModuleRef) -> ctxt {
    let llshapetablesty = trans::common::T_named_struct(~"shapes");
    let _llshapetables = str::as_c_str(~"shapes", |buf| {
        lib::llvm::llvm::LLVMAddGlobal(llmod, llshapetablesty, buf)
    });

    return {mut next_tag_id: 0u16, pad: 0u16, pad2: 0u32};
}

fn add_u16(&dest: ~[u8], val: u16) {
    dest += ~[(val & 0xffu16) as u8, (val >> 8u16) as u8];
}

fn add_substr(&dest: ~[u8], src: ~[u8]) {
    add_u16(dest, vec::len(src) as u16);
    dest += src;
}

<commit_msg>De-mode mysterious unused functions in middle::trans::shape (see comment)<commit_after>// A "shape" is a compact encoding of a type that is used by interpreted glue.
// This substitutes for the runtime tags used by e.g. MLs.

use lib::llvm::llvm;
use lib::llvm::{True, False, ModuleRef, TypeRef, ValueRef};
use driver::session;
use driver::session::session;
use trans::base;
use middle::trans::common::*;
use middle::trans::machine::*;
use back::abi;
use middle::ty;
use middle::ty::field;
use syntax::ast;
use syntax::ast_util::dummy_sp;
use syntax::util::interner;
use util::ppaux::ty_to_str;
use syntax::codemap::span;
use dvec::DVec;

use std::map::HashMap;
use option::is_some;

use ty_ctxt = middle::ty::ctxt;

type ctxt = {mut next_tag_id: u16, pad: u16, pad2: u32};

fn mk_global(ccx: @crate_ctxt, name: ~str, llval: ValueRef, internal: bool) ->
   ValueRef {
    let llglobal = do str::as_c_str(name) |buf| {
        lib::llvm::llvm::LLVMAddGlobal(ccx.llmod, val_ty(llval), buf)
    };
    lib::llvm::llvm::LLVMSetInitializer(llglobal, llval);
    lib::llvm::llvm::LLVMSetGlobalConstant(llglobal, True);

    if internal {
        lib::llvm::SetLinkage(llglobal, lib::llvm::InternalLinkage);
    }

    return llglobal;
}

fn mk_ctxt(llmod: ModuleRef) -> ctxt {
    let llshapetablesty = trans::common::T_named_struct(~"shapes");
    let _llshapetables = str::as_c_str(~"shapes", |buf| {
        lib::llvm::llvm::LLVMAddGlobal(llmod, llshapetablesty, buf)
    });

    return {mut next_tag_id: 0u16, pad: 0u16, pad2: 0u32};
}

/*
Although these two functions are never called, they are here
for a VERY GOOD REASON. See #3670
*/
fn add_u16(dest: &mut ~[u8], val: u16) {
    *dest += ~[(val & 0xffu16) as u8, (val >> 8u16) as u8];
}

fn add_substr(dest: &mut ~[u8], src: ~[u8]) {
    add_u16(dest, vec::len(src) as u16);
    *dest += src;
}

<|endoftext|>


<commit_before>// Copyright Ethereum Classic Contributors
//
// 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.
// Rust Bitcoin Library
// Written in 2014 by
//     Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Big unsigned integer types
//!
//! Implementation of a various large-but-fixed sized unsigned integer types.
//! The functions here are designed to be fast.
//!

// #![no_std]

use ::std::convert::{From, Into};
use ::std::ops::{Add, Sub, Not, Mul, Div, Shr, Shl};
use ::std::cmp::Ordering;

#[repr(C)]
#[derive(Eq, PartialEq, Debug, Copy, Clone)]
pub struct U256([u64; 4]);

impl U256 {
    pub fn zero() -> U256 { U256([0; 4]) }
    pub fn one() -> U256 { U256([1u64, 0u64, 0u64, 0u64]) }

    pub fn overflowing_add(self, other: U256) -> (U256, bool) {
        let U256(ref me) = self;
        let U256(ref you) = other;

        let mut ret = [0u64; 4];
        let mut carry = false;
        for i in 0..4 {
            let (v, o1) = me[i].overflowing_add(you[i]);
            let (v, o2) = v.overflowing_add(if carry { 1 } else { 0 });
            ret[i] = v;
            carry = o1 || o2;
        }

        (U256(ret), carry)
    }

    pub fn low_u32(&self) -> u32 {
        let &U256(ref arr) = self;
        arr[0] as u32
    }

    pub fn mul_u32(self, other: u32) -> U256 {
        let U256(ref arr) = self;
        let mut carry = [0u64; 4];
        let mut ret = [0u64; 4];
        for i in 0..4 {
            let upper = other as u64 * (arr[i] >> 32);
            let lower = other as u64 * (arr[i] & 0xFFFFFFFF);
            if i < 3 {
                carry[i + 1] += upper >> 32;
            }
            ret[i] = lower + (upper << 32);
        }
        U256(ret) + U256(carry)
    }

    pub fn bits(&self) -> usize {
        let &U256(ref arr) = self;
        for i in 1..4 {
            if arr[4 - i] > 0 { return (0x40 * (4 - i + 1)) - arr[4 - i].leading_zeros() as usize; }
        }
        0x40 - arr[0].leading_zeros() as usize
    }
}

impl From<u64> for U256 {
    fn from(val: u64) -> U256 {
        U256([0, 0, 0, val])
    }
}

impl Into<u64> for U256 {
    fn into(self) -> u64 {
        assert!(self.0[0] == 0 && self.0[1] == 0 && self.0[2] == 0);
        self.0[3]
    }
}

impl From<usize> for U256 {
    fn from(val: usize) -> U256 {
        (val as u64).into()
    }
}

impl Into<usize> for U256 {
    fn into(self) -> usize {
        let v64: u64 = self.into();
        v64 as usize
    }
}

impl From<i32> for U256 {
    fn from(val: i32) -> U256 {
        (val as u64).into()
    }
}

impl<'a> From<&'a [u8]> for U256 {
    fn from(val: &'a [u8]) -> U256 {
        assert!(val.len() <= 256 / 8);
        let mut u256 = U256::zero();

        for i in 0..val.len() {
            let rev = val.len() - 1 - i;
            let pos = rev / 8;
            u256.0[pos] += (val[i] as u64) << ((rev % 8) * 8);
        }

        u256
    }
}

impl Not for U256 {
    type Output = U256;

    fn not(self) -> U256 {
        let U256(ref arr) = self;
        let mut ret = [0u64; 4];
        for i in 0..4 {
            ret[i] = !arr[i];
        }
        U256(ret)
    }
}

impl Add for U256 {
    type Output = U256;

    fn add(self, other: U256) -> U256 {
        let (o, v) = self.overflowing_add(other);
        assert!(v == false);
        o
    }
}

impl Sub for U256 {
    type Output = U256;

    #[inline]
    fn sub(self, other: U256) -> U256 {
        let (o, v) = self.overflowing_add(!other);
        assert!(v == true);
        o + U256::one()
    }
}

impl Mul for U256 {
    type Output = U256;

    fn mul(self, other: U256) -> U256 {
        let mut me = self;
        // TODO: be more efficient about this
        for i in 0..(2 * 4) {
            me = (me + me.mul_u32((other >> (32 * i)).low_u32())) << (32 * i);
        }
        me
    }
}

impl Div for U256 {
    type Output = U256;

    fn div(self, other: U256) -> U256 {
        let mut sub_copy = self;
        let mut shift_copy = other;
        let mut ret = [0u64; 4];

        let my_bits = self.bits();
        let your_bits = other.bits();

        // Check for division by 0
        assert!(your_bits != 0);

        // Early return in case we are dividing by a larger number than us
        if my_bits < your_bits {
            return U256(ret);
        }

        // Bitwise long division
        let mut shift = my_bits - your_bits;
        shift_copy = shift_copy << shift;
        loop {
            if sub_copy >= shift_copy {
                ret[shift / 64] |= 1 << (shift % 64);
                sub_copy = sub_copy - shift_copy;
            }
            shift_copy = shift_copy >> 1;
            if shift == 0 { break; }
            shift -= 1;
        }

        U256(ret)
    }
}

impl Ord for U256 {
    fn cmp(&self, other: &U256) -> Ordering {
	let &U256(ref me) = self;
	let &U256(ref you) = other;
	let mut i = 4;
	while i > 0 {
	    i -= 1;
	    if me[i] < you[i] { return Ordering::Less; }
	    if me[i] > you[i] { return Ordering::Greater; }
	}
	Ordering::Equal
    }
}

impl PartialOrd for U256 {
    fn partial_cmp(&self, other: &U256) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Shl<usize> for U256 {
    type Output = U256;

    fn shl(self, shift: usize) -> U256 {
        let U256(ref original) = self;
        let mut ret = [0u64; 4];
        let word_shift = shift / 64;
        let bit_shift = shift % 64;
        for i in 0..4 {
            // Shift
            if bit_shift < 64 && i + word_shift < 4 {
                ret[i + word_shift] += original[i] << bit_shift;
            }
            // Carry
            if bit_shift > 0 && i + word_shift + 1 < 4 {
                ret[i + word_shift + 1] += original[i] >> (64 - bit_shift);
            }
        }
        U256(ret)
    }
}

impl Shr<usize> for U256 {
    type Output = U256;

    fn shr(self, shift: usize) -> U256 {
        let U256(ref original) = self;
        let mut ret = [0u64; 4];
        let word_shift = shift / 64;
        let bit_shift = shift % 64;
        for i in word_shift..4 {
            // Shift
            ret[i - word_shift] += original[i] >> bit_shift;
            // Carry
            if bit_shift > 0 && i < 4 - 1 {
                ret[i - word_shift] += original[i + 1] << (64 - bit_shift);
            }
        }
        U256(ret)
    }
}

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

    #[test]
    fn u256_add() {
        assert_eq!(
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64]) +
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64]),
            U256([0xfffffffffffffffeu64, 1u64, 0u64, 0u64])
        );
    }

    #[test]
    fn u256_sub() {
        assert_eq!(
            U256([0xfffffffffffffffeu64, 1u64, 0u64, 0u64]) -
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64]),
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64])
        );
    }
}
<commit_msg>Problem: no log2floor implementation Solution: implement it!<commit_after>// Copyright Ethereum Classic Contributors
//
// 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.
// Rust Bitcoin Library
// Written in 2014 by
//     Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Big unsigned integer types
//!
//! Implementation of a various large-but-fixed sized unsigned integer types.
//! The functions here are designed to be fast.
//!

// #![no_std]

use ::std::convert::{From, Into};
use ::std::ops::{Add, Sub, Not, Mul, Div, Shr, Shl};
use ::std::cmp::Ordering;

#[repr(C)]
#[derive(Eq, PartialEq, Debug, Copy, Clone)]
pub struct U256([u64; 4]);

impl U256 {
    pub fn zero() -> U256 { U256([0; 4]) }
    pub fn one() -> U256 { U256([1u64, 0u64, 0u64, 0u64]) }

    pub fn overflowing_add(self, other: U256) -> (U256, bool) {
        let U256(ref me) = self;
        let U256(ref you) = other;

        let mut ret = [0u64; 4];
        let mut carry = false;
        for i in 0..4 {
            let (v, o1) = me[i].overflowing_add(you[i]);
            let (v, o2) = v.overflowing_add(if carry { 1 } else { 0 });
            ret[i] = v;
            carry = o1 || o2;
        }

        (U256(ret), carry)
    }

    pub fn low_u32(&self) -> u32 {
        let &U256(ref arr) = self;
        arr[0] as u32
    }

    pub fn mul_u32(self, other: u32) -> U256 {
        let U256(ref arr) = self;
        let mut carry = [0u64; 4];
        let mut ret = [0u64; 4];
        for i in 0..4 {
            let upper = other as u64 * (arr[i] >> 32);
            let lower = other as u64 * (arr[i] & 0xFFFFFFFF);
            if i < 3 {
                carry[i + 1] += upper >> 32;
            }
            ret[i] = lower + (upper << 32);
        }
        U256(ret) + U256(carry)
    }

    pub fn bits(&self) -> usize {
        let &U256(ref arr) = self;
        for i in 1..4 {
            if arr[4 - i] > 0 { return (0x40 * (4 - i + 1)) - arr[4 - i].leading_zeros() as usize; }
        }
        0x40 - arr[0].leading_zeros() as usize
    }

    pub fn log2floor(&self) -> isize {
        assert!(*self != U256::zero());
        let mut l: isize = 256;
        for i in 0..4 {
            if self.0[3 - i] == 0u64 {
                l -= 64;
            } else {
                l -= self.0[3 - i].leading_zeros() as isize;
                return l;
            }
        }
        return l;
    }
}

impl From<u64> for U256 {
    fn from(val: u64) -> U256 {
        U256([0, 0, 0, val])
    }
}

impl Into<u64> for U256 {
    fn into(self) -> u64 {
        assert!(self.0[0] == 0 && self.0[1] == 0 && self.0[2] == 0);
        self.0[3]
    }
}

impl From<usize> for U256 {
    fn from(val: usize) -> U256 {
        (val as u64).into()
    }
}

impl Into<usize> for U256 {
    fn into(self) -> usize {
        let v64: u64 = self.into();
        v64 as usize
    }
}

impl From<i32> for U256 {
    fn from(val: i32) -> U256 {
        (val as u64).into()
    }
}

impl<'a> From<&'a [u8]> for U256 {
    fn from(val: &'a [u8]) -> U256 {
        assert!(val.len() <= 256 / 8);
        let mut u256 = U256::zero();

        for i in 0..val.len() {
            let rev = val.len() - 1 - i;
            let pos = rev / 8;
            u256.0[pos] += (val[i] as u64) << ((rev % 8) * 8);
        }

        u256
    }
}

impl Not for U256 {
    type Output = U256;

    fn not(self) -> U256 {
        let U256(ref arr) = self;
        let mut ret = [0u64; 4];
        for i in 0..4 {
            ret[i] = !arr[i];
        }
        U256(ret)
    }
}

impl Add for U256 {
    type Output = U256;

    fn add(self, other: U256) -> U256 {
        let (o, v) = self.overflowing_add(other);
        assert!(v == false);
        o
    }
}

impl Sub for U256 {
    type Output = U256;

    #[inline]
    fn sub(self, other: U256) -> U256 {
        let (o, v) = self.overflowing_add(!other);
        assert!(v == true);
        o + U256::one()
    }
}

impl Mul for U256 {
    type Output = U256;

    fn mul(self, other: U256) -> U256 {
        let mut me = self;
        // TODO: be more efficient about this
        for i in 0..(2 * 4) {
            me = (me + me.mul_u32((other >> (32 * i)).low_u32())) << (32 * i);
        }
        me
    }
}

impl Div for U256 {
    type Output = U256;

    fn div(self, other: U256) -> U256 {
        let mut sub_copy = self;
        let mut shift_copy = other;
        let mut ret = [0u64; 4];

        let my_bits = self.bits();
        let your_bits = other.bits();

        // Check for division by 0
        assert!(your_bits != 0);

        // Early return in case we are dividing by a larger number than us
        if my_bits < your_bits {
            return U256(ret);
        }

        // Bitwise long division
        let mut shift = my_bits - your_bits;
        shift_copy = shift_copy << shift;
        loop {
            if sub_copy >= shift_copy {
                ret[shift / 64] |= 1 << (shift % 64);
                sub_copy = sub_copy - shift_copy;
            }
            shift_copy = shift_copy >> 1;
            if shift == 0 { break; }
            shift -= 1;
        }

        U256(ret)
    }
}

impl Ord for U256 {
    fn cmp(&self, other: &U256) -> Ordering {
	let &U256(ref me) = self;
	let &U256(ref you) = other;
	let mut i = 4;
	while i > 0 {
	    i -= 1;
	    if me[i] < you[i] { return Ordering::Less; }
	    if me[i] > you[i] { return Ordering::Greater; }
	}
	Ordering::Equal
    }
}

impl PartialOrd for U256 {
    fn partial_cmp(&self, other: &U256) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Shl<usize> for U256 {
    type Output = U256;

    fn shl(self, shift: usize) -> U256 {
        let U256(ref original) = self;
        let mut ret = [0u64; 4];
        let word_shift = shift / 64;
        let bit_shift = shift % 64;
        for i in 0..4 {
            // Shift
            if bit_shift < 64 && i + word_shift < 4 {
                ret[i + word_shift] += original[i] << bit_shift;
            }
            // Carry
            if bit_shift > 0 && i + word_shift + 1 < 4 {
                ret[i + word_shift + 1] += original[i] >> (64 - bit_shift);
            }
        }
        U256(ret)
    }
}

impl Shr<usize> for U256 {
    type Output = U256;

    fn shr(self, shift: usize) -> U256 {
        let U256(ref original) = self;
        let mut ret = [0u64; 4];
        let word_shift = shift / 64;
        let bit_shift = shift % 64;
        for i in word_shift..4 {
            // Shift
            ret[i - word_shift] += original[i] >> bit_shift;
            // Carry
            if bit_shift > 0 && i < 4 - 1 {
                ret[i - word_shift] += original[i + 1] << (64 - bit_shift);
            }
        }
        U256(ret)
    }
}

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

    #[test]
    fn u256_add() {
        assert_eq!(
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64]) +
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64]),
            U256([0xfffffffffffffffeu64, 1u64, 0u64, 0u64])
        );
    }

    #[test]
    fn u256_sub() {
        assert_eq!(
            U256([0xfffffffffffffffeu64, 1u64, 0u64, 0u64]) -
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64]),
            U256([0xffffffffffffffffu64, 0u64, 0u64, 0u64])
        );
    }

    #[test]
    fn u256_log2floor() {
        assert_eq!(
            U256([u64::max_value(), u64::max_value(), u64::max_value(), u64::max_value()])
                .log2floor(),
            256, "testing log2floor for max value");
        assert_eq!(
            U256([0x1u64, 0u64, 0u64, 0u64]).log2floor(),
            1, "testing log2floor for 1");
    }
}
<|endoftext|>


<commit_before>use std::io;

use tag;
use err;

macro_rules! asn1_info {
  ($rs_type:ty, $asn1_ty:expr, $class:expr, $tagnum:expr, $constructed:expr) => (
    impl $crate::serial::traits::Asn1Info for $rs_type {
      fn asn1_type() -> tag::Type {
        tag::Type::from($asn1_ty)
      }

      fn asn1_class() -> tag::Class {
        tag::class::from($class)
      }

      fn asn1_tagnum() -> tag::TagNum {
        tag::TagNum::from($tagnum)
      }

      fn asn1_constructed() -> bool {
        $constructed
      }
    }
  )
}

macro_rules! asn1_sequence {
  ($rs_type:ty, $asn1_ty:expr, $($items:ident),*) => (
    impl $crate::serial::traits::Asn1Info for $rs_type {
      fn asn1_type() -> tag::Type {
        $asn1_ty.into()
      }

      fn asn1_class() -> tag::Class {
      }

      fn asn1_tagnum() -> tag::TagNum {

      }

      fn asn1_constructed() -> bool {
        true
      }
    }
  )
}
<commit_msg>Progress on macros<commit_after>use std::io;

use tag;
use err;

macro_rules! asn1_info {
  ($rs_type:ty, $asn1_ty:expr, $class:expr, $tagnum:expr, $constructed:expr) => (
    impl $crate::serial::traits::Asn1Info for $rs_type {
      fn asn1_type() -> tag::Type {
        tag::Type::from($asn1_ty)
      }

      fn asn1_class() -> tag::Class {
        tag::class::from($class)
      }

      fn asn1_tagnum() -> tag::TagNum {
        tag::TagNum::from($tagnum)
      }

      fn asn1_constructed() -> bool {
        $constructed
      }
    }
  )
}

macro_rules! asn1_sequence_info {
  ($rs_type:ty, $asn1_ty:expr, $($items:ident),*) => (
    impl $crate::serial::traits::Asn1Info for $rs_type {
      fn asn1_type() -> tag::Type {
        $asn1_ty.into()
      }

      fn asn1_class() -> tag::Class {
        tag::Class::Class::Universal
      }

      fn asn1_tagnum() -> tag::TagNum {
        tag::TagNum::from(0x10)
      }

      fn asn1_constructed() -> bool {
        true
      }
    }
  )
}

macro_rules! asn1_sequence_serialize {
  ($rs_type:ty, $($item:ident),*) => (
    impl serial::traits::Asn1Serialize for $rs_type {
      fn serialize<W: io::Write>(&self, writer: &mut W) -> Result<(), err::EncodeError> {
        let tag = tag::Tag {
          class: self.asn1_class(),
          tagnum: self.asn1_tagnum(),
          constructed: self.asn1_constructed(),
          len: tag::Len::Indef,
        };
        try!(
          // Write initial tag.
          tag.encode_tag(writer);
          // For each declared sequence member, serialize it onto the stream.
          $(
            serial::traits::Asn1Serialize::serialize(self.$item, writer);
          )*
          // Write the end of indef length marker.
          byte::write_byte(writer, 0x00);
          byte::write_byte(writer, 0x00);
        );
      }
    }
  )
}

macro_rules! asn1_sequence_deserialize {
  ($rs_type:ty) => (
    impl serial::traits::Asn1Deserialize for u64 {
      fn deserialize<I: Iterator<Item=io::Result<u8>>>(reader: I) -> Result<Self, err::DecodeError> {
        unimplemented!();
      }
    }
  )
}
<|endoftext|>


<commit_before>
use std::old_io::File;
use std::old_io::{Reader, Writer, Seek, SeekSet, SeekEnd};
use std::iter::range_inclusive;
use error::ZipError;
use maybe_utf8::{MaybeUtf8Slice, MaybeUtf8Buf, IntoMaybeUtf8};
use flate;
use crc32;
use format;
use fileinfo::{CompressionMethod, FileInfo};

pub struct ZipReader<R> {
    reader: R,
    end_record: format::EndOfCentralDirectoryRecord,
}

pub struct RawFiles<'a, R:'a> {
    zip_reader: &'a mut ZipReader<R>,
    current_entry: u16,
    current_offset: u64,
}

impl<'a, R: Reader+Seek> Iterator for RawFiles<'a, R> {
    type Item = Result<FileInfo, ZipError>;
    fn next(&mut self) -> Option<Result<FileInfo, ZipError>> {
        if self.current_entry < self.zip_reader.end_record.total_entry_count {
            match self.zip_reader.reader.seek(self.current_offset as i64, SeekSet) {
                Ok(()) => {}
                Err(err) => { return Some(Err(ZipError::IoError(err))); }
            }
            let h = match format::CentralDirectoryHeader::read(&mut self.zip_reader.reader) {
                Ok(h) => h,
                Err(err) => { return Some(Err(err)); }
            };
            let info = FileInfo::from_cdh(&h);
            self.current_entry += 1;
            self.current_offset += h.total_size() as u64;
            Some(Ok(info))
        } else {
            None
        }
    }
}

pub struct Files<'a, R:'a> {
    base: RawFiles<'a, R>,
}

impl<'a, R: Reader+Seek> Iterator for Files<'a, R> {
    type Item = FileInfo;
    fn next(&mut self) -> Option<FileInfo> { self.base.next().map(|i| i.ok().unwrap()) }
    fn size_hint(&self) -> (usize, Option<usize>) { self.base.size_hint() }
}

pub struct FileNames<'a, R:'a> {
    base: RawFiles<'a, R>,
}

impl<'a, R: Reader+Seek> Iterator for FileNames<'a, R> {
    type Item = MaybeUtf8Buf;
    fn next(&mut self) -> Option<MaybeUtf8Buf> {
        self.base.next().map(|i| i.ok().unwrap().name)
    }
    fn size_hint(&self) -> (usize, Option<usize>) { self.base.size_hint() }
}

impl ZipReader<File> {
    pub fn open(path: &Path) -> Result<ZipReader<File>, ZipError> {
        ZipReader::new(try_io!(File::open(path)))
    }
}

impl<R:Reader+Seek> ZipReader<R> {
    pub fn new(reader: R) -> Result<ZipReader<R>, ZipError> {
        // find the End of Central Directory record, looking backwards from the end of the file
        let mut r = reader;
        try_io!(r.seek(0, SeekEnd));
        let file_size = try_io!(r.tell());
        let mut end_record_offset : Option<u64> = None;
        for i in range_inclusive(4, file_size) {
            let offset = file_size - i;
            try_io!(r.seek(offset as i64, SeekSet));

            let sig = try_io!(r.read_le_u32());

            // TODO: check for false positives here
            if sig == format::EOCDR_SIGNATURE {
                end_record_offset = Some(offset);
                break;
            }

        }

        match end_record_offset {
            Some(offset) => {
                try_io!(r.seek(offset as i64, SeekSet));
                let e = try!(format::EndOfCentralDirectoryRecord::read(&mut r));
                Ok(ZipReader {reader: r, end_record: e})
            },
            None => Err(ZipError::NotAZipFile)
        }
    }

    pub fn files_raw<'a>(&'a mut self) -> RawFiles<'a, R> {
        let cdr_offset = self.end_record.central_directory_offset;
        RawFiles {
            zip_reader: self,
            current_entry: 0,
            current_offset: cdr_offset as u64
        }
    }

    pub fn files<'a>(&'a mut self) -> Files<'a, R> {
        Files { base: self.files_raw() }
    }

    pub fn file_names<'a>(&'a mut self) -> FileNames<'a, R> {
        FileNames { base: self.files_raw() }
    }

    pub fn info<'a, T>(&mut self, name: T) -> Result<FileInfo, ZipError>
            where T: IntoMaybeUtf8<MaybeUtf8Slice<'a>> {
        let name = name.into_maybe_utf8();
        for i in self.files() {
            if i.name == name.as_bytes() {
                return Ok(i);
            }
        }
        Err(ZipError::FileNotFoundInArchive)
    }

    pub fn extract_file<T:Writer>(&mut self, f: &FileInfo, writer: &mut T) -> Result<(), ZipError> {
        match self.read(f, -1 as usize) {
            Ok(bytes) => { try_io!(writer.write_all(&bytes[..])); Ok(()) },
            Err(x) => Err(x)
        }
    }

    pub fn extract_first<T:Writer>(&mut self, f: &FileInfo, len: usize, writer: &mut T) -> Result<(), ZipError> {
        match self.read(f, len) {
            Ok(bytes) => { try_io!(writer.write_all(&bytes[..])); Ok(()) },
            Err(x) => Err(x)
        }
    }

    fn read(&mut self, f: &FileInfo, wish_len: usize) -> Result<Vec<u8>, ZipError> {
        try_io!(self.reader.seek(f.local_file_header_offset as i64, SeekSet));
        let header = try!(format::LocalFileHeader::read(&mut self.reader));
        let file_pos = f.local_file_header_offset as i64 + header.total_size() as i64;
        let file_len = header.compressed_size as usize;
        if wish_len > file_len {
            self.extract_block(file_pos, file_len, header.compression_method, header.crc32)
        } else {
            self.extract_block(file_pos, wish_len, header.compression_method, 0)
        }
    }

    fn extract_block(&mut self, pos: i64, len: usize, method: u16, crc32: u32) -> Result<Vec<u8>, ZipError> {
        try_io!(self.reader.seek(pos, SeekSet));
        let compressed = try_io!(self.reader.read_exact(len));
        match CompressionMethod::from_u16(method) {
                CompressionMethod::Store   => Ok(compressed),
                CompressionMethod::Deflate => self.decompress(compressed, len, crc32),
                _ => panic!("Usupported CompressionMethod")
        }
    }

    fn decompress(&mut self, data: Vec<u8>, len:usize, crc32: u32) -> Result<Vec<u8>, ZipError>
    {
        let bytes = match flate::inflate_bytes(&data[..])
        {
            Some(ok) => ok,
            None => return Err(ZipError::DecompressionFailure)
        };
        if crc32 != 0 && crc32 != crc32::crc32(&bytes){
            return Err(ZipError::CrcError);
        }
        Ok(bytes[0..len].to_vec())
    }
}

<commit_msg>Fix build errors with 'rustc 1.0.0-nightly (ea8b82e90 2015-03-17) (built 2015-03-18)'<commit_after>
use std::old_io::File;
use std::old_io::{Reader, Writer, Seek, SeekSet, SeekEnd};
use std::iter::range_inclusive;
use error::ZipError;
use maybe_utf8::{MaybeUtf8Slice, MaybeUtf8Buf, IntoMaybeUtf8};
use flate;
use crc32;
use format;
use fileinfo::{CompressionMethod, FileInfo};

pub struct ZipReader<R> {
    reader: R,
    end_record: format::EndOfCentralDirectoryRecord,
}

pub struct RawFiles<'a, R:'a> {
    zip_reader: &'a mut ZipReader<R>,
    current_entry: u16,
    current_offset: u64,
}

impl<'a, R: Reader+Seek> Iterator for RawFiles<'a, R> {
    type Item = Result<FileInfo, ZipError>;
    fn next(&mut self) -> Option<Result<FileInfo, ZipError>> {
        if self.current_entry < self.zip_reader.end_record.total_entry_count {
            match self.zip_reader.reader.seek(self.current_offset as i64, SeekSet) {
                Ok(()) => {}
                Err(err) => { return Some(Err(ZipError::IoError(err))); }
            }
            let h = match format::CentralDirectoryHeader::read(&mut self.zip_reader.reader) {
                Ok(h) => h,
                Err(err) => { return Some(Err(err)); }
            };
            let info = FileInfo::from_cdh(&h);
            self.current_entry += 1;
            self.current_offset += h.total_size() as u64;
            Some(Ok(info))
        } else {
            None
        }
    }
}

pub struct Files<'a, R:'a> {
    base: RawFiles<'a, R>,
}

impl<'a, R: Reader+Seek> Iterator for Files<'a, R> {
    type Item = FileInfo;
    fn next(&mut self) -> Option<FileInfo> { self.base.next().map(|i| i.ok().unwrap()) }
    fn size_hint(&self) -> (usize, Option<usize>) { self.base.size_hint() }
}

pub struct FileNames<'a, R:'a> {
    base: RawFiles<'a, R>,
}

impl<'a, R: Reader+Seek> Iterator for FileNames<'a, R> {
    type Item = MaybeUtf8Buf;
    fn next(&mut self) -> Option<MaybeUtf8Buf> {
        self.base.next().map(|i| i.ok().unwrap().name)
    }
    fn size_hint(&self) -> (usize, Option<usize>) { self.base.size_hint() }
}

impl ZipReader<File> {
    pub fn open(path: &Path) -> Result<ZipReader<File>, ZipError> {
        ZipReader::new(try_io!(File::open(path)))
    }
}

impl<R:Reader+Seek> ZipReader<R> {
    pub fn new(reader: R) -> Result<ZipReader<R>, ZipError> {
        // find the End of Central Directory record, looking backwards from the end of the file
        let mut r = reader;
        try_io!(r.seek(0, SeekEnd));
        let file_size = try_io!(r.tell());
        let mut end_record_offset : Option<u64> = None;
        for i in range_inclusive(4, file_size) {
            let offset = file_size - i;
            try_io!(r.seek(offset as i64, SeekSet));

            let sig = try_io!(r.read_le_u32());

            // TODO: check for false positives here
            if sig == format::EOCDR_SIGNATURE {
                end_record_offset = Some(offset);
                break;
            }

        }

        match end_record_offset {
            Some(offset) => {
                try_io!(r.seek(offset as i64, SeekSet));
                let e = try!(format::EndOfCentralDirectoryRecord::read(&mut r));
                Ok(ZipReader {reader: r, end_record: e})
            },
            None => Err(ZipError::NotAZipFile)
        }
    }

    pub fn files_raw<'a>(&'a mut self) -> RawFiles<'a, R> {
        let cdr_offset = self.end_record.central_directory_offset;
        RawFiles {
            zip_reader: self,
            current_entry: 0,
            current_offset: cdr_offset as u64
        }
    }

    pub fn files<'a>(&'a mut self) -> Files<'a, R> {
        Files { base: self.files_raw() }
    }

    pub fn file_names<'a>(&'a mut self) -> FileNames<'a, R> {
        FileNames { base: self.files_raw() }
    }

    pub fn info<'a, T>(&mut self, name: T) -> Result<FileInfo, ZipError>
            where T: IntoMaybeUtf8<MaybeUtf8Slice<'a>> {
        let name = name.into_maybe_utf8();
        for i in self.files() {
            if i.name == name.as_bytes() {
                return Ok(i);
            }
        }
        Err(ZipError::FileNotFoundInArchive)
    }

    pub fn extract_file<T:Writer>(&mut self, f: &FileInfo, writer: &mut T) -> Result<(), ZipError> {
        match self.read(f, -1 as usize) {
            Ok(bytes) => { try_io!(writer.write_all(&bytes[..])); Ok(()) },
            Err(x) => Err(x)
        }
    }

    pub fn extract_first<T:Writer>(&mut self, f: &FileInfo, len: usize, writer: &mut T) -> Result<(), ZipError> {
        match self.read(f, len) {
            Ok(bytes) => { try_io!(writer.write_all(&bytes[..])); Ok(()) },
            Err(x) => Err(x)
        }
    }

    fn read(&mut self, f: &FileInfo, wish_len: usize) -> Result<Vec<u8>, ZipError> {
        try_io!(self.reader.seek(f.local_file_header_offset as i64, SeekSet));
        let header = try!(format::LocalFileHeader::read(&mut self.reader));
        let file_pos = f.local_file_header_offset as i64 + header.total_size() as i64;
        let file_len = header.compressed_size as usize;
        if wish_len > file_len {
            self.extract_block(file_pos, file_len, header.compression_method, header.crc32)
        } else {
            self.extract_block(file_pos, wish_len, header.compression_method, 0)
        }
    }

    fn extract_block(&mut self, pos: i64, len: usize, method: u16, crc32: u32) -> Result<Vec<u8>, ZipError> {
        try_io!(self.reader.seek(pos, SeekSet));
        let compressed = try_io!(self.reader.read_exact(len));
        match CompressionMethod::from_u16(method) {
                CompressionMethod::Store   => Ok(compressed),
                CompressionMethod::Deflate => self.decompress(compressed, len, crc32),
                _ => panic!("Usupported CompressionMethod")
        }
    }

    fn decompress(&mut self, data: Vec<u8>, len:usize, crc32: u32) -> Result<Vec<u8>, ZipError>
    {

        //let bytes :u8 = flate::inflate_bytes(&data[..]);
        let bytes = match flate::inflate_bytes(&data[..])
        {
            Ok(decompressed) => decompressed,
            Err(what) => return Err(ZipError::DecompressionFailure)
        };
        if crc32 != 0 && crc32 != crc32::crc32(&bytes){
            return Err(ZipError::CrcError);
        }
        Ok(bytes[0..len].to_vec())
    }
}

<|endoftext|>


<commit_before>/* Copyright 2017 Joel Pedraza
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

use boggle_util;
use bitset::BitSet;
use bitset::IndexIter;

use super::SimpleBoggleBoard;

/*
 * Can this be done cleaner with Enums and some sort of EnumSet
 */

const FLAG_NORTHWEST : u8  = 0b10000000;
const FLAG_NORTH     : u8  = 0b01000000;
const FLAG_NORTHEAST : u8  = 0b00100000;
const FLAG_WEST      : u8  = 0b00010000;
const FLAG_EAST      : u8  = 0b00001000;
const FLAG_SOUTHWEST : u8  = 0b00000100;
const FLAG_SOUTH     : u8  = 0b00000010;
const FLAG_SOUTHEAST : u8  = 0b00000001;

type RadixBoggleCell = [u8; boggle_util::ALPHABET_SIZE];

pub struct RadixBoggleBoard {
    width: usize,
    height: usize,
    /// Top level navigation by value
    /// (Used to quickly find all the cells on the board of a specified value)
    alpha: [BitSet; boggle_util::ALPHABET_SIZE],
    /// Serves as a precomupted adjacency matrix filtered by value
    /// (Used to quickly find all the neighbors of a cell of a specified value)
    cells: Box<[RadixBoggleCell]>,
}

impl RadixBoggleBoard {
    pub fn new(width: usize, height: usize) -> Self {
        use std::mem;
        use std::ptr;

        let mut alpha: [BitSet; boggle_util::ALPHABET_SIZE];
        unsafe {
            alpha = mem::uninitialized();

            for element in alpha.iter_mut() {
                let bs = BitSet::new();
                ptr::write(element, bs);
            }
        }

        RadixBoggleBoard {
            width: width,
            height: height,
            alpha: alpha,
            cells: vec![ Default::default(); width * height ].into_boxed_slice(),
        }
    }

    /// Creates a new radix board from a filled simple board
    /// this is mostly due to laziness, as SimpleBoggleBoard::read
    /// could be cleaner and I dont want to duplicate ugly code xD
    /// Maybe use a generified builder?
    pub fn from(src: &SimpleBoggleBoard) -> Self {
        let mut dst = Self::new(src.width(), src.height());
        for (i, v) in src.iter().enumerate() {
            dst.set(i, *v);
        }
        dst
    }

    pub fn any(&self, v: u8) -> IndexIter {
        self.alpha[v as usize].iter_ones()
    }

    pub fn neighbors(&self, i: usize, v: u8) -> RadixNeighborIter {
        RadixNeighborIter {
            v: self.cells[i][v as usize],
            i: i,
            w: self.width
        }
    }

    #[inline]
    fn mask_cell(&mut self, v: usize, i: usize, mask: u8) {
        self.cells[i][v] |= mask;
    }

    pub fn set(&mut self, i: usize, v: u8) {
        self.alpha[v as usize].add(i);

        let w = self.width;
        let v = v as usize;

        // Mask each neighbor of i with the flag for v (relative to i)
        match i {
            // northwest corner
            0 => {
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            },

            // northeast corner
            x if x == self.width -1 => {
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
            },

            // southwest corner
            x if x == self.width * (self.height - 1) => {
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  +1, FLAG_WEST);
            },

            // southeast corner
            x if x == self.width * self.height - 1 => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i  -1, FLAG_EAST);
            },

            // north edge
            x if x < self.width => {
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            },

            // south edge
            x if x > self.width * (self.height - 1) => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i  +1, FLAG_WEST);
            },

            // west edge
            x if x % self.width == 0 => {
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            },

            // east edge
            x if x % self.width == self.width - 1 => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
            },

            // interior
            _ => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            }
        }
    }
}

pub struct RadixNeighborIter {
    value: u8,
    // the index of the cell
    idx: usize,
    // the width of the board
    width: usize,
}

impl Iterator for RadixNeighborIter {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        match self.value.leading_zeros() {
            0 => { self.value &=0b01111111; Some(self.idx - self.width - 1) }
            1 => { self.value &=0b00111111; Some(self.idx - self.width    ) },
            2 => { self.value &=0b00011111; Some(self.idx - self.width + 1) },
            3 => { self.value &=0b00001111; Some(self.idx              - 1) },
            4 => { self.value &=0b00000111; Some(self.idx              + 1) },
            5 => { self.value &=0b00000011; Some(self.idx + self.width - 1) },
            6 => { self.value &=0b00000001; Some(self.idx + self.width    ) },
            7 => { self.value  =0b00000000; Some(self.idx + self.width + 1) },
            _ => None
        }
    }
}
<commit_msg>fix typo bug<commit_after>/* Copyright 2017 Joel Pedraza
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

use boggle_util;
use bitset::BitSet;
use bitset::IndexIter;

use super::SimpleBoggleBoard;

/*
 * Can this be done cleaner with Enums and some sort of EnumSet
 */

const FLAG_NORTHWEST : u8  = 0b10000000;
const FLAG_NORTH     : u8  = 0b01000000;
const FLAG_NORTHEAST : u8  = 0b00100000;
const FLAG_WEST      : u8  = 0b00010000;
const FLAG_EAST      : u8  = 0b00001000;
const FLAG_SOUTHWEST : u8  = 0b00000100;
const FLAG_SOUTH     : u8  = 0b00000010;
const FLAG_SOUTHEAST : u8  = 0b00000001;

type RadixBoggleCell = [u8; boggle_util::ALPHABET_SIZE];

pub struct RadixBoggleBoard {
    width: usize,
    height: usize,
    /// Top level navigation by value
    /// (Used to quickly find all the cells on the board of a specified value)
    alpha: [BitSet; boggle_util::ALPHABET_SIZE],
    /// Serves as a precomupted adjacency matrix filtered by value
    /// (Used to quickly find all the neighbors of a cell of a specified value)
    cells: Box<[RadixBoggleCell]>,
}

impl RadixBoggleBoard {
    pub fn new(width: usize, height: usize) -> Self {
        use std::mem;
        use std::ptr;

        let mut alpha: [BitSet; boggle_util::ALPHABET_SIZE];
        unsafe {
            alpha = mem::uninitialized();

            for element in alpha.iter_mut() {
                let bs = BitSet::new();
                ptr::write(element, bs);
            }
        }

        RadixBoggleBoard {
            width: width,
            height: height,
            alpha: alpha,
            cells: vec![ Default::default(); width * height ].into_boxed_slice(),
        }
    }

    /// Creates a new radix board from a filled simple board
    /// this is mostly due to laziness, as SimpleBoggleBoard::read
    /// could be cleaner and I dont want to duplicate ugly code xD
    /// Maybe use a generified builder?
    pub fn from(src: &SimpleBoggleBoard) -> Self {
        let mut dst = Self::new(src.width(), src.height());
        for (i, v) in src.iter().enumerate() {
            dst.set(i, *v);
        }
        dst
    }

    pub fn any(&self, v: u8) -> IndexIter {
        self.alpha[v as usize].iter_ones()
    }

    pub fn neighbors(&self, i: usize, v: u8) -> RadixNeighborIter {
        RadixNeighborIter {
            value: self.cells[i][v as usize],
            idx: i,
            width: self.width
        }
    }

    #[inline]
    fn mask_cell(&mut self, v: usize, i: usize, mask: u8) {
        self.cells[i][v] |= mask;
    }

    pub fn set(&mut self, i: usize, v: u8) {
        self.alpha[v as usize].add(i);

        let w = self.width;
        let v = v as usize;

        // Mask each neighbor of i with the flag for v (relative to i)
        match i {
            // northwest corner
            0 => {
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            },

            // northeast corner
            x if x == self.width -1 => {
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
            },

            // southwest corner
            x if x == self.width * (self.height - 1) => {
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  +1, FLAG_WEST);
            },

            // southeast corner
            x if x == self.width * self.height - 1 => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i  -1, FLAG_EAST);
            },

            // north edge
            x if x < self.width => {
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            },

            // south edge
            x if x > self.width * (self.height - 1) => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i  +1, FLAG_WEST);
            },

            // west edge
            x if x % self.width == 0 => {
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            },

            // east edge
            x if x % self.width == self.width - 1 => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
            },

            // interior
            _ => {
                self.mask_cell(v, i-w-1, FLAG_SOUTHEAST);
                self.mask_cell(v, i-w  , FLAG_SOUTH);
                self.mask_cell(v, i-w+1, FLAG_SOUTHWEST);
                self.mask_cell(v, i  -1, FLAG_EAST);
                self.mask_cell(v, i  +1, FLAG_WEST);
                self.mask_cell(v, i+w-1, FLAG_NORTHEAST);
                self.mask_cell(v, i+w  , FLAG_NORTH);
                self.mask_cell(v, i+w+1, FLAG_NORTHWEST);
            }
        }
    }
}

pub struct RadixNeighborIter {
    value: u8,
    // the index of the cell
    idx: usize,
    // the width of the board
    width: usize,
}

impl Iterator for RadixNeighborIter {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        match self.value.leading_zeros() {
            0 => { self.value &=0b01111111; Some(self.idx - self.width - 1) }
            1 => { self.value &=0b00111111; Some(self.idx - self.width    ) },
            2 => { self.value &=0b00011111; Some(self.idx - self.width + 1) },
            3 => { self.value &=0b00001111; Some(self.idx              - 1) },
            4 => { self.value &=0b00000111; Some(self.idx              + 1) },
            5 => { self.value &=0b00000011; Some(self.idx + self.width - 1) },
            6 => { self.value &=0b00000001; Some(self.idx + self.width    ) },
            7 => { self.value  =0b00000000; Some(self.idx + self.width + 1) },
            _ => None
        }
    }
}
<|endoftext|>


<commit_before>use crate::bundle::layout::macos::MacOsKbdLayer;

pub fn layer_attributes(layer: MacOsKbdLayer) -> String {
    match layer {
        MacOsKbdLayer::Default => "command?",
        MacOsKbdLayer::Shift => "anyShift caps? command?",
        MacOsKbdLayer::Caps => "caps",
        MacOsKbdLayer::Alt => "anyOption command?",
        MacOsKbdLayer::AltAndShift => "anyOption anyShift caps? command?",
        MacOsKbdLayer::CapsAndAlt => "caps anyOption command?",
        MacOsKbdLayer::Ctrl => "anyShift? caps? anyOption? anyControl",
        MacOsKbdLayer::Cmd => "command",
        MacOsKbdLayer::CmdAndShift => "command anyShift",
        MacOsKbdLayer::CmdAndAlt => todo!(),
        MacOsKbdLayer::CmdAndAltAndShift => todo!(),
    }
    .to_string()
}
<commit_msg>Being creative<commit_after>use crate::bundle::layout::macos::MacOsKbdLayer;

pub fn layer_attributes(layer: MacOsKbdLayer) -> String {
    match layer {
        MacOsKbdLayer::Default => "command?",
        MacOsKbdLayer::Shift => "anyShift caps? command?",
        MacOsKbdLayer::Caps => "caps",
        MacOsKbdLayer::Alt => "anyOption command?",
        MacOsKbdLayer::AltAndShift => "anyOption anyShift caps? command?",
        MacOsKbdLayer::CapsAndAlt => "caps anyOption command?",
        MacOsKbdLayer::Ctrl => "anyShift? caps? anyOption? anyControl",
        MacOsKbdLayer::Cmd => "command",
        MacOsKbdLayer::CmdAndShift => "command anyShift",
        MacOsKbdLayer::CmdAndAlt => "command anyOption",
        MacOsKbdLayer::CmdAndAltAndShift => "command anyOption anyShift",
    }
    .to_string()
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, ReadBytesExt};
use std::convert::TryFrom;
use std::io::{self, Cursor, Read, Seek, SeekFrom};
use std::marker::PhantomData;
use std::{error, fmt, mem};

use crate::color::ColorType;
use crate::error::{DecodingError, ImageError, ImageResult, UnsupportedError, UnsupportedErrorKind};
use crate::image::{self, ImageDecoder, ImageFormat};

use self::InnerDecoder::*;
use crate::bmp::BmpDecoder;
use crate::png::PngDecoder;

// http://www.w3.org/TR/PNG-Structure.html
// The first eight bytes of a PNG file always contain the following (decimal) values:
const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// Errors that can occur during decoding and parsing an ICO image or one of its enclosed images.
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum DecoderError {
    /// The ICO directory is empty
    NoEntries,
    /// The number of color planes (0 or 1), or the horizontal coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyPlanesOrHotspot,
    /// The bit depth (may be 0 meaning unspecified), or the vertical coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyBitsPerPixelOrHotspot,

    /// The entry is in PNG format and specified a length that is shorter than PNG header.
    PngShorterThanHeader,
    /// The enclosed PNG is not in RGBA, which is invalid: https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/.
    PngNotRgba,

    /// The optional mask row, containing 1 bit per pixel, padded to 4 bytes, was too short for this image.
    BmpIcoMaskTooShortForImage,

    /// The dimensions specified by the entry does not match the dimensions in the header of the enclosed image.
    ImageEntryDimensionMismatch {
        /// The mismatched subimage's type
        format: IcoEntryImageFormat,
        /// The dimensions specified by the entry
        entry: (u16, u16),
        /// The dimensions of the image itself
        image: (u32, u32)
    },
}

impl fmt::Display for DecoderError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecoderError::NoEntries =>
                f.write_str("ICO directory contains no image"),
            DecoderError::IcoEntryTooManyPlanesOrHotspot =>
                f.write_str("ICO image entry has too many color planes or too large hotspot value"),
            DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot =>
                f.write_str("ICO image entry has too many bits per pixel or too large hotspot value"),
            DecoderError::PngShorterThanHeader =>
                f.write_str("Entry specified a length that is shorter than PNG header!"),
            DecoderError::PngNotRgba =>
                f.write_str("The PNG is not in RGBA format!"),
            DecoderError::BmpIcoMaskTooShortForImage =>
                f.write_str("ICO mask too short for the image"),
            DecoderError::ImageEntryDimensionMismatch { format, entry, image } =>
                f.write_fmt(format_args!("Entry{:?} and {}{:?} dimensions do not match!", entry, format, image)),
        }
    }
}

impl From<DecoderError> for ImageError {
    fn from(e: DecoderError) -> ImageError {
        ImageError::Decoding(DecodingError::new(ImageFormat::Ico.into(), e))
    }
}

impl error::Error for DecoderError {}

/// The image formats an ICO may contain
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum IcoEntryImageFormat {
    /// PNG in ARGB
    Png,
    /// BMP with optional alpha mask
    Bmp,
}

impl fmt::Display for IcoEntryImageFormat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            IcoEntryImageFormat::Png => "PNG",
            IcoEntryImageFormat::Bmp => "BMP",
        })
    }
}

impl Into<ImageFormat> for IcoEntryImageFormat {
    fn into(self) -> ImageFormat {
        match self {
            IcoEntryImageFormat::Png => ImageFormat::Png,
            IcoEntryImageFormat::Bmp => ImageFormat::Bmp,
        }
    }
}

/// An ico decoder
pub struct IcoDecoder<R: Read> {
    selected_entry: DirEntry,
    inner_decoder: InnerDecoder<R>,
}

enum InnerDecoder<R: Read> {
    BMP(BmpDecoder<R>),
    PNG(PngDecoder<R>),
}

#[derive(Clone, Copy, Default)]
struct DirEntry {
    width: u8,
    height: u8,
    color_count: u8,
    reserved: u8,

    num_color_planes: u16,
    bits_per_pixel: u16,

    image_length: u32,
    image_offset: u32,
}

impl<R: Read + Seek> IcoDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(mut r: R) -> ImageResult<IcoDecoder<R>> {
        let entries = read_entries(&mut r)?;
        let entry = best_entry(entries)?;
        let decoder = entry.decoder(r)?;

        Ok(IcoDecoder {
            selected_entry: entry,
            inner_decoder: decoder,
        })
    }
}

fn read_entries<R: Read>(r: &mut R) -> ImageResult<Vec<DirEntry>> {
    let _reserved = r.read_u16::<LittleEndian>()?;
    let _type = r.read_u16::<LittleEndian>()?;
    let count = r.read_u16::<LittleEndian>()?;
    (0..count).map(|_| read_entry(r)).collect()
}

fn read_entry<R: Read>(r: &mut R) -> ImageResult<DirEntry> {
    let mut entry = DirEntry::default();

    entry.width = r.read_u8()?;
    entry.height = r.read_u8()?;
    entry.color_count = r.read_u8()?;
    // Reserved value (not used)
    entry.reserved = r.read_u8()?;

    // This may be either the number of color planes (0 or 1), or the horizontal coordinate
    // of the hotspot for CUR files.
    entry.num_color_planes = r.read_u16::<LittleEndian>()?;
    if entry.num_color_planes > 256 {
        return Err(DecoderError::IcoEntryTooManyPlanesOrHotspot.into());
    }

    // This may be either the bit depth (may be 0 meaning unspecified),
    // or the vertical coordinate of the hotspot for CUR files.
    entry.bits_per_pixel = r.read_u16::<LittleEndian>()?;
    if entry.bits_per_pixel > 256 {
        return Err(DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot.into());
    }

    entry.image_length = r.read_u32::<LittleEndian>()?;
    entry.image_offset = r.read_u32::<LittleEndian>()?;

    Ok(entry)
}

/// Find the entry with the highest (color depth, size).
fn best_entry(mut entries: Vec<DirEntry>) -> ImageResult<DirEntry> {
    let mut best = entries.pop().ok_or(DecoderError::NoEntries)?;

    let mut best_score = (
        best.bits_per_pixel,
        u32::from(best.real_width()) * u32::from(best.real_height()),
    );

    for entry in entries {
        let score = (
            entry.bits_per_pixel,
            u32::from(entry.real_width()) * u32::from(entry.real_height()),
        );
        if score > best_score {
            best = entry;
            best_score = score;
        }
    }
    Ok(best)
}

impl DirEntry {
    fn real_width(&self) -> u16 {
        match self.width {
            0 => 256,
            w => u16::from(w),
        }
    }

    fn real_height(&self) -> u16 {
        match self.height {
            0 => 256,
            h => u16::from(h),
        }
    }

    fn matches_dimensions(&self, width: u32, height: u32) -> bool {
        u32::from(self.real_width()) == width && u32::from(self.real_height()) == height
    }

    fn seek_to_start<R: Read + Seek>(&self, r: &mut R) -> ImageResult<()> {
        r.seek(SeekFrom::Start(u64::from(self.image_offset)))?;
        Ok(())
    }

    fn is_png<R: Read + Seek>(&self, r: &mut R) -> ImageResult<bool> {
        self.seek_to_start(r)?;

        // Read the first 8 bytes to sniff the image.
        let mut signature = [0u8; 8];
        r.read_exact(&mut signature)?;

        Ok(signature == PNG_SIGNATURE)
    }

    fn decoder<R: Read + Seek>(&self, mut r: R) -> ImageResult<InnerDecoder<R>> {
        let is_png = self.is_png(&mut r)?;
        self.seek_to_start(&mut r)?;

        if is_png {
            Ok(PNG(PngDecoder::new(r)?))
        } else {
            Ok(BMP(BmpDecoder::new_with_ico_format(r)?))
        }
    }
}

/// Wrapper struct around a `Cursor<Vec<u8>>`
pub struct IcoReader<R>(Cursor<Vec<u8>>, PhantomData<R>);
impl<R> Read for IcoReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        if self.0.position() == 0 && buf.is_empty() {
            mem::swap(buf, self.0.get_mut());
            Ok(buf.len())
        } else {
            self.0.read_to_end(buf)
        }
    }
}

impl<'a, R: 'a + Read + Seek> ImageDecoder<'a> for IcoDecoder<R> {
    type Reader = IcoReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        match self.inner_decoder {
            BMP(ref decoder) => decoder.dimensions(),
            PNG(ref decoder) => decoder.dimensions(),
        }
    }

    fn color_type(&self) -> ColorType {
        match self.inner_decoder {
            BMP(ref decoder) => decoder.color_type(),
            PNG(ref decoder) => decoder.color_type(),
        }
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        Ok(IcoReader(Cursor::new(image::decoder_to_vec(self)?), PhantomData))
    }

    fn read_image(self, buf: &mut [u8]) -> ImageResult<()> {
        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
        match self.inner_decoder {
            PNG(decoder) => {
                if self.selected_entry.image_length < PNG_SIGNATURE.len() as u32 {
                    return Err(DecoderError::PngShorterThanHeader.into());
                }

                // Check if the image dimensions match the ones in the image data.
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Png,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // Embedded PNG images can only be of the 32BPP RGBA format.
                // https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(DecoderError::PngNotRgba.into());
                }

                decoder.read_image(buf)
            }
            BMP(mut decoder) => {
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Bmp,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // The ICO decoder needs an alpha channel to apply the AND mask.
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                        ImageFormat::Bmp.into(),
                        UnsupportedErrorKind::Color(decoder.color_type().into()),
                    )));
                }

                decoder.read_image_data(buf)?;

                // If there's an AND mask following the image, read and apply it.
                let r = decoder.reader();
                let mask_start = r.seek(SeekFrom::Current(0))?;
                let mask_end =
                    u64::from(self.selected_entry.image_offset + self.selected_entry.image_length);
                let mask_length = mask_end - mask_start;

                if mask_length > 0 {
                    // A mask row contains 1 bit per pixel, padded to 4 bytes.
                    let mask_row_bytes = ((width + 31) / 32) * 4;
                    let expected_length = u64::from(mask_row_bytes) * u64::from(height);
                    if mask_length < expected_length {
                        return Err(DecoderError::BmpIcoMaskTooShortForImage.into());
                    }

                    for y in 0..height {
                        let mut x = 0;
                        for _ in 0..mask_row_bytes {
                            // Apply the bits of each byte until we reach the end of the row.
                            let mask_byte = r.read_u8()?;
                            for bit in (0..8).rev() {
                                if x >= width {
                                    break;
                                }
                                if mask_byte & (1 << bit) != 0 {
                                    // Set alpha channel to transparent.
                                    buf[((height - y - 1) * width + x) as usize * 4 + 3] = 0;
                                }
                                x += 1;
                            }
                        }
                    }
                }
                Ok(())
            }
        }
    }
}
<commit_msg>Improve safety of ICO BMP mask length check<commit_after>use byteorder::{LittleEndian, ReadBytesExt};
use std::convert::TryFrom;
use std::io::{self, Cursor, Read, Seek, SeekFrom};
use std::marker::PhantomData;
use std::{error, fmt, mem};

use crate::color::ColorType;
use crate::error::{DecodingError, ImageError, ImageResult, UnsupportedError, UnsupportedErrorKind};
use crate::image::{self, ImageDecoder, ImageFormat};

use self::InnerDecoder::*;
use crate::bmp::BmpDecoder;
use crate::png::PngDecoder;

// http://www.w3.org/TR/PNG-Structure.html
// The first eight bytes of a PNG file always contain the following (decimal) values:
const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// Errors that can occur during decoding and parsing an ICO image or one of its enclosed images.
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum DecoderError {
    /// The ICO directory is empty
    NoEntries,
    /// The number of color planes (0 or 1), or the horizontal coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyPlanesOrHotspot,
    /// The bit depth (may be 0 meaning unspecified), or the vertical coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyBitsPerPixelOrHotspot,

    /// The entry is in PNG format and specified a length that is shorter than PNG header.
    PngShorterThanHeader,
    /// The enclosed PNG is not in RGBA, which is invalid: https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/.
    PngNotRgba,

    /// The optional mask row, containing 1 bit per pixel, padded to 4 bytes, was too short for this image.
    BmpIcoMaskTooShortForImage,

    /// The dimensions specified by the entry does not match the dimensions in the header of the enclosed image.
    ImageEntryDimensionMismatch {
        /// The mismatched subimage's type
        format: IcoEntryImageFormat,
        /// The dimensions specified by the entry
        entry: (u16, u16),
        /// The dimensions of the image itself
        image: (u32, u32)
    },
}

impl fmt::Display for DecoderError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecoderError::NoEntries =>
                f.write_str("ICO directory contains no image"),
            DecoderError::IcoEntryTooManyPlanesOrHotspot =>
                f.write_str("ICO image entry has too many color planes or too large hotspot value"),
            DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot =>
                f.write_str("ICO image entry has too many bits per pixel or too large hotspot value"),
            DecoderError::PngShorterThanHeader =>
                f.write_str("Entry specified a length that is shorter than PNG header!"),
            DecoderError::PngNotRgba =>
                f.write_str("The PNG is not in RGBA format!"),
            DecoderError::BmpIcoMaskTooShortForImage =>
                f.write_str("ICO mask too short for the image"),
            DecoderError::ImageEntryDimensionMismatch { format, entry, image } =>
                f.write_fmt(format_args!("Entry{:?} and {}{:?} dimensions do not match!", entry, format, image)),
        }
    }
}

impl From<DecoderError> for ImageError {
    fn from(e: DecoderError) -> ImageError {
        ImageError::Decoding(DecodingError::new(ImageFormat::Ico.into(), e))
    }
}

impl error::Error for DecoderError {}

/// The image formats an ICO may contain
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum IcoEntryImageFormat {
    /// PNG in ARGB
    Png,
    /// BMP with optional alpha mask
    Bmp,
}

impl fmt::Display for IcoEntryImageFormat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            IcoEntryImageFormat::Png => "PNG",
            IcoEntryImageFormat::Bmp => "BMP",
        })
    }
}

impl Into<ImageFormat> for IcoEntryImageFormat {
    fn into(self) -> ImageFormat {
        match self {
            IcoEntryImageFormat::Png => ImageFormat::Png,
            IcoEntryImageFormat::Bmp => ImageFormat::Bmp,
        }
    }
}

/// An ico decoder
pub struct IcoDecoder<R: Read> {
    selected_entry: DirEntry,
    inner_decoder: InnerDecoder<R>,
}

enum InnerDecoder<R: Read> {
    BMP(BmpDecoder<R>),
    PNG(PngDecoder<R>),
}

#[derive(Clone, Copy, Default)]
struct DirEntry {
    width: u8,
    height: u8,
    color_count: u8,
    reserved: u8,

    num_color_planes: u16,
    bits_per_pixel: u16,

    image_length: u32,
    image_offset: u32,
}

impl<R: Read + Seek> IcoDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(mut r: R) -> ImageResult<IcoDecoder<R>> {
        let entries = read_entries(&mut r)?;
        let entry = best_entry(entries)?;
        let decoder = entry.decoder(r)?;

        Ok(IcoDecoder {
            selected_entry: entry,
            inner_decoder: decoder,
        })
    }
}

fn read_entries<R: Read>(r: &mut R) -> ImageResult<Vec<DirEntry>> {
    let _reserved = r.read_u16::<LittleEndian>()?;
    let _type = r.read_u16::<LittleEndian>()?;
    let count = r.read_u16::<LittleEndian>()?;
    (0..count).map(|_| read_entry(r)).collect()
}

fn read_entry<R: Read>(r: &mut R) -> ImageResult<DirEntry> {
    let mut entry = DirEntry::default();

    entry.width = r.read_u8()?;
    entry.height = r.read_u8()?;
    entry.color_count = r.read_u8()?;
    // Reserved value (not used)
    entry.reserved = r.read_u8()?;

    // This may be either the number of color planes (0 or 1), or the horizontal coordinate
    // of the hotspot for CUR files.
    entry.num_color_planes = r.read_u16::<LittleEndian>()?;
    if entry.num_color_planes > 256 {
        return Err(DecoderError::IcoEntryTooManyPlanesOrHotspot.into());
    }

    // This may be either the bit depth (may be 0 meaning unspecified),
    // or the vertical coordinate of the hotspot for CUR files.
    entry.bits_per_pixel = r.read_u16::<LittleEndian>()?;
    if entry.bits_per_pixel > 256 {
        return Err(DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot.into());
    }

    entry.image_length = r.read_u32::<LittleEndian>()?;
    entry.image_offset = r.read_u32::<LittleEndian>()?;

    Ok(entry)
}

/// Find the entry with the highest (color depth, size).
fn best_entry(mut entries: Vec<DirEntry>) -> ImageResult<DirEntry> {
    let mut best = entries.pop().ok_or(DecoderError::NoEntries)?;

    let mut best_score = (
        best.bits_per_pixel,
        u32::from(best.real_width()) * u32::from(best.real_height()),
    );

    for entry in entries {
        let score = (
            entry.bits_per_pixel,
            u32::from(entry.real_width()) * u32::from(entry.real_height()),
        );
        if score > best_score {
            best = entry;
            best_score = score;
        }
    }
    Ok(best)
}

impl DirEntry {
    fn real_width(&self) -> u16 {
        match self.width {
            0 => 256,
            w => u16::from(w),
        }
    }

    fn real_height(&self) -> u16 {
        match self.height {
            0 => 256,
            h => u16::from(h),
        }
    }

    fn matches_dimensions(&self, width: u32, height: u32) -> bool {
        u32::from(self.real_width()) == width && u32::from(self.real_height()) == height
    }

    fn seek_to_start<R: Read + Seek>(&self, r: &mut R) -> ImageResult<()> {
        r.seek(SeekFrom::Start(u64::from(self.image_offset)))?;
        Ok(())
    }

    fn is_png<R: Read + Seek>(&self, r: &mut R) -> ImageResult<bool> {
        self.seek_to_start(r)?;

        // Read the first 8 bytes to sniff the image.
        let mut signature = [0u8; 8];
        r.read_exact(&mut signature)?;

        Ok(signature == PNG_SIGNATURE)
    }

    fn decoder<R: Read + Seek>(&self, mut r: R) -> ImageResult<InnerDecoder<R>> {
        let is_png = self.is_png(&mut r)?;
        self.seek_to_start(&mut r)?;

        if is_png {
            Ok(PNG(PngDecoder::new(r)?))
        } else {
            Ok(BMP(BmpDecoder::new_with_ico_format(r)?))
        }
    }
}

/// Wrapper struct around a `Cursor<Vec<u8>>`
pub struct IcoReader<R>(Cursor<Vec<u8>>, PhantomData<R>);
impl<R> Read for IcoReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        if self.0.position() == 0 && buf.is_empty() {
            mem::swap(buf, self.0.get_mut());
            Ok(buf.len())
        } else {
            self.0.read_to_end(buf)
        }
    }
}

impl<'a, R: 'a + Read + Seek> ImageDecoder<'a> for IcoDecoder<R> {
    type Reader = IcoReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        match self.inner_decoder {
            BMP(ref decoder) => decoder.dimensions(),
            PNG(ref decoder) => decoder.dimensions(),
        }
    }

    fn color_type(&self) -> ColorType {
        match self.inner_decoder {
            BMP(ref decoder) => decoder.color_type(),
            PNG(ref decoder) => decoder.color_type(),
        }
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        Ok(IcoReader(Cursor::new(image::decoder_to_vec(self)?), PhantomData))
    }

    fn read_image(self, buf: &mut [u8]) -> ImageResult<()> {
        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
        match self.inner_decoder {
            PNG(decoder) => {
                if self.selected_entry.image_length < PNG_SIGNATURE.len() as u32 {
                    return Err(DecoderError::PngShorterThanHeader.into());
                }

                // Check if the image dimensions match the ones in the image data.
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Png,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // Embedded PNG images can only be of the 32BPP RGBA format.
                // https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(DecoderError::PngNotRgba.into());
                }

                decoder.read_image(buf)
            }
            BMP(mut decoder) => {
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Bmp,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // The ICO decoder needs an alpha channel to apply the AND mask.
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                        ImageFormat::Bmp.into(),
                        UnsupportedErrorKind::Color(decoder.color_type().into()),
                    )));
                }

                decoder.read_image_data(buf)?;

                // If there's an AND mask following the image, read and apply it.
                let r = decoder.reader();
                let mask_start = r.seek(SeekFrom::Current(0))?;
                let mask_end =
                    u64::from(self.selected_entry.image_offset + self.selected_entry.image_length);

                if mask_end > mask_start {
                    // A mask row contains 1 bit per pixel, padded to 4 bytes.
                    let mask_row_bytes = ((width + 31) / 32) * 4;
                    let expected_length = u64::from(mask_row_bytes) * u64::from(height);

                    let actual_length = mask_end - mask_start;

                    if actual_length < expected_length {
                        return Err(DecoderError::BmpIcoMaskTooShortForImage.into());
                    }

                    for y in 0..height {
                        let mut x = 0;
                        for _ in 0..mask_row_bytes {
                            // Apply the bits of each byte until we reach the end of the row.
                            let mask_byte = r.read_u8()?;
                            for bit in (0..8).rev() {
                                if x >= width {
                                    break;
                                }
                                if mask_byte & (1 << bit) != 0 {
                                    // Set alpha channel to transparent.
                                    buf[((height - y - 1) * width + x) as usize * 4 + 3] = 0;
                                }
                                x += 1;
                            }
                        }
                    }
                }
                Ok(())
            }
        }
    }
}
<|endoftext|>