<commit_before>use super::primitive::{Point2D, RectHV};

use std::iter;
use std::fmt;
use std::cmp::Ordering;
use super::super::symbol_tables::ST;

pub trait Point: Copy {
    type ValueType: Copy + PartialOrd = f64;
    // const DIMENSION: usize = 2;

    #[inline]
    fn get(&self, d: usize) -> Self::ValueType;

    #[inline]
    fn dimension() -> usize { 2 }
}

impl Point for Point2D {
    #[inline]
    fn get(&self, d: usize) -> f64 {
        if d == 0 {
            self.x
        } else if d == 1 {
            self.y
        } else {
            panic!("dimension not supported")
        }
    }
}

pub struct Node<K: Point, V> {
    pub key: K,
    pub val: V,
    pub left:  Option<Box<Node<K, V>>>,
    pub right: Option<Box<Node<K, V>>>,
    pub depth: usize
}

impl<K: Point, V> Node<K, V> {
    pub fn new(key: K, val: V, depth: usize) -> Node<K, V> {
        Node {
            key: key,
            val: val,
            left: None,
            right: None,
            // depth use (depth % k)-th dimension
            depth: depth
        }
    }

    fn size(&self) -> usize {
        let mut ret = 1;
        if self.left.is_some() {
            ret += self.left.as_ref().unwrap().size()
        }
        if self.right.is_some() {
            ret += self.right.as_ref().unwrap().size()
        }
        ret
    }
}

impl<K: Point + fmt::Debug, V: fmt::Debug> Node<K, V> {
    fn dump(&self, depth: usize, f: &mut fmt::Formatter, symbol: char) {
        if depth == 0 {
            writeln!(f, "\n{:?}[{:?}]", self.key, self.val).unwrap();
        } else {
            writeln!(f, "{}{}--{:?}[{:?}]", iter::repeat("|  ").take(depth-1).collect::<Vec<&str>>().concat(),
                     symbol, self.key, self.val).unwrap();
        }
        if self.left.is_some() {
            self.left.as_ref().unwrap().dump(depth + 1, f, '+');
        }
        if self.right.is_some() {
            self.right.as_ref().unwrap().dump(depth + 1, f, '`');
        }
    }
}

fn put<K: Point, V>(x: Option<Box<Node<K,V>>>, key: K, val: V, depth: usize) -> Option<Box<Node<K,V>>> {
    let mut x = x;
    if x.is_none() {
        return Some(Box::new(Node::new(key, val, depth)));
    }
    let depth = x.as_ref().unwrap().depth;
    let dim = x.as_ref().unwrap().depth % <K as Point>::dimension();

    let cmp = key.get(dim).partial_cmp(&x.as_ref().unwrap().key.get(dim)).unwrap();
    match cmp {
        Ordering::Less => {
            let left = x.as_mut().unwrap().left.take();
            x.as_mut().unwrap().left = put(left, key, val, depth + 1)
        },
        Ordering::Greater => {
            let right = x.as_mut().unwrap().right.take();
            x.as_mut().unwrap().right = put(right, key, val, depth + 1)
        },
        Ordering::Equal => {
            x.as_mut().unwrap().val = val
        }
    }
    x
}

fn delete_min<K: Point, V>(x: Option<Box<Node<K,V>>>) -> (Option<Box<Node<K,V>>>, Option<Box<Node<K,V>>>) {
    let mut x = x;
    if x.is_none() {
        return (None, None);
    }
    match x.as_mut().unwrap().left.take() {
        None           => (x.as_mut().unwrap().right.take(), x),
        left @ Some(_) => {
            let (t, deleted) = delete_min(left);
            x.as_mut().unwrap().left = t;
            (x, deleted)
        }
    }
}

fn delete<K: Point, V>(x: Option<Box<Node<K,V>>>, key: &K) -> Option<Box<Node<K,V>>> {
    if x.is_none() {
        return None;
    }

    let mut x = x;
    let dim = x.as_ref().unwrap().depth % <K as Point>::dimension();

    match key.get(dim).partial_cmp(&x.as_ref().unwrap().key.get(dim)).unwrap() {
        Ordering::Less => {
            let left = x.as_mut().unwrap().left.take();
            x.as_mut().unwrap().left = delete(left, key);
            return x;
        },
        Ordering::Greater => {
            let right = x.as_mut().unwrap().right.take();
            x.as_mut().unwrap().right = delete(right, key);
            return x;
        },
        Ordering::Equal => {
            if x.as_ref().unwrap().right.is_none() {
                return x.as_mut().unwrap().left.take();
            }
            if x.as_ref().unwrap().left.is_none() {
                return x.as_mut().unwrap().right.take();
            }

            // Save top
            let mut t = x;

            // split right into right without min, and the min
            let (right, right_min) = delete_min(t.as_mut().unwrap().right.take());
            x = right_min;
            x.as_mut().unwrap().right = right;
            x.as_mut().unwrap().left = t.as_mut().unwrap().left.take();
            x
        }
    }
}

pub struct KdTree<K: Point, V> {
    pub root: Option<Box<Node<K, V>>>
}

impl<K: Point, V> ST<K, V> for KdTree<K, V> {
    fn new() -> KdTree<K, V> {
        KdTree { root: None }
    }

    fn get(&self, key: &K) -> Option<&V> {
        let mut x = self.root.as_ref();
        while x.is_some() {
            let dim = x.as_ref().unwrap().depth % <K as Point>::dimension();
            match key.get(dim).partial_cmp(&x.unwrap().key.get(dim)).unwrap() {
                Ordering::Less => {
                    x = x.unwrap().left.as_ref();
                },
                Ordering::Greater => {
                    x = x.unwrap().right.as_ref();
                },
                Ordering::Equal  => {
                    return Some(&x.unwrap().val)
                }
            }
        }
        None
    }

    fn put(&mut self, key: K, val: V) {
        self.root = put(self.root.take(), key, val, 0);
    }

    fn delete(&mut self, key: &K) {
        self.root = delete(self.root.take(), key);
    }

    fn is_empty(&self) -> bool {
        self.root.is_none()
    }

    /// number of key-value pairs in the table
    fn size(&self) -> usize {
        if self.is_empty() {
            0
        } else {
            self.root.as_ref().unwrap().size()
        }
    }
}

impl<K: Point, V> KdTree<K, V> {
    pub fn keys<'a>(&'a self) -> ::std::vec::IntoIter<&'a K> {
        let mut queue: Vec<&'a K> = Vec::new();
        fn inorder<'a, K: Point, V>(x: Option<&'a Box<Node<K,V>>>, queue: &mut Vec<&'a K>) {
            if x.is_none() {
                return;
            }
            inorder(x.unwrap().left.as_ref(), queue);
            queue.push(&x.unwrap().key);
            inorder(x.unwrap().right.as_ref(), queue);
        };
        inorder(self.root.as_ref(), &mut queue);
        queue.into_iter()
    }
}

impl<V> KdTree<Point2D, V> {
    /// find all Point2D keys that lie in a 2d range
    pub fn range_search<T: AsRef<RectHV>>(&self, rect: T) -> Option<&Point2D> {
        unimplemented!()
    }

    /// number of keys that lie in a 2d range
    pub fn range_count<T: AsRef<RectHV>>(&self, rect: T) -> usize {
        unimplemented!()
    }
}

impl<K: Point + fmt::Debug, V: fmt::Debug> fmt::Debug for KdTree<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.root.is_none() {
            write!(f, "<empty tree>")
        } else {
            self.root.as_ref().unwrap().dump(0, f, ' ');
            Ok(())
        }
    }
}

#[test]
fn test_kd_tree_with_point_2d() {
    let mut t = KdTree::<Point2D, ()>::new();

    t.put(Point2D::new(0.7, 0.2), ());
    t.put(Point2D::new(0.5, 0.4), ());
    t.put(Point2D::new(0.2, 0.3), ());
    t.put(Point2D::new(0.4, 0.7), ());
    t.put(Point2D::new(0.9, 0.6), ());

    println!("got => {:?}", t);
}
<commit_msg>geometric_search/kd_tree: add range_search/range_count for KdTree<commit_after>use super::primitive::{Point2D, RectHV};

use std::iter;
use std::fmt;
use std::vec::IntoIter;
use std::cmp::Ordering;
use super::super::symbol_tables::ST;

pub trait Point: Copy {
    type ValueType: Copy + PartialOrd = f64;
    // const DIMENSION: usize = 2;

    #[inline]
    fn get(&self, d: usize) -> Self::ValueType;

    #[inline]
    fn dimension() -> usize { 2 }
}

impl Point for Point2D {
    #[inline]
    fn get(&self, d: usize) -> f64 {
        if d == 0 {
            self.x
        } else if d == 1 {
            self.y
        } else {
            panic!("dimension not supported")
        }
    }
}

pub struct Node<K: Point, V> {
    pub key: K,
    pub val: V,
    pub left:  Option<Box<Node<K, V>>>,
    pub right: Option<Box<Node<K, V>>>,
    pub depth: usize
}

impl<K: Point, V> Node<K, V> {
    pub fn new(key: K, val: V, depth: usize) -> Node<K, V> {
        Node {
            key: key,
            val: val,
            left: None,
            right: None,
            // depth use (depth % k)-th dimension
            depth: depth
        }
    }

    fn size(&self) -> usize {
        let mut ret = 1;
        if self.left.is_some() {
            ret += self.left.as_ref().unwrap().size()
        }
        if self.right.is_some() {
            ret += self.right.as_ref().unwrap().size()
        }
        ret
    }

    #[inline]
    fn comparator_for_current_dim(&self) -> K::ValueType {
        // let dim = self.depth % <K as Point>::dimension();
        self.key.get(self.depth % <K as Point>::dimension())
    }
}

impl<K: Point + fmt::Debug, V: fmt::Debug> Node<K, V> {
    fn dump(&self, depth: usize, f: &mut fmt::Formatter, symbol: char) {
        if depth == 0 {
            writeln!(f, "\n{:?}[{:?}]", self.key, self.val).unwrap();
        } else {
            writeln!(f, "{}{}--{:?}[{:?}]", iter::repeat("|  ").take(depth-1).collect::<Vec<&str>>().concat(),
                     symbol, self.key, self.val).unwrap();
        }
        if self.left.is_some() {
            self.left.as_ref().unwrap().dump(depth + 1, f, '+');
        }
        if self.right.is_some() {
            self.right.as_ref().unwrap().dump(depth + 1, f, '`');
        }
    }
}

fn put<K: Point, V>(x: Option<Box<Node<K,V>>>, key: K, val: V, depth: usize) -> Option<Box<Node<K,V>>> {
    let mut x = x;
    if x.is_none() {
        return Some(Box::new(Node::new(key, val, depth)));
    }
    let depth = x.as_ref().unwrap().depth;
    let dim = x.as_ref().unwrap().depth % <K as Point>::dimension();

    let cmp = key.get(dim).partial_cmp(&x.as_ref().unwrap().key.get(dim)).unwrap();
    match cmp {
        Ordering::Less => {
            let left = x.as_mut().unwrap().left.take();
            x.as_mut().unwrap().left = put(left, key, val, depth + 1)
        },
        Ordering::Greater => {
            let right = x.as_mut().unwrap().right.take();
            x.as_mut().unwrap().right = put(right, key, val, depth + 1)
        },
        Ordering::Equal => {
            x.as_mut().unwrap().val = val
        }
    }
    x
}

fn delete_min<K: Point, V>(x: Option<Box<Node<K,V>>>) -> (Option<Box<Node<K,V>>>, Option<Box<Node<K,V>>>) {
    let mut x = x;
    if x.is_none() {
        return (None, None);
    }
    match x.as_mut().unwrap().left.take() {
        None           => (x.as_mut().unwrap().right.take(), x),
        left @ Some(_) => {
            let (t, deleted) = delete_min(left);
            x.as_mut().unwrap().left = t;
            (x, deleted)
        }
    }
}

fn delete<K: Point, V>(x: Option<Box<Node<K,V>>>, key: &K) -> Option<Box<Node<K,V>>> {
    if x.is_none() {
        return None;
    }

    let mut x = x;
    let dim = x.as_ref().unwrap().depth % <K as Point>::dimension();

    match key.get(dim).partial_cmp(&x.as_ref().unwrap().key.get(dim)).unwrap() {
        Ordering::Less => {
            let left = x.as_mut().unwrap().left.take();
            x.as_mut().unwrap().left = delete(left, key);
            return x;
        },
        Ordering::Greater => {
            let right = x.as_mut().unwrap().right.take();
            x.as_mut().unwrap().right = delete(right, key);
            return x;
        },
        Ordering::Equal => {
            if x.as_ref().unwrap().right.is_none() {
                return x.as_mut().unwrap().left.take();
            }
            if x.as_ref().unwrap().left.is_none() {
                return x.as_mut().unwrap().right.take();
            }

            // Save top
            let mut t = x;

            // split right into right without min, and the min
            let (right, right_min) = delete_min(t.as_mut().unwrap().right.take());
            x = right_min;
            x.as_mut().unwrap().right = right;
            x.as_mut().unwrap().left = t.as_mut().unwrap().left.take();
            x
        }
    }
}

pub struct KdTree<K: Point, V> {
    pub root: Option<Box<Node<K, V>>>
}

impl<K: Point, V> ST<K, V> for KdTree<K, V> {
    fn new() -> KdTree<K, V> {
        KdTree { root: None }
    }

    fn get(&self, key: &K) -> Option<&V> {
        let mut x = self.root.as_ref();
        while x.is_some() {
            let dim = x.as_ref().unwrap().depth % <K as Point>::dimension();
            match key.get(dim).partial_cmp(&x.unwrap().key.get(dim)).unwrap() {
                Ordering::Less => {
                    x = x.unwrap().left.as_ref();
                },
                Ordering::Greater => {
                    x = x.unwrap().right.as_ref();
                },
                Ordering::Equal  => {
                    return Some(&x.unwrap().val)
                }
            }
        }
        None
    }

    fn put(&mut self, key: K, val: V) {
        self.root = put(self.root.take(), key, val, 0);
    }

    fn delete(&mut self, key: &K) {
        self.root = delete(self.root.take(), key);
    }

    fn is_empty(&self) -> bool {
        self.root.is_none()
    }

    /// number of key-value pairs in the table
    fn size(&self) -> usize {
        if self.is_empty() {
            0
        } else {
            self.root.as_ref().unwrap().size()
        }
    }
}

impl<K: Point, V> KdTree<K, V> {
    pub fn keys<'a>(&'a self) -> ::std::vec::IntoIter<&'a K> {
        let mut queue: Vec<&'a K> = Vec::new();
        fn inorder<'a, K: Point, V>(x: Option<&'a Box<Node<K,V>>>, queue: &mut Vec<&'a K>) {
            if x.is_none() {
                return;
            }
            inorder(x.unwrap().left.as_ref(), queue);
            queue.push(&x.unwrap().key);
            inorder(x.unwrap().right.as_ref(), queue);
        };
        inorder(self.root.as_ref(), &mut queue);
        queue.into_iter()
    }
}

impl KdTree<Point2D, ()> {
    // add the point to the KdTree
    pub fn insert(&mut self, p: Point2D) {
        self.put(p, ());
    }

    /// find all Point2D keys that lie in a 2d range
    pub fn range_search<T: AsRef<RectHV>>(&self, rect: T) -> IntoIter<&Point2D> {
        let mut result = Vec::new();

        // use stack approach
        let mut stack = Vec::new();
        stack.push(self.root.as_ref());
        while !stack.is_empty() {
            let x = stack.pop().unwrap();

            if x.is_none() {
                continue;
            }

            let dim = x.as_ref().unwrap().depth % 2;

            // Check if point in node lies in given rectangle
            if rect.as_ref().contains(x.as_ref().unwrap().key) {
                result.push(&x.as_ref().unwrap().key)
            }
            // Recursively search left/bottom (if any could fall in rectangle)
            // Recursively search right/top (if any could fall in rectangle)
            if dim == 0 {
                if rect.as_ref().xmin < x.as_ref().unwrap().comparator_for_current_dim() {
                    stack.push(x.unwrap().left.as_ref())
                }
                if rect.as_ref().xmax > x.as_ref().unwrap().comparator_for_current_dim() {
                    stack.push(x.unwrap().right.as_ref())
                }
            } else {        // dim == 1: y
                if rect.as_ref().ymin < x.as_ref().unwrap().comparator_for_current_dim() {
                    stack.push(x.unwrap().left.as_ref())
                }
                if rect.as_ref().ymax > x.as_ref().unwrap().comparator_for_current_dim() {
                    stack.push(x.unwrap().right.as_ref())
                }
            }
        }
        result.into_iter()
    }

    /// number of keys that lie in a 2d range
    pub fn range_count<T: AsRef<RectHV>>(&self, rect: T) -> usize {
        self.range_search(rect).count()
    }
}

impl<K: Point + fmt::Debug, V: fmt::Debug> fmt::Debug for KdTree<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.root.is_none() {
            write!(f, "<empty tree>")
        } else {
            self.root.as_ref().unwrap().dump(0, f, ' ');
            Ok(())
        }
    }
}

#[test]
fn test_kd_tree_with_point_2d() {
    let mut t = KdTree::<Point2D, ()>::new();

    t.put(Point2D::new(0.7, 0.2), ());
    t.put(Point2D::new(0.5, 0.4), ());
    t.put(Point2D::new(0.2, 0.3), ());
    t.put(Point2D::new(0.4, 0.7), ());
    t.put(Point2D::new(0.9, 0.6), ());

    println!("got => {:?}", t);

    assert_eq!(5, t.range_search(RectHV::new(0.1, 0.1, 0.9, 0.9)).count());
    assert_eq!(1, t.range_search(RectHV::new(0.1, 0.1, 0.4, 0.4)).count());
}
<|endoftext|>


<commit_before>use std::mem;

/// Slice from `vec[vec.len()..vec.capacity()]`
pub unsafe fn remaining_capacity_as_slice_mut<A>(vec: &mut Vec<A>) -> &mut [A] {
    let range = vec.len()..vec.capacity();
    vec.get_unchecked_mut(range)
}

pub unsafe fn remove_lifetime_mut<A: ?Sized>(a: &mut A) -> &'static mut A {
    mem::transmute(a)
}

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

    #[test]
    fn test_remaining_capacity_as_slice_mut() {
        let mut v = Vec::with_capacity(5);
        v.push(10);
        v.push(11);
        v.push(12);
        unsafe {
            {
                let s = remaining_capacity_as_slice_mut(&mut v);
                assert_eq!(2, s.len());
                s[0] = 13;
                s[1] = 14;
            }
            v.set_len(5);
        }
        assert_eq!(vec![10, 11, 12, 13, 14], v);
    }
}
<commit_msg>Map some of well known types to Rust equivalents<commit_after>use std::{mem, time::Duration};

use crate::well_known_types;

/// Slice from `vec[vec.len()..vec.capacity()]`
pub unsafe fn remaining_capacity_as_slice_mut<A>(vec: &mut Vec<A>) -> &mut [A] {
    let range = vec.len()..vec.capacity();
    vec.get_unchecked_mut(range)
}

pub unsafe fn remove_lifetime_mut<A: ?Sized>(a: &mut A) -> &'static mut A {
    mem::transmute(a)
}

// bool <-> BoolValue

impl From<well_known_types::BoolValue> for bool {
    fn from(inner: well_known_types::BoolValue) -> Self {
        inner.value
    }
}

impl From<bool> for well_known_types::BoolValue {
    fn from(inner: bool) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// Vec<u8> <-> BytesValue

impl From<well_known_types::BytesValue> for Vec<u8> {
    fn from(inner: well_known_types::BytesValue) -> Self {
        inner.value
    }
}

impl From<Vec<u8>> for well_known_types::BytesValue {
    fn from(inner: Vec<u8>) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// f64 <-> DoubleValue

impl From<well_known_types::DoubleValue> for f64 {
    fn from(inner: well_known_types::DoubleValue) -> Self {
        inner.value
    }
}

impl From<f64> for well_known_types::DoubleValue {
    fn from(inner: f64) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// f32 <-> FloatValue

impl From<well_known_types::FloatValue> for f32 {
    fn from(inner: well_known_types::FloatValue) -> Self {
        inner.value
    }
}

impl From<f32> for well_known_types::FloatValue {
    fn from(inner: f32) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// i32 <-> Int32Value

impl From<well_known_types::Int32Value> for i32 {
    fn from(inner: well_known_types::Int32Value) -> Self {
        inner.value
    }
}

impl From<i32> for well_known_types::Int32Value {
    fn from(inner: i32) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// i64 <-> Int64Value

impl From<well_known_types::Int64Value> for i64 {
    fn from(inner: well_known_types::Int64Value) -> Self {
        inner.value
    }
}

impl From<i64> for well_known_types::Int64Value {
    fn from(inner: i64) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// u32 <-> UInt32Value

impl From<well_known_types::UInt32Value> for u32 {
    fn from(inner: well_known_types::UInt32Value) -> Self {
        inner.value
    }
}

impl From<u32> for well_known_types::UInt32Value {
    fn from(inner: u32) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// u64 <-> UInt64Value

impl From<well_known_types::UInt64Value> for u64 {
    fn from(inner: well_known_types::UInt64Value) -> Self {
        inner.value
    }
}

impl From<u64> for well_known_types::UInt64Value {
    fn from(inner: u64) -> Self {
        let mut value = Self::new();
        value.value = inner;
        value
    }
}

// () <-> Empty

impl From<well_known_types::Empty> for () {
    fn from(_inner: well_known_types::Empty) -> Self {}
}

impl From<()> for well_known_types::Empty {
    fn from(_inner: ()) -> Self {
        Self::new()
    }
}

// TODO Think about `std::time::Duration` and `std::time::SystemTime` conversions.

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

    #[test]
    fn test_remaining_capacity_as_slice_mut() {
        let mut v = Vec::with_capacity(5);
        v.push(10);
        v.push(11);
        v.push(12);
        unsafe {
            {
                let s = remaining_capacity_as_slice_mut(&mut v);
                assert_eq!(2, s.len());
                s[0] = 13;
                s[1] = 14;
            }
            v.set_len(5);
        }
        assert_eq!(vec![10, 11, 12, 13, 14], v);
    }
}
<|endoftext|>


<commit_before>#![crate_id = "jit_macros#0.1.2"]
#![comment = "Macros for LibJIT Bindings"]
#![crate_type = "dylib"]
#![allow(raw_pointer_deriving, dead_code, non_camel_case_types)]
#![deny(unnecessary_parens, unrecognized_lint, unreachable_code, unnecessary_allocation, unnecessary_typecast, unnecessary_allocation, uppercase_variables, unused_must_use)]
#![feature(globs, plugin_registrar, quote)]
#![experimental]
//! This crate provides some macros for better LibJIT interop
//! For example, here's a quick example of automatic struct compilation to a LibJIT value:
//!
//! ```rust
//! extern crate jit;
//! #[phase(syntax)]
//! extern crate jit_macros;
//! use jit::{C_DECL, Context, Function, Type, Types, get_type};
//! struct Pos {
//!     x: f64,
//!     y: f64,
//!     z: f64
//! }
//! fn main() {
//!     let cx = Context::new();
//!     cx.build(|| {
//!         // build the IR
//!         let sig = get_type(fn() -> Pos);
//!         let func = Function::new(cx, sig);
//!         let result = Pos {
//!                x: -10,
//!                y: 0,
//!                z: 10
//!            }.compile();
//!         func.insn_return(&result);
//!         /// run the IR
//!         func.compile();
//!         let rfunc:fn() -> Pos = func.closure0();
//!         assert_eq(rfunc().x, -10)
//!     });
//! }
//! ```
extern crate rustc;
extern crate syntax;

use rustc::plugin::Registry;
use std::gc::{Gc, GC};
use syntax::ast::*;
use syntax::codemap::*;
use syntax::ext::base::*;
use syntax::ext::build::*;
use syntax::ext::quote::rt::*;
use syntax::parse::token::*;
use syntax::owned_slice::OwnedSlice;

fn gen_compile(cx:&mut ExtCtxt, pos:Span, _:Gc<MetaItem>, item:Gc<Item>, cb:|Gc<Item>|) {
    let fields:Vec<(Option<String>, P<Expr>)> = match item.node {
        ItemStruct(def, _) if def.is_virtual => {
            cx.span_err(pos, "Virtual structs cannot be JIT compiled");
            fail!("...")
        },
        ItemStruct(def, _) => {
            def.fields.iter().map(|field| {
                let name:Option<String> = match field.node.kind {
                    NamedField(id, _) => Some(id.to_source()),
                    UnnamedField(_) => None
                };
                let ty = field.node.ty;
                let ty_expr = quote_expr!(&*cx, &::jit::get_type::<$ty>());
                (name, ty_expr)
            }).collect()
        },
        _ => {
            cx.span_err(pos, "Compile can only be automatically implemented for structs");
            fail!("...")
        }
    };
    let jit_type_t = cx.ty_path(cx.path_global(pos, vec!(cx.ident_of("jit"), cx.ident_of("Type"))), None);
    let create_struct = cx.expr_path(cx.path_global(pos, vec!(cx.ident_of("jit"), cx.ident_of("Type"), cx.ident_of("create_struct"))));
    let jit_type_body = cx.block(pos, vec!(
        cx.stmt_let(pos, false, cx.ident_of("ty"), {
            let mut fields = cx.expr_vec(pos, fields.iter().map(|&(_, ex)| ex).collect());
            fields = cx.expr_method_call(pos, fields, cx.ident_of("as_mut_slice"), vec!());
            cx.expr_call(pos, create_struct, vec!(fields))
        }),
    ), Some(cx.expr_ident(pos, cx.ident_of("ty"))));
    let curr_item = cx.ty_ident(pos, item.ident);
    let jit_type = Method {
        ident: cx.ident_of("jit_type"),
        attrs: vec!(
            cx.attribute(pos, cx.meta_word(pos, InternedString::new("inline")))
        ),
        generics: Generics {
            lifetimes: vec!(),
            ty_params: OwnedSlice::empty()
        },
        explicit_self: Spanned {
            node: SelfStatic,
            span: pos
        },
        fn_style: NormalFn,
        decl: cx.fn_decl(
            vec!(
                cx.arg(pos, cx.ident_of("_"), cx.ty_option(curr_item))
            ),
            jit_type_t
        ),
        body: jit_type_body,
        span: pos,
        id: item.id,
        vis: Inherited
    };
    let lifetime = cx.lifetime(pos, cx.ident_of("'a").name);
    let compile_body = cx.block(pos, vec!(
        cx.stmt_let(pos, false, cx.ident_of("ty"), cx.expr_call(pos, quote_expr!(&*cx, ::jit::get_type::<$curr_item>), vec!())),
        cx.stmt_let(pos, false, cx.ident_of("val"), cx.expr_call(pos, quote_expr!(&*cx, ::jit::Value::new), vec!(
            cx.expr_ident(pos, cx.ident_of("func")),
            cx.expr_addr_of(pos, cx.expr_ident(pos, cx.ident_of("ty")))
        )))
    ), Some(cx.expr_ident(pos, cx.ident_of("val"))));
    let jit_func_t = cx.ty_path(cx.path_all(pos, true, vec!(cx.ident_of("jit"), cx.ident_of("Function")), vec!(lifetime.clone()), vec!()), None);
    let jit_val_t = cx.ty_path(cx.path_all(pos, true, vec!(cx.ident_of("jit"), cx.ident_of("Value")), vec!(lifetime.clone()), vec!()), None);
    let compile = Method {
        ident: cx.ident_of("compile"),
        attrs: vec!(
            cx.attribute(pos, cx.meta_word(pos, InternedString::new("inline")))
        ),
        generics: Generics {
            lifetimes: vec!(
                lifetime.clone()
            ),
            ty_params: OwnedSlice::empty()
        },
        explicit_self: Spanned {
            node: SelfRegion(None, MutImmutable),
            span: pos
        },
        fn_style: NormalFn,
        decl: cx.fn_decl(
            vec!(
                Arg::new_self(pos, MutImmutable),
                cx.arg(pos, cx.ident_of("func"), cx.ty_rptr(pos, jit_func_t, None, MutImmutable))
            ),
            jit_val_t
        ),
        body: compile_body,
        span: pos,
        id: item.id,
        vis: Inherited
    };
    let methods = vec!(
        box(GC) jit_type,
        box(GC) compile
    );
    let node = ItemImpl(
        Generics {
            lifetimes: vec!(),
            ty_params: OwnedSlice::empty()
        },
        Some(cx.trait_ref(cx.path_global(pos, vec!(cx.ident_of("jit"), cx.ident_of("Compile"))))),
        cx.ty_path(cx.path_ident(pos, item.ident), None),
        methods
    );
    cb(cx.item(pos, cx.ident_of("Compile"), vec!(), node));
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(intern("jit_compile"), ItemDecorator(gen_compile));
}<commit_msg>Update jit_macro<commit_after>#![crate_id = "jit_macros#0.1.2"]
#![comment = "Macros for LibJIT Bindings"]
#![crate_type = "dylib"]
#![allow(raw_pointer_deriving, dead_code, non_camel_case_types)]
#![deny(unnecessary_parens, unrecognized_lint, unreachable_code, unnecessary_allocation, unnecessary_typecast, unnecessary_allocation, uppercase_variables, unused_must_use)]
#![feature(globs, plugin_registrar, quote)]
#![experimental]
//! This crate provides some macros for better LibJIT interop
//! For example, here's a quick example of automatic struct compilation to a LibJIT value:
//!
//! ```rust
//! extern crate jit;
//! #[phase(syntax)]
//! extern crate jit_macros;
//! use jit::{C_DECL, Context, Function, Type, Types, get_type};
//!
//! #[jit_compile]
//! struct Pos {
//!     x: f64,
//!     y: f64,
//!     z: f64
//! }
//! fn main() {
//!     let cx = Context::new();
//!     cx.build(|| {
//!         // build the IR
//!         let sig = get_type(fn() -> Pos);
//!         let func = Function::new(cx, sig);
//!         let result = Pos {
//!             x: -10,
//!             y: 0,
//!             z: 10
//!         }.compile();
//!         func.insn_return(&result);
//!         /// run the IR
//!         func.compile();
//!         let rfunc:fn() -> Pos = func.closure0();
//!         assert_eq(rfunc().x, -10)
//!     });
//! }
//! ```
extern crate rustc;
extern crate syntax;

use rustc::plugin::Registry;
use std::gc::{Gc, GC};
use syntax::ast::*;
use syntax::codemap::*;
use syntax::ext::base::*;
use syntax::ext::build::*;
use syntax::ext::quote::rt::*;
use syntax::parse::token::*;
use syntax::owned_slice::OwnedSlice;

fn gen_compile(cx:&mut ExtCtxt, pos:Span, _:Gc<MetaItem>, item:Gc<Item>, cb:|Gc<Item>|) {
    let fields:Vec<(Option<String>, P<Expr>)> = match item.node {
        ItemStruct(def, _) if def.is_virtual => {
            cx.span_err(pos, "Virtual structs cannot be JIT compiled");
            fail!("...")
        },
        ItemStruct(def, _) => {
            def.fields.iter().map(|field| {
                let name:Option<String> = match field.node.kind {
                    NamedField(id, _) => Some(id.to_source()),
                    UnnamedField(_) => None
                };
                let ty = field.node.ty;
                let ty_expr = quote_expr!(&*cx, &::jit::get_type::<$ty>());
                (name, ty_expr)
            }).collect()
        },
        _ => {
            cx.span_err(pos, "Compile can only be automatically implemented for structs");
            fail!("...")
        }
    };
    let jit_type_t = cx.ty_path(cx.path_global(pos, vec!(cx.ident_of("jit"), cx.ident_of("Type"))), None);
    let create_struct = cx.expr_path(cx.path_global(pos, vec!(cx.ident_of("jit"), cx.ident_of("Type"), cx.ident_of("create_struct"))));
    let jit_type_body = cx.block(pos, vec!(
        cx.stmt_let(pos, false, cx.ident_of("ty"), {
            let mut fields = cx.expr_vec(pos, fields.iter().map(|&(_, ex)| ex).collect());
            fields = cx.expr_method_call(pos, fields, cx.ident_of("as_mut_slice"), vec!());
            cx.expr_call(pos, create_struct, vec!(fields))
        }),
    ), Some(cx.expr_ident(pos, cx.ident_of("ty"))));
    let curr_item = cx.ty_ident(pos, item.ident);
    let jit_type = Method {
        ident: cx.ident_of("jit_type"),
        attrs: vec!(
            cx.attribute(pos, cx.meta_word(pos, InternedString::new("inline")))
        ),
        generics: Generics {
            lifetimes: vec!(),
            ty_params: OwnedSlice::empty()
        },
        explicit_self: Spanned {
            node: SelfStatic,
            span: pos
        },
        fn_style: NormalFn,
        decl: cx.fn_decl(
            vec!(
                cx.arg(pos, cx.ident_of("_"), cx.ty_option(curr_item))
            ),
            jit_type_t
        ),
        body: jit_type_body,
        span: pos,
        id: item.id,
        vis: Inherited
    };
    let lifetime = cx.lifetime(pos, cx.ident_of("'a").name);
    let compile_body = cx.block(pos, vec!(
        cx.stmt_let(pos, false, cx.ident_of("ty"), cx.expr_call(pos, quote_expr!(&*cx, ::jit::get_type::<$curr_item>), vec!())),
        cx.stmt_let(pos, false, cx.ident_of("val"), cx.expr_call(pos, quote_expr!(&*cx, ::jit::Value::new), vec!(
            cx.expr_ident(pos, cx.ident_of("func")),
            cx.expr_addr_of(pos, cx.expr_ident(pos, cx.ident_of("ty")))
        )))
    ), Some(cx.expr_ident(pos, cx.ident_of("val"))));
    let jit_func_t = cx.ty_path(cx.path_all(pos, true, vec!(cx.ident_of("jit"), cx.ident_of("Function")), vec!(lifetime.clone()), vec!()), None);
    let jit_val_t = cx.ty_path(cx.path_all(pos, true, vec!(cx.ident_of("jit"), cx.ident_of("Value")), vec!(lifetime.clone()), vec!()), None);
    let compile = Method {
        ident: cx.ident_of("compile"),
        attrs: vec!(
            cx.attribute(pos, cx.meta_word(pos, InternedString::new("inline")))
        ),
        generics: Generics {
            lifetimes: vec!(
                lifetime.clone()
            ),
            ty_params: OwnedSlice::empty()
        },
        explicit_self: Spanned {
            node: SelfRegion(None, MutImmutable),
            span: pos
        },
        fn_style: NormalFn,
        decl: cx.fn_decl(
            vec!(
                Arg::new_self(pos, MutImmutable),
                cx.arg(pos, cx.ident_of("func"), cx.ty_rptr(pos, jit_func_t, None, MutImmutable))
            ),
            jit_val_t
        ),
        body: compile_body,
        span: pos,
        id: item.id,
        vis: Inherited
    };
    let methods = vec!(
        box(GC) jit_type,
        box(GC) compile
    );
    let node = ItemImpl(
        Generics {
            lifetimes: vec!(),
            ty_params: OwnedSlice::empty()
        },
        Some(cx.trait_ref(cx.path_global(pos, vec!(cx.ident_of("jit"), cx.ident_of("Compile"))))),
        cx.ty_path(cx.path_ident(pos, item.ident), None),
        methods
    );
    cb(cx.item(pos, cx.ident_of("Compile"), vec!(), node));
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(intern("jit_compile"), ItemDecorator(gen_compile));
}<|endoftext|>


<commit_before>use std::{collections::HashMap,
          env,
          ops::Deref,
          path::PathBuf,
          result};

use crate::hcore::{fs,
                   os::{process::{ShutdownSignal,
                                  ShutdownTimeout},
                        users},
                   package::{PackageIdent,
                             PackageInstall},
                   util::{deserialize_using_from_str,
                          serialize_using_to_string}};
use serde::{ser::SerializeStruct,
            Serialize,
            Serializer};

use crate::{error::{Error,
                    Result},
            util::path};

const DEFAULT_USER: &str = "hab";
const DEFAULT_GROUP: &str = "hab";

const PATH_KEY: &str = "PATH";

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Env(HashMap<String, String>);

impl Deref for Env {
    type Target = HashMap<String, String>;

    fn deref(&self) -> &Self::Target { &self.0 }
}

impl From<HashMap<String, String>> for Env {
    fn from(inner_map: HashMap<String, String>) -> Self { Env(inner_map) }
}

impl Env {
    /// Modifies PATH env with the full run path for this package. This path is composed of any
    /// binary paths specified by this package, or its TDEPS, plus a path to a BusyBox(non-windows),
    /// plus the existing value of the PATH variable.
    ///
    /// This means we work on any operating system, as long as you can invoke the Supervisor,
    /// without having to worry much about context.
    pub fn new(package: &PackageInstall) -> Result<Self> {
        let mut env = package.environment_for_command()?;
        let path = Self::transform_path(env.get(PATH_KEY))?;
        env.insert(PATH_KEY.to_string(), path);
        Ok(Env(env))
    }

    fn transform_path(path: Option<&String>) -> Result<String> {
        let mut paths: Vec<PathBuf> = match path {
            Some(path) => env::split_paths(&path).collect(),
            None => vec![],
        };

        path::append_interpreter_and_path(&mut paths)
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pkg {
    #[serde(deserialize_with = "deserialize_using_from_str",
            serialize_with = "serialize_using_to_string")]
    pub ident: PackageIdent,
    pub origin: String,
    pub name: String,
    pub version: String,
    pub release: String,
    pub deps: Vec<PackageIdent>,
    pub env: Env,
    pub exposes: Vec<String>,
    pub exports: HashMap<String, String>,
    pub path: PathBuf,
    pub svc_path: PathBuf,
    pub svc_config_path: PathBuf,
    pub svc_config_install_path: PathBuf,
    pub svc_data_path: PathBuf,
    pub svc_files_path: PathBuf,
    pub svc_static_path: PathBuf,
    pub svc_var_path: PathBuf,
    pub svc_pid_file: PathBuf,
    pub svc_run: PathBuf,
    pub svc_user: String,
    pub svc_group: String,
    pub shutdown_signal: ShutdownSignal,
    pub shutdown_timeout: ShutdownTimeout,
}

impl Pkg {
    pub fn from_install(package: &PackageInstall) -> Result<Self> {
        let (svc_user, svc_group) = get_user_and_group(&package)?;
        let pkg =
            Pkg { svc_path: fs::svc_path(&package.ident.name),
                  svc_config_path: fs::svc_config_path(&package.ident.name),
                  svc_config_install_path: fs::svc_config_install_path(&package.ident.name),
                  svc_data_path: fs::svc_data_path(&package.ident.name),
                  svc_files_path: fs::svc_files_path(&package.ident.name),
                  svc_run: fs::svc_path(&package.ident.name).join("run"),
                  svc_static_path: fs::svc_static_path(&package.ident.name),
                  svc_var_path: fs::svc_var_path(&package.ident.name),
                  svc_pid_file: fs::svc_pid_file(&package.ident.name),
                  svc_user,
                  svc_group,
                  env: Env::new(&package)?,
                  deps: package.tdeps()?,
                  exposes: package.exposes()?,
                  exports: package.exports()?,
                  path: package.installed_path.clone(),
                  ident: package.ident.clone(),
                  origin: package.ident.origin.clone(),
                  name: package.ident.name.clone(),
                  version: package.ident
                                  .version
                                  .clone()
                                  .expect("No package version in PackageInstall"),
                  release: package.ident
                                  .release
                                  .clone()
                                  .expect("No package release in PackageInstall"),
                  shutdown_signal: package.shutdown_signal()?.unwrap_or_else(Default::default),
                  shutdown_timeout: package.shutdown_timeout()?.unwrap_or_else(Default::default) };
        Ok(pkg)
    }
}

/// This is a proxy struct to represent the data about a Pkg that we actually want to be
/// serialized, similar to ServiceProxy
pub struct PkgProxy<'a> {
    pkg: &'a Pkg,
}

impl<'a> PkgProxy<'a> {
    pub fn new(p: &'a Pkg) -> Self { PkgProxy { pkg: &p } }

    pub fn dependencies(&self) -> Vec<String> {
        self.pkg.deps.iter().map(PackageIdent::to_string).collect()
    }
}

impl<'a> Serialize for PkgProxy<'a> {
    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
        where S: Serializer
    {
        let p = &self.pkg;
        let mut strukt = serializer.serialize_struct("pkg", 21)?;
        strukt.serialize_field("ident", &p.ident.to_string())?;
        strukt.serialize_field("origin", &p.origin)?;
        strukt.serialize_field("name", &p.name)?;
        strukt.serialize_field("version", &p.version)?;
        strukt.serialize_field("release", &p.release)?;
        strukt.serialize_field("deps", &p.deps)?;
        strukt.serialize_field("dependencies", &self.dependencies())?;
        strukt.serialize_field("env", &p.env)?;
        strukt.serialize_field("exposes", &p.exposes)?;
        strukt.serialize_field("exports", &p.exports)?;
        strukt.serialize_field("path", &p.path)?;
        strukt.serialize_field("svc_path", &p.svc_path)?;
        strukt.serialize_field("svc_config_path", &p.svc_config_path)?;
        strukt.serialize_field("svc_config_install_path", &p.svc_config_install_path)?;
        strukt.serialize_field("svc_data_path", &p.svc_data_path)?;
        strukt.serialize_field("svc_files_path", &p.svc_files_path)?;
        strukt.serialize_field("svc_static_path", &p.svc_static_path)?;
        strukt.serialize_field("svc_var_path", &p.svc_var_path)?;
        strukt.serialize_field("svc_pid_file", &p.svc_pid_file)?;
        strukt.serialize_field("svc_run", &p.svc_run)?;
        strukt.serialize_field("svc_user", &p.svc_user)?;
        strukt.serialize_field("svc_group", &p.svc_group)?;
        strukt.end()
    }
}

/// check and see if a user/group is specified in package metadata.
/// if not, we'll try and use hab/hab.
/// If hab/hab doesn't exist, try to use (current username, current group).
/// If that doesn't work, then give up.
#[cfg(unix)]
fn get_user_and_group(pkg_install: &PackageInstall) -> Result<(String, String)> {
    if let Some((user, group)) = get_pkg_user_and_group(&pkg_install)? {
        Ok((user, group))
    } else {
        let defaults = default_user_and_group()?;
        Ok(defaults)
    }
}

/// check and see if a user/group is specified in package metadata.
/// if not, we'll try and use hab/hab.
/// If hab/hab doesn't exist, try to use (current username, current group).
/// If that doesn't work, then give up.
/// Windows will also check if hab exists if it was the given user name
/// If it does not exist then fall back to the current username
/// This is because historically windows plans defaulted to
/// the hab pkg_svc_user even if not explicitly provided
#[cfg(windows)]
fn get_user_and_group(pkg_install: &PackageInstall) -> Result<(String, String)> {
    match get_pkg_user_and_group(&pkg_install)? {
        Some((ref user, ref _group)) if user == DEFAULT_USER => Ok(default_user_and_group()?),
        Some((user, group)) => Ok((user, group)),
        _ => Ok(default_user_and_group()?),
    }
}

/// This function checks to see if a custom SVC_USER and SVC_GROUP has
/// been specified as part of the package metadata.
/// If pkg_svc_user and pkg_svc_group have NOT been defined, return None.
fn get_pkg_user_and_group(pkg_install: &PackageInstall) -> Result<Option<(String, String)>> {
    let svc_user = pkg_install.svc_user()?;
    let svc_group = pkg_install.svc_group()?;
    match (svc_user, svc_group) {
        (Some(user), Some(group)) => Ok(Some((user, group))),
        _ => {
            debug!("User/group not specified in package, running with default");
            Ok(None)
        }
    }
}

/// checks to see if hab/hab exists, if not, fall back to
/// current user/group. If that fails, then return an error.
fn default_user_and_group() -> Result<(String, String)> {
    let uid = users::get_uid_by_name(DEFAULT_USER);
    let gid = users::get_gid_by_name(DEFAULT_GROUP);
    match (uid, gid) {
        (Some(_), Some(_)) => Ok((DEFAULT_USER.to_string(), DEFAULT_GROUP.to_string())),
        _ => {
            debug!("hab:hab does NOT exist");
            let user = users::get_current_username();
            let group = users::get_current_groupname();
            match (user, group) {
                (Some(user), Some(group)) => {
                    debug!("Running as {}/{}", user, group);
                    Ok((user, group))
                }
                _ => {
                    Err(Error::PermissionFailed("Can't determine current \
                                                 user:group"
                                                            .to_string()))
                }
            }
        }
    }
}
<commit_msg>Use unwrap_or_default<commit_after>use std::{collections::HashMap,
          env,
          ops::Deref,
          path::PathBuf,
          result};

use crate::hcore::{fs,
                   os::{process::{ShutdownSignal,
                                  ShutdownTimeout},
                        users},
                   package::{PackageIdent,
                             PackageInstall},
                   util::{deserialize_using_from_str,
                          serialize_using_to_string}};
use serde::{ser::SerializeStruct,
            Serialize,
            Serializer};

use crate::{error::{Error,
                    Result},
            util::path};

const DEFAULT_USER: &str = "hab";
const DEFAULT_GROUP: &str = "hab";

const PATH_KEY: &str = "PATH";

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Env(HashMap<String, String>);

impl Deref for Env {
    type Target = HashMap<String, String>;

    fn deref(&self) -> &Self::Target { &self.0 }
}

impl From<HashMap<String, String>> for Env {
    fn from(inner_map: HashMap<String, String>) -> Self { Env(inner_map) }
}

impl Env {
    /// Modifies PATH env with the full run path for this package. This path is composed of any
    /// binary paths specified by this package, or its TDEPS, plus a path to a BusyBox(non-windows),
    /// plus the existing value of the PATH variable.
    ///
    /// This means we work on any operating system, as long as you can invoke the Supervisor,
    /// without having to worry much about context.
    pub fn new(package: &PackageInstall) -> Result<Self> {
        let mut env = package.environment_for_command()?;
        let path = Self::transform_path(env.get(PATH_KEY))?;
        env.insert(PATH_KEY.to_string(), path);
        Ok(Env(env))
    }

    fn transform_path(path: Option<&String>) -> Result<String> {
        let mut paths: Vec<PathBuf> = match path {
            Some(path) => env::split_paths(&path).collect(),
            None => vec![],
        };

        path::append_interpreter_and_path(&mut paths)
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pkg {
    #[serde(deserialize_with = "deserialize_using_from_str",
            serialize_with = "serialize_using_to_string")]
    pub ident: PackageIdent,
    pub origin: String,
    pub name: String,
    pub version: String,
    pub release: String,
    pub deps: Vec<PackageIdent>,
    pub env: Env,
    pub exposes: Vec<String>,
    pub exports: HashMap<String, String>,
    pub path: PathBuf,
    pub svc_path: PathBuf,
    pub svc_config_path: PathBuf,
    pub svc_config_install_path: PathBuf,
    pub svc_data_path: PathBuf,
    pub svc_files_path: PathBuf,
    pub svc_static_path: PathBuf,
    pub svc_var_path: PathBuf,
    pub svc_pid_file: PathBuf,
    pub svc_run: PathBuf,
    pub svc_user: String,
    pub svc_group: String,
    pub shutdown_signal: ShutdownSignal,
    pub shutdown_timeout: ShutdownTimeout,
}

impl Pkg {
    pub fn from_install(package: &PackageInstall) -> Result<Self> {
        let (svc_user, svc_group) = get_user_and_group(&package)?;
        let pkg = Pkg { svc_path: fs::svc_path(&package.ident.name),
                        svc_config_path: fs::svc_config_path(&package.ident.name),
                        svc_config_install_path: fs::svc_config_install_path(&package.ident
                                                                                .name),
                        svc_data_path: fs::svc_data_path(&package.ident.name),
                        svc_files_path: fs::svc_files_path(&package.ident.name),
                        svc_run: fs::svc_path(&package.ident.name).join("run"),
                        svc_static_path: fs::svc_static_path(&package.ident.name),
                        svc_var_path: fs::svc_var_path(&package.ident.name),
                        svc_pid_file: fs::svc_pid_file(&package.ident.name),
                        svc_user,
                        svc_group,
                        env: Env::new(&package)?,
                        deps: package.tdeps()?,
                        exposes: package.exposes()?,
                        exports: package.exports()?,
                        path: package.installed_path.clone(),
                        ident: package.ident.clone(),
                        origin: package.ident.origin.clone(),
                        name: package.ident.name.clone(),
                        version: package.ident
                                        .version
                                        .clone()
                                        .expect("No package version in PackageInstall"),
                        release: package.ident
                                        .release
                                        .clone()
                                        .expect("No package release in PackageInstall"),
                        shutdown_signal: package.shutdown_signal()?.unwrap_or_default(),
                        shutdown_timeout: package.shutdown_timeout()?.unwrap_or_default() };
        Ok(pkg)
    }
}

/// This is a proxy struct to represent the data about a Pkg that we actually want to be
/// serialized, similar to ServiceProxy
pub struct PkgProxy<'a> {
    pkg: &'a Pkg,
}

impl<'a> PkgProxy<'a> {
    pub fn new(p: &'a Pkg) -> Self { PkgProxy { pkg: &p } }

    pub fn dependencies(&self) -> Vec<String> {
        self.pkg.deps.iter().map(PackageIdent::to_string).collect()
    }
}

impl<'a> Serialize for PkgProxy<'a> {
    fn serialize<S>(&self, serializer: S) -> result::Result<S::Ok, S::Error>
        where S: Serializer
    {
        let p = &self.pkg;
        let mut strukt = serializer.serialize_struct("pkg", 21)?;
        strukt.serialize_field("ident", &p.ident.to_string())?;
        strukt.serialize_field("origin", &p.origin)?;
        strukt.serialize_field("name", &p.name)?;
        strukt.serialize_field("version", &p.version)?;
        strukt.serialize_field("release", &p.release)?;
        strukt.serialize_field("deps", &p.deps)?;
        strukt.serialize_field("dependencies", &self.dependencies())?;
        strukt.serialize_field("env", &p.env)?;
        strukt.serialize_field("exposes", &p.exposes)?;
        strukt.serialize_field("exports", &p.exports)?;
        strukt.serialize_field("path", &p.path)?;
        strukt.serialize_field("svc_path", &p.svc_path)?;
        strukt.serialize_field("svc_config_path", &p.svc_config_path)?;
        strukt.serialize_field("svc_config_install_path", &p.svc_config_install_path)?;
        strukt.serialize_field("svc_data_path", &p.svc_data_path)?;
        strukt.serialize_field("svc_files_path", &p.svc_files_path)?;
        strukt.serialize_field("svc_static_path", &p.svc_static_path)?;
        strukt.serialize_field("svc_var_path", &p.svc_var_path)?;
        strukt.serialize_field("svc_pid_file", &p.svc_pid_file)?;
        strukt.serialize_field("svc_run", &p.svc_run)?;
        strukt.serialize_field("svc_user", &p.svc_user)?;
        strukt.serialize_field("svc_group", &p.svc_group)?;
        strukt.end()
    }
}

/// check and see if a user/group is specified in package metadata.
/// if not, we'll try and use hab/hab.
/// If hab/hab doesn't exist, try to use (current username, current group).
/// If that doesn't work, then give up.
#[cfg(unix)]
fn get_user_and_group(pkg_install: &PackageInstall) -> Result<(String, String)> {
    if let Some((user, group)) = get_pkg_user_and_group(&pkg_install)? {
        Ok((user, group))
    } else {
        let defaults = default_user_and_group()?;
        Ok(defaults)
    }
}

/// check and see if a user/group is specified in package metadata.
/// if not, we'll try and use hab/hab.
/// If hab/hab doesn't exist, try to use (current username, current group).
/// If that doesn't work, then give up.
/// Windows will also check if hab exists if it was the given user name
/// If it does not exist then fall back to the current username
/// This is because historically windows plans defaulted to
/// the hab pkg_svc_user even if not explicitly provided
#[cfg(windows)]
fn get_user_and_group(pkg_install: &PackageInstall) -> Result<(String, String)> {
    match get_pkg_user_and_group(&pkg_install)? {
        Some((ref user, ref _group)) if user == DEFAULT_USER => Ok(default_user_and_group()?),
        Some((user, group)) => Ok((user, group)),
        _ => Ok(default_user_and_group()?),
    }
}

/// This function checks to see if a custom SVC_USER and SVC_GROUP has
/// been specified as part of the package metadata.
/// If pkg_svc_user and pkg_svc_group have NOT been defined, return None.
fn get_pkg_user_and_group(pkg_install: &PackageInstall) -> Result<Option<(String, String)>> {
    let svc_user = pkg_install.svc_user()?;
    let svc_group = pkg_install.svc_group()?;
    match (svc_user, svc_group) {
        (Some(user), Some(group)) => Ok(Some((user, group))),
        _ => {
            debug!("User/group not specified in package, running with default");
            Ok(None)
        }
    }
}

/// checks to see if hab/hab exists, if not, fall back to
/// current user/group. If that fails, then return an error.
fn default_user_and_group() -> Result<(String, String)> {
    let uid = users::get_uid_by_name(DEFAULT_USER);
    let gid = users::get_gid_by_name(DEFAULT_GROUP);
    match (uid, gid) {
        (Some(_), Some(_)) => Ok((DEFAULT_USER.to_string(), DEFAULT_GROUP.to_string())),
        _ => {
            debug!("hab:hab does NOT exist");
            let user = users::get_current_username();
            let group = users::get_current_groupname();
            match (user, group) {
                (Some(user), Some(group)) => {
                    debug!("Running as {}/{}", user, group);
                    Ok((user, group))
                }
                _ => {
                    Err(Error::PermissionFailed("Can't determine current \
                                                 user:group"
                                                            .to_string()))
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 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::fs;
use std::path::Path;
use std::str::FromStr;

use super::{ExecutionStrategy, Scope};
use common::package_graph::PackageGraph;
use common::ui::{Status, UIWriter, UI};
use error::{Error, Result};
use hcore::error as herror;
use hcore::fs as hfs;
use hcore::package::{all_packages, Identifiable, PackageIdent, PackageInstall};

use hcore::package::list::temp_package_directory;

///
/// Delete a package and all dependencies which are not used by other packages.
/// We do an ordered traverse of the dependencies, updating the graph as we delete a
/// package. This lets us use simple logic where we continually check if the package
/// we're trying to delete currently has no packages depending on it.
///
/// The full logic is:
/// 1. We find the fully qualified package ident and all its dependencies
/// 2. We find all packages on the filesystem and convert them into a graph
/// 3. We do a BFS on the graph to get the dependencies in order
/// 4. We check if the specified package has any reverse deps
///     4a. If there are, we throw an error
///     4b. If not, we delete the package
/// 5. For each dependency we check if there are any packages which depend on it
///     5a. If there are, we skip it
///     5b. If there are not, we delete it from disk and the graph
///
pub fn start(
    ui: &mut UI,
    ident: &PackageIdent,
    fs_root_path: &Path,
    execution_strategy: ExecutionStrategy,
    scope: Scope,
    excludes: Vec<PackageIdent>,
    services: Vec<PackageIdent>,
) -> Result<()> {
    // 1.
    let pkg_install = PackageInstall::load(ident, Some(fs_root_path))?;
    let ident = pkg_install.ident();
    ui.begin(format!("Uninstalling {}", &ident))?;

    if !services.is_empty() {
        ui.status(
            Status::Determining,
            "list of running services in supervisor",
        )?;
        for s in &services {
            ui.status(Status::Found, format!("running service {}", s))?;
        }
    }

    let packages = load_all_packages(&fs_root_path)?;
    let mut graph = PackageGraph::new();
    graph.build(&packages)?;

    // 3.
    let deps = graph.ordered_deps(&ident);

    // 4.
    match graph.count_rdeps(&ident) {
        None => {
            // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
            // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
            // dependency tree anyway
            ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &ident))?;
        }
        Some(0) => {
            maybe_delete(
                ui,
                &fs_root_path,
                &pkg_install,
                &execution_strategy,
                &excludes,
                &services,
            )?;
            graph.remove(&ident);
        }
        Some(c) => {
            return Err(Error::CannotRemovePackage(ident.clone(), c));
        }
    }

    // 5.
    let mut count = 0;
    match scope {
        Scope::Package => {
            ui.status(Status::Skipping, "dependencies (--no-deps specified)")?;
        }
        Scope::PackageAndDependencies => {
            for p in &deps {
                match graph.count_rdeps(&p) {
                    None => {
                        // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
                        // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
                        // dependency tree anyway
                        ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &p))?;
                    }
                    Some(0) => {
                        let install = packages.iter().find(|&i| i.ident() == p).unwrap();
                        maybe_delete(
                            ui,
                            &fs_root_path,
                            &install,
                            &execution_strategy,
                            &excludes,
                            &services,
                        )?;

                        graph.remove(&p);
                        count = count + 1;
                    }
                    Some(c) => {
                        ui.status(
                            Status::Skipping,
                            format!("{}. It is a dependency of {} packages", &p, c),
                        )?;
                    }
                }
            }
        }
    }

    match execution_strategy {
        ExecutionStrategy::DryRun => {
            ui.end(format!(
                "Would uninstall {} and {} dependencies (Dry run)",
                &ident, count
            ))?;
        }
        ExecutionStrategy::Run => {
            ui.end(format!(
                "Uninstall of {} and {} dependencies complete",
                &ident, count
            ))?;
        }
    };
    Ok(())
}

/// Delete a package from disk, depending upon the ExecutionStrategy supplied
///
/// Returns:
///   Ok(true) - package is deleted
///   Ok(false) - package would be deleted but it's a dry run
///   Err(_) -  IO problem deleting package from filesystem
fn maybe_delete(
    ui: &mut UI,
    fs_root_path: &Path,
    install: &PackageInstall,
    strategy: &ExecutionStrategy,
    excludes: &Vec<PackageIdent>,
    services: &Vec<PackageIdent>,
) -> Result<bool> {
    let ident = install.ident();
    let pkg_root_path = hfs::pkg_root_path(Some(fs_root_path));

    let hab = PackageIdent::from_str("core/hab")?;
    if ident.satisfies(&hab) {
        ui.status(
            Status::Skipping,
            format!("{}. You can't uninstall core/hab", &ident),
        )?;
        return Ok(false);
    }

    let is_running = services.iter().any(|i| i.satisfies(ident));
    if is_running {
        ui.status(
            Status::Skipping,
            format!("{}. It is currently running in the supervisor", &ident),
        )?;
        return Ok(false);
    }

    // The excludes list could be looser than the fully qualified idents.  E.g. if core/redis is on the
    // exclude list then we should exclude core/redis/1.1.0/20180608091936.  We use the `Identifiable`
    // trait which supplies this logic for PackageIdents
    let should_exclude = excludes.iter().any(|i| i.satisfies(ident));
    if should_exclude {
        ui.status(
            Status::Skipping,
            format!("{}. It is on the exclusion list", &ident),
        )?;
        return Ok(false);
    }

    match strategy {
        ExecutionStrategy::DryRun => {
            ui.status(Status::DryRunDeleting, &ident)?;
            Ok(false)
        }
        ExecutionStrategy::Run => {
            ui.status(Status::Deleting, &ident)?;
            let pkg_dir = install.installed_path();
            do_clean_delete(&pkg_root_path, &pkg_dir)
        }
    }
}

/// Delete empty parent directories from a given path. don't traverse above
/// the `pkg_root_path`
fn do_clean_delete(pkg_root_path: &Path, real_install_path: &Path) -> Result<bool> {
    // This match will always return Ok(Path) as the install path is always 4 levels
    // below the pkg_root_path
    match real_install_path.parent() {
        Some(real_install_base) => {
            let temp_install_path = temp_package_directory(real_install_path)?;
            fs::rename(&real_install_path, &temp_install_path)?;
            fs::remove_dir_all(&temp_install_path)?;

            for p in real_install_base.ancestors() {
                // Let's be safe and not rm below the package directories
                if p == pkg_root_path {
                    break;
                }
                match p.read_dir() {
                    Ok(contents) => {
                        // This will calculate the amount of items in the directory
                        match contents.collect::<Vec<_>>().len() {
                            0 => fs::remove_dir(&p)?,
                            _ => break,
                        }
                    }
                    Err(e) => return Err(Error::HabitatCore(herror::Error::IO(e))),
                }
            }
            Ok(true)
        }
        None => unreachable!("Install path doesn't have a parent"),
    }
}

fn load_all_packages(fs_root_path: &Path) -> Result<Vec<PackageInstall>> {
    let package_path = hfs::pkg_root_path(Some(&fs_root_path));
    let idents = all_packages(&package_path)?;

    let mut result = Vec::with_capacity(idents.len());
    for i in idents {
        let pkg_install = PackageInstall::load(&i, Some(fs_root_path))?;
        result.push(pkg_install);
    }
    Ok(result)
}
<commit_msg>review comments<commit_after>// Copyright (c) 2018 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::fs;
use std::path::Path;
use std::str::FromStr;

use super::{ExecutionStrategy, Scope};
use common::package_graph::PackageGraph;
use common::ui::{Status, UIWriter, UI};
use error::{Error, Result};
use hcore::error as herror;
use hcore::fs as hfs;
use hcore::package::{all_packages, Identifiable, PackageIdent, PackageInstall};

use hcore::package::list::temp_package_directory;

///
/// Delete a package and all dependencies which are not used by other packages.
/// We do an ordered traverse of the dependencies, updating the graph as we delete a
/// package. This lets us use simple logic where we continually check if the package
/// we're trying to delete currently has no packages depending on it.
///
/// The full logic is:
/// 1. We find the fully qualified package ident and all its dependencies
/// 2. We find all packages on the filesystem and convert them into a graph
/// 3. We do a BFS on the graph to get the dependencies in order
/// 4. We check if the specified package has any reverse deps
///     4a. If there are, we throw an error
///     4b. If not, we delete the package
/// 5. For each dependency we check if there are any packages which depend on it
///     5a. If there are, we skip it
///     5b. If there are not, we delete it from disk and the graph
///
/// `excludes` is a list of user-supplied `PackageIdent`s.
/// `services` is a list of fully-qualified `PackageIdent`s which are currently
///    running in a supervisor out of the `fs_root_path`.
///
pub fn start(
    ui: &mut UI,
    ident: &PackageIdent,
    fs_root_path: &Path,
    execution_strategy: ExecutionStrategy,
    scope: Scope,
    excludes: Vec<PackageIdent>,
    services: Vec<PackageIdent>,
) -> Result<()> {
    // 1.
    let pkg_install = PackageInstall::load(ident, Some(fs_root_path))?;
    let ident = pkg_install.ident();
    ui.begin(format!("Uninstalling {}", &ident))?;

    if !services.is_empty() {
        ui.status(
            Status::Determining,
            "list of running services in supervisor",
        )?;
        for s in &services {
            ui.status(Status::Found, format!("running service {}", s))?;
        }
    }

    // 2.
    let packages = load_all_packages(&fs_root_path)?;
    let mut graph = PackageGraph::new();
    graph.build(&packages)?;

    // 3.
    let deps = graph.ordered_deps(&ident);

    // 4.
    match graph.count_rdeps(&ident) {
        None => {
            // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
            // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
            // dependency tree anyway
            ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &ident))?;
        }
        Some(0) => {
            maybe_delete(
                ui,
                &fs_root_path,
                &pkg_install,
                &execution_strategy,
                &excludes,
                &services,
            )?;
            graph.remove(&ident);
        }
        Some(c) => {
            return Err(Error::CannotRemovePackage(ident.clone(), c));
        }
    }

    // 5.
    let mut count = 0;
    match scope {
        Scope::Package => {
            ui.status(Status::Skipping, "dependencies (--no-deps specified)")?;
        }
        Scope::PackageAndDependencies => {
            for p in &deps {
                match graph.count_rdeps(&p) {
                    None => {
                        // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
                        // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
                        // dependency tree anyway
                        ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &p))?;
                    }
                    Some(0) => {
                        let install = packages.iter().find(|&i| i.ident() == p).unwrap();
                        maybe_delete(
                            ui,
                            &fs_root_path,
                            &install,
                            &execution_strategy,
                            &excludes,
                            &services,
                        )?;

                        graph.remove(&p);
                        count = count + 1;
                    }
                    Some(c) => {
                        ui.status(
                            Status::Skipping,
                            format!("{}. It is a dependency of {} packages", &p, c),
                        )?;
                    }
                }
            }
        }
    }

    match execution_strategy {
        ExecutionStrategy::DryRun => {
            ui.end(format!(
                "Would uninstall {} and {} dependencies (Dry run)",
                &ident, count
            ))?;
        }
        ExecutionStrategy::Run => {
            ui.end(format!(
                "Uninstall of {} and {} dependencies complete",
                &ident, count
            ))?;
        }
    };
    Ok(())
}

/// Delete a package from disk, depending upon the ExecutionStrategy supplied
///
/// Returns:
///   Ok(true) - package is deleted
///   Ok(false) - package would be deleted but it's a dry run
///   Err(_) -  IO problem deleting package from filesystem
fn maybe_delete(
    ui: &mut UI,
    fs_root_path: &Path,
    install: &PackageInstall,
    strategy: &ExecutionStrategy,
    excludes: &Vec<PackageIdent>,
    services: &Vec<PackageIdent>,
) -> Result<bool> {
    let ident = install.ident();
    let pkg_root_path = hfs::pkg_root_path(Some(fs_root_path));

    let hab = PackageIdent::from_str("core/hab")?;
    if ident.satisfies(&hab) {
        ui.status(
            Status::Skipping,
            format!("{}. You can't uninstall core/hab", &ident),
        )?;
        return Ok(false);
    }

    let is_running = services.iter().any(|i| i.satisfies(ident));
    if is_running {
        ui.status(
            Status::Skipping,
            format!("{}. It is currently running in the supervisor", &ident),
        )?;
        return Ok(false);
    }

    // The excludes list could be looser than the fully qualified idents.  E.g. if core/redis is on the
    // exclude list then we should exclude core/redis/1.1.0/20180608091936.  We use the `Identifiable`
    // trait which supplies this logic for PackageIdents
    let should_exclude = excludes.iter().any(|i| i.satisfies(ident));
    if should_exclude {
        ui.status(
            Status::Skipping,
            format!("{}. It is on the exclusion list", &ident),
        )?;
        return Ok(false);
    }

    match strategy {
        ExecutionStrategy::DryRun => {
            ui.status(Status::DryRunDeleting, &ident)?;
            Ok(false)
        }
        ExecutionStrategy::Run => {
            ui.status(Status::Deleting, &ident)?;
            let pkg_dir = install.installed_path();
            do_clean_delete(&pkg_root_path, &pkg_dir)
        }
    }
}

/// Delete empty parent directories from a given path. don't traverse above
/// the `pkg_root_path`
fn do_clean_delete(pkg_root_path: &Path, real_install_path: &Path) -> Result<bool> {
    // This match will always return Ok(Path) as the install path is always 4 levels
    // below the pkg_root_path
    match real_install_path.parent() {
        Some(real_install_base) => {
            let temp_install_path = temp_package_directory(real_install_path)?;
            fs::rename(&real_install_path, &temp_install_path)?;
            fs::remove_dir_all(&temp_install_path)?;

            for p in real_install_base.ancestors() {
                // Let's be safe and not rm below the package directories
                if p == pkg_root_path {
                    break;
                }
                match p.read_dir() {
                    Ok(contents) => {
                        // This will calculate the amount of items in the directory
                        match contents.collect::<Vec<_>>().len() {
                            0 => fs::remove_dir(&p)?,
                            _ => break,
                        }
                    }
                    Err(e) => return Err(Error::HabitatCore(herror::Error::IO(e))),
                }
            }
            Ok(true)
        }
        None => unreachable!("Install path doesn't have a parent"),
    }
}

fn load_all_packages(fs_root_path: &Path) -> Result<Vec<PackageInstall>> {
    let package_path = hfs::pkg_root_path(Some(&fs_root_path));
    let idents = all_packages(&package_path)?;

    let mut result = Vec::with_capacity(idents.len());
    for i in idents {
        let pkg_install = PackageInstall::load(&i, Some(fs_root_path))?;
        result.push(pkg_install);
    }
    Ok(result)
}
<|endoftext|>


<commit_before>//! See `AssistContext`

use algo::find_covering_element;
use hir::Semantics;
use ra_db::{FileId, FileRange};
use ra_fmt::{leading_indent, reindent};
use ra_ide_db::{
    source_change::{SourceChange, SourceFileEdit},
    RootDatabase,
};
use ra_syntax::{
    algo::{self, find_node_at_offset, SyntaxRewriter},
    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextSize,
    TokenAtOffset,
};
use ra_text_edit::TextEditBuilder;

use crate::{
    assist_config::{AssistConfig, SnippetCap},
    Assist, AssistId, GroupLabel, ResolvedAssist,
};
use rustc_hash::FxHashMap;

/// `AssistContext` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat over-engineered approach, given the current needs.
/// The assists workflow consists of two phases. In the first phase, a user asks
/// for the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then necessary to answer "is
///   assist applicable" in the first phase.
/// * second, when we are applying assist, we don't have a guarantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach reminiscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `resolve`. It first check if an edit is applicable (potentially computing
/// info required to compute the actual edit). If it is applicable, and
/// `resolve` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `resolve = false`, and then applying the selected edit again, with
/// `resolve = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagerly :-)
pub(crate) struct AssistContext<'a> {
    pub(crate) config: &'a AssistConfig,
    pub(crate) sema: Semantics<'a, RootDatabase>,
    pub(crate) db: &'a RootDatabase,
    pub(crate) frange: FileRange,
    source_file: SourceFile,
}

impl<'a> AssistContext<'a> {
    pub(crate) fn new(
        sema: Semantics<'a, RootDatabase>,
        config: &'a AssistConfig,
        frange: FileRange,
    ) -> AssistContext<'a> {
        let source_file = sema.parse(frange.file_id);
        let db = sema.db;
        AssistContext { config, sema, db, frange, source_file }
    }

    // NB, this ignores active selection.
    pub(crate) fn offset(&self) -> TextSize {
        self.frange.range.start()
    }

    pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
        self.source_file.syntax().token_at_offset(self.offset())
    }
    pub(crate) fn find_token_at_offset(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
        self.token_at_offset().find(|it| it.kind() == kind)
    }
    pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
        find_node_at_offset(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn find_node_at_offset_with_descend<N: AstNode>(&self) -> Option<N> {
        self.sema.find_node_at_offset_with_descend(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn covering_element(&self) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), self.frange.range)
    }
    // FIXME: remove
    pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), range)
    }
}

pub(crate) struct Assists {
    resolve: bool,
    file: FileId,
    buf: Vec<(Assist, Option<SourceChange>)>,
}

impl Assists {
    pub(crate) fn new_resolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: true, file: ctx.frange.file_id, buf: Vec::new() }
    }
    pub(crate) fn new_unresolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: false, file: ctx.frange.file_id, buf: Vec::new() }
    }

    pub(crate) fn finish_unresolved(self) -> Vec<Assist> {
        assert!(!self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| {
                assert!(edit.is_none());
                label
            })
            .collect()
    }

    pub(crate) fn finish_resolved(self) -> Vec<ResolvedAssist> {
        assert!(self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| ResolvedAssist { assist: label, source_change: edit.unwrap() })
            .collect()
    }

    pub(crate) fn add(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl(label, f)
    }
    pub(crate) fn add_in_multiple_files(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl_multiple_files(label, f)
    }
    pub(crate) fn add_group(
        &mut self,
        group: &GroupLabel,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), Some(group.clone()), target);
        self.add_impl(label, f)
    }
    fn add_impl(&mut self, label: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
        let source_change = if self.resolve {
            let mut builder = AssistBuilder::new(self.file);
            f(&mut builder);
            Some(builder.finish())
        } else {
            None
        };

        self.buf.push((label, source_change));
        Some(())
    }

    fn add_impl_multiple_files(
        &mut self,
        label: Assist,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        if !self.resolve {
            return None;
        }
        let mut director = AssistDirector::new();
        f(&mut director);
        let changes = director.finish();
        let file_edits: Vec<SourceFileEdit> =
            changes.into_iter().map(|mut change| change.source_file_edits.pop().unwrap()).collect();

        let source_change = SourceChange {
            source_file_edits: file_edits,
            file_system_edits: vec![],
            is_snippet: false,
        };

        self.buf.push((label, Some(source_change)));
        Some(())
    }

    fn finish(mut self) -> Vec<(Assist, Option<SourceChange>)> {
        self.buf.sort_by_key(|(label, _edit)| label.target.len());
        self.buf
    }
}

pub(crate) struct AssistBuilder {
    edit: TextEditBuilder,
    file: FileId,
    is_snippet: bool,
}

impl AssistBuilder {
    pub(crate) fn new(file: FileId) -> AssistBuilder {
        AssistBuilder { edit: TextEditBuilder::default(), file, is_snippet: false }
    }

    /// Remove specified `range` of text.
    pub(crate) fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    /// Append specified `text` at the given `offset`
    pub(crate) fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    /// Append specified `snippet` at the given `offset`
    pub(crate) fn insert_snippet(
        &mut self,
        _cap: SnippetCap,
        offset: TextSize,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.insert(offset, snippet);
    }
    /// Replaces specified `range` of text with a given string.
    pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    /// Replaces specified `range` of text with a given `snippet`.
    pub(crate) fn replace_snippet(
        &mut self,
        _cap: SnippetCap,
        range: TextRange,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.replace(range, snippet);
    }
    pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }
    /// Replaces specified `node` of text with a given string, reindenting the
    /// string to maintain `node`'s existing indent.
    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
    pub(crate) fn replace_node_and_indent(
        &mut self,
        node: &SyntaxNode,
        replace_with: impl Into<String>,
    ) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, &indent)
        }
        self.replace(node.text_range(), replace_with)
    }
    pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
        let node = rewriter.rewrite_root().unwrap();
        let new = rewriter.rewrite(&node);
        algo::diff(&node, &new).into_text_edit(&mut self.edit)
    }

    // FIXME: better API
    pub(crate) fn set_file(&mut self, assist_file: FileId) {
        self.file = assist_file;
    }

    // FIXME: kill this API
    /// Get access to the raw `TextEditBuilder`.
    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
        &mut self.edit
    }

    fn finish(self) -> SourceChange {
        let edit = self.edit.finish();
        let source_file_edit = SourceFileEdit { file_id: self.file, edit };
        let mut res: SourceChange = source_file_edit.into();
        if self.is_snippet {
            res.is_snippet = true;
        }
        res
    }
}

pub(crate) struct AssistDirector {
    source_changes: Vec<SourceChange>,
    builders: FxHashMap<FileId, AssistBuilder>,
}

impl AssistDirector {
    fn new() -> AssistDirector {
        AssistDirector { source_changes: vec![], builders: FxHashMap::default() }
    }

    pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) {
        let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id));
        f(&mut builder);
    }

    fn finish(mut self) -> Vec<SourceChange> {
        for (_, builder) in
            self.builders.into_iter().collect::<Vec<(FileId, AssistBuilder)>>()
        {
            self.source_changes.push(builder.finish());
        }
        self.source_changes
    }
}
<commit_msg>More formatting<commit_after>//! See `AssistContext`

use algo::find_covering_element;
use hir::Semantics;
use ra_db::{FileId, FileRange};
use ra_fmt::{leading_indent, reindent};
use ra_ide_db::{
    source_change::{SourceChange, SourceFileEdit},
    RootDatabase,
};
use ra_syntax::{
    algo::{self, find_node_at_offset, SyntaxRewriter},
    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextSize,
    TokenAtOffset,
};
use ra_text_edit::TextEditBuilder;

use crate::{
    assist_config::{AssistConfig, SnippetCap},
    Assist, AssistId, GroupLabel, ResolvedAssist,
};
use rustc_hash::FxHashMap;

/// `AssistContext` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat over-engineered approach, given the current needs.
/// The assists workflow consists of two phases. In the first phase, a user asks
/// for the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then necessary to answer "is
///   assist applicable" in the first phase.
/// * second, when we are applying assist, we don't have a guarantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach reminiscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `resolve`. It first check if an edit is applicable (potentially computing
/// info required to compute the actual edit). If it is applicable, and
/// `resolve` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `resolve = false`, and then applying the selected edit again, with
/// `resolve = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagerly :-)
pub(crate) struct AssistContext<'a> {
    pub(crate) config: &'a AssistConfig,
    pub(crate) sema: Semantics<'a, RootDatabase>,
    pub(crate) db: &'a RootDatabase,
    pub(crate) frange: FileRange,
    source_file: SourceFile,
}

impl<'a> AssistContext<'a> {
    pub(crate) fn new(
        sema: Semantics<'a, RootDatabase>,
        config: &'a AssistConfig,
        frange: FileRange,
    ) -> AssistContext<'a> {
        let source_file = sema.parse(frange.file_id);
        let db = sema.db;
        AssistContext { config, sema, db, frange, source_file }
    }

    // NB, this ignores active selection.
    pub(crate) fn offset(&self) -> TextSize {
        self.frange.range.start()
    }

    pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
        self.source_file.syntax().token_at_offset(self.offset())
    }
    pub(crate) fn find_token_at_offset(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
        self.token_at_offset().find(|it| it.kind() == kind)
    }
    pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
        find_node_at_offset(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn find_node_at_offset_with_descend<N: AstNode>(&self) -> Option<N> {
        self.sema.find_node_at_offset_with_descend(self.source_file.syntax(), self.offset())
    }
    pub(crate) fn covering_element(&self) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), self.frange.range)
    }
    // FIXME: remove
    pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), range)
    }
}

pub(crate) struct Assists {
    resolve: bool,
    file: FileId,
    buf: Vec<(Assist, Option<SourceChange>)>,
}

impl Assists {
    pub(crate) fn new_resolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: true, file: ctx.frange.file_id, buf: Vec::new() }
    }
    pub(crate) fn new_unresolved(ctx: &AssistContext) -> Assists {
        Assists { resolve: false, file: ctx.frange.file_id, buf: Vec::new() }
    }

    pub(crate) fn finish_unresolved(self) -> Vec<Assist> {
        assert!(!self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| {
                assert!(edit.is_none());
                label
            })
            .collect()
    }

    pub(crate) fn finish_resolved(self) -> Vec<ResolvedAssist> {
        assert!(self.resolve);
        self.finish()
            .into_iter()
            .map(|(label, edit)| ResolvedAssist { assist: label, source_change: edit.unwrap() })
            .collect()
    }

    pub(crate) fn add(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl(label, f)
    }
    pub(crate) fn add_in_multiple_files(
        &mut self,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), None, target);
        self.add_impl_multiple_files(label, f)
    }
    pub(crate) fn add_group(
        &mut self,
        group: &GroupLabel,
        id: AssistId,
        label: impl Into<String>,
        target: TextRange,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<()> {
        let label = Assist::new(id, label.into(), Some(group.clone()), target);
        self.add_impl(label, f)
    }
    fn add_impl(&mut self, label: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
        let source_change = if self.resolve {
            let mut builder = AssistBuilder::new(self.file);
            f(&mut builder);
            Some(builder.finish())
        } else {
            None
        };

        self.buf.push((label, source_change));
        Some(())
    }

    fn add_impl_multiple_files(
        &mut self,
        label: Assist,
        f: impl FnOnce(&mut AssistDirector),
    ) -> Option<()> {
        if !self.resolve {
            return None;
        }
        let mut director = AssistDirector::new();
        f(&mut director);
        let changes = director.finish();
        let file_edits: Vec<SourceFileEdit> =
            changes.into_iter().map(|mut change| change.source_file_edits.pop().unwrap()).collect();

        let source_change = SourceChange {
            source_file_edits: file_edits,
            file_system_edits: vec![],
            is_snippet: false,
        };

        self.buf.push((label, Some(source_change)));
        Some(())
    }

    fn finish(mut self) -> Vec<(Assist, Option<SourceChange>)> {
        self.buf.sort_by_key(|(label, _edit)| label.target.len());
        self.buf
    }
}

pub(crate) struct AssistBuilder {
    edit: TextEditBuilder,
    file: FileId,
    is_snippet: bool,
}

impl AssistBuilder {
    pub(crate) fn new(file: FileId) -> AssistBuilder {
        AssistBuilder { edit: TextEditBuilder::default(), file, is_snippet: false }
    }

    /// Remove specified `range` of text.
    pub(crate) fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    /// Append specified `text` at the given `offset`
    pub(crate) fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    /// Append specified `snippet` at the given `offset`
    pub(crate) fn insert_snippet(
        &mut self,
        _cap: SnippetCap,
        offset: TextSize,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.insert(offset, snippet);
    }
    /// Replaces specified `range` of text with a given string.
    pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    /// Replaces specified `range` of text with a given `snippet`.
    pub(crate) fn replace_snippet(
        &mut self,
        _cap: SnippetCap,
        range: TextRange,
        snippet: impl Into<String>,
    ) {
        self.is_snippet = true;
        self.replace(range, snippet);
    }
    pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }
    /// Replaces specified `node` of text with a given string, reindenting the
    /// string to maintain `node`'s existing indent.
    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
    pub(crate) fn replace_node_and_indent(
        &mut self,
        node: &SyntaxNode,
        replace_with: impl Into<String>,
    ) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, &indent)
        }
        self.replace(node.text_range(), replace_with)
    }
    pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
        let node = rewriter.rewrite_root().unwrap();
        let new = rewriter.rewrite(&node);
        algo::diff(&node, &new).into_text_edit(&mut self.edit)
    }

    // FIXME: better API
    pub(crate) fn set_file(&mut self, assist_file: FileId) {
        self.file = assist_file;
    }

    // FIXME: kill this API
    /// Get access to the raw `TextEditBuilder`.
    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
        &mut self.edit
    }

    fn finish(self) -> SourceChange {
        let edit = self.edit.finish();
        let source_file_edit = SourceFileEdit { file_id: self.file, edit };
        let mut res: SourceChange = source_file_edit.into();
        if self.is_snippet {
            res.is_snippet = true;
        }
        res
    }
}

pub(crate) struct AssistDirector {
    source_changes: Vec<SourceChange>,
    builders: FxHashMap<FileId, AssistBuilder>,
}

impl AssistDirector {
    fn new() -> AssistDirector {
        AssistDirector { source_changes: vec![], builders: FxHashMap::default() }
    }

    pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) {
        let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id));
        f(&mut builder);
    }

    fn finish(mut self) -> Vec<SourceChange> {
        for (_, builder) in self.builders.into_iter().collect::<Vec<(FileId, AssistBuilder)>>() {
            self.source_changes.push(builder.finish());
        }
        self.source_changes
    }
}
<|endoftext|>