<commit_before>use ra_fmt::unwrap_trivial_block;
use ra_syntax::{ast, AstNode, TextRange, T};

use crate::{AssistContext, AssistId, Assists};

// Assist: unwrap_block
//
// This assist removes if...else, for, while and loop control statements to just keep the body.
//
// ```
// fn foo() {
//     if true {<|>
//         println!("foo");
//     }
// }
// ```
// ->
// ```
// fn foo() {
//     println!("foo");
// }
// ```
pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
    let block = ast::BlockExpr::cast(l_curly_token.parent())?;
    let parent = block.syntax().parent()?;
    let assist_id = AssistId("unwrap_block");
    let assist_label = "Unwrap block";
    let parent = ast::Expr::cast(parent)?;

    match parent.clone() {
        ast::Expr::ForExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::LoopExpr(_) => (),
        ast::Expr::IfExpr(if_expr) => {
            let then_branch = if_expr.then_branch()?;
            if then_branch == block {
                if let Some(ancestor) = if_expr.syntax().parent().and_then(ast::IfExpr::cast) {
                    // For `else if` blocks
                    let ancestor_then_branch = ancestor.then_branch()?;

                    let target = then_branch.syntax().text_range();
                    return acc.add(assist_id, assist_label, target, |edit| {
                        let range_to_del_else_if = TextRange::new(
                            ancestor_then_branch.syntax().text_range().end(),
                            l_curly_token.text_range().start(),
                        );
                        let range_to_del_rest = TextRange::new(
                            then_branch.syntax().text_range().end(),
                            if_expr.syntax().text_range().end(),
                        );

                        edit.delete(range_to_del_rest);
                        edit.delete(range_to_del_else_if);
                        edit.replace(
                            target,
                            update_expr_string(then_branch.to_string(), &[' ', '{']),
                        );
                    });
                }
            } else {
                let target = block.syntax().text_range();
                return acc.add(assist_id, assist_label, target, |edit| {
                    let range_to_del = TextRange::new(
                        then_branch.syntax().text_range().end(),
                        l_curly_token.text_range().start(),
                    );

                    edit.delete(range_to_del);
                    edit.replace(target, update_expr_string(block.to_string(), &[' ', '{']));
                });
            }
        }
        _ => return None,
    };

    let unwrapped = unwrap_trivial_block(block);
    let target = unwrapped.syntax().text_range();
    acc.add(assist_id, assist_label, target, |builder| {
        builder.replace(
            parent.syntax().text_range(),
            update_expr_string(unwrapped.to_string(), &[' ', '{', '\n']),
        );
    })
}

fn update_expr_string(expr_str: String, trim_start_pat: &[char]) -> String {
    let expr_string = expr_str.trim_start_matches(trim_start_pat);
    let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
    expr_string_lines.pop(); // Delete last line

    expr_string_lines
        .into_iter()
        .map(|line| line.replacen("    ", "", 1)) // Delete indentation
        .collect::<Vec<String>>()
        .join("\n")
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn simple_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {<|>
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                foo();

                //comment
                bar();
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                } else {<|>
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                }
                println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {<|>
                    println!("bar");
                } else {
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                }
                println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {<|>
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                }
                println!("foo");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {
                    println!("foo");
                } else {<|>
                    println!("else");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {
                    println!("foo");
                }
                println!("else");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested_middle() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {<|>
                    println!("foo");
                } else {
                    println!("else");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                }
                println!("foo");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                bar();<|>
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {
                    if true {<|>
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                for i in 0..5 {
                    foo();

                    //comment
                    bar();
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_loop() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                loop {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_while() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                while true {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_while_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                while true {
                    if true {
                        foo();<|>

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
        );
    }
}
<commit_msg>Unwrap block works with match arms<commit_after>use ra_fmt::unwrap_trivial_block;
use ra_syntax::{
    ast::{
        self,
        edit::{AstNodeEdit, IndentLevel},
    },
    AstNode, TextRange, T,
};

use crate::{AssistContext, AssistId, Assists};

// Assist: unwrap_block
//
// This assist removes if...else, for, while and loop control statements to just keep the body.
//
// ```
// fn foo() {
//     if true {<|>
//         println!("foo");
//     }
// }
// ```
// ->
// ```
// fn foo() {
//     println!("foo");
// }
// ```
pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let assist_id = AssistId("unwrap_block");
    let assist_label = "Unwrap block";

    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
    let mut block = ast::BlockExpr::cast(l_curly_token.parent())?;
    let mut parent = block.syntax().parent()?;
    if ast::MatchArm::can_cast(parent.kind()) {
        parent = parent.ancestors().find(|it| ast::MatchExpr::can_cast(it.kind()))?
    }

    let parent = ast::Expr::cast(parent)?;

    match parent.clone() {
        ast::Expr::ForExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::LoopExpr(_) => (),
        ast::Expr::MatchExpr(_) => block = block.dedent(IndentLevel(1)),
        ast::Expr::IfExpr(if_expr) => {
            let then_branch = if_expr.then_branch()?;
            if then_branch == block {
                if let Some(ancestor) = if_expr.syntax().parent().and_then(ast::IfExpr::cast) {
                    // For `else if` blocks
                    let ancestor_then_branch = ancestor.then_branch()?;

                    let target = then_branch.syntax().text_range();
                    return acc.add(assist_id, assist_label, target, |edit| {
                        let range_to_del_else_if = TextRange::new(
                            ancestor_then_branch.syntax().text_range().end(),
                            l_curly_token.text_range().start(),
                        );
                        let range_to_del_rest = TextRange::new(
                            then_branch.syntax().text_range().end(),
                            if_expr.syntax().text_range().end(),
                        );

                        edit.delete(range_to_del_rest);
                        edit.delete(range_to_del_else_if);
                        edit.replace(
                            target,
                            update_expr_string(then_branch.to_string(), &[' ', '{']),
                        );
                    });
                }
            } else {
                let target = block.syntax().text_range();
                return acc.add(assist_id, assist_label, target, |edit| {
                    let range_to_del = TextRange::new(
                        then_branch.syntax().text_range().end(),
                        l_curly_token.text_range().start(),
                    );

                    edit.delete(range_to_del);
                    edit.replace(target, update_expr_string(block.to_string(), &[' ', '{']));
                });
            }
        }
        _ => return None,
    };

    let unwrapped = unwrap_trivial_block(block);
    let target = unwrapped.syntax().text_range();
    acc.add(assist_id, assist_label, target, |builder| {
        builder.replace(
            parent.syntax().text_range(),
            update_expr_string(unwrapped.to_string(), &[' ', '{', '\n']),
        );
    })
}

fn update_expr_string(expr_str: String, trim_start_pat: &[char]) -> String {
    let expr_string = expr_str.trim_start_matches(trim_start_pat);
    let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
    expr_string_lines.pop(); // Delete last line

    expr_string_lines
        .into_iter()
        .map(|line| line.replacen("    ", "", 1)) // Delete indentation
        .collect::<Vec<String>>()
        .join("\n")
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn simple_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {<|>
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                foo();

                //comment
                bar();
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                } else {<|>
                    println!("bar");
                }
            }
            "#,
            r#"
            fn main() {
                bar();
                if true {
                    foo();

                    //comment
                    bar();
                }
                println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {<|>
                    println!("bar");
                } else {
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                }
                println!("bar");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {<|>
                    println!("foo");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                }
                println!("foo");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested_else() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {
                    println!("foo");
                } else {<|>
                    println!("else");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {
                    println!("foo");
                }
                println!("else");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_else_if_nested_middle() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                } else if true {<|>
                    println!("foo");
                } else {
                    println!("else");
                }
            }
            "#,
            r#"
            fn main() {
                //bar();
                if true {
                    println!("true");

                    //comment
                    //bar();
                } else if false {
                    println!("bar");
                }
                println!("foo");
            }
            "#,
        );
    }

    #[test]
    fn simple_if_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                bar();<|>
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_if_in_for() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                for i in 0..5 {
                    if true {<|>
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                for i in 0..5 {
                    foo();

                    //comment
                    bar();
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_loop() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                loop {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn simple_while() {
        check_assist(
            unwrap_block,
            r#"
            fn main() {
                while true {<|>
                    if true {
                        foo();

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
            r#"
            fn main() {
                if true {
                    foo();

                    //comment
                    bar();
                } else {
                    println!("bar");
                }
            }
            "#,
        );
    }

    #[test]
    fn unwrap_match_arm() {
        check_assist(
            unwrap_block,
            r#"
fn main() {
    match rel_path {
        Ok(rel_path) => {<|>
            let rel_path = RelativePathBuf::from_path(rel_path).ok()?;
            Some((*id, rel_path))
        }
        Err(_) => None,
    }
}
"#,
            r#"
fn main() {
    let rel_path = RelativePathBuf::from_path(rel_path).ok()?;
    Some((*id, rel_path))
}
"#,
        );
    }

    #[test]
    fn simple_if_in_while_bad_cursor_position() {
        check_assist_not_applicable(
            unwrap_block,
            r#"
            fn main() {
                while true {
                    if true {
                        foo();<|>

                        //comment
                        bar();
                    } else {
                        println!("bar");
                    }
                }
            }
            "#,
        );
    }
}
<|endoftext|>


<commit_before>use calx_alg::{compact_bits_by_2, noise, spread_bits_by_2};
use calx_grid::{Dir6, GridNode, HexGeom};
use euclid::{Vector2D, vec2};
use std::num::Wrapping;
use std::ops::{Add, Sub};
use terraform::TerrainQuery;

pub const SECTOR_WIDTH: i32 = 40;
pub const SECTOR_HEIGHT: i32 = 22;

/// Unambiguous location in the game world.
#[derive(Copy, Eq, PartialEq, Clone, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
pub struct Location {
    pub x: i8,
    pub y: i8,
    pub z: i8,
}

/// The type for a unique location in the game world.
///
/// IMPORTANT: Be careful where you use the simple "location + vec" algebra. That does not take
/// portals into account, and will usually cause unwanted effects near them in high-level code.
/// Most high-level logic should use `Location::jump` to displace locations. This will correctly
/// traverse portals.
impl Location {
    pub fn origin() -> Location { Location { x: 0, y: 0, z: 0 } }

    pub fn new(x: i8, y: i8, z: i8) -> Location { Location { x, y, z } }

    /// Construct a Location from a Morton code representation.
    ///
    /// Use the representation generated with `to_morton`. The odd low 16 bits are compacted to x
    /// value, the even low 16 bits to z and the first 8 of the high 16 bits become z.
    pub fn from_morton(morton_code: u32) -> Location {
        let xy = morton_code & 0xffff_ffff;
        let x = compact_bits_by_2(xy) as u8;
        let y = compact_bits_by_2(xy >> 1) as u8;
        let z = (morton_code >> 16) as u8;

        unsafe {
            Location {
                x: ::std::mem::transmute(x),
                y: ::std::mem::transmute(y),
                z: ::std::mem::transmute(z),
            }
        }
    }

    /// Turn the Location to a Morton code value.
    ///
    /// Spatially close locations are often numerically close in Morton codes, these are useful for
    /// quadtree-like structures.
    pub fn to_morton(&self) -> u32 {
        let mut ret = 0;
        let x: u8 = unsafe { ::std::mem::transmute(self.x) };
        let y: u8 = unsafe { ::std::mem::transmute(self.y) };
        let z: u8 = unsafe { ::std::mem::transmute(self.z) };
        ret ^= spread_bits_by_2(x as u32);
        ret ^= spread_bits_by_2(y as u32) << 1;
        ret ^= (z as u32) << 16;
        ret
    }

    /// Vector pointing from this location into the other one if the locations
    /// are on the same Euclidean plane.
    pub fn v2_at(&self, other: Location) -> Option<Vector2D<i32>> {
        if self.z != other.z {
            return None;
        }
        Some(
            vec2(other.x as i32, other.y as i32) - vec2(self.x as i32, self.y as i32),
        )
    }

    /// Hex distance from this location to the other one, if applicable.
    pub fn distance_from(&self, other: Location) -> Option<i32> {
        if let Some(v) = self.v2_at(other) {
            Some(v.hex_dist())
        } else {
            None
        }
    }

    /// Distance that defaults to max integer value for separate zones.
    ///
    /// Can be used for situations that want a straightforward metric function like A* search.
    pub fn metric_distance(&self, other: Location) -> i32 {
        self.distance_from(other).unwrap_or(i32::max_value())
    }

    pub fn dir6_towards(&self, other: Location) -> Option<Dir6> {
        if let Some(v) = self.v2_at(other) {
            Some(Dir6::from_v2(v))
        } else {
            None
        }
    }

    /// A pseudorandom value in [-1.0, 1.0] corresponding to this specific location.
    ///
    /// Is always the same for the same `Location`.
    pub fn noise(&self) -> f32 { noise(self.x as i32 + self.y as i32 * 59 + self.z as i32 * 919) }

    /// Offset location and follow any portals in target site.
    pub fn jump<T: TerrainQuery, V: Into<Vector2D<i32>> + Sized>(
        self,
        ctx: &T,
        offset: V,
    ) -> Location {
        let loc = self + offset.into();
        ctx.portal(loc).unwrap_or(loc)
    }

    /// Return `Sector` this location is in.
    pub fn sector(self) -> Sector {
        let (u, v) = self.to_rect_coords();

        Sector::new(
            (u as f32 / SECTOR_WIDTH as f32).floor() as i8,
            (v as f32 / SECTOR_HEIGHT as f32).floor() as i8,
            self.z,
        )
    }

    /// Map location's x, y to rectangular (offset) coordinates.
    pub fn to_rect_coords(self) -> (i32, i32) {
        let u = self.x as i32 - self.y as i32;
        let v = ((self.x as f32 + self.y as f32) / 2.0).floor() as i32;
        (u, v)
    }

    pub fn from_rect_coords(u: i32, v: i32, z: i8) -> Location {
        // Yeah I don't know either how you're supposed to come up with the right ceil/floor
        // juggling, just tweaked it around until it passed all the unit tests.
        let half_u = u as f32 / 2.0;
        Location::new(
            (half_u.ceil() as i32 + v) as i8,
            (v - half_u.floor() as i32) as i8,
            z,
        )
    }
}

impl<V: Into<Vector2D<i32>>> Add<V> for Location {
    type Output = Location;
    fn add(self, other: V) -> Location {
        let other = other.into();
        Location {
            x: (self.x as i32 + other.x) as i8,
            y: (self.y as i32 + other.y) as i8,
            z: self.z,
        }
    }
}

impl Add<Portal> for Location {
    type Output = Location;
    fn add(self, other: Portal) -> Location {
        Location {
            x: (Wrapping(self.x) + Wrapping(other.dx)).0,
            y: (Wrapping(self.y) + Wrapping(other.dy)).0,
            z: other.z,
        }
    }
}

impl<V: Into<Vector2D<i32>>> Sub<V> for Location {
    type Output = Location;
    fn sub(self, other: V) -> Location {
        let other = other.into();
        Location {
            x: (self.x as i32 - other.x) as i8,
            y: (self.y as i32 - other.y) as i8,
            z: self.z,
        }
    }
}

impl GridNode for Location {
    fn neighbors(&self) -> Vec<Location> { Dir6::iter().map(|d| *self + d.to_v2()).collect() }
}

#[derive(Copy, Eq, PartialEq, Clone, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
pub struct Portal {
    pub dx: i8,
    pub dy: i8,
    pub z: i8,
}

impl Portal {
    pub fn new(from: Location, to: Location) -> Portal {
        Portal {
            dx: (Wrapping(to.x) - Wrapping(from.x)).0,
            dy: (Wrapping(to.y) - Wrapping(from.y)).0,
            z: to.z,
        }
    }
}

impl Add<Portal> for Portal {
    type Output = Portal;
    fn add(self, other: Portal) -> Portal {
        Portal {
            dx: (Wrapping(self.dx) + Wrapping(other.dx)).0,
            dy: (Wrapping(self.dy) + Wrapping(other.dy)).0,
            z: other.z,
        }
    }
}

/// Non-scrolling screen.
///
/// A sector represents a rectangular chunk of locations that fit on the visual screen. Sector
/// coordinates form their own sector space that tiles the location space with sectors.
#[derive(Copy, Eq, PartialEq, Clone, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
pub struct Sector {
    pub x: i8,
    pub y: i8,
    pub z: i8,
}

impl Sector {
    pub fn new(x: i8, y: i8, z: i8) -> Sector { Sector { x, y, z } }

    pub fn origin(self) -> Location { self.rect_coord_loc(0, 0) }

    pub fn rect_coord_loc(self, u: i32, v: i32) -> Location {
        Location::from_rect_coords(
            self.x as i32 * SECTOR_WIDTH + u,
            self.y as i32 * SECTOR_HEIGHT + v,
            self.z,
        )
    }

    /// Center location for this sector.
    ///
    /// Usually you want the camera positioned here.
    pub fn center(self) -> Location {
        // XXX: If the width/height are even (as they currently are), there isn't a centered cell.
        self.rect_coord_loc(SECTOR_WIDTH / 2, SECTOR_HEIGHT / 2)
    }

    // TODO: Use impl Trait instead of box for return type once it's stable.
    pub fn iter(self) -> Box<Iterator<Item = Location>> {
        let n = SECTOR_WIDTH * SECTOR_HEIGHT;
        let pitch = SECTOR_WIDTH;
        Box::new((0..n).map(
            move |i| self.rect_coord_loc(i % pitch, i / pitch),
        ))
    }
}

#[cfg(test)]
mod test {
    use super::{Location, Sector};
    use euclid::vec2;

    #[test]
    fn test_wraparound() {
        let l1 = Location::new(0, 0, 0);
        let l2 = l1 + vec2(300, 300);
        assert_eq!((44, 44), (l2.x, l2.y));
    }

    #[test]
    fn test_morton() {
        use rand::{self, Rng};
        let mut rng = rand::thread_rng();

        for _ in 0..1000 {
            let x = rng.gen::<u32>() & 0xff_ffff;
            assert_eq!(x, Location::from_morton(x).to_morton());
        }
    }

    #[test]
    fn test_location_to_sector() {
        let s = Sector::new(0, 0, 0);
        assert_eq!(s.origin(), Location::new(0, 0, 0));

        // Sector division near origin
        assert_eq!(Location::new(0, 0, 0).sector(), Sector::new(0, 0, 0));
        assert_eq!(Location::new(-1, -1, 0).sector(), Sector::new(0, -1, 0));
        assert_eq!(Location::new(0, 1, 0).sector(), Sector::new(-1, 0, 0));
        assert_eq!(Location::new(-1, 0, 0).sector(), Sector::new(-1, -1, 0));

        for y in -100..100 {
            for x in -100..100 {
                let loc = Location::new(x, y, 0);
                let (u, v) = loc.to_rect_coords();
                assert_eq!(
                    loc,
                    Location::from_rect_coords(u, v, loc.z),
                    "u: {}, v: {}",
                    u,
                    v
                );

                assert!(
                    loc.sector().iter().find(|&x| x == loc).is_some(),
                    format!("{:?} not found in sector {:?}", loc, loc.sector())
                );
            }
        }
    }

    #[test]
    fn test_sector_iter() {
        let s = Sector::new(0, 0, 0);

        for loc in s.iter() {
            assert_eq!(s, loc.sector(), "Location: {:?}", loc);
        }
    }
}
<commit_msg>Need to adjust it like this for the current camera<commit_after>use calx_alg::{compact_bits_by_2, noise, spread_bits_by_2};
use calx_grid::{Dir6, GridNode, HexGeom};
use euclid::{Vector2D, vec2};
use std::num::Wrapping;
use std::ops::{Add, Sub};
use terraform::TerrainQuery;

pub const SECTOR_WIDTH: i32 = 40;
pub const SECTOR_HEIGHT: i32 = 22;

/// Unambiguous location in the game world.
#[derive(Copy, Eq, PartialEq, Clone, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
pub struct Location {
    pub x: i8,
    pub y: i8,
    pub z: i8,
}

/// The type for a unique location in the game world.
///
/// IMPORTANT: Be careful where you use the simple "location + vec" algebra. That does not take
/// portals into account, and will usually cause unwanted effects near them in high-level code.
/// Most high-level logic should use `Location::jump` to displace locations. This will correctly
/// traverse portals.
impl Location {
    pub fn origin() -> Location { Location { x: 0, y: 0, z: 0 } }

    pub fn new(x: i8, y: i8, z: i8) -> Location { Location { x, y, z } }

    /// Construct a Location from a Morton code representation.
    ///
    /// Use the representation generated with `to_morton`. The odd low 16 bits are compacted to x
    /// value, the even low 16 bits to z and the first 8 of the high 16 bits become z.
    pub fn from_morton(morton_code: u32) -> Location {
        let xy = morton_code & 0xffff_ffff;
        let x = compact_bits_by_2(xy) as u8;
        let y = compact_bits_by_2(xy >> 1) as u8;
        let z = (morton_code >> 16) as u8;

        unsafe {
            Location {
                x: ::std::mem::transmute(x),
                y: ::std::mem::transmute(y),
                z: ::std::mem::transmute(z),
            }
        }
    }

    /// Turn the Location to a Morton code value.
    ///
    /// Spatially close locations are often numerically close in Morton codes, these are useful for
    /// quadtree-like structures.
    pub fn to_morton(&self) -> u32 {
        let mut ret = 0;
        let x: u8 = unsafe { ::std::mem::transmute(self.x) };
        let y: u8 = unsafe { ::std::mem::transmute(self.y) };
        let z: u8 = unsafe { ::std::mem::transmute(self.z) };
        ret ^= spread_bits_by_2(x as u32);
        ret ^= spread_bits_by_2(y as u32) << 1;
        ret ^= (z as u32) << 16;
        ret
    }

    /// Vector pointing from this location into the other one if the locations
    /// are on the same Euclidean plane.
    pub fn v2_at(&self, other: Location) -> Option<Vector2D<i32>> {
        if self.z != other.z {
            return None;
        }
        Some(
            vec2(other.x as i32, other.y as i32) - vec2(self.x as i32, self.y as i32),
        )
    }

    /// Hex distance from this location to the other one, if applicable.
    pub fn distance_from(&self, other: Location) -> Option<i32> {
        if let Some(v) = self.v2_at(other) {
            Some(v.hex_dist())
        } else {
            None
        }
    }

    /// Distance that defaults to max integer value for separate zones.
    ///
    /// Can be used for situations that want a straightforward metric function like A* search.
    pub fn metric_distance(&self, other: Location) -> i32 {
        self.distance_from(other).unwrap_or(i32::max_value())
    }

    pub fn dir6_towards(&self, other: Location) -> Option<Dir6> {
        if let Some(v) = self.v2_at(other) {
            Some(Dir6::from_v2(v))
        } else {
            None
        }
    }

    /// A pseudorandom value in [-1.0, 1.0] corresponding to this specific location.
    ///
    /// Is always the same for the same `Location`.
    pub fn noise(&self) -> f32 { noise(self.x as i32 + self.y as i32 * 59 + self.z as i32 * 919) }

    /// Offset location and follow any portals in target site.
    pub fn jump<T: TerrainQuery, V: Into<Vector2D<i32>> + Sized>(
        self,
        ctx: &T,
        offset: V,
    ) -> Location {
        let loc = self + offset.into();
        ctx.portal(loc).unwrap_or(loc)
    }

    /// Return `Sector` this location is in.
    pub fn sector(self) -> Sector {
        let (u, v) = self.to_rect_coords();

        Sector::new(
            (u as f32 / SECTOR_WIDTH as f32).floor() as i8,
            (v as f32 / SECTOR_HEIGHT as f32).floor() as i8,
            self.z,
        )
    }

    /// Map location's x, y to rectangular (offset) coordinates.
    pub fn to_rect_coords(self) -> (i32, i32) {
        let u = self.x as i32 - self.y as i32;
        let v = ((self.x as f32 + self.y as f32) / 2.0).floor() as i32;
        (u, v)
    }

    pub fn from_rect_coords(u: i32, v: i32, z: i8) -> Location {
        // Yeah I don't know either how you're supposed to come up with the right ceil/floor
        // juggling, just tweaked it around until it passed all the unit tests.
        let half_u = u as f32 / 2.0;
        Location::new(
            (half_u.ceil() as i32 + v) as i8,
            (v - half_u.floor() as i32) as i8,
            z,
        )
    }
}

impl<V: Into<Vector2D<i32>>> Add<V> for Location {
    type Output = Location;
    fn add(self, other: V) -> Location {
        let other = other.into();
        Location {
            x: (self.x as i32 + other.x) as i8,
            y: (self.y as i32 + other.y) as i8,
            z: self.z,
        }
    }
}

impl Add<Portal> for Location {
    type Output = Location;
    fn add(self, other: Portal) -> Location {
        Location {
            x: (Wrapping(self.x) + Wrapping(other.dx)).0,
            y: (Wrapping(self.y) + Wrapping(other.dy)).0,
            z: other.z,
        }
    }
}

impl<V: Into<Vector2D<i32>>> Sub<V> for Location {
    type Output = Location;
    fn sub(self, other: V) -> Location {
        let other = other.into();
        Location {
            x: (self.x as i32 - other.x) as i8,
            y: (self.y as i32 - other.y) as i8,
            z: self.z,
        }
    }
}

impl GridNode for Location {
    fn neighbors(&self) -> Vec<Location> { Dir6::iter().map(|d| *self + d.to_v2()).collect() }
}

#[derive(Copy, Eq, PartialEq, Clone, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
pub struct Portal {
    pub dx: i8,
    pub dy: i8,
    pub z: i8,
}

impl Portal {
    pub fn new(from: Location, to: Location) -> Portal {
        Portal {
            dx: (Wrapping(to.x) - Wrapping(from.x)).0,
            dy: (Wrapping(to.y) - Wrapping(from.y)).0,
            z: to.z,
        }
    }
}

impl Add<Portal> for Portal {
    type Output = Portal;
    fn add(self, other: Portal) -> Portal {
        Portal {
            dx: (Wrapping(self.dx) + Wrapping(other.dx)).0,
            dy: (Wrapping(self.dy) + Wrapping(other.dy)).0,
            z: other.z,
        }
    }
}

/// Non-scrolling screen.
///
/// A sector represents a rectangular chunk of locations that fit on the visual screen. Sector
/// coordinates form their own sector space that tiles the location space with sectors.
#[derive(Copy, Eq, PartialEq, Clone, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
pub struct Sector {
    pub x: i8,
    pub y: i8,
    pub z: i8,
}

impl Sector {
    pub fn new(x: i8, y: i8, z: i8) -> Sector { Sector { x, y, z } }

    pub fn origin(self) -> Location { self.rect_coord_loc(0, 0) }

    pub fn rect_coord_loc(self, u: i32, v: i32) -> Location {
        Location::from_rect_coords(
            self.x as i32 * SECTOR_WIDTH + u,
            self.y as i32 * SECTOR_HEIGHT + v,
            self.z,
        )
    }

    /// Center location for this sector.
    ///
    /// Usually you want the camera positioned here.
    pub fn center(self) -> Location {
        // XXX: If the width/height are even (as they currently are), there isn't a centered cell.
        self.rect_coord_loc(SECTOR_WIDTH / 2 - 1, SECTOR_HEIGHT / 2 - 1)
    }

    // TODO: Use impl Trait instead of box for return type once it's stable.
    pub fn iter(self) -> Box<Iterator<Item = Location>> {
        let n = SECTOR_WIDTH * SECTOR_HEIGHT;
        let pitch = SECTOR_WIDTH;
        Box::new((0..n).map(
            move |i| self.rect_coord_loc(i % pitch, i / pitch),
        ))
    }
}

#[cfg(test)]
mod test {
    use super::{Location, Sector};
    use euclid::vec2;

    #[test]
    fn test_wraparound() {
        let l1 = Location::new(0, 0, 0);
        let l2 = l1 + vec2(300, 300);
        assert_eq!((44, 44), (l2.x, l2.y));
    }

    #[test]
    fn test_morton() {
        use rand::{self, Rng};
        let mut rng = rand::thread_rng();

        for _ in 0..1000 {
            let x = rng.gen::<u32>() & 0xff_ffff;
            assert_eq!(x, Location::from_morton(x).to_morton());
        }
    }

    #[test]
    fn test_location_to_sector() {
        let s = Sector::new(0, 0, 0);
        assert_eq!(s.origin(), Location::new(0, 0, 0));

        // Sector division near origin
        assert_eq!(Location::new(0, 0, 0).sector(), Sector::new(0, 0, 0));
        assert_eq!(Location::new(-1, -1, 0).sector(), Sector::new(0, -1, 0));
        assert_eq!(Location::new(0, 1, 0).sector(), Sector::new(-1, 0, 0));
        assert_eq!(Location::new(-1, 0, 0).sector(), Sector::new(-1, -1, 0));

        for y in -100..100 {
            for x in -100..100 {
                let loc = Location::new(x, y, 0);
                let (u, v) = loc.to_rect_coords();
                assert_eq!(
                    loc,
                    Location::from_rect_coords(u, v, loc.z),
                    "u: {}, v: {}",
                    u,
                    v
                );

                assert!(
                    loc.sector().iter().find(|&x| x == loc).is_some(),
                    format!("{:?} not found in sector {:?}", loc, loc.sector())
                );
            }
        }
    }

    #[test]
    fn test_sector_iter() {
        let s = Sector::new(0, 0, 0);

        for loc in s.iter() {
            assert_eq!(s, loc.sector(), "Location: {:?}", loc);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! Raspberry Pi system-related tools.
//!
//! Use [`DeviceInfo`] to identify what Raspberry Pi model and SoC the software is
//! running on. This information is used internally to calculate the correct memory
//! locations for the various BCM283x peripherals.
//!
//! [`DeviceInfo`]: struct.DeviceInfo.html

use std::fmt;
use std::fs;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::result;

const PERIPHERAL_BASE_RPI: u32 = 0x2000_0000;
const PERIPHERAL_BASE_RPI2: u32 = 0x3f00_0000;
const GPIO_OFFSET: u32 = 0x20_0000;

quick_error! {
/// Errors that can occur when trying to identify the Raspberry Pi hardware.
    #[derive(Debug)]
    pub enum Error {
/// Unknown model.
///
/// `DeviceInfo` was unable to identify the Raspberry Pi model based on the
/// contents of `/proc/cpuinfo`, `/sys/firmware/devicetree/base/compatible`
/// and `/sys/firmware/devicetree/base/model`.
        UnknownModel { description("unknown Raspberry Pi model") }
    }
}

/// Result type returned from methods that can have `system::Error`s.
pub type Result<T> = result::Result<T, Error>;

/// Identifiable Raspberry Pi models.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Model {
    RaspberryPiA,
    RaspberryPiAPlus,
    RaspberryPiBRev1,
    RaspberryPiBRev2,
    RaspberryPiBPlus,
    RaspberryPi2B,
    RaspberryPi3B,
    RaspberryPi3BPlus,
    RaspberryPiComputeModule,
    RaspberryPiComputeModule3,
    RaspberryPiZero,
    RaspberryPiZeroW,
}

impl fmt::Display for Model {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Model::RaspberryPiA => write!(f, "Raspberry Pi A"),
            Model::RaspberryPiAPlus => write!(f, "Raspberry Pi A+"),
            Model::RaspberryPiBRev1 => write!(f, "Raspberry Pi B Rev 1"),
            Model::RaspberryPiBRev2 => write!(f, "Raspberry Pi B Rev 2"),
            Model::RaspberryPiBPlus => write!(f, "Raspberry Pi B+"),
            Model::RaspberryPi2B => write!(f, "Raspberry Pi 2 B"),
            Model::RaspberryPi3B => write!(f, "Raspberry Pi 3 B"),
            Model::RaspberryPi3BPlus => write!(f, "Raspberry Pi 3 B+"),
            Model::RaspberryPiComputeModule => write!(f, "Raspberry Pi Compute Module"),
            Model::RaspberryPiComputeModule3 => write!(f, "Raspberry Pi Compute Module 3"),
            Model::RaspberryPiZero => write!(f, "Raspberry Pi Zero"),
            Model::RaspberryPiZeroW => write!(f, "Raspberry Pi Zero W"),
        }
    }
}

/// Identifiable Raspberry Pi SoCs.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum SoC {
    Bcm2835,
    Bcm2836,
    Bcm2837A1,
    Bcm2837B0,
}

impl fmt::Display for SoC {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            SoC::Bcm2835 => write!(f, "BCM2835"),
            SoC::Bcm2836 => write!(f, "BCM2836"),
            SoC::Bcm2837A1 => write!(f, "BCM2837A1"),
            SoC::Bcm2837B0 => write!(f, "BCM2837B0"),
        }
    }
}

// Identify Pi model based on /proc/cpuinfo
fn parse_proc_cpuinfo() -> Result<Model> {
    let proc_cpuinfo = BufReader::new(match File::open("/proc/cpuinfo") {
        Err(_) => return Err(Error::UnknownModel),
        Ok(file) => file,
    });

    let mut hardware: String = String::new();
    let mut revision: String = String::new();
    for line_result in proc_cpuinfo.lines() {
        if let Ok(line) = line_result {
            if line.starts_with("Hardware\t: ") {
                hardware = String::from(&line[11..]);
            } else if line.starts_with("Revision\t: ") {
                revision = String::from(&line[11..]).to_lowercase();
            }
        }
    }

    // Return an error if we don't recognize the SoC. This check is
    // done to prevent accidentally identifying a non-Pi SBC as a Pi
    // solely based on the revision field.
    match &hardware[..] {
        "BCM2708" | "BCM2835" | "BCM2709" | "BCM2836" | "BCM2710" | "BCM2837" | "BCM2837A1"
        | "BCM2837B0" => {}
        _ => return Err(Error::UnknownModel),
    }

    let model = if (revision.len() == 4) || (revision.len() == 8) {
        // Older revisions are 4 characters long, or 8 if they've been over-volted
        match &revision[revision.len() - 4..] {
            "0007" | "0008" | "0009" | "0015" => Model::RaspberryPiA,
            "Beta" | "0002" | "0003" => Model::RaspberryPiBRev1,
            "0004" | "0005" | "0006" | "000d" | "000e" | "000f" => Model::RaspberryPiBRev2,
            "0012" => Model::RaspberryPiAPlus,
            "0010" | "0013" => Model::RaspberryPiBPlus,
            "0011" | "0014" => Model::RaspberryPiComputeModule,
            _ => return Err(Error::UnknownModel),
        }
    } else if revision.len() >= 6 {
        // Newer revisions consist of at least 6 characters
        match &revision[revision.len() - 3..revision.len() - 1] {
            "00" => Model::RaspberryPiA,
            "01" => Model::RaspberryPiBRev2,
            "02" => Model::RaspberryPiAPlus,
            "03" => Model::RaspberryPiBPlus,
            "04" => Model::RaspberryPi2B,
            "06" => Model::RaspberryPiComputeModule,
            "08" => Model::RaspberryPi3B,
            "09" => Model::RaspberryPiZero,
            "0a" => Model::RaspberryPiComputeModule3,
            "0c" => Model::RaspberryPiZeroW,
            "0d" => Model::RaspberryPi3BPlus,
            _ => return Err(Error::UnknownModel),
        }
    } else {
        return Err(Error::UnknownModel);
    };

    Ok(model)
}

// Identify Pi model based on /sys/firmware/devicetree/base/compatible
fn parse_firmware_compatible() -> Result<Model> {
    let compatible = match fs::read_to_string("/sys/firmware/devicetree/base/compatible") {
        Err(_) => return Err(Error::UnknownModel),
        Ok(s) => s,
    };

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    for compid in compatible.split("\0") {
        match compid {
            "raspberrypi,model-b-i2c0" => return Ok(Model::RaspberryPiBRev1),
            "raspberrypi,model-b" => return Ok(Model::RaspberryPiBRev1),
            "raspberrypi,model-a" => return Ok(Model::RaspberryPiA),
            "raspberrypi,model-b-rev2" => return Ok(Model::RaspberryPiBRev2),
            "raspberrypi,model-a-plus" => return Ok(Model::RaspberryPiAPlus),
            "raspberrypi,model-b-plus" => return Ok(Model::RaspberryPiBPlus),
            "raspberrypi,2-model-b" => return Ok(Model::RaspberryPi2B),
            "raspberrypi,compute-module" => return Ok(Model::RaspberryPiComputeModule),
            "raspberrypi,3-model-b" => return Ok(Model::RaspberryPi3B),
            "raspberrypi,model-zero" => return Ok(Model::RaspberryPiZero),
            "raspberrypi,3-compute-module" => return Ok(Model::RaspberryPiComputeModule3),
            "raspberrypi,model-zero-w" => return Ok(Model::RaspberryPiZeroW),
            "raspberrypi,3-model-b-plus" => return Ok(Model::RaspberryPi3BPlus),
            _ => (),
        }
    }

    Err(Error::UnknownModel)
}

// Identify Pi model based on /sys/firmware/devicetree/base/model
fn parse_firmware_model() -> Result<Model> {
    let model = match fs::read_to_string("/sys/firmware/devicetree/base/model") {
        Err(_) => return Err(Error::UnknownModel),
        Ok(buffer) => if let Some(rev_idx) = buffer.find(" Rev ") {
            // We don't want to strip rev2 here
            buffer[0..rev_idx].to_owned()
        } else {
            buffer
        },
    };

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    match &model[..] {
        "Raspberry Pi Model B (no P5)" => return Ok(Model::RaspberryPiBRev1),
        "Raspberry Pi Model B" => return Ok(Model::RaspberryPiBRev1),
        "Raspberry Pi Model A" => return Ok(Model::RaspberryPiA),
        "Raspberry Pi Model B rev2" => return Ok(Model::RaspberryPiBRev2),
        "Raspberry Pi Model A+" => return Ok(Model::RaspberryPiAPlus),
        "Raspberry Pi Model B+" => return Ok(Model::RaspberryPiBPlus),
        "Raspberry Pi 2 Model B" => return Ok(Model::RaspberryPi2B),
        "Raspberry Pi Compute Module" => return Ok(Model::RaspberryPiComputeModule),
        "Raspberry Pi 3 Model B" => return Ok(Model::RaspberryPi3B),
        "Raspberry Pi Zero" => return Ok(Model::RaspberryPiZero),
        "Raspberry Pi Compute Module 3" => return Ok(Model::RaspberryPiComputeModule3),
        "Raspberry Pi Zero W" => return Ok(Model::RaspberryPiZeroW),
        "Raspberry Pi 3 Model B+" => return Ok(Model::RaspberryPi3BPlus),
        _ => return Err(Error::UnknownModel),
    }
}

/// Retrieves Raspberry Pi device information.
#[derive(Debug, PartialEq, Copy, Clone)]
pub struct DeviceInfo {
    model: Model,
    soc: SoC,
    peripheral_base: u32,
    gpio_offset: u32,
}

impl DeviceInfo {
    /// Constructs a new `DeviceInfo`.
    ///
    /// `new` parses the contents of `/proc/cpuinfo` to identify the Raspberry
    /// Pi's model and SoC.
    pub fn new() -> Result<DeviceInfo> {
        let model =
            parse_proc_cpuinfo().or(parse_firmware_compatible().or(parse_firmware_model()))?;

        // Set SoC and memory offsets based on model
        match model {
            Model::RaspberryPiA
            | Model::RaspberryPiAPlus
            | Model::RaspberryPiBRev1
            | Model::RaspberryPiBRev2
            | Model::RaspberryPiBPlus
            | Model::RaspberryPiComputeModule
            | Model::RaspberryPiZero
            | Model::RaspberryPiZeroW => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2835,
                peripheral_base: PERIPHERAL_BASE_RPI,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi2B => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2836,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3B | Model::RaspberryPiComputeModule3 => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837A1,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3BPlus => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837B0,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
        }
    }

    /// Returns the Raspberry Pi's model.
    pub fn model(&self) -> Model {
        self.model
    }

    /// Returns the Raspberry Pi's SoC.
    pub fn soc(&self) -> SoC {
        self.soc
    }

    /// Returns the base memory address for the BCM283x peripherals.
    pub fn peripheral_base(&self) -> u32 {
        self.peripheral_base
    }

    /// Returns the offset from the base memory address for the GPIO section.
    pub fn gpio_offset(&self) -> u32 {
        self.gpio_offset
    }
}
<commit_msg>Improve code formatting<commit_after>// Copyright (c) 2017-2018 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! Raspberry Pi system-related tools.
//!
//! Use [`DeviceInfo`] to identify what Raspberry Pi model and SoC the software is
//! running on. This information is used internally to calculate the correct memory
//! locations for the various BCM283x peripherals.
//!
//! [`DeviceInfo`]: struct.DeviceInfo.html

use std::fmt;
use std::fs;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::result;

const PERIPHERAL_BASE_RPI: u32 = 0x2000_0000;
const PERIPHERAL_BASE_RPI2: u32 = 0x3f00_0000;
const GPIO_OFFSET: u32 = 0x20_0000;

quick_error! {
/// Errors that can occur when trying to identify the Raspberry Pi hardware.
    #[derive(Debug)]
    pub enum Error {
/// Unknown model.
///
/// `DeviceInfo` was unable to identify the Raspberry Pi model based on the
/// contents of `/proc/cpuinfo`, `/sys/firmware/devicetree/base/compatible`
/// and `/sys/firmware/devicetree/base/model`.
        UnknownModel { description("unknown Raspberry Pi model") }
    }
}

/// Result type returned from methods that can have `system::Error`s.
pub type Result<T> = result::Result<T, Error>;

/// Identifiable Raspberry Pi models.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Model {
    RaspberryPiA,
    RaspberryPiAPlus,
    RaspberryPiBRev1,
    RaspberryPiBRev2,
    RaspberryPiBPlus,
    RaspberryPi2B,
    RaspberryPi3B,
    RaspberryPi3BPlus,
    RaspberryPiComputeModule,
    RaspberryPiComputeModule3,
    RaspberryPiZero,
    RaspberryPiZeroW,
}

impl fmt::Display for Model {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Model::RaspberryPiA => write!(f, "Raspberry Pi A"),
            Model::RaspberryPiAPlus => write!(f, "Raspberry Pi A+"),
            Model::RaspberryPiBRev1 => write!(f, "Raspberry Pi B Rev 1"),
            Model::RaspberryPiBRev2 => write!(f, "Raspberry Pi B Rev 2"),
            Model::RaspberryPiBPlus => write!(f, "Raspberry Pi B+"),
            Model::RaspberryPi2B => write!(f, "Raspberry Pi 2 B"),
            Model::RaspberryPi3B => write!(f, "Raspberry Pi 3 B"),
            Model::RaspberryPi3BPlus => write!(f, "Raspberry Pi 3 B+"),
            Model::RaspberryPiComputeModule => write!(f, "Raspberry Pi Compute Module"),
            Model::RaspberryPiComputeModule3 => write!(f, "Raspberry Pi Compute Module 3"),
            Model::RaspberryPiZero => write!(f, "Raspberry Pi Zero"),
            Model::RaspberryPiZeroW => write!(f, "Raspberry Pi Zero W"),
        }
    }
}

/// Identifiable Raspberry Pi SoCs.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum SoC {
    Bcm2835,
    Bcm2836,
    Bcm2837A1,
    Bcm2837B0,
}

impl fmt::Display for SoC {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            SoC::Bcm2835 => write!(f, "BCM2835"),
            SoC::Bcm2836 => write!(f, "BCM2836"),
            SoC::Bcm2837A1 => write!(f, "BCM2837A1"),
            SoC::Bcm2837B0 => write!(f, "BCM2837B0"),
        }
    }
}

// Identify Pi model based on /proc/cpuinfo
fn parse_proc_cpuinfo() -> Result<Model> {
    let proc_cpuinfo = BufReader::new(match File::open("/proc/cpuinfo") {
        Err(_) => return Err(Error::UnknownModel),
        Ok(file) => file,
    });

    let mut hardware: String = String::new();
    let mut revision: String = String::new();
    for line_result in proc_cpuinfo.lines() {
        if let Ok(line) = line_result {
            if line.starts_with("Hardware\t: ") {
                hardware = String::from(&line[11..]);
            } else if line.starts_with("Revision\t: ") {
                revision = String::from(&line[11..]).to_lowercase();
            }
        }
    }

    // Return an error if we don't recognize the SoC. This check is
    // done to prevent accidentally identifying a non-Pi SBC as a Pi
    // solely based on the revision field.
    match &hardware[..] {
        "BCM2708" | "BCM2835" | "BCM2709" | "BCM2836" | "BCM2710" | "BCM2837" | "BCM2837A1"
        | "BCM2837B0" => {}
        _ => return Err(Error::UnknownModel),
    }

    let model = if (revision.len() == 4) || (revision.len() == 8) {
        // Older revisions are 4 characters long, or 8 if they've been over-volted
        match &revision[revision.len() - 4..] {
            "0007" | "0008" | "0009" | "0015" => Model::RaspberryPiA,
            "Beta" | "0002" | "0003" => Model::RaspberryPiBRev1,
            "0004" | "0005" | "0006" | "000d" | "000e" | "000f" => Model::RaspberryPiBRev2,
            "0012" => Model::RaspberryPiAPlus,
            "0010" | "0013" => Model::RaspberryPiBPlus,
            "0011" | "0014" => Model::RaspberryPiComputeModule,
            _ => return Err(Error::UnknownModel),
        }
    } else if revision.len() >= 6 {
        // Newer revisions consist of at least 6 characters
        match &revision[revision.len() - 3..revision.len() - 1] {
            "00" => Model::RaspberryPiA,
            "01" => Model::RaspberryPiBRev2,
            "02" => Model::RaspberryPiAPlus,
            "03" => Model::RaspberryPiBPlus,
            "04" => Model::RaspberryPi2B,
            "06" => Model::RaspberryPiComputeModule,
            "08" => Model::RaspberryPi3B,
            "09" => Model::RaspberryPiZero,
            "0a" => Model::RaspberryPiComputeModule3,
            "0c" => Model::RaspberryPiZeroW,
            "0d" => Model::RaspberryPi3BPlus,
            _ => return Err(Error::UnknownModel),
        }
    } else {
        return Err(Error::UnknownModel);
    };

    Ok(model)
}

// Identify Pi model based on /sys/firmware/devicetree/base/compatible
fn parse_firmware_compatible() -> Result<Model> {
    let compatible = match fs::read_to_string("/sys/firmware/devicetree/base/compatible") {
        Err(_) => return Err(Error::UnknownModel),
        Ok(s) => s,
    };

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    for compid in compatible.split('\0') {
        match compid {
            "raspberrypi,model-b-i2c0" => return Ok(Model::RaspberryPiBRev1),
            "raspberrypi,model-b" => return Ok(Model::RaspberryPiBRev1),
            "raspberrypi,model-a" => return Ok(Model::RaspberryPiA),
            "raspberrypi,model-b-rev2" => return Ok(Model::RaspberryPiBRev2),
            "raspberrypi,model-a-plus" => return Ok(Model::RaspberryPiAPlus),
            "raspberrypi,model-b-plus" => return Ok(Model::RaspberryPiBPlus),
            "raspberrypi,2-model-b" => return Ok(Model::RaspberryPi2B),
            "raspberrypi,compute-module" => return Ok(Model::RaspberryPiComputeModule),
            "raspberrypi,3-model-b" => return Ok(Model::RaspberryPi3B),
            "raspberrypi,model-zero" => return Ok(Model::RaspberryPiZero),
            "raspberrypi,3-compute-module" => return Ok(Model::RaspberryPiComputeModule3),
            "raspberrypi,model-zero-w" => return Ok(Model::RaspberryPiZeroW),
            "raspberrypi,3-model-b-plus" => return Ok(Model::RaspberryPi3BPlus),
            _ => (),
        }
    }

    Err(Error::UnknownModel)
}

// Identify Pi model based on /sys/firmware/devicetree/base/model
fn parse_firmware_model() -> Result<Model> {
    let model = match fs::read_to_string("/sys/firmware/devicetree/base/model") {
        Err(_) => return Err(Error::UnknownModel),
        Ok(buffer) => if let Some(rev_idx) = buffer.find(" Rev ") {
            // We don't want to strip rev2 here
            buffer[0..rev_idx].to_owned()
        } else {
            buffer
        },
    };

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    match &model[..] {
        "Raspberry Pi Model B (no P5)" => Ok(Model::RaspberryPiBRev1),
        "Raspberry Pi Model B" => Ok(Model::RaspberryPiBRev1),
        "Raspberry Pi Model A" => Ok(Model::RaspberryPiA),
        "Raspberry Pi Model B rev2" => Ok(Model::RaspberryPiBRev2),
        "Raspberry Pi Model A+" => Ok(Model::RaspberryPiAPlus),
        "Raspberry Pi Model B+" => Ok(Model::RaspberryPiBPlus),
        "Raspberry Pi 2 Model B" => Ok(Model::RaspberryPi2B),
        "Raspberry Pi Compute Module" => Ok(Model::RaspberryPiComputeModule),
        "Raspberry Pi 3 Model B" => Ok(Model::RaspberryPi3B),
        "Raspberry Pi Zero" => Ok(Model::RaspberryPiZero),
        "Raspberry Pi Compute Module 3" => Ok(Model::RaspberryPiComputeModule3),
        "Raspberry Pi Zero W" => Ok(Model::RaspberryPiZeroW),
        "Raspberry Pi 3 Model B+" => Ok(Model::RaspberryPi3BPlus),
        _ => Err(Error::UnknownModel),
    }
}

/// Retrieves Raspberry Pi device information.
#[derive(Debug, PartialEq, Copy, Clone)]
pub struct DeviceInfo {
    model: Model,
    soc: SoC,
    peripheral_base: u32,
    gpio_offset: u32,
}

impl DeviceInfo {
    /// Constructs a new `DeviceInfo`.
    ///
    /// `new` parses the contents of `/proc/cpuinfo` to identify the Raspberry
    /// Pi's model and SoC.
    pub fn new() -> Result<DeviceInfo> {
        let model = parse_proc_cpuinfo()
            .or_else(|_| parse_firmware_compatible().or_else(|_| parse_firmware_model()))?;

        // Set SoC and memory offsets based on model
        match model {
            Model::RaspberryPiA
            | Model::RaspberryPiAPlus
            | Model::RaspberryPiBRev1
            | Model::RaspberryPiBRev2
            | Model::RaspberryPiBPlus
            | Model::RaspberryPiComputeModule
            | Model::RaspberryPiZero
            | Model::RaspberryPiZeroW => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2835,
                peripheral_base: PERIPHERAL_BASE_RPI,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi2B => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2836,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3B | Model::RaspberryPiComputeModule3 => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837A1,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3BPlus => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837B0,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
        }
    }

    /// Returns the Raspberry Pi's model.
    pub fn model(&self) -> Model {
        self.model
    }

    /// Returns the Raspberry Pi's SoC.
    pub fn soc(&self) -> SoC {
        self.soc
    }

    /// Returns the base memory address for the BCM283x peripherals.
    pub fn peripheral_base(&self) -> u32 {
        self.peripheral_base
    }

    /// Returns the offset from the base memory address for the GPIO section.
    pub fn gpio_offset(&self) -> u32 {
        self.gpio_offset
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use {Word, Result};
use private::layout::{StructReader, StructBuilder, StructSize, PointerBuilder, PointerReader};

pub trait FromStructReader<'a> {
    fn new(reader : StructReader<'a>) -> Self;
}

pub trait HasStructSize {
    fn struct_size() -> StructSize;
}

pub trait FromStructBuilder<'a> {
    fn new(structBuilder : StructBuilder<'a>) -> Self;
}

pub trait FromPointerReader<'a> {
    fn get_from_pointer(reader : &PointerReader<'a>) -> Result<Self>;
}

/// Associated types hackery that allows us to reason about Cap'n Proto types
/// without needing to give them a lifetime `'a`.
///
/// If `Foo` is a Cap'n Proto struct and `Bar` is a Rust-native struct, then
/// `foo::Reader<'a>` is to `foo::Owned` as `&'a Bar` is to `Bar`, and
/// `foo::Builder<'a>` is to `foo::Owned` as `&'a mut Bar` is to `Bar`.
/// The relationship is formalized by an `impl <'a> capnp::traits::Owned<'a> for foo::Owned`.
/// Because Cap'n Proto struct layout differs from Rust struct layout,
/// `foo::Owned` cannot be an inhabited type; it is useful nonetheless
/// as a type parameter, e.g. for a generic container that owns a Cap'n Proto message
/// of type `T: for<'a> capnp::traits::Owned<'a>`.
pub trait Owned<'a> {
    type Reader: FromPointerReader<'a>;
    type Builder: FromPointerBuilder<'a>;
    type Pipeline;
}

pub trait OwnedStruct<'a> {
   type Reader: FromStructReader<'a>;
   type Builder: FromStructBuilder<'a> + HasStructSize;
   type Pipeline;
}

pub trait FromPointerReaderRefDefault<'a> {
    fn get_from_pointer(reader : &PointerReader<'a>, default_value : *const Word) -> Self;
}

pub trait FromPointerBuilder<'a> {
    fn init_pointer(PointerBuilder<'a>, u32) -> Self;
    fn get_from_pointer(builder : PointerBuilder<'a>) -> Result<Self>;
}

pub trait FromPointerBuilderRefDefault<'a> {
    fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *const Word) -> Self;
}

pub trait SetPointerBuilder<To> {
    fn set_pointer_builder<'a>(PointerBuilder<'a>, Self) -> Result<()>;
}

pub trait HasTypeId {
    fn type_id() -> u64;
}

pub trait ToU16 {
    fn to_u16(self) -> u16;
}

pub trait FromU16 {
    fn from_u16(value : u16) -> ::std::result::Result<Self, ::NotInSchema>;
}

pub trait IndexMove<I, T> {
    fn index_move(&self, index : I) -> T;
}

pub struct ListIter<T, U> {
    marker : ::std::marker::PhantomData<U>,
    list : T,
    index : u32,
    size : u32,
}

impl <T, U> ListIter<T, U> {
    pub fn new(list : T, size : u32) -> ListIter<T, U> {
        ListIter { list : list, index : 0, size : size, marker : ::std::marker::PhantomData }
    }
}

impl <U, T : IndexMove<u32, U>> ::std::iter::Iterator for ListIter<T, U> {
    type Item = U;
    fn next(&mut self) -> ::std::option::Option<U> {
        if self.index < self.size {
            let result = self.list.index_move(self.index);
            self.index += 1;
            return Some(result);
        } else {
            return None;
        }
    }
}
<commit_msg>Added SetPointerBuilder to Owned<'a>::Reader to maximise usefulness of Owned.<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use {Word, Result};
use private::layout::{StructReader, StructBuilder, StructSize, PointerBuilder, PointerReader};

pub trait FromStructReader<'a> {
    fn new(reader : StructReader<'a>) -> Self;
}

pub trait HasStructSize {
    fn struct_size() -> StructSize;
}

pub trait FromStructBuilder<'a> {
    fn new(structBuilder : StructBuilder<'a>) -> Self;
}

pub trait FromPointerReader<'a> {
    fn get_from_pointer(reader : &PointerReader<'a>) -> Result<Self>;
}

/// Associated types hackery that allows us to reason about Cap'n Proto types
/// without needing to give them a lifetime `'a`.
///
/// If `Foo` is a Cap'n Proto struct and `Bar` is a Rust-native struct, then
/// `foo::Reader<'a>` is to `foo::Owned` as `&'a Bar` is to `Bar`, and
/// `foo::Builder<'a>` is to `foo::Owned` as `&'a mut Bar` is to `Bar`.
/// The relationship is formalized by an `impl <'a> capnp::traits::Owned<'a> for foo::Owned`.
/// Because Cap'n Proto struct layout differs from Rust struct layout,
/// `foo::Owned` cannot be an inhabited type; it is useful nonetheless
/// as a type parameter, e.g. for a generic container that owns a Cap'n Proto message
/// of type `T: for<'a> capnp::traits::Owned<'a>`.
pub trait Owned<'a> {
    type Reader: FromPointerReader<'a> + SetPointerBuilder<Self::Builder>;
    type Builder: FromPointerBuilder<'a>;
    type Pipeline;
}

pub trait OwnedStruct<'a> {
   type Reader: FromStructReader<'a>;
   type Builder: FromStructBuilder<'a> + HasStructSize;
   type Pipeline;
}

pub trait FromPointerReaderRefDefault<'a> {
    fn get_from_pointer(reader : &PointerReader<'a>, default_value : *const Word) -> Self;
}

pub trait FromPointerBuilder<'a> {
    fn init_pointer(PointerBuilder<'a>, u32) -> Self;
    fn get_from_pointer(builder : PointerBuilder<'a>) -> Result<Self>;
}

pub trait FromPointerBuilderRefDefault<'a> {
    fn get_from_pointer(builder : PointerBuilder<'a>, default_value : *const Word) -> Self;
}

pub trait SetPointerBuilder<To> {
    fn set_pointer_builder<'a>(PointerBuilder<'a>, Self) -> Result<()>;
}

pub trait HasTypeId {
    fn type_id() -> u64;
}

pub trait ToU16 {
    fn to_u16(self) -> u16;
}

pub trait FromU16 {
    fn from_u16(value : u16) -> ::std::result::Result<Self, ::NotInSchema>;
}

pub trait IndexMove<I, T> {
    fn index_move(&self, index : I) -> T;
}

pub struct ListIter<T, U> {
    marker : ::std::marker::PhantomData<U>,
    list : T,
    index : u32,
    size : u32,
}

impl <T, U> ListIter<T, U> {
    pub fn new(list : T, size : u32) -> ListIter<T, U> {
        ListIter { list : list, index : 0, size : size, marker : ::std::marker::PhantomData }
    }
}

impl <U, T : IndexMove<u32, U>> ::std::iter::Iterator for ListIter<T, U> {
    type Item = U;
    fn next(&mut self) -> ::std::option::Option<U> {
        if self.index < self.size {
            let result = self.list.index_move(self.index);
            self.index += 1;
            return Some(result);
        } else {
            return None;
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{File, Open, Write};
use std::io::BufferedReader;
use std::cmp::min;
use std::iter::AdditiveIterator;
use std::io::process::{ExitStatus, Command, Ignored, InheritFd};

use libc::funcs::posix88::unistd::{geteuid, getegid};
use libc::pid_t;

use super::config::Range;
use super::linux::get_user_name;

fn read_uid_map() -> Result<Vec<Range>,String> {
    let uid = unsafe { geteuid() };
    let user = try!(get_user_name(uid)
        .map_err(|e| format!("Error getting user name for {}: {}", uid, e)));
    let file = File::open(&Path::new("/etc/subuid"));
    let mut res = Vec::new();
    let mut reader = BufferedReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subuid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = from_str(parts[1]);
        let count = from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_none() || count.is_none() {
            return Err(format!("/etc/subuid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(&user.as_slice()) {
            let start: uint = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

fn read_gid_map() -> Result<Vec<Range>,String> {
    let uid = unsafe { geteuid() };
    let user = try!(get_user_name(uid)
        .map_err(|e| format!("Error getting user name for {}: {}", uid, e)));
    let file = File::open(&Path::new("/etc/subgid"));
    let mut res = Vec::new();
    let mut reader = BufferedReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subgid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = from_str(parts[1]);
        let count = from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_none() || count.is_none() {
            return Err(format!("/etc/subgid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(&user.as_slice()) {
            let start: uint = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

pub fn match_ranges(req: &Vec<Range>, allowed: &Vec<Range>, own_id: uint)
    -> Vec<(uint, uint, uint)>
{
    let mut res = vec!((0, own_id, 1));
    let mut reqiter = req.iter();
    let mut reqval = *reqiter.next().unwrap();
    let mut allowiter = allowed.iter();
    let mut allowval = *allowiter.next().unwrap();
    loop {
        if reqval.start == 0 {
            reqval = reqval.shift(1);
            continue;
        }
        let clen = min(reqval.len(), allowval.len());
        res.push((reqval.start, allowval.start, clen));
        reqval = reqval.shift(clen);
        allowval = allowval.shift(clen);
        if reqval.len() == 0 {
            reqval = match reqiter.next() {
                Some(val) => *val,
                None => break,
            };
        }
        if allowval.len() == 0 {
            allowval = match allowiter.next() {
                Some(val) => *val,
                None => unreachable!(),
            };
        }
    }
    return res;
}

pub fn write_uid_map(pid: pid_t, uid_req: &Vec<Range>, gid_req: &Vec<Range>)
    -> Result<(), String>
{
    let uid = unsafe { geteuid() };
    let gid = unsafe { getegid() };
    if uid_req.len() == 0 {
        //  No required mapping, just write our effective uid/gid to mapping
        let uid_map = format!("0 {} 1", uid);
        debug!("Writing uid_map: {}", uid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("uid_map"), Open, Write)
                .write_str(uid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing uid mapping: {}", e)),
        }
    } else {
        let allowed_uids = try!(read_uid_map());
        if uid_req.len() > 5 {
            warn!(concat!("Probably too many uid ranges required. ",
                          "Kernel might support no more than 5"));
        }
        let usum = uid_req.iter().map(|r| r.len()).sum();
        let ausum = allowed_uids.iter().map(|r| r.len()).sum();

        if usum > ausum {
            error!(concat!("Container requires more users than you ",
                           "are allowed to use (see /etc/subuid)"));
        }

        let uid_map = match_ranges(uid_req, &allowed_uids, uid as uint);

        let mut cmd = Command::new("newuidmap");
        cmd.stdin(Ignored).stdout(InheritFd(0)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        for &(req, allowed, count) in uid_map.iter() {
            cmd
                .arg(req.to_string())
                .arg(allowed.to_string())
                .arg(count.to_string());
        }
        info!("Uid map command: {}", cmd);
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newuidmap: {}", x)),
        }
    }

    if gid_req.len() == 0 {
        let gid_map = format!("0 {} 1", gid);
        debug!("Writing gid_map: {}", gid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("gid_map"), Open, Write)
                .write_str(gid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing gid mapping: {}", e)),
        }
    } else {
        let allowed_gids = try!(read_gid_map());

        if gid_req.len() > 5 {
            warn!(concat!("Probably too many gid ranges required. ",
                          "Kernel might support no more than 5"));
        }

        let gsum = gid_req.iter().map(|r| r.len()).sum();
        let agsum = allowed_gids.iter().map(|r| r.len()).sum();
        if gsum > agsum {
            error!(concat!("Container requires more groups than you are ",
                           "allowed to use (see /etc/subgid)"));
        }

        let gid_map = match_ranges(gid_req, &allowed_gids, gid as uint);

        let mut cmd = Command::new("newgidmap");
        cmd.stdin(Ignored).stdout(InheritFd(0)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        for &(req, allowed, count) in gid_map.iter() {
            cmd
                .arg(req.to_string())
                .arg(allowed.to_string())
                .arg(count.to_string());
        }
        info!("Gid map command: {}", cmd);
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newgidmap: {}", x)),
        }
    }
    return Ok(());
}

pub fn write_max_map(pid: pid_t) -> Result<(), String>
{
    let uid_map = read_uid_map().ok();
    let gid_map = read_gid_map().ok();

    let uid = unsafe { geteuid() };
    let gid = unsafe { getegid() };
    if uid_map.is_none() {
        //  No required mapping, just write our effective uid/gid to mapping
        let uid_map = format!("0 {} 1", uid);
        debug!("Writing uid_map: {}", uid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("uid_map"), Open, Write)
                .write_str(uid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing uid mapping: {}", e)),
        }
    } else {
        let uid_map = uid_map.unwrap();
        let mut cmd = Command::new("newuidmap");
        cmd.stdin(Ignored).stdout(InheritFd(0)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        cmd.arg("0");
        cmd.arg(uid.to_string());
        cmd.arg("1");
        for &rng in uid_map.iter() {
            let mut rng = rng;
            if uid as uint >= rng.start && uid as uint <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(uid as uint == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.len().to_string());
        }
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newuidmap: {}", x)),
        }
    }

    if gid_map.is_none() {
        let gid_map = format!("0 {} 1", gid);
        debug!("Writing gid_map: {}", gid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("gid_map"), Open, Write)
                .write_str(gid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing gid mapping: {}", e)),
        }
    } else {
        let gid_map = gid_map.unwrap();
        let mut cmd = Command::new("newgidmap");
        cmd.stdin(Ignored).stdout(InheritFd(0)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        cmd.arg("0");
        cmd.arg(gid.to_string());
        cmd.arg("1");
        for &rng in gid_map.iter() {
            let mut rng = rng;
            if gid as uint >= rng.start && gid as uint <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(gid as uint == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.len().to_string());
        }
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newgidmap: {}", x)),
        }
    }
    return Ok(());
}
<commit_msg>Fix fd inheritance for newuidmap and newgidmap<commit_after>use std::io::{File, Open, Write};
use std::io::BufferedReader;
use std::cmp::min;
use std::iter::AdditiveIterator;
use std::io::process::{ExitStatus, Command, Ignored, InheritFd};

use libc::funcs::posix88::unistd::{geteuid, getegid};
use libc::pid_t;

use super::config::Range;
use super::linux::get_user_name;

fn read_uid_map() -> Result<Vec<Range>,String> {
    let uid = unsafe { geteuid() };
    let user = try!(get_user_name(uid)
        .map_err(|e| format!("Error getting user name for {}: {}", uid, e)));
    let file = File::open(&Path::new("/etc/subuid"));
    let mut res = Vec::new();
    let mut reader = BufferedReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subuid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = from_str(parts[1]);
        let count = from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_none() || count.is_none() {
            return Err(format!("/etc/subuid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(&user.as_slice()) {
            let start: uint = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

fn read_gid_map() -> Result<Vec<Range>,String> {
    let uid = unsafe { geteuid() };
    let user = try!(get_user_name(uid)
        .map_err(|e| format!("Error getting user name for {}: {}", uid, e)));
    let file = File::open(&Path::new("/etc/subgid"));
    let mut res = Vec::new();
    let mut reader = BufferedReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subgid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = from_str(parts[1]);
        let count = from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_none() || count.is_none() {
            return Err(format!("/etc/subgid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(&user.as_slice()) {
            let start: uint = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

pub fn match_ranges(req: &Vec<Range>, allowed: &Vec<Range>, own_id: uint)
    -> Vec<(uint, uint, uint)>
{
    let mut res = vec!((0, own_id, 1));
    let mut reqiter = req.iter();
    let mut reqval = *reqiter.next().unwrap();
    let mut allowiter = allowed.iter();
    let mut allowval = *allowiter.next().unwrap();
    loop {
        if reqval.start == 0 {
            reqval = reqval.shift(1);
            continue;
        }
        let clen = min(reqval.len(), allowval.len());
        res.push((reqval.start, allowval.start, clen));
        reqval = reqval.shift(clen);
        allowval = allowval.shift(clen);
        if reqval.len() == 0 {
            reqval = match reqiter.next() {
                Some(val) => *val,
                None => break,
            };
        }
        if allowval.len() == 0 {
            allowval = match allowiter.next() {
                Some(val) => *val,
                None => unreachable!(),
            };
        }
    }
    return res;
}

pub fn write_uid_map(pid: pid_t, uid_req: &Vec<Range>, gid_req: &Vec<Range>)
    -> Result<(), String>
{
    let uid = unsafe { geteuid() };
    let gid = unsafe { getegid() };
    if uid_req.len() == 0 {
        //  No required mapping, just write our effective uid/gid to mapping
        let uid_map = format!("0 {} 1", uid);
        debug!("Writing uid_map: {}", uid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("uid_map"), Open, Write)
                .write_str(uid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing uid mapping: {}", e)),
        }
    } else {
        let allowed_uids = try!(read_uid_map());
        if uid_req.len() > 5 {
            warn!(concat!("Probably too many uid ranges required. ",
                          "Kernel might support no more than 5"));
        }
        let usum = uid_req.iter().map(|r| r.len()).sum();
        let ausum = allowed_uids.iter().map(|r| r.len()).sum();

        if usum > ausum {
            error!(concat!("Container requires more users than you ",
                           "are allowed to use (see /etc/subuid)"));
        }

        let uid_map = match_ranges(uid_req, &allowed_uids, uid as uint);

        let mut cmd = Command::new("newuidmap");
        cmd.stdin(Ignored).stdout(InheritFd(1)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        for &(req, allowed, count) in uid_map.iter() {
            cmd
                .arg(req.to_string())
                .arg(allowed.to_string())
                .arg(count.to_string());
        }
        info!("Uid map command: {}", cmd);
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newuidmap: {}", x)),
        }
    }

    if gid_req.len() == 0 {
        let gid_map = format!("0 {} 1", gid);
        debug!("Writing gid_map: {}", gid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("gid_map"), Open, Write)
                .write_str(gid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing gid mapping: {}", e)),
        }
    } else {
        let allowed_gids = try!(read_gid_map());

        if gid_req.len() > 5 {
            warn!(concat!("Probably too many gid ranges required. ",
                          "Kernel might support no more than 5"));
        }

        let gsum = gid_req.iter().map(|r| r.len()).sum();
        let agsum = allowed_gids.iter().map(|r| r.len()).sum();
        if gsum > agsum {
            error!(concat!("Container requires more groups than you are ",
                           "allowed to use (see /etc/subgid)"));
        }

        let gid_map = match_ranges(gid_req, &allowed_gids, gid as uint);

        let mut cmd = Command::new("newgidmap");
        cmd.stdin(Ignored).stdout(InheritFd(1)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        for &(req, allowed, count) in gid_map.iter() {
            cmd
                .arg(req.to_string())
                .arg(allowed.to_string())
                .arg(count.to_string());
        }
        info!("Gid map command: {}", cmd);
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newgidmap: {}", x)),
        }
    }
    return Ok(());
}

pub fn write_max_map(pid: pid_t) -> Result<(), String>
{
    let uid_map = read_uid_map().ok();
    let gid_map = read_gid_map().ok();

    let uid = unsafe { geteuid() };
    let gid = unsafe { getegid() };
    if uid_map.is_none() {
        //  No required mapping, just write our effective uid/gid to mapping
        let uid_map = format!("0 {} 1", uid);
        debug!("Writing uid_map: {}", uid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("uid_map"), Open, Write)
                .write_str(uid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing uid mapping: {}", e)),
        }
    } else {
        let uid_map = uid_map.unwrap();
        let mut cmd = Command::new("newuidmap");
        cmd.stdin(Ignored).stdout(InheritFd(1)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        cmd.arg("0");
        cmd.arg(uid.to_string());
        cmd.arg("1");
        for &rng in uid_map.iter() {
            let mut rng = rng;
            if uid as uint >= rng.start && uid as uint <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(uid as uint == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.len().to_string());
        }
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newuidmap: {}", x)),
        }
    }

    if gid_map.is_none() {
        let gid_map = format!("0 {} 1", gid);
        debug!("Writing gid_map: {}", gid_map);
        match File::open_mode(&Path::new("/proc")
                          .join(pid.to_string())
                          .join("gid_map"), Open, Write)
                .write_str(gid_map.as_slice()) {
            Ok(()) => {}
            Err(e) => return Err(format!(
                "Error writing gid mapping: {}", e)),
        }
    } else {
        let gid_map = gid_map.unwrap();
        let mut cmd = Command::new("newgidmap");
        cmd.stdin(Ignored).stdout(InheritFd(1)).stderr(InheritFd(2));
        cmd.arg(pid.to_string());
        cmd.arg("0");
        cmd.arg(gid.to_string());
        cmd.arg("1");
        for &rng in gid_map.iter() {
            let mut rng = rng;
            if gid as uint >= rng.start && gid as uint <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(gid as uint == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.start.to_string());
            cmd.arg(rng.len().to_string());
        }
        match cmd.status() {
            Ok(ExitStatus(0)) => {},
            x => return Err(format!("Error running newgidmap: {}", x)),
        }
    }
    return Ok(());
}
<|endoftext|>