<commit_before>//! FIXME: write short doc here

use std::{
    fmt,
    ops::{self, RangeInclusive},
};

use itertools::Itertools;
use ra_text_edit::TextEditBuilder;
use rustc_hash::FxHashMap;

use crate::{
    AstNode, Direction, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxNodePtr,
    SyntaxToken, TextRange, TextSize,
};

/// Returns ancestors of the node at the offset, sorted by length. This should
/// do the right thing at an edge, e.g. when searching for expressions at `{
/// <|>foo }` we will get the name reference instead of the whole block, which
/// we would get if we just did `find_token_at_offset(...).flat_map(|t|
/// t.parent().ancestors())`.
pub fn ancestors_at_offset(
    node: &SyntaxNode,
    offset: TextSize,
) -> impl Iterator<Item = SyntaxNode> {
    node.token_at_offset(offset)
        .map(|token| token.parent().ancestors())
        .kmerge_by(|node1, node2| node1.text_range().len() < node2.text_range().len())
}

/// Finds a node of specific Ast type at offset. Note that this is slightly
/// imprecise: if the cursor is strictly between two nodes of the desired type,
/// as in
///
/// ```no-run
/// struct Foo {}|struct Bar;
/// ```
///
/// then the shorter node will be silently preferred.
pub fn find_node_at_offset<N: AstNode>(syntax: &SyntaxNode, offset: TextSize) -> Option<N> {
    ancestors_at_offset(syntax, offset).find_map(N::cast)
}

/// Skip to next non `trivia` token
pub fn skip_trivia_token(mut token: SyntaxToken, direction: Direction) -> Option<SyntaxToken> {
    while token.kind().is_trivia() {
        token = match direction {
            Direction::Next => token.next_token()?,
            Direction::Prev => token.prev_token()?,
        }
    }
    Some(token)
}

/// Finds the first sibling in the given direction which is not `trivia`
pub fn non_trivia_sibling(element: SyntaxElement, direction: Direction) -> Option<SyntaxElement> {
    return match element {
        NodeOrToken::Node(node) => node.siblings_with_tokens(direction).skip(1).find(not_trivia),
        NodeOrToken::Token(token) => token.siblings_with_tokens(direction).skip(1).find(not_trivia),
    };

    fn not_trivia(element: &SyntaxElement) -> bool {
        match element {
            NodeOrToken::Node(_) => true,
            NodeOrToken::Token(token) => !token.kind().is_trivia(),
        }
    }
}

pub fn find_covering_element(root: &SyntaxNode, range: TextRange) -> SyntaxElement {
    root.covering_element(range)
}

pub fn least_common_ancestor(u: &SyntaxNode, v: &SyntaxNode) -> Option<SyntaxNode> {
    if u == v {
        return Some(u.clone());
    }

    let u_depth = u.ancestors().count();
    let v_depth = v.ancestors().count();
    let keep = u_depth.min(v_depth);

    let u_candidates = u.ancestors().skip(u_depth - keep);
    let v_canidates = v.ancestors().skip(v_depth - keep);
    let (res, _) = u_candidates.zip(v_canidates).find(|(x, y)| x == y)?;
    Some(res)
}

pub fn neighbor<T: AstNode>(me: &T, direction: Direction) -> Option<T> {
    me.syntax().siblings(direction).skip(1).find_map(T::cast)
}

pub fn has_errors(node: &SyntaxNode) -> bool {
    node.children().any(|it| it.kind() == SyntaxKind::ERROR)
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum InsertPosition<T> {
    First,
    Last,
    Before(T),
    After(T),
}

pub struct TreeDiff {
    replacements: FxHashMap<SyntaxElement, SyntaxElement>,
}

impl TreeDiff {
    pub fn into_text_edit(&self, builder: &mut TextEditBuilder) {
        for (from, to) in self.replacements.iter() {
            builder.replace(from.text_range(), to.to_string())
        }
    }

    pub fn is_empty(&self) -> bool {
        self.replacements.is_empty()
    }
}

/// Finds minimal the diff, which, applied to `from`, will result in `to`.
///
/// Specifically, returns a map whose keys are descendants of `from` and values
/// are descendants of `to`, such that  `replace_descendants(from, map) == to`.
///
/// A trivial solution is a singleton map `{ from: to }`, but this function
/// tries to find a more fine-grained diff.
pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
    let mut buf = FxHashMap::default();
    // FIXME: this is both horrible inefficient and gives larger than
    // necessary diff. I bet there's a cool algorithm to diff trees properly.
    go(&mut buf, from.clone().into(), to.clone().into());
    return TreeDiff { replacements: buf };

    fn go(
        buf: &mut FxHashMap<SyntaxElement, SyntaxElement>,
        lhs: SyntaxElement,
        rhs: SyntaxElement,
    ) {
        if lhs.kind() == rhs.kind()
            && lhs.text_range().len() == rhs.text_range().len()
            && match (&lhs, &rhs) {
                (NodeOrToken::Node(lhs), NodeOrToken::Node(rhs)) => {
                    lhs.green() == rhs.green() || lhs.text() == rhs.text()
                }
                (NodeOrToken::Token(lhs), NodeOrToken::Token(rhs)) => lhs.text() == rhs.text(),
                _ => false,
            }
        {
            return;
        }
        if let (Some(lhs), Some(rhs)) = (lhs.as_node(), rhs.as_node()) {
            if lhs.children_with_tokens().count() == rhs.children_with_tokens().count() {
                for (lhs, rhs) in lhs.children_with_tokens().zip(rhs.children_with_tokens()) {
                    go(buf, lhs, rhs)
                }
                return;
            }
        }
        buf.insert(lhs, rhs);
    }
}

/// Adds specified children (tokens or nodes) to the current node at the
/// specific position.
///
/// This is a type-unsafe low-level editing API, if you need to use it,
/// prefer to create a type-safe abstraction on top of it instead.
pub fn insert_children(
    parent: &SyntaxNode,
    position: InsertPosition<SyntaxElement>,
    to_insert: impl IntoIterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let mut to_insert = to_insert.into_iter();
    _insert_children(parent, position, &mut to_insert)
}

fn _insert_children(
    parent: &SyntaxNode,
    position: InsertPosition<SyntaxElement>,
    to_insert: &mut dyn Iterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let mut delta = TextSize::default();
    let to_insert = to_insert.map(|element| {
        delta += element.text_range().len();
        to_green_element(element)
    });

    let mut old_children = parent.green().children().map(|it| match it {
        NodeOrToken::Token(it) => NodeOrToken::Token(it.clone()),
        NodeOrToken::Node(it) => NodeOrToken::Node(it.clone()),
    });

    let new_children = match &position {
        InsertPosition::First => to_insert.chain(old_children).collect::<Vec<_>>(),
        InsertPosition::Last => old_children.chain(to_insert).collect::<Vec<_>>(),
        InsertPosition::Before(anchor) | InsertPosition::After(anchor) => {
            let take_anchor = if let InsertPosition::After(_) = position { 1 } else { 0 };
            let split_at = position_of_child(parent, anchor.clone()) + take_anchor;
            let before = old_children.by_ref().take(split_at).collect::<Vec<_>>();
            before.into_iter().chain(to_insert).chain(old_children).collect::<Vec<_>>()
        }
    };

    with_children(parent, new_children)
}

/// Replaces all nodes in `to_delete` with nodes from `to_insert`
///
/// This is a type-unsafe low-level editing API, if you need to use it,
/// prefer to create a type-safe abstraction on top of it instead.
pub fn replace_children(
    parent: &SyntaxNode,
    to_delete: RangeInclusive<SyntaxElement>,
    to_insert: impl IntoIterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let mut to_insert = to_insert.into_iter();
    _replace_children(parent, to_delete, &mut to_insert)
}

fn _replace_children(
    parent: &SyntaxNode,
    to_delete: RangeInclusive<SyntaxElement>,
    to_insert: &mut dyn Iterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let start = position_of_child(parent, to_delete.start().clone());
    let end = position_of_child(parent, to_delete.end().clone());
    let mut old_children = parent.green().children().map(|it| match it {
        NodeOrToken::Token(it) => NodeOrToken::Token(it.clone()),
        NodeOrToken::Node(it) => NodeOrToken::Node(it.clone()),
    });

    let before = old_children.by_ref().take(start).collect::<Vec<_>>();
    let new_children = before
        .into_iter()
        .chain(to_insert.map(to_green_element))
        .chain(old_children.skip(end + 1 - start))
        .collect::<Vec<_>>();
    with_children(parent, new_children)
}

#[derive(Default)]
pub struct SyntaxRewriter<'a> {
    f: Option<Box<dyn Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a>>,
    //FIXME: add debug_assertions that all elements are in fact from the same file.
    replacements: FxHashMap<SyntaxElement, Replacement>,
}

impl fmt::Debug for SyntaxRewriter<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SyntaxRewriter").field("replacements", &self.replacements).finish()
    }
}

impl<'a> SyntaxRewriter<'a> {
    pub fn from_fn(f: impl Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a) -> SyntaxRewriter<'a> {
        SyntaxRewriter { f: Some(Box::new(f)), replacements: FxHashMap::default() }
    }
    pub fn delete<T: Clone + Into<SyntaxElement>>(&mut self, what: &T) {
        let what = what.clone().into();
        let replacement = Replacement::Delete;
        self.replacements.insert(what, replacement);
    }
    pub fn replace<T: Clone + Into<SyntaxElement>>(&mut self, what: &T, with: &T) {
        let what = what.clone().into();
        let replacement = Replacement::Single(with.clone().into());
        self.replacements.insert(what, replacement);
    }
    pub fn replace_with_many<T: Clone + Into<SyntaxElement>>(
        &mut self,
        what: &T,
        with: Vec<SyntaxElement>,
    ) {
        let what = what.clone().into();
        let replacement = Replacement::Many(with);
        self.replacements.insert(what, replacement);
    }
    pub fn replace_ast<T: AstNode>(&mut self, what: &T, with: &T) {
        self.replace(what.syntax(), with.syntax())
    }

    pub fn rewrite(&self, node: &SyntaxNode) -> SyntaxNode {
        if self.f.is_none() && self.replacements.is_empty() {
            return node.clone();
        }
        self.rewrite_children(node)
    }

    pub fn rewrite_ast<N: AstNode>(self, node: &N) -> N {
        N::cast(self.rewrite(node.syntax())).unwrap()
    }

    pub fn rewrite_root(&self) -> Option<SyntaxNode> {
        assert!(self.f.is_none());
        self.replacements
            .keys()
            .map(|element| match element {
                SyntaxElement::Node(it) => it.clone(),
                SyntaxElement::Token(it) => it.parent(),
            })
            .fold1(|a, b| least_common_ancestor(&a, &b).unwrap())
    }

    fn replacement(&self, element: &SyntaxElement) -> Option<Replacement> {
        if let Some(f) = &self.f {
            assert!(self.replacements.is_empty());
            return f(element).map(Replacement::Single);
        }
        self.replacements.get(element).cloned()
    }

    fn rewrite_children(&self, node: &SyntaxNode) -> SyntaxNode {
        //  FIXME: this could be made much faster.
        let mut new_children = Vec::new();
        for child in node.children_with_tokens() {
            self.rewrite_self(&mut new_children, &child);
        }
        with_children(node, new_children)
    }

    fn rewrite_self(
        &self,
        acc: &mut Vec<NodeOrToken<rowan::GreenNode, rowan::GreenToken>>,
        element: &SyntaxElement,
    ) {
        if let Some(replacement) = self.replacement(&element) {
            match replacement {
                Replacement::Single(NodeOrToken::Node(it)) => {
                    acc.push(NodeOrToken::Node(it.green().clone()))
                }
                Replacement::Single(NodeOrToken::Token(it)) => {
                    acc.push(NodeOrToken::Token(it.green().clone()))
                }
                Replacement::Many(replacements) => {
                    acc.extend(replacements.iter().map(|it| match it {
                        NodeOrToken::Node(it) => NodeOrToken::Node(it.green().clone()),
                        NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
                    }))
                }
                Replacement::Delete => (),
            };
            return;
        }
        let res = match element {
            NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
            NodeOrToken::Node(it) => NodeOrToken::Node(self.rewrite_children(it).green().clone()),
        };
        acc.push(res)
    }
}

impl ops::AddAssign for SyntaxRewriter<'_> {
    fn add_assign(&mut self, rhs: SyntaxRewriter) {
        assert!(rhs.f.is_none());
        self.replacements.extend(rhs.replacements)
    }
}

#[derive(Clone, Debug)]
enum Replacement {
    Delete,
    Single(SyntaxElement),
    Many(Vec<SyntaxElement>),
}

fn with_children(
    parent: &SyntaxNode,
    new_children: Vec<NodeOrToken<rowan::GreenNode, rowan::GreenToken>>,
) -> SyntaxNode {
    let len = new_children.iter().map(|it| it.text_len()).sum::<TextSize>();
    let new_node = rowan::GreenNode::new(rowan::SyntaxKind(parent.kind() as u16), new_children);
    let new_root_node = parent.replace_with(new_node);
    let new_root_node = SyntaxNode::new_root(new_root_node);

    // FIXME: use a more elegant way to re-fetch the node (#1185), make
    // `range` private afterwards
    let mut ptr = SyntaxNodePtr::new(parent);
    ptr.range = TextRange::at(ptr.range.start(), len);
    ptr.to_node(&new_root_node)
}

fn position_of_child(parent: &SyntaxNode, child: SyntaxElement) -> usize {
    parent
        .children_with_tokens()
        .position(|it| it == child)
        .expect("element is not a child of current element")
}

fn to_green_element(element: SyntaxElement) -> NodeOrToken<rowan::GreenNode, rowan::GreenToken> {
    match element {
        NodeOrToken::Node(it) => it.green().clone().into(),
        NodeOrToken::Token(it) => it.green().clone().into(),
    }
}
<commit_msg>Fix `rewrite_root` when there's only 1 replacement<commit_after>//! FIXME: write short doc here

use std::{
    fmt,
    ops::{self, RangeInclusive},
};

use itertools::Itertools;
use ra_text_edit::TextEditBuilder;
use rustc_hash::FxHashMap;

use crate::{
    AstNode, Direction, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxNodePtr,
    SyntaxToken, TextRange, TextSize,
};

/// Returns ancestors of the node at the offset, sorted by length. This should
/// do the right thing at an edge, e.g. when searching for expressions at `{
/// <|>foo }` we will get the name reference instead of the whole block, which
/// we would get if we just did `find_token_at_offset(...).flat_map(|t|
/// t.parent().ancestors())`.
pub fn ancestors_at_offset(
    node: &SyntaxNode,
    offset: TextSize,
) -> impl Iterator<Item = SyntaxNode> {
    node.token_at_offset(offset)
        .map(|token| token.parent().ancestors())
        .kmerge_by(|node1, node2| node1.text_range().len() < node2.text_range().len())
}

/// Finds a node of specific Ast type at offset. Note that this is slightly
/// imprecise: if the cursor is strictly between two nodes of the desired type,
/// as in
///
/// ```no-run
/// struct Foo {}|struct Bar;
/// ```
///
/// then the shorter node will be silently preferred.
pub fn find_node_at_offset<N: AstNode>(syntax: &SyntaxNode, offset: TextSize) -> Option<N> {
    ancestors_at_offset(syntax, offset).find_map(N::cast)
}

/// Skip to next non `trivia` token
pub fn skip_trivia_token(mut token: SyntaxToken, direction: Direction) -> Option<SyntaxToken> {
    while token.kind().is_trivia() {
        token = match direction {
            Direction::Next => token.next_token()?,
            Direction::Prev => token.prev_token()?,
        }
    }
    Some(token)
}

/// Finds the first sibling in the given direction which is not `trivia`
pub fn non_trivia_sibling(element: SyntaxElement, direction: Direction) -> Option<SyntaxElement> {
    return match element {
        NodeOrToken::Node(node) => node.siblings_with_tokens(direction).skip(1).find(not_trivia),
        NodeOrToken::Token(token) => token.siblings_with_tokens(direction).skip(1).find(not_trivia),
    };

    fn not_trivia(element: &SyntaxElement) -> bool {
        match element {
            NodeOrToken::Node(_) => true,
            NodeOrToken::Token(token) => !token.kind().is_trivia(),
        }
    }
}

pub fn find_covering_element(root: &SyntaxNode, range: TextRange) -> SyntaxElement {
    root.covering_element(range)
}

pub fn least_common_ancestor(u: &SyntaxNode, v: &SyntaxNode) -> Option<SyntaxNode> {
    if u == v {
        return Some(u.clone());
    }

    let u_depth = u.ancestors().count();
    let v_depth = v.ancestors().count();
    let keep = u_depth.min(v_depth);

    let u_candidates = u.ancestors().skip(u_depth - keep);
    let v_canidates = v.ancestors().skip(v_depth - keep);
    let (res, _) = u_candidates.zip(v_canidates).find(|(x, y)| x == y)?;
    Some(res)
}

pub fn neighbor<T: AstNode>(me: &T, direction: Direction) -> Option<T> {
    me.syntax().siblings(direction).skip(1).find_map(T::cast)
}

pub fn has_errors(node: &SyntaxNode) -> bool {
    node.children().any(|it| it.kind() == SyntaxKind::ERROR)
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum InsertPosition<T> {
    First,
    Last,
    Before(T),
    After(T),
}

pub struct TreeDiff {
    replacements: FxHashMap<SyntaxElement, SyntaxElement>,
}

impl TreeDiff {
    pub fn into_text_edit(&self, builder: &mut TextEditBuilder) {
        for (from, to) in self.replacements.iter() {
            builder.replace(from.text_range(), to.to_string())
        }
    }

    pub fn is_empty(&self) -> bool {
        self.replacements.is_empty()
    }
}

/// Finds minimal the diff, which, applied to `from`, will result in `to`.
///
/// Specifically, returns a map whose keys are descendants of `from` and values
/// are descendants of `to`, such that  `replace_descendants(from, map) == to`.
///
/// A trivial solution is a singleton map `{ from: to }`, but this function
/// tries to find a more fine-grained diff.
pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
    let mut buf = FxHashMap::default();
    // FIXME: this is both horrible inefficient and gives larger than
    // necessary diff. I bet there's a cool algorithm to diff trees properly.
    go(&mut buf, from.clone().into(), to.clone().into());
    return TreeDiff { replacements: buf };

    fn go(
        buf: &mut FxHashMap<SyntaxElement, SyntaxElement>,
        lhs: SyntaxElement,
        rhs: SyntaxElement,
    ) {
        if lhs.kind() == rhs.kind()
            && lhs.text_range().len() == rhs.text_range().len()
            && match (&lhs, &rhs) {
                (NodeOrToken::Node(lhs), NodeOrToken::Node(rhs)) => {
                    lhs.green() == rhs.green() || lhs.text() == rhs.text()
                }
                (NodeOrToken::Token(lhs), NodeOrToken::Token(rhs)) => lhs.text() == rhs.text(),
                _ => false,
            }
        {
            return;
        }
        if let (Some(lhs), Some(rhs)) = (lhs.as_node(), rhs.as_node()) {
            if lhs.children_with_tokens().count() == rhs.children_with_tokens().count() {
                for (lhs, rhs) in lhs.children_with_tokens().zip(rhs.children_with_tokens()) {
                    go(buf, lhs, rhs)
                }
                return;
            }
        }
        buf.insert(lhs, rhs);
    }
}

/// Adds specified children (tokens or nodes) to the current node at the
/// specific position.
///
/// This is a type-unsafe low-level editing API, if you need to use it,
/// prefer to create a type-safe abstraction on top of it instead.
pub fn insert_children(
    parent: &SyntaxNode,
    position: InsertPosition<SyntaxElement>,
    to_insert: impl IntoIterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let mut to_insert = to_insert.into_iter();
    _insert_children(parent, position, &mut to_insert)
}

fn _insert_children(
    parent: &SyntaxNode,
    position: InsertPosition<SyntaxElement>,
    to_insert: &mut dyn Iterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let mut delta = TextSize::default();
    let to_insert = to_insert.map(|element| {
        delta += element.text_range().len();
        to_green_element(element)
    });

    let mut old_children = parent.green().children().map(|it| match it {
        NodeOrToken::Token(it) => NodeOrToken::Token(it.clone()),
        NodeOrToken::Node(it) => NodeOrToken::Node(it.clone()),
    });

    let new_children = match &position {
        InsertPosition::First => to_insert.chain(old_children).collect::<Vec<_>>(),
        InsertPosition::Last => old_children.chain(to_insert).collect::<Vec<_>>(),
        InsertPosition::Before(anchor) | InsertPosition::After(anchor) => {
            let take_anchor = if let InsertPosition::After(_) = position { 1 } else { 0 };
            let split_at = position_of_child(parent, anchor.clone()) + take_anchor;
            let before = old_children.by_ref().take(split_at).collect::<Vec<_>>();
            before.into_iter().chain(to_insert).chain(old_children).collect::<Vec<_>>()
        }
    };

    with_children(parent, new_children)
}

/// Replaces all nodes in `to_delete` with nodes from `to_insert`
///
/// This is a type-unsafe low-level editing API, if you need to use it,
/// prefer to create a type-safe abstraction on top of it instead.
pub fn replace_children(
    parent: &SyntaxNode,
    to_delete: RangeInclusive<SyntaxElement>,
    to_insert: impl IntoIterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let mut to_insert = to_insert.into_iter();
    _replace_children(parent, to_delete, &mut to_insert)
}

fn _replace_children(
    parent: &SyntaxNode,
    to_delete: RangeInclusive<SyntaxElement>,
    to_insert: &mut dyn Iterator<Item = SyntaxElement>,
) -> SyntaxNode {
    let start = position_of_child(parent, to_delete.start().clone());
    let end = position_of_child(parent, to_delete.end().clone());
    let mut old_children = parent.green().children().map(|it| match it {
        NodeOrToken::Token(it) => NodeOrToken::Token(it.clone()),
        NodeOrToken::Node(it) => NodeOrToken::Node(it.clone()),
    });

    let before = old_children.by_ref().take(start).collect::<Vec<_>>();
    let new_children = before
        .into_iter()
        .chain(to_insert.map(to_green_element))
        .chain(old_children.skip(end + 1 - start))
        .collect::<Vec<_>>();
    with_children(parent, new_children)
}

#[derive(Default)]
pub struct SyntaxRewriter<'a> {
    f: Option<Box<dyn Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a>>,
    //FIXME: add debug_assertions that all elements are in fact from the same file.
    replacements: FxHashMap<SyntaxElement, Replacement>,
}

impl fmt::Debug for SyntaxRewriter<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SyntaxRewriter").field("replacements", &self.replacements).finish()
    }
}

impl<'a> SyntaxRewriter<'a> {
    pub fn from_fn(f: impl Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a) -> SyntaxRewriter<'a> {
        SyntaxRewriter { f: Some(Box::new(f)), replacements: FxHashMap::default() }
    }
    pub fn delete<T: Clone + Into<SyntaxElement>>(&mut self, what: &T) {
        let what = what.clone().into();
        let replacement = Replacement::Delete;
        self.replacements.insert(what, replacement);
    }
    pub fn replace<T: Clone + Into<SyntaxElement>>(&mut self, what: &T, with: &T) {
        let what = what.clone().into();
        let replacement = Replacement::Single(with.clone().into());
        self.replacements.insert(what, replacement);
    }
    pub fn replace_with_many<T: Clone + Into<SyntaxElement>>(
        &mut self,
        what: &T,
        with: Vec<SyntaxElement>,
    ) {
        let what = what.clone().into();
        let replacement = Replacement::Many(with);
        self.replacements.insert(what, replacement);
    }
    pub fn replace_ast<T: AstNode>(&mut self, what: &T, with: &T) {
        self.replace(what.syntax(), with.syntax())
    }

    pub fn rewrite(&self, node: &SyntaxNode) -> SyntaxNode {
        if self.f.is_none() && self.replacements.is_empty() {
            return node.clone();
        }
        self.rewrite_children(node)
    }

    pub fn rewrite_ast<N: AstNode>(self, node: &N) -> N {
        N::cast(self.rewrite(node.syntax())).unwrap()
    }

    /// Returns a node that encompasses all replacements to be done by this rewriter.
    ///
    /// Passing the returned node to `rewrite` will apply all replacements queued up in `self`.
    ///
    /// Returns `None` when there are no replacements.
    pub fn rewrite_root(&self) -> Option<SyntaxNode> {
        assert!(self.f.is_none());
        self.replacements
            .keys()
            .map(|element| match element {
                SyntaxElement::Node(it) => it.clone(),
                SyntaxElement::Token(it) => it.parent(),
            })
            // If we only have one replacement, we must return its parent node, since `rewrite` does
            // not replace the node passed to it.
            .map(|it| it.parent().unwrap_or(it))
            .fold1(|a, b| least_common_ancestor(&a, &b).unwrap())
    }

    fn replacement(&self, element: &SyntaxElement) -> Option<Replacement> {
        if let Some(f) = &self.f {
            assert!(self.replacements.is_empty());
            return f(element).map(Replacement::Single);
        }
        self.replacements.get(element).cloned()
    }

    fn rewrite_children(&self, node: &SyntaxNode) -> SyntaxNode {
        //  FIXME: this could be made much faster.
        let mut new_children = Vec::new();
        for child in node.children_with_tokens() {
            self.rewrite_self(&mut new_children, &child);
        }
        with_children(node, new_children)
    }

    fn rewrite_self(
        &self,
        acc: &mut Vec<NodeOrToken<rowan::GreenNode, rowan::GreenToken>>,
        element: &SyntaxElement,
    ) {
        if let Some(replacement) = self.replacement(&element) {
            match replacement {
                Replacement::Single(NodeOrToken::Node(it)) => {
                    acc.push(NodeOrToken::Node(it.green().clone()))
                }
                Replacement::Single(NodeOrToken::Token(it)) => {
                    acc.push(NodeOrToken::Token(it.green().clone()))
                }
                Replacement::Many(replacements) => {
                    acc.extend(replacements.iter().map(|it| match it {
                        NodeOrToken::Node(it) => NodeOrToken::Node(it.green().clone()),
                        NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
                    }))
                }
                Replacement::Delete => (),
            };
            return;
        }
        let res = match element {
            NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
            NodeOrToken::Node(it) => NodeOrToken::Node(self.rewrite_children(it).green().clone()),
        };
        acc.push(res)
    }
}

impl ops::AddAssign for SyntaxRewriter<'_> {
    fn add_assign(&mut self, rhs: SyntaxRewriter) {
        assert!(rhs.f.is_none());
        self.replacements.extend(rhs.replacements)
    }
}

#[derive(Clone, Debug)]
enum Replacement {
    Delete,
    Single(SyntaxElement),
    Many(Vec<SyntaxElement>),
}

fn with_children(
    parent: &SyntaxNode,
    new_children: Vec<NodeOrToken<rowan::GreenNode, rowan::GreenToken>>,
) -> SyntaxNode {
    let len = new_children.iter().map(|it| it.text_len()).sum::<TextSize>();
    let new_node = rowan::GreenNode::new(rowan::SyntaxKind(parent.kind() as u16), new_children);
    let new_root_node = parent.replace_with(new_node);
    let new_root_node = SyntaxNode::new_root(new_root_node);

    // FIXME: use a more elegant way to re-fetch the node (#1185), make
    // `range` private afterwards
    let mut ptr = SyntaxNodePtr::new(parent);
    ptr.range = TextRange::at(ptr.range.start(), len);
    ptr.to_node(&new_root_node)
}

fn position_of_child(parent: &SyntaxNode, child: SyntaxElement) -> usize {
    parent
        .children_with_tokens()
        .position(|it| it == child)
        .expect("element is not a child of current element")
}

fn to_green_element(element: SyntaxElement) -> NodeOrToken<rowan::GreenNode, rowan::GreenToken> {
    match element {
        NodeOrToken::Node(it) => it.green().clone().into(),
        NodeOrToken::Token(it) => it.green().clone().into(),
    }
}
<|endoftext|>


<commit_before>use core::prelude::*;

use io;

const BACKSPACE: u8 = 0x08;
const TAB: u8 = 0x09;
const NEWLINE: u8 = 0x0A;
const CR: u8 = 0x0D;
const WHITESPACE: u8 = 0x20;

const VGA_ADDRESS: int = 0xB8000;
const VGA_HEIGHT: u16 = 25;
const VGA_WIDTH: u16 = 80;

enum Color {
    Black = 0,
    Blue = 1,
    Green = 2,
    Cyan = 3,
    Red = 4,
    Pink = 5,
    Brown = 6,
    LightGray = 7,
    DarkGray = 8,
    LightBlue = 9,
    LightGreen = 10,
    LightCyan = 11,
    LightRed = 12,
    LightPink = 13,
    Yellow = 14,
    White = 15,
}

static mut VGA: VGA = VGA { height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0 };

struct VGA {
    height: u16,
    width: u16,
    x: u16,
    y: u16,
}

impl VGA {

    fn new() -> VGA {
        VGA {height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0}
    }

    fn back(&mut self) {

        if self.x > 0 {
            self.x -= 1;

            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, WHITESPACE);
        }
    }

    fn forward(&mut self) {
        self.x += 1;

        if self.x >= self.width {
            self.newline();
        }
    }

    fn cr(&mut self) {
        self.x = 0;
    }

    fn newline(&mut self) {
        self.x = 0;
        self.y += 1;
    }

    fn offset(&self) -> u16 {
        self.y * self.width + self.x
    }

    fn reset(&mut self) {
        self.x = 0;
        self.y = 0;
        self.mov();
    }

    // handle a tab by increasing the cursor's X, but only to a point
    // where it is divisible by 8
    fn tab(&mut self) {
        self.x = (self.x + 8) & !(8 - 1);
    }

    fn mov(&mut self) {
        let offset = self.offset() as u8;

        io::port::write(0x3D4, 14);  // tell the VGA board we are setting the high cursor byte
        io::port::write(0x3D5, offset >> 8); // send the high cursor byte
        io::port::write(0x3D4, 15);  // tell the VGA board we are setting the low cursor byte
        io::port::write(0x3D5, offset);
    }

    fn clear (&mut self) {
        let mut x: u16 = 0;

        loop {
            if x > 80 * 200 {
                break;
            }

            self.put(x, Black as u16, White as u16, WHITESPACE);
            x += 1;
        }
    }

    fn put(&mut self, offset: u16, background: u16, foreground: u16, character: u8) {
        let pixel: u16 = (background << 12) | (foreground << 8) | character as u16;

        unsafe {
                *((VGA_ADDRESS + offset as int * 2) as *mut u16) = pixel;
            }
    }

    fn puti(&mut self, integer: uint) {  // FIXME
        let mut integer = integer;

        loop {
            self.putc((integer % 10) as u8 + '0' as u8);
            integer = integer / 10;

            if integer == 0 {
                break;
            }
        }
    }

    fn putc(&mut self, character: u8) {

        if character == BACKSPACE {
            self.back();
        }

        else if character == TAB {
            self.tab();
        }

        else if character == NEWLINE {
            self.newline();
        }

        else if character == CR {
            self.cr();
        }

        else if character >= WHITESPACE {
            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, character);
            self.forward();
        }

        self.mov();
   }

    fn puth(&mut self, integer: uint) {
        self.puts("0x");

        let mut nibbles = 1u;

        while (integer >> nibbles * 4) > 0 {
                nibbles += 1
        }

        for i in range(0, nibbles) {
            let nibble = ((integer >> (nibbles - i - 1) * 4) & 0xF) as u8;
            let character = if nibble < 10 { '0' as u8 + nibble } else { 'a' as u8 + nibble - 10 };
            self.putc(character);
        }
    }

    fn puts(&mut self, string: &str) {
        for character in string.bytes() {
            self.putc(character);
        }

    }
}

pub fn clear() {
    unsafe {
        VGA.clear();
    }
}

pub fn puth(integer: uint) {
    unsafe {
        VGA.puth(integer);
    }
}

pub fn puti(integer: uint) {
    unsafe {
        VGA.puti(integer);
    }
}

pub fn putc(character: u8) {
    unsafe {
        VGA.putc(character);
    }
}

pub fn puts(string: &str) {
    unsafe {
        VGA.puts(string);
    }
}
<commit_msg>Fix print integers<commit_after>use core::prelude::*;

use io;

const BACKSPACE: u8 = 0x08;
const TAB: u8 = 0x09;
const NEWLINE: u8 = 0x0A;
const CR: u8 = 0x0D;
const WHITESPACE: u8 = 0x20;

const VGA_ADDRESS: int = 0xB8000;
const VGA_HEIGHT: u16 = 25;
const VGA_WIDTH: u16 = 80;

enum Color {
    Black = 0,
    Blue = 1,
    Green = 2,
    Cyan = 3,
    Red = 4,
    Pink = 5,
    Brown = 6,
    LightGray = 7,
    DarkGray = 8,
    LightBlue = 9,
    LightGreen = 10,
    LightCyan = 11,
    LightRed = 12,
    LightPink = 13,
    Yellow = 14,
    White = 15,
}

static mut VGA: VGA = VGA { height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0 };

struct VGA {
    height: u16,
    width: u16,
    x: u16,
    y: u16,
}

impl VGA {

    fn new() -> VGA {
        VGA {height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0}
    }

    fn back(&mut self) {

        if self.x > 0 {
            self.x -= 1;

            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, WHITESPACE);
        }
    }

    fn forward(&mut self) {
        self.x += 1;

        if self.x >= self.width {
            self.newline();
        }
    }

    fn cr(&mut self) {
        self.x = 0;
    }

    fn newline(&mut self) {
        self.x = 0;
        self.y += 1;
    }

    fn offset(&self) -> u16 {
        self.y * self.width + self.x
    }

    fn reset(&mut self) {
        self.x = 0;
        self.y = 0;
        self.mov();
    }

    // handle a tab by increasing the cursor's X, but only to a point
    // where it is divisible by 8
    fn tab(&mut self) {
        self.x = (self.x + 8) & !(8 - 1);
    }

    fn mov(&mut self) {
        let offset = self.offset() as u8;

        io::port::write(0x3D4, 14);  // tell the VGA board we are setting the high cursor byte
        io::port::write(0x3D5, offset >> 8); // send the high cursor byte
        io::port::write(0x3D4, 15);  // tell the VGA board we are setting the low cursor byte
        io::port::write(0x3D5, offset);
    }

    fn clear (&mut self) {
        let mut x: u16 = 0;

        loop {
            if x > 80 * 200 {
                break;
            }

            self.put(x, Black as u16, White as u16, WHITESPACE);
            x += 1;
        }
    }

    fn put(&mut self, offset: u16, background: u16, foreground: u16, character: u8) {
        let pixel: u16 = (background << 12) | (foreground << 8) | character as u16;

        unsafe {
                *((VGA_ADDRESS + offset as int * 2) as *mut u16) = pixel;
            }
    }

    fn putc(&mut self, character: u8) {

        if character == BACKSPACE {
            self.back();
        }

        else if character == TAB {
            self.tab();
        }

        else if character == NEWLINE {
            self.newline();
        }

        else if character == CR {
            self.cr();
        }

        else if character >= WHITESPACE {
            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, character);
            self.forward();
        }

        self.mov();
    }

    fn puti(&mut self, integer: uint) {
        if integer == 0 {
            self.puts("0");
        }
        else {
            let mut integer = integer;
            let mut reversed = 0u;

            while integer > 0 {
                reversed *= 10;
                reversed += integer % 10;
                integer /= 10;
            }

            while reversed > 0 {
                let character = (reversed % 10) as u8 + '0' as u8;
                self.putc(character);
                reversed /= 10;
            }
        }
    }

    fn puth(&mut self, integer: uint) {
        self.puts("0x");

        let mut nibbles = 1u;

        while (integer >> nibbles * 4) > 0 {
                nibbles += 1
        }

        for i in range(0, nibbles) {
            let nibble = ((integer >> (nibbles - i - 1) * 4) & 0xF) as u8;
            let character = if nibble < 10 { '0' as u8 + nibble } else { 'a' as u8 + nibble - 10 };
            self.putc(character);
        }
    }

    fn puts(&mut self, string: &str) {
        for character in string.bytes() {
            self.putc(character);
        }

    }
}

pub fn clear() {
    unsafe {
        VGA.clear();
    }
}

pub fn puth(integer: uint) {
    unsafe {
        VGA.puth(integer);
    }
}

pub fn puti(integer: uint) {
    unsafe {
        VGA.puti(integer);
    }
}

pub fn putc(character: u8) {
    unsafe {
        VGA.putc(character);
    }
}

pub fn puts(string: &str) {
    unsafe {
        VGA.puts(string);
    }
}
<|endoftext|>


<commit_before>use serde_json;
use types::*;
use utils::*;
use languageclient::*;

pub trait IVim {
    fn echo(&self, message: &str) -> Result<()>;
    fn echo_ellipsis(&self, message: &str) -> Result<()>;
    fn echomsg(&self, message: &str) -> Result<()>;
    fn echoerr(&self, message: &str) -> Result<()>;
    fn echowarn(&self, message: &str) -> Result<()>;
    fn eval<E, T>(&self, exp: E) -> Result<T>
    where
        E: ToVimExp,
        T: DeserializeOwned;
    fn command(&self, cmd: &str) -> Result<()>;
    fn getbufline(&self, bufexp: Option<&str>) -> Result<Vec<String>>;
    fn goto_location(&self, filename: &str, line: u64, character: u64) -> Result<()>;
}

// Whether should it be Mutex or RwLock.
// Even though RwLock allows several readers at the same time, it won't bring too much good in this
// use case. As in this project, read and write are almost same amount of short operations. For
// RwLock to work, a writer still needs to wait for all readers finish their work before making the
// change.

impl IVim for Arc<Mutex<State>> {
    fn echo(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echo '{}'", message);
        self.command(cmd.as_str())
    }

    fn echo_ellipsis(&self, message: &str) -> Result<()> {
        let columns: usize = self.eval("&columns")?;
        let mut message = message.replace('\n', ". ");
        if message.len() > columns - 12 {
            message = message[..columns - 15].to_owned();
            message += "...";
        }
        self.echo(message.as_str())
    }

    fn echomsg(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echomsg '{}'", message);
        self.command(cmd.as_str())
    }

    fn echoerr(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl Error | echomsg '{}' | echohl None", message);
        self.command(cmd.as_str())
    }

    fn echowarn(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl WarningMsg | echomsg '{}' | echohl None", message);
        self.command(cmd.as_str())
    }

    fn eval<E, T>(&self, exp: E) -> Result<T>
    where
        E: ToVimExp,
        T: DeserializeOwned,
    {
        let result = self.call(None, "eval", exp.to_exp())?;
        Ok(serde_json::from_value(result)?)
    }

    fn command(&self, cmd: &str) -> Result<()> {
        self.notify(None, "execute", cmd)
    }

    fn getbufline(&self, bufexp: Option<&str>) -> Result<Vec<String>> {
        let bufexp = if let Some(bufexp) = bufexp {
            json!(bufexp)
        } else {
            json!(0)
        };

        let result = self.call(None, "getbufline", json!([bufexp, '.', '$']))?;
        Ok(serde_json::from_value(result)?)
    }

    fn goto_location(&self, filename: &str, line: u64, character: u64) -> Result<()> {
        let bufname: Option<String> = self.eval(format!("bufname('{}')", filename).as_str())?;
        let bufname = bufname.unwrap_or("".to_owned());

        let action = if bufname.is_empty() { "edit" } else { "buffer" };
        let cmd = format!(
            "echo | execute '{} +:call\\ cursor({},{}) ' . fnameescape('{}')",
            action,
            line + 1,
            character + 1,
            filename
        );

        self.command(cmd.as_str())
    }
}
<commit_msg>Fix getbufline.<commit_after>use serde_json;
use types::*;
use utils::*;
use languageclient::*;

pub trait IVim {
    fn echo(&self, message: &str) -> Result<()>;
    fn echo_ellipsis(&self, message: &str) -> Result<()>;
    fn echomsg(&self, message: &str) -> Result<()>;
    fn echoerr(&self, message: &str) -> Result<()>;
    fn echowarn(&self, message: &str) -> Result<()>;
    fn eval<E, T>(&self, exp: E) -> Result<T>
    where
        E: ToVimExp,
        T: DeserializeOwned;
    fn command(&self, cmd: &str) -> Result<()>;
    fn getbufline(&self, bufexp: Option<&str>) -> Result<Vec<String>>;
    fn goto_location(&self, filename: &str, line: u64, character: u64) -> Result<()>;
}

// Whether should it be Mutex or RwLock.
// Even though RwLock allows several readers at the same time, it won't bring too much good in this
// use case. As in this project, read and write are almost same amount of short operations. For
// RwLock to work, a writer still needs to wait for all readers finish their work before making the
// change.

impl IVim for Arc<Mutex<State>> {
    fn echo(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echo '{}'", message);
        self.command(cmd.as_str())
    }

    fn echo_ellipsis(&self, message: &str) -> Result<()> {
        let columns: usize = self.eval("&columns")?;
        let mut message = message.replace('\n', ". ");
        if message.len() > columns - 12 {
            message = message[..columns - 15].to_owned();
            message += "...";
        }
        self.echo(message.as_str())
    }

    fn echomsg(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echomsg '{}'", message);
        self.command(cmd.as_str())
    }

    fn echoerr(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl Error | echomsg '{}' | echohl None", message);
        self.command(cmd.as_str())
    }

    fn echowarn(&self, message: &str) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl WarningMsg | echomsg '{}' | echohl None", message);
        self.command(cmd.as_str())
    }

    fn eval<E, T>(&self, exp: E) -> Result<T>
    where
        E: ToVimExp,
        T: DeserializeOwned,
    {
        let result = self.call(None, "eval", exp.to_exp())?;
        Ok(serde_json::from_value(result)?)
    }

    fn command(&self, cmd: &str) -> Result<()> {
        self.notify(None, "execute", cmd)
    }

    fn getbufline(&self, bufexp: Option<&str>) -> Result<Vec<String>> {
        let bufexp = if let Some(bufexp) = bufexp {
            json!(bufexp)
        } else {
            json!(0)
        };

        let result = self.call(None, "getbufline", json!([bufexp, 1, '$']))?;
        Ok(serde_json::from_value(result)?)
    }

    fn goto_location(&self, filename: &str, line: u64, character: u64) -> Result<()> {
        let bufname: Option<String> = self.eval(format!("bufname('{}')", filename).as_str())?;
        let bufname = bufname.unwrap_or("".to_owned());

        let action = if bufname.is_empty() { "edit" } else { "buffer" };
        let cmd = format!(
            "echo | execute '{} +:call\\ cursor({},{}) ' . fnameescape('{}')",
            action,
            line + 1,
            character + 1,
            filename
        );

        self.command(cmd.as_str())
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::env;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};

use strum_macros::EnumString;
use walkdir;

use crate::template::{self, RenderTemplateError, TemplateInput};
use crate::util;
use crate::version::{self, Version};

#[derive(Debug, EnumString)]
#[strum(serialize_all = "lowercase")]
pub enum BumpVersion {
    Major,
    Minor,
    Patch,
    Prerelease,
    Build,
}

pub enum VutError {
    VersionFileOpen(util::FileError),
    VersionFileRead(io::Error),
    VersionFileWrite(io::Error),
    TemplateGenerate(RenderTemplateError),
    Other(Cow<'static, str>),
}

pub struct Vut {
    root_path: PathBuf,
    version_file_path: PathBuf,
}

impl Vut {
    const VERSION_FILENAME: &'static str = "VERSION";

    pub fn new(path: impl AsRef<Path>) -> Self {
        let path = path.as_ref();

        Self {
            root_path: path.to_path_buf(),
            version_file_path: path.join(Self::VERSION_FILENAME),
        }
    }

    pub fn from_path(path: impl AsRef<Path>) -> Option<Self> {
        if let Some(version_file_path) = util::locate_config_file(path, Self::VERSION_FILENAME) {
            // Can this actually fail?
            let root_path = version_file_path.parent().unwrap();

            Some(Self {
                root_path: root_path.to_path_buf(),
                version_file_path,
            })
        } else {
            None
        }
    }

    pub fn from_current_dir() -> Option<Self> {
        let current_dir = env::current_dir().unwrap();

        Self::from_path(current_dir)
    }

    pub fn exists(&self) -> bool {
        self.version_file_path.exists()
    }

    pub fn get_root_path(&self) -> &Path {
        &self.root_path
    }

    pub fn get_version_file_path(&self) -> &Path {
        &self.version_file_path
    }

    pub fn get_version(&self) -> Result<Version, VutError> {
        let version_str = {
            let mut file = util::open_file(&self.version_file_path)
                .map_err(|err| VutError::VersionFileOpen(err))?;

            let mut version_str = String::new();

            file.read_to_string(&mut version_str)
                .map_err(|err| VutError::VersionFileRead(err))?;

            version_str
        };

        let version = version_str.parse()
            .map_err(|err| VutError::Other(Cow::Owned(err)))?;

        Ok(version)
    }

    pub fn set_version(&self, version: &Version) -> Result<(), VutError> {
        let mut file = util::create_file(&self.version_file_path)
            .map_err(|err| VutError::VersionFileOpen(err))?;

        file.write(version.to_string().as_bytes())
            .map_err(|err| VutError::VersionFileWrite(err))?;

        Ok(())
    }

    pub fn bump_version(&self, bump_version: BumpVersion) -> Result<Version, VutError> {
        let version = self.get_version()?;

        let version = match bump_version {
            BumpVersion::Major => version.bump_major(),
            BumpVersion::Minor => version.bump_minor(),
            BumpVersion::Patch => version.bump_patch(),
            BumpVersion::Prerelease => version.bump_prerelease(),
            BumpVersion::Build => version.bump_build(),
        };

        self.set_version(&version)?;

        Ok(version)
    }

    pub fn generate_template_input(&self) -> Result<TemplateInput, VutError> {
        let version = self.get_version()?;

        let mut template_input = TemplateInput::new();

        let split_prerelease = version::split_numbered_prerelease(&version.prerelease);
        let split_build = version::split_numbered_prerelease(&version.build);

        template_input.values.insert("FullVersion".to_owned(), version.to_string());
        template_input.values.insert("Version".to_owned(), Version { build: "".to_owned(), ..version.clone() }.to_string());
        template_input.values.insert("MajorMinorPatch".to_owned(), format!("{}.{}.{}", version.major, version.minor, version.patch));
        template_input.values.insert("MajorMinor".to_owned(), format!("{}.{}", version.major, version.minor));
        template_input.values.insert("Major".to_owned(), format!("{}", version.major));
        template_input.values.insert("Minor".to_owned(), format!("{}", version.minor));
        template_input.values.insert("Patch".to_owned(), format!("{}", version.patch));
        template_input.values.insert("Prerelease".to_owned(), version.prerelease.to_owned());
        template_input.values.insert("PrereleasePrefix".to_owned(), split_prerelease.and_then(|sp| Some(sp.0.to_owned())).unwrap_or_else(|| "".to_owned()));
        template_input.values.insert("PrereleaseNumber".to_owned(), split_prerelease.and_then(|sp| Some(format!("{}", sp.1))).unwrap_or_else(|| "".to_owned()));
        template_input.values.insert("Build".to_owned(), version.build.to_owned());
        template_input.values.insert("BuildPrefix".to_owned(), split_build.and_then(|sp| Some(sp.0.to_owned())).unwrap_or_else(|| "".to_owned()));
        template_input.values.insert("BuildNumber".to_owned(), split_build.and_then(|sp| Some(format!("{}", sp.1))).unwrap_or_else(|| "".to_owned()));

        Ok(template_input)
    }

    pub fn generate_output(&self) -> Result<(), VutError> {
        let root_path = &self.root_path;
        let template_input = self.generate_template_input()?;

        let files: Vec<PathBuf> = walkdir::WalkDir::new(root_path).into_iter()
            // Filter known VCS metadata directories
            .filter_entry(|entry| !entry.file_name().to_str().map(|s| s == ".git" || s == ".hg").unwrap_or(false))
            .filter_map(|entry| entry.ok())
            .map(|entry| entry.into_path())
            // Only include template files
            .filter(|path| {
                match path.extension() {
                    Some(ext) => ext.to_str().unwrap() == "vutemplate",
                    None => false,
                }
            })
            // Make paths absolute
            .map(|path| util::normalize_path(&path).into_owned())
            // Exclude paths outside the root path
            .filter(|path| path.starts_with(&root_path))
            .collect();

        let mut processed_files: Vec<PathBuf> = Vec::new();
        let mut generated_files: Vec<PathBuf> = Vec::new();

        for file in files {
            let generated_file = template::generate_template::<template::processor::ClassicProcessor>(&file, &template_input, None)
                .map_err(|err| VutError::TemplateGenerate(err))?;

            processed_files.push(file);
            generated_files.push(generated_file);
        }

        Ok(())
    }
}
<commit_msg>refactor: Avoid using .to_str() in extension string when checking extension<commit_after>use std::borrow::Cow;
use std::env;
use std::ffi::OsStr;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};

use lazy_static::lazy_static;
use strum_macros::EnumString;
use walkdir;

use crate::template::{self, RenderTemplateError, TemplateInput};
use crate::util;
use crate::version::{self, Version};

lazy_static! {
    static ref VUTEMPLATE_EXTENSION: &'static OsStr = OsStr::new("vutemplate");
}

#[derive(Debug, EnumString)]
#[strum(serialize_all = "lowercase")]
pub enum BumpVersion {
    Major,
    Minor,
    Patch,
    Prerelease,
    Build,
}

pub enum VutError {
    VersionFileOpen(util::FileError),
    VersionFileRead(io::Error),
    VersionFileWrite(io::Error),
    TemplateGenerate(RenderTemplateError),
    Other(Cow<'static, str>),
}

pub struct Vut {
    root_path: PathBuf,
    version_file_path: PathBuf,
}

impl Vut {
    const VERSION_FILENAME: &'static str = "VERSION";

    pub fn new(path: impl AsRef<Path>) -> Self {
        let path = path.as_ref();

        Self {
            root_path: path.to_path_buf(),
            version_file_path: path.join(Self::VERSION_FILENAME),
        }
    }

    pub fn from_path(path: impl AsRef<Path>) -> Option<Self> {
        if let Some(version_file_path) = util::locate_config_file(path, Self::VERSION_FILENAME) {
            // Can this actually fail?
            let root_path = version_file_path.parent().unwrap();

            Some(Self {
                root_path: root_path.to_path_buf(),
                version_file_path,
            })
        } else {
            None
        }
    }

    pub fn from_current_dir() -> Option<Self> {
        let current_dir = env::current_dir().unwrap();

        Self::from_path(current_dir)
    }

    pub fn exists(&self) -> bool {
        self.version_file_path.exists()
    }

    pub fn get_root_path(&self) -> &Path {
        &self.root_path
    }

    pub fn get_version_file_path(&self) -> &Path {
        &self.version_file_path
    }

    pub fn get_version(&self) -> Result<Version, VutError> {
        let version_str = {
            let mut file = util::open_file(&self.version_file_path)
                .map_err(|err| VutError::VersionFileOpen(err))?;

            let mut version_str = String::new();

            file.read_to_string(&mut version_str)
                .map_err(|err| VutError::VersionFileRead(err))?;

            version_str
        };

        let version = version_str.parse()
            .map_err(|err| VutError::Other(Cow::Owned(err)))?;

        Ok(version)
    }

    pub fn set_version(&self, version: &Version) -> Result<(), VutError> {
        let mut file = util::create_file(&self.version_file_path)
            .map_err(|err| VutError::VersionFileOpen(err))?;

        file.write(version.to_string().as_bytes())
            .map_err(|err| VutError::VersionFileWrite(err))?;

        Ok(())
    }

    pub fn bump_version(&self, bump_version: BumpVersion) -> Result<Version, VutError> {
        let version = self.get_version()?;

        let version = match bump_version {
            BumpVersion::Major => version.bump_major(),
            BumpVersion::Minor => version.bump_minor(),
            BumpVersion::Patch => version.bump_patch(),
            BumpVersion::Prerelease => version.bump_prerelease(),
            BumpVersion::Build => version.bump_build(),
        };

        self.set_version(&version)?;

        Ok(version)
    }

    pub fn generate_template_input(&self) -> Result<TemplateInput, VutError> {
        let version = self.get_version()?;

        let mut template_input = TemplateInput::new();

        let split_prerelease = version::split_numbered_prerelease(&version.prerelease);
        let split_build = version::split_numbered_prerelease(&version.build);

        template_input.values.insert("FullVersion".to_owned(), version.to_string());
        template_input.values.insert("Version".to_owned(), Version { build: "".to_owned(), ..version.clone() }.to_string());
        template_input.values.insert("MajorMinorPatch".to_owned(), format!("{}.{}.{}", version.major, version.minor, version.patch));
        template_input.values.insert("MajorMinor".to_owned(), format!("{}.{}", version.major, version.minor));
        template_input.values.insert("Major".to_owned(), format!("{}", version.major));
        template_input.values.insert("Minor".to_owned(), format!("{}", version.minor));
        template_input.values.insert("Patch".to_owned(), format!("{}", version.patch));
        template_input.values.insert("Prerelease".to_owned(), version.prerelease.to_owned());
        template_input.values.insert("PrereleasePrefix".to_owned(), split_prerelease.and_then(|sp| Some(sp.0.to_owned())).unwrap_or_else(|| "".to_owned()));
        template_input.values.insert("PrereleaseNumber".to_owned(), split_prerelease.and_then(|sp| Some(format!("{}", sp.1))).unwrap_or_else(|| "".to_owned()));
        template_input.values.insert("Build".to_owned(), version.build.to_owned());
        template_input.values.insert("BuildPrefix".to_owned(), split_build.and_then(|sp| Some(sp.0.to_owned())).unwrap_or_else(|| "".to_owned()));
        template_input.values.insert("BuildNumber".to_owned(), split_build.and_then(|sp| Some(format!("{}", sp.1))).unwrap_or_else(|| "".to_owned()));

        Ok(template_input)
    }

    pub fn generate_output(&self) -> Result<(), VutError> {
        let root_path = &self.root_path;
        let template_input = self.generate_template_input()?;

        let files: Vec<PathBuf> = walkdir::WalkDir::new(root_path).into_iter()
            // Filter known VCS metadata directories
            .filter_entry(|entry| !entry.file_name().to_str().map(|s| s == ".git" || s == ".hg").unwrap_or(false))
            .filter_map(|entry| entry.ok())
            .map(|entry| entry.into_path())
            // Only include template files
            .filter(|path| {
                match path.extension() {
                    Some(ext) => ext == *VUTEMPLATE_EXTENSION,
                    None => false,
                }
            })
            // Make paths absolute
            .map(|path| util::normalize_path(&path).into_owned())
            // Exclude paths outside the root path
            .filter(|path| path.starts_with(&root_path))
            .collect();

        let mut processed_files: Vec<PathBuf> = Vec::new();
        let mut generated_files: Vec<PathBuf> = Vec::new();

        for file in files {
            let generated_file = template::generate_template::<template::processor::ClassicProcessor>(&file, &template_input, None)
                .map_err(|err| VutError::TemplateGenerate(err))?;

            processed_files.push(file);
            generated_files.push(generated_file);
        }

        Ok(())
    }
}
<|endoftext|>


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

//! This module contains a wrapper around the libsodium implementation of ed25519.
//! It reduces the size of signatures to 64 byte.

use sodiumoxide;
use sodiumoxide::crypto::sign::ed25519;

/// This is the length of a ed25519 private key.
pub const PRIVATE_KEY_LEN: usize = 64;

/// This is the length of a ed25519 public key.
pub const PUBLIC_KEY_LEN: usize = 32;

/// This is the length of a ed25519 signature.
pub const SIGNATURE_LEN: usize = 64;

static mut inited: bool = false;

/// This method generates a random ed25519 keypair from a cryptographically secure source
/// (on unix this is /dev/urandom). Returns (public_key, private_key).
pub fn generate_keypair() -> ([u8; PUBLIC_KEY_LEN], [u8; PRIVATE_KEY_LEN]) {

    unsafe {

        // we cann do this simple unsafe "lazy init",
        // because it would be OK to call init() twice.

        if !inited {
            inited = true;
            sodiumoxide::init();
        }
    }

    let (pk, sk) = ed25519::gen_keypair();

    let public_key = pk.0;
    let private_key = sk.0;

    assert_eq!(PUBLIC_KEY_LEN, public_key.len());
    assert_eq!(PRIVATE_KEY_LEN, private_key.len());

    (public_key, private_key)
}

/// This method takes a data vector and a private key and computes the signature which can be
/// verified using the public key.
pub fn sign(data: &[u8], private_key: &[u8]) -> Vec<u8> {
    assert_eq!(private_key.len(), PRIVATE_KEY_LEN);

    let sk = ed25519::SecretKey::from_slice(private_key);
    let sk = sk.as_ref().unwrap();

    let s = ed25519::sign(data, &sk);

    let mut v = Vec::new();
    v.extend_from_slice(&s[0..64]);
    v
}

/// This method takes a data vector, a signature and a public key and returns true, if the
/// signature has been created using the correct private key.
pub fn verify(data: &[u8], signature: &[u8], public_key: &[u8]) -> bool {

    let pk = ed25519::PublicKey::from_slice(public_key);
    let pk = pk.as_ref().unwrap();

    let mut vi = Vec::with_capacity(SIGNATURE_LEN + data.len());
    vi.extend_from_slice(&signature);
    vi.extend_from_slice(data);

    let r = ed25519::verify(&vi, &pk);

    if r.is_err() {
        false
    } else {
        let bytes = r.unwrap();
        bytes == data
    }
}

#[test]
fn test_ed25519_simple() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 128][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_ed25519_shortmsg() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 32][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_testvectors() {
    use rustc_serialize::json;
    use bytescontainer::BytesContainer;

    let decode = |t| {
        let t = json::decode::<BytesContainer>(&format!("\"{}\"", t)).unwrap();
        t.get().clone()
    };

    // test vectors from
    // https://github.com/cryptosphere/rbnacl/blob/master/lib/rbnacl/test_vectors.rb

    let prv = decode("b18e1d0045995ec3d010c387ccfeb984d783af8fbb0f40fa7db126d889f6dadd77f48b59cae\
                      da77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let pbl = decode("77f48b59caeda77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let msg = decode("916c7d1d268fc0e77c1bef238432573c39be577bbea099893\
                      6add2b50a653171ce18a542b0b7f96c1691a3be6031522894a8\
                      634183eda38798a0c5d5d79fbd01dd04a8646d71873b77b2219\
                      98a81922d8105f892316369d5224c9983372d2313c6b1f4556e\
                      a26ba49d46e8b561e0fc76633ac9766e68e21fba7edca93c4c7\
                      460376d7f3ac22ff372c18f613f2ae2e856af40");

    let sig = decode("6bd710a368c1249923fc7a1610747403040f0cc30815a00f9ff548a896bbda0b4eb2ca19ebc\
                      f917f0f34200a9edbad3901b64ab09cc5ef7b9bcc3c40c0ff7509");

    let sig_my = sign(&msg, &prv);

    // both signatures must be equal
    assert_eq!(&sig_my, &sig);

    // the signature must be valid
    assert_eq!(true, verify(&msg, &sig_my, &pbl));
}
<commit_msg>small changes<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2016 Marvin Böcker
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! This module contains a wrapper around the libsodium implementation of ed25519.
//! It reduces the size of signatures to 64 byte.

use sodiumoxide;
use sodiumoxide::crypto::sign::ed25519;

/// This is the length of a ed25519 private key.
pub const PRIVATE_KEY_LEN: usize = 64;

/// This is the length of a ed25519 public key.
pub const PUBLIC_KEY_LEN: usize = 32;

/// This is the length of a ed25519 signature.
pub const SIGNATURE_LEN: usize = 64;

static mut inited: bool = false;

/// This method generates a random ed25519 keypair from a cryptographically secure source
/// (on unix this is /dev/urandom). Returns (public_key, private_key).
pub fn generate_keypair() -> ([u8; PUBLIC_KEY_LEN], [u8; PRIVATE_KEY_LEN]) {

    unsafe {

        // we cann do this simple unsafe "lazy init",
        // because it would be OK to call init() twice.

        if !inited {
            inited = true;
            sodiumoxide::init();
        }
    }

    let (pk, sk) = ed25519::gen_keypair();

    let public_key = pk.0;
    let private_key = sk.0;

    assert_eq!(PUBLIC_KEY_LEN, public_key.len());
    assert_eq!(PRIVATE_KEY_LEN, private_key.len());

    (public_key, private_key)
}

/// This method takes a data vector and a private key and computes the signature which can be
/// verified using the public key.
pub fn sign(data: &[u8], private_key: &[u8]) -> Vec<u8> {
    assert_eq!(private_key.len(), PRIVATE_KEY_LEN);

    let sk = ed25519::SecretKey::from_slice(private_key);
    let sk = sk.as_ref().unwrap();

    let s = ed25519::sign(data, &sk);

    let mut v = Vec::new();
    v.extend_from_slice(&s[0..64]);
    v
}

/// This method takes a data vector, a signature and a public key and returns true, if the
/// signature has been created using the correct private key.
pub fn verify(data: &[u8], signature: &[u8], public_key: &[u8]) -> bool {

    let pk = ed25519::PublicKey::from_slice(public_key);
    let pk = pk.as_ref().unwrap();

    let mut vi = Vec::with_capacity(SIGNATURE_LEN + data.len());
    vi.extend_from_slice(&signature);
    vi.extend_from_slice(data);

    let r = ed25519::verify(&vi, &pk);

    if r.is_err() {
        false
    } else {
        let bytes = r.unwrap();
        bytes == data
    }
}

#[test]
fn test_ed25519_simple() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 128][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_ed25519_shortmsg() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 32][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_testvectors() {
    use rustc_serialize::json;
    use bytescontainer::BytesContainer;

    let decode = |t| {
        let t = json::decode::<BytesContainer>(&format!("\"{}\"", t)).unwrap();
        t.get().clone()
    };

    // test vectors from
    // https://github.com/cryptosphere/rbnacl/blob/master/lib/rbnacl/test_vectors.rb

    let prv = decode("b18e1d0045995ec3d010c387ccfeb984d783af8fbb0f40fa7db126d889f6dadd77f48b59cae\
                      da77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let pbl = decode("77f48b59caeda77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let msg = decode("916c7d1d268fc0e77c1bef238432573c39be577bbea099893\
                      6add2b50a653171ce18a542b0b7f96c1691a3be6031522894a8\
                      634183eda38798a0c5d5d79fbd01dd04a8646d71873b77b2219\
                      98a81922d8105f892316369d5224c9983372d2313c6b1f4556e\
                      a26ba49d46e8b561e0fc76633ac9766e68e21fba7edca93c4c7\
                      460376d7f3ac22ff372c18f613f2ae2e856af40");

    let sig = decode("6bd710a368c1249923fc7a1610747403040f0cc30815a00f9ff548a896bbda0b4eb2ca19ebc\
                      f917f0f34200a9edbad3901b64ab09cc5ef7b9bcc3c40c0ff7509");

    let sig_my = sign(&msg, &prv);

    // both signatures must be equal
    assert_eq!(&sig_my, &sig);

    // the signature must be valid
    assert_eq!(true, verify(&msg, &sig_my, &pbl));
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::fs::File;
use std::f64;

use image::{ DynamicImage, GenericImage, ImageFormat };
use image;

use bitreader::BitReader;

use num_traits::float::Float;

pub struct CoverImage {
    cover_image: DynamicImage
}

impl CoverImage {
    pub fn new(file_path: &str) -> CoverImage {
        CoverImage {
            cover_image: image::open(&Path::new(&file_path)).unwrap()
        }
    }

    pub fn encode_with(&self, message: &str) {
        let bit_vector = get_bit_vec(message);

        encode(&self.cover_image, &bit_vector);

        let ref mut fout = File::create(&Path::new("/home/hugh/Pictures/yurt.png")).unwrap();

        let _ = self.cover_image.save(fout, ImageFormat::PNG).unwrap();
    }
}

// Get a Vector of Bits to Encode the Image with
fn get_bit_vec(message: &str) -> Vec<u8> {
    let mut bit_vector = Vec::new();
    let mut reader = BitReader::new(&message.as_bytes());

    // Multiplied by 8 because it's a Vec of bits not bytes
    for _ in 0..message.len() * 8 {
        bit_vector.push(reader.read_u8(1).unwrap());
    }

    bit_vector
}

// Encode the image with the Vector of Bits
fn encode(c_image: &DynamicImage, bit_vec: &Vec<u8>) {
    let mut tiled_image_vec = tile_image(&c_image);

    let (width, height) = c_image.dimensions();

    let encoded_image = encode_image(&mut tiled_image_vec, bit_vec, width, height);
}

fn tile_image(c_image: &DynamicImage) -> Vec<image::Rgba<u8>> {
    let (width, height) = c_image.dimensions();

    let mut image_blocks = Vec::new();

    for row_index in 0..(height / 8) as u32 {
        for col_index in 0..(width / 8) as u32 {
            for row in 0..8 {
                for column in 0..8 {
                    image_blocks.push(c_image.get_pixel(column + (col_index * 8), row + (row_index * 8)));
                }
            }
        }
    }

    image_blocks
}

fn encode_image(tiled_image: &mut Vec<image::Rgba<u8>>, message: &Vec<u8>, width: u32, height: u32) {
    /*let mut cu = 0.0;
    let mut cv = 0.0;*/
    let mut total = 0.0;
    let mut index: usize = 0;
    let mut count = 0;
    let num_of_iterations = if message.len() / 3 == 0 {
                                    (message.len() / 3) as u32
                                } else {
                                    ((message.len() / 3) + 1) as u32
                                };

    for _ in 0..num_of_iterations {
        for channel in 0..3 {
            index = 36 + (num_of_iterations * 64) as usize;

            let colour_value = tiled_image[index].data[channel];

            // Forward transform
            for y in 0..8 {
                for x in 0..8 {
                    // These are scaling factors, needed in future
                    /*if u == 0 {
                        cu = 0.0
                    } else {
                        cu = 1.0 / 2.0.sqrt()
                    }

                    if v == 0 {
                        cv = 0.0
                    } else {
                        cv = 1.0 / 2.0.sqrt()
                    }*/

                    total = total + (4.0 * f64::consts::PI * ((2.0 * y as f64) + 1.0) / 16.0).cos() *
                                    (4.0 * f64::consts::PI * ((2.0 * x as f64) + 1.0) / 16.0).cos() *
                                    colour_value as f64;
                }
            }

            // 0.25 is another scaling factor
            total = total * 0.25;

            // Encode the message
            if total.trunc() % 2.0 == 0.0 {
                if message[count] == 1 {
                    total = total + 1.0
                }
            }
            else {
                if message[count] == 0 {
                    total = total + 1.0
                }
            }

            // Inverse transform
            for y in 0..8 {
                for x in 0..8 {
                    // These are scaling factors, needed in future
                    /*if u == 0 {
                        cu = 0.0
                    } else {
                        cu = 1.0 / 2.0.sqrt()
                    }

                    if v == 0 {
                        cv = 0.0
                    } else {
                        cv = 1.0 / 2.0.sqrt()
                    }*/

                    total = total + (4.0 * f64::consts::PI * ((2.0 * y as f64) + 1.0) / 16.0).cos() *
                        (4.0 * f64::consts::PI * ((2.0 * x as f64) + 1.0) / 16.0).cos() *
                        colour_value as f64;
                }
            }
        }
    }
}

/**************************************************************************************************
 *                                                                                *
 *                                         Tests                                                  *
 *                                                                                *
 **************************************************************************************************/

#[cfg(test)]
#[test]
fn test_get_bit_vec_len() {
    let bits = get_bit_vec("h");

    assert_eq!(8, bits.len());
}

#[cfg(test)]
#[test]
fn test_get_bit_vec() {
    let test_bits = vec![0, 1, 1, 0, 1, 0, 0, 0];
    let bits = get_bit_vec("h");

    assert_eq!(bits, test_bits);
}

#[cfg(test)]
#[test]
fn test_encode_red_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[0] % 2, bit_vec[0]);
}

#[cfg(test)]
#[test]
fn test_encode_green_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[1] % 2, bit_vec[1]);
}

#[cfg(test)]
#[test]
fn test_encode_blue_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[2] % 2 as u8, bit_vec[2]);
}

#[cfg(test)]
#[test]
fn test_full_byte_is_encoded() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let mut encoded_bit_vec = Vec::new();
    let mut count = 0;

    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    'outer: for x_co_ord in 0..3 {
        let pixel = img.get_pixel(x_co_ord, 0);

        for channel in 0..3 {
            if count >= 8 {
                break 'outer;
            }

            encoded_bit_vec.push(pixel.data[channel] % 2);

            count += 1;
        }
    }

    assert_eq!(encoded_bit_vec, bit_vec);
}

#[cfg(test)]
#[test]
fn test_tile_image_length() {
    let img = image::open(&Path::new("/home/hugh/Pictures/colour.jpg")).unwrap();

    let vec = tile_image(&img);

    assert_eq!(vec.len(), 91200);
}

#[cfg(test)]
#[test]
fn test_tile_image_first_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(0, 0);

    let vec = tile_image(&img);

    assert_eq!(vec[0], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel1() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel2() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(40, 16);

    let vec = tile_image(&img);

    assert_eq!(vec[6720], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel3() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_last_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(399, 399);

    let vec = tile_image(&img);

    assert_eq!(vec[159999], pixel);
}
<commit_msg>Fix up for DCT<commit_after>use std::path::Path;
use std::fs::File;
use std::f64;

use image::{ DynamicImage, GenericImage, ImageFormat };
use image;

use bitreader::BitReader;

use num_traits::float::Float;

pub struct CoverImage {
    cover_image: DynamicImage
}

impl CoverImage {
    pub fn new(file_path: &str) -> CoverImage {
        CoverImage {
            cover_image: image::open(&Path::new(&file_path)).unwrap()
        }
    }

    pub fn encode_with(&self, message: &str) {
        let bit_vector = get_bit_vec(message);

        encode(&self.cover_image, &bit_vector);

        let ref mut fout = File::create(&Path::new("/home/hugh/Pictures/yurt.png")).unwrap();

        let _ = self.cover_image.save(fout, ImageFormat::PNG).unwrap();
    }
}

// Get a Vector of Bits to Encode the Image with
fn get_bit_vec(message: &str) -> Vec<u8> {
    let mut bit_vector = Vec::new();
    let mut reader = BitReader::new(&message.as_bytes());

    // Multiplied by 8 because it's a Vec of bits not bytes
    for _ in 0..message.len() * 8 {
        bit_vector.push(reader.read_u8(1).unwrap());
    }

    bit_vector
}

// Encode the image with the Vector of Bits
fn encode(c_image: &DynamicImage, bit_vec: &Vec<u8>) {
    let mut tiled_image_vec = tile_image(&c_image);

    let (width, height) = c_image.dimensions();

    let encoded_image = encode_image(&mut tiled_image_vec, bit_vec, width, height);
}

fn tile_image(c_image: &DynamicImage) -> Vec<image::Rgba<u8>> {
    let (width, height) = c_image.dimensions();

    let mut image_blocks = Vec::new();

    for row_index in 0..(height / 8) as u32 {
        for col_index in 0..(width / 8) as u32 {
            for row in 0..8 {
                for column in 0..8 {
                    image_blocks.push(c_image.get_pixel(column + (col_index * 8), row + (row_index * 8)));
                }
            }
        }
    }

    image_blocks
}

fn encode_image(tiled_image: &mut Vec<image::Rgba<u8>>, message: &Vec<u8>, width: u32, height: u32) {
    let mut cu = 0.0;
    let mut cv = 0.0;
    let mut total = 0.0;
    let mut inverse_total = 0.0;
    let mut index: usize = 0;
    let mut count = 0;
    let mut rgb_coeffs = Vec::new();

    let num_of_iterations = if message.len() / 3 == 0 {
                                    (message.len() / 3) as u32
                                } else {
                                    ((message.len() / 3) + 1) as u32
                                };

    for _ in 0..(num_of_iterations * 64) {
        for channel in 0..3 {
            for v in 0..8 {
                for u in 0..8 {

                    // Forward transform
                    for y in 0..8 {
                        for x in 0..8 {
                            index = u + (v * 8) + (num_of_iterations * 64) as usize;

                            let colour_value = tiled_image[index].data[channel];

                            if u == 0 {
                                cu = 1.0 / 2.0.sqrt()
                            } else {
                                cu = 1.0
                            }

                            if v == 0 {
                                cv = 1.0 / 2.0.sqrt()
                            } else {
                                cv = 0.0
                            }

                            total = total + (4.0 * f64::consts::PI * ((2.0 * y as f64) + 1.0) / 16.0).cos() *
                                            (4.0 * f64::consts::PI * ((2.0 * x as f64) + 1.0) / 16.0).cos() *
                                            colour_value as f64;
                        }
                    }

                    // 0.25, Cu and Cv are scaling factors
                    total = total * 0.25 * cu * cv;

                    // Encode the message
                    if total.trunc() % 2.0 == 0.0 {
                        if message[count] == 1 {
                            total = total + 1.0
                        }
                    } else {
                        if message[count] == 0 {
                            total = total + 1.0
                        }
                    }

                    /*// Inverse transform
                    for v in 0..8 {
                        for u in 0..8 {
                            if u == 0 {
                                cu = 0.0
                            } else {
                                cu = 1.0 / 2.0.sqrt()
                            }

                            if v == 0 {
                                cv = 0.0
                            } else {
                                cv = 1.0 / 2.0.sqrt()
                            }

                            inverse_total = inverse_total + (v * f64::consts::PI * ((2.0 * 4.0) + 1.0) / 16.0).cos() *
                                                            (u * f64::consts::PI * ((2.0 * 4.0) + 1.0) / 16.0).cos() *
                                                            colour_value as f64;
                        }
                    }*/
                }
            }
        }
    }
}

/**************************************************************************************************
 *                                                                                *
 *                                         Tests                                                  *
 *                                                                                *
 **************************************************************************************************/

#[cfg(test)]
#[test]
fn test_get_bit_vec_len() {
    let bits = get_bit_vec("h");

    assert_eq!(8, bits.len());
}

#[cfg(test)]
#[test]
fn test_get_bit_vec() {
    let test_bits = vec![0, 1, 1, 0, 1, 0, 0, 0];
    let bits = get_bit_vec("h");

    assert_eq!(bits, test_bits);
}

#[cfg(test)]
#[test]
fn test_encode_red_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[0] % 2, bit_vec[0]);
}

#[cfg(test)]
#[test]
fn test_encode_green_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[1] % 2, bit_vec[1]);
}

#[cfg(test)]
#[test]
fn test_encode_blue_channel_lsb_set() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    let pixel = img.get_pixel(0, 0);

    assert_eq!(pixel.data[2] % 2 as u8, bit_vec[2]);
}

#[cfg(test)]
#[test]
fn test_full_byte_is_encoded() {
    let mut img = image::open(&Path::new("/home/hugh/Pictures/scenery.jpg")).unwrap();
    let mut encoded_bit_vec = Vec::new();
    let mut count = 0;

    let bit_vec = vec![0, 1, 1, 0, 1, 0, 0, 0];

    encode(&mut img, &bit_vec);

    'outer: for x_co_ord in 0..3 {
        let pixel = img.get_pixel(x_co_ord, 0);

        for channel in 0..3 {
            if count >= 8 {
                break 'outer;
            }

            encoded_bit_vec.push(pixel.data[channel] % 2);

            count += 1;
        }
    }

    assert_eq!(encoded_bit_vec, bit_vec);
}

#[cfg(test)]
#[test]
fn test_tile_image_length() {
    let img = image::open(&Path::new("/home/hugh/Pictures/colour.jpg")).unwrap();

    let vec = tile_image(&img);

    assert_eq!(vec.len(), 91200);
}

#[cfg(test)]
#[test]
fn test_tile_image_first_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(0, 0);

    let vec = tile_image(&img);

    assert_eq!(vec[0], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel1() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel2() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(40, 16);

    let vec = tile_image(&img);

    assert_eq!(vec[6720], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_random_pixel3() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(15, 7);

    let vec = tile_image(&img);

    assert_eq!(vec[127], pixel);
}

#[cfg(test)]
#[test]
fn test_tile_image_last_pixel() {
    let img = image::open(&Path::new("/home/hugh/Pictures/project.jpg")).unwrap();

    let pixel = img.get_pixel(399, 399);

    let vec = tile_image(&img);

    assert_eq!(vec[159999], pixel);
}
<|endoftext|>