<commit_before>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It is also responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//!
//! TODO: Some LSP clients support watching the disk, so this crate should to
//! support custom watcher events (related to
//! <https://github.com/rust-analyzer/rust-analyzer/issues/131>)
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! which are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod roots;
mod io;

use std::{
    fmt, fs, mem,
    path::{Path, PathBuf},
    sync::Arc,
};

use crossbeam_channel::Receiver;
use ra_arena::{impl_arena_id, Arena, RawId, map::ArenaMap};
use relative_path::{RelativePath, RelativePathBuf};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{
    io::{TaskResult, Worker},
    roots::Roots,
};

pub use crate::{
    io::TaskResult as VfsTask,
    roots::VfsRoot,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct VfsFile(pub RawId);
impl_arena_id!(VfsFile);

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    is_overlayed: bool,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arc<Roots>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: ArenaMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: Worker,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Vfs")
            .field("n_roots", &self.roots.len())
            .field("n_files", &self.files.len())
            .field("n_pending_changes", &self.pending_changes.len())
            .finish()
    }
}

impl Vfs {
    pub fn new(roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let roots = Arc::new(Roots::new(roots));
        let worker = io::start(Arc::clone(&roots));
        let mut root2files = ArenaMap::default();

        for root in roots.iter() {
            root2files.insert(root, Default::default());
            worker.sender().send(io::Task::AddRoot { root }).unwrap();
        }
        let res =
            Vfs { roots, files: Arena::default(), root2files, worker, pending_changes: Vec::new() };
        let vfs_roots = res.roots.iter().collect();
        (res, vfs_roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots.path(root).to_path_buf()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots.path(self.files[file].root);
        rel_path.to_path(root_path)
    }

    pub fn n_roots(&self) -> usize {
        self.roots.len()
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text), false);
                let change = VfsChange::AddFile { file, text, root, path: rel_path };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        if let Some(file) = file {
            self.do_change_file(file, text, true);
            Some(file)
        } else {
            self.do_add_file(root, rel_path, text, true)
        }
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            self.do_change_file(file, new_text, true);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        let file = file.expect("can't remove a file which wasn't added");
        let full_path = rel_path.to_path(&self.roots.path(root));
        if let Ok(text) = fs::read_to_string(&full_path) {
            self.do_change_file(file, text, false);
        } else {
            self.do_remove_file(root, rel_path, file);
        }
        Some(file)
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        self.worker.receiver()
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        match task {
            TaskResult::BulkLoadRoot { root, files } => {
                let mut cur_files = Vec::new();
                // While we were scanning the root in the background, a file might have
                // been open in the editor, so we need to account for that.
                let existing = self.root2files[root]
                    .iter()
                    .map(|&file| (self.files[file].path.clone(), file))
                    .collect::<FxHashMap<_, _>>();
                for (path, text) in files {
                    if let Some(&file) = existing.get(&path) {
                        let text = Arc::clone(&self.files[file].text);
                        cur_files.push((file, path, text));
                        continue;
                    }
                    let text = Arc::new(text);
                    let file = self.add_file(root, path.clone(), Arc::clone(&text), false);
                    cur_files.push((file, path, text));
                }

                let change = VfsChange::AddRoot { root, files: cur_files };
                self.pending_changes.push(change);
            }
            TaskResult::SingleFile { root, path, text } => {
                let existing_file = self.find_file(root, &path);
                if existing_file.map(|file| self.files[file].is_overlayed) == Some(true) {
                    return;
                }
                match (existing_file, text) {
                    (Some(file), None) => {
                        self.do_remove_file(root, path, file);
                    }
                    (None, Some(text)) => {
                        self.do_add_file(root, path, text, false);
                    }
                    (Some(file), Some(text)) => {
                        self.do_change_file(file, text, false);
                    }
                    (None, None) => (),
                }
            }
        }
    }

    fn do_add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
        is_overlay: bool,
    ) -> Option<VfsFile> {
        let text = Arc::new(text);
        let file = self.add_file(root, path.clone(), text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::AddFile { file, root, path, text });
        Some(file)
    }

    fn do_change_file(&mut self, file: VfsFile, text: String, is_overlay: bool) {
        let text = Arc::new(text);
        self.change_file(file, text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::ChangeFile { file, text });
    }

    fn do_remove_file(&mut self, root: VfsRoot, path: RelativePathBuf, file: VfsFile) {
        self.remove_file(file);
        self.pending_changes.push(VfsChange::RemoveFile { root, path, file });
    }

    fn add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: Arc<String>,
        is_overlayed: bool,
    ) -> VfsFile {
        let data = VfsFileData { root, path, text, is_overlayed };
        let file = self.files.alloc(data);
        self.root2files.get_mut(root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>, is_overlayed: bool) {
        let mut file_data = &mut self.files[file];
        file_data.text = new_text;
        file_data.is_overlayed = is_overlayed;
    }

    fn remove_file(&mut self, file: VfsFile) {
        // FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self.roots.find(&path)?;
        let file = self.find_file(root, &path);
        Some((root, path, file))
    }

    fn find_file(&self, root: VfsRoot, path: &RelativePath) -> Option<VfsFile> {
        self.root2files[root].iter().map(|&it| it).find(|&file| self.files[file].path == path)
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot { root: VfsRoot, files: Vec<(VfsFile, RelativePathBuf, Arc<String>)> },
    AddFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf, text: Arc<String> },
    RemoveFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf },
    ChangeFile { file: VfsFile, text: Arc<String> },
}
<commit_msg>marginally better names<commit_after>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It is also responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//!
//! TODO: Some LSP clients support watching the disk, so this crate should to
//! support custom watcher events (related to
//! <https://github.com/rust-analyzer/rust-analyzer/issues/131>)
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! which are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod roots;
mod io;

use std::{
    fmt, fs, mem,
    path::{Path, PathBuf},
    sync::Arc,
};

use crossbeam_channel::Receiver;
use ra_arena::{impl_arena_id, Arena, RawId, map::ArenaMap};
use relative_path::{RelativePath, RelativePathBuf};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{
    io::{TaskResult, Worker},
    roots::Roots,
};

pub use crate::{
    io::TaskResult as VfsTask,
    roots::VfsRoot,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct VfsFile(pub RawId);
impl_arena_id!(VfsFile);

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    is_overlayed: bool,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arc<Roots>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: ArenaMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: Worker,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Vfs")
            .field("n_roots", &self.roots.len())
            .field("n_files", &self.files.len())
            .field("n_pending_changes", &self.pending_changes.len())
            .finish()
    }
}

impl Vfs {
    pub fn new(roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let roots = Arc::new(Roots::new(roots));
        let worker = io::start(Arc::clone(&roots));
        let mut root2files = ArenaMap::default();

        for root in roots.iter() {
            root2files.insert(root, Default::default());
            worker.sender().send(io::Task::AddRoot { root }).unwrap();
        }
        let res =
            Vfs { roots, files: Arena::default(), root2files, worker, pending_changes: Vec::new() };
        let vfs_roots = res.roots.iter().collect();
        (res, vfs_roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots.path(root).to_path_buf()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots.path(self.files[file].root);
        rel_path.to_path(root_path)
    }

    pub fn n_roots(&self) -> usize {
        self.roots.len()
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.raw_add_file(root, rel_path.clone(), Arc::clone(&text), false);
                let change = VfsChange::AddFile { file, text, root, path: rel_path };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        if let Some(file) = file {
            self.change_file_event(file, text, true);
            Some(file)
        } else {
            self.add_file_event(root, rel_path, text, true)
        }
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            self.change_file_event(file, new_text, true);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let (root, rel_path, file) = self.find_root(path)?;
        let file = file.expect("can't remove a file which wasn't added");
        let full_path = rel_path.to_path(&self.roots.path(root));
        if let Ok(text) = fs::read_to_string(&full_path) {
            self.change_file_event(file, text, false);
        } else {
            self.remove_file_event(root, rel_path, file);
        }
        Some(file)
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        self.worker.receiver()
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        match task {
            TaskResult::BulkLoadRoot { root, files } => {
                let mut cur_files = Vec::new();
                // While we were scanning the root in the background, a file might have
                // been open in the editor, so we need to account for that.
                let existing = self.root2files[root]
                    .iter()
                    .map(|&file| (self.files[file].path.clone(), file))
                    .collect::<FxHashMap<_, _>>();
                for (path, text) in files {
                    if let Some(&file) = existing.get(&path) {
                        let text = Arc::clone(&self.files[file].text);
                        cur_files.push((file, path, text));
                        continue;
                    }
                    let text = Arc::new(text);
                    let file = self.raw_add_file(root, path.clone(), Arc::clone(&text), false);
                    cur_files.push((file, path, text));
                }

                let change = VfsChange::AddRoot { root, files: cur_files };
                self.pending_changes.push(change);
            }
            TaskResult::SingleFile { root, path, text } => {
                let existing_file = self.find_file(root, &path);
                if existing_file.map(|file| self.files[file].is_overlayed) == Some(true) {
                    return;
                }
                match (existing_file, text) {
                    (Some(file), None) => {
                        self.remove_file_event(root, path, file);
                    }
                    (None, Some(text)) => {
                        self.add_file_event(root, path, text, false);
                    }
                    (Some(file), Some(text)) => {
                        self.change_file_event(file, text, false);
                    }
                    (None, None) => (),
                }
            }
        }
    }

    // *_event calls change the state of VFS and push a change onto pending
    // changes array.

    fn add_file_event(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
        is_overlay: bool,
    ) -> Option<VfsFile> {
        let text = Arc::new(text);
        let file = self.raw_add_file(root, path.clone(), text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::AddFile { file, root, path, text });
        Some(file)
    }

    fn change_file_event(&mut self, file: VfsFile, text: String, is_overlay: bool) {
        let text = Arc::new(text);
        self.raw_change_file(file, text.clone(), is_overlay);
        self.pending_changes.push(VfsChange::ChangeFile { file, text });
    }

    fn remove_file_event(&mut self, root: VfsRoot, path: RelativePathBuf, file: VfsFile) {
        self.raw_remove_file(file);
        self.pending_changes.push(VfsChange::RemoveFile { root, path, file });
    }

    // raw_* calls change the state of VFS, but **do not** emit events.

    fn raw_add_file(
        &mut self,
        root: VfsRoot,
        path: RelativePathBuf,
        text: Arc<String>,
        is_overlayed: bool,
    ) -> VfsFile {
        let data = VfsFileData { root, path, text, is_overlayed };
        let file = self.files.alloc(data);
        self.root2files.get_mut(root).unwrap().insert(file);
        file
    }

    fn raw_change_file(&mut self, file: VfsFile, new_text: Arc<String>, is_overlayed: bool) {
        let mut file_data = &mut self.files[file];
        file_data.text = new_text;
        file_data.is_overlayed = is_overlayed;
    }

    fn raw_remove_file(&mut self, file: VfsFile) {
        // FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self.roots.find(&path)?;
        let file = self.find_file(root, &path);
        Some((root, path, file))
    }

    fn find_file(&self, root: VfsRoot, path: &RelativePath) -> Option<VfsFile> {
        self.root2files[root].iter().map(|&it| it).find(|&file| self.files[file].path == path)
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot { root: VfsRoot, files: Vec<(VfsFile, RelativePathBuf, Arc<String>)> },
    AddFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf, text: Arc<String> },
    RemoveFile { root: VfsRoot, file: VfsFile, path: RelativePathBuf },
    ChangeFile { file: VfsFile, text: Arc<String> },
}
<|endoftext|>


<commit_before>use std::task::{Context, Poll};
use std::{cmp, future::Future, marker::PhantomData, net, pin::Pin, rc::Rc};

use actix_codec::{AsyncRead, AsyncWrite};
use actix_service::Service;
use bytes::{Bytes, BytesMut};
use futures_core::ready;
use h2::{
    server::{Connection, SendResponse},
    SendStream,
};
use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING};
use log::{error, trace};

use crate::body::{Body, BodySize, MessageBody};
use crate::config::ServiceConfig;
use crate::error::{DispatchError, Error};
use crate::message::ResponseHead;
use crate::payload::Payload;
use crate::request::Request;
use crate::response::Response;
use crate::service::HttpFlow;
use crate::OnConnectData;

const CHUNK_SIZE: usize = 16_384;

/// Dispatcher for HTTP/2 protocol.
#[pin_project::pin_project]
pub struct Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    B: MessageBody,
{
    flow: Rc<HttpFlow<S, X, U>>,
    connection: Connection<T, Bytes>,
    on_connect_data: OnConnectData,
    config: ServiceConfig,
    peer_addr: Option<net::SocketAddr>,
    _phantom: PhantomData<B>,
}

impl<T, S, B, X, U> Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,
    S: Service<Request>,
    S::Error: Into<Error>,
    S::Response: Into<Response<B>>,
    B: MessageBody,
{
    pub(crate) fn new(
        flow: Rc<HttpFlow<S, X, U>>,
        connection: Connection<T, Bytes>,
        on_connect_data: OnConnectData,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
    ) -> Self {
        Dispatcher {
            flow,
            config,
            peer_addr,
            connection,
            on_connect_data,
            _phantom: PhantomData,
        }
    }
}

impl<T, S, B, X, U> Future for Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,

    S: Service<Request>,
    S::Error: Into<Error> + 'static,
    S::Future: 'static,
    S::Response: Into<Response<B>> + 'static,

    B: MessageBody + 'static,
    B::Error: Into<Error>,
{
    type Output = Result<(), DispatchError>;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        loop {
            match ready!(Pin::new(&mut this.connection).poll_accept(cx)) {
                None => return Poll::Ready(Ok(())),

                Some(Err(err)) => return Poll::Ready(Err(err.into())),

                Some(Ok((req, res))) => {
                    let (parts, body) = req.into_parts();
                    let pl = crate::h2::Payload::new(body);
                    let pl = Payload::<crate::payload::PayloadStream>::H2(pl);
                    let mut req = Request::with_payload(pl);

                    let head = req.head_mut();
                    head.uri = parts.uri;
                    head.method = parts.method;
                    head.version = parts.version;
                    head.headers = parts.headers.into();
                    head.peer_addr = this.peer_addr;

                    // merge on_connect_ext data into request extensions
                    this.on_connect_data.merge_into(&mut req);

                    let svc = ServiceResponse {
                        state: ServiceResponseState::ServiceCall(
                            this.flow.service.call(req),
                            Some(res),
                        ),
                        config: this.config.clone(),
                        buffer: None,
                        _phantom: PhantomData,
                    };

                    actix_rt::spawn(svc);
                }
            }
        }
    }
}

#[pin_project::pin_project]
struct ServiceResponse<F, I, E, B> {
    #[pin]
    state: ServiceResponseState<F, B>,
    config: ServiceConfig,
    buffer: Option<Bytes>,
    _phantom: PhantomData<(I, E)>,
}

#[pin_project::pin_project(project = ServiceResponseStateProj)]
enum ServiceResponseState<F, B> {
    ServiceCall(#[pin] F, Option<SendResponse<Bytes>>),
    SendPayload(SendStream<Bytes>, #[pin] B),
    SendErrorPayload(SendStream<Bytes>, #[pin] Body),
}

impl<F, I, E, B> ServiceResponse<F, I, E, B>
where
    F: Future<Output = Result<I, E>>,
    E: Into<Error>,
    I: Into<Response<B>>,

    B: MessageBody,
    B::Error: Into<Error>,
{
    fn prepare_response(
        &self,
        head: &ResponseHead,
        size: &mut BodySize,
    ) -> http::Response<()> {
        let mut has_date = false;
        let mut skip_len = size != &BodySize::Stream;

        let mut res = http::Response::new(());
        *res.status_mut() = head.status;
        *res.version_mut() = http::Version::HTTP_2;

        // Content length
        match head.status {
            http::StatusCode::NO_CONTENT
            | http::StatusCode::CONTINUE
            | http::StatusCode::PROCESSING => *size = BodySize::None,
            http::StatusCode::SWITCHING_PROTOCOLS => {
                skip_len = true;
                *size = BodySize::Stream;
            }
            _ => {}
        }

        let _ = match size {
            BodySize::None | BodySize::Stream => None,
            BodySize::Empty => res
                .headers_mut()
                .insert(CONTENT_LENGTH, HeaderValue::from_static("0")),
            BodySize::Sized(len) => {
                let mut buf = itoa::Buffer::new();

                res.headers_mut().insert(
                    CONTENT_LENGTH,
                    HeaderValue::from_str(buf.format(*len)).unwrap(),
                )
            }
        };

        // copy headers
        for (key, value) in head.headers.iter() {
            match *key {
                // TODO: consider skipping other headers according to:
                //       https://tools.ietf.org/html/rfc7540#section-8.1.2.2
                // omit HTTP/1.x only headers
                CONNECTION | TRANSFER_ENCODING => continue,
                CONTENT_LENGTH if skip_len => continue,
                DATE => has_date = true,
                _ => {}
            }

            res.headers_mut().append(key, value.clone());
        }

        // set date header
        if !has_date {
            let mut bytes = BytesMut::with_capacity(29);
            self.config.set_date_header(&mut bytes);
            res.headers_mut().insert(
                DATE,
                // SAFETY: serialized date-times are known ASCII strings
                unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) },
            );
        }

        res
    }
}

impl<F, I, E, B> Future for ServiceResponse<F, I, E, B>
where
    F: Future<Output = Result<I, E>>,
    E: Into<Error>,
    I: Into<Response<B>>,

    B: MessageBody,
    B::Error: Into<Error>,
{
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.as_mut().project();

        match this.state.project() {
            ServiceResponseStateProj::ServiceCall(call, send) => {
                match ready!(call.poll(cx)) {
                    Ok(res) => {
                        let (res, body) = res.into().replace_body(());

                        let mut send = send.take().unwrap();
                        let mut size = body.size();
                        let h2_res =
                            self.as_mut().prepare_response(res.head(), &mut size);
                        this = self.as_mut().project();

                        let stream = match send.send_response(h2_res, size.is_eof()) {
                            Err(e) => {
                                trace!("Error sending HTTP/2 response: {:?}", e);
                                return Poll::Ready(());
                            }
                            Ok(stream) => stream,
                        };

                        if size.is_eof() {
                            Poll::Ready(())
                        } else {
                            this.state
                                .set(ServiceResponseState::SendPayload(stream, body));
                            self.poll(cx)
                        }
                    }

                    Err(err) => {
                        let res = Response::from_error(err.into());
                        let (res, body) = res.replace_body(());

                        let mut send = send.take().unwrap();
                        let mut size = body.size();
                        let h2_res =
                            self.as_mut().prepare_response(res.head(), &mut size);
                        this = self.as_mut().project();

                        let stream = match send.send_response(h2_res, size.is_eof()) {
                            Err(e) => {
                                trace!("Error sending HTTP/2 response: {:?}", e);
                                return Poll::Ready(());
                            }
                            Ok(stream) => stream,
                        };

                        if size.is_eof() {
                            Poll::Ready(())
                        } else {
                            this.state.set(ServiceResponseState::SendErrorPayload(
                                stream, body,
                            ));
                            self.poll(cx)
                        }
                    }
                }
            }

            ServiceResponseStateProj::SendPayload(ref mut stream, ref mut body) => {
                loop {
                    match this.buffer {
                        Some(ref mut buffer) => match ready!(stream.poll_capacity(cx)) {
                            None => return Poll::Ready(()),

                            Some(Ok(cap)) => {
                                let len = buffer.len();
                                let bytes = buffer.split_to(cmp::min(cap, len));

                                if let Err(e) = stream.send_data(bytes, false) {
                                    warn!("{:?}", e);
                                    return Poll::Ready(());
                                } else if !buffer.is_empty() {
                                    let cap = cmp::min(buffer.len(), CHUNK_SIZE);
                                    stream.reserve_capacity(cap);
                                } else {
                                    this.buffer.take();
                                }
                            }

                            Some(Err(e)) => {
                                warn!("{:?}", e);
                                return Poll::Ready(());
                            }
                        },

                        None => match ready!(body.as_mut().poll_next(cx)) {
                            None => {
                                if let Err(e) = stream.send_data(Bytes::new(), true) {
                                    warn!("{:?}", e);
                                }
                                return Poll::Ready(());
                            }

                            Some(Ok(chunk)) => {
                                stream
                                    .reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));
                                *this.buffer = Some(chunk);
                            }

                            Some(Err(err)) => {
                                error!(
                                    "Response payload stream error: {:?}",
                                    err.into()
                                );

                                return Poll::Ready(());
                            }
                        },
                    }
                }
            }

            ServiceResponseStateProj::SendErrorPayload(ref mut stream, ref mut body) => {
                // TODO: de-dupe impl with SendPayload

                loop {
                    match this.buffer {
                        Some(ref mut buffer) => match ready!(stream.poll_capacity(cx)) {
                            None => return Poll::Ready(()),

                            Some(Ok(cap)) => {
                                let len = buffer.len();
                                let bytes = buffer.split_to(cmp::min(cap, len));

                                if let Err(e) = stream.send_data(bytes, false) {
                                    warn!("{:?}", e);
                                    return Poll::Ready(());
                                } else if !buffer.is_empty() {
                                    let cap = cmp::min(buffer.len(), CHUNK_SIZE);
                                    stream.reserve_capacity(cap);
                                } else {
                                    this.buffer.take();
                                }
                            }

                            Some(Err(e)) => {
                                warn!("{:?}", e);
                                return Poll::Ready(());
                            }
                        },

                        None => match ready!(body.as_mut().poll_next(cx)) {
                            None => {
                                if let Err(e) = stream.send_data(Bytes::new(), true) {
                                    warn!("{:?}", e);
                                }
                                return Poll::Ready(());
                            }

                            Some(Ok(chunk)) => {
                                stream
                                    .reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));
                                *this.buffer = Some(chunk);
                            }

                            Some(Err(err)) => {
                                error!("Response payload stream error: {:?}", err);

                                return Poll::Ready(());
                            }
                        },
                    }
                }
            }
        }
    }
}
<commit_msg>refactor h2 dispatcher to async/await.reduce duplicate code (#2211)<commit_after>use std::{
    cmp,
    future::Future,
    marker::PhantomData,
    net,
    pin::Pin,
    rc::Rc,
    task::{Context, Poll},
};

use actix_codec::{AsyncRead, AsyncWrite};
use actix_service::Service;
use actix_utils::future::poll_fn;
use bytes::{Bytes, BytesMut};
use futures_core::ready;
use h2::server::{Connection, SendResponse};
use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING};
use log::{error, trace};
use pin_project_lite::pin_project;

use crate::body::{BodySize, MessageBody};
use crate::config::ServiceConfig;
use crate::error::Error;
use crate::message::ResponseHead;
use crate::payload::Payload;
use crate::request::Request;
use crate::response::Response;
use crate::service::HttpFlow;
use crate::OnConnectData;

const CHUNK_SIZE: usize = 16_384;

pin_project! {
    /// Dispatcher for HTTP/2 protocol.
    pub struct Dispatcher<T, S, B, X, U> {
        flow: Rc<HttpFlow<S, X, U>>,
        connection: Connection<T, Bytes>,
        on_connect_data: OnConnectData,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
        _phantom: PhantomData<B>,
    }
}

impl<T, S, B, X, U> Dispatcher<T, S, B, X, U> {
    pub(crate) fn new(
        flow: Rc<HttpFlow<S, X, U>>,
        connection: Connection<T, Bytes>,
        on_connect_data: OnConnectData,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
    ) -> Self {
        Self {
            flow,
            config,
            peer_addr,
            connection,
            on_connect_data,
            _phantom: PhantomData,
        }
    }
}

impl<T, S, B, X, U> Future for Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,

    S: Service<Request>,
    S::Error: Into<Error>,
    S::Future: 'static,
    S::Response: Into<Response<B>>,

    B: MessageBody,
    B::Error: Into<Error>,
{
    type Output = Result<(), crate::error::DispatchError>;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        while let Some((req, tx)) =
            ready!(Pin::new(&mut this.connection).poll_accept(cx)?)
        {
            let (parts, body) = req.into_parts();
            let pl = crate::h2::Payload::new(body);
            let pl = Payload::<crate::payload::PayloadStream>::H2(pl);
            let mut req = Request::with_payload(pl);

            let head = req.head_mut();
            head.uri = parts.uri;
            head.method = parts.method;
            head.version = parts.version;
            head.headers = parts.headers.into();
            head.peer_addr = this.peer_addr;

            // merge on_connect_ext data into request extensions
            this.on_connect_data.merge_into(&mut req);

            let fut = this.flow.service.call(req);
            let config = this.config.clone();

            // multiplex request handling with spawn task
            actix_rt::spawn(async move {
                // resolve service call and send response.
                let res = match fut.await {
                    Ok(res) => handle_response(res.into(), tx, config).await,
                    Err(err) => {
                        let res = Response::from_error(err.into());
                        handle_response(res, tx, config).await
                    }
                };

                // log error.
                if let Err(err) = res {
                    match err {
                        DispatchError::SendResponse(err) => {
                            trace!("Error sending HTTP/2 response: {:?}", err)
                        }
                        DispatchError::SendData(err) => warn!("{:?}", err),
                        DispatchError::ResponseBody(err) => {
                            error!("Response payload stream error: {:?}", err)
                        }
                    }
                }
            });
        }

        Poll::Ready(Ok(()))
    }
}

enum DispatchError {
    SendResponse(h2::Error),
    SendData(h2::Error),
    ResponseBody(Error),
}

async fn handle_response<B>(
    res: Response<B>,
    mut tx: SendResponse<Bytes>,
    config: ServiceConfig,
) -> Result<(), DispatchError>
where
    B: MessageBody,
    B::Error: Into<Error>,
{
    let (res, body) = res.replace_body(());

    // prepare response.
    let mut size = body.size();
    let res = prepare_response(config, res.head(), &mut size);
    let eof = size.is_eof();

    // send response head and return on eof.
    let mut stream = tx
        .send_response(res, eof)
        .map_err(DispatchError::SendResponse)?;

    if eof {
        return Ok(());
    }

    // poll response body and send chunks to client.
    actix_rt::pin!(body);

    while let Some(res) = poll_fn(|cx| body.as_mut().poll_next(cx)).await {
        let mut chunk = res.map_err(|err| DispatchError::ResponseBody(err.into()))?;

        'send: loop {
            // reserve enough space and wait for stream ready.
            stream.reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));

            match poll_fn(|cx| stream.poll_capacity(cx)).await {
                // No capacity left. drop body and return.
                None => return Ok(()),
                Some(res) => {
                    // Split chuck to writeable size and send to client.
                    let cap = res.map_err(DispatchError::SendData)?;

                    let len = chunk.len();
                    let bytes = chunk.split_to(cmp::min(cap, len));

                    stream
                        .send_data(bytes, false)
                        .map_err(DispatchError::SendData)?;

                    // Current chuck completely sent. break send loop and poll next one.
                    if chunk.is_empty() {
                        break 'send;
                    }
                }
            }
        }
    }

    // response body streaming finished. send end of stream and return.
    stream
        .send_data(Bytes::new(), true)
        .map_err(DispatchError::SendData)?;

    Ok(())
}

fn prepare_response(
    config: ServiceConfig,
    head: &ResponseHead,
    size: &mut BodySize,
) -> http::Response<()> {
    let mut has_date = false;
    let mut skip_len = size != &BodySize::Stream;

    let mut res = http::Response::new(());
    *res.status_mut() = head.status;
    *res.version_mut() = http::Version::HTTP_2;

    // Content length
    match head.status {
        http::StatusCode::NO_CONTENT
        | http::StatusCode::CONTINUE
        | http::StatusCode::PROCESSING => *size = BodySize::None,
        http::StatusCode::SWITCHING_PROTOCOLS => {
            skip_len = true;
            *size = BodySize::Stream;
        }
        _ => {}
    }

    let _ = match size {
        BodySize::None | BodySize::Stream => None,
        BodySize::Empty => res
            .headers_mut()
            .insert(CONTENT_LENGTH, HeaderValue::from_static("0")),
        BodySize::Sized(len) => {
            let mut buf = itoa::Buffer::new();

            res.headers_mut().insert(
                CONTENT_LENGTH,
                HeaderValue::from_str(buf.format(*len)).unwrap(),
            )
        }
    };

    // copy headers
    for (key, value) in head.headers.iter() {
        match *key {
            // TODO: consider skipping other headers according to:
            //       https://tools.ietf.org/html/rfc7540#section-8.1.2.2
            // omit HTTP/1.x only headers
            CONNECTION | TRANSFER_ENCODING => continue,
            CONTENT_LENGTH if skip_len => continue,
            DATE => has_date = true,
            _ => {}
        }

        res.headers_mut().append(key, value.clone());
    }

    // set date header
    if !has_date {
        let mut bytes = BytesMut::with_capacity(29);
        config.set_date_header(&mut bytes);
        res.headers_mut().insert(
            DATE,
            // SAFETY: serialized date-times are known ASCII strings
            unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) },
        );
    }

    res
}
<|endoftext|>


<commit_before>use std::char;

use common::{is_name_start_char, is_name_char, is_whitespace_char};

use reader::lexer::Token;

use super::{Result, PullParser, State};

impl PullParser {
    pub fn inside_reference(&mut self, t: Token, prev_st: State) -> Option<Result> {
        match t {
            Token::Character(c) if !self.data.ref_data.is_empty() && is_name_char(c) ||
                             self.data.ref_data.is_empty() && (is_name_start_char(c) || c == '#') => {
                self.data.ref_data.push(c);
                None
            }

            Token::ReferenceEnd => {
                // TODO: check for unicode correctness
                let name = self.data.take_ref_data();
                let name_len = name.len();  // compute once
                let c = match &name[..] {
                    "lt"   => Ok('<'),
                    "gt"   => Ok('>'),
                    "amp"  => Ok('&'),
                    "apos" => Ok('\''),
                    "quot" => Ok('"'),
                    ""     => Err(self_error!(self; "Encountered empty entity")),
                    _ if name_len > 2 && name.starts_with("#x") => {
                        let num_str = &name[2..name_len];
                        if num_str == "0" {
                            Err(self_error!(self; "Null character entity is not allowed"))
                        } else {
                            match u32::from_str_radix(num_str, 16).ok().and_then(char::from_u32) {
                                Some(c) => Ok(c),
                                None    => Err(self_error!(self; "Invalid hexadecimal character number in an entity: {}", name))
                            }
                        }
                    }
                    _ if name_len > 1 && name.starts_with('#') => {
                        let num_str = &name[1..name_len];
                        if num_str == "0" {
                            Err(self_error!(self; "Null character entity is not allowed"))
                        } else {
                            match u32::from_str_radix(num_str, 10).ok().and_then(char::from_u32) {
                                Some(c) => Ok(c),
                                None    => Err(self_error!(self; "Invalid decimal character number in an entity: {}", name))
                            }
                        }
                    },
                    _ => {
                        if let Some(v) = self.config.extra_entities.get(&name) {
                            Ok(v.clone())
                        } else {
                            Err(self_error!(self; "Unexpected entity: {}", name))
                        }
                    }
                };
                match c {
                    Ok(c) => {
                        self.buf.push(c);
                        if !is_whitespace_char(c) {
                            self.inside_whitespace = false;
                        }
                        self.into_state_continue(prev_st)
                    }
                    Err(e) => Some(e)
                }
            }

            _ => Some(self_error!(self; "Unexpected token inside an entity: {}", t))
        }
    }
}
<commit_msg>Fixed incomplete fix for parsing entities as whitespace<commit_after>use std::char;

use common::{is_name_start_char, is_name_char, is_whitespace_char};

use reader::lexer::Token;

use super::{Result, PullParser, State};

impl PullParser {
    pub fn inside_reference(&mut self, t: Token, prev_st: State) -> Option<Result> {
        match t {
            Token::Character(c) if !self.data.ref_data.is_empty() && is_name_char(c) ||
                             self.data.ref_data.is_empty() && (is_name_start_char(c) || c == '#') => {
                self.data.ref_data.push(c);
                None
            }

            Token::ReferenceEnd => {
                // TODO: check for unicode correctness
                let name = self.data.take_ref_data();
                let name_len = name.len();  // compute once
                let c = match &name[..] {
                    "lt"   => Ok('<'),
                    "gt"   => Ok('>'),
                    "amp"  => Ok('&'),
                    "apos" => Ok('\''),
                    "quot" => Ok('"'),
                    ""     => Err(self_error!(self; "Encountered empty entity")),
                    _ if name_len > 2 && name.starts_with("#x") => {
                        let num_str = &name[2..name_len];
                        if num_str == "0" {
                            Err(self_error!(self; "Null character entity is not allowed"))
                        } else {
                            match u32::from_str_radix(num_str, 16).ok().and_then(char::from_u32) {
                                Some(c) => Ok(c),
                                None    => Err(self_error!(self; "Invalid hexadecimal character number in an entity: {}", name))
                            }
                        }
                    }
                    _ if name_len > 1 && name.starts_with('#') => {
                        let num_str = &name[1..name_len];
                        if num_str == "0" {
                            Err(self_error!(self; "Null character entity is not allowed"))
                        } else {
                            match u32::from_str_radix(num_str, 10).ok().and_then(char::from_u32) {
                                Some(c) => Ok(c),
                                None    => Err(self_error!(self; "Invalid decimal character number in an entity: {}", name))
                            }
                        }
                    },
                    _ => {
                        if let Some(v) = self.config.extra_entities.get(&name) {
                            Ok(v.clone())
                        } else {
                            Err(self_error!(self; "Unexpected entity: {}", name))
                        }
                    }
                };
                match c {
                    Ok(c) => {
                        self.buf.push(c);
                        if prev_st == State::OutsideTag && !is_whitespace_char(c) {
                            self.inside_whitespace = false;
                        }
                        self.into_state_continue(prev_st)
                    }
                    Err(e) => Some(e)
                }
            }

            _ => Some(self_error!(self; "Unexpected token inside an entity: {}", t))
        }
    }
}
<|endoftext|>


<commit_before>
pub enum NoParams {}
impl ::syn::synom::Synom for NoParams {
    named!(parse -> Self, reject!() );
}

pub enum StrOption {
    Str( ::syn::LitStr ),
    None
}
impl ::syn::synom::Synom for StrOption {
    named!(parse -> Self, alt!(
        syn!( ::syn::LitStr ) => { |s| StrOption::Str( s ) }
        |
        custom_keyword!( None ) => { |_| StrOption::None }
    ) );
}

macro_rules! intercom_attribute {

    ( $attr:ident < $attr_param:ident, $params:ident > { $( $name:ident : $type:ident, )* } ) => {

        enum $attr_param {
            $( $name ( $type ), )*
            args( $params ),
        }

        impl ::syn::synom::Synom for $attr_param {
            named!( parse -> Self, alt!(
                $(
                    do_parse!(
                        custom_keyword!( $name ) >>
                        eq : punct!(=) >>
                        value : syn!( $type ) >>
                        ( $attr_param::$name( value ) )
                    )
                    |
                )*
                do_parse!( a : syn!( $params ) >> ( $attr_param::args( a ) ) )
            ) );
        }

        struct $attr( Vec<$attr_param> );
        impl ::syn::synom::Synom for $attr {
            named!(parse -> Self, do_parse!(
                p: call!( ::syn::punctuated::Punctuated::<$attr_param, ::syn::token::Comma>::parse_separated_nonempty ) >>
                ( $attr( p.into_iter().collect() ) )
            ) );
        }

        impl $attr {
            $(
                pub fn $name( &self ) -> Result< Option< &$type >, String >
                {
                    let v = self.0.iter().filter_map( |p| match p {
                        $attr_param::$name( v ) => Some( v ),
                        _ => None
                    } ).take( 2 ).collect::<Vec<_>>();

                    match v.len() {
                        0 => Ok( None ),
                        1 => Ok( Some( &v[0] ) ),
                        _ => Err( format!(
                                "Multiple {} arguments", stringify!( $name ) ) )
                    }
                }
            )*

            pub fn args<'a>( &'a self ) -> Vec<&'a $params>
            {
                self.0.iter().filter_map( |p| match p {
                    $attr_param::args( v ) => Some( v ),
                    _ => None
                } ).collect()
            }
        }
    }
}

<commit_msg>Document macros<commit_after>
/// Type that implements `Synom` but fails parsing.
///
/// Use this for defining attribute parsing that accepts only named arguments.
pub enum NoParams {}
impl ::syn::synom::Synom for NoParams {
    named!(parse -> Self, reject!() );
}

/// Literal string or 'None' if there should be no value.
pub enum StrOption {
    Str( ::syn::LitStr ),
    None
}
impl ::syn::synom::Synom for StrOption {
    named!(parse -> Self, alt!(
        syn!( ::syn::LitStr ) => { |s| StrOption::Str( s ) }
        |
        custom_keyword!( None ) => { |_| StrOption::None }
    ) );
}

/// Defines intercom attribute parameter parsing.
///
/// ```
/// intercom_attribute!(
///     SomeAttr< SomeAttrParam, Ident > {
///         param1: Ident,
///         param2: LitStr,
///         param3: Expr,
///     }
/// );
/// ```
///
/// Will define structures to parse attribute params such as:
///
/// ```
/// #[some_attr( Ident1, Ident2, Ident3 )]
/// #[some_attr( param1 = SomeIdent, List, Idents )]
/// #[some_attr( param2 = "literal", param3 = expression() + 1 )]
/// ```
macro_rules! intercom_attribute {

    ( $attr:ident < $attr_param:ident, $params:ident > { $( $name:ident : $type:ident, )* } ) => {

        enum $attr_param {
            $( $name ( $type ), )*
            args( $params ),
        }

        impl ::syn::synom::Synom for $attr_param {
            named!( parse -> Self, alt!(
                $(
                    do_parse!(
                        custom_keyword!( $name ) >>
                        eq : punct!(=) >>
                        value : syn!( $type ) >>
                        ( $attr_param::$name( value ) )
                    )
                    |
                )*
                do_parse!( a : syn!( $params ) >> ( $attr_param::args( a ) ) )
            ) );
        }

        struct $attr( Vec<$attr_param> );
        impl ::syn::synom::Synom for $attr {
            named!(parse -> Self, do_parse!(
                p: call!( ::syn::punctuated::Punctuated::<$attr_param, ::syn::token::Comma>::parse_separated_nonempty ) >>
                ( $attr( p.into_iter().collect() ) )
            ) );
        }

        impl $attr {
            $(
                pub fn $name( &self ) -> Result< Option< &$type >, String >
                {
                    let v = self.0.iter().filter_map( |p| match p {
                        $attr_param::$name( v ) => Some( v ),
                        _ => None
                    } ).take( 2 ).collect::<Vec<_>>();

                    match v.len() {
                        0 => Ok( None ),
                        1 => Ok( Some( &v[0] ) ),
                        _ => Err( format!(
                                "Multiple {} arguments", stringify!( $name ) ) )
                    }
                }
            )*

            pub fn args<'a>( &'a self ) -> Vec<&'a $params>
            {
                self.0.iter().filter_map( |p| match p {
                    $attr_param::args( v ) => Some( v ),
                    _ => None
                } ).collect()
            }
        }
    }
}

<|endoftext|>


<commit_before>extern crate ring;
extern crate rustc_serialize;

use self::rustc_serialize::{ Encodable, Encoder, Decodable, Decoder };
use self::ring::digest::{ Algorithm, SHA256 };

pub mod hashutils;
pub mod tree;
pub mod proof;
pub mod merkletree;

use self::hashutils::*;
use self::tree::*;
use self::merkletree::*;

static DIGEST: &'static Algorithm = &SHA256;

impl MerkleTree {
    fn count_bits(v: usize) -> usize {
        let mut ret = 0;
        let mut val = v;
        while val != 0 {
            val &= val - 1;
            ret += 1;
        }
        return ret;
    }

    pub fn append(&mut self, node: TreeLeafData) {
        if self.count == 0 {
            // empty tree
            self.root = Tree::new_leaf(self.algorithm, node);
        }
        else if Self::count_bits(self.count) != 1 {
            // add to right subtree
            match self.root.clone() {
                Tree::Node { ref left, ref right, .. } => {
                    let mut iter = right.iter().map(|x| (*x).clone()).collect::<Vec<TreeLeafData>>();
                    iter.push(node);
                    let new_right = MerkleTree::from_vec(self.algorithm, iter);
                    let combined_hash = self.algorithm.hash_nodes(
                        left.hash(),
                        new_right.root_hash() as &Vec<u8>
                    );
                    self.root = Tree::Node {
                        left: (*left).clone(),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
        }
        else
        {
            // add tree layer
            let new_right = MerkleTree::from_vec(self.algorithm, vec![ node ]);
            match self.root.clone() {
                Tree::Node { ref left, ref right, ref hash } => {
                    let combined_hash = self.algorithm.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(self.root.clone()),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                Tree::Leaf { ref hash, ref value } => {
                    let combined_hash = self.algorithm.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(Tree::new_leaf(self.algorithm, (*value).clone())),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
            self.height += 1;
        }
        self.count += 1;
    }
}

impl Encodable for MerkleTree {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_enum("MerkleTree", |s| {
            self.root.encode(s)
        })
    }
}

impl Decodable for MerkleTree {
    fn decode<D: Decoder>(d: &mut D) -> Result<MerkleTree, D::Error> {
        let r = Tree::decode(d);
        match r {
            Ok(root) => {
                let h = root.get_height();
                let c = root.get_count();
                Ok(MerkleTree {
                    algorithm: DIGEST,
                    root: root,
                    height: h,
                    count: c
                })
            }
            Err(e) => {
                return Err(e);
            }
        }
    }
}

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

    #[test]
    fn test_merkletree_append() {
        let values = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9" ];
        let mut mt = MerkleTree::from_vec(DIGEST, vec![]);
        println!("root(0)={}", mt.root_hash_hex());
        let mut r = 1;
        for i in values {
            mt.append(String::from(i));
            println!("root({})={}", r, mt.root_hash_hex());
            r+=1;
        }
    }

    #[test]
    fn test_valid_proof() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let tree      = MerkleTree::from_vec(DIGEST, values.clone());
        let root_hash = tree.root_hash();

        for value in values {
            let proof    = tree.gen_proof(value);
            let is_valid = proof.map(|p| p.validate(&root_hash)).unwrap_or(false);

            assert!(is_valid);
        }
    }

    #[test]
    fn test_merkletree_serialize() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let mt = MerkleTree::from_vec(DIGEST, values.clone());
        println!("serialize mt: h={}, c={}, rhash={}", mt.height, mt.count, json::encode(&mt.root_hash()).unwrap());
        let serialized = json::encode(&mt).unwrap();
        println!("serialize: {}", serialized);
        let newmt: MerkleTree = json::decode(serialized.as_str()).unwrap();
        println!("serialize newmt: h={}, c={}", newmt.height, newmt.count);

        let mut collected = Vec::new();
        for value in &newmt {
            collected.push(value);
        }
        let refs = values.iter().collect::<Vec<_>>();
        assert_eq!(refs, collected);

        assert_eq!(mt.root_hash(), newmt.root_hash());
    }
}
<commit_msg>add real data to test and check root hash<commit_after>extern crate ring;
extern crate rustc_serialize;

use self::rustc_serialize::{ Encodable, Encoder, Decodable, Decoder };
use self::ring::digest::{ Algorithm, SHA256 };

pub mod hashutils;
pub mod tree;
pub mod proof;
pub mod merkletree;

use self::hashutils::*;
use self::tree::*;
use self::merkletree::*;

static DIGEST: &'static Algorithm = &SHA256;

impl MerkleTree {
    fn count_bits(v: usize) -> usize {
        let mut ret = 0;
        let mut val = v;
        while val != 0 {
            val &= val - 1;
            ret += 1;
        }
        return ret;
    }

    pub fn append(&mut self, node: TreeLeafData) {
        if self.count == 0 {
            // empty tree
            self.root = Tree::new_leaf(self.algorithm, node);
        }
        else if Self::count_bits(self.count) != 1 {
            // add to right subtree
            match self.root.clone() {
                Tree::Node { ref left, ref right, .. } => {
                    let mut iter = right.iter().map(|x| (*x).clone()).collect::<Vec<TreeLeafData>>();
                    iter.push(node);
                    let new_right = MerkleTree::from_vec(self.algorithm, iter);
                    let combined_hash = self.algorithm.hash_nodes(
                        left.hash(),
                        new_right.root_hash() as &Vec<u8>
                    );
                    self.root = Tree::Node {
                        left: (*left).clone(),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
        }
        else
        {
            // add tree layer
            let new_right = MerkleTree::from_vec(self.algorithm, vec![ node ]);
            match self.root.clone() {
                Tree::Node { ref left, ref right, ref hash } => {
                    let combined_hash = self.algorithm.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(self.root.clone()),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                Tree::Leaf { ref hash, ref value } => {
                    let combined_hash = self.algorithm.hash_nodes(
                        hash,
                        new_right.root_hash()
                    );
                    self.root = Tree::Node {
                        left: Box::new(Tree::new_leaf(self.algorithm, (*value).clone())),
                        right: Box::new(new_right.root),
                        hash: combined_hash.as_ref().into()
                    }
                }
                _ => {
                    assert!(false);
                }
            }
            self.height += 1;
        }
        self.count += 1;
    }
}

impl Encodable for MerkleTree {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_enum("MerkleTree", |s| {
            self.root.encode(s)
        })
    }
}

impl Decodable for MerkleTree {
    fn decode<D: Decoder>(d: &mut D) -> Result<MerkleTree, D::Error> {
        let r = Tree::decode(d);
        match r {
            Ok(root) => {
                let h = root.get_height();
                let c = root.get_count();
                Ok(MerkleTree {
                    algorithm: DIGEST,
                    root: root,
                    height: h,
                    count: c
                })
            }
            Err(e) => {
                return Err(e);
            }
        }
    }
}

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

    #[test]
    fn test_merkletree_append() {
        let values = vec![
            "{\"data\":{\"alias\":\"Node1\",\"client_ip\":\"192.168.1.35\",\"client_port\":9702,\"node_ip\":\"192.168.1.35\",\"node_port\":9701,\"services\":[\"VALIDATOR\"]},\"dest\":\"Gw6pDLhcBcoQesN72qfotTgFa7cbuqZpkX3Xo6pLhPhv\",\"identifier\":\"FYmoFw55GeQH7SRFa37dkx1d2dZ3zUF8ckg7wmL7ofN4\",\"txnId\":\"fea82e10e894419fe2bea7d96296a6d46f50f93f9eeda954ec461b2ed2950b62\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node2\",\"client_ip\":\"192.168.1.35\",\"client_port\":9704,\"node_ip\":\"192.168.1.35\",\"node_port\":9703,\"services\":[\"VALIDATOR\"]},\"dest\":\"8ECVSk179mjsjKRLWiQtssMLgp6EPhWXtaYyStWPSGAb\",\"identifier\":\"8QhFxKxyaFsJy4CyxeYX34dFH8oWqyBv1P4HLQCsoeLy\",\"txnId\":\"1ac8aece2a18ced660fef8694b61aac3af08ba875ce3026a160acbc3a3af35fc\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node3\",\"client_ip\":\"192.168.1.35\",\"client_port\":9706,\"node_ip\":\"192.168.1.35\",\"node_port\":9705,\"services\":[\"VALIDATOR\"]},\"dest\":\"DKVxG2fXXTU8yT5N7hGEbXB3dfdAnYv1JczDUHpmDxya\",\"identifier\":\"2yAeV5ftuasWNgQwVYzeHeTuM7LwwNtPR3Zg9N4JiDgF\",\"txnId\":\"7e9f355dffa78ed24668f0e0e369fd8c224076571c51e2ea8be5f26479edebe4\",\"type\":\"0\"}",
            "{\"data\":{\"alias\":\"Node4\",\"client_ip\":\"192.168.1.35\",\"client_port\":9708,\"node_ip\":\"192.168.1.35\",\"node_port\":9707,\"services\":[\"VALIDATOR\"]},\"dest\":\"4PS3EDQ3dW1tci1Bp6543CfuuebjFrg36kLAUcskGfaA\",\"identifier\":\"FTE95CVthRtrBnK2PYCBbC9LghTcGwi9Zfi1Gz2dnyNx\",\"txnId\":\"aa5e817d7cc626170eca175822029339a444eb0ee8f0bd20d3b0b76e566fb008\",\"type\":\"0\"}" ];
        let mut mt = MerkleTree::from_vec(DIGEST, vec![]);
        println!("root(0)={}", mt.root_hash_hex());
        let mut r = 1;
        for i in values {
            mt.append(String::from(i));
            println!("root({})={}", r, mt.root_hash_hex());
            r+=1;
        }
        assert_eq!(mt.root_hash_hex(), "1285070cf01debc1155cef8dfd5ba54c05abb919a4c08c8632b079fb1e1e5e7c");
    }

    #[test]
    fn test_valid_proof() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let tree      = MerkleTree::from_vec(DIGEST, values.clone());
        let root_hash = tree.root_hash();

        for value in values {
            let proof    = tree.gen_proof(value);
            let is_valid = proof.map(|p| p.validate(&root_hash)).unwrap_or(false);

            assert!(is_valid);
        }
    }

    #[test]
    fn test_merkletree_serialize() {
        let strvals   = vec![ "1", "2", "3", "4", "5", "6", "7", "8", "9", "10" ];
        let values    = strvals.iter().map(|x| String::from(*x)).collect::<Vec<_>>();
        let mt = MerkleTree::from_vec(DIGEST, values.clone());
        println!("serialize mt: h={}, c={}, rhash={}", mt.height, mt.count, json::encode(&mt.root_hash()).unwrap());
        let serialized = json::encode(&mt).unwrap();
        println!("serialize: {}", serialized);
        let newmt: MerkleTree = json::decode(serialized.as_str()).unwrap();
        println!("serialize newmt: h={}, c={}", newmt.height, newmt.count);

        let mut collected = Vec::new();
        for value in &newmt {
            collected.push(value);
        }
        let refs = values.iter().collect::<Vec<_>>();
        assert_eq!(refs, collected);

        assert_eq!(mt.root_hash(), newmt.root_hash());
    }
}
<|endoftext|>


<commit_before>extern crate chrono;

#[cfg(test)]
mod test {
    use chrono::{Datelike, Duration, NaiveDate, Weekday};

    #[test]
    fn can_detect_day_of_week() {
        let date = NaiveDate::from_ymd(2016, 3, 24);

        assert!(Weekday::Thu == date.weekday());
    }

}<commit_msg>added documentation example and build it<commit_after>//! Library for generating relative datestream
//!
//! This is some long documentation

extern crate chrono;

#[cfg(test)]
mod test {
    use chrono::{Datelike, Duration, NaiveDate, Weekday};

    #[test]
    fn can_detect_day_of_week() {
        let date = NaiveDate::from_ymd(2016, 3, 24);

        assert!(Weekday::Thu == date.weekday());
    }

}<|endoftext|>


<commit_before>#[macro_use]
extern crate quote;
extern crate reqwest;
extern crate serde_json;
#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate lazy_static;

use std::fs::File;
use serde_json::{ Value, Map };
use std::collections::{ HashSet, HashMap };
use quote::{ Tokens, Ident };

mod util;

use util::*;

mod errors {
    error_chain! {
        types {
            Error, ErrorKind, ResultExt, Result;
        }

        links {
        }

        foreign_links {
            ReqwestError(::reqwest::Error);
            JsonError(::serde_json::Error);
            IoError(::std::io::Error);
        }

        errors {
            MissingSource {
                description("No source for sample specified")
            }
            ExistingType(t: String) {
                display("No code generated, JSON matches existing type {}", t)
            }
        }
    }
}

pub use errors::*;

const RUST_KEYWORDS_ARR: &'static [&'static str] = &["abstract", "alignof", "as", "become", "box",
    "break", "const", "continue", "crate", "do", "else", "enum", "extern", "false", "final", "fn",
    "for", "if", "impl", "in", "let", "loop", "macro", "match", "mod", "move", "mut", "offsetof",
    "override", "priv", "proc", "pub", "pure", "ref", "return", "Self", "self", "sizeof", "static",
    "struct", "super", "trait", "true", "type", "typeof", "unsafe", "unsized", "use", "virtual",
    "where", "while", "yield"];

lazy_static! {
    static ref RUST_KEYWORDS: HashSet<&'static str> = {
        RUST_KEYWORDS_ARR.iter().cloned().collect()
    };
}

pub enum SampleSource<'a> {
    Url(&'a str),
    File(&'a str),
    Text(&'a str),
}

pub struct Options {
    pub use_serde: bool,
    pub extern_crate: bool,
    pub runnable: bool
}

impl Default for Options {
    fn default() -> Options {
        Options {
            use_serde: true,
            extern_crate: false,
            runnable: false
        }
    }
}

struct Ctxt {
    options: Options
}

macro_rules! some_if {
    ($cond:expr, $then:expr) => ({
        if $cond {
            Some($then)
        } else {
            None
        }
    })
}

#[allow(dead_code)]
#[derive(Debug, PartialEq, Clone)]
enum InferredType {
    Null,
    Any,
    Bool,
    StringT,
    Integer,
    Floating,
    VecT { elem_type: Box<InferredType> },
    Struct { fields: HashMap<String, InferredType> },
    Optional(Box<InferredType>)
}

#[allow(dead_code)]
fn unify(a: InferredType, b: InferredType) -> InferredType {
    if a == b {
        return a;
    }
    use InferredType::*;
    match (a, b) {
        (Integer, Floating) |
        (Floating, Integer) => Floating,
        (a, Null) | (Null, a) => make_optional(a),
        (a, Optional(b)) | (Optional(b), a) => make_optional(unify(a, *b)),
        (VecT { elem_type: e1 }, VecT { elem_type: e2 }) => {
            VecT { elem_type: Box::new(unify(*e1, *e2)) }
        }
        (Struct { fields: f1 }, Struct { fields: f2 }) => {
            Struct { fields: unify_struct_fields(f1, f2) }
        }
        _ => Any,
    }
}

fn make_optional(a: InferredType) -> InferredType {
    use InferredType::*;
    match a {
        Null | Any | Optional(_) => a,
        non_nullable => Optional(Box::new(non_nullable)),
    }
}

fn unify_struct_fields(mut f1: HashMap<String, InferredType>,
                       mut f2: HashMap<String, InferredType>)
                       -> HashMap<String, InferredType> {
    if f1 == f2 {
        return f1;
    }
    let mut unified = HashMap::new();
    for (key, val) in f1.drain() {
        match f2.remove(&key) {
            Some(val2) => {
                unified.insert(key, unify(val, val2));
            },
            None => {
                unified.insert(key, make_optional(val));
            }
        }
    }
    for (key, val) in f2.drain() {
        unified.insert(key, make_optional(val));
    }
    unified
}

pub fn from_str_with_defaults(name: &str, json: &str) -> Result<Tokens> {
    codegen(name, &SampleSource::Text(json), Options::default())
}

pub fn codegen(name: &str, source: &SampleSource, options: Options) -> Result<Tokens> {
    let sample = get_and_parse_sample(source)?;
    let mut ctxt = Ctxt {
        options: options
    };
    let (type_name, type_def) = generate_type_from_value(&mut ctxt, name, &sample);

    let example = some_if!(ctxt.options.runnable, {
        ctxt.options.extern_crate = true;
        usage_example(&type_name)
    });

    let crates = some_if!(ctxt.options.extern_crate, quote! {
        #[macro_use]
        extern crate serde_derive;
        extern crate serde_json;
    });

    if type_def.is_none() && !ctxt.options.runnable {
        return Err(ErrorKind::ExistingType(String::from(type_name.as_str())).into());
    }

    Ok(quote! {
        #crates

        #type_def

        #example
    })
}

fn usage_example(type_id: &Tokens) -> Tokens {
    let var_id = Ident::from(snake_case(type_id.as_str()));

    quote! {
        fn main() {
            let #var_id: #type_id = Default::default();

            let serialized = serde_json::to_string(&#var_id).unwrap();

            println!("serialized = {}", serialized);

            let deserialized: #type_id = serde_json::from_str(&serialized).unwrap();

            println!("deserialized = {:?}", deserialized);
        }
    }
}

fn generate_type_from_value(ctxt: &mut Ctxt, path: &str, value: &Value) -> (Tokens, Option<Tokens>) {
    match *value {
        Value::Null => (quote! { Option<::serde_json::Value> }, None),
        Value::Bool(_) => (quote! { bool }, None),
        Value::Number(ref n) => {
            if n.is_i64() {
                (quote! { i64 }, None)
            } else {
                (quote! { f64 }, None)
            }
        },
        Value::String(_) => (quote! { String }, None),
        Value::Array(ref values) => {
            generate_type_for_array(ctxt, path, values)
        },
        Value::Object(ref map) => {
            generate_struct_from_object(ctxt, path, map)
        }
    }
}

fn generate_type_for_array(ctxt: &mut Ctxt, path: &str, values: &[Value]) -> (Tokens, Option<Tokens>) {
    let mut defs = Vec::new();
    let mut types = HashSet::new();

    for value in values.iter() {
        let (elemtype, elemtype_def) = generate_type_from_value(ctxt, path, value);
        types.insert(elemtype.into_string());
        if let Some(def) = elemtype_def {
            defs.push(def);
        }
    }

    if types.len() == 1 {
        let ident = Ident::new(types.into_iter().next().unwrap());
        (quote! { Vec<#ident> }, defs.into_iter().next())
    } else {
        (quote! { Vec<::serde_json::Value> }, None)
    }
}

fn generate_struct_from_object(ctxt: &mut Ctxt, path: &str, map: &Map<String, Value>) -> (Tokens, Option<Tokens>) {
    let type_name = type_case(path);
    let ident = Ident::from(type_name);
    let mut defs = Vec::new();

    let fields: Vec<Tokens> = map.iter()
        .map(|(name, value)| {
            let mut field_name = snake_case(name);
            if RUST_KEYWORDS.contains::<str>(&field_name) {
                field_name.push_str("_field")
            }
            let rename = some_if!(&field_name != name,
                quote! { #[serde(rename = #name)] });

            let field_ident = Ident::from(field_name);
            let (fieldtype, fieldtype_def) = generate_type_from_value(ctxt, name, value);
            if let Some(def) = fieldtype_def {
                defs.push(def);
            }
            quote! {
                #rename
                #field_ident: #fieldtype
            }
        })
        .collect();

    let derives = if ctxt.options.use_serde {
        quote! { #[derive(Default, Debug, Clone, Serialize, Deserialize)] }
    } else {
        quote! { #[derive(Default, Debug, Clone)] }
    };

    let code = quote! {
        #derives
        struct #ident {
            #(#fields),*
        }

        #(#defs)*
    };

    (quote! { #ident }, Some(code))
}

pub fn infer_source_type(s: &str) -> SampleSource {
    let s = s.trim();
    if s.starts_with('{') || s.starts_with('[') {
        return SampleSource::Text(s);
    }
    if s.starts_with("http://") || s.starts_with("https://") {
        return SampleSource::Url(s);
    }
    SampleSource::File(s)
}

fn get_and_parse_sample(source: &SampleSource) -> Result<Value> {
    let parse_result = match *source {
        SampleSource::Url(url) => serde_json::de::from_reader(reqwest::get(url)?),
        SampleSource::File(path) => serde_json::de::from_reader(File::open(path)?),
        SampleSource::Text(text) => serde_json::from_str(text),
    };
    Ok(parse_result.chain_err(|| "Unable to parse JSON sample")?)
}

#[test]
fn test_unify() {
    use InferredType::*;
    assert_eq!(unify(Bool, Bool), Bool);
    assert_eq!(unify(Bool, Integer), Any);
    assert_eq!(unify(Integer, Floating), Floating);
    assert_eq!(unify(Null, Any), Any);
    assert_eq!(unify(Null, Bool), Optional(Box::new(Bool)));
    assert_eq!(unify(Null, Optional(Box::new(Integer))), Optional(Box::new(Integer)));
    assert_eq!(unify(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(unify(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(unify(Optional(Box::new(Integer)), Optional(Box::new(Floating))),
               Optional(Box::new(Floating)));
    assert_eq!(unify(Optional(Box::new(StringT)), Optional(Box::new(Integer))), Any);
}

// based on hashmap! macro from maplit crate
macro_rules! string_hashmap {
    ($($key:expr => $value:expr,)+) => { string_hashmap!($($key => $value),+) };
    ($($key:expr => $value:expr),*) => {
        {
            let mut _map = ::std::collections::HashMap::new();
            $(
                _map.insert($key.to_string(), $value);
            )*
            _map
        }
    };
}

#[test]
fn test_unify_struct_fields() {
    use InferredType::*;
    {
        let f1 = string_hashmap!{
            "a" => Integer,
            "b" => Bool,
            "c" => Integer,
            "d" => StringT,
        };
        let f2 = string_hashmap!{
            "a" => Integer,
            "c" => Floating,
            "d" => Null,
            "e" => Any,
        };
        assert_eq!(unify_struct_fields(f1, f2), string_hashmap!{
            "a" => Integer,
            "b" => Optional(Box::new(Bool)),
            "c" => Floating,
            "d" => Optional(Box::new(StringT)),
            "e" => Any,
        });
    }
}
<commit_msg>shared: Impl inference functions<commit_after>#[macro_use]
extern crate quote;
extern crate reqwest;
extern crate serde_json;
#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate lazy_static;

use std::fs::File;
use serde_json::{ Value, Map };
use std::collections::{ HashSet, HashMap };
use quote::{ Tokens, Ident };

mod util;

use util::*;

mod errors {
    error_chain! {
        types {
            Error, ErrorKind, ResultExt, Result;
        }

        links {
        }

        foreign_links {
            ReqwestError(::reqwest::Error);
            JsonError(::serde_json::Error);
            IoError(::std::io::Error);
        }

        errors {
            MissingSource {
                description("No source for sample specified")
            }
            ExistingType(t: String) {
                display("No code generated, JSON matches existing type {}", t)
            }
        }
    }
}

pub use errors::*;

const RUST_KEYWORDS_ARR: &'static [&'static str] = &["abstract", "alignof", "as", "become", "box",
    "break", "const", "continue", "crate", "do", "else", "enum", "extern", "false", "final", "fn",
    "for", "if", "impl", "in", "let", "loop", "macro", "match", "mod", "move", "mut", "offsetof",
    "override", "priv", "proc", "pub", "pure", "ref", "return", "Self", "self", "sizeof", "static",
    "struct", "super", "trait", "true", "type", "typeof", "unsafe", "unsized", "use", "virtual",
    "where", "while", "yield"];

lazy_static! {
    static ref RUST_KEYWORDS: HashSet<&'static str> = {
        RUST_KEYWORDS_ARR.iter().cloned().collect()
    };
}

pub enum SampleSource<'a> {
    Url(&'a str),
    File(&'a str),
    Text(&'a str),
}

pub struct Options {
    pub use_serde: bool,
    pub extern_crate: bool,
    pub runnable: bool
}

impl Default for Options {
    fn default() -> Options {
        Options {
            use_serde: true,
            extern_crate: false,
            runnable: false
        }
    }
}

struct Ctxt {
    options: Options
}

macro_rules! some_if {
    ($cond:expr, $then:expr) => ({
        if $cond {
            Some($then)
        } else {
            None
        }
    })
}

#[allow(dead_code)]
#[derive(Debug, PartialEq, Clone)]
enum InferredType {
    Null,
    Any,
    Bool,
    StringT,
    Integer,
    Floating,
    EmptyVec,
    VecT { elem_type: Box<InferredType> },
    Struct { fields: HashMap<String, InferredType> },
    Optional(Box<InferredType>)
}

#[allow(dead_code)]
fn unify(a: InferredType, b: InferredType) -> InferredType {
    if a == b {
        return a;
    }
    use InferredType::*;
    match (a, b) {
        (Integer, Floating) |
        (Floating, Integer) => Floating,
        (a, Null) | (Null, a) => make_optional(a),
        (a, Optional(b)) | (Optional(b), a) => make_optional(unify(a, *b)),
        (EmptyVec, VecT { elem_type: e }) |
        (VecT { elem_type: e }, EmptyVec) => VecT { elem_type: e },
        (VecT { elem_type: e1 }, VecT { elem_type: e2 }) => {
            VecT { elem_type: Box::new(unify(*e1, *e2)) }
        }
        (Struct { fields: f1 }, Struct { fields: f2 }) => {
            Struct { fields: unify_struct_fields(f1, f2) }
        }
        _ => Any,
    }
}

fn make_optional(a: InferredType) -> InferredType {
    use InferredType::*;
    match a {
        Null | Any | Optional(_) => a,
        non_nullable => Optional(Box::new(non_nullable)),
    }
}

fn unify_struct_fields(mut f1: HashMap<String, InferredType>,
                       mut f2: HashMap<String, InferredType>)
                       -> HashMap<String, InferredType> {
    if f1 == f2 {
        return f1;
    }
    let mut unified = HashMap::new();
    for (key, val) in f1.drain() {
        match f2.remove(&key) {
            Some(val2) => {
                unified.insert(key, unify(val, val2));
            },
            None => {
                unified.insert(key, make_optional(val));
            }
        }
    }
    for (key, val) in f2.drain() {
        unified.insert(key, make_optional(val));
    }
    unified
}

#[allow(dead_code)]
fn infer_type_from_value(ctxt: &mut Ctxt, value: &Value) -> InferredType {
    match *value {
        Value::Null => InferredType::Null,
        Value::Bool(_) => InferredType::Bool,
        Value::Number(ref n) => {
            if n.is_i64() {
                InferredType::Integer
            } else {
                InferredType::Floating
            }
        },
        Value::String(_) => InferredType::StringT,
        Value::Array(ref values) => {
            infer_type_for_array(ctxt, values)
        },
        Value::Object(ref map) => {
            InferredType::Struct { fields: infer_types_for_fields(ctxt, map) }
        }
    }
}

#[allow(dead_code)]
fn infer_type_for_array(ctxt: &mut Ctxt, values: &[Value]) -> InferredType {
    match values.split_first() {
        None => InferredType::EmptyVec,
        Some((first, rest)) => {
            let first_type = infer_type_from_value(ctxt, first);
            let inner = rest.iter().fold(first_type, |typ, val| {
                let new_type = infer_type_from_value(ctxt, val);
                unify(typ, new_type)
            });
            InferredType::VecT { elem_type: Box::new(inner) }
        }
    }
}

#[allow(dead_code)]
fn infer_types_for_fields(ctxt: &mut Ctxt, map: &Map<String, Value>) -> HashMap<String, InferredType> {
    map.iter()
        .map(|(name, value)| (name.clone(), infer_type_from_value(ctxt, value)))
        .collect()
}

pub fn from_str_with_defaults(name: &str, json: &str) -> Result<Tokens> {
    codegen(name, &SampleSource::Text(json), Options::default())
}

pub fn codegen(name: &str, source: &SampleSource, options: Options) -> Result<Tokens> {
    let sample = get_and_parse_sample(source)?;
    let mut ctxt = Ctxt {
        options: options
    };
    let (type_name, type_def) = generate_type_from_value(&mut ctxt, name, &sample);

    let example = some_if!(ctxt.options.runnable, {
        ctxt.options.extern_crate = true;
        usage_example(&type_name)
    });

    let crates = some_if!(ctxt.options.extern_crate, quote! {
        #[macro_use]
        extern crate serde_derive;
        extern crate serde_json;
    });

    if type_def.is_none() && !ctxt.options.runnable {
        return Err(ErrorKind::ExistingType(String::from(type_name.as_str())).into());
    }

    Ok(quote! {
        #crates

        #type_def

        #example
    })
}

fn usage_example(type_id: &Tokens) -> Tokens {
    let var_id = Ident::from(snake_case(type_id.as_str()));

    quote! {
        fn main() {
            let #var_id: #type_id = Default::default();

            let serialized = serde_json::to_string(&#var_id).unwrap();

            println!("serialized = {}", serialized);

            let deserialized: #type_id = serde_json::from_str(&serialized).unwrap();

            println!("deserialized = {:?}", deserialized);
        }
    }
}

fn generate_type_from_value(ctxt: &mut Ctxt, path: &str, value: &Value) -> (Tokens, Option<Tokens>) {
    match *value {
        Value::Null => (quote! { Option<::serde_json::Value> }, None),
        Value::Bool(_) => (quote! { bool }, None),
        Value::Number(ref n) => {
            if n.is_i64() {
                (quote! { i64 }, None)
            } else {
                (quote! { f64 }, None)
            }
        },
        Value::String(_) => (quote! { String }, None),
        Value::Array(ref values) => {
            generate_type_for_array(ctxt, path, values)
        },
        Value::Object(ref map) => {
            generate_struct_from_object(ctxt, path, map)
        }
    }
}

fn generate_type_for_array(ctxt: &mut Ctxt, path: &str, values: &[Value]) -> (Tokens, Option<Tokens>) {
    let mut defs = Vec::new();
    let mut types = HashSet::new();

    for value in values.iter() {
        let (elemtype, elemtype_def) = generate_type_from_value(ctxt, path, value);
        types.insert(elemtype.into_string());
        if let Some(def) = elemtype_def {
            defs.push(def);
        }
    }

    if types.len() == 1 {
        let ident = Ident::new(types.into_iter().next().unwrap());
        (quote! { Vec<#ident> }, defs.into_iter().next())
    } else {
        (quote! { Vec<::serde_json::Value> }, None)
    }
}

fn generate_struct_from_object(ctxt: &mut Ctxt, path: &str, map: &Map<String, Value>) -> (Tokens, Option<Tokens>) {
    let type_name = type_case(path);
    let ident = Ident::from(type_name);
    let mut defs = Vec::new();

    let fields: Vec<Tokens> = map.iter()
        .map(|(name, value)| {
            let mut field_name = snake_case(name);
            if RUST_KEYWORDS.contains::<str>(&field_name) {
                field_name.push_str("_field")
            }
            let rename = some_if!(&field_name != name,
                quote! { #[serde(rename = #name)] });

            let field_ident = Ident::from(field_name);
            let (fieldtype, fieldtype_def) = generate_type_from_value(ctxt, name, value);
            if let Some(def) = fieldtype_def {
                defs.push(def);
            }
            quote! {
                #rename
                #field_ident: #fieldtype
            }
        })
        .collect();

    let derives = if ctxt.options.use_serde {
        quote! { #[derive(Default, Debug, Clone, Serialize, Deserialize)] }
    } else {
        quote! { #[derive(Default, Debug, Clone)] }
    };

    let code = quote! {
        #derives
        struct #ident {
            #(#fields),*
        }

        #(#defs)*
    };

    (quote! { #ident }, Some(code))
}

pub fn infer_source_type(s: &str) -> SampleSource {
    let s = s.trim();
    if s.starts_with('{') || s.starts_with('[') {
        return SampleSource::Text(s);
    }
    if s.starts_with("http://") || s.starts_with("https://") {
        return SampleSource::Url(s);
    }
    SampleSource::File(s)
}

fn get_and_parse_sample(source: &SampleSource) -> Result<Value> {
    let parse_result = match *source {
        SampleSource::Url(url) => serde_json::de::from_reader(reqwest::get(url)?),
        SampleSource::File(path) => serde_json::de::from_reader(File::open(path)?),
        SampleSource::Text(text) => serde_json::from_str(text),
    };
    Ok(parse_result.chain_err(|| "Unable to parse JSON sample")?)
}

#[test]
fn test_unify() {
    use InferredType::*;
    assert_eq!(unify(Bool, Bool), Bool);
    assert_eq!(unify(Bool, Integer), Any);
    assert_eq!(unify(Integer, Floating), Floating);
    assert_eq!(unify(Null, Any), Any);
    assert_eq!(unify(Null, Bool), Optional(Box::new(Bool)));
    assert_eq!(unify(Null, Optional(Box::new(Integer))), Optional(Box::new(Integer)));
    assert_eq!(unify(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(unify(Any, Optional(Box::new(Integer))), Any);
    assert_eq!(unify(Optional(Box::new(Integer)), Optional(Box::new(Floating))),
               Optional(Box::new(Floating)));
    assert_eq!(unify(Optional(Box::new(StringT)), Optional(Box::new(Integer))), Any);
}

// based on hashmap! macro from maplit crate
macro_rules! string_hashmap {
    ($($key:expr => $value:expr,)+) => { string_hashmap!($($key => $value),+) };
    ($($key:expr => $value:expr),*) => {
        {
            let mut _map = ::std::collections::HashMap::new();
            $(
                _map.insert($key.to_string(), $value);
            )*
            _map
        }
    };
}

#[test]
fn test_unify_struct_fields() {
    use InferredType::*;
    {
        let f1 = string_hashmap!{
            "a" => Integer,
            "b" => Bool,
            "c" => Integer,
            "d" => StringT,
        };
        let f2 = string_hashmap!{
            "a" => Integer,
            "c" => Floating,
            "d" => Null,
            "e" => Any,
        };
        assert_eq!(unify_struct_fields(f1, f2), string_hashmap!{
            "a" => Integer,
            "b" => Optional(Box::new(Bool)),
            "c" => Floating,
            "d" => Optional(Box::new(StringT)),
            "e" => Any,
        });
    }
}
<|endoftext|>


<commit_before>#![allow(stable_features)]
// ignore-windows - this is a unix-specific test
// ignore-cloudabi no processes
// ignore-emscripten no processes

#![feature(process_exec, rustc_private)]

extern crate libc;

use std::env;
use std::io::Error;
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

fn main() {
    if let Some(arg) = env::args().nth(1) {
        match &arg[..] {
            "test1" => println!("hello2"),
            "test2" => assert_eq!(env::var("FOO").unwrap(), "BAR"),
            "test3" => assert_eq!(env::current_dir().unwrap()
                                      .to_str().unwrap(), "/"),
            "empty" => {}
            _ => panic!("unknown argument: {}", arg),
        }
        return
    }

    let me = env::current_exe().unwrap();

    let output = unsafe { Command::new(&me).arg("test1").pre_exec(|| {
        println!("hello");
        Ok(())
    }).output().unwrap() };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert_eq!(output.stdout, b"hello\nhello2\n");

    let output = unsafe { Command::new(&me).arg("test2").pre_exec(|| {
        env::set_var("FOO", "BAR");
        Ok(())
    }).output().unwrap() };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());

    let output = unsafe { Command::new(&me).arg("test3").pre_exec(|| {
        env::set_current_dir("/").unwrap();
        Ok(())
    }).output().unwrap() };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());

    let output = unsafe { Command::new(&me).arg("bad").pre_exec(|| {
        Err(Error::from_raw_os_error(102))
    }).output().unwrap_err() };
    assert_eq!(output.raw_os_error(), Some(102));

    let pid = unsafe { libc::getpid() };
    assert!(pid >= 0);
    let output = unsafe { Command::new(&me).arg("empty").pre_exec(move || {
        let child = libc::getpid();
        assert!(child >= 0);
        assert!(pid != child);
        Ok(())
    }).output().unwrap() };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());

    let mem = Arc::new(AtomicUsize::new(0));
    let mem2 = mem.clone();
    let output = unsafe { Command::new(&me).arg("empty").pre_exec(move || {
        assert_eq!(mem2.fetch_add(1, Ordering::SeqCst), 0);
        Ok(())
    }).output().unwrap() };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());
    assert_eq!(mem.load(Ordering::SeqCst), 0);
}
<commit_msg>rustfmt the test<commit_after>#![allow(stable_features)]
// ignore-windows - this is a unix-specific test
// ignore-cloudabi no processes
// ignore-emscripten no processes
#![feature(process_exec, rustc_private)]

extern crate libc;

use std::env;
use std::io::Error;
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

fn main() {
    if let Some(arg) = env::args().nth(1) {
        match &arg[..] {
            "test1" => println!("hello2"),
            "test2" => assert_eq!(env::var("FOO").unwrap(), "BAR"),
            "test3" => assert_eq!(env::current_dir().unwrap().to_str().unwrap(), "/"),
            "empty" => {}
            _ => panic!("unknown argument: {}", arg),
        }
        return;
    }

    let me = env::current_exe().unwrap();

    let output = unsafe {
        Command::new(&me)
            .arg("test1")
            .pre_exec(|| {
                println!("hello");
                Ok(())
            })
            .output()
            .unwrap()
    };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert_eq!(output.stdout, b"hello\nhello2\n");

    let output = unsafe {
        Command::new(&me)
            .arg("test2")
            .pre_exec(|| {
                env::set_var("FOO", "BAR");
                Ok(())
            })
            .output()
            .unwrap()
    };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());

    let output = unsafe {
        Command::new(&me)
            .arg("test3")
            .pre_exec(|| {
                env::set_current_dir("/").unwrap();
                Ok(())
            })
            .output()
            .unwrap()
    };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());

    let output = unsafe {
        Command::new(&me)
            .arg("bad")
            .pre_exec(|| Err(Error::from_raw_os_error(102)))
            .output()
            .unwrap_err()
    };
    assert_eq!(output.raw_os_error(), Some(102));

    let pid = unsafe { libc::getpid() };
    assert!(pid >= 0);
    let output = unsafe {
        Command::new(&me)
            .arg("empty")
            .pre_exec(move || {
                let child = libc::getpid();
                assert!(child >= 0);
                assert!(pid != child);
                Ok(())
            })
            .output()
            .unwrap()
    };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());

    let mem = Arc::new(AtomicUsize::new(0));
    let mem2 = mem.clone();
    let output = unsafe {
        Command::new(&me)
            .arg("empty")
            .pre_exec(move || {
                assert_eq!(mem2.fetch_add(1, Ordering::SeqCst), 0);
                Ok(())
            })
            .output()
            .unwrap()
    };
    assert!(output.status.success());
    assert!(output.stderr.is_empty());
    assert!(output.stdout.is_empty());
    assert_eq!(mem.load(Ordering::SeqCst), 0);
}
<|endoftext|>