<commit_before>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::ascii::AsciiExt;
use std::ffi::OsString;
use std::ops::Range;
use std::str;

use super::Input;

/// A `struct` useful for matching borrowed `str`s.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct StrInput<'a> {
    str_ref: &'a str
}

impl<'a> StrInput<'a> {
    /// Creates a new `StrInput` from a `&str`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use pest::inputs::{Input, StrInput};
    /// let input = StrInput::new("asd");
    ///
    /// assert_eq!(input.len(), 3);
    /// ```
    pub fn new(source: &'a str) -> StrInput<'a> {
        StrInput { str_ref: source }
    }
}

#[inline]
unsafe fn line_col(source: &str, pos: usize) -> (usize, usize) {
    if pos > source.len() {
        panic!("position out of bounds");
    }

    let mut pos = pos;
    let slice = &source[..pos];
    let mut chars = slice.chars().peekable();

    let mut line_col = (1, 1);

    while pos != 0 {
        match chars.next() {
            Some('\r') => {
                if let Some(&'\n') = chars.peek() {
                    chars.next();

                    if pos == 1 {
                        pos -= 1;
                    } else {
                        pos -= 2;
                    }

                    line_col = (line_col.0 + 1, 1);
                } else {
                    pos -= 1;
                    line_col = (line_col.0, line_col.1 + 1);
                }
            }
            Some('\n') => {
                pos -= 1;
                line_col = (line_col.0 + 1, 1);
            }
            Some(c) => {
                pos -= c.len_utf8();
                line_col = (line_col.0, line_col.1 + 1);
            }
            None => unreachable!()
        }
    }

    line_col
}

#[inline]
unsafe fn line_of(source: &str, pos: usize) -> &str {
    if pos > source.len() {
        panic!("position out of bounds");
    }

    let start = if pos == 0 {
        0
    } else {
        let start = source.char_indices()
                          .rev()
                          .skip_while(|&(i, _)| i >= pos)
                          .find(|&(_, c)| c == '\n');
        match start {
            Some((i, _)) => i + 1,
            None => 0
        }
    };

    let end = if source.len() == 0 {
        0
    } else if pos == source.len() - 1 {
        let mut end = source.len();

        if end > 0 && source.slice_unchecked(end - 1, end) == "\n" {
            end -= 1;
        }
        if end > 0 && source.slice_unchecked(end - 1, end) == "\r" {
            end -= 1;
        }

        end
    } else {
        let end = source.char_indices()
                        .skip_while(|&(i, _)| i < pos)
                        .find(|&(_, c)| c == '\n');
        let mut end = match end {
            Some((i, _)) => i,
            None => source.len()
        };

        if end > 0 && source.slice_unchecked(end - 1, end) == "\r" {
            end -= 1;
        }

        end
    };

    source.slice_unchecked(start, end)
}

#[inline]
unsafe fn skip(source: &str, n: usize, pos: usize) -> Option<usize> {
    let mut len = 0;
    let mut chars = source.slice_unchecked(pos, source.len()).chars();

    for _ in 0..n {
        if let Some(c) = chars.next() {
            len += c.len_utf8();
        } else {
            return None;
        }
    }

    Some(len)
}

#[inline]
unsafe fn match_string(source: &str, string: &str, pos: usize) -> bool {
    let to = pos + string.len();

    if to <= source.len() {
        let slice = source.slice_unchecked(pos, to);
        slice == string
    } else {
        false
    }
}

#[inline]
unsafe fn match_insensitive(source: &str, string: &str, pos: usize) -> bool {
    let slice = source.slice_unchecked(pos, source.len());

    if slice.is_char_boundary(string.len()) {
        let slice = slice.slice_unchecked(0, string.len());
        slice.eq_ignore_ascii_case(string)
    } else {
        false
    }
}

#[inline]
unsafe fn match_range(source: &str, range: Range<char>, pos: usize) -> Option<usize> {
    let slice = source.slice_unchecked(pos, source.len());

    if let Some(char) = slice.chars().next() {
        if range.start <= char && char <= range.end {
            Some(char.len_utf8())
        } else {
            None
        }
    } else {
        None
    }
}

impl<'a> Input<'a> for StrInput<'a> {
    #[inline]
    fn len(&self) -> usize {
        self.str_ref.len()
    }

    #[inline]
    fn is_empty(&self) -> bool {
        self.str_ref.is_empty()
    }

    #[inline]
    fn file_name(&self) -> Option<OsString> {
        None
    }

    #[inline]
    unsafe fn slice(&self, start: usize, end: usize) -> &'a str {
        self.str_ref.slice_unchecked(start, end)
    }

    unsafe fn line_col(&self, pos: usize) -> (usize, usize) {
        line_col(self.str_ref, pos)
    }

    unsafe fn line_of(&self, pos: usize) -> &'a str {
        line_of(self.str_ref, pos)
    }

    #[inline]
    unsafe fn skip(&self, n: usize, pos: usize) -> Option<usize> {
        skip(self.str_ref, n, pos)
    }

    #[inline]
    unsafe fn match_string(&self, string: &str, pos: usize) -> bool {
        match_string(self.str_ref, string, pos)
    }

    #[inline]
    unsafe fn match_insensitive(&self, string: &str, pos: usize) -> bool {
        match_insensitive(self.str_ref, string, pos)
    }

    #[inline]
    unsafe fn match_range(&self, range: Range<char>, pos: usize) -> Option<usize> {
        match_range(self.str_ref, range, pos)
    }
}

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

    #[test]
    fn empty() {
        let input2 = StrInput::new("");

        unsafe {
            assert!(input2.is_empty());
            assert!(input2.match_string("", 0));
            assert!(!input2.match_string("a", 0));
        }
    }

    #[test]
    fn parts() {
        let input2 = StrInput::new("asdasdf");

        unsafe {
            assert!(!input2.is_empty());
            assert!(input2.match_string("asd", 0));
            assert!(input2.match_string("asdf", 3));
        }
    }

    #[test]
    fn len() {
        assert_eq!(StrInput::new("asdasdf").len(), 7);
    }

    #[test]
    fn slice() {
        let input2 = StrInput::new("asdasdf");

        unsafe {
            assert_eq!(input2.slice(1, 3), "sd");
        }
    }

    #[test]
    fn line_col() {
        let input2 = StrInput::new("a\rb\nc\r\nd嗨");

        unsafe {
            assert_eq!(input2.line_col(0), (1, 1));
            assert_eq!(input2.line_col(1), (1, 2));
            assert_eq!(input2.line_col(2), (1, 3));
            assert_eq!(input2.line_col(3), (1, 4));
            assert_eq!(input2.line_col(4), (2, 1));
            assert_eq!(input2.line_col(5), (2, 2));
            assert_eq!(input2.line_col(6), (2, 3));
            assert_eq!(input2.line_col(7), (3, 1));
            assert_eq!(input2.line_col(8), (3, 2));
            assert_eq!(input2.line_col(11), (3, 3));
        }
    }

    #[test]
    fn line_of() {
        let input2 = StrInput::new("a\rb\nc\r\nd嗨");

        unsafe {
            assert_eq!(input2.line_of(0), "a\rb");
            assert_eq!(input2.line_of(1), "a\rb");
            assert_eq!(input2.line_of(2), "a\rb");
            assert_eq!(input2.line_of(3), "a\rb");
            assert_eq!(input2.line_of(4), "c");
            assert_eq!(input2.line_of(5), "c");
            assert_eq!(input2.line_of(6), "c");
            assert_eq!(input2.line_of(7), "d嗨");
            assert_eq!(input2.line_of(8), "d嗨");
            assert_eq!(input2.line_of(11), "d嗨");
        }
    }

    #[test]
    fn line_of_empty() {
        let input2 = StrInput::new("");

        unsafe {
            assert_eq!(input2.line_of(0), "");
        }
    }

    #[test]
    fn line_of_new_line() {
        let input2 = StrInput::new("\n");

        unsafe {
            assert_eq!(input2.line_of(0), "");
        }
    }

    #[test]
    fn line_of_between_new_line() {
        let input2 = StrInput::new("\n\n");

        unsafe {
            assert_eq!(input2.line_of(1), "");
        }
    }

    #[test]
    fn skip_empty() {
        let input2 = StrInput::new("");

        unsafe {
            assert_eq!(input2.skip(0, 0), Some(0));
            assert_eq!(input2.skip(1, 0), None);
        }
    }

    #[test]
    fn skip() {
        let input2 = StrInput::new("d嗨");

        unsafe {
            assert_eq!(input2.skip(0, 0), Some(0));
            assert_eq!(input2.skip(1, 0), Some(1));
            assert_eq!(input2.skip(1, 1), Some(3));
        }
    }

    #[test]
    fn match_range() {
        let input2 = StrInput::new("b");

        unsafe {
            assert!(input2.match_range('a'..'c', 0).is_some());
            assert!(input2.match_range('b'..'b', 0).is_some());
            assert!(input2.match_range('a'..'a', 0).is_none());
            assert!(input2.match_range('c'..'c', 0).is_none());
            assert!(input2.match_range('a'..'嗨', 0).is_some());
        }
    }

    #[test]
    fn match_insensitive() {
        let input2 = StrInput::new("AsdASdF");

        unsafe {
            assert!(input2.match_insensitive("asd", 0));
            assert!(input2.match_insensitive("asdf", 3));
        }
    }
}
<commit_msg>add test re: slice lifetime<commit_after>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::ascii::AsciiExt;
use std::ffi::OsString;
use std::ops::Range;
use std::str;

use super::Input;

/// A `struct` useful for matching borrowed `str`s.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct StrInput<'a> {
    str_ref: &'a str
}

impl<'a> StrInput<'a> {
    /// Creates a new `StrInput` from a `&str`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use pest::inputs::{Input, StrInput};
    /// let input = StrInput::new("asd");
    ///
    /// assert_eq!(input.len(), 3);
    /// ```
    pub fn new(source: &'a str) -> StrInput<'a> {
        StrInput { str_ref: source }
    }
}

#[inline]
unsafe fn line_col(source: &str, pos: usize) -> (usize, usize) {
    if pos > source.len() {
        panic!("position out of bounds");
    }

    let mut pos = pos;
    let slice = &source[..pos];
    let mut chars = slice.chars().peekable();

    let mut line_col = (1, 1);

    while pos != 0 {
        match chars.next() {
            Some('\r') => {
                if let Some(&'\n') = chars.peek() {
                    chars.next();

                    if pos == 1 {
                        pos -= 1;
                    } else {
                        pos -= 2;
                    }

                    line_col = (line_col.0 + 1, 1);
                } else {
                    pos -= 1;
                    line_col = (line_col.0, line_col.1 + 1);
                }
            }
            Some('\n') => {
                pos -= 1;
                line_col = (line_col.0 + 1, 1);
            }
            Some(c) => {
                pos -= c.len_utf8();
                line_col = (line_col.0, line_col.1 + 1);
            }
            None => unreachable!()
        }
    }

    line_col
}

#[inline]
unsafe fn line_of(source: &str, pos: usize) -> &str {
    if pos > source.len() {
        panic!("position out of bounds");
    }

    let start = if pos == 0 {
        0
    } else {
        let start = source.char_indices()
                          .rev()
                          .skip_while(|&(i, _)| i >= pos)
                          .find(|&(_, c)| c == '\n');
        match start {
            Some((i, _)) => i + 1,
            None => 0
        }
    };

    let end = if source.len() == 0 {
        0
    } else if pos == source.len() - 1 {
        let mut end = source.len();

        if end > 0 && source.slice_unchecked(end - 1, end) == "\n" {
            end -= 1;
        }
        if end > 0 && source.slice_unchecked(end - 1, end) == "\r" {
            end -= 1;
        }

        end
    } else {
        let end = source.char_indices()
                        .skip_while(|&(i, _)| i < pos)
                        .find(|&(_, c)| c == '\n');
        let mut end = match end {
            Some((i, _)) => i,
            None => source.len()
        };

        if end > 0 && source.slice_unchecked(end - 1, end) == "\r" {
            end -= 1;
        }

        end
    };

    source.slice_unchecked(start, end)
}

#[inline]
unsafe fn skip(source: &str, n: usize, pos: usize) -> Option<usize> {
    let mut len = 0;
    let mut chars = source.slice_unchecked(pos, source.len()).chars();

    for _ in 0..n {
        if let Some(c) = chars.next() {
            len += c.len_utf8();
        } else {
            return None;
        }
    }

    Some(len)
}

#[inline]
unsafe fn match_string(source: &str, string: &str, pos: usize) -> bool {
    let to = pos + string.len();

    if to <= source.len() {
        let slice = source.slice_unchecked(pos, to);
        slice == string
    } else {
        false
    }
}

#[inline]
unsafe fn match_insensitive(source: &str, string: &str, pos: usize) -> bool {
    let slice = source.slice_unchecked(pos, source.len());

    if slice.is_char_boundary(string.len()) {
        let slice = slice.slice_unchecked(0, string.len());
        slice.eq_ignore_ascii_case(string)
    } else {
        false
    }
}

#[inline]
unsafe fn match_range(source: &str, range: Range<char>, pos: usize) -> Option<usize> {
    let slice = source.slice_unchecked(pos, source.len());

    if let Some(char) = slice.chars().next() {
        if range.start <= char && char <= range.end {
            Some(char.len_utf8())
        } else {
            None
        }
    } else {
        None
    }
}

impl<'a> Input<'a> for StrInput<'a> {
    #[inline]
    fn len(&self) -> usize {
        self.str_ref.len()
    }

    #[inline]
    fn is_empty(&self) -> bool {
        self.str_ref.is_empty()
    }

    #[inline]
    fn file_name(&self) -> Option<OsString> {
        None
    }

    #[inline]
    unsafe fn slice(&self, start: usize, end: usize) -> &'a str {
        self.str_ref.slice_unchecked(start, end)
    }

    unsafe fn line_col(&self, pos: usize) -> (usize, usize) {
        line_col(self.str_ref, pos)
    }

    unsafe fn line_of(&self, pos: usize) -> &'a str {
        line_of(self.str_ref, pos)
    }

    #[inline]
    unsafe fn skip(&self, n: usize, pos: usize) -> Option<usize> {
        skip(self.str_ref, n, pos)
    }

    #[inline]
    unsafe fn match_string(&self, string: &str, pos: usize) -> bool {
        match_string(self.str_ref, string, pos)
    }

    #[inline]
    unsafe fn match_insensitive(&self, string: &str, pos: usize) -> bool {
        match_insensitive(self.str_ref, string, pos)
    }

    #[inline]
    unsafe fn match_range(&self, range: Range<char>, pos: usize) -> Option<usize> {
        match_range(self.str_ref, range, pos)
    }
}

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

    #[test]
    fn empty() {
        let input2 = StrInput::new("");

        unsafe {
            assert!(input2.is_empty());
            assert!(input2.match_string("", 0));
            assert!(!input2.match_string("a", 0));
        }
    }

    #[test]
    fn parts() {
        let input2 = StrInput::new("asdasdf");

        unsafe {
            assert!(!input2.is_empty());
            assert!(input2.match_string("asd", 0));
            assert!(input2.match_string("asdf", 3));
        }
    }

    #[test]
    fn len() {
        assert_eq!(StrInput::new("asdasdf").len(), 7);
    }

    #[test]
    fn slice() {
        let input2 = StrInput::new("asdasdf");

        unsafe {
            assert_eq!(input2.slice(1, 3), "sd");
        }
    }

    #[test]
    fn slice_lifetime() {
        let string: String = "asdasdf".to_owned();

        let slice: &str = {
            let input2 = StrInput::new(&string);
            unsafe { input2.slice(1, 3) }

        };

        assert_eq!(slice, "sd");
    }

    #[test]
    fn line_col() {
        let input2 = StrInput::new("a\rb\nc\r\nd嗨");

        unsafe {
            assert_eq!(input2.line_col(0), (1, 1));
            assert_eq!(input2.line_col(1), (1, 2));
            assert_eq!(input2.line_col(2), (1, 3));
            assert_eq!(input2.line_col(3), (1, 4));
            assert_eq!(input2.line_col(4), (2, 1));
            assert_eq!(input2.line_col(5), (2, 2));
            assert_eq!(input2.line_col(6), (2, 3));
            assert_eq!(input2.line_col(7), (3, 1));
            assert_eq!(input2.line_col(8), (3, 2));
            assert_eq!(input2.line_col(11), (3, 3));
        }
    }

    #[test]
    fn line_of() {
        let input2 = StrInput::new("a\rb\nc\r\nd嗨");

        unsafe {
            assert_eq!(input2.line_of(0), "a\rb");
            assert_eq!(input2.line_of(1), "a\rb");
            assert_eq!(input2.line_of(2), "a\rb");
            assert_eq!(input2.line_of(3), "a\rb");
            assert_eq!(input2.line_of(4), "c");
            assert_eq!(input2.line_of(5), "c");
            assert_eq!(input2.line_of(6), "c");
            assert_eq!(input2.line_of(7), "d嗨");
            assert_eq!(input2.line_of(8), "d嗨");
            assert_eq!(input2.line_of(11), "d嗨");
        }
    }

    #[test]
    fn line_of_empty() {
        let input2 = StrInput::new("");

        unsafe {
            assert_eq!(input2.line_of(0), "");
        }
    }

    #[test]
    fn line_of_new_line() {
        let input2 = StrInput::new("\n");

        unsafe {
            assert_eq!(input2.line_of(0), "");
        }
    }

    #[test]
    fn line_of_between_new_line() {
        let input2 = StrInput::new("\n\n");

        unsafe {
            assert_eq!(input2.line_of(1), "");
        }
    }

    #[test]
    fn skip_empty() {
        let input2 = StrInput::new("");

        unsafe {
            assert_eq!(input2.skip(0, 0), Some(0));
            assert_eq!(input2.skip(1, 0), None);
        }
    }

    #[test]
    fn skip() {
        let input2 = StrInput::new("d嗨");

        unsafe {
            assert_eq!(input2.skip(0, 0), Some(0));
            assert_eq!(input2.skip(1, 0), Some(1));
            assert_eq!(input2.skip(1, 1), Some(3));
        }
    }

    #[test]
    fn match_range() {
        let input2 = StrInput::new("b");

        unsafe {
            assert!(input2.match_range('a'..'c', 0).is_some());
            assert!(input2.match_range('b'..'b', 0).is_some());
            assert!(input2.match_range('a'..'a', 0).is_none());
            assert!(input2.match_range('c'..'c', 0).is_none());
            assert!(input2.match_range('a'..'嗨', 0).is_some());
        }
    }

    #[test]
    fn match_insensitive() {
        let input2 = StrInput::new("AsdASdF");

        unsafe {
            assert!(input2.match_insensitive("asd", 0));
            assert!(input2.match_insensitive("asdf", 3));
        }
    }
}
<|endoftext|>


<commit_before>use std::convert;
use std::ops::{Add, Sub};
use std::default::Default;
use internals::Node;

pub trait Semigroup {
    /// The sum operation of the semigroup
    fn sum(self, other: Self) -> Self;
}

/// The class of monoids
///
/// The class of monoids (types with an accumulative binary operation that has
/// an default).
///
/// An instance _M_ should satisfy the following laws:
///
///  + _x_`.sum(`_M_`::default())` = _x_
///  + _M_`::default().sum(`_x_`)` = _x_
///  + _x_`.sum(`_y_`.sum(`_z_`))` = _z_`.sum(`_x_`.sum(`_y_`))`
///  + _M_`::accumulate(`_a_`)` = _a_`.fold(`_M_`::default,`_M_`::sum)`
///
pub trait Monoid: Semigroup + Default + Sized {

    #[inline]
    fn accumulate<F>(xs: F) -> Self
    where F: Iterator<Item=Self>
        , Self: Sized {
        xs.fold(Self::default(), Self::sum)
    }
}

/// A monoid that can be applied to a `Node` as a measurement
pub trait Metric: Monoid + Eq + Add<usize, Output=Self>
                              + Sub<usize, Output=Self>
                              + Sized {
    /// Apply this `Metric` to the given `Node`
    fn measure(node: &Node) -> Self;

    /// Convert the `Metric` into a byte index into the given `Node`
    ///
    /// # Returns
    /// - `Some` with the byte index of the beginning of the `n`th  element
    ///    in `node` measured by this `Metric`, if there is an `n`th element
    /// - `None` if there is no `n`th element in `node`
    fn to_byte_index(&self, node: &Node) -> Option<usize>;

    /// Returns the byte index of the next element of this metric in `Node`
    #[inline] fn next(self, node: &Node) -> Option<usize> {
        (self + 1).to_byte_index(node)
    }

    /// Returns the byte index of the previous element of this metric in `Node`
    #[inline] fn back(self, node: &Node) -> Option<usize> {
        (self - 1).to_byte_index(node)
    }

    /// Returns true if index `i` in `node` is a boundary along this `Metric`
    // TODO: should this be a method on `Node`s instead?
    fn is_boundary(node: &Node, i: usize) -> bool;
}

#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
pub struct Grapheme(usize);
impl Default for Grapheme {
    #[inline] fn default() -> Self { Grapheme(0) }
}
impl Semigroup for Grapheme {
    #[inline] fn sum(self, other: Self) -> Self { Grapheme(self.0 + other.0) }
}
impl convert::From<usize> for Grapheme {
    #[inline] fn from(u: usize) -> Self { Grapheme(u) }
}
impl Add<usize> for Grapheme {
    type Output = Self;
    #[inline] fn add(self, rhs: usize) -> Self { Grapheme(self.0 + rhs) }
}
impl Sub<usize> for Grapheme {
    type Output = Self;
    #[inline] fn sub(self, rhs: usize) -> Self { Grapheme(self.0 - rhs) }
}

#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
pub struct Line(usize);
impl Default for Line {
    #[inline] fn default() -> Self { Line(0) }
}
impl Semigroup for Line {
    #[inline] fn sum(self, other: Self) -> Self { Line(self.0 + other.0) }
}

impl convert::From<usize> for Line {
    #[inline] fn from(u: usize) -> Self { Line(u) }
}
impl Add<usize> for Line {
    type Output = Self;
    #[inline] fn add(self, rhs: usize) -> Self { Line(self.0 + rhs) }
}
impl Sub<usize> for Line {
    type Output = Self;
    #[inline] fn sub(self, rhs: usize) -> Self { Line(self.0 - rhs) }
}
<commit_msg>docs: add some "documentation" to Metrics<commit_after>use std::convert;
use std::ops::{Add, Sub};
use std::default::Default;
use internals::Node;

pub trait Semigroup {
    /// The sum operation of the semigroup
    fn sum(self, other: Self) -> Self;
}

/// The class of monoids
///
/// The class of monoids (types with an accumulative binary operation that has
/// an default).
///
/// An instance _M_ should satisfy the following laws:
///
///  + _x_`.sum(`_M_`::default())` = _x_
///  + _M_`::default().sum(`_x_`)` = _x_
///  + _x_`.sum(`_y_`.sum(`_z_`))` = _z_`.sum(`_x_`.sum(`_y_`))`
///  + _M_`::accumulate(`_a_`)` = _a_`.fold(`_M_`::default,`_M_`::sum)`
///
pub trait Monoid: Semigroup + Default + Sized {

    #[inline]
    fn accumulate<F>(xs: F) -> Self
    where F: Iterator<Item=Self>
        , Self: Sized {
        xs.fold(Self::default(), Self::sum)
    }
}

/// A monoid that can be applied to a `Node` as a measurement
pub trait Metric: Monoid + Eq + Add<usize, Output=Self>
                              + Sub<usize, Output=Self>
                              + Sized {
    /// Apply this `Metric` to the given `Node`
    ///
    /// Although we aren't currently enforcing this, `measure`ing a `Node` with
    /// two children should produce the same result as `measure`ing both
    /// children and `Monoid::sum`ming the result. That is to say, `measure`
    /// should be a [_monoid homomorphism_]
    /// (https://en.wikipedia.org/wiki/Monoid#Monoid_homomorphisms).
    fn measure(node: &Node) -> Self;

    /// Convert the `Metric` into a byte index into the given `Node`
    ///
    /// # Returns
    /// - `Some` with the byte index of the beginning of the `n`th  element
    ///    in `node` measured by this `Metric`, if there is an `n`th element
    /// - `None` if there is no `n`th element in `node`
    fn to_byte_index(&self, node: &Node) -> Option<usize>;

    /// Returns the byte index of the next element of this metric in `Node`
    #[inline] fn next(self, node: &Node) -> Option<usize> {
        (self + 1).to_byte_index(node)
    }

    /// Returns the byte index of the previous element of this metric in `Node`
    #[inline] fn back(self, node: &Node) -> Option<usize> {
        (self - 1).to_byte_index(node)
    }

    /// Returns true if index `i` in `node` is a boundary along this `Metric`
    // TODO: should this be a method on `Node`s instead?
    fn is_boundary(node: &Node, i: usize) -> bool;
}

#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
pub struct Grapheme(usize);
impl Default for Grapheme {
    #[inline] fn default() -> Self { Grapheme(0) }
}
impl Semigroup for Grapheme {
    #[inline] fn sum(self, other: Self) -> Self { Grapheme(self.0 + other.0) }
}
impl convert::From<usize> for Grapheme {
    #[inline] fn from(u: usize) -> Self { Grapheme(u) }
}
impl Add<usize> for Grapheme {
    type Output = Self;
    #[inline] fn add(self, rhs: usize) -> Self { Grapheme(self.0 + rhs) }
}
impl Sub<usize> for Grapheme {
    type Output = Self;
    #[inline] fn sub(self, rhs: usize) -> Self { Grapheme(self.0 - rhs) }
}

#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
pub struct Line(usize);
impl Default for Line {
    #[inline] fn default() -> Self { Line(0) }
}
impl Semigroup for Line {
    #[inline] fn sum(self, other: Self) -> Self { Line(self.0 + other.0) }
}

impl convert::From<usize> for Line {
    #[inline] fn from(u: usize) -> Self { Line(u) }
}
impl Add<usize> for Line {
    type Output = Self;
    #[inline] fn add(self, rhs: usize) -> Self { Line(self.0 + rhs) }
}
impl Sub<usize> for Line {
    type Output = Self;
    #[inline] fn sub(self, rhs: usize) -> Self { Line(self.0 - rhs) }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::fs::File;
use std::io::Read;

use sdl2;

// use core_compat::entity::resource_file::ResourceFile;
use core_compat::entity::entry::Entry;
use core_compat::entity::sprite::Sprite;
use core_compat::entity::sprite_type::SpriteType::{self, Bullet, Character, Interface, Icon, Tile, Object};
use core_compat::parser::rle::parse_rle;

use crate::error::Error;
use crate::sdl::Sdl;

pub struct SpriteEntry {
    pub sprite: Sprite,
    pub texture: sdl2::render::Texture
}

pub struct SpriteManager {
    db_path: PathBuf,
    bul_map: HashMap<Entry, Rc<SpriteEntry>>,
    ico_map: HashMap<Entry, Rc<SpriteEntry>>,
    chr_map: HashMap<Entry, Rc<SpriteEntry>>,
    obj_map: HashMap<Entry, Rc<SpriteEntry>>,
    tle_map: HashMap<Entry, Rc<SpriteEntry>>,
    int_map: HashMap<Entry, Rc<SpriteEntry>>,
}

impl SpriteManager {
    pub fn new(db_path: &Path) -> SpriteManager {
        SpriteManager {
            db_path: db_path.into(),
            bul_map: HashMap::new(),
            ico_map: HashMap::new(),
            chr_map: HashMap::new(),
            obj_map: HashMap::new(),
            tle_map: HashMap::new(),
            int_map: HashMap::new(),
        }
    }

    pub fn get_sprite_entry(
        &mut self,
        req_entry: &Entry,
        sprite_type: SpriteType,
        sdl: &mut Sdl
    ) -> Result<&SpriteEntry, Error> {

        // TODO: Work around borrowing rules to only request the sprite once...
        {
            let mut need_load = false;

            if self.req_sprite(req_entry, sprite_type).is_none() {
                need_load = true;
            }

            if need_load {
                self.load_sprite(req_entry.file(), sprite_type, sdl)?;
            }
        }

        if let Some(entry) = self.req_sprite(req_entry, sprite_type) {
            Ok(entry)
        } else {
            Err(Error::SpriteLoad)
        }
    }

    fn req_sprite(
        &self,
        entry: &Entry,
        sprite_type: SpriteType
    ) -> Option<&SpriteEntry> {
        if let Some(entry) = match sprite_type {
            Bullet    => { self.bul_map.get(entry) },
            Icon      => { self.ico_map.get(entry) },
            Character => { self.chr_map.get(entry) },
            Object    => { self.obj_map.get(entry) },
            Tile      => { self.tle_map.get(entry) },
            Interface => { self.int_map.get(entry) },
        } {
            Some(entry)
        } else {
            None
        }
    }

    fn load_sprite(
        &mut self,
        number: u32,
        sprite_type: SpriteType,
        sdl: &mut Sdl
    ) -> Result<(), Error> {
        // generate correct path for the sprite
        let folder_str = match sprite_type {
            Bullet    => {"Bul"},
            Icon      => {"Ico"},
            Character => {"Chr"},
            Object    => {"Obj"},
            Tile      => {"Tle"},
            Interface => {"Int"},
        };
        let file_str = match sprite_type {
            Bullet    => format!("bul{:05}.rle", number),
            Icon      => format!("ico{:05}.rle", number),
            Character => format!("chr{:05}.rle", number),
            Object    => format!("obj{:05}.rle", number),
            Tile      => format!("tle{:05}.rle", number),
            Interface => format!("int{:05}.rle", number),
        };
        let mut path: PathBuf = self.db_path.clone();
        path.push(folder_str);
        path.push(file_str);
        // load data
        let mut file = match File::open(&path) {
            Ok(file) => file,
            Err(e) => {
                println!("Failed to load sprite file: {:?}", &path);
                return Err(Error::Io(e))
            }
        };
        let mut data = Vec::<u8>::new();
        file.read_to_end(&mut data)?;
        // parse rle file and insert into resource_manager
        let resource_file = parse_rle(number, &data)?;
        for resource in resource_file.resources {
            let entry = Entry::new(number, resource.index() );
            let sprite = Sprite {
                class: sprite_type,
                rle_entry: entry,
                x_dim: resource.width,
                y_dim: resource.height,
                x_off: resource.offset_x,
                y_off: resource.offset_y,
                image_raw: resource.image_raw,
            };

            let mut texture = sdl.texture_creator.create_texture(
                Some(sdl2::pixels::PixelFormatEnum::ABGR8888),
                sdl2::render::TextureAccess::Static,
                resource.width as u32,
                resource.height as u32)?;
            texture.set_blend_mode(sdl2::render::BlendMode::Blend);
            let pitch = resource.width as usize * 4;
            texture.update(None, &sprite.image_raw, pitch).unwrap();

            let sprite_entry = Rc::new(SpriteEntry { sprite, texture });

            match sprite_type {
                Bullet    => { self.bul_map.insert(entry, sprite_entry); },
                Icon      => { self.ico_map.insert(entry, sprite_entry); },
                Character => { self.chr_map.insert(entry, sprite_entry); },
                Object    => { self.obj_map.insert(entry, sprite_entry); },
                Tile      => { self.tle_map.insert(entry, sprite_entry); },
                Interface => { self.int_map.insert(entry, sprite_entry); },
            }
        }
        Ok(())
    }

    pub fn get_count(&self) -> usize {
          self.bul_map.len()
        + self.ico_map.len()
        + self.chr_map.len()
        + self.obj_map.len()
        + self.tle_map.len()
        + self.int_map.len()
    }
}
<commit_msg>fix missed method -> property change<commit_after>use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::fs::File;
use std::io::Read;

use sdl2;

// use core_compat::entity::resource_file::ResourceFile;
use core_compat::entity::entry::Entry;
use core_compat::entity::sprite::Sprite;
use core_compat::entity::sprite_type::SpriteType::{self, Bullet, Character, Interface, Icon, Tile, Object};
use core_compat::parser::rle::parse_rle;

use crate::error::Error;
use crate::sdl::Sdl;

pub struct SpriteEntry {
    pub sprite: Sprite,
    pub texture: sdl2::render::Texture
}

pub struct SpriteManager {
    db_path: PathBuf,
    bul_map: HashMap<Entry, Rc<SpriteEntry>>,
    ico_map: HashMap<Entry, Rc<SpriteEntry>>,
    chr_map: HashMap<Entry, Rc<SpriteEntry>>,
    obj_map: HashMap<Entry, Rc<SpriteEntry>>,
    tle_map: HashMap<Entry, Rc<SpriteEntry>>,
    int_map: HashMap<Entry, Rc<SpriteEntry>>,
}

impl SpriteManager {
    pub fn new(db_path: &Path) -> SpriteManager {
        SpriteManager {
            db_path: db_path.into(),
            bul_map: HashMap::new(),
            ico_map: HashMap::new(),
            chr_map: HashMap::new(),
            obj_map: HashMap::new(),
            tle_map: HashMap::new(),
            int_map: HashMap::new(),
        }
    }

    pub fn get_sprite_entry(
        &mut self,
        req_entry: &Entry,
        sprite_type: SpriteType,
        sdl: &mut Sdl
    ) -> Result<&SpriteEntry, Error> {

        // TODO: Work around borrowing rules to only request the sprite once...
        {
            let mut need_load = false;

            if self.req_sprite(req_entry, sprite_type).is_none() {
                need_load = true;
            }

            if need_load {
                self.load_sprite(req_entry.file(), sprite_type, sdl)?;
            }
        }

        if let Some(entry) = self.req_sprite(req_entry, sprite_type) {
            Ok(entry)
        } else {
            Err(Error::SpriteLoad)
        }
    }

    fn req_sprite(
        &self,
        entry: &Entry,
        sprite_type: SpriteType
    ) -> Option<&SpriteEntry> {
        if let Some(entry) = match sprite_type {
            Bullet    => { self.bul_map.get(entry) },
            Icon      => { self.ico_map.get(entry) },
            Character => { self.chr_map.get(entry) },
            Object    => { self.obj_map.get(entry) },
            Tile      => { self.tle_map.get(entry) },
            Interface => { self.int_map.get(entry) },
        } {
            Some(entry)
        } else {
            None
        }
    }

    fn load_sprite(
        &mut self,
        number: u32,
        sprite_type: SpriteType,
        sdl: &mut Sdl
    ) -> Result<(), Error> {
        // generate correct path for the sprite
        let folder_str = match sprite_type {
            Bullet    => {"Bul"},
            Icon      => {"Ico"},
            Character => {"Chr"},
            Object    => {"Obj"},
            Tile      => {"Tle"},
            Interface => {"Int"},
        };
        let file_str = match sprite_type {
            Bullet    => format!("bul{:05}.rle", number),
            Icon      => format!("ico{:05}.rle", number),
            Character => format!("chr{:05}.rle", number),
            Object    => format!("obj{:05}.rle", number),
            Tile      => format!("tle{:05}.rle", number),
            Interface => format!("int{:05}.rle", number),
        };
        let mut path: PathBuf = self.db_path.clone();
        path.push(folder_str);
        path.push(file_str);
        // load data
        let mut file = match File::open(&path) {
            Ok(file) => file,
            Err(e) => {
                println!("Failed to load sprite file: {:?}", &path);
                return Err(Error::Io(e))
            }
        };
        let mut data = Vec::<u8>::new();
        file.read_to_end(&mut data)?;
        // parse rle file and insert into resource_manager
        let resource_file = parse_rle(number, &data)?;
        for resource in resource_file.resources {
            let entry = Entry::new(number, resource.index);
            let sprite = Sprite {
                class: sprite_type,
                rle_entry: entry,
                x_dim: resource.width,
                y_dim: resource.height,
                x_off: resource.offset_x,
                y_off: resource.offset_y,
                image_raw: resource.image_raw,
            };

            let mut texture = sdl.texture_creator.create_texture(
                Some(sdl2::pixels::PixelFormatEnum::ABGR8888),
                sdl2::render::TextureAccess::Static,
                resource.width as u32,
                resource.height as u32)?;
            texture.set_blend_mode(sdl2::render::BlendMode::Blend);
            let pitch = resource.width as usize * 4;
            texture.update(None, &sprite.image_raw, pitch).unwrap();

            let sprite_entry = Rc::new(SpriteEntry { sprite, texture });

            match sprite_type {
                Bullet    => { self.bul_map.insert(entry, sprite_entry); },
                Icon      => { self.ico_map.insert(entry, sprite_entry); },
                Character => { self.chr_map.insert(entry, sprite_entry); },
                Object    => { self.obj_map.insert(entry, sprite_entry); },
                Tile      => { self.tle_map.insert(entry, sprite_entry); },
                Interface => { self.int_map.insert(entry, sprite_entry); },
            }
        }
        Ok(())
    }

    pub fn get_count(&self) -> usize {
          self.bul_map.len()
        + self.ico_map.len()
        + self.chr_map.len()
        + self.obj_map.len()
        + self.tle_map.len()
        + self.int_map.len()
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/// A struct that makes it easier to print out a pretty tree of data, which
/// can be visually scanned more easily.
pub struct PrintTree {
    /// The current level of recursion.
    level: u32,

    /// An item which is queued up, so that we can determine if we need
    /// a mid-tree prefix or a branch ending prefix.
    queued_item: Option<String>,
}

impl PrintTree {
    pub fn new(title: String) -> PrintTree {
        println!("\u{9484} {}", title);
        PrintTree {
            level: 1,
            queued_item: None,
        }
    }

    /// Descend one level in the tree with the given title.
    pub fn new_level(&mut self, title: String) {
        self.flush_queued_item("\u{9500}\u{9472}");

        self.print_level_prefix();
        println!("\u{9500}\u{9472} {}", title);

        self.level = self.level + 1;
    }

    /// Ascend one level in the tree.
    pub fn end_level(&mut self) {
        self.flush_queued_item("\u{9492}\u{9472}");
        self.level = self.level - 1;
    }

    /// Add an item to the current level in the tree.
    pub fn add_item(&mut self, text: String) {
        self.flush_queued_item("\u{9500}\u{9472}");
        self.queued_item = Some(text);
    }

    fn print_level_prefix(&self) {
        for _ in 0..self.level {
            print!("\u{9474}  ");
        }
    }

    fn flush_queued_item(&mut self, prefix: &str) {
        if let Some(queued_item) = self.queued_item.take() {
            self.print_level_prefix();
            println!("{} {}", prefix, queued_item);
        }
    }
}

impl Drop for PrintTree {
    fn drop(&mut self) {
        self.flush_queued_item("\u{9492}\u{9472}");
    }
}
<commit_msg>Correct the unicode codes used for tree printing<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/// A struct that makes it easier to print out a pretty tree of data, which
/// can be visually scanned more easily.
pub struct PrintTree {
    /// The current level of recursion.
    level: u32,

    /// An item which is queued up, so that we can determine if we need
    /// a mid-tree prefix or a branch ending prefix.
    queued_item: Option<String>,
}

impl PrintTree {
    pub fn new(title: String) -> PrintTree {
        println!("\u{250c} {}", title);
        PrintTree {
            level: 1,
            queued_item: None,
        }
    }

    /// Descend one level in the tree with the given title.
    pub fn new_level(&mut self, title: String) {
        self.flush_queued_item("\u{251C}\u{2500}");

        self.print_level_prefix();
        println!("\u{251C}\u{2500} {}", title);

        self.level = self.level + 1;
    }

    /// Ascend one level in the tree.
    pub fn end_level(&mut self) {
        self.flush_queued_item("\u{2514}\u{2500}");
        self.level = self.level - 1;
    }

    /// Add an item to the current level in the tree.
    pub fn add_item(&mut self, text: String) {
        self.flush_queued_item("\u{251C}\u{2500}");
        self.queued_item = Some(text);
    }

    fn print_level_prefix(&self) {
        for _ in 0..self.level {
            print!("\u{2502}  ");
        }
    }

    fn flush_queued_item(&mut self, prefix: &str) {
        if let Some(queued_item) = self.queued_item.take() {
            self.print_level_prefix();
            println!("{} {}", prefix, queued_item);
        }
    }
}

impl Drop for PrintTree {
    fn drop(&mut self) {
        self.flush_queued_item("\u{9492}\u{9472}");
    }
}
<|endoftext|>


<commit_before>//! This pass adds validation calls (AcquireValid, ReleaseValid) where appropriate.
//! It has to be run really early, before transformations like inlining, because
//! introducing these calls *adds* UB -- so, conceptually, this pass is actually part
//! of MIR building, and only after this pass we think of the program has having the
//! normal MIR semantics.

use crate::transform::MirPass;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, Ty, TyCtxt};

pub struct AddRetag;

/// Determines whether this place is "stable": Whether, if we evaluate it again
/// after the assignment, we can be sure to obtain the same place value.
/// (Concurrent accesses by other threads are no problem as these are anyway non-atomic
/// copies.  Data races are UB.)
fn is_stable(place: PlaceRef<'_>) -> bool {
    place.projection.iter().all(|elem| {
        match elem {
            // Which place this evaluates to can change with any memory write,
            // so cannot assume this to be stable.
            ProjectionElem::Deref => false,
            // Array indices are interesting, but MIR building generates a *fresh*
            // temporary for every array access, so the index cannot be changed as
            // a side-effect.
            ProjectionElem::Index { .. } |
            // The rest is completely boring, they just offset by a constant.
            ProjectionElem::Field { .. } |
            ProjectionElem::ConstantIndex { .. } |
            ProjectionElem::Subslice { .. } |
            ProjectionElem::Downcast { .. } => true,
        }
    })
}

/// Determine whether this type may be a reference (or box), and thus needs retagging.
fn may_be_reference(ty: Ty<'tcx>) -> bool {
    match ty.kind() {
        // Primitive types that are not references
        ty::Bool
        | ty::Char
        | ty::Float(_)
        | ty::Int(_)
        | ty::Uint(_)
        | ty::RawPtr(..)
        | ty::FnPtr(..)
        | ty::Str
        | ty::FnDef(..)
        | ty::Never => false,
        // References
        ty::Ref(..) => true,
        ty::Adt(..) if ty.is_box() => true,
        // Compound types are not references
        ty::Array(..) | ty::Slice(..) | ty::Tuple(..) | ty::Adt(..) => false,
        // Conservative fallback
        _ => true,
    }
}

impl<'tcx> MirPass<'tcx> for AddRetag {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        if !tcx.sess.opts.debugging_opts.mir_emit_retag {
            return;
        }

        // We need an `AllCallEdges` pass before we can do any work.
        super::add_call_guards::AllCallEdges.run_pass(tcx, body);

        let (span, arg_count) = (body.span, body.arg_count);
        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
        let needs_retag = |place: &Place<'tcx>| {
            // FIXME: Instead of giving up for unstable places, we should introduce
            // a temporary and retag on that.
            is_stable(place.as_ref()) && may_be_reference(place.ty(&*local_decls, tcx).ty)
        };

        // PART 1
        // Retag arguments at the beginning of the start block.
        {
            // FIXME: Consider using just the span covering the function
            // argument declaration.
            let source_info = SourceInfo::outermost(span);
            // Gather all arguments, skip return value.
            let places = local_decls
                .iter_enumerated()
                .skip(1)
                .take(arg_count)
                .map(|(local, _)| Place::from(local))
                .filter(needs_retag);
            // Emit their retags.
            basic_blocks[START_BLOCK].statements.splice(
                0..0,
                places.map(|place| Statement {
                    source_info,
                    kind: StatementKind::Retag(RetagKind::FnEntry, box (place)),
                }),
            );
        }

        // PART 2
        // Retag return values of functions.  Also escape-to-raw the argument of `drop`.
        // We collect the return destinations because we cannot mutate while iterating.
        let returns = basic_blocks
            .iter_mut()
            .filter_map(|block_data| {
                match block_data.terminator().kind {
                    TerminatorKind::Call { destination: Some(ref destination), .. }
                        if needs_retag(&destination.0) =>
                    {
                        // Remember the return destination for later
                        Some((block_data.terminator().source_info, destination.0, destination.1))
                    }

                    // `Drop` is also a call, but it doesn't return anything so we are good.
                    TerminatorKind::Drop { .. } | TerminatorKind::DropAndReplace { .. } => None,
                    // Not a block ending in a Call -> ignore.
                    _ => None,
                }
            })
            .collect::<Vec<_>>();
        // Now we go over the returns we collected to retag the return values.
        for (source_info, dest_place, dest_block) in returns {
            basic_blocks[dest_block].statements.insert(
                0,
                Statement {
                    source_info,
                    kind: StatementKind::Retag(RetagKind::Default, box (dest_place)),
                },
            );
        }

        // PART 3
        // Add retag after assignment.
        for block_data in basic_blocks {
            // We want to insert statements as we iterate.  To this end, we
            // iterate backwards using indices.
            for i in (0..block_data.statements.len()).rev() {
                let (retag_kind, place) = match block_data.statements[i].kind {
                    // Retag-as-raw after escaping to a raw pointer.
                    StatementKind::Assign(box (place, Rvalue::AddressOf(..))) => {
                        (RetagKind::Raw, place)
                    }
                    // Assignments of reference or ptr type are the ones where we may have
                    // to update tags.  This includes `x = &[mut] ...` and hence
                    // we also retag after taking a reference!
                    StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => {
                        let kind = match rvalue {
                            Rvalue::Ref(_, borrow_kind, _)
                                if borrow_kind.allows_two_phase_borrow() =>
                            {
                                RetagKind::TwoPhase
                            }
                            _ => RetagKind::Default,
                        };
                        (kind, *place)
                    }
                    // Do nothing for the rest
                    _ => continue,
                };
                // Insert a retag after the statement.
                let source_info = block_data.statements[i].source_info;
                block_data.statements.insert(
                    i + 1,
                    Statement { source_info, kind: StatementKind::Retag(retag_kind, box (place)) },
                );
            }
        }
    }
}
<commit_msg>Auto merge of #78597 - RalfJung:raw-retag, r=oli-obk<commit_after>//! This pass adds validation calls (AcquireValid, ReleaseValid) where appropriate.
//! It has to be run really early, before transformations like inlining, because
//! introducing these calls *adds* UB -- so, conceptually, this pass is actually part
//! of MIR building, and only after this pass we think of the program has having the
//! normal MIR semantics.

use crate::transform::MirPass;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, Ty, TyCtxt};

pub struct AddRetag;

/// Determines whether this place is "stable": Whether, if we evaluate it again
/// after the assignment, we can be sure to obtain the same place value.
/// (Concurrent accesses by other threads are no problem as these are anyway non-atomic
/// copies.  Data races are UB.)
fn is_stable(place: PlaceRef<'_>) -> bool {
    place.projection.iter().all(|elem| {
        match elem {
            // Which place this evaluates to can change with any memory write,
            // so cannot assume this to be stable.
            ProjectionElem::Deref => false,
            // Array indices are interesting, but MIR building generates a *fresh*
            // temporary for every array access, so the index cannot be changed as
            // a side-effect.
            ProjectionElem::Index { .. } |
            // The rest is completely boring, they just offset by a constant.
            ProjectionElem::Field { .. } |
            ProjectionElem::ConstantIndex { .. } |
            ProjectionElem::Subslice { .. } |
            ProjectionElem::Downcast { .. } => true,
        }
    })
}

/// Determine whether this type may be a reference (or box), and thus needs retagging.
fn may_be_reference(ty: Ty<'tcx>) -> bool {
    match ty.kind() {
        // Primitive types that are not references
        ty::Bool
        | ty::Char
        | ty::Float(_)
        | ty::Int(_)
        | ty::Uint(_)
        | ty::RawPtr(..)
        | ty::FnPtr(..)
        | ty::Str
        | ty::FnDef(..)
        | ty::Never => false,
        // References
        ty::Ref(..) => true,
        ty::Adt(..) if ty.is_box() => true,
        // Compound types are not references
        ty::Array(..) | ty::Slice(..) | ty::Tuple(..) | ty::Adt(..) => false,
        // Conservative fallback
        _ => true,
    }
}

impl<'tcx> MirPass<'tcx> for AddRetag {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        if !tcx.sess.opts.debugging_opts.mir_emit_retag {
            return;
        }

        // We need an `AllCallEdges` pass before we can do any work.
        super::add_call_guards::AllCallEdges.run_pass(tcx, body);

        let (span, arg_count) = (body.span, body.arg_count);
        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
        let needs_retag = |place: &Place<'tcx>| {
            // FIXME: Instead of giving up for unstable places, we should introduce
            // a temporary and retag on that.
            is_stable(place.as_ref()) && may_be_reference(place.ty(&*local_decls, tcx).ty)
        };
        let place_base_raw = |place: &Place<'tcx>| {
            // If this is a `Deref`, get the type of what we are deref'ing.
            let deref_base =
                place.projection.iter().rposition(|p| matches!(p, ProjectionElem::Deref));
            if let Some(deref_base) = deref_base {
                let base_proj = &place.projection[..deref_base];
                let ty = Place::ty_from(place.local, base_proj, &*local_decls, tcx).ty;
                ty.is_unsafe_ptr()
            } else {
                // Not a deref, and thus not raw.
                false
            }
        };

        // PART 1
        // Retag arguments at the beginning of the start block.
        {
            // FIXME: Consider using just the span covering the function
            // argument declaration.
            let source_info = SourceInfo::outermost(span);
            // Gather all arguments, skip return value.
            let places = local_decls
                .iter_enumerated()
                .skip(1)
                .take(arg_count)
                .map(|(local, _)| Place::from(local))
                .filter(needs_retag);
            // Emit their retags.
            basic_blocks[START_BLOCK].statements.splice(
                0..0,
                places.map(|place| Statement {
                    source_info,
                    kind: StatementKind::Retag(RetagKind::FnEntry, box (place)),
                }),
            );
        }

        // PART 2
        // Retag return values of functions.  Also escape-to-raw the argument of `drop`.
        // We collect the return destinations because we cannot mutate while iterating.
        let returns = basic_blocks
            .iter_mut()
            .filter_map(|block_data| {
                match block_data.terminator().kind {
                    TerminatorKind::Call { destination: Some(ref destination), .. }
                        if needs_retag(&destination.0) =>
                    {
                        // Remember the return destination for later
                        Some((block_data.terminator().source_info, destination.0, destination.1))
                    }

                    // `Drop` is also a call, but it doesn't return anything so we are good.
                    TerminatorKind::Drop { .. } | TerminatorKind::DropAndReplace { .. } => None,
                    // Not a block ending in a Call -> ignore.
                    _ => None,
                }
            })
            .collect::<Vec<_>>();
        // Now we go over the returns we collected to retag the return values.
        for (source_info, dest_place, dest_block) in returns {
            basic_blocks[dest_block].statements.insert(
                0,
                Statement {
                    source_info,
                    kind: StatementKind::Retag(RetagKind::Default, box (dest_place)),
                },
            );
        }

        // PART 3
        // Add retag after assignment.
        for block_data in basic_blocks {
            // We want to insert statements as we iterate.  To this end, we
            // iterate backwards using indices.
            for i in (0..block_data.statements.len()).rev() {
                let (retag_kind, place) = match block_data.statements[i].kind {
                    // Retag-as-raw after escaping to a raw pointer, if the referent
                    // is not already a raw pointer.
                    StatementKind::Assign(box (lplace, Rvalue::AddressOf(_, ref rplace)))
                        if !place_base_raw(rplace) =>
                    {
                        (RetagKind::Raw, lplace)
                    }
                    // Retag after assignments of reference type.
                    StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => {
                        let kind = match rvalue {
                            Rvalue::Ref(_, borrow_kind, _)
                                if borrow_kind.allows_two_phase_borrow() =>
                            {
                                RetagKind::TwoPhase
                            }
                            _ => RetagKind::Default,
                        };
                        (kind, *place)
                    }
                    // Do nothing for the rest
                    _ => continue,
                };
                // Insert a retag after the statement.
                let source_info = block_data.statements[i].source_info;
                block_data.statements.insert(
                    i + 1,
                    Statement { source_info, kind: StatementKind::Retag(retag_kind, box (place)) },
                );
            }
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::hashmap::HashMap;
use std::os;
use rust_url;
use rust_url::{Url, UrlParser};

/**
Create a URL object from a string. Does various helpful browsery things like

* If there's no current url and the path looks like a file then it will
  create a file url based of the current working directory
* If there's a current url and the new path is relative then the new url
  is based off the current url

*/
// TODO: about:failure->
pub fn try_parse_url(str_url: &str, base_url: Option<Url>) -> Result<Url, &'static str> {
    let mut parser = UrlParser::new();
    let parser = match base_url {
        Some(ref base) => &*parser.base_url(base),
        None => &parser,
    };
    let str_url = match parser.parse(str_url) {
        Err(err) => {
            if base_url.is_none() {
                // Assume we've been given a file path. If it's absolute just return
                // it, otherwise make it absolute with the cwd.
                if str_url.as_slice().starts_with("/") {
                    format!("file://{}", str_url)
                } else {
                    let mut path = os::getcwd();
                    path.push(str_url);
                    // FIXME (#1094): not the right way to transform a path
                    format!("file://{}", path.display().to_str())
                }
            } else {
                return Err(err)
            }
        },
        Ok(url) => {
            match (url.scheme.as_slice(), url.scheme_data.clone()) {
                ("data", _) => {
                    // Drop whitespace within data: URLs, e.g. newlines within a base64
                    // src="..." block.  Whitespace intended as content should be
                    // %-encoded or base64'd.
                    str_url.as_slice().chars().filter(|&c| !c.is_whitespace()).collect()
                },
                _ => return Ok(url)
            }
        }
    };
    parser.parse(str_url.as_slice())
}

pub fn parse_url(str_url: &str, base_url: Option<rust_url::Url>) -> rust_url::Url {
    // FIXME: Need to handle errors
    try_parse_url(str_url, base_url).ok().expect("URL parsing failed")
}

pub type UrlMap<T> = HashMap<rust_url::Url, T>;

pub fn url_map<T: Clone + 'static>() -> UrlMap<T> {
    HashMap::new()
}

pub fn is_image_data(uri: &str) -> bool {
    static types: &'static [&'static str] = &["data:image/png", "data:image/gif", "data:image/jpeg"];
    types.iter().any(|&type_| uri.starts_with(type_))
}

<commit_msg>Don’t drop whitespace in data: URLs. FromBase64 already skips newlines.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::hashmap::HashMap;
use std::os;
use rust_url;
use rust_url::{Url, UrlParser};

/**
Create a URL object from a string. Does various helpful browsery things like

* If there's no current url and the path looks like a file then it will
  create a file url based of the current working directory
* If there's a current url and the new path is relative then the new url
  is based off the current url

*/
// TODO: about:failure->
pub fn try_parse_url(str_url: &str, base_url: Option<Url>) -> Result<Url, &'static str> {
    let mut parser = UrlParser::new();
    let parser = match base_url {
        Some(ref base) => &*parser.base_url(base),
        None => &parser,
    };
    let str_url = match parser.parse(str_url) {
        Err(err) => {
            if base_url.is_none() {
                // Assume we've been given a file path. If it's absolute just return
                // it, otherwise make it absolute with the cwd.
                if str_url.as_slice().starts_with("/") {
                    format!("file://{}", str_url)
                } else {
                    let mut path = os::getcwd();
                    path.push(str_url);
                    // FIXME (#1094): not the right way to transform a path
                    format!("file://{}", path.display().to_str())
                }
            } else {
                return Err(err)
            }
        },
        Ok(url) => return Ok(url)
    };
    parser.parse(str_url.as_slice())
}

pub fn parse_url(str_url: &str, base_url: Option<rust_url::Url>) -> rust_url::Url {
    // FIXME: Need to handle errors
    try_parse_url(str_url, base_url).ok().expect("URL parsing failed")
}

pub type UrlMap<T> = HashMap<rust_url::Url, T>;

pub fn url_map<T: Clone + 'static>() -> UrlMap<T> {
    HashMap::new()
}

pub fn is_image_data(uri: &str) -> bool {
    static types: &'static [&'static str] = &["data:image/png", "data:image/gif", "data:image/jpeg"];
    types.iter().any(|&type_| uri.starts_with(type_))
}

<|endoftext|>


<commit_before><commit_msg>Avoid string copying by end-trimming in place<commit_after><|endoftext|>


<commit_before>use std::ffi::{CStr, CString};
use std::marker::PhantomData;

use tensorflow_sys as tf;

use crate::eager::Context;
use crate::{AnyTensor, DataType, Result, Status, Tensor, TensorType};

/// A handle to a tensor on a device.
///
/// Constructing a TensorHandle requires a reference to an execute context so that the
/// generated handle will not out live the context.
/// ```
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts).unwrap();
///
/// let t = Tensor::from(&[3i32]);
/// let h = TensorHandle::new(&ctx, &t).unwrap();
/// let v: Tensor<i32> = h.resolve().unwrap();
/// assert_eq!(&v[..], &[3i32]);
/// ```
///
/// Since TensorHandle cannot be alive beyond the lifetime of the context, the following
/// code will not compile.
/// ```compile_fail
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let h = {
///     let opts = ContextOptions::new();
///     let ctx = Context::new(opts).unwrap();
///
///     let t = Tensor::from(&[3i32]);
///     let h = TensorHandle::new(&ctx, &t).unwrap();
///     h
/// };
/// ```
#[derive(Debug)]
pub struct TensorHandle<'a> {
    inner: *mut tf::TFE_TensorHandle,
    // TensorHandle should not live longer than a given context.
    ctx: PhantomData<&'a Context>,
}

impl<'a> Drop for TensorHandle<'a> {
    fn drop(&mut self) {
        unsafe {
            tf::TFE_DeleteTensorHandle(self.inner);
        }
    }
}

impl<'a> TensorHandle<'a> {
    /// Crate a TensorHandle from the input Tensor
    pub fn new<T: TensorType>(_ctx: &'a Context, t: &Tensor<T>) -> Result<TensorHandle<'a>> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_NewTensorHandle(t.inner()?, status.inner) };

        if inner.is_null() {
            Err(status)
        } else {
            Ok(TensorHandle {
                inner,
                ctx: PhantomData,
            })
        }
    }

    /// Return the DataType that corresponds to this type.
    pub fn data_type(&self) -> DataType {
        unsafe { DataType::from_c(tf::TFE_TensorHandleDataType(self.inner)) }
    }

    /// Return the number of dimensions.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn num_dims(&self) -> Result<usize> {
        let status = Status::new();
        let num_dims = unsafe { tf::TFE_TensorHandleNumDims(self.inner, status.inner) };
        if status.is_ok() {
            // num_dims >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_dims as usize)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements
    pub fn num_elements(&self) -> Result<u64> {
        let status = Status::new();
        let num_elements = unsafe { tf::TFE_TensorHandleNumElements(self.inner, status.inner) };
        if status.is_ok() {
            // num_elements >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_elements as u64)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements for a given dim_index.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn dim(&self, dim_index: i32) -> Result<u64> {
        let status = Status::new();
        let dim = unsafe { tf::TFE_TensorHandleDim(self.inner, dim_index, status.inner) };
        if status.is_ok() {
            // dim >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(dim as u64)
        } else {
            Err(status)
        }
    }

    /// Return the device of the operation that produced the current TensorHandle.
    ///
    /// If the TensorHandle was produced by a copy, returns the destination device of the copy.
    /// Note that the returned device name is not always the device holding the tensor handle's memory.
    /// If you want the latter, use backing_device_name.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Returns the name of the device in whose memory underlying the current TensorHandle resides.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn backing_device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleBackingDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Return a new TensorHandle that shares the underlying tensor with the current TensorHandle.
    pub fn copy_sharing_tensor(&self) -> Result<Self> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_TensorHandleCopySharingTensor(self.inner, status.inner) };
        if status.is_ok() {
            Ok(Self {
                inner,
                ctx: self.ctx,
            })
        } else {
            Err(status)
        }
    }

    /// This function will block till the operation that produces the current TensorHandle has completed.
    /// The memory returned might alias the internal memory used by TensorFlow.
    /// Hence, callers should not mutate this memory.
    pub fn resolve<T: TensorType>(&self) -> Result<Tensor<T>> {
        let mut status = Status::new();
        let tf_tensor = unsafe { tf::TFE_TensorHandleResolve(self.inner, status.inner) };
        if !status.is_ok() {
            return Err(status);
        }
        if self.data_type() != T::data_type() {
            let msg = format!(
                "The expected data type ({}) and underlying data type ({}) did not match.",
                T::data_type(),
                self.data_type()
            );

            status.set_lossy(crate::Code::InvalidArgument, &msg);
            return Err(status);
        }

        // Safely unwrap since data_type was checked beforehand.
        unsafe { Ok(Tensor::from_tf_tensor(tf_tensor).unwrap()) }
    }

    /// Create a new TensorHandle with the same contents as the current TensorHandle but placed
    /// in the memory of the device name 'device_name'.
    /// If source and destination are the same device, then this creates a new handle
    /// that shares the underlying buffer. Otherwise, it currently requires at least
    /// one of the source or destination devices to be CPU (i.e., for the source or
    /// destination tensor to be placed in host memory).
    /// If async execution is enabled, the copy may be enqueued and the call will
    /// return "non-ready" TensorHandle. Else, this function returns after the copy has
    /// been done.
    pub fn copy_to_device<'b>(
        &self,
        ctx: &'b Context,
        device_name: &str,
    ) -> Result<TensorHandle<'b>> {
        let status = Status::new();

        let device_name = CString::new(device_name)?;
        unsafe {
            let inner = tf::TFE_TensorHandleCopyToDevice(
                self.inner,
                ctx.inner,
                device_name.as_ptr(),
                status.inner,
            );

            if status.is_ok() {
                Ok(TensorHandle {
                    inner,
                    ctx: PhantomData,
                })
            } else {
                Err(status)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eager::ContextOptions;

    #[test]
    fn test_tensor_handle() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();

        assert_eq!(h.data_type(), DataType::Int32);
        assert_eq!(h.num_elements().unwrap(), 6);
        assert_eq!(h.num_dims().unwrap(), 2);
        assert_eq!(h.dim(0).unwrap(), 2);
        assert_eq!(h.dim(1).unwrap(), 3);
    }

    #[test]
    fn test_copy_sharing_tensor() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_copy = h.copy_sharing_tensor().unwrap();
        let t2 = h_copy.resolve::<i32>().unwrap();

        // t and t2 may share the same memory, but it's difficuly to check
        // since the `resolve` does not guarantee that.
        assert_eq!(&t[..], &t2[..]);
    }

    #[cfg(feature = "tensorflow_gpu")]
    #[test]
    fn test_copy_to_device() {
        let values = [0_i32, 1, 2, 3];
        let target_device = "/job:localhost/replica:0/task:0/device:GPU:0";

        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();
        let devices = ctx.device_list().unwrap();
        assert!(
            devices.iter().any(|d| d.device_type == "GPU"),
            "Skip test_copy_to_device because no GPU device was found"
        );

        let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_gpu = TensorHandle::copy_to_device(&h, &ctx, target_device).unwrap();
        assert_eq!(h_gpu.device_name().unwrap(), target_device);
        let t2 = h_gpu.resolve::<i32>().unwrap();

        assert_eq!(&t[..], &t2[..]);
    }

    #[cfg(feature = "tensorflow_gpu")]
    #[test]
    fn test_copy_to_device_lifetime() {
        let values = [0_i32, 1, 2, 3];
        let target_device = "/job:localhost/replica:0/task:0/device:GPU:0";

        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();
        let devices = ctx.device_list().unwrap();
        assert!(
            devices.iter().any(|d| d.device_type == "GPU"),
            "Skip test_copy_to_device because no GPU device was found"
        );

        let h_gpu = {
            // Create a temporal Context
            let opts = ContextOptions::new();
            let ctx2 = Context::new(opts).unwrap();
            let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();

            // Create a TensorHandle managed by the context `ctx2`.
            let h = TensorHandle::new(&ctx2, &t).unwrap();

            // Copy to GPU. This creates a new handle managed by the context `ctx`.
            h.copy_to_device(&ctx, target_device).unwrap()
        };
        assert_eq!(h_gpu.device_name().unwrap(), target_device);
        let t2 = h_gpu.resolve::<i32>().unwrap();

        assert_eq!(&values[..], &t2[..]);
    }
}
<commit_msg>Fix typo<commit_after>use std::ffi::{CStr, CString};
use std::marker::PhantomData;

use tensorflow_sys as tf;

use crate::eager::Context;
use crate::{AnyTensor, DataType, Result, Status, Tensor, TensorType};

/// A handle to a tensor on a device.
///
/// Constructing a TensorHandle requires a reference to an execute context so that the
/// generated handle will not out live the context.
/// ```
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let opts = ContextOptions::new();
/// let ctx = Context::new(opts).unwrap();
///
/// let t = Tensor::from(&[3i32]);
/// let h = TensorHandle::new(&ctx, &t).unwrap();
/// let v: Tensor<i32> = h.resolve().unwrap();
/// assert_eq!(&v[..], &[3i32]);
/// ```
///
/// Since TensorHandle cannot be alive beyond the lifetime of the context, the following
/// code will not compile.
/// ```compile_fail
/// # use tensorflow::Tensor;
/// use tensorflow::eager::{ContextOptions, Context, TensorHandle};
///
/// let h = {
///     let opts = ContextOptions::new();
///     let ctx = Context::new(opts).unwrap();
///
///     let t = Tensor::from(&[3i32]);
///     let h = TensorHandle::new(&ctx, &t).unwrap();
///     h
/// };
/// ```
#[derive(Debug)]
pub struct TensorHandle<'a> {
    inner: *mut tf::TFE_TensorHandle,
    // TensorHandle should not live longer than a given context.
    ctx: PhantomData<&'a Context>,
}

impl<'a> Drop for TensorHandle<'a> {
    fn drop(&mut self) {
        unsafe {
            tf::TFE_DeleteTensorHandle(self.inner);
        }
    }
}

impl<'a> TensorHandle<'a> {
    /// Create a TensorHandle from the input Tensor
    pub fn new<T: TensorType>(_ctx: &'a Context, t: &Tensor<T>) -> Result<TensorHandle<'a>> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_NewTensorHandle(t.inner()?, status.inner) };

        if inner.is_null() {
            Err(status)
        } else {
            Ok(TensorHandle {
                inner,
                ctx: PhantomData,
            })
        }
    }

    /// Return the DataType that corresponds to this type.
    pub fn data_type(&self) -> DataType {
        unsafe { DataType::from_c(tf::TFE_TensorHandleDataType(self.inner)) }
    }

    /// Return the number of dimensions.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn num_dims(&self) -> Result<usize> {
        let status = Status::new();
        let num_dims = unsafe { tf::TFE_TensorHandleNumDims(self.inner, status.inner) };
        if status.is_ok() {
            // num_dims >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_dims as usize)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements
    pub fn num_elements(&self) -> Result<u64> {
        let status = Status::new();
        let num_elements = unsafe { tf::TFE_TensorHandleNumElements(self.inner, status.inner) };
        if status.is_ok() {
            // num_elements >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(num_elements as u64)
        } else {
            Err(status)
        }
    }

    /// Return the number of elements for a given dim_index.
    ///
    /// This function will block till the operation that produces the TensorHandle has completed.
    pub fn dim(&self, dim_index: i32) -> Result<u64> {
        let status = Status::new();
        let dim = unsafe { tf::TFE_TensorHandleDim(self.inner, dim_index, status.inner) };
        if status.is_ok() {
            // dim >= 0 when the status is ok, so we can safely cast it to u64.
            Ok(dim as u64)
        } else {
            Err(status)
        }
    }

    /// Return the device of the operation that produced the current TensorHandle.
    ///
    /// If the TensorHandle was produced by a copy, returns the destination device of the copy.
    /// Note that the returned device name is not always the device holding the tensor handle's memory.
    /// If you want the latter, use backing_device_name.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Returns the name of the device in whose memory underlying the current TensorHandle resides.
    ///
    /// This function will block till the operation that produces the current TensorHandle has completed.
    pub fn backing_device_name(&self) -> Result<String> {
        let status = Status::new();
        unsafe {
            let device_name = tf::TFE_TensorHandleBackingDeviceName(self.inner, status.inner);
            if status.is_ok() {
                Ok(CStr::from_ptr(device_name).to_str()?.to_string())
            } else {
                Err(status)
            }
        }
    }

    /// Return a new TensorHandle that shares the underlying tensor with the current TensorHandle.
    pub fn copy_sharing_tensor(&self) -> Result<Self> {
        let status = Status::new();
        let inner = unsafe { tf::TFE_TensorHandleCopySharingTensor(self.inner, status.inner) };
        if status.is_ok() {
            Ok(Self {
                inner,
                ctx: self.ctx,
            })
        } else {
            Err(status)
        }
    }

    /// This function will block till the operation that produces the current TensorHandle has completed.
    /// The memory returned might alias the internal memory used by TensorFlow.
    /// Hence, callers should not mutate this memory.
    pub fn resolve<T: TensorType>(&self) -> Result<Tensor<T>> {
        let mut status = Status::new();
        let tf_tensor = unsafe { tf::TFE_TensorHandleResolve(self.inner, status.inner) };
        if !status.is_ok() {
            return Err(status);
        }
        if self.data_type() != T::data_type() {
            let msg = format!(
                "The expected data type ({}) and underlying data type ({}) did not match.",
                T::data_type(),
                self.data_type()
            );

            status.set_lossy(crate::Code::InvalidArgument, &msg);
            return Err(status);
        }

        // Safely unwrap since data_type was checked beforehand.
        unsafe { Ok(Tensor::from_tf_tensor(tf_tensor).unwrap()) }
    }

    /// Create a new TensorHandle with the same contents as the current TensorHandle but placed
    /// in the memory of the device name 'device_name'.
    /// If source and destination are the same device, then this creates a new handle
    /// that shares the underlying buffer. Otherwise, it currently requires at least
    /// one of the source or destination devices to be CPU (i.e., for the source or
    /// destination tensor to be placed in host memory).
    /// If async execution is enabled, the copy may be enqueued and the call will
    /// return "non-ready" TensorHandle. Else, this function returns after the copy has
    /// been done.
    pub fn copy_to_device<'b>(
        &self,
        ctx: &'b Context,
        device_name: &str,
    ) -> Result<TensorHandle<'b>> {
        let status = Status::new();

        let device_name = CString::new(device_name)?;
        unsafe {
            let inner = tf::TFE_TensorHandleCopyToDevice(
                self.inner,
                ctx.inner,
                device_name.as_ptr(),
                status.inner,
            );

            if status.is_ok() {
                Ok(TensorHandle {
                    inner,
                    ctx: PhantomData,
                })
            } else {
                Err(status)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eager::ContextOptions;

    #[test]
    fn test_tensor_handle() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();

        assert_eq!(h.data_type(), DataType::Int32);
        assert_eq!(h.num_elements().unwrap(), 6);
        assert_eq!(h.num_dims().unwrap(), 2);
        assert_eq!(h.dim(0).unwrap(), 2);
        assert_eq!(h.dim(1).unwrap(), 3);
    }

    #[test]
    fn test_copy_sharing_tensor() {
        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();

        let t = Tensor::new(&[2, 3])
            .with_values(&[0_i32, 1, 2, 3, 4, 5])
            .unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_copy = h.copy_sharing_tensor().unwrap();
        let t2 = h_copy.resolve::<i32>().unwrap();

        // t and t2 may share the same memory, but it's difficuly to check
        // since the `resolve` does not guarantee that.
        assert_eq!(&t[..], &t2[..]);
    }

    #[cfg(feature = "tensorflow_gpu")]
    #[test]
    fn test_copy_to_device() {
        let values = [0_i32, 1, 2, 3];
        let target_device = "/job:localhost/replica:0/task:0/device:GPU:0";

        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();
        let devices = ctx.device_list().unwrap();
        assert!(
            devices.iter().any(|d| d.device_type == "GPU"),
            "Skip test_copy_to_device because no GPU device was found"
        );

        let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();
        let h = TensorHandle::new(&ctx, &t).unwrap();
        let h_gpu = TensorHandle::copy_to_device(&h, &ctx, target_device).unwrap();
        assert_eq!(h_gpu.device_name().unwrap(), target_device);
        let t2 = h_gpu.resolve::<i32>().unwrap();

        assert_eq!(&t[..], &t2[..]);
    }

    #[cfg(feature = "tensorflow_gpu")]
    #[test]
    fn test_copy_to_device_lifetime() {
        let values = [0_i32, 1, 2, 3];
        let target_device = "/job:localhost/replica:0/task:0/device:GPU:0";

        let opts = ContextOptions::new();
        let ctx = Context::new(opts).unwrap();
        let devices = ctx.device_list().unwrap();
        assert!(
            devices.iter().any(|d| d.device_type == "GPU"),
            "Skip test_copy_to_device because no GPU device was found"
        );

        let h_gpu = {
            // Create a temporal Context
            let opts = ContextOptions::new();
            let ctx2 = Context::new(opts).unwrap();
            let t = Tensor::new(&[2, 2]).with_values(&values).unwrap();

            // Create a TensorHandle managed by the context `ctx2`.
            let h = TensorHandle::new(&ctx2, &t).unwrap();

            // Copy to GPU. This creates a new handle managed by the context `ctx`.
            h.copy_to_device(&ctx, target_device).unwrap()
        };
        assert_eq!(h_gpu.device_name().unwrap(), target_device);
        let t2 = h_gpu.resolve::<i32>().unwrap();

        assert_eq!(&values[..], &t2[..]);
    }
}
<|endoftext|>