<commit_before>use std::ptr;

use cairo::{ImageSurface, Status};
use cairo_sys;
use gdk_pixbuf::{PixbufLoader, PixbufLoaderExt};
use glib;
use glib::translate::*;
use glib_sys;
use libc;

use css::{self, CssStyles, RsvgCssStyles};
use defs::{Defs, RsvgDefs};
use error::LoadingError;
use surface_utils::shared_surface::SharedImageSurface;
use util::utf8_cstr;

pub enum RsvgHandle {}

#[allow(improper_ctypes)]
extern "C" {
    fn rsvg_handle_get_defs(handle: *const RsvgHandle) -> *const RsvgDefs;

    fn rsvg_handle_resolve_uri(
        handle: *const RsvgHandle,
        uri: *const libc::c_char,
    ) -> *const libc::c_char;

    fn rsvg_handle_load_extern(
        handle: *const RsvgHandle,
        href: *const libc::c_char,
    ) -> *const RsvgHandle;

    fn rsvg_handle_get_css_styles(handle: *const RsvgHandle) -> *mut RsvgCssStyles;

    fn _rsvg_handle_acquire_data(
        handle: *mut RsvgHandle,
        href: *const libc::c_char,
        out_content_type: *mut *mut libc::c_char,
        out_len: *mut usize,
        error: *mut *mut glib_sys::GError,
    ) -> *mut u8;

    fn rsvg_handle_keep_image_data(handle: *const RsvgHandle) -> glib_sys::gboolean;

    fn rsvg_load_handle_xml_xinclude(
        handle: *mut RsvgHandle,
        href: *const libc::c_char,
    ) -> glib_sys::gboolean;
}

pub fn get_defs<'a>(handle: *const RsvgHandle) -> &'a mut Defs {
    unsafe {
        let d = rsvg_handle_get_defs(handle);
        &mut *(d as *mut Defs)
    }
}

pub fn resolve_uri(handle: *const RsvgHandle, uri: &str) -> Option<String> {
    unsafe {
        let resolved = rsvg_handle_resolve_uri(handle, uri.to_glib_none().0);
        if resolved.is_null() {
            None
        } else {
            Some(from_glib_full(resolved))
        }
    }
}

pub fn load_extern(handle: *const RsvgHandle, uri: &str) -> *const RsvgHandle {
    unsafe { rsvg_handle_load_extern(handle, uri.to_glib_none().0) }
}

pub fn get_css_styles<'a>(handle: *const RsvgHandle) -> &'a CssStyles {
    unsafe { &*(rsvg_handle_get_css_styles(handle) as *const CssStyles) }
}

pub fn get_css_styles_mut<'a>(handle: *const RsvgHandle) -> &'a mut CssStyles {
    unsafe { &mut *(rsvg_handle_get_css_styles(handle) as *mut CssStyles) }
}

pub struct BinaryData {
    pub data: Vec<u8>,
    pub content_type: Option<String>,
}

pub fn acquire_data(handle: *mut RsvgHandle, href: &str) -> Result<BinaryData, glib::Error> {
    unsafe {
        let mut content_type: *mut libc::c_char = ptr::null_mut();
        let mut len = 0;
        let mut error = ptr::null_mut();

        let buf = _rsvg_handle_acquire_data(
            handle,
            href.to_glib_none().0,
            &mut content_type as *mut *mut _,
            &mut len,
            &mut error,
        );

        if buf.is_null() {
            if error.is_null() && len == 0 {
                Ok(BinaryData {
                    data: Vec::new(),
                    content_type: None,
                })
            } else {
                Err(from_glib_full(error))
            }
        } else {
            Ok(BinaryData {
                data: FromGlibContainer::from_glib_full_num(buf as *mut u8, len),
                content_type: from_glib_full(content_type),
            })
        }
    }
}

fn keep_image_data(handle: *const RsvgHandle) -> bool {
    unsafe { from_glib(rsvg_handle_keep_image_data(handle)) }
}

pub fn image_surface_new_from_href(
    handle: *mut RsvgHandle,
    href: &str,
) -> Result<ImageSurface, LoadingError> {
    let data = acquire_data(handle, href)?;

    if data.data.len() == 0 {
        return Err(LoadingError::EmptyData);
    }

    let loader = if let Some(ref content_type) = data.content_type {
        PixbufLoader::new_with_mime_type(content_type)?
    } else {
        PixbufLoader::new()
    };

    loader.write(&data.data)?;
    loader.close()?;

    let pixbuf = loader.get_pixbuf().ok_or(LoadingError::Unknown)?;

    let surface = SharedImageSurface::from_pixbuf(&pixbuf)?.into_image_surface()?;

    if keep_image_data(handle) {
        if let Some(mime_type) = data.content_type {
            extern "C" {
                fn cairo_surface_set_mime_data(
                    surface: *mut cairo_sys::cairo_surface_t,
                    mime_type: *const libc::c_char,
                    data: *mut libc::c_char,
                    length: libc::c_ulong,
                    destroy: cairo_sys::cairo_destroy_func_t,
                    closure: *mut libc::c_void,
                ) -> Status;
            }

            let data_ptr = ToGlibContainerFromSlice::to_glib_full_from_slice(&data.data);

            unsafe {
                let status = cairo_surface_set_mime_data(
                    surface.to_glib_none().0,
                    mime_type.to_glib_none().0,
                    data_ptr as *mut _,
                    data.data.len() as libc::c_ulong,
                    Some(glib_sys::g_free),
                    data_ptr as *mut _,
                );

                if status != Status::Success {
                    return Err(LoadingError::Cairo(status));
                }
            }
        }
    }

    Ok(surface)
}

// FIXME: distinguish between "file not found" and "invalid XML"
pub fn load_xml_xinclude(handle: *mut RsvgHandle, href: &str) -> bool {
    unsafe { from_glib(rsvg_load_handle_xml_xinclude(handle, href.to_glib_none().0)) }
}

pub fn load_css(handle: *mut RsvgHandle, href: &str) {
    if let Ok(data) = acquire_data(handle, href) {
        let BinaryData {
            data: bytes,
            content_type,
        } = data;

        if content_type.as_ref().map(String::as_ref) != Some("text/css") {
            rsvg_log!("\"{}\" is not of type text/css; ignoring", href);
            // FIXME: report errors
            return;
        }

        if let Ok(utf8) = String::from_utf8(bytes) {
            css::parse_into_handle(handle, &utf8);
        } else {
            rsvg_log!(
                "\"{}\" does not contain valid UTF-8 CSS data; ignoring",
                href
            );
            // FIXME: report errors
            return;
        }
    } else {
        rsvg_log!("Could not load \"{}\" for CSS data", href);
        // FIXME: report errors from not being to acquire data; this should be a fatal error
    }
}

// This function just slurps CSS data from a possibly-relative href
// and parses it.  We'll move it to a better place in the end.
#[no_mangle]
pub unsafe extern "C" fn rsvg_handle_load_css(handle: *mut RsvgHandle, href: *const libc::c_char) {
    assert!(!handle.is_null());
    assert!(!href.is_null());

    let href = utf8_cstr(href);
    load_css(handle, href);
}
<commit_msg>handle::acquire_stream(): Bind to Rust<commit_after>use std::ptr;

use cairo::{ImageSurface, Status};
use cairo_sys;
use gdk_pixbuf::{PixbufLoader, PixbufLoaderExt};
use gio::InputStream;
use gio_sys;
use glib;
use glib::translate::*;
use glib_sys;
use libc;

use css::{self, CssStyles, RsvgCssStyles};
use defs::{Defs, RsvgDefs};
use error::LoadingError;
use surface_utils::shared_surface::SharedImageSurface;
use util::utf8_cstr;

pub enum RsvgHandle {}

#[allow(improper_ctypes)]
extern "C" {
    fn rsvg_handle_get_defs(handle: *const RsvgHandle) -> *const RsvgDefs;

    fn rsvg_handle_resolve_uri(
        handle: *const RsvgHandle,
        uri: *const libc::c_char,
    ) -> *const libc::c_char;

    fn rsvg_handle_load_extern(
        handle: *const RsvgHandle,
        href: *const libc::c_char,
    ) -> *const RsvgHandle;

    fn rsvg_handle_get_css_styles(handle: *const RsvgHandle) -> *mut RsvgCssStyles;

    fn _rsvg_handle_acquire_data(
        handle: *mut RsvgHandle,
        href: *const libc::c_char,
        out_content_type: *mut *mut libc::c_char,
        out_len: *mut usize,
        error: *mut *mut glib_sys::GError,
    ) -> *mut u8;

    fn _rsvg_handle_acquire_stream(
        handle: *mut RsvgHandle,
        href: *const libc::c_char,
        error: *mut *mut glib_sys::GError,
    ) -> *mut gio_sys::GInputStream;

    fn rsvg_handle_keep_image_data(handle: *const RsvgHandle) -> glib_sys::gboolean;

    fn rsvg_load_handle_xml_xinclude(
        handle: *mut RsvgHandle,
        href: *const libc::c_char,
    ) -> glib_sys::gboolean;
}

pub fn get_defs<'a>(handle: *const RsvgHandle) -> &'a mut Defs {
    unsafe {
        let d = rsvg_handle_get_defs(handle);
        &mut *(d as *mut Defs)
    }
}

pub fn resolve_uri(handle: *const RsvgHandle, uri: &str) -> Option<String> {
    unsafe {
        let resolved = rsvg_handle_resolve_uri(handle, uri.to_glib_none().0);
        if resolved.is_null() {
            None
        } else {
            Some(from_glib_full(resolved))
        }
    }
}

pub fn load_extern(handle: *const RsvgHandle, uri: &str) -> *const RsvgHandle {
    unsafe { rsvg_handle_load_extern(handle, uri.to_glib_none().0) }
}

pub fn get_css_styles<'a>(handle: *const RsvgHandle) -> &'a CssStyles {
    unsafe { &*(rsvg_handle_get_css_styles(handle) as *const CssStyles) }
}

pub fn get_css_styles_mut<'a>(handle: *const RsvgHandle) -> &'a mut CssStyles {
    unsafe { &mut *(rsvg_handle_get_css_styles(handle) as *mut CssStyles) }
}

pub struct BinaryData {
    pub data: Vec<u8>,
    pub content_type: Option<String>,
}

pub fn acquire_data(handle: *mut RsvgHandle, href: &str) -> Result<BinaryData, glib::Error> {
    unsafe {
        let mut content_type: *mut libc::c_char = ptr::null_mut();
        let mut len = 0;
        let mut error = ptr::null_mut();

        let buf = _rsvg_handle_acquire_data(
            handle,
            href.to_glib_none().0,
            &mut content_type as *mut *mut _,
            &mut len,
            &mut error,
        );

        if buf.is_null() {
            if error.is_null() && len == 0 {
                Ok(BinaryData {
                    data: Vec::new(),
                    content_type: None,
                })
            } else {
                Err(from_glib_full(error))
            }
        } else {
            Ok(BinaryData {
                data: FromGlibContainer::from_glib_full_num(buf as *mut u8, len),
                content_type: from_glib_full(content_type),
            })
        }
    }
}

pub fn acquire_stream(handle: *mut RsvgHandle, href: &str) -> Result<InputStream, glib::Error> {
    unsafe {
        let mut error = ptr::null_mut();

        let stream = _rsvg_handle_acquire_stream(
            handle,
            href.to_glib_none().0,
            &mut error,
        );

        if stream.is_null() {
            Err(from_glib_full(error))
        } else {
            Ok(from_glib_full(stream))
        }
    }
}

fn keep_image_data(handle: *const RsvgHandle) -> bool {
    unsafe { from_glib(rsvg_handle_keep_image_data(handle)) }
}

pub fn image_surface_new_from_href(
    handle: *mut RsvgHandle,
    href: &str,
) -> Result<ImageSurface, LoadingError> {
    let data = acquire_data(handle, href)?;

    if data.data.len() == 0 {
        return Err(LoadingError::EmptyData);
    }

    let loader = if let Some(ref content_type) = data.content_type {
        PixbufLoader::new_with_mime_type(content_type)?
    } else {
        PixbufLoader::new()
    };

    loader.write(&data.data)?;
    loader.close()?;

    let pixbuf = loader.get_pixbuf().ok_or(LoadingError::Unknown)?;

    let surface = SharedImageSurface::from_pixbuf(&pixbuf)?.into_image_surface()?;

    if keep_image_data(handle) {
        if let Some(mime_type) = data.content_type {
            extern "C" {
                fn cairo_surface_set_mime_data(
                    surface: *mut cairo_sys::cairo_surface_t,
                    mime_type: *const libc::c_char,
                    data: *mut libc::c_char,
                    length: libc::c_ulong,
                    destroy: cairo_sys::cairo_destroy_func_t,
                    closure: *mut libc::c_void,
                ) -> Status;
            }

            let data_ptr = ToGlibContainerFromSlice::to_glib_full_from_slice(&data.data);

            unsafe {
                let status = cairo_surface_set_mime_data(
                    surface.to_glib_none().0,
                    mime_type.to_glib_none().0,
                    data_ptr as *mut _,
                    data.data.len() as libc::c_ulong,
                    Some(glib_sys::g_free),
                    data_ptr as *mut _,
                );

                if status != Status::Success {
                    return Err(LoadingError::Cairo(status));
                }
            }
        }
    }

    Ok(surface)
}

// FIXME: distinguish between "file not found" and "invalid XML"
pub fn load_xml_xinclude(handle: *mut RsvgHandle, href: &str) -> bool {
    unsafe { from_glib(rsvg_load_handle_xml_xinclude(handle, href.to_glib_none().0)) }
}

pub fn load_css(handle: *mut RsvgHandle, href: &str) {
    if let Ok(data) = acquire_data(handle, href) {
        let BinaryData {
            data: bytes,
            content_type,
        } = data;

        if content_type.as_ref().map(String::as_ref) != Some("text/css") {
            rsvg_log!("\"{}\" is not of type text/css; ignoring", href);
            // FIXME: report errors
            return;
        }

        if let Ok(utf8) = String::from_utf8(bytes) {
            css::parse_into_handle(handle, &utf8);
        } else {
            rsvg_log!(
                "\"{}\" does not contain valid UTF-8 CSS data; ignoring",
                href
            );
            // FIXME: report errors
            return;
        }
    } else {
        rsvg_log!("Could not load \"{}\" for CSS data", href);
        // FIXME: report errors from not being to acquire data; this should be a fatal error
    }
}

// This function just slurps CSS data from a possibly-relative href
// and parses it.  We'll move it to a better place in the end.
#[no_mangle]
pub unsafe extern "C" fn rsvg_handle_load_css(handle: *mut RsvgHandle, href: *const libc::c_char) {
    assert!(!handle.is_null());
    assert!(!href.is_null());

    let href = utf8_cstr(href);
    load_css(handle, href);
}
<|endoftext|>


<commit_before>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Decodes the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn decode_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Encodes the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn encode_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use crate::wtf8;

    #[test]
    fn decode_empty() {
        assert_eq!(wtf8::decode_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn decode_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn decode_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn decode_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn decode_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn decode_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }
}
<commit_msg>Write tests for the WTF-8 to UTF-16 conversion.<commit_after>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Decodes the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn decode_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Encodes the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn encode_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use crate::wtf8;

    #[test]
    fn decode_empty() {
        assert_eq!(wtf8::decode_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn decode_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn decode_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn decode_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn decode_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn decode_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn encode_empty() {
        assert_eq!(wtf8::encode_ill_formed_utf16(std::iter::empty()), vec![]);
    }

    #[test]
    fn encode_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::encode_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[test]
    fn encode_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::encode_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    data_notification::{DataClientRequest, EpochEndingLedgerInfosRequest, SentDataNotification},
    error::Error,
    stream_progress_tracker::StreamProgressTracker,
    streaming_client::{GetAllEpochEndingLedgerInfosRequest, StreamRequest},
};
use claim::assert_matches;
use diem_data_client::{DataClientPayload, DataClientResponse, GlobalDataSummary};
use std::cmp;
use storage_service_types::CompleteDataRange;

#[test]
fn test_epoch_ending_requests() {
    // Create a new data stream progress tracker
    let stream_progress_tracker = create_epoch_ending_progress_tracker(0, 900);
    let StreamProgressTracker::EpochEndingStreamTracker(mut stream_tracker) =
        stream_progress_tracker;

    // Create a batch of large client requests and verify the result
    let client_requests = stream_tracker
        .create_epoch_ending_client_requests(5, 10000)
        .unwrap();
    let expected_requests = vec![DataClientRequest::EpochEndingLedgerInfos(
        EpochEndingLedgerInfosRequest {
            start_epoch: 0,
            end_epoch: 899,
        },
    )];
    assert_eq!(client_requests, expected_requests);

    // Create a batch of regular client requests and verify the result
    let client_requests = stream_tracker
        .create_epoch_ending_client_requests(3, 50)
        .unwrap();
    for (i, client_request) in client_requests.iter().enumerate() {
        let i = i as u64;
        let expected_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch: i * 50,
                end_epoch: ((i + 1) * 50) - 1,
            });
        assert_eq!(*client_request, expected_request);
    }

    // Create a large batch of small client requests and verify the result
    let client_requests = stream_tracker
        .create_epoch_ending_client_requests(100, 14)
        .unwrap();
    for (i, client_request) in client_requests.iter().enumerate() {
        let i = i as u64;
        let expected_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch: i * 14,
                end_epoch: cmp::min(((i + 1) * 14) - 1, 899),
            });
        assert_eq!(*client_request, expected_request);
    }
}

#[test]
fn test_epoch_ending_requests_dynamic() {
    // Create a new data stream progress tracker
    let stream_progress_tracker = create_epoch_ending_progress_tracker(0, 1000);
    let StreamProgressTracker::EpochEndingStreamTracker(mut stream_tracker) =
        stream_progress_tracker;

    // Update the tracker with a new next request epoch
    stream_tracker.next_request_epoch = 150;

    // Create a batch of client requests and verify the result
    let client_requests = stream_tracker
        .create_epoch_ending_client_requests(5, 700)
        .unwrap();
    let expected_requests = vec![
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 150,
            end_epoch: 849,
        }),
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 850,
            end_epoch: 999,
        }),
    ];
    assert_eq!(client_requests, expected_requests);

    // Update the tracker with a new next request epoch
    stream_tracker.next_request_epoch = 700;

    // Create a batch of client requests and verify the result
    let client_requests = stream_tracker
        .create_epoch_ending_client_requests(10, 50)
        .unwrap();
    for (i, client_request) in client_requests.iter().enumerate() {
        let i = i as u64;
        let expected_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch: 700 + (i * 50),
                end_epoch: cmp::min(700 + ((i + 1) * 50) - 1, 999),
            });
        assert_eq!(*client_request, expected_request);
    }

    // Update the tracker with a new next request epoch that matches the stream end
    stream_tracker.next_request_epoch = 999;

    // Create a batch of client requests and verify the result
    let client_requests = stream_tracker
        .create_epoch_ending_client_requests(5, 700)
        .unwrap();
    let expected_requests = vec![DataClientRequest::EpochEndingLedgerInfos(
        EpochEndingLedgerInfosRequest {
            start_epoch: 999,
            end_epoch: 999,
        },
    )];
    assert_eq!(client_requests, expected_requests);

    // Update the tracker with a new next request epoch that is at the end
    stream_tracker.next_request_epoch = 1000;

    // Create a batch of client requests and verify an overflow error
    let client_requests = stream_tracker.create_epoch_ending_client_requests(10, 50);
    assert_matches!(client_requests, Err(Error::IntegerOverflow(_)));
}

#[test]
fn test_epoch_ending_stream_tracker() {
    // Create an epoch ending stream request
    let stream_request =
        StreamRequest::GetAllEpochEndingLedgerInfos(GetAllEpochEndingLedgerInfosRequest {
            start_epoch: 0,
        });

    // Try to create a progress stream tracker where there is no advertised data
    // and verify an error is returned.
    let result =
        StreamProgressTracker::new(&stream_request, &GlobalDataSummary::empty().advertised_data);
    assert_matches!(result, Err(Error::DataIsUnavailable(_)));

    // Create a data summary with various advertised epoch ranges (common highest is zero)
    let mut global_data_summary = GlobalDataSummary::empty();
    global_data_summary
        .advertised_data
        .epoch_ending_ledger_infos = vec![
        CompleteDataRange::new(0, 0),
        CompleteDataRange::new(0, 0),
        CompleteDataRange::new(0, 1),
    ];

    // Try to create a progress stream tracker where the highest epoch is zero
    // and verify an error is returned.
    let result = StreamProgressTracker::new(&stream_request, &global_data_summary.advertised_data);
    assert_matches!(result, Err(Error::NoDataToFetch(_)));

    // Create a global data summary with non-zero advertised epoch ranges
    let mut global_data_summary = GlobalDataSummary::empty();
    global_data_summary
        .advertised_data
        .epoch_ending_ledger_infos = vec![
        CompleteDataRange::new(0, 1),
        CompleteDataRange::new(0, 100),
        CompleteDataRange::new(0, 99999999),
        CompleteDataRange::new(0, 100),
    ];

    // Create a new data stream progress tracker and verify the most common highest
    // epoch is chosen.
    let stream_progress_tracker =
        StreamProgressTracker::new(&stream_request, &global_data_summary.advertised_data).unwrap();
    let StreamProgressTracker::EpochEndingStreamTracker(stream_tracker) = stream_progress_tracker;
    assert_eq!(stream_tracker.end_epoch, 99); // End epoch is highest - 1
}

#[test]
fn test_epoch_ending_request_progress() {
    // Create a new data stream progress tracker
    let mut stream_progress_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the progress tracker using valid sent request notifications
    for i in 0..10 {
        let start_epoch = i * 100;
        let end_epoch = (i * 100) + 99;
        let client_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch,
                end_epoch,
            });
        stream_progress_tracker
            .update_request_progress(&client_request)
            .unwrap();

        // Verify internal state
        let StreamProgressTracker::EpochEndingStreamTracker(stream_tracker) =
            &stream_progress_tracker;
        assert_eq!(stream_tracker.next_request_epoch, end_epoch + 1);
    }
}

#[test]
fn test_epoch_ending_stream_progress() {
    // Create a new data stream progress tracker
    let mut stream_progress_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the progress tracker using valid sent data notifications
    let epoch_ending_chunk_size = 100;
    for i in 0..10 {
        let start_epoch = i * epoch_ending_chunk_size;
        let end_epoch = ((i + 1) * epoch_ending_chunk_size) - 1;
        let sent_data_notification = create_sent_data_notification(
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch,
                end_epoch,
            }),
        );
        stream_progress_tracker
            .update_notification_progress(&sent_data_notification)
            .unwrap();

        // Verify internal state
        let StreamProgressTracker::EpochEndingStreamTracker(stream_tracker) =
            &stream_progress_tracker;
        assert_eq!(stream_tracker.next_stream_epoch, end_epoch + 1);
    }
}

#[test]
#[should_panic(expected = "Updating an epoch ending tracker with an old request!")]
fn test_epoch_ending_request_progress_panic() {
    // Create a new data stream progress tracker
    let mut stream_progress_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the tracker with a valid request
    let sent_data_notification =
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 0,
            end_epoch: 100,
        });
    stream_progress_tracker
        .update_request_progress(&sent_data_notification)
        .unwrap();

    // Update the tracker with a request that misses data and verify a panic
    let sent_data_notification =
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 102,
            end_epoch: 200,
        });
    stream_progress_tracker
        .update_request_progress(&sent_data_notification)
        .unwrap();
}

#[test]
#[should_panic(expected = "Updating an epoch ending tracker with an old notification!")]
fn test_epoch_ending_stream_progress_panic() {
    // Create a new data stream progress tracker
    let mut stream_progress_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the tracker with a valid notification
    let sent_data_notification = create_sent_data_notification(
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 0,
            end_epoch: 100,
        }),
    );
    stream_progress_tracker
        .update_notification_progress(&sent_data_notification)
        .unwrap();

    // Update the tracker with an old notification and verify a panic
    let sent_data_notification = create_sent_data_notification(
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 50,
            end_epoch: 1100,
        }),
    );
    stream_progress_tracker
        .update_notification_progress(&sent_data_notification)
        .unwrap();
}

fn create_epoch_ending_progress_tracker(
    start_epoch: u64,
    max_advertised_epoch: u64,
) -> StreamProgressTracker {
    // Create an epoch ending stream request
    let stream_request =
        StreamRequest::GetAllEpochEndingLedgerInfos(GetAllEpochEndingLedgerInfosRequest {
            start_epoch,
        });

    // Create a global data summary with a single epoch range
    let mut global_data_summary = GlobalDataSummary::empty();
    global_data_summary
        .advertised_data
        .epoch_ending_ledger_infos =
        vec![CompleteDataRange::new(start_epoch, max_advertised_epoch)];

    // Create a new data stream progress tracker
    StreamProgressTracker::new(&stream_request, &global_data_summary.advertised_data).unwrap()
}

fn create_sent_data_notification(client_request: DataClientRequest) -> SentDataNotification {
    SentDataNotification {
        client_request,
        client_response: DataClientResponse {
            response_id: 0,
            response_payload: DataClientPayload::EpochEndingLedgerInfos(vec![]),
        },
    }
}
<commit_msg>[State Sync] Update the stream progress tracker unit tests.<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    data_notification::{DataClientRequest, EpochEndingLedgerInfosRequest, SentDataNotification},
    error::Error,
    stream_progress_tracker::{DataStreamTracker, EpochEndingStreamTracker, StreamProgressTracker},
    streaming_client::{GetAllEpochEndingLedgerInfosRequest, StreamRequest},
};
use claim::assert_matches;
use diem_data_client::{
    DataClientPayload, DataClientResponse, GlobalDataSummary, OptimalChunkSizes,
};
use std::cmp;
use storage_service_types::CompleteDataRange;

#[test]
fn test_create_epoch_ending_requests() {
    // Create a new data stream progress tracker
    let stream_tracker = create_epoch_ending_progress_tracker(0, 900);

    // Create a batch of large client requests and verify the result
    let client_requests = stream_tracker
        .create_data_client_requests(5, &create_epoch_ending_chunk_sizes(10000))
        .unwrap();
    let expected_requests = vec![DataClientRequest::EpochEndingLedgerInfos(
        EpochEndingLedgerInfosRequest {
            start_epoch: 0,
            end_epoch: 899,
        },
    )];
    assert_eq!(client_requests, expected_requests);

    // Create a batch of regular client requests and verify the result
    let client_requests = stream_tracker
        .create_data_client_requests(3, &create_epoch_ending_chunk_sizes(50))
        .unwrap();
    for (i, client_request) in client_requests.iter().enumerate() {
        let i = i as u64;
        let expected_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch: i * 50,
                end_epoch: ((i + 1) * 50) - 1,
            });
        assert_eq!(*client_request, expected_request);
    }

    // Create a batch of small client requests and verify the result
    let client_requests = stream_tracker
        .create_data_client_requests(100, &create_epoch_ending_chunk_sizes(14))
        .unwrap();
    for (i, client_request) in client_requests.iter().enumerate() {
        let i = i as u64;
        let expected_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch: i * 14,
                end_epoch: cmp::min(((i + 1) * 14) - 1, 899),
            });
        assert_eq!(*client_request, expected_request);
    }
}

#[test]
fn test_create_epoch_ending_requests_dynamic() {
    // Create a new data stream progress tracker
    let mut stream_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the tracker with a new next request epoch
    stream_tracker.next_request_epoch = 150;

    // Create a batch of client requests and verify the result
    let client_requests = stream_tracker
        .create_data_client_requests(5, &create_epoch_ending_chunk_sizes(700))
        .unwrap();
    let expected_requests = vec![
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 150,
            end_epoch: 849,
        }),
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 850,
            end_epoch: 999,
        }),
    ];
    assert_eq!(client_requests, expected_requests);

    // Update the tracker with a new next request epoch
    stream_tracker.next_request_epoch = 700;

    // Create a batch of client requests and verify the result
    let client_requests = stream_tracker
        .create_data_client_requests(10, &create_epoch_ending_chunk_sizes(50))
        .unwrap();
    for (i, client_request) in client_requests.iter().enumerate() {
        let i = i as u64;
        let expected_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch: 700 + (i * 50),
                end_epoch: cmp::min(700 + ((i + 1) * 50) - 1, 999),
            });
        assert_eq!(*client_request, expected_request);
    }

    // Update the tracker with a new next request epoch that matches the stream end
    stream_tracker.next_request_epoch = 999;

    // Create a batch of client requests and verify the result
    let client_requests = stream_tracker
        .create_data_client_requests(5, &create_epoch_ending_chunk_sizes(700))
        .unwrap();
    let expected_requests = vec![DataClientRequest::EpochEndingLedgerInfos(
        EpochEndingLedgerInfosRequest {
            start_epoch: 999,
            end_epoch: 999,
        },
    )];
    assert_eq!(client_requests, expected_requests);

    // Update the tracker with a new next request epoch that is at the end
    stream_tracker.next_request_epoch = 1000;

    // Create a batch of client requests and verify an overflow error
    let client_requests =
        stream_tracker.create_data_client_requests(10, &create_epoch_ending_chunk_sizes(50));
    assert_matches!(client_requests, Err(Error::IntegerOverflow(_)));
}

#[test]
fn test_epoch_ending_stream_tracker() {
    // Create an epoch ending stream request
    let stream_request =
        StreamRequest::GetAllEpochEndingLedgerInfos(GetAllEpochEndingLedgerInfosRequest {
            start_epoch: 0,
        });

    // Try to create a progress stream tracker where there is no advertised data
    // and verify an error is returned.
    let result =
        StreamProgressTracker::new(&stream_request, &GlobalDataSummary::empty().advertised_data);
    assert_matches!(result, Err(Error::DataIsUnavailable(_)));

    // Create a data summary with various advertised epoch ranges (common highest is zero)
    let mut global_data_summary = GlobalDataSummary::empty();
    global_data_summary
        .advertised_data
        .epoch_ending_ledger_infos = vec![
        CompleteDataRange::new(0, 0),
        CompleteDataRange::new(0, 0),
        CompleteDataRange::new(0, 1),
    ];

    // Try to create a progress stream tracker where the highest epoch is zero
    // and verify an error is returned.
    let result = StreamProgressTracker::new(&stream_request, &global_data_summary.advertised_data);
    assert_matches!(result, Err(Error::NoDataToFetch(_)));

    // Create a global data summary with non-zero advertised epoch ranges
    let mut global_data_summary = GlobalDataSummary::empty();
    global_data_summary
        .advertised_data
        .epoch_ending_ledger_infos = vec![
        CompleteDataRange::new(0, 1),
        CompleteDataRange::new(0, 100),
        CompleteDataRange::new(0, 99999999),
        CompleteDataRange::new(0, 100),
    ];

    // Create a new data stream progress tracker and verify the most common highest
    // epoch is chosen.
    match StreamProgressTracker::new(&stream_request, &global_data_summary.advertised_data).unwrap()
    {
        StreamProgressTracker::EpochEndingStreamTracker(stream_tracker) => {
            assert_eq!(stream_tracker.end_epoch, 99); // End epoch is highest - 1
        }
        unexpected_tracker => {
            panic!(
                "Expected epoch ending stream tracker but got {:?}",
                unexpected_tracker
            );
        }
    }
}

#[test]
fn test_update_epoch_ending_request_progress() {
    // Create a new data stream progress tracker
    let mut stream_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the progress tracker using valid sent request notifications
    for i in 0..10 {
        let start_epoch = i * 100;
        let end_epoch = (i * 100) + 99;
        let client_request =
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch,
                end_epoch,
            });
        stream_tracker
            .update_request_tracking(&client_request)
            .unwrap();

        // Verify internal state
        assert_eq!(stream_tracker.next_request_epoch, end_epoch + 1);
    }
}

#[test]
#[should_panic(expected = "The start index did not match the expected next index!")]
fn test_update_epoch_ending_request_panic() {
    // Create a new data stream progress tracker
    let mut stream_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the tracker with a valid request
    let sent_data_notification =
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 0,
            end_epoch: 100,
        });
    stream_tracker
        .update_request_tracking(&sent_data_notification)
        .unwrap();

    // Update the tracker with a request that misses data and verify a panic
    let sent_data_notification =
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 102,
            end_epoch: 200,
        });
    stream_tracker
        .update_request_tracking(&sent_data_notification)
        .unwrap();
}

#[test]
fn test_update_epoch_ending_stream_progress() {
    // Create a new data stream progress tracker
    let mut stream_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the progress tracker using valid sent data notifications
    for i in 0..10 {
        let start_epoch = i * 100;
        let end_epoch = (i * 100) + 99;
        let sent_data_notification = create_sent_data_notification(
            DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
                start_epoch,
                end_epoch,
            }),
        );
        stream_tracker
            .update_notification_tracking(&sent_data_notification)
            .unwrap();

        // Verify internal state
        assert_eq!(stream_tracker.next_stream_epoch, end_epoch + 1);
    }
}

#[test]
#[should_panic(expected = "The start index did not match the expected next index!")]
fn test_update_epoch_ending_stream_panic() {
    // Create a new data stream progress tracker
    let mut stream_tracker = create_epoch_ending_progress_tracker(0, 1000);

    // Update the tracker with a valid notification
    let sent_data_notification = create_sent_data_notification(
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 0,
            end_epoch: 100,
        }),
    );
    stream_tracker
        .update_notification_tracking(&sent_data_notification)
        .unwrap();

    // Update the tracker with an old notification and verify a panic
    let sent_data_notification = create_sent_data_notification(
        DataClientRequest::EpochEndingLedgerInfos(EpochEndingLedgerInfosRequest {
            start_epoch: 50,
            end_epoch: 1100,
        }),
    );
    stream_tracker
        .update_notification_tracking(&sent_data_notification)
        .unwrap();
}

fn create_epoch_ending_progress_tracker(
    start_epoch: u64,
    max_advertised_epoch: u64,
) -> EpochEndingStreamTracker {
    // Create an epoch ending stream request
    let stream_request =
        StreamRequest::GetAllEpochEndingLedgerInfos(GetAllEpochEndingLedgerInfosRequest {
            start_epoch,
        });

    // Create a global data summary with a single epoch range
    let mut global_data_summary = GlobalDataSummary::empty();
    global_data_summary
        .advertised_data
        .epoch_ending_ledger_infos =
        vec![CompleteDataRange::new(start_epoch, max_advertised_epoch)];

    // Create a new epoch ending stream progress tracker
    match StreamProgressTracker::new(&stream_request, &global_data_summary.advertised_data).unwrap()
    {
        StreamProgressTracker::EpochEndingStreamTracker(stream_tracker) => stream_tracker,
        unexpected_tracker => {
            panic!(
                "Expected epoch ending stream tracker but got {:?}",
                unexpected_tracker
            );
        }
    }
}

fn create_epoch_ending_chunk_sizes(epoch_chunk_size: u64) -> OptimalChunkSizes {
    let mut optimal_chunk_sizes = OptimalChunkSizes::empty();
    optimal_chunk_sizes.epoch_chunk_size = epoch_chunk_size;
    optimal_chunk_sizes
}

fn create_sent_data_notification(client_request: DataClientRequest) -> SentDataNotification {
    SentDataNotification {
        client_request,
        client_response: DataClientResponse {
            response_id: 0,
            response_payload: DataClientPayload::EpochEndingLedgerInfos(vec![]),
        },
    }
}
<|endoftext|>


<commit_before>use crate::os::process::can_run_services_as_svc_user;
use nix::unistd::{setgid,
                  setuid,
                  Gid,
                  Uid};
use std::{ffi::OsStr,
          io,
          os::unix::process::CommandExt,
          process::{Command,
                    Stdio},
          result};

/// Prepare a `Command` to execute a lifecycle hook.
// TODO (CM): Ideally, `ids` would not be an `Option`, but separate
// `Uid` and `Gid` inputs. However, the `Option` interface provides
// the least disruption to other existing code for the time being.
pub fn hook_command<X, I, K, V>(executable: X, env: I, ids: Option<(Uid, Gid)>) -> Command
    where X: AsRef<OsStr>,
          I: IntoIterator<Item = (K, V)>,
          K: AsRef<OsStr>,
          V: AsRef<OsStr>
{
    let mut cmd = Command::new(executable);

    // NOTE: CommandExt::uid and CommandExt::guid should *not* be
    // called here! They are set in `with_user_and_group_information`;
    // see there for further details.
    cmd.stdin(Stdio::null())
       .stdout(Stdio::piped())
       .stderr(Stdio::piped())
       .envs(env);

    with_own_process_group(&mut cmd);
    if let Some((uid, gid)) = ids {
        with_user_and_group_information(&mut cmd, uid, gid);
    }

    cmd
}

/// Ensures that the `Command` is executed within its own process
/// group, and not that of its parent process.
///
/// This should be used when spawning all hooks. This ensures that
/// they are not the same process group as the Launcher (for `run`
/// hooks) or the Supervisor (for all other hooks). Otherwise, if a
/// child process were to send `SIGTERM`, the Launcher could be
/// terminated. Similarly, it prevents a `^C` sent to a foregrounded
/// Supervisor from terminating any hooks prematurely.
///
/// This basically ensures that all hooks are properly isolated,
/// without signaling cross-talk between them and the Launcher /
/// Supervisor.
fn with_own_process_group(cmd: &mut Command) -> &mut Command {
    unsafe {
        cmd.pre_exec(set_own_process_group);
    }
    cmd
}

/// Set the process group of the calling process to be the same as its
/// PID.
///
/// Intended for use in a
/// `std::os::unix::process::CommandExt::pre_exec` callback.
fn set_own_process_group() -> result::Result<(), io::Error> {
    unsafe {
        if libc::setpgid(0, 0) == 0 {
            Ok(())
        } else {
            Err(io::Error::last_os_error())
        }
    }
}
/// Sets uid, gid, and supplementary groups on command.
///
/// DO NOT call `CommandExt#uid` or `CommandExt#gid` on this command,
/// either before or after calling this function, or it will probably
/// not work like you want it to.
fn with_user_and_group_information(cmd: &mut Command, uid: Uid, gid: Gid) -> &mut Command {
    unsafe {
        cmd.pre_exec(set_supplementary_groups(uid, gid));
    }
    cmd
}

/// Stupid little private helper macro to make mapping `Nix` errors to
/// IO errors for our `pre_exec` hooks.
///
/// The format string should have a single variable placeholder for
/// the actual error.
///
/// e.g. `result.map_err(io_err!("blah blah {:?}"))`
macro_rules! io_error {
    ($format_string:tt) => {
        move |e| io::Error::new(io::ErrorKind::Other, format!($format_string, e))
    };
}

/// Returns a function that sets the supplementary group IDs of the
/// process to those that `user_id` belongs to.
///
/// Also sets the uid and gid of the process. We must do that here,
/// rather than using the `CommandExt::uid` and `CommandExt::gid`
/// methods to ensure that all the IDs are set on the process in the
/// correct order (that is, supplementary groups, gid, and finally uid).
///
/// Once https://github.com/rust-lang/rust/pull/72160 merges, we can
/// use all `CommandExt` methods, and thus simplify things a (little)
/// bit.
fn set_supplementary_groups(user_id: Uid,
                            group_id: Gid)
                            -> impl Fn() -> result::Result<(), io::Error> {
    // Note: since this function will be run a separate process that doesn't
    // inherit RUST_LOG, none of the log! macros will work actually
    // work here.

    move || {
        // Note that if we *can't* run services as another user,
        // that's OK; not an error. We just won't set supplementary
        // groups, and run all hooks as the user we currently are.
        if can_run_services_as_svc_user() {
            // These calls don't have a macOS counterpart (well, not a
            // direct one in nix, at least) and we don't need to
            // execute hooks on macOS, so it's not important to
            // implement this. Our crates aren't really factored well
            // enough to let us cut out larger chunks of code on macOS
            // at the moment, so we just won't compile this particular
            // bit for now.
            #[cfg(not(target_os = "macos"))]
            {
                use nix::unistd::{getgrouplist,
                                  setgroups,
                                  User};
                use std::ffi::CString;

                if let Some(user) = User::from_uid(user_id).map_err(io_error!("Error resolving \
                                                                               user from ID: \
                                                                               {:?}"))?
                {
                    let user = CString::new(user.name).map_err(io_error!("User name cannot \
                                                                          convert to CString!: \
                                                                          {:?}"))?;

                    // There are some platforms (ex. SUSE 12 sp5) that may return
                    // EINVAL from getgrouplist. This only appears to occur from
                    // statically compiled (MUSL) executables like the hab CLI. The
                    // error has not been reproducible from a dynamic executable like
                    // the supervisor.
                    let groups = getgrouplist(&user, group_id).unwrap_or_else(|e| {
                                                                  warn!("unable to get \
                                                                         supplementary groups \
                                                                         with getgrouplist: {}",
                                                                        e);
                                                                  vec![group_id]
                                                              });
                    setgroups(&groups).map_err(io_error!("setgroups failed! {:?}"))?; // CAP_SETGID
                } else {
                    return Err(io::Error::new(io::ErrorKind::Other,
                                              "Could not find user from user ID"));
                }
            }

            // These calls replace `CommandExt::uid` and `CommandExt::gid`
            setgid(group_id).map_err(io_error!("setgid failed! {:?}"))?; // CAP_SETGID
            setuid(user_id).map_err(io_error!("setuid failed! {:?}"))?; // CAP_SETUID
        }

        Ok(())
    }
}
<commit_msg>use eprintln! to output specific error messages in pre_exec<commit_after>use crate::os::process::can_run_services_as_svc_user;
use nix::unistd::{setgid,
                  setuid,
                  Gid,
                  Uid};
use std::{ffi::OsStr,
          io,
          os::unix::process::CommandExt,
          process::{Command,
                    Stdio},
          result};

/// Prepare a `Command` to execute a lifecycle hook.
// TODO (CM): Ideally, `ids` would not be an `Option`, but separate
// `Uid` and `Gid` inputs. However, the `Option` interface provides
// the least disruption to other existing code for the time being.
pub fn hook_command<X, I, K, V>(executable: X, env: I, ids: Option<(Uid, Gid)>) -> Command
    where X: AsRef<OsStr>,
          I: IntoIterator<Item = (K, V)>,
          K: AsRef<OsStr>,
          V: AsRef<OsStr>
{
    let mut cmd = Command::new(executable);

    // NOTE: CommandExt::uid and CommandExt::guid should *not* be
    // called here! They are set in `with_user_and_group_information`;
    // see there for further details.
    cmd.stdin(Stdio::null())
       .stdout(Stdio::piped())
       .stderr(Stdio::piped())
       .envs(env);

    with_own_process_group(&mut cmd);
    if let Some((uid, gid)) = ids {
        with_user_and_group_information(&mut cmd, uid, gid);
    }

    cmd
}

/// Ensures that the `Command` is executed within its own process
/// group, and not that of its parent process.
///
/// This should be used when spawning all hooks. This ensures that
/// they are not the same process group as the Launcher (for `run`
/// hooks) or the Supervisor (for all other hooks). Otherwise, if a
/// child process were to send `SIGTERM`, the Launcher could be
/// terminated. Similarly, it prevents a `^C` sent to a foregrounded
/// Supervisor from terminating any hooks prematurely.
///
/// This basically ensures that all hooks are properly isolated,
/// without signaling cross-talk between them and the Launcher /
/// Supervisor.
fn with_own_process_group(cmd: &mut Command) -> &mut Command {
    unsafe {
        cmd.pre_exec(set_own_process_group);
    }
    cmd
}

/// Set the process group of the calling process to be the same as its
/// PID.
///
/// Intended for use in a
/// `std::os::unix::process::CommandExt::pre_exec` callback.
fn set_own_process_group() -> result::Result<(), io::Error> {
    unsafe {
        if libc::setpgid(0, 0) == 0 {
            Ok(())
        } else {
            Err(io::Error::last_os_error())
        }
    }
}
/// Sets uid, gid, and supplementary groups on command.
///
/// DO NOT call `CommandExt#uid` or `CommandExt#gid` on this command,
/// either before or after calling this function, or it will probably
/// not work like you want it to.
fn with_user_and_group_information(cmd: &mut Command, uid: Uid, gid: Gid) -> &mut Command {
    unsafe {
        cmd.pre_exec(set_supplementary_groups(uid, gid));
    }
    cmd
}

/// Returns a function that sets the supplementary group IDs of the
/// process to those that `user_id` belongs to.
///
/// Also sets the uid and gid of the process. We must do that here,
/// rather than using the `CommandExt::uid` and `CommandExt::gid`
/// methods to ensure that all the IDs are set on the process in the
/// correct order (that is, supplementary groups, gid, and finally uid).
///
/// Once https://github.com/rust-lang/rust/pull/72160 merges, we can
/// use all `CommandExt` methods, and thus simplify things a (little)
/// bit.
fn set_supplementary_groups(user_id: Uid,
                            group_id: Gid)
                            -> impl Fn() -> result::Result<(), io::Error> {
    // Note: since this function will be run a separate process that doesn't
    // inherit RUST_LOG, none of the log! macros will work actually
    // work here.

    move || {
        // Note that if we *can't* run services as another user,
        // that's OK; not an error. We just won't set supplementary
        // groups, and run all hooks as the user we currently are.
        if can_run_services_as_svc_user() {
            // These calls don't have a macOS counterpart (well, not a
            // direct one in nix, at least) and we don't need to
            // execute hooks on macOS, so it's not important to
            // implement this. Our crates aren't really factored well
            // enough to let us cut out larger chunks of code on macOS
            // at the moment, so we just won't compile this particular
            // bit for now.
            #[cfg(not(target_os = "macos"))]
            {
                use nix::unistd::{getgrouplist,
                                  setgroups,
                                  User};
                use std::ffi::CString;

                if let Some(user) =
                    User::from_uid(user_id).map_err(|e| {
                                               eprintln!("Error resolving user from ID: {:?}", e);
                                               io::Error::last_os_error()
                                           })?
                {
                    let user = CString::new(user.name).map_err(|e| {
                                                          eprintln!("User name cannot convert to \
                                                                     CString!: {:?}",
                                                                    e);
                                                          e
                                                      })?;

                    // There are some platforms (ex. SUSE 12 sp5) that may return
                    // EINVAL from getgrouplist. This only appears to occur from
                    // statically compiled (MUSL) executables like the hab CLI. The
                    // error has not been reproducible from a dynamic executable like
                    // the supervisor.
                    let groups = getgrouplist(&user, group_id).unwrap_or_else(|e| {
                                                                  warn!("unable to get \
                                                                         supplementary groups \
                                                                         with getgrouplist: {}",
                                                                        e);
                                                                  vec![group_id]
                                                              });
                    setgroups(&groups).map_err(|e| {
                                          eprintln!("setgroups failed! {:?}", e);
                                          io::Error::last_os_error()
                                      })?; // CAP_SETGID
                } else {
                    eprintln!("Could not find user from user ID. Wil not set supplementary \
                               groups.");
                }
            }

            // These calls replace `CommandExt::uid` and `CommandExt::gid`
            setgid(group_id).map_err(|e| {
                                eprintln!("setgid failed! {:?}", e);
                                io::Error::last_os_error()
                            })?; // CAP_SETGID
            setuid(user_id).map_err(|e| {
                               eprintln!("setuid failed! {:?}", e);
                               io::Error::last_os_error()
                           })?; // CAP_SETUID
        }

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


<commit_before>#![recursion_limit = "128"]

use proc_macro::TokenStream;

mod context;
mod ex_struct;
mod init;
mod map;
mod nif;
mod record;
mod tuple;
mod unit_enum;
mod untagged_enum;

#[derive(Debug)]
enum RustlerAttr {
    Encode,
    Decode,
    Module(String),
    Tag(String),
}

/// Implementation of a Native Implemented Function (NIF) macro that lets the user annotate
/// a function that will be wrapped in higer-level NIF implementation.
///
/// ```ignore
/// #[rustler::nif]
/// fn add(a: i64, b: i64) -> i64 {
///     a + b
/// }
///
/// #[rustler::nif]
/// fn sub(a: i64, b: i64) -> i64 {
///     a - b
/// }
///
/// #[rustler::nif]
/// fn mul(a: i64, b: i64) -> i64 {
///     a * b
/// }
///
/// #[rustler::nif]
/// fn div(a: i64, b: i64) -> i64 {
///     a / b
/// }
///
/// rustler::init!("Elixir.Math", [add, sub, mul, div], Some(load));
/// ```
#[proc_macro]
pub fn init(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as init::InitMacroInput);
    let output: proc_macro2::TokenStream = input.into();
    output.into()
}

/// Implementation of a Native Implemented Function (NIF) macro that lets the user annotate
/// a function that will be wrapped in higer-level NIF implementation.
///
/// ```ignore
/// #[nif]
/// fn add(a: i64, b: i64) -> i64 {
///     a + b
/// }
/// ```
#[proc_macro_attribute]
pub fn nif(args: TokenStream, input: TokenStream) -> TokenStream {
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
    let input = syn::parse_macro_input!(input as syn::ItemFn);

    nif::transcoder_decorator(args, input).into()
}

/// Implementation of the `NifStruct` macro that lets the user annotate a struct that will
/// be translated directly from an Elixir struct to a Rust struct. For example, the following
/// struct, annotated as such:
///
/// ```ignore
/// #[derive(Debug, NifStruct)]
/// #[module = "AddStruct"]
/// struct AddStruct {
///    lhs: i32,
///    rhs: i32,
/// }
/// ```
///
/// This would be translated by Rustler into:
///
/// ```elixir
/// defmodule AddStruct do
///   defstruct lhs: 0, rhs: 0
/// end
/// ```
#[proc_macro_derive(NifStruct, attributes(module, rustler))]
pub fn nif_struct(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    ex_struct::transcoder_decorator(&ast, false).into()
}

/// Implementation of the `NifException` macro that lets the user annotate a struct that will
/// be translated directly from an Elixir exception to a Rust struct. For example, the following
/// struct, annotated as such:
///
/// ```ignore
/// #[derive(Debug, NifException)]
/// #[module = "AddException"]
/// pub struct AddException {
///     message: String,
/// }
/// ```
///
/// This would be translated by Rustler into:
///
/// ```elixir
/// defmodule AddException do
///   defexception message: ""
/// end
/// ```
#[proc_macro_derive(NifException, attributes(module, rustler))]
pub fn nif_exception(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    ex_struct::transcoder_decorator(&ast, true).into()
}

/// Implementation of a macro that lets the user annotate a struct with `NifMap` so that the
/// struct can be encoded or decoded from an Elixir map. For example, the following struct
/// annotated as such:
///
/// ```ignore
/// #[derive(NifMap)]
/// struct AddMap {
///     lhs: i32,
///     rhs: i32,
/// }
/// ```
///
/// Given the values 33 and 21 for this struct, this would result, when encoded, in an elixir
/// map with two elements like:
///
/// ```elixir
/// %{lhs: 33, rhs: 21}
/// ```
#[proc_macro_derive(NifMap, attributes(rustler))]
pub fn nif_map(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    map::transcoder_decorator(&ast).into()
}

/// Implementation of a macro that lets the user annotate a struct with `NifTuple` so that the
/// struct can be encoded or decoded from an Elixir tuple. For example, the following struct
/// annotated as such:
///
/// ```ignore
/// #[derive(NifTuple)]
/// struct AddTuple {
///     lhs: i32,
///     rhs: i32,
/// }
/// ```
///
/// Given the values 33 and 21 for this struct, this would result, when encoded, in an elixir
/// tuple with two elements like:
///
/// ```elixir
/// {33, 21}
/// ```
///
/// The size of the tuple will depend on the number of elements in the struct.
#[proc_macro_derive(NifTuple, attributes(rustler))]
pub fn nif_tuple(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    tuple::transcoder_decorator(&ast).into()
}

/// Implementation of the `NifRecord` macro that lets the user annotate a struct that will
/// be translated directly from an Elixir struct to a Rust struct. For example, the following
/// struct, annotated as such:
///
/// ```ignore
/// #[derive(Debug, NifRecord)]
/// #[tag = "record"]
/// struct AddRecord {
///    lhs: i32,
///    rhs: i32,
/// }
/// ```
///
/// This would be translated by Rustler into:
///
/// ```elixir
/// defmodule AddRecord do
///   import Record
///   defrecord :record, [lhs: 1, rhs: 2]
/// end
/// ```
#[proc_macro_derive(NifRecord, attributes(tag, rustler))]
pub fn nif_record(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    record::transcoder_decorator(&ast).into()
}

/// Implementation of the `NifUnitEnum` macro that lets the user annotate an enum with a unit type
/// that will generate elixir atoms when encoded
///
/// ```ignore
/// #[derive(NifUnitEnum)]
/// enum UnitEnum {
///    FooBar,
///    Baz,
/// }
/// ```
///
/// An example usage in elixir would look like the following.
///
/// ```elixir
/// test "unit enum transcoder" do
///   assert :foo_bar == RustlerTest.unit_enum_echo(:foo_bar)
///   assert :baz == RustlerTest.unit_enum_echo(:baz)
///   assert :invalid_variant == RustlerTest.unit_enum_echo(:somethingelse)
/// end
/// ```
///
/// Note that the `:invalid_variant` atom is returned if the user tries to encode something
/// that isn't in the Rust enum.
#[proc_macro_derive(NifUnitEnum, attributes(rustler))]
pub fn nif_unit_enum(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    unit_enum::transcoder_decorator(&ast).into()
}

/// Implementation of the `NifUntaggedEnum` macro that lets the user annotate an enum that will
/// generate elixir values when decoded. This can be used for rust enums that contain data and
/// will generate a value based on the kind of data encoded. For example from the test code:
///
/// ```ignore
/// #[derive(NifUntaggedEnum)]
/// enum UntaggedEnum {
///     Foo(u32),
///     Bar(String),
///     Baz(AddStruct),
/// }
///
/// pub fn untagged_enum_echo<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
///     let untagged_enum: UntaggedEnum = args[0].decode()?;
///     Ok(untagged_enum.encode(env))
/// }
/// ```
///
/// This can be used from elixir in the following manner.
///
/// ```elixir
/// test "untagged enum transcoder" do
///   assert 123 == RustlerTest.untagged_enum_echo(123)
///   assert "Hello" == RustlerTest.untagged_enum_echo("Hello")
///   assert %AddStruct{lhs: 45, rhs: 123} = RustlerTest.untagged_enum_echo(%AddStruct{lhs: 45, rhs: 123})
///   assert :invalid_variant == RustlerTest.untagged_enum_echo([1,2,3,4])
/// end
/// ```
///
/// Note that the type of elixir return is dependent on the data in the enum and the actual enum
/// type is lost in the translation because Elixir has no such concept.
#[proc_macro_derive(NifUntaggedEnum, attributes(rustler))]
pub fn nif_untagged_enum(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    untagged_enum::transcoder_decorator(&ast).into()
}
<commit_msg>Improve documentation of derive macros<commit_after>#![recursion_limit = "128"]

use proc_macro::TokenStream;

mod context;
mod ex_struct;
mod init;
mod map;
mod nif;
mod record;
mod tuple;
mod unit_enum;
mod untagged_enum;

#[derive(Debug)]
enum RustlerAttr {
    Encode,
    Decode,
    Module(String),
    Tag(String),
}

/// Initialise the Native Implemented Function (NIF) environment
/// and register NIF functions in an Elixir module.
///
/// ```ignore
/// #[rustler::nif]
/// fn add(a: i64, b: i64) -> i64 {
///     a + b
/// }
///
/// #[rustler::nif]
/// fn sub(a: i64, b: i64) -> i64 {
///     a - b
/// }
///
/// #[rustler::nif]
/// fn mul(a: i64, b: i64) -> i64 {
///     a * b
/// }
///
/// #[rustler::nif]
/// fn div(a: i64, b: i64) -> i64 {
///     a / b
/// }
///
/// rustler::init!("Elixir.Math", [add, sub, mul, div], Some(load));
/// ```
#[proc_macro]
pub fn init(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as init::InitMacroInput);
    let output: proc_macro2::TokenStream = input.into();
    output.into()
}

/// Wrap a function in a Native Implemented Function (NIF) implementation,
/// so that it can be called from Elixir,
/// with all encoding and decoding steps done automatically.
///
/// ```ignore
/// #[nif]
/// fn add(a: i64, b: i64) -> i64 {
///     a + b
/// }
/// ```
#[proc_macro_attribute]
pub fn nif(args: TokenStream, input: TokenStream) -> TokenStream {
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
    let input = syn::parse_macro_input!(input as syn::ItemFn);

    nif::transcoder_decorator(args, input).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between an Elixir struct and a Rust struct.
///
/// For example, annotate the following Rust struct:
///
/// ```ignore
/// #[derive(Debug, NifStruct)]
/// #[module = "AddStruct"]
/// struct AddStruct {
///    lhs: i32,
///    rhs: i32,
/// }
/// ```
///
/// Write the following corresponding Elixir struct definition:
///
/// ```elixir
/// defmodule AddStruct do
///   defstruct lhs: 0, rhs: 0
/// end
/// ```
///
/// Then the traits `Encoder` and `Decoder` are derived automatically for your Rust struct
/// such that you can use the Elixir struct definition for it.

#[proc_macro_derive(NifStruct, attributes(module, rustler))]
pub fn nif_struct(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    ex_struct::transcoder_decorator(&ast, false).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between an Elixir exception and a Rust struct.
///
/// For example, annotate the following struct:
///
/// ```ignore
/// #[derive(Debug, NifException)]
/// #[module = "AddException"]
/// pub struct AddException {
///     message: String,
/// }
/// ```
///
/// Write the corresponding Elixir exception definition:
///
/// ```elixir
/// defmodule AddException do
///   defexception message: ""
/// end
/// ```
///
/// Then the traits `Encoder` and `Decoder` are derived automatically for your Rust struct
/// such that you can use the Elixir exception definition for it.
#[proc_macro_derive(NifException, attributes(module, rustler))]
pub fn nif_exception(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    ex_struct::transcoder_decorator(&ast, true).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between Rust struct and an Elixir map.
///
/// For example, annotate the following struct:
///
/// ```ignore
/// #[derive(NifMap)]
/// struct AddMap {
///     lhs: i32,
///     rhs: i32,
/// }
/// ```
///
/// Create a value of that type:
///
/// ```ignore
/// let value = AddMap { lhs: 33, rhs: 21 };
/// ```
///
/// Then the traits `Encoder` and `Decoder` are derived automatically for your Rust struct
/// such that encoding `value` would result in an elixir
/// map with two elements like:
///
/// ```elixir
/// %{lhs: 33, rhs: 21}
/// ```
///
/// And vice versa, decoding this map would result in `value`.
#[proc_macro_derive(NifMap, attributes(rustler))]
pub fn nif_map(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    map::transcoder_decorator(&ast).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between a Rust struct and an Elixir tuple.
///
/// For example, annotate the following struct:
///
/// ```ignore
/// #[derive(NifTuple)]
/// struct AddTuple {
///     lhs: i32,
///     rhs: i32,
/// }
/// ```
///
/// Create a value of that type:
///
/// ```ignore
/// let value = AddMap { lhs: 33, rhs: 21 };
/// ```
///
/// Then the traits `Encoder` and `Decoder` are derived automatically for your Rust struct
/// such that encoding `value` would result in an elixir
/// tuple with two elements like:
///
/// ```elixir
/// {33, 21}
/// ```
///
/// And vice versa, decoding this map would result in `value`.
///
/// The size of the tuple will depend on the number of elements in the struct.
#[proc_macro_derive(NifTuple, attributes(rustler))]
pub fn nif_tuple(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    tuple::transcoder_decorator(&ast).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between a Rust struct and an Elixir record.
///
/// For example, annotate the following struct:
///
/// ```ignore
/// #[derive(Debug, NifRecord)]
/// #[tag = "record"]
/// struct AddRecord {
///    lhs: i32,
///    rhs: i32,
/// }
/// ```
///
/// Create a value of that type:
///
/// ```ignore
/// let value = AddRecord { lhs: 33, rhs: 21 };
/// ```
///
/// Then the traits `Encoder` and `Decoder` are derived automatically for your Rust struct
/// such that `value` would be encoded into the following elixir value:
///
/// ```elixir
/// {:record, 33, 21}
/// ```
///
/// If you supply the following matching Elixir record definition:
///
/// ```elixir
/// defmodule AddRecord do
///   import Record
///   defrecord :record, [lhs: 1, rhs: 2]
/// end
/// ```
///
/// Then you can use record functions such as `AddRecord.record/0`, `AddRecord.record/1`, `AddRecord.record/2`,
/// to work with the encoded data,
/// and to create data that can be decoded back into your Rust struct.
#[proc_macro_derive(NifRecord, attributes(tag, rustler))]
pub fn nif_record(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    record::transcoder_decorator(&ast).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between an enum and a union of elixir atoms.
///
/// For example:
///
/// ```ignore
/// #[derive(NifUnitEnum)]
/// enum UnitEnum {
///    FooBar,
///    Baz,
/// }
/// ```
///
/// Then the traits `Encoder` and `Decoder` are derived automatically for your Rust struct
/// such that `FooBar` is encoded to, and decoded from, `:foo_bar`.
/// - The variant name is translated from camel case to snake case for the atom name.
/// - Each constructor is required not to have arguments, i.e. to be of unit type.
///
/// An example usage in Rust and Elixir would look like the following.
///
/// ```ignore
/// #[rustler::nif]
/// pub fn unit_enum_echo(unit_enum: UnitEnum) -> UnitEnum {
///     unit_enum
/// }
/// ```
///
/// (We are leaving out some boiler plate code to connect the rust code to elixir functions.)
///
/// ```elixir
/// test "unit enum transcoder" do
///   assert :foo_bar == unit_enum_echo(:foo_bar)
///   assert :baz == unit_enum_echo(:baz)
///   assert :invalid_variant == unit_enum_echo(:somethingelse)
/// end
/// ```
///
/// Note that the `:invalid_variant` atom is returned if the user tries to encode something
/// that isn't in the Rust enum.
#[proc_macro_derive(NifUnitEnum, attributes(rustler))]
pub fn nif_unit_enum(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    unit_enum::transcoder_decorator(&ast).into()
}

/// Derives implementations for the `Encoder` and `Decoder` traits
/// which convert between a Rust enum and a union of Elixir types.
///
/// This can be used for Rust enums that contain several constructors containing different types of data,
/// each implementing the `Encoder` and `Decoder` traits.
/// An enum value will be encoded based on the constructor used,
/// and an Elixir value will be decoded based on the value.
///
/// For example from the test code:
///
/// ```ignore
/// #[derive(NifUntaggedEnum)]
/// enum UntaggedEnum {
///     Foo(u32),
///     Bar(String),
///     Baz(AddStruct),
/// }
///
/// #[rustler::nif]
/// pub fn untagged_enum_echo(untagged_enum: UntaggedEnum) -> UntaggedEnum {
///     untagged_enum
/// }
/// ```
///
/// Adding boiler plate code to connect Rust code to elixir functions,
/// this can be used from elixir in the following manner.
///
/// ```elixir
/// test "untagged enum transcoder" do
///   assert 123 == untagged_enum_echo(123)
///   assert "Hello" == untagged_enum_echo("Hello")
///   assert %AddStruct{lhs: 45, rhs: 123} = untagged_enum_echo(%AddStruct{lhs: 45, rhs: 123})
///   assert :invalid_variant == untagged_enum_echo([1,2,3,4])
/// end
/// ```
///
/// Note that the type of elixir return is dependent on the data in the enum and the actual enum
/// type is lost in the translation because Elixir has no such concept.
#[proc_macro_derive(NifUntaggedEnum, attributes(rustler))]
pub fn nif_untagged_enum(input: TokenStream) -> TokenStream {
    let ast = syn::parse(input).unwrap();
    untagged_enum::transcoder_decorator(&ast).into()
}
<|endoftext|>