<commit_before>use std::f64::consts::PI;
use std::cmp::Ordering;
use cgmath::prelude::*;
use cgmath::Point3;
use point::{Point, SinCosCache};
use events::{Events, EventKind, Circle};
use beach::{Beach, Arc};
use diagram::{Diagram, Kind, Vertex, Edge, Face};

pub trait Position : From<Point> {
    fn point(&self) -> &Point;

    fn position(&self) -> &Point3<f64> {
        &self.point().position
    }
}

fn in_range(phi: f64, phi_start: f64, phi_end: f64) -> Ordering {
    if phi_start <= phi_end {
        if phi < phi_start {
            Ordering::Less
        } else if phi > phi_end {
            Ordering::Greater
        } else {
            Ordering::Equal
        }
    } else {
        if phi < phi_end || phi > phi_start {
            Ordering::Equal
        } else if phi >= phi_end {
            Ordering::Greater
        } else {
            Ordering::Less
        }
    }
}

fn are_clockwise(n: Point3<f64>, p0: Point3<f64>, p1: Point3<f64>) -> bool {
    (p0 - n).cross(p1 - n).dot(n.to_vec()) < 0.0
}

fn wrap(phi: f64) -> f64 {
    if phi > PI {
        phi - 2.0 * PI
    } else if phi < -PI {
        phi + 2.0 * PI
    } else {
        phi
    }
}

struct Builder<K: Kind> where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    events: Events<K>,
    beach: Beach<K>,
    diagram: Diagram<K>,
    scan_theta: SinCosCache,
}

impl<K: Kind> Builder<K> where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    fn new(points: &[Point]) -> Result<Self, Error> {
        if points.len() < 2 {
            return Err(Error::FewPoints);
        }
        let mut builder = Builder {
            events: Events::default(),
            beach: Beach::default(),
            diagram: Diagram::default(),
            scan_theta: 0.0.into(),
        };
        for &point in points {
            let face = builder.diagram.add_face(point.into());
            builder.events.add_site(face, point);
        }
        Ok(builder)
    }

    fn build(mut self) -> Result<Diagram<K>, Error> {
        while let Some(event) = self.events.pop() {
            self.scan_theta = event.point.theta;
            match event.kind {
                EventKind::Site(face) => self.site_event(face, event.point),
                EventKind::Circle(event) => self.circle_event(event),
            }
        }
        self.cleanup_vertices();
        self.finish_faces();
        Ok(self.diagram)
    }

    fn arc_point(&self, arc: Arc<K>) -> &Point {
        &self.diagram.face_data(self.beach.face(arc)).point()
    }

    fn create_vertex(&mut self, point: Point, faces: &[Face<K>]) -> Vertex<K> {
        let vertex = self.diagram.add_vertex(point.into());
        for &face in faces {
            self.diagram.add_vertex_face(vertex, face);
        }
        vertex
    }

    fn create_edge(&mut self, vertex0: Vertex<K>, vertex1: Vertex<K>) -> Edge<K> {
        let edge = self.diagram.add_edge(Default::default());
        self.diagram.set_edge_vertices(edge, vertex0, vertex1);
        edge
    }

    fn site_event(&mut self, face: Face<K>, point: Point) {
        if let Some(mut arc) = self.beach.root() {
            if self.beach.len() == 1 {
                self.beach.insert_after(Some(arc), face);
                let arc0 = self.beach.first();
                let arc1 = self.beach.last();
                let point = self.phi_to_point(point.phi, self.arc_point(arc0));
                let faces = [face, self.beach.face(arc0)];
                let vertex = self.create_vertex(point, &faces);
                self.beach.set_start(arc0, Some(vertex));
                self.beach.set_start(arc1, Some(vertex));
                return;
            }
            let mut use_tree = true;
            loop {
                let prev_arc = self.beach.prev(arc);
                let next_arc = self.beach.next(arc);
                let phi_start = self.arcs_intersection(prev_arc, arc);
                let phi_end = self.arcs_intersection(arc, next_arc);
                match in_range(point.phi(), phi_start, phi_end) {
                    Ordering::Less => {
                        if use_tree {
                            if let Some(left) = self.beach.left(arc) {
                                arc = left;
                            } else {
                                // the tree has failed us, do the linear search from now on.
                                arc = self.beach.last();
                                use_tree = false;
                            }
                        } else {
                            arc = self.beach.prev(arc);
                        }
                    },
                    Ordering::Greater => {
                        if use_tree {
                            if let Some(right) = self.beach.right(arc) {
                                arc = right;
                            } else {
                                // the tree has failed us, do the linear search from now on.
                                arc = self.beach.first();
                                use_tree = false;
                            }
                        } else {
                            arc = self.beach.next(arc);
                        }
                    },
                    Ordering::Equal => {
                        self.try_remove_circle(arc);
                        let arc2 = {
                            let face = self.beach.face(arc);
                            let a = if prev_arc == self.beach.last() {
                                None
                            } else {
                                Some(prev_arc)
                            };
                            self.beach.insert_after(a, face)
                        };
                        let new_arc = self.beach.insert_after(Some(arc2), face);
                        let point = self.phi_to_point(point.phi, self.arc_point(arc));
                        let faces = [face, self.beach.face(arc)];
                        let vertex = self.create_vertex(point, &faces);
                        self.beach.set_start(arc2, Some(vertex));
                        self.beach.set_start(new_arc, Some(vertex));
                        self.try_add_circle(prev_arc, arc2, new_arc, point.theta());
                        self.try_add_circle(new_arc, arc, next_arc, point.theta());
                        if self.try_remove_circle(prev_arc) {
                            let prev_prev = self.beach.prev(prev_arc);
                            self.try_add_circle(prev_prev, prev_arc, arc2, -2.0 * PI);
                        }
                        break;
                    }
                }
            }
        } else {
            self.beach.insert_after(None, face);
        }
    }

    fn merge_arcs(&mut self, arc0: Arc<K>, arc1: Arc<K>, arc2: Arc<K>, vertex: Option<Vertex<K>>) {
        let (face0, face1, face2) = (self.beach.face(arc0), self.beach.face(arc1), self.beach.face(arc2));
        if face0 != face1 && face1 != face2 && face2 != face0 {
            let theta = self.scan_theta.value;
            if self.try_add_circle(arc0, arc1, arc2, theta) {
                if vertex.is_some() {
                    self.beach.set_start(arc1, vertex);
                }
            }
        }
    }

    fn edge_from_arc(&mut self, arc: Arc<K>, vertex: Vertex<K>) {
        if let Some(start) = self.beach.start(arc) {
            self.create_edge(start, vertex);
        }
    }

    fn circle_event(&mut self, event: Circle<K>) {
        if self.events.is_invalid(event) {
            return;
        }
        let (arc0, arc1, arc2) = self.events.arcs(event);
        assert_eq!(self.beach.circle(arc1), Some(event));
        self.beach.set_circle(arc1, None);
        self.try_remove_circle(arc0);
        self.try_remove_circle(arc2);
        let faces = [self.beach.face(arc0), self.beach.face(arc1), self.beach.face(arc2)];
        let point = self.events.center(event);
        let vertex = self.create_vertex(point, &faces);
        self.edge_from_arc(arc0, vertex);
        self.edge_from_arc(arc1, vertex);
        self.beach.remove(arc1);
        if self.beach.prev(arc0) == arc2 {
            self.edge_from_arc(arc2, vertex);
            self.beach.remove(arc0);
            self.beach.remove(arc2);
        } else {
            let prev = self.beach.prev(arc0);
            let next = self.beach.next(arc2);
            self.merge_arcs(prev, arc0, arc2, Some(vertex));
            self.merge_arcs(arc0, arc2, next, None);
        }
    }

    fn arcs_intersection(&mut self, arc0: Arc<K>, arc1: Arc<K>) -> f64 {
        let point0 = self.arc_point(arc0);
        let point1 = self.arc_point(arc1);
        let theta0 = point0.theta;
        let phi0 = point0.phi;
        let theta1 = point1.theta;
        let phi1 = point1.phi;
        if theta0.value >= self.scan_theta.value {
            point0.phi.value
        } else if theta1.value >= self.scan_theta.value {
            point1.phi.value
        } else {
            let u1 = (self.scan_theta.cos - theta1.cos) * theta0.sin;
            let u2 = (self.scan_theta.cos - theta0.cos) * theta1.sin;
            let a1 = u1 * phi0.cos;
            let a2 = u2 * phi1.cos;
            let a = a1 - a2;
            let b1 = u1 * phi0.sin;
            let b2 = u2 * phi1.sin;
            let b = b1 - b2;
            let c = (theta0.cos - theta1.cos) * self.scan_theta.sin;
            let length = (a * a + b * b).sqrt();
            if a.abs() > length || c.abs() > length {
                unreachable!()
            } else {
                let gamma = a.atan2(b);
                let phi_int_plus_gamma1 = (c / length).asin();
                wrap(phi_int_plus_gamma1 - gamma)
            }
        }
    }

    fn phi_to_point(&self, phi: SinCosCache, point: &Point) -> Point {
        let mut phi = phi;
        phi.value = wrap(phi.value);
        if point.theta() >= self.scan_theta.value {
            Point::from_cache(self.scan_theta, phi) // could be any point on the line segment
        } else {
            let a = self.scan_theta.sin - point.theta.sin * (phi.value - point.phi()).cos();
            let b = point.theta.cos - self.scan_theta.cos;
            let theta = SinCosCache::from(b.atan2(a));
            Point::from_cache(theta, phi)
        }
    }

    fn try_add_circle(&mut self, arc0: Arc<K>, arc1: Arc<K>, arc2: Arc<K>, theta: f64) -> bool {
        let p01 = self.arc_point(arc0).position - self.arc_point(arc1).position;
        let p21 = self.arc_point(arc2).position - self.arc_point(arc1).position;
        let cross = p01.cross(p21);
        let center = Point::from_cartesian(cross.x, cross.y, cross.z);
        let radius = center.distance(&self.arc_point(arc0));
        let point = Point::from_spherical(center.theta() + radius, center.phi());
        if point.theta() >= theta {
            let circle = self.events.add_circle((arc0, arc1, arc2), center, radius, point);
            self.beach.set_circle(arc1, Some(circle));
            true
        } else {
            false
        }
    }

    fn try_remove_circle(&mut self, arc: Arc<K>) -> bool {
        if let Some(event) = self.beach.circle(arc) {
            self.events.set_invalid(event, true);
            self.beach.set_circle(arc, None);
            true
        } else {
            false
        }
    }

    fn cleanup_vertices(&mut self) {
        let mut bad_vertices = Vec::new();
        for vertex in self.diagram.vertices() {
            let num_faces = self.diagram.vertex_faces(vertex).len();
            if num_faces < 3 {
                if num_faces == 2 {
                    let (vertex0, vertex1) = {
                        let neighbors: Vec<_> = self.diagram.vertex_neighbors(vertex).collect();
                        assert_eq!(neighbors.len(), 2);
                        (neighbors[0], neighbors[1])
                    };
                    self.create_edge(vertex0, vertex1);
                } else {
                    assert_eq!(num_faces, 0);
                }
                bad_vertices.push(vertex);
            }
        }
        self.diagram.remove_vertices(&bad_vertices);
    }

    fn finish_faces(&mut self) {
        for edge in self.diagram.edges() {
            let mut common = Vec::new();
            let (vertex0, vertex1) = self.diagram.edge_vertices(edge);
            self.diagram.add_vertex_edge(vertex0, edge);
            self.diagram.add_vertex_edge(vertex1, edge);
            for face0 in self.diagram.vertex_faces(vertex0) {
                for face1 in self.diagram.vertex_faces(vertex1) {
                    if face0 == face1 {
                        common.push(face0);
                    }
                }
            }
            assert_eq!(common.len(), 2);
            self.diagram.set_edge_faces(edge, common[0], common[1]);
        }
        for face in self.diagram.faces() {
            let n = *self.diagram.face_data(face).position();
            let mut edge = self.diagram.face_edges(face).next().unwrap();
            let (v0, v1) = self.diagram.edge_vertices(edge);
            let (prev, v) = if are_clockwise(n, *self.diagram[v0].position(), *self.diagram[v1].position()) {
                (v0, v1)
            } else {
                (v1, v0)
            };
            self.diagram.add_face_vertex(face, prev);
            let mut vertex = v;
            for _ in 0..self.diagram.face_edges(face).len() - 1 {
                self.diagram.add_face_vertex(face, vertex);
                for e in self.diagram.face_edges(face) {
                    if e != edge {
                        if let Some(v) = self.diagram.other_edge_vertex(e, vertex) {
                            vertex = v;
                            edge = e;
                            break;
                        }
                    }
                }
            }
        }
    }
}

#[derive(PartialEq)]
pub enum Error {
    FewPoints,
}

pub fn build<K: Kind>(points: &[Point]) -> Result<Diagram<K>, Error>
    where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    let builder = try!(Builder::new(points));
    builder.build()
}

pub fn build_relaxed<K: Kind>(points: &[Point], relaxations: usize) -> Result<Diagram<K>, Error>
    where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    let mut diagram = try!(build(points));
    for _ in 0..relaxations {
        let new_points: Vec<_> = diagram.faces().
            map(|face| {
                let face_points: Vec<_> = diagram.
                    face_vertices(face).
                    map(|vertex| {
                        let data: &K::Vertex = &diagram[vertex];
                        *data.position()
                    }).collect();
                let p = Point3::centroid(&face_points);
                Point::from_cartesian(p.x, p.y, p.z)
            }).
            collect();
        diagram = try!(build(&new_points));
    }
    Ok(diagram)
}
<commit_msg>Fix cleanup_vertices<commit_after>use std::f64::consts::PI;
use std::cmp::Ordering;
use cgmath::prelude::*;
use cgmath::Point3;
use point::{Point, SinCosCache};
use events::{Events, EventKind, Circle};
use beach::{Beach, Arc};
use diagram::{Diagram, Kind, Vertex, Edge, Face};

pub trait Position : From<Point> {
    fn point(&self) -> &Point;

    fn position(&self) -> &Point3<f64> {
        &self.point().position
    }
}

fn in_range(phi: f64, phi_start: f64, phi_end: f64) -> Ordering {
    if phi_start <= phi_end {
        if phi < phi_start {
            Ordering::Less
        } else if phi > phi_end {
            Ordering::Greater
        } else {
            Ordering::Equal
        }
    } else {
        if phi < phi_end || phi > phi_start {
            Ordering::Equal
        } else if phi >= phi_end {
            Ordering::Greater
        } else {
            Ordering::Less
        }
    }
}

fn are_clockwise(n: Point3<f64>, p0: Point3<f64>, p1: Point3<f64>) -> bool {
    (p0 - n).cross(p1 - n).dot(n.to_vec()) < 0.0
}

fn wrap(phi: f64) -> f64 {
    if phi > PI {
        phi - 2.0 * PI
    } else if phi < -PI {
        phi + 2.0 * PI
    } else {
        phi
    }
}

struct Builder<K: Kind> where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    events: Events<K>,
    beach: Beach<K>,
    diagram: Diagram<K>,
    scan_theta: SinCosCache,
}

impl<K: Kind> Builder<K> where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    fn new(points: &[Point]) -> Result<Self, Error> {
        if points.len() < 2 {
            return Err(Error::FewPoints);
        }
        let mut builder = Builder {
            events: Events::default(),
            beach: Beach::default(),
            diagram: Diagram::default(),
            scan_theta: 0.0.into(),
        };
        for &point in points {
            let face = builder.diagram.add_face(point.into());
            builder.events.add_site(face, point);
        }
        Ok(builder)
    }

    fn build(mut self) -> Result<Diagram<K>, Error> {
        while let Some(event) = self.events.pop() {
            self.scan_theta = event.point.theta;
            match event.kind {
                EventKind::Site(face) => self.site_event(face, event.point),
                EventKind::Circle(event) => self.circle_event(event),
            }
        }
        self.cleanup_vertices();
        self.finish_faces();
        Ok(self.diagram)
    }

    fn arc_point(&self, arc: Arc<K>) -> &Point {
        &self.diagram.face_data(self.beach.face(arc)).point()
    }

    fn create_vertex(&mut self, point: Point, faces: &[Face<K>]) -> Vertex<K> {
        let vertex = self.diagram.add_vertex(point.into());
        for &face in faces {
            self.diagram.add_vertex_face(vertex, face);
        }
        vertex
    }

    fn create_edge(&mut self, vertex0: Vertex<K>, vertex1: Vertex<K>) -> Edge<K> {
        let edge = self.diagram.add_edge(Default::default());
        self.diagram.set_edge_vertices(edge, vertex0, vertex1);
        edge
    }

    fn site_event(&mut self, face: Face<K>, point: Point) {
        if let Some(mut arc) = self.beach.root() {
            if self.beach.len() == 1 {
                self.beach.insert_after(Some(arc), face);
                let arc0 = self.beach.first();
                let arc1 = self.beach.last();
                let point = self.phi_to_point(point.phi, self.arc_point(arc0));
                let faces = [face, self.beach.face(arc0)];
                let vertex = self.create_vertex(point, &faces);
                self.beach.set_start(arc0, Some(vertex));
                self.beach.set_start(arc1, Some(vertex));
                return;
            }
            let mut use_tree = true;
            loop {
                let prev_arc = self.beach.prev(arc);
                let next_arc = self.beach.next(arc);
                let phi_start = self.arcs_intersection(prev_arc, arc);
                let phi_end = self.arcs_intersection(arc, next_arc);
                match in_range(point.phi(), phi_start, phi_end) {
                    Ordering::Less => {
                        if use_tree {
                            if let Some(left) = self.beach.left(arc) {
                                arc = left;
                            } else {
                                // the tree has failed us, do the linear search from now on.
                                arc = self.beach.last();
                                use_tree = false;
                            }
                        } else {
                            arc = self.beach.prev(arc);
                        }
                    },
                    Ordering::Greater => {
                        if use_tree {
                            if let Some(right) = self.beach.right(arc) {
                                arc = right;
                            } else {
                                // the tree has failed us, do the linear search from now on.
                                arc = self.beach.first();
                                use_tree = false;
                            }
                        } else {
                            arc = self.beach.next(arc);
                        }
                    },
                    Ordering::Equal => {
                        self.try_remove_circle(arc);
                        let arc2 = {
                            let face = self.beach.face(arc);
                            let a = if prev_arc == self.beach.last() {
                                None
                            } else {
                                Some(prev_arc)
                            };
                            self.beach.insert_after(a, face)
                        };
                        let new_arc = self.beach.insert_after(Some(arc2), face);
                        let point = self.phi_to_point(point.phi, self.arc_point(arc));
                        let faces = [face, self.beach.face(arc)];
                        let vertex = self.create_vertex(point, &faces);
                        self.beach.set_start(arc2, Some(vertex));
                        self.beach.set_start(new_arc, Some(vertex));
                        self.try_add_circle(prev_arc, arc2, new_arc, point.theta());
                        self.try_add_circle(new_arc, arc, next_arc, point.theta());
                        if self.try_remove_circle(prev_arc) {
                            let prev_prev = self.beach.prev(prev_arc);
                            self.try_add_circle(prev_prev, prev_arc, arc2, -2.0 * PI);
                        }
                        break;
                    }
                }
            }
        } else {
            self.beach.insert_after(None, face);
        }
    }

    fn merge_arcs(&mut self, arc0: Arc<K>, arc1: Arc<K>, arc2: Arc<K>, vertex: Option<Vertex<K>>) {
        let (face0, face1, face2) = (self.beach.face(arc0), self.beach.face(arc1), self.beach.face(arc2));
        if face0 != face1 && face1 != face2 && face2 != face0 {
            let theta = self.scan_theta.value;
            if self.try_add_circle(arc0, arc1, arc2, theta) {
                if vertex.is_some() {
                    self.beach.set_start(arc1, vertex);
                }
            }
        }
    }

    fn edge_from_arc(&mut self, arc: Arc<K>, vertex: Vertex<K>) {
        if let Some(start) = self.beach.start(arc) {
            self.create_edge(start, vertex);
        }
    }

    fn circle_event(&mut self, event: Circle<K>) {
        if self.events.is_invalid(event) {
            return;
        }
        let (arc0, arc1, arc2) = self.events.arcs(event);
        assert_eq!(self.beach.circle(arc1), Some(event));
        self.beach.set_circle(arc1, None);
        self.try_remove_circle(arc0);
        self.try_remove_circle(arc2);
        let faces = [self.beach.face(arc0), self.beach.face(arc1), self.beach.face(arc2)];
        let point = self.events.center(event);
        let vertex = self.create_vertex(point, &faces);
        self.edge_from_arc(arc0, vertex);
        self.edge_from_arc(arc1, vertex);
        self.beach.remove(arc1);
        if self.beach.prev(arc0) == arc2 {
            self.edge_from_arc(arc2, vertex);
            self.beach.remove(arc0);
            self.beach.remove(arc2);
        } else {
            let prev = self.beach.prev(arc0);
            let next = self.beach.next(arc2);
            self.merge_arcs(prev, arc0, arc2, Some(vertex));
            self.merge_arcs(arc0, arc2, next, None);
        }
    }

    fn arcs_intersection(&mut self, arc0: Arc<K>, arc1: Arc<K>) -> f64 {
        let point0 = self.arc_point(arc0);
        let point1 = self.arc_point(arc1);
        let theta0 = point0.theta;
        let phi0 = point0.phi;
        let theta1 = point1.theta;
        let phi1 = point1.phi;
        if theta0.value >= self.scan_theta.value {
            point0.phi.value
        } else if theta1.value >= self.scan_theta.value {
            point1.phi.value
        } else {
            let u1 = (self.scan_theta.cos - theta1.cos) * theta0.sin;
            let u2 = (self.scan_theta.cos - theta0.cos) * theta1.sin;
            let a1 = u1 * phi0.cos;
            let a2 = u2 * phi1.cos;
            let a = a1 - a2;
            let b1 = u1 * phi0.sin;
            let b2 = u2 * phi1.sin;
            let b = b1 - b2;
            let c = (theta0.cos - theta1.cos) * self.scan_theta.sin;
            let length = (a * a + b * b).sqrt();
            if a.abs() > length || c.abs() > length {
                unreachable!()
            } else {
                let gamma = a.atan2(b);
                let phi_int_plus_gamma1 = (c / length).asin();
                wrap(phi_int_plus_gamma1 - gamma)
            }
        }
    }

    fn phi_to_point(&self, phi: SinCosCache, point: &Point) -> Point {
        let mut phi = phi;
        phi.value = wrap(phi.value);
        if point.theta() >= self.scan_theta.value {
            Point::from_cache(self.scan_theta, phi) // could be any point on the line segment
        } else {
            let a = self.scan_theta.sin - point.theta.sin * (phi.value - point.phi()).cos();
            let b = point.theta.cos - self.scan_theta.cos;
            let theta = SinCosCache::from(b.atan2(a));
            Point::from_cache(theta, phi)
        }
    }

    fn try_add_circle(&mut self, arc0: Arc<K>, arc1: Arc<K>, arc2: Arc<K>, theta: f64) -> bool {
        let p01 = self.arc_point(arc0).position - self.arc_point(arc1).position;
        let p21 = self.arc_point(arc2).position - self.arc_point(arc1).position;
        let cross = p01.cross(p21);
        let center = Point::from_cartesian(cross.x, cross.y, cross.z);
        let radius = center.distance(&self.arc_point(arc0));
        let point = Point::from_spherical(center.theta() + radius, center.phi());
        if point.theta() >= theta {
            let circle = self.events.add_circle((arc0, arc1, arc2), center, radius, point);
            self.beach.set_circle(arc1, Some(circle));
            true
        } else {
            false
        }
    }

    fn try_remove_circle(&mut self, arc: Arc<K>) -> bool {
        if let Some(event) = self.beach.circle(arc) {
            self.events.set_invalid(event, true);
            self.beach.set_circle(arc, None);
            true
        } else {
            false
        }
    }

    fn cleanup_vertices(&mut self) {
        let mut bad_vertices = Vec::new();
        for vertex in self.diagram.vertices() {
            let num_faces = self.diagram.vertex_faces(vertex).len();
            if num_faces == 2 {
                let (vertex0, vertex1) = {
                    let neighbors: Vec<_> = self.diagram.vertex_neighbors(vertex).collect();
                    assert_eq!(neighbors.len(), 2);
                    (neighbors[0], neighbors[1])
                };
                self.create_edge(vertex0, vertex1);
                bad_vertices.push(vertex);
            }
        }
        self.diagram.remove_vertices(&bad_vertices);
    }

    fn finish_faces(&mut self) {
        for edge in self.diagram.edges() {
            let mut common = Vec::new();
            let (vertex0, vertex1) = self.diagram.edge_vertices(edge);
            self.diagram.add_vertex_edge(vertex0, edge);
            self.diagram.add_vertex_edge(vertex1, edge);
            for face0 in self.diagram.vertex_faces(vertex0) {
                for face1 in self.diagram.vertex_faces(vertex1) {
                    if face0 == face1 {
                        common.push(face0);
                    }
                }
            }
            assert_eq!(common.len(), 2);
            self.diagram.set_edge_faces(edge, common[0], common[1]);
        }
        for face in self.diagram.faces() {
            let n = *self.diagram.face_data(face).position();
            let mut edge = self.diagram.face_edges(face).next().unwrap();
            let (v0, v1) = self.diagram.edge_vertices(edge);
            let (prev, v) = if are_clockwise(n, *self.diagram[v0].position(), *self.diagram[v1].position()) {
                (v0, v1)
            } else {
                (v1, v0)
            };
            self.diagram.add_face_vertex(face, prev);
            let mut vertex = v;
            for _ in 0..self.diagram.face_edges(face).len() - 1 {
                self.diagram.add_face_vertex(face, vertex);
                for e in self.diagram.face_edges(face) {
                    if e != edge {
                        if let Some(v) = self.diagram.other_edge_vertex(e, vertex) {
                            vertex = v;
                            edge = e;
                            break;
                        }
                    }
                }
            }
        }
    }
}

#[derive(PartialEq)]
pub enum Error {
    FewPoints,
}

pub fn build<K: Kind>(points: &[Point]) -> Result<Diagram<K>, Error>
    where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    let builder = try!(Builder::new(points));
    builder.build()
}

pub fn build_relaxed<K: Kind>(points: &[Point], relaxations: usize) -> Result<Diagram<K>, Error>
    where K::Vertex: Position, K::Edge: Default, K::Face: Position {
    let mut diagram = try!(build(points));
    for _ in 0..relaxations {
        let new_points: Vec<_> = diagram.faces().
            map(|face| {
                let face_points: Vec<_> = diagram.
                    face_vertices(face).
                    map(|vertex| {
                        let data: &K::Vertex = &diagram[vertex];
                        *data.position()
                    }).collect();
                let p = Point3::centroid(&face_points);
                Point::from_cartesian(p.x, p.y, p.z)
            }).
            collect();
        diagram = try!(build(&new_points));
    }
    Ok(diagram)
}
<|endoftext|>


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

use std::os::unix::process::CommandExt;
use std::path::PathBuf;
use std::process::{Command, ExitStatus, Stdio};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

use hab_core::channel::{BLDR_CHANNEL_ENVVAR, STABLE_CHANNEL};
use hab_core::env;
use hab_core::fs;
use hab_core::url::BLDR_URL_ENVVAR;
use hab_core::AUTH_TOKEN_ENVVAR;

use error::{Error, Result};
use runner::log_pipe::LogPipe;
use runner::{NONINTERACTIVE_ENVVAR, RUNNER_DEBUG_ENVVAR};
use runner::workspace::Workspace;

pub static STUDIO_UID: AtomicUsize = ATOMIC_USIZE_INIT;
pub static STUDIO_GID: AtomicUsize = ATOMIC_USIZE_INIT;
pub const STUDIO_USER: &'static str = "krangschnak";
pub const STUDIO_GROUP: &'static str = "krangschnak";

lazy_static! {
    /// Absolute path to the Studio program
    static ref STUDIO_PROGRAM: PathBuf = fs::resolve_cmd_in_pkg(
        "hab-studio",
        include_str!(concat!(env!("OUT_DIR"), "/STUDIO_PKG_IDENT")),
    );

    pub static ref STUDIO_HOME: Mutex<PathBuf> = {
        Mutex::new(PathBuf::new())
    };
}

pub struct Studio<'a> {
    workspace: &'a Workspace,
    bldr_url: &'a str,
    auth_token: &'a str,
}

impl<'a> Studio<'a> {
    /// Creates a new Studio runner for a given `Workspace` and Builder URL.
    pub fn new(workspace: &'a Workspace, bldr_url: &'a str, auth_token: &'a str) -> Self {
        Studio {
            workspace,
            bldr_url,
            auth_token,
        }
    }

    /// Spawns a Studio build command, pipes output streams to the given `LogPipe` and returns the
    /// process' `ExitStatus`.
    ///
    /// # Errors
    ///
    /// * If the child process can't be spawned
    /// * If the calling thread can't wait on the child process
    /// * If the `LogPipe` fails to pipe output
    pub fn build(&self, log_pipe: &mut LogPipe) -> Result<ExitStatus> {
        let channel = if self.workspace.job.has_channel() {
            self.workspace.job.get_channel()
        } else {
            STABLE_CHANNEL
        };

        let mut cmd = Command::new("airlock");
        cmd.current_dir(self.workspace.src());
        cmd.uid(studio_uid());
        cmd.gid(studio_gid());
        cmd.env_clear();
        if let Some(val) = env::var_os(RUNNER_DEBUG_ENVVAR) {
            cmd.env("DEBUG", val);
        }
        cmd.env("HOME", &*STUDIO_HOME.lock().unwrap()); // Sets `$HOME` for build user
        cmd.env("USER", STUDIO_USER); // Sets `$USER` for build user
        cmd.env("PATH", env::var("PATH").unwrap_or(String::from(""))); // Sets `$PATH`
        cmd.env(NONINTERACTIVE_ENVVAR, "true"); // Disables progress bars
        cmd.env("TERM", "xterm-256color"); // Emits ANSI color codes
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
        cmd.arg("run");
        cmd.arg(&*STUDIO_PROGRAM);
        cmd.arg("-k"); // Origin key
        cmd.arg(self.workspace.job.origin());
        cmd.arg("build");
        cmd.arg(build_path(self.workspace.job.get_project().get_plan_path()));
        debug!("building studio build command, cmd={:?}", &cmd);
        cmd.env("AIRLOCK_FS_ROOT", self.workspace.studio());
        debug!(
            "setting studio build command env, AIRLOCK_FS_ROOT={}",
            self.workspace.studio().display()
        );
        debug!(
            "setting studio build command env, {}={}",
            BLDR_CHANNEL_ENVVAR,
            &channel
        );
        cmd.env(BLDR_CHANNEL_ENVVAR, channel);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_URL_ENVVAR,
            self.bldr_url
        );
        cmd.env(BLDR_URL_ENVVAR, self.bldr_url);
        cmd.env(AUTH_TOKEN_ENVVAR, self.auth_token);

        debug!("spawning studio build command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        log_pipe.pipe(&mut child)?;
        let exit_status = child.wait().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        debug!("completed studio build command, status={:?}", exit_status);
        Ok(exit_status)
    }
}

pub fn studio_gid() -> u32 {
    STUDIO_GID.load(Ordering::Relaxed) as u32
}

pub fn studio_uid() -> u32 {
    STUDIO_UID.load(Ordering::Relaxed) as u32
}

pub fn set_studio_gid(gid: u32) {
    STUDIO_GID.store(gid as usize, Ordering::Relaxed);
}

pub fn set_studio_uid(uid: u32) {
    STUDIO_UID.store(uid as usize, Ordering::Relaxed);
}

/// Returns a path argument suitable to pass to a Studio build command.
fn build_path(plan_path: &str) -> String {
    let mut parts: Vec<_> = plan_path.split("/").collect();
    if parts.last().map_or("", |p| *p) == "plan.sh" {
        parts.pop();
    }
    if parts.last().map_or("", |p| *p) == "habitat" {
        parts.pop();
    }

    if parts.is_empty() {
        String::from(".")
    } else {
        parts.join("/")
    }
}

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

    #[test]
    fn build_path_with_plan_sh() {
        assert_eq!(".", build_path("plan.sh"));
    }

    #[test]
    fn build_path_with_habitat_plan_sh() {
        assert_eq!(".", build_path("habitat/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_plan_sh() {
        assert_eq!("haaay", build_path("haaay/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_habitat_plan_sh() {
        assert_eq!(
            "components/yep",
            build_path("components/yep/habitat/plan.sh")
        );
    }
}
<commit_msg>[builder-worker] Update `airlock` invocation with new CLI options.<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::os::unix::process::CommandExt;
use std::path::PathBuf;
use std::process::{Command, ExitStatus, Stdio};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

use hab_core::channel::{BLDR_CHANNEL_ENVVAR, STABLE_CHANNEL};
use hab_core::env;
use hab_core::fs;
use hab_core::url::BLDR_URL_ENVVAR;
use hab_core::AUTH_TOKEN_ENVVAR;

use error::{Error, Result};
use runner::log_pipe::LogPipe;
use runner::{NONINTERACTIVE_ENVVAR, RUNNER_DEBUG_ENVVAR};
use runner::workspace::Workspace;

pub static STUDIO_UID: AtomicUsize = ATOMIC_USIZE_INIT;
pub static STUDIO_GID: AtomicUsize = ATOMIC_USIZE_INIT;
pub const STUDIO_USER: &'static str = "krangschnak";
pub const STUDIO_GROUP: &'static str = "krangschnak";

lazy_static! {
    /// Absolute path to the Studio program
    static ref STUDIO_PROGRAM: PathBuf = fs::resolve_cmd_in_pkg(
        "hab-studio",
        include_str!(concat!(env!("OUT_DIR"), "/STUDIO_PKG_IDENT")),
    );

    pub static ref STUDIO_HOME: Mutex<PathBuf> = {
        Mutex::new(PathBuf::new())
    };
}

pub struct Studio<'a> {
    workspace: &'a Workspace,
    bldr_url: &'a str,
    auth_token: &'a str,
}

impl<'a> Studio<'a> {
    /// Creates a new Studio runner for a given `Workspace` and Builder URL.
    pub fn new(workspace: &'a Workspace, bldr_url: &'a str, auth_token: &'a str) -> Self {
        Studio {
            workspace,
            bldr_url,
            auth_token,
        }
    }

    /// Spawns a Studio build command, pipes output streams to the given `LogPipe` and returns the
    /// process' `ExitStatus`.
    ///
    /// # Errors
    ///
    /// * If the child process can't be spawned
    /// * If the calling thread can't wait on the child process
    /// * If the `LogPipe` fails to pipe output
    pub fn build(&self, log_pipe: &mut LogPipe) -> Result<ExitStatus> {
        let channel = if self.workspace.job.has_channel() {
            self.workspace.job.get_channel()
        } else {
            STABLE_CHANNEL
        };

        let mut cmd = Command::new("airlock");
        cmd.current_dir(self.workspace.src());
        cmd.uid(studio_uid());
        cmd.gid(studio_gid());
        cmd.env_clear();
        if let Some(val) = env::var_os(RUNNER_DEBUG_ENVVAR) {
            cmd.env("DEBUG", val);
        }
        cmd.env("HOME", &*STUDIO_HOME.lock().unwrap()); // Sets `$HOME` for build user
        cmd.env("USER", STUDIO_USER); // Sets `$USER` for build user
        cmd.env("PATH", env::var("PATH").unwrap_or(String::from(""))); // Sets `$PATH`
        cmd.env(NONINTERACTIVE_ENVVAR, "true"); // Disables progress bars
        cmd.env("TERM", "xterm-256color"); // Emits ANSI color codes
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
        cmd.arg("run");
        cmd.arg("--fs-root");
        cmd.arg(self.workspace.studio());
        cmd.arg("--no-rm");
        cmd.arg(&*STUDIO_PROGRAM);
        cmd.arg("-k"); // Origin key
        cmd.arg(self.workspace.job.origin());
        cmd.arg("build");
        cmd.arg(build_path(self.workspace.job.get_project().get_plan_path()));
        debug!("building studio build command, cmd={:?}", &cmd);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_CHANNEL_ENVVAR,
            &channel
        );
        cmd.env(BLDR_CHANNEL_ENVVAR, channel);
        debug!(
            "setting studio build command env, {}={}",
            BLDR_URL_ENVVAR,
            self.bldr_url
        );
        cmd.env(BLDR_URL_ENVVAR, self.bldr_url);
        cmd.env(AUTH_TOKEN_ENVVAR, self.auth_token);

        debug!("spawning studio build command");
        let mut child = cmd.spawn().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        log_pipe.pipe(&mut child)?;
        let exit_status = child.wait().map_err(|e| {
            Error::StudioBuild(self.workspace.studio().to_path_buf(), e)
        })?;
        debug!("completed studio build command, status={:?}", exit_status);
        Ok(exit_status)
    }
}

pub fn studio_gid() -> u32 {
    STUDIO_GID.load(Ordering::Relaxed) as u32
}

pub fn studio_uid() -> u32 {
    STUDIO_UID.load(Ordering::Relaxed) as u32
}

pub fn set_studio_gid(gid: u32) {
    STUDIO_GID.store(gid as usize, Ordering::Relaxed);
}

pub fn set_studio_uid(uid: u32) {
    STUDIO_UID.store(uid as usize, Ordering::Relaxed);
}

/// Returns a path argument suitable to pass to a Studio build command.
fn build_path(plan_path: &str) -> String {
    let mut parts: Vec<_> = plan_path.split("/").collect();
    if parts.last().map_or("", |p| *p) == "plan.sh" {
        parts.pop();
    }
    if parts.last().map_or("", |p| *p) == "habitat" {
        parts.pop();
    }

    if parts.is_empty() {
        String::from(".")
    } else {
        parts.join("/")
    }
}

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

    #[test]
    fn build_path_with_plan_sh() {
        assert_eq!(".", build_path("plan.sh"));
    }

    #[test]
    fn build_path_with_habitat_plan_sh() {
        assert_eq!(".", build_path("habitat/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_plan_sh() {
        assert_eq!("haaay", build_path("haaay/plan.sh"));
    }

    #[test]
    fn build_path_with_subdir_habitat_plan_sh() {
        assert_eq!(
            "components/yep",
            build_path("components/yep/habitat/plan.sh")
        );
    }
}
<|endoftext|>


<commit_before>// Filebuffer -- Fast and simple file reading
// Copyright 2016 Ruud van Asseldonk
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// A copy of the License has been included in the root of the repository.

//! This mod contains the platform-specific implementations for Windows based on the winapi crate.

use std::fs;
use std::io;
use std::os::windows::io::AsRawHandle;
use std::ptr;

extern crate kernel32;
extern crate winapi;

pub struct PlatformData {
    // On Windows, the file must be kept open for the lifetime of the mapping.
    #[allow(dead_code)] // The field is not dead, the destructor is effectful.
    file: fs::File,
    mapping_handle: winapi::winnt::HANDLE,
}

impl Drop for PlatformData {
    fn drop (&mut self) {
        if self.mapping_handle != ptr::null_mut() {
            let success = unsafe { kernel32::CloseHandle(self.mapping_handle) };
            assert!(success != 0);
        }
    }
}

pub fn map_file(file: fs::File) -> io::Result<(*const u8, usize, PlatformData)> {
    let file_handle = file.as_raw_handle();
    let length = try!(file.metadata()).len();

    if length > usize::max_value() as u64 {
        return Err(io::Error::new(io::ErrorKind::Other, "file is larger than address space"));
    }

    let mut platform_data = PlatformData {
        file: file,
        mapping_handle: ptr::null_mut(),
    };

    // Don't try to map anything if the file is empty.
    if length == 0 {
        return Ok((ptr::null(), 0, platform_data));
    }

    platform_data.mapping_handle = unsafe { kernel32::CreateFileMappingW(
        file_handle,
        ptr::null_mut(),              // Use default security policy.
        winapi::winnt::PAGE_READONLY, // The memory will be read-only.
        0, 0,                         // The mapping size is the size of the file.
        ptr::null_mut())              // The mapping does not have a name.
    };

    if platform_data.mapping_handle == ptr::null_mut() {
        return Err(io::Error::last_os_error());
    }

    // TODO: Create the view.
    Ok((ptr::null(), 0, platform_data))
}

pub fn unmap_file(buffer: *const u8, length: usize) {
    // TODO: Implement this.
}

/// See also `unix::get_resident`.
pub fn get_resident(buffer: *const u8, length: usize, residency: &mut [bool]) {
    // TODO: Implement this.
}

/// See also `unix::prefetch`.
pub fn prefetch(buffer: *const u8, length: usize) {
    // TODO: Implement this.
}

pub fn get_page_size() -> usize {
    // Fill the `SYSTEM_INFO` struct with zeroes. It will be filled by
    // `GetSystemInfo` later but Rust requires it to be initialized.
    let mut sysinfo = winapi::sysinfoapi::SYSTEM_INFO {
        wProcessorArchitecture: 0,
        wReserved: 0,
        dwPageSize: 0,
        lpMinimumApplicationAddress: ptr::null_mut(),
        lpMaximumApplicationAddress: ptr::null_mut(),
        dwActiveProcessorMask: 0,
        dwNumberOfProcessors: 0,
        dwProcessorType: 0,
        dwAllocationGranularity: 0,
        wProcessorLevel: 0,
        wProcessorRevision: 0
    };
    unsafe { kernel32::GetSystemInfo(&mut sysinfo); }
    sysinfo.dwPageSize as usize
}
<commit_msg>Implement map_file() for Windows<commit_after>// Filebuffer -- Fast and simple file reading
// Copyright 2016 Ruud van Asseldonk
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// A copy of the License has been included in the root of the repository.

//! This mod contains the platform-specific implementations for Windows based on the winapi crate.

use std::fs;
use std::io;
use std::os::windows::io::AsRawHandle;
use std::ptr;

extern crate kernel32;
extern crate winapi;

pub struct PlatformData {
    // On Windows, the file must be kept open for the lifetime of the mapping.
    #[allow(dead_code)] // The field is not dead, the destructor is effectful.
    file: fs::File,
    mapping_handle: winapi::winnt::HANDLE,
}

impl Drop for PlatformData {
    fn drop (&mut self) {
        if self.mapping_handle != ptr::null_mut() {
            let success = unsafe { kernel32::CloseHandle(self.mapping_handle) };
            assert!(success != 0);
        }
    }
}

pub fn map_file(file: fs::File) -> io::Result<(*const u8, usize, PlatformData)> {
    let file_handle = file.as_raw_handle();
    let length = try!(file.metadata()).len();

    if length > usize::max_value() as u64 {
        return Err(io::Error::new(io::ErrorKind::Other, "file is larger than address space"));
    }

    let mut platform_data = PlatformData {
        file: file,
        mapping_handle: ptr::null_mut(),
    };

    // Don't try to map anything if the file is empty.
    if length == 0 {
        return Ok((ptr::null(), 0, platform_data));
    }

    // Memory-mapping a file on Windows is a two-step process: first we create a file mapping
    // object, and then we create a view of that mapping in the virtual address space.
    platform_data.mapping_handle = unsafe { kernel32::CreateFileMappingW(
        file_handle,
        ptr::null_mut(),              // Use default security policy.
        winapi::winnt::PAGE_READONLY, // The memory will be read-only.
        0, 0,                         // The mapping size is the size of the file.
        ptr::null_mut())              // The mapping does not have a name.
    };

    if platform_data.mapping_handle == ptr::null_mut() {
        return Err(io::Error::last_os_error());
    }

    let result = unsafe { kernel32::MapViewOfFile(
        platform_data.mapping_handle,
        winapi::memoryapi::FILE_MAP_READ, // The memory mapping will be read-only.
        0, 0,                             // Start offset of the mapping is 0.
        length)
    };

    if result == ptr::null_mut() {
        Err(io::Error::last_os_error())
    } else {
        Ok((result as *const u8, length as usize, platform_data))
    }
}

pub fn unmap_file(buffer: *const u8, length: usize) {
    // TODO: Implement this.
}

/// See also `unix::get_resident`.
pub fn get_resident(buffer: *const u8, length: usize, residency: &mut [bool]) {
    // TODO: Implement this.
}

/// See also `unix::prefetch`.
pub fn prefetch(buffer: *const u8, length: usize) {
    // TODO: Implement this.
}

pub fn get_page_size() -> usize {
    // Fill the `SYSTEM_INFO` struct with zeroes. It will be filled by
    // `GetSystemInfo` later but Rust requires it to be initialized.
    let mut sysinfo = winapi::sysinfoapi::SYSTEM_INFO {
        wProcessorArchitecture: 0,
        wReserved: 0,
        dwPageSize: 0,
        lpMinimumApplicationAddress: ptr::null_mut(),
        lpMaximumApplicationAddress: ptr::null_mut(),
        dwActiveProcessorMask: 0,
        dwNumberOfProcessors: 0,
        dwProcessorType: 0,
        dwAllocationGranularity: 0,
        wProcessorLevel: 0,
        wProcessorRevision: 0
    };
    unsafe { kernel32::GetSystemInfo(&mut sysinfo); }
    sysinfo.dwPageSize as usize
}
<|endoftext|>


<commit_before>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

extern crate exif;

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_dynamodb::model::AttributeValue;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The S3 bucket.
    #[structopt(short, long)]
    bucket: String,

    /// The filename.
    #[structopt(short, long)]
    filename: String,

    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// The DynamoDB table.
    #[structopt(short, long)]
    table: String,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

#[derive(Clone, Debug)]
struct Litem {
    name: String,
    confidence: f32,
}

#[derive(Debug)]
struct Edata {
    height: String,
    width: String,
    created: String,
}

#[derive(Debug)]
struct Ldata {
    label1: String,
    value1: String,
    label2: String,
    value2: String,
    label3: String,
    value3: String,
}

async fn add_file_to_bucket(client: &aws_sdk_s3::Client, bucket: &str, filename: &str) {
    let body = aws_sdk_s3::ByteStream::from_path(std::path::Path::new(filename)).await;

    match body {
        Ok(b) => {
            match client
                .put_object()
                .bucket(bucket)
                .key(filename)
                .body(b)
                .send()
                .await
            {
                Err(e) => {
                    println!("Got an error uploading file to bucket:");
                    println!("{}", e);
                    process::exit(1);
                }
                Ok(_) => {
                    println!("Uploaded file to bucket.");
                }
            }
        }
        Err(e) => {
            println!("Got an error uploading file to bucket:");
            println!("{}", e);
            process::exit(1);
        }
    }
}

fn get_exif_data(filename: &str) -> Edata {
    let height: String = "".to_owned();
    let width: String = "".to_owned();
    let created: String = "".to_owned();
    let mut edata = Edata {
        height,
        width,
        created,
    };

    let file = std::fs::File::open(&filename).unwrap();
    let mut bufreader = std::io::BufReader::new(&file);
    let exifreader = exif::Reader::new();

    match exifreader.read_from_container(&mut bufreader) {
        Ok(exif) => {
            println!("{}", &filename);

            for f in exif.fields() {
                // Get EXIF values for image width, height, and when image was created.
                match &*f.tag.to_string() {
                    "ImageWidth" => edata.height.push_str(&*f.display_value().to_string()),
                    "ImageLength" => edata.width.push_str(&*f.display_value().to_string()),
                    "DateTimeOriginal" => edata.created.push_str(&*f.display_value().to_string()),
                    _ => {}
                }
            }
        }
        Err(_) => {
            println!();
            println!("File does not contain ELIF data");
        }
    };

    edata
}

async fn add_data_to_table(
    verbose: bool,
    client: &aws_sdk_dynamodb::Client,
    table: &str,
    filename: &str,
    edata: Edata,
    labels: Vec<Litem>,
) {
    if verbose {
        println!("Added ")
    }
    let filename_string = AttributeValue::S(filename.to_string());
    let height_string = AttributeValue::S(edata.height);
    let created_string = AttributeValue::S(edata.created);
    let width_string = AttributeValue::S(edata.width);
    let label1_label = AttributeValue::S(labels[0].name.to_string());
    let label1_value = AttributeValue::S(labels[0].confidence.to_string());
    let label2_label = AttributeValue::S(labels[1].name.to_string());
    let label2_value = AttributeValue::S(labels[1].confidence.to_string());
    let label3_label = AttributeValue::S(labels[2].name.to_string());
    let label3_value = AttributeValue::S(labels[2].confidence.to_string());

    match client
        .put_item()
        .table_name(table)
        .item("filename", filename_string) // Table key.
        .item("height", height_string)
        .item("width", width_string)
        .item("created", created_string)
        .item("Label1", label1_label)
        .item("Value1", label1_value)
        .item("Label2", label2_label)
        .item("Value2", label2_value)
        .item("Label3", label3_label)
        .item("Value3", label3_value)
        .send()
        .await
    {
        Err(e) => {
            println!("Got an error adding data to table:");
            println!("{}", e);
            process::exit(1);
        }
        Ok(_) => {
            println!("Added info to table.");
        }
    }
}

async fn get_label_data(
    rekog_client: &aws_sdk_rekognition::Client,
    bucket: &str,
    key: &str,
) -> Vec<Litem> {
    let s3_obj = aws_sdk_rekognition::model::S3Object::builder()
        .bucket(bucket)
        .name(key)
        .build();

    let s3_img = aws_sdk_rekognition::model::Image::builder()
        .s3_object(s3_obj)
        .build();

    let resp = rekog_client.detect_labels().image(s3_img).send().await;

    let labels = resp.unwrap().labels.unwrap_or_default();

    // Create vector of Labels.
    let mut label_vec: Vec<Litem> = vec![];

    for label in labels {
        let name = label.name.as_deref().unwrap_or_default();
        let confidence = label.confidence.unwrap();

        let label = Litem {
            name: name.to_string(),
            confidence,
        };
        label_vec.push(label);
    }

    // Sort label items by confidence.
    label_vec.sort_by(|b, a| a.confidence.partial_cmp(&b.confidence).unwrap());

    // Return the first three items.
    label_vec[0..3].to_vec()
}

/// Gets EXIF information from a a JPG, JPEG, or PNG file,
/// uploads it to an S3 bucket,
/// uses Rekognition to identify the three top attributes (labels in Rekognition) in the file,
/// and adds the EXIF and label information to a DynamoDB table in the Region.
///
/// # Arguments
///
/// * `-b BUCKET` - The S3 bucket to which the file is uploaded.aws_sdk_dynamodb
/// * `-t TABLE` - The DynamoDB table in which the EXIF and label information is stored.
///   It must use the primary key `filename`.
/// * `[-r REGION]` - The Region in which the client is created.
///   If not supplied, uses the value of the **AWS_REGION** environment variable.
///   If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display information.
#[tokio::main]
async fn main() -> Result<(), exif::Error> {
    tracing_subscriber::fmt::init();

    let Opt {
        bucket,
        filename,
        region,
        table,
        verbose,
    } = Opt::from_args();

    let dynamo_region = region.clone();
    let s3_region = region.clone();
    let rek_region = region.clone();

    let dynamo_region_provider =
        RegionProviderChain::first_try(dynamo_region.map(aws_sdk_dynamodb::Region::new))
            .or_default_provider()
            .or_else(aws_sdk_dynamodb::Region::new("us-west-2"));

    let rek_region_provider =
        RegionProviderChain::first_try(s3_region.map(aws_sdk_rekognition::Region::new))
            .or_default_provider()
            .or_else(aws_sdk_rekognition::Region::new("us-west-2"));

    let s3_region_provider =
        RegionProviderChain::first_try(rek_region.map(aws_sdk_s3::Region::new))
            .or_default_provider()
            .or_else(aws_sdk_s3::Region::new("us-west-2"));

    println!();

    if verbose {
        println!(
            "DynamoDB client version:    {}",
            aws_sdk_dynamodb::PKG_VERSION
        );
        println!(
            "Rekognition client version: {}",
            aws_sdk_dynamodb::PKG_VERSION
        );
        println!("S3 client version:          {}", aws_sdk_s3::PKG_VERSION);
        println!("Filename:                   {}", &filename);
        println!("Bucket:                     {}", &bucket);
        println!("Table:                      {}", &table);
        println!();
        println!(
            "Region:                     {}",
            dynamo_region_provider.region().await.unwrap().as_ref()
        );
        println!();
    }

    let s3_shared_config = aws_config::from_env()
        .region(s3_region_provider)
        .load()
        .await;
    let s3_client = aws_sdk_s3::Client::new(&s3_shared_config);

    add_file_to_bucket(&s3_client, &bucket, &filename).await;

    let edata = get_exif_data(&filename);

    let rek_shared_config = aws_config::from_env()
        .region(rek_region_provider)
        .load()
        .await;
    let rek_client = aws_sdk_rekognition::Client::new(&rek_shared_config);

    let labels = get_label_data(&rek_client, &bucket, &filename).await;

    let dynamo_shared_config = aws_config::from_env()
        .region(dynamo_region_provider)
        .load()
        .await;
    let dynamo_client = aws_sdk_dynamodb::Client::new(&dynamo_shared_config);

    // Add data to table.
    add_data_to_table(verbose, &dynamo_client, &table, &filename, edata, labels).await;

    Ok(())
}
<commit_msg>Fixed RegionProviderChain args in detect faces Rust code example<commit_after>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

extern crate exif;

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_dynamodb::model::AttributeValue;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The S3 bucket.
    #[structopt(short, long)]
    bucket: String,

    /// The filename.
    #[structopt(short, long)]
    filename: String,

    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// The DynamoDB table.
    #[structopt(short, long)]
    table: String,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

#[derive(Clone, Debug)]
struct Litem {
    name: String,
    confidence: f32,
}

#[derive(Debug)]
struct Edata {
    height: String,
    width: String,
    created: String,
}

#[derive(Debug)]
struct Ldata {
    label1: String,
    value1: String,
    label2: String,
    value2: String,
    label3: String,
    value3: String,
}

async fn add_file_to_bucket(client: &aws_sdk_s3::Client, bucket: &str, filename: &str) {
    let body = aws_sdk_s3::ByteStream::from_path(std::path::Path::new(filename)).await;

    match body {
        Ok(b) => {
            match client
                .put_object()
                .bucket(bucket)
                .key(filename)
                .body(b)
                .send()
                .await
            {
                Err(e) => {
                    println!("Got an error uploading file to bucket:");
                    println!("{}", e);
                    process::exit(1);
                }
                Ok(_) => {
                    println!("Uploaded file to bucket.");
                }
            }
        }
        Err(e) => {
            println!("Got an error uploading file to bucket:");
            println!("{}", e);
            process::exit(1);
        }
    }
}

fn get_exif_data(filename: &str) -> Edata {
    let height: String = "".to_owned();
    let width: String = "".to_owned();
    let created: String = "".to_owned();
    let mut edata = Edata {
        height,
        width,
        created,
    };

    let file = std::fs::File::open(&filename).unwrap();
    let mut bufreader = std::io::BufReader::new(&file);
    let exifreader = exif::Reader::new();

    match exifreader.read_from_container(&mut bufreader) {
        Ok(exif) => {
            println!("{}", &filename);

            for f in exif.fields() {
                // Get EXIF values for image width, height, and when image was created.
                match &*f.tag.to_string() {
                    "ImageWidth" => edata.height.push_str(&*f.display_value().to_string()),
                    "ImageLength" => edata.width.push_str(&*f.display_value().to_string()),
                    "DateTimeOriginal" => edata.created.push_str(&*f.display_value().to_string()),
                    _ => {}
                }
            }
        }
        Err(_) => {
            println!();
            println!("File does not contain ELIF data");
        }
    };

    edata
}

async fn add_data_to_table(
    verbose: bool,
    client: &aws_sdk_dynamodb::Client,
    table: &str,
    filename: &str,
    edata: Edata,
    labels: Vec<Litem>,
) {
    if verbose {
        println!("Added ")
    }
    let filename_string = AttributeValue::S(filename.to_string());
    let height_string = AttributeValue::S(edata.height);
    let created_string = AttributeValue::S(edata.created);
    let width_string = AttributeValue::S(edata.width);
    let label1_label = AttributeValue::S(labels[0].name.to_string());
    let label1_value = AttributeValue::S(labels[0].confidence.to_string());
    let label2_label = AttributeValue::S(labels[1].name.to_string());
    let label2_value = AttributeValue::S(labels[1].confidence.to_string());
    let label3_label = AttributeValue::S(labels[2].name.to_string());
    let label3_value = AttributeValue::S(labels[2].confidence.to_string());

    match client
        .put_item()
        .table_name(table)
        .item("filename", filename_string) // Table key.
        .item("height", height_string)
        .item("width", width_string)
        .item("created", created_string)
        .item("Label1", label1_label)
        .item("Value1", label1_value)
        .item("Label2", label2_label)
        .item("Value2", label2_value)
        .item("Label3", label3_label)
        .item("Value3", label3_value)
        .send()
        .await
    {
        Err(e) => {
            println!("Got an error adding data to table:");
            println!("{}", e);
            process::exit(1);
        }
        Ok(_) => {
            println!("Added info to table.");
        }
    }
}

async fn get_label_data(
    rekog_client: &aws_sdk_rekognition::Client,
    bucket: &str,
    key: &str,
) -> Vec<Litem> {
    let s3_obj = aws_sdk_rekognition::model::S3Object::builder()
        .bucket(bucket)
        .name(key)
        .build();

    let s3_img = aws_sdk_rekognition::model::Image::builder()
        .s3_object(s3_obj)
        .build();

    let resp = rekog_client.detect_labels().image(s3_img).send().await;

    let labels = resp.unwrap().labels.unwrap_or_default();

    // Create vector of Labels.
    let mut label_vec: Vec<Litem> = vec![];

    for label in labels {
        let name = label.name.as_deref().unwrap_or_default();
        let confidence = label.confidence.unwrap();

        let label = Litem {
            name: name.to_string(),
            confidence,
        };
        label_vec.push(label);
    }

    // Sort label items by confidence.
    label_vec.sort_by(|b, a| a.confidence.partial_cmp(&b.confidence).unwrap());

    // Return the first three items.
    label_vec[0..3].to_vec()
}

/// Gets EXIF information from a a JPG, JPEG, or PNG file,
/// uploads it to an S3 bucket,
/// uses Rekognition to identify the three top attributes (labels in Rekognition) in the file,
/// and adds the EXIF and label information to a DynamoDB table in the Region.
///
/// # Arguments
///
/// * `-b BUCKET` - The S3 bucket to which the file is uploaded.aws_sdk_dynamodb
/// * `-t TABLE` - The DynamoDB table in which the EXIF and label information is stored.
///   It must use the primary key `filename`.
/// * `[-r REGION]` - The Region in which the client is created.
///   If not supplied, uses the value of the **AWS_REGION** environment variable.
///   If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display information.
#[tokio::main]
async fn main() -> Result<(), exif::Error> {
    tracing_subscriber::fmt::init();

    let Opt {
        bucket,
        filename,
        region,
        table,
        verbose,
    } = Opt::from_args();

    let dynamo_region = region.clone();
    let s3_region = region.clone();
    let rek_region = region.clone();

    let dynamo_region_provider =
        RegionProviderChain::first_try(dynamo_region.map(aws_sdk_dynamodb::Region::new))
            .or_default_provider()
            .or_else(aws_sdk_dynamodb::Region::new("us-west-2"));

    let rek_region_provider =
        RegionProviderChain::first_try(rek_region.map(aws_sdk_rekognition::Region::new))
            .or_default_provider()
            .or_else(aws_sdk_rekognition::Region::new("us-west-2"));

    let s3_region_provider =
        RegionProviderChain::first_try(s3_region.map(aws_sdk_s3::Region::new))
            .or_default_provider()
            .or_else(aws_sdk_s3::Region::new("us-west-2"));

    println!();

    if verbose {
        println!(
            "DynamoDB client version:    {}",
            aws_sdk_dynamodb::PKG_VERSION
        );
        println!(
            "Rekognition client version: {}",
            aws_sdk_dynamodb::PKG_VERSION
        );
        println!("S3 client version:          {}", aws_sdk_s3::PKG_VERSION);
        println!("Filename:                   {}", &filename);
        println!("Bucket:                     {}", &bucket);
        println!("Table:                      {}", &table);
        println!();
        println!(
            "Region:                     {}",
            dynamo_region_provider.region().await.unwrap().as_ref()
        );
        println!();
    }

    let s3_shared_config = aws_config::from_env()
        .region(s3_region_provider)
        .load()
        .await;
    let s3_client = aws_sdk_s3::Client::new(&s3_shared_config);

    add_file_to_bucket(&s3_client, &bucket, &filename).await;

    let edata = get_exif_data(&filename);

    let rek_shared_config = aws_config::from_env()
        .region(rek_region_provider)
        .load()
        .await;
    let rek_client = aws_sdk_rekognition::Client::new(&rek_shared_config);

    let labels = get_label_data(&rek_client, &bucket, &filename).await;

    let dynamo_shared_config = aws_config::from_env()
        .region(dynamo_region_provider)
        .load()
        .await;
    let dynamo_client = aws_sdk_dynamodb::Client::new(&dynamo_shared_config);

    // Add data to table.
    add_data_to_table(verbose, &dynamo_client, &table, &filename, edata, labels).await;

    Ok(())
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use std::sync::Mutex;

use either::Either;
use pinyin::{Pinyin, ToPinyin};
use tbot::{
    contexts::{Command, Text},
    types::{
        input_file,
        parameters::{self, WebPagePreviewState},
    },
    Bot,
};

use crate::client::pull_feed;
use crate::data::Database;
use crate::messages::{format_large_msg, Escape};

mod opml;

macro_rules! reject_cmd_from_channel {
    ($cmd: tt, $target: tt) => {{
        use tbot::contexts::fields::Message;
        if $cmd.chat().kind.is_channel() {
            let msg = "请在私聊中使用命令为频道管理订阅";
            update_response(&$cmd.bot, $target, parameters::Text::plain(&msg)).await?;
            return Ok(());
        }
    }};
}

#[derive(Debug, Copy, Clone)]
struct MsgTarget {
    chat_id: tbot::types::chat::Id,
    message_id: tbot::types::message::Id,
    first_time: bool,
}

impl MsgTarget {
    fn new(chat_id: tbot::types::chat::Id, message_id: tbot::types::message::Id) -> Self {
        MsgTarget {
            chat_id,
            message_id,
            first_time: true,
        }
    }
    fn update(&mut self, message_id: tbot::types::message::Id) {
        self.message_id = message_id;
        self.first_time = false;
    }
}

pub async fn start(
    _db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let target = &mut MsgTarget::new(cmd.chat.id, cmd.message_id);
    reject_cmd_from_channel!(cmd, target);
    let msg = "命令列表：\n\
               /rss       - 显示当前订阅的 RSS 列表\n\
               /sub       - 订阅一个 RSS: /sub http://example.com/feed.xml\n\
               /unsub     - 退订一个 RSS: /unsub http://example.com/feed.xml\n\
               /export    - 导出为 OPML";
    update_response(&cmd.bot, target, parameters::Text::plain(&msg)).await?;
    Ok(())
}

pub async fn rss(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let channel = &cmd.text.value;
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    reject_cmd_from_channel!(cmd, target);

    if !channel.is_empty() {
        let user_id = cmd.from.as_ref().unwrap().id;
        let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
        if channel_id.is_none() {
            return Ok(());
        }
        target_id = channel_id.unwrap();
    }

    let feeds = db.lock().unwrap().subscribed_feeds(target_id.0);
    let msgs = if let Some(mut feeds) = feeds {
        feeds.sort_by_cached_key(|feed| {
            feed.title
                .chars()
                .map(|c| {
                    c.to_pinyin()
                        .map(Pinyin::plain)
                        .map(Either::Right)
                        .unwrap_or_else(|| Either::Left(c))
                })
                .collect::<Vec<Either<char, &str>>>()
        });
        format_large_msg("订阅列表：".to_string(), &feeds, |feed| {
            format!(
                "{}",
                Escape(&feed.link),
                Escape(&feed.title)
            )
        })
    } else {
        vec!["订阅列表为空".to_string()]
    };

    let mut prev_msg = cmd.message_id;
    for msg in msgs {
        let text = parameters::Text::html(&msg);
        let msg = cmd
            .bot
            .send_message(chat_id, text)
            .reply_to_message_id(prev_msg)
            .web_page_preview(WebPagePreviewState::Disabled)
            .call()
            .await?;
        prev_msg = msg.id;
    }
    Ok(())
}

pub async fn sub(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let text = &cmd.text.value;
    let args = text.split_whitespace().collect::<Vec<_>>();
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    let feed_url;
    reject_cmd_from_channel!(cmd, target);

    match &*args {
        [url] => feed_url = url,
        [channel, url] => {
            let user_id = cmd.from.as_ref().unwrap().id;
            let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
            if channel_id.is_none() {
                return Ok(());
            }
            target_id = channel_id.unwrap();
            feed_url = url;
        }
        [..] => {
            let msg = "使用方法: /sub [Channel ID] <RSS URL>";
            update_response(&cmd.bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(());
        }
    };
    if db.lock().unwrap().is_subscribed(target_id.0, feed_url) {
        update_response(&cmd.bot, target, parameters::Text::plain("已订阅过的 RSS")).await?;
        return Ok(());
    }

    if cfg!(feature = "hosted-by-iovxw") && db.lock().unwrap().all_feeds().len() >= 1500 {
        let msg = "已达到全局最大订阅数量, \
                   为防止服务器压力过大请退订不需要的 RSS 或者\
                   [自己搭建服务](https://github.com/iovxw/rssbot)\n\
                   注: 本机器人主要用于提供即时提醒功能, 例如服务器状态监控和社区论坛提醒\n\
                   默认更新频率为 5 分钟, 不建议用于其他类型的 RSS 订阅\n\
                   如有相关需求推荐使用其他 RSS 机器人实现";
        update_response(&cmd.bot, target, parameters::Text::markdown(msg)).await?;
        return Ok(());
    }
    update_response(&cmd.bot, target, parameters::Text::plain("处理中，请稍候")).await?;
    let msg = match pull_feed(feed_url).await {
        Ok(feed) => {
            if db.lock().unwrap().subscribe(target_id.0, feed_url, &feed) {
                format!(
                    "《{}》 订阅成功",
                    Escape(&feed.link),
                    Escape(&feed.title)
                )
            } else {
                "已订阅过的 RSS".into()
            }
        }
        Err(e) => format!("订阅失败：{}", Escape(&e.to_user_friendly())),
    };
    update_response(&cmd.bot, target, parameters::Text::html(&msg)).await?;
    Ok(())
}

pub async fn unsub(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let text = &cmd.text.value;
    let args = text.split_whitespace().collect::<Vec<_>>();
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    let feed_url;
    reject_cmd_from_channel!(cmd, target);

    match &*args {
        [url] => feed_url = url,
        [channel, url] => {
            let user_id = cmd.from.as_ref().unwrap().id;
            let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
            if channel_id.is_none() {
                return Ok(());
            }
            target_id = channel_id.unwrap();
            feed_url = url;
        }
        [..] => {
            let msg = "使用方法: /unsub [Channel ID] <RSS URL>";
            update_response(&cmd.bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(());
        }
    };
    let msg = if let Some(feed) = db.lock().unwrap().unsubscribe(target_id.0, feed_url) {
        format!(
            "《{}》 退订成功",
            Escape(&feed.link),
            Escape(&feed.title)
        )
    } else {
        "未订阅过的 RSS".into()
    };
    update_response(&cmd.bot, target, parameters::Text::html(&msg)).await?;
    Ok(())
}

pub async fn export(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let channel = &cmd.text.value;
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    reject_cmd_from_channel!(cmd, target);

    if !channel.is_empty() {
        let user_id = cmd.from.as_ref().unwrap().id;
        let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
        if channel_id.is_none() {
            return Ok(());
        }
        target_id = channel_id.unwrap();
    }

    let feeds = db.lock().unwrap().subscribed_feeds(target_id.0);
    if feeds.is_none() {
        update_response(&cmd.bot, target, parameters::Text::plain("订阅列表为空")).await?;
        return Ok(());
    }
    let opml = opml::into_opml(feeds.unwrap());

    cmd.bot
        .send_document(
            chat_id,
            input_file::Document::bytes("feeds.opml", opml.as_bytes()),
        )
        .reply_to_message_id(cmd.message_id)
        .call()
        .await?;
    Ok(())
}

async fn update_response(
    bot: &Bot,
    target: &mut MsgTarget,
    message: parameters::Text<'_>,
) -> Result<(), tbot::errors::MethodCall> {
    let msg = if target.first_time {
        bot.send_message(target.chat_id, message)
            .reply_to_message_id(target.message_id)
            .web_page_preview(WebPagePreviewState::Disabled)
            .call()
            .await?
    } else {
        bot.edit_message_text(target.chat_id, target.message_id, message)
            .web_page_preview(WebPagePreviewState::Disabled)
            .call()
            .await?
    };
    target.update(msg.id);
    Ok(())
}

async fn check_channel_permission(
    bot: &Bot,
    channel: &str,
    target: &mut MsgTarget,
    user_id: tbot::types::user::Id,
) -> Result<Option<tbot::types::chat::Id>, tbot::errors::MethodCall> {
    use tbot::errors::MethodCall;
    let channel_id = channel
        .parse::<i64>()
        .map(|id| parameters::ChatId::Id(id.into()))
        .unwrap_or_else(|_| parameters::ChatId::Username(channel));
    update_response(bot, target, parameters::Text::plain("正在验证 Channel")).await?;

    let chat = match bot.get_chat(channel_id).call().await {
        Err(MethodCall::RequestError {
            description,
            error_code: 400,
            ..
        }) => {
            let msg = format!("无法找到目标 Channel：{}", description);
            update_response(bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(None);
        }
        other => other?,
    };
    if !chat.kind.is_channel() {
        update_response(bot, target, parameters::Text::plain("目标需为 Channel")).await?;
        return Ok(None);
    }
    let admins = match bot.get_chat_administrators(channel_id).call().await {
        Err(MethodCall::RequestError {
            description,
            error_code: 400,
            ..
        }) => {
            let msg = format!("无法获取频道信息（{}），请将本 Bot 设为管理员", description);
            update_response(bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(None);
        }
        other => other?,
    };
    let user_is_admin = admins
        .iter()
        .find(|member| member.user.id == user_id)
        .is_some();
    if !user_is_admin {
        update_response(
            bot,
            target,
            parameters::Text::plain("该命令只能由 Channel 管理员使用"),
        )
        .await?;
        return Ok(None);
    }
    let bot_is_admin = admins
        .iter()
        .find(|member| member.user.id == *crate::BOT_ID.get().unwrap())
        .is_some();
    if !bot_is_admin {
        update_response(
            bot,
            target,
            parameters::Text::plain("请将本 Bot 设为管理员"),
        )
        .await?;
        return Ok(None);
    }
    Ok(Some(chat.id))
}
<commit_msg>Add notes about Telegram Channels<commit_after>use std::sync::Arc;
use std::sync::Mutex;

use either::Either;
use pinyin::{Pinyin, ToPinyin};
use tbot::{
    contexts::{Command, Text},
    types::{
        input_file,
        parameters::{self, WebPagePreviewState},
    },
    Bot,
};

use crate::client::pull_feed;
use crate::data::Database;
use crate::messages::{format_large_msg, Escape};

mod opml;

macro_rules! reject_cmd_from_channel {
    ($cmd: tt, $target: tt) => {{
        use tbot::contexts::fields::Message;
        if $cmd.chat().kind.is_channel() {
            let msg = "请在私聊中使用命令为频道管理订阅";
            update_response(&$cmd.bot, $target, parameters::Text::plain(&msg)).await?;
            return Ok(());
        }
    }};
}

#[derive(Debug, Copy, Clone)]
struct MsgTarget {
    chat_id: tbot::types::chat::Id,
    message_id: tbot::types::message::Id,
    first_time: bool,
}

impl MsgTarget {
    fn new(chat_id: tbot::types::chat::Id, message_id: tbot::types::message::Id) -> Self {
        MsgTarget {
            chat_id,
            message_id,
            first_time: true,
        }
    }
    fn update(&mut self, message_id: tbot::types::message::Id) {
        self.message_id = message_id;
        self.first_time = false;
    }
}

pub async fn start(
    _db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let target = &mut MsgTarget::new(cmd.chat.id, cmd.message_id);
    reject_cmd_from_channel!(cmd, target);
    let msg = "命令列表：\n\
               /rss       - 显示当前订阅的 RSS 列表\n\
               /sub       - 订阅一个 RSS：`/sub http://example.com/feed.xml`\n\
               /unsub     - 退订一个 RSS：`/unsub http://example.com/feed.xml`\n\
               /export    - 导出为 OPML\n\
               所有命令均可在后面跟上频道 ID 来管理频道订阅\n\
               例如 `/sub @BotNews http://example.com/feed.xml`";
    update_response(&cmd.bot, target, parameters::Text::markdown(&msg)).await?;
    Ok(())
}

pub async fn rss(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let channel = &cmd.text.value;
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    reject_cmd_from_channel!(cmd, target);

    if !channel.is_empty() {
        let user_id = cmd.from.as_ref().unwrap().id;
        let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
        if channel_id.is_none() {
            return Ok(());
        }
        target_id = channel_id.unwrap();
    }

    let feeds = db.lock().unwrap().subscribed_feeds(target_id.0);
    let msgs = if let Some(mut feeds) = feeds {
        feeds.sort_by_cached_key(|feed| {
            feed.title
                .chars()
                .map(|c| {
                    c.to_pinyin()
                        .map(Pinyin::plain)
                        .map(Either::Right)
                        .unwrap_or_else(|| Either::Left(c))
                })
                .collect::<Vec<Either<char, &str>>>()
        });
        format_large_msg("订阅列表：".to_string(), &feeds, |feed| {
            format!(
                "{}",
                Escape(&feed.link),
                Escape(&feed.title)
            )
        })
    } else {
        vec!["订阅列表为空".to_string()]
    };

    let mut prev_msg = cmd.message_id;
    for msg in msgs {
        let text = parameters::Text::html(&msg);
        let msg = cmd
            .bot
            .send_message(chat_id, text)
            .reply_to_message_id(prev_msg)
            .web_page_preview(WebPagePreviewState::Disabled)
            .call()
            .await?;
        prev_msg = msg.id;
    }
    Ok(())
}

pub async fn sub(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let text = &cmd.text.value;
    let args = text.split_whitespace().collect::<Vec<_>>();
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    let feed_url;
    reject_cmd_from_channel!(cmd, target);

    match &*args {
        [url] => feed_url = url,
        [channel, url] => {
            let user_id = cmd.from.as_ref().unwrap().id;
            let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
            if channel_id.is_none() {
                return Ok(());
            }
            target_id = channel_id.unwrap();
            feed_url = url;
        }
        [..] => {
            let msg = "使用方法: /sub [Channel ID] <RSS URL>";
            update_response(&cmd.bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(());
        }
    };
    if db.lock().unwrap().is_subscribed(target_id.0, feed_url) {
        update_response(&cmd.bot, target, parameters::Text::plain("已订阅过的 RSS")).await?;
        return Ok(());
    }

    if cfg!(feature = "hosted-by-iovxw") && db.lock().unwrap().all_feeds().len() >= 1500 {
        let msg = "已达到全局最大订阅数量, \
                   为防止服务器压力过大请退订不需要的 RSS 或者\
                   [自己搭建服务](https://github.com/iovxw/rssbot)\n\
                   注: 本机器人主要用于提供即时提醒功能, 例如服务器状态监控和社区论坛提醒\n\
                   默认更新频率为 5 分钟, 不建议用于其他类型的 RSS 订阅\n\
                   如有相关需求推荐使用其他 RSS 机器人实现";
        update_response(&cmd.bot, target, parameters::Text::markdown(msg)).await?;
        return Ok(());
    }
    update_response(&cmd.bot, target, parameters::Text::plain("处理中，请稍候")).await?;
    let msg = match pull_feed(feed_url).await {
        Ok(feed) => {
            if db.lock().unwrap().subscribe(target_id.0, feed_url, &feed) {
                format!(
                    "《{}》 订阅成功",
                    Escape(&feed.link),
                    Escape(&feed.title)
                )
            } else {
                "已订阅过的 RSS".into()
            }
        }
        Err(e) => format!("订阅失败：{}", Escape(&e.to_user_friendly())),
    };
    update_response(&cmd.bot, target, parameters::Text::html(&msg)).await?;
    Ok(())
}

pub async fn unsub(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let text = &cmd.text.value;
    let args = text.split_whitespace().collect::<Vec<_>>();
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    let feed_url;
    reject_cmd_from_channel!(cmd, target);

    match &*args {
        [url] => feed_url = url,
        [channel, url] => {
            let user_id = cmd.from.as_ref().unwrap().id;
            let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
            if channel_id.is_none() {
                return Ok(());
            }
            target_id = channel_id.unwrap();
            feed_url = url;
        }
        [..] => {
            let msg = "使用方法: /unsub [Channel ID] <RSS URL>";
            update_response(&cmd.bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(());
        }
    };
    let msg = if let Some(feed) = db.lock().unwrap().unsubscribe(target_id.0, feed_url) {
        format!(
            "《{}》 退订成功",
            Escape(&feed.link),
            Escape(&feed.title)
        )
    } else {
        "未订阅过的 RSS".into()
    };
    update_response(&cmd.bot, target, parameters::Text::html(&msg)).await?;
    Ok(())
}

pub async fn export(
    db: Arc<Mutex<Database>>,
    cmd: Arc<Command<Text>>,
) -> Result<(), tbot::errors::MethodCall> {
    let chat_id = cmd.chat.id;
    let channel = &cmd.text.value;
    let mut target_id = chat_id;
    let target = &mut MsgTarget::new(chat_id, cmd.message_id);
    reject_cmd_from_channel!(cmd, target);

    if !channel.is_empty() {
        let user_id = cmd.from.as_ref().unwrap().id;
        let channel_id = check_channel_permission(&cmd.bot, channel, target, user_id).await?;
        if channel_id.is_none() {
            return Ok(());
        }
        target_id = channel_id.unwrap();
    }

    let feeds = db.lock().unwrap().subscribed_feeds(target_id.0);
    if feeds.is_none() {
        update_response(&cmd.bot, target, parameters::Text::plain("订阅列表为空")).await?;
        return Ok(());
    }
    let opml = opml::into_opml(feeds.unwrap());

    cmd.bot
        .send_document(
            chat_id,
            input_file::Document::bytes("feeds.opml", opml.as_bytes()),
        )
        .reply_to_message_id(cmd.message_id)
        .call()
        .await?;
    Ok(())
}

async fn update_response(
    bot: &Bot,
    target: &mut MsgTarget,
    message: parameters::Text<'_>,
) -> Result<(), tbot::errors::MethodCall> {
    let msg = if target.first_time {
        bot.send_message(target.chat_id, message)
            .reply_to_message_id(target.message_id)
            .web_page_preview(WebPagePreviewState::Disabled)
            .call()
            .await?
    } else {
        bot.edit_message_text(target.chat_id, target.message_id, message)
            .web_page_preview(WebPagePreviewState::Disabled)
            .call()
            .await?
    };
    target.update(msg.id);
    Ok(())
}

async fn check_channel_permission(
    bot: &Bot,
    channel: &str,
    target: &mut MsgTarget,
    user_id: tbot::types::user::Id,
) -> Result<Option<tbot::types::chat::Id>, tbot::errors::MethodCall> {
    use tbot::errors::MethodCall;
    let channel_id = channel
        .parse::<i64>()
        .map(|id| parameters::ChatId::Id(id.into()))
        .unwrap_or_else(|_| parameters::ChatId::Username(channel));
    update_response(bot, target, parameters::Text::plain("正在验证 Channel")).await?;

    let chat = match bot.get_chat(channel_id).call().await {
        Err(MethodCall::RequestError {
            description,
            error_code: 400,
            ..
        }) => {
            let msg = format!("无法找到目标 Channel：{}", description);
            update_response(bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(None);
        }
        other => other?,
    };
    if !chat.kind.is_channel() {
        update_response(bot, target, parameters::Text::plain("目标需为 Channel")).await?;
        return Ok(None);
    }
    let admins = match bot.get_chat_administrators(channel_id).call().await {
        Err(MethodCall::RequestError {
            description,
            error_code: 400,
            ..
        }) => {
            let msg = format!("无法获取频道信息（{}），请将本 Bot 设为管理员", description);
            update_response(bot, target, parameters::Text::plain(&msg)).await?;
            return Ok(None);
        }
        other => other?,
    };
    let user_is_admin = admins
        .iter()
        .find(|member| member.user.id == user_id)
        .is_some();
    if !user_is_admin {
        update_response(
            bot,
            target,
            parameters::Text::plain("该命令只能由 Channel 管理员使用"),
        )
        .await?;
        return Ok(None);
    }
    let bot_is_admin = admins
        .iter()
        .find(|member| member.user.id == *crate::BOT_ID.get().unwrap())
        .is_some();
    if !bot_is_admin {
        update_response(
            bot,
            target,
            parameters::Text::plain("请将本 Bot 设为管理员"),
        )
        .await?;
        return Ok(None);
    }
    Ok(Some(chat.id))
}
<|endoftext|>