<commit_before>use libc::{c_char,  c_uint};
use ffi::target_machine::{self, LLVMTargetRef};
use ffi::target::{self, LLVMTargetDataRef, LLVMOpaqueTargetData};
use std::ffi::CString;
use std::fmt;
use util::{self, CBox, DisposeRef};

pub struct TargetData;
native_ref!(&TargetData = LLVMTargetDataRef);

impl TargetData {
    /// Creates target data from a target layout string
    pub fn new(rep: &str) -> CBox<TargetData> {
        let c_rep = CString::new(rep).unwrap();
        CBox::new(unsafe {
            target::LLVMCreateTargetData(c_rep.as_ptr())
        }.into())
    }
    /// Returns true if the target is big endian
    pub fn is_big_endian(&self) -> bool {
        let order = unsafe { target::LLVMByteOrder(self.into()) } as c_uint;
        order == 0
    }
    /// Returns the size of a pointer on the target
    pub fn get_pointer_size(&self) -> usize {
        unsafe { target::LLVMPointerSize(self.into()) as usize }
    }
    /// Convers this to a target layout string
    pub fn as_str(&self) -> CBox<str> {
        unsafe {
            CBox::new(target::LLVMCopyStringRepOfTargetData(self.into()))
        }
    }
}
impl fmt::Display for TargetData {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(&self.as_str())
    }
}

impl DisposeRef for TargetData {
    type RefTo = LLVMOpaqueTargetData;
    unsafe fn dispose(ptr: LLVMTargetDataRef) {
        target::LLVMDisposeTargetData(ptr)
    }
}

pub struct Target;
native_ref!(&Target = LLVMTargetRef);
impl Target {
    /// Get the name of this target
    pub fn get_name(&self) -> &str {
        unsafe { util::to_str(target_machine::LLVMGetTargetName(self.into()) as *mut c_char) }
    }
    /// Get the description of this target
    pub fn get_description(&self) -> &str {
        unsafe { util::to_str(target_machine::LLVMGetTargetDescription(self.into()) as *mut c_char) }
    }
    pub fn has_asm_backend(&self) -> bool {
        unsafe { target_machine::LLVMTargetHasAsmBackend(self.into()) != 0 }
    }
    pub fn has_jit(&self) -> bool {
        unsafe { target_machine::LLVMTargetHasJIT(self.into()) != 0 }
    }
    pub fn has_target_machine(&self) -> bool {
        unsafe { target_machine::LLVMTargetHasTargetMachine(self.into()) != 0 }
    }
}
<commit_msg>Add `size_of...` and `alignment_of` methods to `TargetData` for getting the ABI details of types<commit_after>use libc::{c_char,  c_uint};
use ffi::target_machine::{self, LLVMTargetRef};
use ffi::target::{self, LLVMTargetDataRef, LLVMOpaqueTargetData};
use std::ffi::CString;
use std::fmt;
use ty::Type;
use util::{self, CBox, DisposeRef};

pub struct TargetData;
native_ref!(&TargetData = LLVMTargetDataRef);

impl TargetData {
    /// Creates target data from a target layout string
    pub fn new(rep: &str) -> CBox<TargetData> {
        let c_rep = CString::new(rep).unwrap();
        CBox::new(unsafe {
            target::LLVMCreateTargetData(c_rep.as_ptr())
        }.into())
    }
    /// Returns true if the target is big endian
    pub fn is_big_endian(&self) -> bool {
        let order = unsafe { target::LLVMByteOrder(self.into()) } as c_uint;
        order == 0
    }
    /// Returns the size of a pointer on the target
    pub fn get_pointer_size(&self) -> usize {
        unsafe { target::LLVMPointerSize(self.into()) as usize }
    }
    /// Returns the size of the type given in bits
    pub fn size_of_in_bits(&self, ty: &Type) -> u64 {
        unsafe { target::LLVMSizeOfTypeInBits(self.into(), ty.into()) }
    }
    /// Returns the size of the type given in bytes
    pub fn size_of(&self, ty: &Type) -> u64 {
        unsafe { target::LLVMStoreSizeOfType(self.into(), ty.into()) }
    }
    /// Returns the alignment of the type given in bytes
    pub fn alignment_of(&self, ty: &Type) -> usize {
        unsafe { target::LLVMABIAlignmentOfType(self.into(), ty.into()) as usize }
    }
    /// Convers this to a target layout string
    pub fn as_str(&self) -> CBox<str> {
        unsafe {
            CBox::new(target::LLVMCopyStringRepOfTargetData(self.into()))
        }
    }
}
impl fmt::Display for TargetData {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(&self.as_str())
    }
}

impl DisposeRef for TargetData {
    type RefTo = LLVMOpaqueTargetData;
    unsafe fn dispose(ptr: LLVMTargetDataRef) {
        target::LLVMDisposeTargetData(ptr)
    }
}

pub struct Target;
native_ref!(&Target = LLVMTargetRef);
impl Target {
    /// Get the name of this target
    pub fn get_name(&self) -> &str {
        unsafe { util::to_str(target_machine::LLVMGetTargetName(self.into()) as *mut c_char) }
    }
    /// Get the description of this target
    pub fn get_description(&self) -> &str {
        unsafe { util::to_str(target_machine::LLVMGetTargetDescription(self.into()) as *mut c_char) }
    }
    pub fn has_asm_backend(&self) -> bool {
        unsafe { target_machine::LLVMTargetHasAsmBackend(self.into()) != 0 }
    }
    pub fn has_jit(&self) -> bool {
        unsafe { target_machine::LLVMTargetHasJIT(self.into()) != 0 }
    }
    pub fn has_target_machine(&self) -> bool {
        unsafe { target_machine::LLVMTargetHasTargetMachine(self.into()) != 0 }
    }
}
<|endoftext|>


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

/// Generic Oak error.
#[derive(Debug)]
pub enum OakError {
    ProtobufError(protobuf::ProtobufError),
    OakStatus(crate::OakStatus),
    IoError(std::io::Error),
}

impl From<protobuf::ProtobufError> for OakError {
    fn from(err: protobuf::ProtobufError) -> Self {
        OakError::ProtobufError(err)
    }
}

impl From<std::io::Error> for OakError {
    fn from(err: std::io::Error) -> Self {
        OakError::IoError(err)
    }
}

impl From<crate::OakStatus> for OakError {
    fn from(status: crate::OakStatus) -> Self {
        OakError::OakStatus(status)
    }
}
<commit_msg>sdk: implement Error trait for oak::OakError<commit_after>//
// Copyright 2020 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

/// Generic Oak error.
#[derive(Debug)]
pub enum OakError {
    ProtobufError(protobuf::ProtobufError),
    OakStatus(crate::OakStatus),
    IoError(std::io::Error),
}

impl std::fmt::Display for OakError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            OakError::ProtobufError(e) => write!(f, "protobuf error: {}", e),
            OakError::OakStatus(e) => write!(f, "Oak status value: {:?}", e),
            OakError::IoError(e) => write!(f, "I/O error: {}", e),
        }
    }
}

impl std::error::Error for OakError {}

impl From<protobuf::ProtobufError> for OakError {
    fn from(err: protobuf::ProtobufError) -> Self {
        OakError::ProtobufError(err)
    }
}

impl From<std::io::Error> for OakError {
    fn from(err: std::io::Error) -> Self {
        OakError::IoError(err)
    }
}

impl From<crate::OakStatus> for OakError {
    fn from(status: crate::OakStatus) -> Self {
        OakError::OakStatus(status)
    }
}
<|endoftext|>


<commit_before>use grain_capnp::{PowerboxCapability, UiView, UiSession};
use web_session_capnp::{WebSession};

use capnp::capability::{ClientHook, FromServer};
use capnp::AnyPointer;
use capnp_rpc::rpc::{RpcConnectionState, SturdyRefRestorer};
use capnp_rpc::capability::{LocalClient};

pub struct UiViewImpl;

impl PowerboxCapability::Server for UiViewImpl {
    fn get_powerbox_info(&mut self, context : PowerboxCapability::GetPowerboxInfoContext) {
        context.done()
    }
}

impl UiView::Server for UiViewImpl {
    fn get_view_info(&mut self, context : UiView::GetViewInfoContext) {
        context.done()
    }

    fn new_session(&mut self, mut context : UiView::NewSessionContext) {
        println!("asked for a new session!");
        let (_, results) = context.get();
        let client : UiSession::Client = FromServer::new(None::<LocalClient>, ~WebSessionImpl);
        results.set_session(client);
        context.done()
    }
}

pub struct WebSessionImpl;

impl UiSession::Server for WebSessionImpl {

}

impl WebSession::Server for WebSessionImpl {
    fn get(&mut self, context : WebSession::GetContext) {
        println!("GET");
        context.done()
    }
    fn post(&mut self, context : WebSession::PostContext) {
        println!("POST");
        context.done()
    }
    fn put(&mut self, context : WebSession::PutContext) {
        println!("PUT");
        context.done()
    }
    fn delete(&mut self, context : WebSession::DeleteContext) {
        println!("DELETE");
        context.done()
    }
    fn open_web_socket(&mut self, context : WebSession::OpenWebSocketContext) {
        println!("OPEN WEB SOCKET");
        context.done()
    }
}

pub struct FdStream {
    inner : ~::std::rt::rtio::RtioFileStream:Send,
}

impl FdStream {
    pub fn new(fd : ::libc::c_int) -> ::std::io::IoResult<FdStream> {
        ::std::rt::rtio::LocalIo::maybe_raise(|io| {
            Ok (FdStream { inner : io.fs_from_raw_fd(fd, ::std::rt::rtio::DontClose) })
        })
    }
}

impl Reader for FdStream {
    fn read(&mut self, buf : &mut [u8]) -> ::std::io::IoResult<uint> {
        self.inner.read(buf).map(|i| i as uint)
    }
}

impl Writer for FdStream {
    fn write(&mut self, buf : &[u8]) -> ::std::io::IoResult<()> {
        self.inner.write(buf)
    }
}

pub struct Restorer;

impl SturdyRefRestorer for Restorer {
    fn restore(&self, obj_id : AnyPointer::Reader) -> Option<~ClientHook:Send> {
        if obj_id.is_null() {
            let client : UiView::Client = FromServer::new(None::<LocalClient>, ~UiViewImpl);
            Some(client.client.hook)
        } else {
            None
        }
    }
}

pub fn main() -> ::std::io::IoResult<()> {
    let ifs = try!(FdStream::new(3));
    let ofs = try!(FdStream::new(3));

    let connection_state = RpcConnectionState::new();
    connection_state.run(ifs, ofs, Restorer);

    Ok(())
}

<commit_msg>downcast dance<commit_after>use grain_capnp::{PowerboxCapability, UiView, UiSession};
use web_session_capnp::{WebSession};

use capnp::capability::{ClientHook, FromServer};
use capnp::AnyPointer;
use capnp_rpc::rpc::{RpcConnectionState, SturdyRefRestorer};
use capnp_rpc::capability::{LocalClient};

pub struct UiViewImpl;

impl PowerboxCapability::Server for UiViewImpl {
    fn get_powerbox_info(&mut self, context : PowerboxCapability::GetPowerboxInfoContext) {
        context.done()
    }
}

impl UiView::Server for UiViewImpl {
    fn get_view_info(&mut self, context : UiView::GetViewInfoContext) {
        context.done()
    }

    fn new_session(&mut self, mut context : UiView::NewSessionContext) {
        println!("asked for a new session!");
        let (_, results) = context.get();
        let client : WebSession::Client = FromServer::new(None::<LocalClient>, ~WebSessionImpl);
        results.set_session(UiSession::Client { client : client.client});
        context.done()
    }
}

pub struct WebSessionImpl;

impl UiSession::Server for WebSessionImpl {

}

impl WebSession::Server for WebSessionImpl {
    fn get(&mut self, context : WebSession::GetContext) {
        println!("GET");
        context.done()
    }
    fn post(&mut self, context : WebSession::PostContext) {
        println!("POST");
        context.done()
    }
    fn put(&mut self, context : WebSession::PutContext) {
        println!("PUT");
        context.done()
    }
    fn delete(&mut self, context : WebSession::DeleteContext) {
        println!("DELETE");
        context.done()
    }
    fn open_web_socket(&mut self, context : WebSession::OpenWebSocketContext) {
        println!("OPEN WEB SOCKET");
        context.done()
    }
}

pub struct FdStream {
    inner : ~::std::rt::rtio::RtioFileStream:Send,
}

impl FdStream {
    pub fn new(fd : ::libc::c_int) -> ::std::io::IoResult<FdStream> {
        ::std::rt::rtio::LocalIo::maybe_raise(|io| {
            Ok (FdStream { inner : io.fs_from_raw_fd(fd, ::std::rt::rtio::DontClose) })
        })
    }
}

impl Reader for FdStream {
    fn read(&mut self, buf : &mut [u8]) -> ::std::io::IoResult<uint> {
        self.inner.read(buf).map(|i| i as uint)
    }
}

impl Writer for FdStream {
    fn write(&mut self, buf : &[u8]) -> ::std::io::IoResult<()> {
        self.inner.write(buf)
    }
}

pub struct Restorer;

impl SturdyRefRestorer for Restorer {
    fn restore(&self, obj_id : AnyPointer::Reader) -> Option<~ClientHook:Send> {
        if obj_id.is_null() {
            let client : UiView::Client = FromServer::new(None::<LocalClient>, ~UiViewImpl);
            Some(client.client.hook)
        } else {
            None
        }
    }
}

pub fn main() -> ::std::io::IoResult<()> {
    let ifs = try!(FdStream::new(3));
    let ofs = try!(FdStream::new(3));

    let connection_state = RpcConnectionState::new();
    connection_state.run(ifs, ofs, Restorer);

    Ok(())
}

<|endoftext|>


<commit_before>use std::net::{TcpStream, TcpListener};
use std::io::Write;
use std::thread;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

pub static PROXY_PORT: AtomicUsize = ATOMIC_USIZE_INIT;

pub fn init() {
    thread::spawn(move || {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port() as usize;

        PROXY_PORT.fetch_add(port, Ordering::Release);

        for stream in listener.incoming() {
            match stream {
                Ok(stream) => handle_client(stream),
                Err(e)     => println!("Error: {}", e)
            }
        }

        drop(listener);
    });

    while PROXY_PORT.load(Ordering::Acquire) == 0 {}
}

fn handle_client(mut stream: TcpStream) {
    let response = "HTTP/1.1 200 OK\n\nHello world";

    stream.write(response.as_bytes()).unwrap();
}
<commit_msg>Aligned<commit_after>use std::net::{TcpStream, TcpListener};
use std::io::Write;
use std::thread;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

pub static PROXY_PORT: AtomicUsize = ATOMIC_USIZE_INIT;

pub fn init() {
    thread::spawn(move || {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port     = listener.local_addr().unwrap().port() as usize;

        PROXY_PORT.fetch_add(port, Ordering::Release);

        for stream in listener.incoming() {
            match stream {
                Ok(stream) => handle_client(stream),
                Err(e)     => println!("Error: {}", e)
            }
        }

        drop(listener);
    });

    while PROXY_PORT.load(Ordering::Acquire) == 0 {}
}

fn handle_client(mut stream: TcpStream) {
    let response = "HTTP/1.1 200 OK\n\nHello world";

    stream.write(response.as_bytes()).unwrap();
}
<|endoftext|>


<commit_before>use std::net::{TcpStream, TcpListener};
use std::io::Write;
use std::thread;
use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering, ATOMIC_USIZE_INIT, ATOMIC_BOOL_INIT};

pub static PORT: AtomicUsize = ATOMIC_USIZE_INIT;
pub static SERVER_THREAD_SPAWNED: AtomicBool = ATOMIC_BOOL_INIT;

pub fn start() {
    // Ensures only one server is running.
    if is_server_thread_spawned() { return }

    thread::spawn(move || {
        SERVER_THREAD_SPAWNED.store(true, Ordering::SeqCst);

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port     = listener.local_addr().unwrap().port() as usize;

        PORT.store(port, Ordering::SeqCst);

        for stream in listener.incoming() {
            match stream {
                Ok(stream) => handle_client(stream),
                Err(e)     => println!("Error: {}", e)
            }
        }

        drop(listener);
    });

    while !is_running() {}
}

pub fn is_server_thread_spawned() -> bool {
    SERVER_THREAD_SPAWNED.load(Ordering::SeqCst)
}

pub fn is_running() -> bool {
    port() != 0
}

pub fn port() -> usize {
    PORT.load(Ordering::SeqCst)
}

pub fn host() -> String {
    format!("127.0.0.1:{}", port())
}

pub fn host_with_protocol() -> String {
    format!("http://127.0.0.1:{}", port())
}

fn handle_client(mut stream: TcpStream) {
    let response = "HTTP/1.1 200 OK\n\nHello world";

    stream.write(response.as_bytes()).unwrap();
}

#[cfg(test)]
mod tests {
    use server;

    #[test]
    fn test_start_server() {
        server::start();

        assert!(server::is_server_thread_spawned());
        assert!(server::is_running());
        assert!(server::port() != 0);
    }
}
<commit_msg>Server tests<commit_after>use std::net::{TcpStream, TcpListener};
use std::io::Write;
use std::thread;
use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering, ATOMIC_USIZE_INIT, ATOMIC_BOOL_INIT};

pub static PORT: AtomicUsize = ATOMIC_USIZE_INIT;
pub static SERVER_THREAD_SPAWNED: AtomicBool = ATOMIC_BOOL_INIT;

pub fn start() {
    // Ensures only one server is running.
    if is_server_thread_spawned() { wait_for_listener(); return }

    thread::spawn(move || {
        SERVER_THREAD_SPAWNED.store(true, Ordering::SeqCst);

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port     = listener.local_addr().unwrap().port() as usize;

        PORT.store(port, Ordering::SeqCst);

        for stream in listener.incoming() {
            match stream {
                Ok(stream) => handle_client(stream),
                Err(e)     => println!("Error: {}", e)
            }
        }

        drop(listener);
    });

    // Wait for the server to start listening.
    wait_for_listener()
}

pub fn is_server_thread_spawned() -> bool {
    SERVER_THREAD_SPAWNED.load(Ordering::SeqCst)
}

pub fn is_listening() -> bool {
    port() != 0
}

pub fn port() -> usize {
    PORT.load(Ordering::SeqCst)
}

pub fn host() -> String {
    format!("127.0.0.1:{}", port())
}

pub fn host_with_protocol() -> String {
    format!("http://127.0.0.1:{}", port())
}

fn wait_for_listener() {
    while !is_listening() {}
}

fn handle_client(mut stream: TcpStream) {
    let response = "HTTP/1.1 200 OK\n\nHello world";

    stream.write(response.as_bytes()).unwrap();
}

#[cfg(test)]
mod tests {
    use server;
    use std::net::TcpStream;

    #[test]
    fn test_start_server() {
        server::start();

        assert!(server::is_server_thread_spawned());
        assert!(server::is_listening());
        assert!(server::port() != 0);
    }

    #[test]
    fn test_started_server_is_listening() {
        server::start();

        let host = server::host();
        let stream = TcpStream::connect(&*host);

        assert!(stream.is_ok());
    }
}
<|endoftext|>


<commit_before>use std::collections::HashSet;

use syn;
use syn::punctuated::{Pair, Punctuated};

use internals::ast::{Container, Data};
use internals::{attr, ungroup};

use proc_macro2::Span;

// Remove the default from every type parameter because in the generated impls
// they look like associated types: "error: associated type bindings are not
// allowed here".
pub fn without_defaults(generics: &syn::Generics) -> syn::Generics {
    syn::Generics {
        params: generics
            .params
            .iter()
            .map(|param| match param {
                syn::GenericParam::Type(param) => syn::GenericParam::Type(syn::TypeParam {
                    eq_token: None,
                    default: None,
                    ..param.clone()
                }),
                _ => param.clone(),
            })
            .collect(),
        ..generics.clone()
    }
}

pub fn with_where_predicates(
    generics: &syn::Generics,
    predicates: &[syn::WherePredicate],
) -> syn::Generics {
    let mut generics = generics.clone();
    generics
        .make_where_clause()
        .predicates
        .extend(predicates.iter().cloned());
    generics
}

pub fn with_where_predicates_from_fields(
    cont: &Container,
    generics: &syn::Generics,
    from_field: fn(&attr::Field) -> Option<&[syn::WherePredicate]>,
) -> syn::Generics {
    let predicates = cont
        .data
        .all_fields()
        .filter_map(|field| from_field(&field.attrs))
        .flat_map(|predicates| predicates.to_vec());

    let mut generics = generics.clone();
    generics.make_where_clause().predicates.extend(predicates);
    generics
}

pub fn with_where_predicates_from_variants(
    cont: &Container,
    generics: &syn::Generics,
    from_variant: fn(&attr::Variant) -> Option<&[syn::WherePredicate]>,
) -> syn::Generics {
    let variants = match &cont.data {
        Data::Enum(variants) => variants,
        Data::Struct(_, _) => {
            return generics.clone();
        }
    };

    let predicates = variants
        .iter()
        .filter_map(|variant| from_variant(&variant.attrs))
        .flat_map(|predicates| predicates.to_vec());

    let mut generics = generics.clone();
    generics.make_where_clause().predicates.extend(predicates);
    generics
}

// Puts the given bound on any generic type parameters that are used in fields
// for which filter returns true.
//
// For example, the following struct needs the bound `A: Serialize, B:
// Serialize`.
//
//     struct S<'b, A, B: 'b, C> {
//         a: A,
//         b: Option<&'b B>
//         #[serde(skip_serializing)]
//         c: C,
//     }
pub fn with_bound(
    cont: &Container,
    generics: &syn::Generics,
    filter: fn(&attr::Field, Option<&attr::Variant>) -> bool,
    bound: &syn::Path,
) -> syn::Generics {
    struct FindTyParams<'ast> {
        // Set of all generic type parameters on the current struct (A, B, C in
        // the example). Initialized up front.
        all_type_params: HashSet<syn::Ident>,

        // Set of generic type parameters used in fields for which filter
        // returns true (A and B in the example). Filled in as the visitor sees
        // them.
        relevant_type_params: HashSet<syn::Ident>,

        // Fields whose type is an associated type of one of the generic type
        // parameters.
        associated_type_usage: Vec<&'ast syn::TypePath>,
    }

    impl<'ast> FindTyParams<'ast> {
        fn visit_field(&mut self, field: &'ast syn::Field) {
            if let syn::Type::Path(ty) = ungroup(&field.ty) {
                if let Some(Pair::Punctuated(t, _)) = ty.path.segments.pairs().next() {
                    if self.all_type_params.contains(&t.ident) {
                        self.associated_type_usage.push(ty);
                    }
                }
            }
            self.visit_type(&field.ty);
        }

        fn visit_path(&mut self, path: &'ast syn::Path) {
            if let Some(seg) = path.segments.last() {
                if seg.ident == "PhantomData" {
                    // Hardcoded exception, because PhantomData<T> implements
                    // Serialize and Deserialize whether or not T implements it.
                    return;
                }
            }
            if path.leading_colon.is_none() && path.segments.len() == 1 {
                let id = &path.segments[0].ident;
                if self.all_type_params.contains(id) {
                    self.relevant_type_params.insert(id.clone());
                }
            }
            for segment in &path.segments {
                self.visit_path_segment(segment);
            }
        }

        // Everything below is simply traversing the syntax tree.

        fn visit_type(&mut self, ty: &'ast syn::Type) {
            match ty {
                syn::Type::Array(ty) => self.visit_type(&ty.elem),
                syn::Type::BareFn(ty) => {
                    for arg in &ty.inputs {
                        self.visit_type(&arg.ty);
                    }
                    self.visit_return_type(&ty.output);
                }
                syn::Type::Group(ty) => self.visit_type(&ty.elem),
                syn::Type::ImplTrait(ty) => {
                    for bound in &ty.bounds {
                        self.visit_type_param_bound(bound);
                    }
                }
                syn::Type::Macro(ty) => self.visit_macro(&ty.mac),
                syn::Type::Paren(ty) => self.visit_type(&ty.elem),
                syn::Type::Path(ty) => {
                    if let Some(qself) = &ty.qself {
                        self.visit_type(&qself.ty);
                    }
                    self.visit_path(&ty.path);
                }
                syn::Type::Ptr(ty) => self.visit_type(&ty.elem),
                syn::Type::Reference(ty) => self.visit_type(&ty.elem),
                syn::Type::Slice(ty) => self.visit_type(&ty.elem),
                syn::Type::TraitObject(ty) => {
                    for bound in &ty.bounds {
                        self.visit_type_param_bound(bound);
                    }
                }
                syn::Type::Tuple(ty) => {
                    for elem in &ty.elems {
                        self.visit_type(elem);
                    }
                }

                syn::Type::Infer(_) | syn::Type::Never(_) | syn::Type::Verbatim(_) => {}

                #[cfg(test)]
                syn::Type::__TestExhaustive(_) => unimplemented!(),
                #[cfg(not(test))]
                _ => {}
            }
        }

        fn visit_path_segment(&mut self, segment: &'ast syn::PathSegment) {
            self.visit_path_arguments(&segment.arguments);
        }

        fn visit_path_arguments(&mut self, arguments: &'ast syn::PathArguments) {
            match arguments {
                syn::PathArguments::None => {}
                syn::PathArguments::AngleBracketed(arguments) => {
                    for arg in &arguments.args {
                        match arg {
                            syn::GenericArgument::Type(arg) => self.visit_type(arg),
                            syn::GenericArgument::Binding(arg) => self.visit_type(&arg.ty),
                            syn::GenericArgument::Lifetime(_)
                            | syn::GenericArgument::Constraint(_)
                            | syn::GenericArgument::Const(_) => {}
                        }
                    }
                }
                syn::PathArguments::Parenthesized(arguments) => {
                    for argument in &arguments.inputs {
                        self.visit_type(argument);
                    }
                    self.visit_return_type(&arguments.output);
                }
            }
        }

        fn visit_return_type(&mut self, return_type: &'ast syn::ReturnType) {
            match return_type {
                syn::ReturnType::Default => {}
                syn::ReturnType::Type(_, output) => self.visit_type(output),
            }
        }

        fn visit_type_param_bound(&mut self, bound: &'ast syn::TypeParamBound) {
            match bound {
                syn::TypeParamBound::Trait(bound) => self.visit_path(&bound.path),
                syn::TypeParamBound::Lifetime(_) => {}
            }
        }

        // Type parameter should not be considered used by a macro path.
        //
        //     struct TypeMacro<T> {
        //         mac: T!(),
        //         marker: PhantomData<T>,
        //     }
        fn visit_macro(&mut self, _mac: &'ast syn::Macro) {}
    }

    let all_type_params = generics
        .type_params()
        .map(|param| param.ident.clone())
        .collect();

    let mut visitor = FindTyParams {
        all_type_params,
        relevant_type_params: HashSet::new(),
        associated_type_usage: Vec::new(),
    };
    match &cont.data {
        Data::Enum(variants) => {
            for variant in variants.iter() {
                let relevant_fields = variant
                    .fields
                    .iter()
                    .filter(|field| filter(&field.attrs, Some(&variant.attrs)));
                for field in relevant_fields {
                    visitor.visit_field(field.original);
                }
            }
        }
        Data::Struct(_, fields) => {
            for field in fields.iter().filter(|field| filter(&field.attrs, None)) {
                visitor.visit_field(field.original);
            }
        }
    }

    let relevant_type_params = visitor.relevant_type_params;
    let associated_type_usage = visitor.associated_type_usage;
    let new_predicates = generics
        .type_params()
        .map(|param| param.ident.clone())
        .filter(|id| relevant_type_params.contains(id))
        .map(|id| syn::TypePath {
            qself: None,
            path: id.into(),
        })
        .chain(associated_type_usage.into_iter().cloned())
        .map(|bounded_ty| {
            syn::WherePredicate::Type(syn::PredicateType {
                lifetimes: None,
                // the type parameter that is being bounded e.g. T
                bounded_ty: syn::Type::Path(bounded_ty),
                colon_token: <Token![:]>::default(),
                // the bound e.g. Serialize
                bounds: vec![syn::TypeParamBound::Trait(syn::TraitBound {
                    paren_token: None,
                    modifier: syn::TraitBoundModifier::None,
                    lifetimes: None,
                    path: bound.clone(),
                })]
                .into_iter()
                .collect(),
            })
        });

    let mut generics = generics.clone();
    generics
        .make_where_clause()
        .predicates
        .extend(new_predicates);
    generics
}

pub fn with_self_bound(
    cont: &Container,
    generics: &syn::Generics,
    bound: &syn::Path,
) -> syn::Generics {
    let mut generics = generics.clone();
    generics
        .make_where_clause()
        .predicates
        .push(syn::WherePredicate::Type(syn::PredicateType {
            lifetimes: None,
            // the type that is being bounded e.g. MyStruct<'a, T>
            bounded_ty: type_of_item(cont),
            colon_token: <Token![:]>::default(),
            // the bound e.g. Default
            bounds: vec![syn::TypeParamBound::Trait(syn::TraitBound {
                paren_token: None,
                modifier: syn::TraitBoundModifier::None,
                lifetimes: None,
                path: bound.clone(),
            })]
            .into_iter()
            .collect(),
        }));
    generics
}

pub fn with_lifetime_bound(generics: &syn::Generics, lifetime: &str) -> syn::Generics {
    let bound = syn::Lifetime::new(lifetime, Span::call_site());
    let def = syn::LifetimeDef {
        attrs: Vec::new(),
        lifetime: bound.clone(),
        colon_token: None,
        bounds: Punctuated::new(),
    };

    let params = Some(syn::GenericParam::Lifetime(def))
        .into_iter()
        .chain(generics.params.iter().cloned().map(|mut param| {
            match &mut param {
                syn::GenericParam::Lifetime(param) => {
                    param.bounds.push(bound.clone());
                }
                syn::GenericParam::Type(param) => {
                    param
                        .bounds
                        .push(syn::TypeParamBound::Lifetime(bound.clone()));
                }
                syn::GenericParam::Const(_) => {}
            }
            param
        }))
        .collect();

    syn::Generics {
        params,
        ..generics.clone()
    }
}

fn type_of_item(cont: &Container) -> syn::Type {
    syn::Type::Path(syn::TypePath {
        qself: None,
        path: syn::Path {
            leading_colon: None,
            segments: vec![syn::PathSegment {
                ident: cont.ident.clone(),
                arguments: syn::PathArguments::AngleBracketed(
                    syn::AngleBracketedGenericArguments {
                        colon2_token: None,
                        lt_token: <Token![<]>::default(),
                        args: cont
                            .generics
                            .params
                            .iter()
                            .map(|param| match param {
                                syn::GenericParam::Type(param) => {
                                    syn::GenericArgument::Type(syn::Type::Path(syn::TypePath {
                                        qself: None,
                                        path: param.ident.clone().into(),
                                    }))
                                }
                                syn::GenericParam::Lifetime(param) => {
                                    syn::GenericArgument::Lifetime(param.lifetime.clone())
                                }
                                syn::GenericParam::Const(_) => {
                                    panic!("Serde does not support const generics yet");
                                }
                            })
                            .collect(),
                        gt_token: <Token![>]>::default(),
                    },
                ),
            }]
            .into_iter()
            .collect(),
        },
    })
}
<commit_msg>Resolve redundant_closure_for_method_calls clippy lints<commit_after>use std::collections::HashSet;

use syn;
use syn::punctuated::{Pair, Punctuated};

use internals::ast::{Container, Data};
use internals::{attr, ungroup};

use proc_macro2::Span;

// Remove the default from every type parameter because in the generated impls
// they look like associated types: "error: associated type bindings are not
// allowed here".
pub fn without_defaults(generics: &syn::Generics) -> syn::Generics {
    syn::Generics {
        params: generics
            .params
            .iter()
            .map(|param| match param {
                syn::GenericParam::Type(param) => syn::GenericParam::Type(syn::TypeParam {
                    eq_token: None,
                    default: None,
                    ..param.clone()
                }),
                _ => param.clone(),
            })
            .collect(),
        ..generics.clone()
    }
}

pub fn with_where_predicates(
    generics: &syn::Generics,
    predicates: &[syn::WherePredicate],
) -> syn::Generics {
    let mut generics = generics.clone();
    generics
        .make_where_clause()
        .predicates
        .extend(predicates.iter().cloned());
    generics
}

pub fn with_where_predicates_from_fields(
    cont: &Container,
    generics: &syn::Generics,
    from_field: fn(&attr::Field) -> Option<&[syn::WherePredicate]>,
) -> syn::Generics {
    let predicates = cont
        .data
        .all_fields()
        .filter_map(|field| from_field(&field.attrs))
        .flat_map(<[syn::WherePredicate]>::to_vec);

    let mut generics = generics.clone();
    generics.make_where_clause().predicates.extend(predicates);
    generics
}

pub fn with_where_predicates_from_variants(
    cont: &Container,
    generics: &syn::Generics,
    from_variant: fn(&attr::Variant) -> Option<&[syn::WherePredicate]>,
) -> syn::Generics {
    let variants = match &cont.data {
        Data::Enum(variants) => variants,
        Data::Struct(_, _) => {
            return generics.clone();
        }
    };

    let predicates = variants
        .iter()
        .filter_map(|variant| from_variant(&variant.attrs))
        .flat_map(<[syn::WherePredicate]>::to_vec);

    let mut generics = generics.clone();
    generics.make_where_clause().predicates.extend(predicates);
    generics
}

// Puts the given bound on any generic type parameters that are used in fields
// for which filter returns true.
//
// For example, the following struct needs the bound `A: Serialize, B:
// Serialize`.
//
//     struct S<'b, A, B: 'b, C> {
//         a: A,
//         b: Option<&'b B>
//         #[serde(skip_serializing)]
//         c: C,
//     }
pub fn with_bound(
    cont: &Container,
    generics: &syn::Generics,
    filter: fn(&attr::Field, Option<&attr::Variant>) -> bool,
    bound: &syn::Path,
) -> syn::Generics {
    struct FindTyParams<'ast> {
        // Set of all generic type parameters on the current struct (A, B, C in
        // the example). Initialized up front.
        all_type_params: HashSet<syn::Ident>,

        // Set of generic type parameters used in fields for which filter
        // returns true (A and B in the example). Filled in as the visitor sees
        // them.
        relevant_type_params: HashSet<syn::Ident>,

        // Fields whose type is an associated type of one of the generic type
        // parameters.
        associated_type_usage: Vec<&'ast syn::TypePath>,
    }

    impl<'ast> FindTyParams<'ast> {
        fn visit_field(&mut self, field: &'ast syn::Field) {
            if let syn::Type::Path(ty) = ungroup(&field.ty) {
                if let Some(Pair::Punctuated(t, _)) = ty.path.segments.pairs().next() {
                    if self.all_type_params.contains(&t.ident) {
                        self.associated_type_usage.push(ty);
                    }
                }
            }
            self.visit_type(&field.ty);
        }

        fn visit_path(&mut self, path: &'ast syn::Path) {
            if let Some(seg) = path.segments.last() {
                if seg.ident == "PhantomData" {
                    // Hardcoded exception, because PhantomData<T> implements
                    // Serialize and Deserialize whether or not T implements it.
                    return;
                }
            }
            if path.leading_colon.is_none() && path.segments.len() == 1 {
                let id = &path.segments[0].ident;
                if self.all_type_params.contains(id) {
                    self.relevant_type_params.insert(id.clone());
                }
            }
            for segment in &path.segments {
                self.visit_path_segment(segment);
            }
        }

        // Everything below is simply traversing the syntax tree.

        fn visit_type(&mut self, ty: &'ast syn::Type) {
            match ty {
                syn::Type::Array(ty) => self.visit_type(&ty.elem),
                syn::Type::BareFn(ty) => {
                    for arg in &ty.inputs {
                        self.visit_type(&arg.ty);
                    }
                    self.visit_return_type(&ty.output);
                }
                syn::Type::Group(ty) => self.visit_type(&ty.elem),
                syn::Type::ImplTrait(ty) => {
                    for bound in &ty.bounds {
                        self.visit_type_param_bound(bound);
                    }
                }
                syn::Type::Macro(ty) => self.visit_macro(&ty.mac),
                syn::Type::Paren(ty) => self.visit_type(&ty.elem),
                syn::Type::Path(ty) => {
                    if let Some(qself) = &ty.qself {
                        self.visit_type(&qself.ty);
                    }
                    self.visit_path(&ty.path);
                }
                syn::Type::Ptr(ty) => self.visit_type(&ty.elem),
                syn::Type::Reference(ty) => self.visit_type(&ty.elem),
                syn::Type::Slice(ty) => self.visit_type(&ty.elem),
                syn::Type::TraitObject(ty) => {
                    for bound in &ty.bounds {
                        self.visit_type_param_bound(bound);
                    }
                }
                syn::Type::Tuple(ty) => {
                    for elem in &ty.elems {
                        self.visit_type(elem);
                    }
                }

                syn::Type::Infer(_) | syn::Type::Never(_) | syn::Type::Verbatim(_) => {}

                #[cfg(test)]
                syn::Type::__TestExhaustive(_) => unimplemented!(),
                #[cfg(not(test))]
                _ => {}
            }
        }

        fn visit_path_segment(&mut self, segment: &'ast syn::PathSegment) {
            self.visit_path_arguments(&segment.arguments);
        }

        fn visit_path_arguments(&mut self, arguments: &'ast syn::PathArguments) {
            match arguments {
                syn::PathArguments::None => {}
                syn::PathArguments::AngleBracketed(arguments) => {
                    for arg in &arguments.args {
                        match arg {
                            syn::GenericArgument::Type(arg) => self.visit_type(arg),
                            syn::GenericArgument::Binding(arg) => self.visit_type(&arg.ty),
                            syn::GenericArgument::Lifetime(_)
                            | syn::GenericArgument::Constraint(_)
                            | syn::GenericArgument::Const(_) => {}
                        }
                    }
                }
                syn::PathArguments::Parenthesized(arguments) => {
                    for argument in &arguments.inputs {
                        self.visit_type(argument);
                    }
                    self.visit_return_type(&arguments.output);
                }
            }
        }

        fn visit_return_type(&mut self, return_type: &'ast syn::ReturnType) {
            match return_type {
                syn::ReturnType::Default => {}
                syn::ReturnType::Type(_, output) => self.visit_type(output),
            }
        }

        fn visit_type_param_bound(&mut self, bound: &'ast syn::TypeParamBound) {
            match bound {
                syn::TypeParamBound::Trait(bound) => self.visit_path(&bound.path),
                syn::TypeParamBound::Lifetime(_) => {}
            }
        }

        // Type parameter should not be considered used by a macro path.
        //
        //     struct TypeMacro<T> {
        //         mac: T!(),
        //         marker: PhantomData<T>,
        //     }
        fn visit_macro(&mut self, _mac: &'ast syn::Macro) {}
    }

    let all_type_params = generics
        .type_params()
        .map(|param| param.ident.clone())
        .collect();

    let mut visitor = FindTyParams {
        all_type_params,
        relevant_type_params: HashSet::new(),
        associated_type_usage: Vec::new(),
    };
    match &cont.data {
        Data::Enum(variants) => {
            for variant in variants.iter() {
                let relevant_fields = variant
                    .fields
                    .iter()
                    .filter(|field| filter(&field.attrs, Some(&variant.attrs)));
                for field in relevant_fields {
                    visitor.visit_field(field.original);
                }
            }
        }
        Data::Struct(_, fields) => {
            for field in fields.iter().filter(|field| filter(&field.attrs, None)) {
                visitor.visit_field(field.original);
            }
        }
    }

    let relevant_type_params = visitor.relevant_type_params;
    let associated_type_usage = visitor.associated_type_usage;
    let new_predicates = generics
        .type_params()
        .map(|param| param.ident.clone())
        .filter(|id| relevant_type_params.contains(id))
        .map(|id| syn::TypePath {
            qself: None,
            path: id.into(),
        })
        .chain(associated_type_usage.into_iter().cloned())
        .map(|bounded_ty| {
            syn::WherePredicate::Type(syn::PredicateType {
                lifetimes: None,
                // the type parameter that is being bounded e.g. T
                bounded_ty: syn::Type::Path(bounded_ty),
                colon_token: <Token![:]>::default(),
                // the bound e.g. Serialize
                bounds: vec![syn::TypeParamBound::Trait(syn::TraitBound {
                    paren_token: None,
                    modifier: syn::TraitBoundModifier::None,
                    lifetimes: None,
                    path: bound.clone(),
                })]
                .into_iter()
                .collect(),
            })
        });

    let mut generics = generics.clone();
    generics
        .make_where_clause()
        .predicates
        .extend(new_predicates);
    generics
}

pub fn with_self_bound(
    cont: &Container,
    generics: &syn::Generics,
    bound: &syn::Path,
) -> syn::Generics {
    let mut generics = generics.clone();
    generics
        .make_where_clause()
        .predicates
        .push(syn::WherePredicate::Type(syn::PredicateType {
            lifetimes: None,
            // the type that is being bounded e.g. MyStruct<'a, T>
            bounded_ty: type_of_item(cont),
            colon_token: <Token![:]>::default(),
            // the bound e.g. Default
            bounds: vec![syn::TypeParamBound::Trait(syn::TraitBound {
                paren_token: None,
                modifier: syn::TraitBoundModifier::None,
                lifetimes: None,
                path: bound.clone(),
            })]
            .into_iter()
            .collect(),
        }));
    generics
}

pub fn with_lifetime_bound(generics: &syn::Generics, lifetime: &str) -> syn::Generics {
    let bound = syn::Lifetime::new(lifetime, Span::call_site());
    let def = syn::LifetimeDef {
        attrs: Vec::new(),
        lifetime: bound.clone(),
        colon_token: None,
        bounds: Punctuated::new(),
    };

    let params = Some(syn::GenericParam::Lifetime(def))
        .into_iter()
        .chain(generics.params.iter().cloned().map(|mut param| {
            match &mut param {
                syn::GenericParam::Lifetime(param) => {
                    param.bounds.push(bound.clone());
                }
                syn::GenericParam::Type(param) => {
                    param
                        .bounds
                        .push(syn::TypeParamBound::Lifetime(bound.clone()));
                }
                syn::GenericParam::Const(_) => {}
            }
            param
        }))
        .collect();

    syn::Generics {
        params,
        ..generics.clone()
    }
}

fn type_of_item(cont: &Container) -> syn::Type {
    syn::Type::Path(syn::TypePath {
        qself: None,
        path: syn::Path {
            leading_colon: None,
            segments: vec![syn::PathSegment {
                ident: cont.ident.clone(),
                arguments: syn::PathArguments::AngleBracketed(
                    syn::AngleBracketedGenericArguments {
                        colon2_token: None,
                        lt_token: <Token![<]>::default(),
                        args: cont
                            .generics
                            .params
                            .iter()
                            .map(|param| match param {
                                syn::GenericParam::Type(param) => {
                                    syn::GenericArgument::Type(syn::Type::Path(syn::TypePath {
                                        qself: None,
                                        path: param.ident.clone().into(),
                                    }))
                                }
                                syn::GenericParam::Lifetime(param) => {
                                    syn::GenericArgument::Lifetime(param.lifetime.clone())
                                }
                                syn::GenericParam::Const(_) => {
                                    panic!("Serde does not support const generics yet");
                                }
                            })
                            .collect(),
                        gt_token: <Token![>]>::default(),
                    },
                ),
            }]
            .into_iter()
            .collect(),
        },
    })
}
<|endoftext|>


<commit_before>use bollard::service::{HostConfig, PortBinding};
use bollard::{
    container::{Config, CreateContainerOptions, ListContainersOptions, StartContainerOptions},
    errors::Error as BollardError,
    image::CreateImageOptions,
    Docker,
};
use elasticsearch::{
    http::transport::BuildError as TransportBuilderError,
    indices::{IndicesDeleteAliasParts, IndicesDeleteIndexTemplateParts, IndicesDeleteParts},
    Error as ElasticsearchError,
};
use futures::stream::TryStreamExt;
use lazy_static::lazy_static;
use snafu::{ResultExt, Snafu};
use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};
use tokio::time::{sleep, Duration};

use crate::adapters::secondary::elasticsearch::remote::{self, Error as ElasticsearchRemoteError};
use crate::domain::ports::secondary::remote::{Error as RemoteError, Remote};

lazy_static! {
    pub static ref AVAILABLE: Arc<Mutex<()>> = Arc::new(Mutex::new(()));
}

// This function takes the mutex related to the availability of the elasticsearch docker container,
// and then creates or reset the container depending on its availability.
// Finally, it tests if elasticsearch is available, before returning
// the mutex guard.
// TODO Document the usage of the mutexguard, how you should avoid running asserts
// while holding the mutex.
pub async fn initialize() -> Result<MutexGuard<'static, ()>, Error> {
    let guard = AVAILABLE.lock().unwrap();
    let mut docker = DockerWrapper::new();
    let is_available = docker.is_container_available().await?;
    if !is_available {
        docker.create_container().await?;
    } else {
        docker.cleanup().await?;
    }
    let is_available = docker.is_container_available().await?;
    if !is_available {
        return Err(Error::Misc {
            msg: format!("Cannot get docker {} available", docker.docker_image),
        });
    }
    let pool = remote::connection_test_pool()
        .await
        .context(ElasticsearchPoolCreation)?;
    let _client = pool.conn().await.context(ElasticsearchConnection)?;

    Ok(guard)
}

#[derive(Debug, Snafu)]
pub enum Error {
    #[snafu(display("Connection to docker socket: {}", source))]
    DockerConnection { source: BollardError },

    #[snafu(display("Creation of elasticsearch pool: {}", source))]
    ElasticsearchPoolCreation { source: ElasticsearchRemoteError },

    #[snafu(display("Connection to elasticsearch: {}", source))]
    ElasticsearchConnection { source: RemoteError },

    #[snafu(display("docker version: {}", source))]
    Version { source: BollardError },

    #[snafu(display("url parsing error: {}", source))]
    UrlParse { source: url::ParseError },

    #[snafu(display("elasticsearch transport error: {}", source))]
    ElasticsearchTransport { source: TransportBuilderError },

    #[snafu(display("elasticsearch client error: {}", source))]
    ElasticsearchClient { source: ElasticsearchError },

    #[snafu(display("docker error: {}", source))]
    DockerEngine { source: BollardError },

    #[snafu(display("error: {}", msg))]
    Misc { msg: String },
}

pub struct DockerWrapper {
    ports: Vec<(u32, u32)>, // list of ports to publish (host port, container port)
    docker_image: String,
    container_name: String, // ip: String,
}

// FIXME: Here we hardcode elasticsearch ports. These should be had from the ELASTICSEARCH_TEST_URL
impl Default for DockerWrapper {
    fn default() -> Self {
        DockerWrapper {
            ports: vec![(9201, 9200), (9301, 9300)],
            docker_image: String::from("docker.elastic.co/elasticsearch/elasticsearch:7.13.0"),
            container_name: String::from("mimir-test-elasticsearch"),
        }
    }
}

impl DockerWrapper {
    pub fn new() -> DockerWrapper {
        DockerWrapper::default()
    }

    // Returns true if the container self.container_name is running
    // TODO Probably should run a check on Elasticsearch status
    pub async fn is_container_available(&mut self) -> Result<bool, Error> {
        let docker = Docker::connect_with_unix(
            "unix:///var/run/docker.sock",
            120,
            &bollard::ClientVersion {
                major_version: 1,
                minor_version: 24,
            },
        )
        .context(DockerConnection)?;

        let docker = &docker.negotiate_version().await.context(Version)?;

        docker.version().await.context(Version)?;

        let mut filters = HashMap::new();
        filters.insert("name", vec![self.container_name.as_str()]);

        let options = Some(ListContainersOptions {
            all: false, // only running containers
            filters,
            ..Default::default()
        });

        let containers = docker
            .list_containers(options)
            .await
            .context(DockerEngine)?;

        Ok(!containers.is_empty())
    }

    // If the container is already created, then start it.
    // If it is not created, then create it and start it.
    pub async fn create_container(&mut self) -> Result<(), Error> {
        let docker = Docker::connect_with_unix(
            "unix:///var/run/docker.sock",
            120,
            &bollard::ClientVersion {
                major_version: 1,
                minor_version: 24,
            },
        )
        .context(DockerConnection)?;

        let docker = docker.negotiate_version().await.context(Version)?;

        let _ = docker.version().await.context(Version);

        let mut filters = HashMap::new();
        filters.insert("name", vec![self.container_name.as_str()]);

        let options = Some(ListContainersOptions {
            all: true, // only running containers
            filters,
            ..Default::default()
        });

        let containers = docker
            .list_containers(options)
            .await
            .context(DockerEngine)?;

        if containers.is_empty() {
            let options = CreateContainerOptions {
                name: &self.container_name,
            };

            let mut port_bindings = HashMap::new();
            for (host_port, container_port) in self.ports.iter() {
                port_bindings.insert(
                    format!("{}/tcp", &container_port),
                    Some(vec![PortBinding {
                        host_ip: Some(String::from("0.0.0.0")),
                        host_port: Some(host_port.to_string()),
                    }]),
                );
            }

            let host_config = HostConfig {
                port_bindings: Some(port_bindings),
                ..Default::default()
            };

            let mut exposed_ports = HashMap::new();
            self.ports.iter().for_each(|(_, container)| {
                let v: HashMap<(), ()> = HashMap::new();
                exposed_ports.insert(format!("{}/tcp", container), v);
            });

            let env_vars = vec![String::from("discovery.type=single-node")];

            let config = Config {
                image: Some(self.docker_image.clone()),
                exposed_ports: Some(exposed_ports),
                host_config: Some(host_config),
                env: Some(env_vars),
                ..Default::default()
            };

            docker
                .create_image(
                    Some(CreateImageOptions {
                        from_image: self.docker_image.clone(),
                        ..Default::default()
                    }),
                    None,
                    None,
                )
                .try_collect::<Vec<_>>()
                .await
                .context(DockerEngine)?;

            let _ = docker
                .create_container(Some(options), config)
                .await
                .context(DockerEngine)?;

            sleep(Duration::from_secs(5)).await;
        }
        let _ = docker
            .start_container(&self.container_name, None::<StartContainerOptions<String>>)
            .await
            .context(DockerEngine)?;

        sleep(Duration::from_secs(15)).await;

        Ok(())
    }

    async fn cleanup(&mut self) -> Result<(), Error> {
        let pool = remote::connection_test_pool()
            .await
            .context(ElasticsearchPoolCreation)?;
        let client = pool.conn().await.context(ElasticsearchConnection)?;

        let _ = client
            .0
            .indices()
            .delete(IndicesDeleteParts::Index(&["*"]))
            .send()
            .await
            .context(ElasticsearchClient)?;

        let _ = client
            .0
            .indices()
            .delete_alias(IndicesDeleteAliasParts::IndexName(&["*"], &["*"]))
            .send()
            .await
            .context(ElasticsearchClient)?;

        let _ = client
            .0
            .indices()
            .delete_index_template(IndicesDeleteIndexTemplateParts::Name("*"))
            .send()
            .await
            .context(ElasticsearchClient)?;

        println!("done with cleanup");
        sleep(Duration::from_secs(1)).await;
        Ok(())
    }

    async fn _drop(&mut self) {
        if std::env::var("DONT_KILL_THE_WHALE") == Ok("1".to_string()) {
            println!(
                "the docker won't be stoped at the end, you can debug it.
                Note: ES has been mapped to the port 9242 in you localhost
                manually stop and rm the container mimirsbrunn_tests after debug"
            );
            return;
        }
        let docker = Docker::connect_with_unix(
            "unix:///var/run/docker.sock",
            120,
            &bollard::ClientVersion {
                major_version: 1,
                minor_version: 24,
            },
        )
        .expect("docker connection");

        let options = Some(bollard::container::StopContainerOptions { t: 0 });
        docker
            .stop_container(&self.container_name, options)
            .await
            .expect("stop container");

        let options = Some(bollard::container::RemoveContainerOptions {
            force: true,
            ..Default::default()
        });

        let _res = docker
            .remove_container(&self.container_name, options)
            .await
            .expect("remove container");
    }
}
<commit_msg>WIP Try longer sleep<commit_after>use bollard::service::{HostConfig, PortBinding};
use bollard::{
    container::{Config, CreateContainerOptions, ListContainersOptions, StartContainerOptions},
    errors::Error as BollardError,
    image::CreateImageOptions,
    Docker,
};
use elasticsearch::{
    http::transport::BuildError as TransportBuilderError,
    indices::{IndicesDeleteAliasParts, IndicesDeleteIndexTemplateParts, IndicesDeleteParts},
    Error as ElasticsearchError,
};
use futures::stream::TryStreamExt;
use lazy_static::lazy_static;
use snafu::{ResultExt, Snafu};
use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};
use tokio::time::{sleep, Duration};

use crate::adapters::secondary::elasticsearch::remote::{self, Error as ElasticsearchRemoteError};
use crate::domain::ports::secondary::remote::{Error as RemoteError, Remote};

lazy_static! {
    pub static ref AVAILABLE: Arc<Mutex<()>> = Arc::new(Mutex::new(()));
}

// This function takes the mutex related to the availability of the elasticsearch docker container,
// and then creates or reset the container depending on its availability.
// Finally, it tests if elasticsearch is available, before returning
// the mutex guard.
// TODO Document the usage of the mutexguard, how you should avoid running asserts
// while holding the mutex.
pub async fn initialize() -> Result<MutexGuard<'static, ()>, Error> {
    let guard = AVAILABLE.lock().unwrap();
    let mut docker = DockerWrapper::new();
    let is_available = docker.is_container_available().await?;
    if !is_available {
        docker.create_container().await?;
    } else {
        docker.cleanup().await?;
    }
    let is_available = docker.is_container_available().await?;
    if !is_available {
        return Err(Error::Misc {
            msg: format!("Cannot get docker {} available", docker.docker_image),
        });
    }
    let pool = remote::connection_test_pool()
        .await
        .context(ElasticsearchPoolCreation)?;
    let _client = pool.conn().await.context(ElasticsearchConnection)?;

    Ok(guard)
}

#[derive(Debug, Snafu)]
pub enum Error {
    #[snafu(display("Connection to docker socket: {}", source))]
    DockerConnection { source: BollardError },

    #[snafu(display("Creation of elasticsearch pool: {}", source))]
    ElasticsearchPoolCreation { source: ElasticsearchRemoteError },

    #[snafu(display("Connection to elasticsearch: {}", source))]
    ElasticsearchConnection { source: RemoteError },

    #[snafu(display("docker version: {}", source))]
    Version { source: BollardError },

    #[snafu(display("url parsing error: {}", source))]
    UrlParse { source: url::ParseError },

    #[snafu(display("elasticsearch transport error: {}", source))]
    ElasticsearchTransport { source: TransportBuilderError },

    #[snafu(display("elasticsearch client error: {}", source))]
    ElasticsearchClient { source: ElasticsearchError },

    #[snafu(display("docker error: {}", source))]
    DockerEngine { source: BollardError },

    #[snafu(display("error: {}", msg))]
    Misc { msg: String },
}

pub struct DockerWrapper {
    ports: Vec<(u32, u32)>, // list of ports to publish (host port, container port)
    docker_image: String,
    container_name: String, // ip: String,
}

// FIXME: Here we hardcode elasticsearch ports. These should be had from the ELASTICSEARCH_TEST_URL
impl Default for DockerWrapper {
    fn default() -> Self {
        DockerWrapper {
            ports: vec![(9201, 9200), (9301, 9300)],
            docker_image: String::from("docker.elastic.co/elasticsearch/elasticsearch:7.13.0"),
            container_name: String::from("mimir-test-elasticsearch"),
        }
    }
}

impl DockerWrapper {
    pub fn new() -> DockerWrapper {
        DockerWrapper::default()
    }

    // Returns true if the container self.container_name is running
    // TODO Probably should run a check on Elasticsearch status
    pub async fn is_container_available(&mut self) -> Result<bool, Error> {
        let docker = Docker::connect_with_unix(
            "unix:///var/run/docker.sock",
            120,
            &bollard::ClientVersion {
                major_version: 1,
                minor_version: 24,
            },
        )
        .context(DockerConnection)?;

        let docker = &docker.negotiate_version().await.context(Version)?;

        docker.version().await.context(Version)?;

        let mut filters = HashMap::new();
        filters.insert("name", vec![self.container_name.as_str()]);

        let options = Some(ListContainersOptions {
            all: false, // only running containers
            filters,
            ..Default::default()
        });

        let containers = docker
            .list_containers(options)
            .await
            .context(DockerEngine)?;

        Ok(!containers.is_empty())
    }

    // If the container is already created, then start it.
    // If it is not created, then create it and start it.
    pub async fn create_container(&mut self) -> Result<(), Error> {
        let docker = Docker::connect_with_unix(
            "unix:///var/run/docker.sock",
            120,
            &bollard::ClientVersion {
                major_version: 1,
                minor_version: 24,
            },
        )
        .context(DockerConnection)?;

        let docker = docker.negotiate_version().await.context(Version)?;

        let _ = docker.version().await.context(Version);

        let mut filters = HashMap::new();
        filters.insert("name", vec![self.container_name.as_str()]);

        let options = Some(ListContainersOptions {
            all: true, // only running containers
            filters,
            ..Default::default()
        });

        let containers = docker
            .list_containers(options)
            .await
            .context(DockerEngine)?;

        if containers.is_empty() {
            let options = CreateContainerOptions {
                name: &self.container_name,
            };

            let mut port_bindings = HashMap::new();
            for (host_port, container_port) in self.ports.iter() {
                port_bindings.insert(
                    format!("{}/tcp", &container_port),
                    Some(vec![PortBinding {
                        host_ip: Some(String::from("0.0.0.0")),
                        host_port: Some(host_port.to_string()),
                    }]),
                );
            }

            let host_config = HostConfig {
                port_bindings: Some(port_bindings),
                ..Default::default()
            };

            let mut exposed_ports = HashMap::new();
            self.ports.iter().for_each(|(_, container)| {
                let v: HashMap<(), ()> = HashMap::new();
                exposed_ports.insert(format!("{}/tcp", container), v);
            });

            let env_vars = vec![String::from("discovery.type=single-node")];

            let config = Config {
                image: Some(self.docker_image.clone()),
                exposed_ports: Some(exposed_ports),
                host_config: Some(host_config),
                env: Some(env_vars),
                ..Default::default()
            };

            docker
                .create_image(
                    Some(CreateImageOptions {
                        from_image: self.docker_image.clone(),
                        ..Default::default()
                    }),
                    None,
                    None,
                )
                .try_collect::<Vec<_>>()
                .await
                .context(DockerEngine)?;

            let _ = docker
                .create_container(Some(options), config)
                .await
                .context(DockerEngine)?;

            sleep(Duration::from_secs(5)).await;
        }
        let _ = docker
            .start_container(&self.container_name, None::<StartContainerOptions<String>>)
            .await
            .context(DockerEngine)?;

        sleep(Duration::from_secs(30)).await;

        Ok(())
    }

    async fn cleanup(&mut self) -> Result<(), Error> {
        let pool = remote::connection_test_pool()
            .await
            .context(ElasticsearchPoolCreation)?;
        let client = pool.conn().await.context(ElasticsearchConnection)?;

        let _ = client
            .0
            .indices()
            .delete(IndicesDeleteParts::Index(&["*"]))
            .send()
            .await
            .context(ElasticsearchClient)?;

        let _ = client
            .0
            .indices()
            .delete_alias(IndicesDeleteAliasParts::IndexName(&["*"], &["*"]))
            .send()
            .await
            .context(ElasticsearchClient)?;

        let _ = client
            .0
            .indices()
            .delete_index_template(IndicesDeleteIndexTemplateParts::Name("*"))
            .send()
            .await
            .context(ElasticsearchClient)?;

        println!("done with cleanup");
        sleep(Duration::from_secs(5)).await;
        Ok(())
    }

    async fn _drop(&mut self) {
        if std::env::var("DONT_KILL_THE_WHALE") == Ok("1".to_string()) {
            println!(
                "the docker won't be stoped at the end, you can debug it.
                Note: ES has been mapped to the port 9242 in you localhost
                manually stop and rm the container mimirsbrunn_tests after debug"
            );
            return;
        }
        let docker = Docker::connect_with_unix(
            "unix:///var/run/docker.sock",
            120,
            &bollard::ClientVersion {
                major_version: 1,
                minor_version: 24,
            },
        )
        .expect("docker connection");

        let options = Some(bollard::container::StopContainerOptions { t: 0 });
        docker
            .stop_container(&self.container_name, options)
            .await
            .expect("stop container");

        let options = Some(bollard::container::RemoveContainerOptions {
            force: true,
            ..Default::default()
        });

        let _res = docker
            .remove_container(&self.container_name, options)
            .await
            .expect("remove container");
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{BufWriter, Result, Write};
use std::path::Path;

use super::Object::*;
use super::{Dictionary, Document, Object, Stream, StringFormat};
use crate::xref::*;

impl Document {
	/// Save PDF document to specified file path.
	#[inline]
	pub fn save<P: AsRef<Path>>(&mut self, path: P) -> Result<File> {
		let mut file = BufWriter::new(File::create(path)?);
		self.save_internal(&mut file)?;
		Ok(file.into_inner()?)
	}

	/// Save PDF to arbitrary target
	#[inline]
	pub fn save_to<W: Write>(&mut self, target: &mut W) -> Result<()> {
		self.save_internal(target)
	}

	fn save_internal<W: Write>(&mut self, target: &mut W) -> Result<()> {
		let mut target = CountingWrite { inner: target, bytes_written: 0 };
		let mut xref = Xref::new(self.max_id + 1);
		write!(target, "%PDF-{}\n", self.version)?;

		for (&(id, generation), object) in &self.objects {
			if object.type_name().map(|name| ["ObjStm", "XRef", "Linearized"].contains(&name)) != Some(true) {
				Writer::write_indirect_object(&mut target, id, generation, object, &mut xref)?;
			}
		}

		let xref_start = target.bytes_written;
		Writer::write_xref(&mut target, &xref)?;
		self.write_trailer(&mut target)?;
		write!(target, "\nstartxref\n{}\n%%EOF", xref_start)?;

		Ok(())
	}

	fn write_trailer(&mut self, file: &mut dyn Write) -> Result<()> {
		self.trailer.set("Size", i64::from(self.max_id + 1));
		file.write_all(b"trailer\n")?;
		Writer::write_dictionary(file, &self.trailer)?;
		Ok(())
	}
}

pub struct Writer;

impl Writer {
	fn need_separator(object: &Object) -> bool {
		match *object {
			Null => true,
			Boolean(_) => true,
			Integer(_) => true,
			Real(_) => true,
			Reference(_) => true,
			_ => false,
		}
	}

	fn need_end_separator(object: &Object) -> bool {
		match *object {
			Null => true,
			Boolean(_) => true,
			Integer(_) => true,
			Real(_) => true,
			Name(_) => true,
			Reference(_) => true,
			Object::Stream(_) => true,
			_ => false,
		}
	}

	fn write_xref(file: &mut dyn Write, xref: &Xref) -> Result<()> {
		file.write_all(b"xref\n")?;
		write!(file, "0 {}\n", xref.size)?;

		let mut write_xref_entry = |offset: u32, generation: u16, kind: char| write!(file, "{:>010} {:>05} {} \n", offset, generation, kind);
		write_xref_entry(0, 65535, 'f')?;

		let mut obj_id = 1;
		while obj_id < xref.size {
			if let Some(entry) = xref.get(obj_id) {
				if let XrefEntry::Normal { offset, generation } = *entry {
					write_xref_entry(offset, generation, 'n')?;
				};
			} else {
				write_xref_entry(0, 65535, 'f')?;
			}
			obj_id += 1;
		}
		Ok(())
	}

	fn write_indirect_object<W: Write>(file: &mut CountingWrite<&mut W>, id: u32, generation: u16, object: &Object, xref: &mut Xref) -> Result<()> {
		let offset = file.bytes_written as u32;
		xref.insert(id, XrefEntry::Normal { offset, generation });
		write!(file, "{} {} obj{}", id, generation, if Writer::need_separator(object) { " " } else { "" })?;
		Writer::write_object(file, object)?;
		write!(file, "{}endobj\n", if Writer::need_end_separator(object) { " " } else { "" })?;
		Ok(())
	}

	pub fn write_object(file: &mut dyn Write, object: &Object) -> Result<()> {
		match *object {
			Null => file.write_all(b"null"),
			Boolean(ref value) => if *value {
				file.write_all(b"true")
			} else {
				file.write_all(b"false")
			},
			Integer(ref value) => {
				let _ = itoa::write(file, *value);
				Ok(())
			}
			Real(ref value) => {
				let _ = dtoa::write(file, *value);
				Ok(())
			}
			Name(ref name) => Writer::write_name(file, name),
			String(ref text, ref format) => Writer::write_string(file, text, format),
			Array(ref array) => Writer::write_array(file, array),
			Object::Dictionary(ref dict) => Writer::write_dictionary(file, dict),
			Object::Stream(ref stream) => Writer::write_stream(file, stream),
			Reference(ref id) => write!(file, "{} {} R", id.0, id.1),
		}
	}

	fn write_name(file: &mut dyn Write, name: &[u8]) -> Result<()> {
		file.write_all(b"/")?;
		for &byte in name {
			// white-space and delimiter chars are encoded to # sequences
			// also encode bytes outside of the range 33 (!) to 126 (~)
			if b" \t\n\r\x0C()<>[]{}/%#".contains(&byte) || byte < 33 || byte > 126 {
				write!(file, "#{:02X}", byte)?;
			} else {
				file.write_all(&[byte])?;
			}
		}
		Ok(())
	}

	fn write_string(file: &mut dyn Write, text: &[u8], format: &StringFormat) -> Result<()> {
		match *format {
			// Within a Literal string, backslash (\) and unbalanced parentheses should be escaped.
			// This rule apply to each individual byte in a string object,
			// whether the string is interpreted as single-byte or multiple-byte character codes.
			// If an end-of-line marker appears within a literal string without a preceding backslash, the result is equivalent to \n.
			// So \r also need be escaped.
			StringFormat::Literal => {
				let mut escape_indice = Vec::new();
				let mut parentheses = Vec::new();
				for (index, &byte) in text.iter().enumerate() {
					match byte {
						b'(' => parentheses.push(index),
						b')' => {
							if !parentheses.is_empty() {
								parentheses.pop();
							} else {
								escape_indice.push(index);
							}
						}
						b'\\' | b'\r' => escape_indice.push(index),
						_ => continue,
					}
				}
				escape_indice.append(&mut parentheses);

				file.write_all(b"(")?;
				if !escape_indice.is_empty() {
					for (index, &byte) in text.iter().enumerate() {
						if escape_indice.contains(&index) {
							file.write_all(b"\\")?;
							file.write_all(&[if byte == b'\r' { b'r' } else { byte }])?;
						} else {
							file.write_all(&[byte])?;
						}
					}
				} else {
					file.write_all(text)?;
				}
				file.write_all(b")")?;
			}
			StringFormat::Hexadecimal => {
				file.write_all(b"<")?;
				for &byte in text {
					write!(file, "{:02X}", byte)?;
				}
				file.write_all(b">")?;
			}
		}
		Ok(())
	}

	fn write_array(file: &mut dyn Write, array: &[Object]) -> Result<()> {
		file.write_all(b"[")?;
		let mut first = true;
		for object in array {
			if first {
				first = false;
			} else if Writer::need_separator(object) {
				file.write_all(b" ")?;
			}
			Writer::write_object(file, object)?;
		}
		file.write_all(b"]")?;
		Ok(())
	}

	fn write_dictionary(file: &mut dyn Write, dictionary: &Dictionary) -> Result<()> {
		file.write_all(b"<<")?;
		for (key, value) in dictionary {
			Writer::write_name(file, key)?;
			if Writer::need_separator(value) {
				file.write_all(b" ")?;
			}
			Writer::write_object(file, value)?;
		}
		file.write_all(b">>")?;
		Ok(())
	}

	fn write_stream(file: &mut dyn Write, stream: &Stream) -> Result<()> {
		Writer::write_dictionary(file, &stream.dict)?;
		file.write_all(b"stream\n")?;
		file.write_all(&stream.content)?;
		file.write_all(b"endstream")?;
		Ok(())
	}
}

pub struct CountingWrite<W: Write> {
	inner: W,
	bytes_written: usize,
}

impl<W: Write> Write for CountingWrite<W> {
	#[inline]
	fn write(&mut self, buffer: &[u8]) -> Result<usize> {
		let result = self.inner.write(buffer);
		if let Ok(bytes) = result {
			self.bytes_written += bytes;
		}
		result
	}

	#[inline]
	fn write_all(&mut self, buffer: &[u8]) -> Result<()> {
		self.bytes_written += buffer.len();
		// If this returns `Err` we can’t know how many bytes were actually written (if any)
		// but that doesn’t matter since we’re gonna abort the entire PDF generation anyway.
		self.inner.write_all(buffer)
	}

	#[inline]
	fn flush(&mut self) -> Result<()> {
		self.inner.flush()
	}
}

#[test]
fn save_document() {
	let mut doc = Document::with_version("1.5");
	doc.objects.insert((1, 0), Null);
	doc.objects.insert((2, 0), Boolean(true));
	doc.objects.insert((3, 0), Integer(3));
	doc.objects.insert((4, 0), Real(0.5));
	doc.objects.insert((5, 0), String("text((\r)".as_bytes().to_vec(), StringFormat::Literal));
	doc.objects.insert((6, 0), String("text((\r)".as_bytes().to_vec(), StringFormat::Hexadecimal));
	doc.objects.insert((7, 0), Name(b"name \t".to_vec()));
	doc.objects.insert((8, 0), Reference((1, 0)));
	doc.objects.insert((9, 2), Array(vec![Integer(1), Integer(2), Integer(3)]));
	doc.objects.insert((11, 0), Stream(Stream::new(Dictionary::new(), vec![0x41, 0x42, 0x43])));
	let mut dict = Dictionary::new();
	dict.set("A", Null);
	dict.set("B", false);
	dict.set("C", Name(b"name".to_vec()));
	doc.objects.insert((12, 0), Object::Dictionary(dict));
	doc.max_id = 12;

	doc.save("test_0_save.pdf").unwrap();
}
<commit_msg>Use writeln where appropriate.<commit_after>use std::fs::File;
use std::io::{BufWriter, Result, Write};
use std::path::Path;

use super::Object::*;
use super::{Dictionary, Document, Object, Stream, StringFormat};
use crate::xref::*;

impl Document {
	/// Save PDF document to specified file path.
	#[inline]
	pub fn save<P: AsRef<Path>>(&mut self, path: P) -> Result<File> {
		let mut file = BufWriter::new(File::create(path)?);
		self.save_internal(&mut file)?;
		Ok(file.into_inner()?)
	}

	/// Save PDF to arbitrary target
	#[inline]
	pub fn save_to<W: Write>(&mut self, target: &mut W) -> Result<()> {
		self.save_internal(target)
	}

	fn save_internal<W: Write>(&mut self, target: &mut W) -> Result<()> {
		let mut target = CountingWrite { inner: target, bytes_written: 0 };
		let mut xref = Xref::new(self.max_id + 1);
		writeln!(target, "%PDF-{}", self.version)?;

		for (&(id, generation), object) in &self.objects {
			if object.type_name().map(|name| ["ObjStm", "XRef", "Linearized"].contains(&name)) != Some(true) {
				Writer::write_indirect_object(&mut target, id, generation, object, &mut xref)?;
			}
		}

		let xref_start = target.bytes_written;
		Writer::write_xref(&mut target, &xref)?;
		self.write_trailer(&mut target)?;
		write!(target, "\nstartxref\n{}\n%%EOF", xref_start)?;

		Ok(())
	}

	fn write_trailer(&mut self, file: &mut dyn Write) -> Result<()> {
		self.trailer.set("Size", i64::from(self.max_id + 1));
		file.write_all(b"trailer\n")?;
		Writer::write_dictionary(file, &self.trailer)?;
		Ok(())
	}
}

pub struct Writer;

impl Writer {
	fn need_separator(object: &Object) -> bool {
		match *object {
			Null => true,
			Boolean(_) => true,
			Integer(_) => true,
			Real(_) => true,
			Reference(_) => true,
			_ => false,
		}
	}

	fn need_end_separator(object: &Object) -> bool {
		match *object {
			Null => true,
			Boolean(_) => true,
			Integer(_) => true,
			Real(_) => true,
			Name(_) => true,
			Reference(_) => true,
			Object::Stream(_) => true,
			_ => false,
		}
	}

	fn write_xref(file: &mut dyn Write, xref: &Xref) -> Result<()> {
		writeln!(file, "xref\n0 {}", xref.size)?;

		let mut write_xref_entry = |offset: u32, generation: u16, kind: char| writeln!(file, "{:>010} {:>05} {} ", offset, generation, kind);
		write_xref_entry(0, 65535, 'f')?;

		let mut obj_id = 1;
		while obj_id < xref.size {
			if let Some(entry) = xref.get(obj_id) {
				if let XrefEntry::Normal { offset, generation } = *entry {
					write_xref_entry(offset, generation, 'n')?;
				};
			} else {
				write_xref_entry(0, 65535, 'f')?;
			}
			obj_id += 1;
		}
		Ok(())
	}

	fn write_indirect_object<W: Write>(file: &mut CountingWrite<&mut W>, id: u32, generation: u16, object: &Object, xref: &mut Xref) -> Result<()> {
		let offset = file.bytes_written as u32;
		xref.insert(id, XrefEntry::Normal { offset, generation });
		write!(file, "{} {} obj{}", id, generation, if Writer::need_separator(object) { " " } else { "" })?;
		Writer::write_object(file, object)?;
		writeln!(file, "{}endobj", if Writer::need_end_separator(object) { " " } else { "" })?;
		Ok(())
	}

	pub fn write_object(file: &mut dyn Write, object: &Object) -> Result<()> {
		match *object {
			Null => file.write_all(b"null"),
			Boolean(ref value) => if *value {
				file.write_all(b"true")
			} else {
				file.write_all(b"false")
			},
			Integer(ref value) => {
				let _ = itoa::write(file, *value);
				Ok(())
			}
			Real(ref value) => {
				let _ = dtoa::write(file, *value);
				Ok(())
			}
			Name(ref name) => Writer::write_name(file, name),
			String(ref text, ref format) => Writer::write_string(file, text, format),
			Array(ref array) => Writer::write_array(file, array),
			Object::Dictionary(ref dict) => Writer::write_dictionary(file, dict),
			Object::Stream(ref stream) => Writer::write_stream(file, stream),
			Reference(ref id) => write!(file, "{} {} R", id.0, id.1),
		}
	}

	fn write_name(file: &mut dyn Write, name: &[u8]) -> Result<()> {
		file.write_all(b"/")?;
		for &byte in name {
			// white-space and delimiter chars are encoded to # sequences
			// also encode bytes outside of the range 33 (!) to 126 (~)
			if b" \t\n\r\x0C()<>[]{}/%#".contains(&byte) || byte < 33 || byte > 126 {
				write!(file, "#{:02X}", byte)?;
			} else {
				file.write_all(&[byte])?;
			}
		}
		Ok(())
	}

	fn write_string(file: &mut dyn Write, text: &[u8], format: &StringFormat) -> Result<()> {
		match *format {
			// Within a Literal string, backslash (\) and unbalanced parentheses should be escaped.
			// This rule apply to each individual byte in a string object,
			// whether the string is interpreted as single-byte or multiple-byte character codes.
			// If an end-of-line marker appears within a literal string without a preceding backslash, the result is equivalent to \n.
			// So \r also need be escaped.
			StringFormat::Literal => {
				let mut escape_indice = Vec::new();
				let mut parentheses = Vec::new();
				for (index, &byte) in text.iter().enumerate() {
					match byte {
						b'(' => parentheses.push(index),
						b')' => {
							if !parentheses.is_empty() {
								parentheses.pop();
							} else {
								escape_indice.push(index);
							}
						}
						b'\\' | b'\r' => escape_indice.push(index),
						_ => continue,
					}
				}
				escape_indice.append(&mut parentheses);

				file.write_all(b"(")?;
				if !escape_indice.is_empty() {
					for (index, &byte) in text.iter().enumerate() {
						if escape_indice.contains(&index) {
							file.write_all(b"\\")?;
							file.write_all(&[if byte == b'\r' { b'r' } else { byte }])?;
						} else {
							file.write_all(&[byte])?;
						}
					}
				} else {
					file.write_all(text)?;
				}
				file.write_all(b")")?;
			}
			StringFormat::Hexadecimal => {
				file.write_all(b"<")?;
				for &byte in text {
					write!(file, "{:02X}", byte)?;
				}
				file.write_all(b">")?;
			}
		}
		Ok(())
	}

	fn write_array(file: &mut dyn Write, array: &[Object]) -> Result<()> {
		file.write_all(b"[")?;
		let mut first = true;
		for object in array {
			if first {
				first = false;
			} else if Writer::need_separator(object) {
				file.write_all(b" ")?;
			}
			Writer::write_object(file, object)?;
		}
		file.write_all(b"]")?;
		Ok(())
	}

	fn write_dictionary(file: &mut dyn Write, dictionary: &Dictionary) -> Result<()> {
		file.write_all(b"<<")?;
		for (key, value) in dictionary {
			Writer::write_name(file, key)?;
			if Writer::need_separator(value) {
				file.write_all(b" ")?;
			}
			Writer::write_object(file, value)?;
		}
		file.write_all(b">>")?;
		Ok(())
	}

	fn write_stream(file: &mut dyn Write, stream: &Stream) -> Result<()> {
		Writer::write_dictionary(file, &stream.dict)?;
		file.write_all(b"stream\n")?;
		file.write_all(&stream.content)?;
		file.write_all(b"endstream")?;
		Ok(())
	}
}

pub struct CountingWrite<W: Write> {
	inner: W,
	bytes_written: usize,
}

impl<W: Write> Write for CountingWrite<W> {
	#[inline]
	fn write(&mut self, buffer: &[u8]) -> Result<usize> {
		let result = self.inner.write(buffer);
		if let Ok(bytes) = result {
			self.bytes_written += bytes;
		}
		result
	}

	#[inline]
	fn write_all(&mut self, buffer: &[u8]) -> Result<()> {
		self.bytes_written += buffer.len();
		// If this returns `Err` we can’t know how many bytes were actually written (if any)
		// but that doesn’t matter since we’re gonna abort the entire PDF generation anyway.
		self.inner.write_all(buffer)
	}

	#[inline]
	fn flush(&mut self) -> Result<()> {
		self.inner.flush()
	}
}

#[test]
fn save_document() {
	let mut doc = Document::with_version("1.5");
	doc.objects.insert((1, 0), Null);
	doc.objects.insert((2, 0), Boolean(true));
	doc.objects.insert((3, 0), Integer(3));
	doc.objects.insert((4, 0), Real(0.5));
	doc.objects.insert((5, 0), String("text((\r)".as_bytes().to_vec(), StringFormat::Literal));
	doc.objects.insert((6, 0), String("text((\r)".as_bytes().to_vec(), StringFormat::Hexadecimal));
	doc.objects.insert((7, 0), Name(b"name \t".to_vec()));
	doc.objects.insert((8, 0), Reference((1, 0)));
	doc.objects.insert((9, 2), Array(vec![Integer(1), Integer(2), Integer(3)]));
	doc.objects.insert((11, 0), Stream(Stream::new(Dictionary::new(), vec![0x41, 0x42, 0x43])));
	let mut dict = Dictionary::new();
	dict.set("A", Null);
	dict.set("B", false);
	dict.set("C", Name(b"name".to_vec()));
	doc.objects.insert((12, 0), Object::Dictionary(dict));
	doc.max_id = 12;

	doc.save("test_0_save.pdf").unwrap();
}
<|endoftext|>