<commit_before>//! C API for mp4parse module.
//!
//! Parses ISO Base Media Format aka video/mp4 streams.
//!
//! # Examples
//!
//! ```rust
//! extern crate mp4parse;
//!
//! // Minimal valid mp4 containing no tracks.
//! let data = b"\0\0\0\x0cftypmp42";
//!
//! let context = mp4parse::mp4parse_new();
//! unsafe {
//!     let rv = mp4parse::mp4parse_read(context, data.as_ptr(), data.len());
//!     assert_eq!(0, rv);
//!     mp4parse::mp4parse_free(context);
//! }
//! ```

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

use std;
use std::io::Cursor;

// Symbols we need from our rust api.
use MediaContext;
use TrackType;
use read_mp4;
use Error;
use media_time_to_ms;
use track_time_to_ms;
use SampleEntry;

const TRACK_TYPE_H264: u32 = 0;
const TRACK_TYPE_AAC:  u32 = 1;

// These structs *must* match those declared in include/mp4parse.h.
#[repr(C)]
pub struct TrackInfo {
    track_type: u32,
    track_id: u32,
    duration: u64,
    media_time: i64, // wants to be u64? understand how elst adjustment works
}

#[repr(C)]
pub struct TrackAudioInfo {
    channels: u16,
    bit_depth: u16,
    sample_rate: u32,
//    profile: i32,
//    extended_profile: i32, // check types
}

#[repr(C)]
pub struct TrackVideoInfo {
    display_width: u32,
    display_height: u32,
    image_width: u16,
    image_height: u16,
}

/// Allocate an opaque rust-side parser context.
#[no_mangle]
pub extern "C" fn mp4parse_new() -> *mut MediaContext {
    let mut context = Box::new(MediaContext::new());
    context.trace(true);
    Box::into_raw(context)
}

/// Free a rust-side parser context.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_free(context: *mut MediaContext) {
    assert!(!context.is_null());
    let _ = Box::from_raw(context);
}

/// Feed a buffer through `read_mp4()` with the given rust-side
/// parser context, returning the number of detected tracks.
///
/// This is safe to call with NULL arguments but will crash
/// if given invalid pointers, as is usual for C.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_read(context: *mut MediaContext, buffer: *const u8, size: usize) -> i32 {
    // Validate arguments from C.
    if context.is_null() || buffer.is_null() || size < 8 {
        return -1;
    }

    let mut context: &mut MediaContext = &mut *context;

    // Wrap the buffer we've been give in a slice.
    let b = std::slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread to catch any panics.
    let task = std::thread::spawn(move || {
        match read_mp4(&mut c, &mut context) {
            Ok(_) => {},
            Err(Error::UnexpectedEOF) => {},
            Err(e) => { panic!(e); },
        }
        // Make sure the track count fits in an i32 so we can use
        // negative values for failure.
        assert!(context.tracks.len() < i32::max_value() as usize);
        context.tracks.len() as i32
    });
    task.join().unwrap_or(-1)
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_info(context: *mut MediaContext, track: i32, info: *mut TrackInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    (*info).track_type = match track.track_type {
        TrackType::Video => TRACK_TYPE_H264,
        TrackType::Audio => TRACK_TYPE_AAC,
        TrackType::Unknown => return -1,
    };

    // Maybe context & track should just have a single simple is_valid() instead?
    if context.timescale.is_none() ||
        track.timescale.is_none() ||
        track.empty_duration.is_none() ||
        track.media_time.is_none() ||
        track.duration.is_none() ||
        track.track_id.is_none() {
            return -1;
        }

    let empty_duration = media_time_to_ms(track.empty_duration.unwrap(), context.timescale.unwrap());
    (*info).media_time = track_time_to_ms(track.media_time.unwrap(), track.timescale.unwrap()) as i64 - empty_duration as i64;
    (*info).duration = track_time_to_ms(track.duration.unwrap(), track.timescale.unwrap());
    (*info).track_id = track.track_id.unwrap();

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_audio_info(context: *mut MediaContext, track: i32, info: *mut TrackAudioInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Audio => {},
        _ => return -1,
    };

    let audio = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let audio = match audio {
        &SampleEntry::Audio(ref x) => x,
        _ => return -1,
    };

    (*info).channels = audio.channelcount;
    (*info).bit_depth = audio.samplesize;
    (*info).sample_rate = audio.samplerate >> 16; // 16.16 fixed point

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_video_info(context: *mut MediaContext, track: i32, info: *mut TrackVideoInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Video => {},
        _ => return -1,
    };

    let video = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let video = match video {
        &SampleEntry::Video(ref x) => x,
        _ => return -1,
    };

    if let Some(ref tkhd) = track.tkhd {
        (*info).display_width = tkhd.width >> 16; // 16.16 fixed point
        (*info).display_height = tkhd.height >> 16; // 16.16 fixed point
    } else {
        return -1
    }
    (*info).image_width = video.width;
    (*info).image_width = video.height;

    0
}

#[test]
fn new_context() {
    let context = mp4parse_new();
    assert!(!context.is_null());
    unsafe { mp4parse_free(context); }
}

#[test]
#[should_panic(expected = "assertion failed")]
fn free_null_context() {
    unsafe { mp4parse_free(std::ptr::null_mut()); }
}

#[test]
fn arg_validation() {
    let null_buffer = std::ptr::null();
    let null_context = std::ptr::null_mut();

    let context = mp4parse_new();
    assert!(!context.is_null());

    let buffer = vec![0u8; 8];

    unsafe {
        assert_eq!(-1, mp4parse_read(null_context, null_buffer, 0));
        assert_eq!(-1, mp4parse_read(context, null_buffer, 0));
    }

    for size in 0..buffer.len() {
        println!("testing buffer length {}", size);
        unsafe {
            assert_eq!(-1, mp4parse_read(context, buffer.as_ptr(), size));
        }
    }

    unsafe { mp4parse_free(context); }
}
<commit_msg>Remove MediaContext tracing via C API enabled during debugging.<commit_after>//! C API for mp4parse module.
//!
//! Parses ISO Base Media Format aka video/mp4 streams.
//!
//! # Examples
//!
//! ```rust
//! extern crate mp4parse;
//!
//! // Minimal valid mp4 containing no tracks.
//! let data = b"\0\0\0\x0cftypmp42";
//!
//! let context = mp4parse::mp4parse_new();
//! unsafe {
//!     let rv = mp4parse::mp4parse_read(context, data.as_ptr(), data.len());
//!     assert_eq!(0, rv);
//!     mp4parse::mp4parse_free(context);
//! }
//! ```

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

use std;
use std::io::Cursor;

// Symbols we need from our rust api.
use MediaContext;
use TrackType;
use read_mp4;
use Error;
use media_time_to_ms;
use track_time_to_ms;
use SampleEntry;

const TRACK_TYPE_H264: u32 = 0;
const TRACK_TYPE_AAC:  u32 = 1;

// These structs *must* match those declared in include/mp4parse.h.
#[repr(C)]
pub struct TrackInfo {
    track_type: u32,
    track_id: u32,
    duration: u64,
    media_time: i64, // wants to be u64? understand how elst adjustment works
}

#[repr(C)]
pub struct TrackAudioInfo {
    channels: u16,
    bit_depth: u16,
    sample_rate: u32,
//    profile: i32,
//    extended_profile: i32, // check types
}

#[repr(C)]
pub struct TrackVideoInfo {
    display_width: u32,
    display_height: u32,
    image_width: u16,
    image_height: u16,
}

/// Allocate an opaque rust-side parser context.
#[no_mangle]
pub extern "C" fn mp4parse_new() -> *mut MediaContext {
    let context = Box::new(MediaContext::new());
    Box::into_raw(context)
}

/// Free a rust-side parser context.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_free(context: *mut MediaContext) {
    assert!(!context.is_null());
    let _ = Box::from_raw(context);
}

/// Feed a buffer through `read_mp4()` with the given rust-side
/// parser context, returning the number of detected tracks.
///
/// This is safe to call with NULL arguments but will crash
/// if given invalid pointers, as is usual for C.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_read(context: *mut MediaContext, buffer: *const u8, size: usize) -> i32 {
    // Validate arguments from C.
    if context.is_null() || buffer.is_null() || size < 8 {
        return -1;
    }

    let mut context: &mut MediaContext = &mut *context;

    // Wrap the buffer we've been give in a slice.
    let b = std::slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread to catch any panics.
    let task = std::thread::spawn(move || {
        match read_mp4(&mut c, &mut context) {
            Ok(_) => {},
            Err(Error::UnexpectedEOF) => {},
            Err(e) => { panic!(e); },
        }
        // Make sure the track count fits in an i32 so we can use
        // negative values for failure.
        assert!(context.tracks.len() < i32::max_value() as usize);
        context.tracks.len() as i32
    });
    task.join().unwrap_or(-1)
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_info(context: *mut MediaContext, track: i32, info: *mut TrackInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    (*info).track_type = match track.track_type {
        TrackType::Video => TRACK_TYPE_H264,
        TrackType::Audio => TRACK_TYPE_AAC,
        TrackType::Unknown => return -1,
    };

    // Maybe context & track should just have a single simple is_valid() instead?
    if context.timescale.is_none() ||
        track.timescale.is_none() ||
        track.empty_duration.is_none() ||
        track.media_time.is_none() ||
        track.duration.is_none() ||
        track.track_id.is_none() {
            return -1;
        }

    let empty_duration = media_time_to_ms(track.empty_duration.unwrap(), context.timescale.unwrap());
    (*info).media_time = track_time_to_ms(track.media_time.unwrap(), track.timescale.unwrap()) as i64 - empty_duration as i64;
    (*info).duration = track_time_to_ms(track.duration.unwrap(), track.timescale.unwrap());
    (*info).track_id = track.track_id.unwrap();

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_audio_info(context: *mut MediaContext, track: i32, info: *mut TrackAudioInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Audio => {},
        _ => return -1,
    };

    let audio = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let audio = match audio {
        &SampleEntry::Audio(ref x) => x,
        _ => return -1,
    };

    (*info).channels = audio.channelcount;
    (*info).bit_depth = audio.samplesize;
    (*info).sample_rate = audio.samplerate >> 16; // 16.16 fixed point

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_video_info(context: *mut MediaContext, track: i32, info: *mut TrackVideoInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Video => {},
        _ => return -1,
    };

    let video = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let video = match video {
        &SampleEntry::Video(ref x) => x,
        _ => return -1,
    };

    if let Some(ref tkhd) = track.tkhd {
        (*info).display_width = tkhd.width >> 16; // 16.16 fixed point
        (*info).display_height = tkhd.height >> 16; // 16.16 fixed point
    } else {
        return -1
    }
    (*info).image_width = video.width;
    (*info).image_width = video.height;

    0
}

#[test]
fn new_context() {
    let context = mp4parse_new();
    assert!(!context.is_null());
    unsafe { mp4parse_free(context); }
}

#[test]
#[should_panic(expected = "assertion failed")]
fn free_null_context() {
    unsafe { mp4parse_free(std::ptr::null_mut()); }
}

#[test]
fn arg_validation() {
    let null_buffer = std::ptr::null();
    let null_context = std::ptr::null_mut();

    let context = mp4parse_new();
    assert!(!context.is_null());

    let buffer = vec![0u8; 8];

    unsafe {
        assert_eq!(-1, mp4parse_read(null_context, null_buffer, 0));
        assert_eq!(-1, mp4parse_read(context, null_buffer, 0));
    }

    for size in 0..buffer.len() {
        println!("testing buffer length {}", size);
        unsafe {
            assert_eq!(-1, mp4parse_read(context, buffer.as_ptr(), size));
        }
    }

    unsafe { mp4parse_free(context); }
}
<|endoftext|>


<commit_before>//! # C API for rav1e
//!
//! [rav1e](https://github.com/xiph/rav1e/) is an [AV1](https://aomediacodec.github.io/av1-spec/)
//! encoder written in [Rust](https://rust-lang.org)
//!
//! This is the C-compatible API

extern crate rav1e;
extern crate libc;

use std::slice;
use std::sync::Arc;

use std::ffi::CStr;
use std::os::raw::c_char;
use std::os::raw::c_int;

use libc::size_t;
use libc::ptrdiff_t;

/// Raw video Frame
///
/// It can be allocated throught rav1e_frame_new(), populated using rav1e_frame_fill_plane()
/// and freed using rav1e_frame_unref().
pub struct Frame(Arc<rav1e::Frame>);

type EncoderStatus=rav1e::EncoderStatus;

/// Encoder configuration
///
/// Instantiate it using rav1e_config_default() and fine-tune it using
/// rav1e_config_parse().
///
/// Use rav1e_config_unref() to free its memory.
pub struct Config {
    cfg: rav1e::Config,
    last_err: Option<EncoderStatus>,
}

/// Encoder context
///
/// Contains the encoding state, it is created by rav1e_context_new() using an
/// Encoder configuration.
///
/// Use rav1e_context_unref() to free its memory.
pub struct Context {
    ctx: rav1e::Context,
    last_err: Option<EncoderStatus>,
}

type FrameType = rav1e::FrameType;

/// Encoded Packet
///
/// The encoded packets are retrieved using rav1e_receive_packet().
///
/// Use rav1e_packet_unref() to free its memory.
#[repr(C)]
pub struct Packet {
    /// Encoded data buffer
    pub data: *const u8,
    /// Encoded data buffer size
    pub len: size_t,
    /// Frame sequence number
    pub number: u64,
    /// Frame type
    pub frame_type: FrameType,
}

type ChromaSamplePosition=rav1e::ChromaSamplePosition;
type ChromaSampling=rav1e::ChromaSampling;
type MatrixCoefficients=rav1e::MatrixCoefficients;
type ColorPrimaries=rav1e::ColorPrimaries;
type TransferCharacteristics=rav1e::TransferCharacteristics;
type Rational=rav1e::Rational;

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_default(
    width: u32,
    height: u32,
    bit_depth: u8,
    chroma_sampling: ChromaSampling,
    chroma_sample_position: ChromaSamplePosition,
    timebase: Rational,
) -> *mut Config {
    let cfg = rav1e::Config {
        video_info: rav1e::VideoDetails {
            width: width as usize,
            height: height as usize,
            bit_depth: bit_depth as usize,
            chroma_sampling,
            chroma_sample_position,
            mono: false,
            time_base: timebase,
        },
        enc: Default::default(),
    };

    let c = Box::new(Config {
        cfg,
        last_err: None,
    });

    Box::into_raw(c)
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_color_description(cfg: *mut Config,
                                                            matrix: MatrixCoefficients,
                                                            primaries: ColorPrimaries,
                                                            transfer: TransferCharacteristics)
{
    (*cfg).cfg.enc.color_description = Some(rav1e::ColorDescription {
        matrix_coefficients: matrix,
        color_primaries: primaries,
        transfer_characteristics: transfer,
    });
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_unref(cfg: *mut Config) {
    let _ = Box::from_raw(cfg);
}

/// Set a configuration parameter using its key and value
///
/// Available keys and values
/// - "quantizer": 0-255, default 100
/// - "speed": 0-10, default 3
/// - "tune": "psnr"-"psychovisual", default "psnr"
///
/// Returns a negative value on non-zero.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_parse(
    cfg: *mut Config,
    key: *const c_char,
    value: *const c_char,
) -> c_int {
    let key = CStr::from_ptr(key).to_string_lossy();
    let value = CStr::from_ptr(value).to_string_lossy();

    (*cfg)
        .cfg
        .parse(&key, &value)
        .map(|_v| {
            (*cfg).last_err = None;
            0
        }).map_err(|e| (*cfg).last_err = Some(e))
        .unwrap_or(-1)
}

/// Generate a new encoding context from a populated encoder configuration
///
/// Multiple contexts can be generated through it.
#[no_mangle]
pub unsafe extern "C" fn rav1e_context_new(cfg: *const Config) -> *mut Context {
    let ctx = Context {
        ctx: (*cfg).cfg.new_context(),
        last_err: None,
    };

    Box::into_raw(Box::new(ctx))
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_context_unref(ctx: *mut Context) {
    let _ = Box::from_raw(ctx);
}

/// Produce a new frame from the encoding context
///
/// It must be populated using rav1e_frame_fill_plane().
///
/// The frame is reference counted and must be released passing it to rav1e_frame_unref(),
/// see rav1e_send_frame().
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_new(ctx: *const Context) -> *mut Frame {
    let f = (*ctx).ctx.new_frame();
    let frame = Box::new(Frame(f));

    Box::into_raw(frame)
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_unref(frame: *mut Frame) {
    let _ = Arc::from_raw(frame);
}

/// Send the frame for encoding
///
/// The function increases the frame internal reference count and it can be passed multiple
/// times to different rav1e_send_frame().
#[no_mangle]
pub unsafe extern "C" fn rav1e_send_frame(ctx: *mut Context, frame: *const Frame) -> c_int {
    let frame = if frame.is_null() {
        None
    } else {
        Some((*frame).0.clone())
    };

    (*ctx)
        .ctx
        .send_frame(frame)
        .map(|_v| {
            (*ctx).last_err = None;
            0
        }).map_err(|e| (*ctx).last_err = Some(e))
        .unwrap_or(-1)
}

/// Receive encoded data
#[no_mangle]
pub unsafe extern "C" fn rav1e_receive_packet(
    ctx: *mut Context,
    pkt: *mut *mut Packet,
) -> c_int {
    (*ctx)
        .ctx
        .receive_packet()
        .map(|p| {
            (*ctx).last_err = None;
            let packet = Packet {
                data: p.data.as_ptr(),
                len: p.data.len(),
                number: p.number,
                frame_type: p.frame_type,
            };
            *pkt = Box::into_raw(Box::new(packet));
            0
        }).map_err(|e| (*ctx).last_err = Some(e))
        .unwrap_or(-1)
}

#[no_mangle]
pub unsafe extern fn rav1e_packet_unref(pkt: *mut Packet) {
    let _ = Box::from_raw(pkt);
}

/// Produce a sequence header matching the current encoding context
///
/// Its format is compatible with the AV1 Matroska and ISOBMFF specification.
///
/// Use rav1e_container_sequence_header_unref() to free it.
#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header(ctx: *mut Context, buf_size: *mut size_t) -> *mut u8 {
    let buf = (*ctx).ctx.container_sequence_header();

    *buf_size = buf.len();
    Box::into_raw(buf.into_boxed_slice()) as *mut u8
}

#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header_unref(sequence: *mut u8) {
    let _ = Box::from_raw(sequence);
}

/// Fill a frame plane
///
/// Currently the frame contains 3 planes, the first is luminance followed by
/// chrominance.
///
/// The data is copied and this function has to be called for each plane.
///
/// frame: A frame provided by rav1e_frame_new()
/// plane: The index of the plane starting from 0
/// data: The data to be copied
/// data_len: Lenght of the buffer
/// stride: Plane line in bytes, including padding
/// bytewidth: Number of bytes per component, either 1 or 2
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_fill_plane(
    frame: *mut Frame,
    plane: c_int,
    data: *const u8,
    data_len: size_t,
    stride: ptrdiff_t,
    bytewidth: c_int,
) {
    let input = Arc::get_mut(&mut (*frame).0).unwrap();
    let data_slice = slice::from_raw_parts(data, data_len as usize);

    input.planes[plane as usize].copy_from_raw_u8(data_slice, stride as usize, bytewidth as usize);
}
<commit_msg>crav1e: Update API to the unified EncoderConfig<commit_after>//! # C API for rav1e
//!
//! [rav1e](https://github.com/xiph/rav1e/) is an [AV1](https://aomediacodec.github.io/av1-spec/)
//! encoder written in [Rust](https://rust-lang.org)
//!
//! This is the C-compatible API

extern crate rav1e;
extern crate libc;

use std::slice;
use std::sync::Arc;

use std::ffi::CStr;
use std::os::raw::c_char;
use std::os::raw::c_int;

use libc::size_t;
use libc::ptrdiff_t;

/// Raw video Frame
///
/// It can be allocated throught rav1e_frame_new(), populated using rav1e_frame_fill_plane()
/// and freed using rav1e_frame_unref().
pub struct Frame(Arc<rav1e::Frame>);

type EncoderStatus=rav1e::EncoderStatus;

/// Encoder configuration
///
/// Instantiate it using rav1e_config_default() and fine-tune it using
/// rav1e_config_parse().
///
/// Use rav1e_config_unref() to free its memory.
pub struct Config {
    cfg: rav1e::Config,
    last_err: Option<EncoderStatus>,
}

/// Encoder context
///
/// Contains the encoding state, it is created by rav1e_context_new() using an
/// Encoder configuration.
///
/// Use rav1e_context_unref() to free its memory.
pub struct Context {
    ctx: rav1e::Context,
    last_err: Option<EncoderStatus>,
}

type FrameType = rav1e::FrameType;

/// Encoded Packet
///
/// The encoded packets are retrieved using rav1e_receive_packet().
///
/// Use rav1e_packet_unref() to free its memory.
#[repr(C)]
pub struct Packet {
    /// Encoded data buffer
    pub data: *const u8,
    /// Encoded data buffer size
    pub len: size_t,
    /// Frame sequence number
    pub number: u64,
    /// Frame type
    pub frame_type: FrameType,
}

type ChromaSamplePosition=rav1e::ChromaSamplePosition;
type ChromaSampling=rav1e::ChromaSampling;
type MatrixCoefficients=rav1e::MatrixCoefficients;
type ColorPrimaries=rav1e::ColorPrimaries;
type TransferCharacteristics=rav1e::TransferCharacteristics;
type Rational=rav1e::Rational;

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_default(
    width: u32,
    height: u32,
    bit_depth: u8,
    chroma_sampling: ChromaSampling,
    chroma_sample_position: ChromaSamplePosition,
    time_base: Rational,
) -> *mut Config {
    let enc = rav1e::EncoderConfig{
      width : width as usize,
      height: height as usize,
      bit_depth: bit_depth as usize,
      chroma_sampling,
      chroma_sample_position,
      time_base,
      ..Default::default()
    };
    let cfg = rav1e::Config {
        enc
    };

    let c = Box::new(Config {
        cfg,
        last_err: None,
    });

    Box::into_raw(c)
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_set_color_description(cfg: *mut Config,
                                                            matrix: MatrixCoefficients,
                                                            primaries: ColorPrimaries,
                                                            transfer: TransferCharacteristics)
{
    (*cfg).cfg.enc.color_description = Some(rav1e::ColorDescription {
        matrix_coefficients: matrix,
        color_primaries: primaries,
        transfer_characteristics: transfer,
    });
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_config_unref(cfg: *mut Config) {
    let _ = Box::from_raw(cfg);
}

/// Set a configuration parameter using its key and value
///
/// Available keys and values
/// - "quantizer": 0-255, default 100
/// - "speed": 0-10, default 3
/// - "tune": "psnr"-"psychovisual", default "psnr"
///
/// Returns a negative value on non-zero.
#[no_mangle]
pub unsafe extern "C" fn rav1e_config_parse(
    cfg: *mut Config,
    key: *const c_char,
    value: *const c_char,
) -> c_int {
    let key = CStr::from_ptr(key).to_string_lossy();
    let value = CStr::from_ptr(value).to_string_lossy();

    (*cfg)
        .cfg
        .parse(&key, &value)
        .map(|_v| {
            (*cfg).last_err = None;
            0
        }).map_err(|e| (*cfg).last_err = Some(e))
        .unwrap_or(-1)
}

/// Generate a new encoding context from a populated encoder configuration
///
/// Multiple contexts can be generated through it.
#[no_mangle]
pub unsafe extern "C" fn rav1e_context_new(cfg: *const Config) -> *mut Context {
    let ctx = Context {
        ctx: (*cfg).cfg.new_context(),
        last_err: None,
    };

    Box::into_raw(Box::new(ctx))
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_context_unref(ctx: *mut Context) {
    let _ = Box::from_raw(ctx);
}

/// Produce a new frame from the encoding context
///
/// It must be populated using rav1e_frame_fill_plane().
///
/// The frame is reference counted and must be released passing it to rav1e_frame_unref(),
/// see rav1e_send_frame().
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_new(ctx: *const Context) -> *mut Frame {
    let f = (*ctx).ctx.new_frame();
    let frame = Box::new(Frame(f));

    Box::into_raw(frame)
}

#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_unref(frame: *mut Frame) {
    let _ = Arc::from_raw(frame);
}

/// Send the frame for encoding
///
/// The function increases the frame internal reference count and it can be passed multiple
/// times to different rav1e_send_frame().
#[no_mangle]
pub unsafe extern "C" fn rav1e_send_frame(ctx: *mut Context, frame: *const Frame) -> c_int {
    let frame = if frame.is_null() {
        None
    } else {
        Some((*frame).0.clone())
    };

    (*ctx)
        .ctx
        .send_frame(frame)
        .map(|_v| {
            (*ctx).last_err = None;
            0
        }).map_err(|e| (*ctx).last_err = Some(e))
        .unwrap_or(-1)
}

/// Receive encoded data
#[no_mangle]
pub unsafe extern "C" fn rav1e_receive_packet(
    ctx: *mut Context,
    pkt: *mut *mut Packet,
) -> c_int {
    (*ctx)
        .ctx
        .receive_packet()
        .map(|p| {
            (*ctx).last_err = None;
            let packet = Packet {
                data: p.data.as_ptr(),
                len: p.data.len(),
                number: p.number,
                frame_type: p.frame_type,
            };
            *pkt = Box::into_raw(Box::new(packet));
            0
        }).map_err(|e| (*ctx).last_err = Some(e))
        .unwrap_or(-1)
}

#[no_mangle]
pub unsafe extern fn rav1e_packet_unref(pkt: *mut Packet) {
    let _ = Box::from_raw(pkt);
}

/// Produce a sequence header matching the current encoding context
///
/// Its format is compatible with the AV1 Matroska and ISOBMFF specification.
///
/// Use rav1e_container_sequence_header_unref() to free it.
#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header(ctx: *mut Context, buf_size: *mut size_t) -> *mut u8 {
    let buf = (*ctx).ctx.container_sequence_header();

    *buf_size = buf.len();
    Box::into_raw(buf.into_boxed_slice()) as *mut u8
}

#[no_mangle]
pub unsafe extern fn rav1e_container_sequence_header_unref(sequence: *mut u8) {
    let _ = Box::from_raw(sequence);
}

/// Fill a frame plane
///
/// Currently the frame contains 3 planes, the first is luminance followed by
/// chrominance.
///
/// The data is copied and this function has to be called for each plane.
///
/// frame: A frame provided by rav1e_frame_new()
/// plane: The index of the plane starting from 0
/// data: The data to be copied
/// data_len: Lenght of the buffer
/// stride: Plane line in bytes, including padding
/// bytewidth: Number of bytes per component, either 1 or 2
#[no_mangle]
pub unsafe extern "C" fn rav1e_frame_fill_plane(
    frame: *mut Frame,
    plane: c_int,
    data: *const u8,
    data_len: size_t,
    stride: ptrdiff_t,
    bytewidth: c_int,
) {
    let input = Arc::get_mut(&mut (*frame).0).unwrap();
    let data_slice = slice::from_raw_parts(data, data_len as usize);

    input.planes[plane as usize].copy_from_raw_u8(data_slice, stride as usize, bytewidth as usize);
}
<|endoftext|>


<commit_before>use activation_function::ActivationFunction;
use acyclic_network::{NodeType, Network};
pub use acyclic_network::NodeIndex as CppnNodeIndex;
use fixedbitset::FixedBitSet;
use std::fmt::Debug;

pub trait CppnNodeType: NodeType + ActivationFunction {
    fn is_input_node(&self) -> bool;
    fn is_output_node(&self) -> bool;
}

#[derive(Clone, Copy, Debug)]
pub enum CppnNodeKind {
    Bias,
    Input,
    Output,
    Hidden,
}

/// A concrete implementation of a CppnNodeType.
#[derive(Clone, Debug)]
pub struct CppnNode {
    kind: CppnNodeKind,
    activation_function: A,
}

impl CppnNode
    where A: ActivationFunction
{
    pub fn new(kind: CppnNodeKind, activation_function: A) -> Self {
        CppnNode {
            kind: kind,
            activation_function: activation_function,
        }
    }

    pub fn input(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Input, activation_function)
    }

    pub fn output(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Output, activation_function)
    }

    pub fn hidden(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Hidden, activation_function)
    }

    pub fn bias(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Bias, activation_function)
    }
}

impl ActivationFunction for CppnNode {
    fn formula_gnuplot(&self, x: String) -> String {
        match self.kind {
            CppnNodeKind::Input |
            CppnNodeKind::Output |
            CppnNodeKind::Hidden |
            CppnNodeKind::Bias => self.activation_function.formula_gnuplot(x),
        }
    }

    fn calculate(&self, input: f64) -> f64 {
        self.activation_function.calculate(input)
    }
}

impl NodeType for CppnNode {
    fn accept_incoming_links(&self) -> bool {
        match self.kind {
            CppnNodeKind::Hidden | CppnNodeKind::Output => true,
            CppnNodeKind::Input | CppnNodeKind::Bias => false,
        }
    }

    fn accept_outgoing_links(&self) -> bool {
        match self.kind {
            CppnNodeKind::Hidden | CppnNodeKind::Input | CppnNodeKind::Bias => true,
            CppnNodeKind::Output => false,
        }
    }
}

impl CppnNodeType for CppnNode {
    fn is_input_node(&self) -> bool {
        match self.kind {
            CppnNodeKind::Input => true,
            _ => false,
        }
    }
    fn is_output_node(&self) -> bool {
        match self.kind {
            CppnNodeKind::Output => true,
            _ => false,
        }
    }
}

pub type CppnGraph<N, L, EXTID>
    where N: CppnNodeType,
          L: Copy + Debug + Send + Sized + Into<f64>,
          EXTID: Copy + Debug + Send + Sized + Ord = Network<N, L, EXTID>;

/// Represents a Compositional Pattern Producing Network (CPPN)
pub struct Cppn<'a, N, L, EXTID>
    where N: CppnNodeType + 'a,
          L: Copy + Debug + Send + Sized + Into<f64> + 'a,
          EXTID: Copy + Debug + Send + Sized + Ord + 'a
{
    graph: &'a CppnGraph<N, L, EXTID>,
    inputs: Vec<CppnNodeIndex>,
    outputs: Vec<CppnNodeIndex>,

    // For each node in `graph` there exists a corresponding field in `incoming_signals` describing
    // the sum of all input signals for that node.  We could store it inline in the `CppnNode`, but
    // this would require to make the whole CppnGraph mutable.
    incoming_signals: Vec<f64>,
}

impl<'a, N, L, EXTID> Cppn<'a, N, L, EXTID>
    where N: CppnNodeType + 'a,
          L: Copy + Debug + Send + Sized + Into<f64> + 'a,
          EXTID: Copy + Debug + Send + Sized + Ord + 'a
{
    pub fn new(graph: &'a CppnGraph<N, L, EXTID>) -> Cppn<'a, N, L, EXTID> {
        let mut inputs = Vec::new();
        let mut outputs = Vec::new();

        graph.each_node_with_index(|node, index| {
            if node.node_type().is_input_node() {
                inputs.push(index);
            }
            if node.node_type().is_output_node() {
                outputs.push(index);
            }
        });

        Cppn {
            graph: graph,
            inputs: inputs,
            outputs: outputs,
            incoming_signals: graph.nodes().iter().map(|_| 0.0).collect(),
        }
    }

    fn set_signal(&mut self, node_idx: CppnNodeIndex, value: f64) {
        self.incoming_signals[node_idx.index()] = value;
    }

    fn reset_signals(&mut self) {
        for value in self.incoming_signals.iter_mut() {
            *value = 0.0;
        }
    }

    /// Forward-propagate the signals starting from `from_nodes`. We use
    /// breadth-first-search (BFS).
    fn propagate_signals(&mut self, mut nodes: Vec<CppnNodeIndex>, mut seen: FixedBitSet) {
        while let Some(node_idx) = nodes.pop() {
            let input = self.incoming_signals[node_idx.index()];
            let output = self.graph.node(node_idx).node_type().calculate(input);

            // propagate output signal to outgoing links.
            self.graph.each_active_forward_link_of_node(node_idx, |out_node_idx, weight| {
                let out_node = out_node_idx.index();
                let weight: f64 = weight.into();
                self.incoming_signals[out_node] += weight * output;
                if !seen.contains(out_node) {
                    seen.insert(out_node);
                    nodes.push(out_node_idx);
                }
            });
        }
    }

    /// Calculate all outputs
    pub fn calculate(&mut self, inputs: &[&[f64]]) -> Vec<f64> {
        assert!(self.incoming_signals.len() == self.graph.nodes().len());
        self.reset_signals();

        // assign all inputs
        let mut i = 0;
        for input_list in inputs.iter() {
            for &input in input_list.iter() {
                let input_idx = self.inputs[i];
                self.set_signal(input_idx, input);
                i += 1;
            }
        }
        assert!(i == self.inputs.len());

        let mut nodes = Vec::new(); // XXX: worst case capacity
        let mut seen = FixedBitSet::with_capacity(self.incoming_signals.len());

        // start from all nodes which have zero in_degree()
        self.graph.each_node_with_index(|node, index| {
            if node.in_degree() == 0 {
                nodes.push(index);
                seen.insert(index.index());
            }
        });

        // propagate the signals starting from the nodes with zero in degree.
        self.propagate_signals(nodes, seen);

        self.outputs
            .iter()
            .map(|&node_idx| {
                let input = self.incoming_signals[node_idx.index()];
                let output = self.graph.node(node_idx).node_type().calculate(input);
                output
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use activation_function::GeometricActivationFunction as AF;
    use super::{Cppn, CppnGraph, CppnNode};
    use acyclic_network::ExternalId;
    use rand;

    #[test]
    fn test_cycle() {
        let mut g = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let h1 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(2));
        let h2 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(3));
        assert_eq!(true, g.valid_link(i1, i1).is_err());
        assert_eq!(true, g.valid_link(h1, h1).is_err());

        assert_eq!(true, g.valid_link(h1, i1).is_err());
        assert_eq!(Ok(()), g.valid_link(i1, h1));
        assert_eq!(Ok(()), g.valid_link(i1, h2));
        assert_eq!(Ok(()), g.valid_link(h1, h2));

        g.add_link(i1, h1, 0.0, ExternalId(1));
        assert_eq!(true, g.link_would_cycle(h1, i1));
        assert_eq!(false, g.link_would_cycle(i1, h1));
        assert_eq!(false, g.link_would_cycle(i1, h2));
        assert_eq!(true, g.link_would_cycle(i1, i1));
        assert_eq!(false, g.link_would_cycle(h1, h2));
        assert_eq!(false, g.link_would_cycle(h2, h1));
        assert_eq!(false, g.link_would_cycle(h2, i1));

        g.add_link(h1, h2, 0.0, ExternalId(2));
        assert_eq!(true, g.link_would_cycle(h2, i1));
        assert_eq!(true, g.link_would_cycle(h1, i1));
        assert_eq!(true, g.link_would_cycle(h2, h1));
        assert_eq!(false, g.link_would_cycle(i1, h2));
    }

    #[test]
    fn test_simple_cppn() {
        let mut g = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let h1 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(2));
        let o1 = g.add_node(CppnNode::output(AF::Linear), ExternalId(3));
        g.add_link(i1, h1, 0.5, ExternalId(1));
        g.add_link(h1, o1, 1.0, ExternalId(2));

        let mut cppn = Cppn::new(&g);

        let f = |x| 0.5 * x * 1.0;
        assert_eq!(vec![f(0.5)], cppn.calculate(&[&[0.5]]));
        assert_eq!(vec![f(4.0)], cppn.calculate(&[&[4.0]]));
        assert_eq!(vec![f(-4.0)], cppn.calculate(&[&[-4.0]]));
    }

    #[test]
    fn test_cppn_with_output_activation_function() {
        let mut g = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let h1 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(2));
        let o1 = g.add_node(CppnNode::output(AF::Constant1), ExternalId(3));
        g.add_link(i1, h1, 0.5, ExternalId(1));
        g.add_link(h1, o1, 1.0, ExternalId(2));

        let mut cppn = Cppn::new(&g);

        assert_eq!(vec![1.0], cppn.calculate(&[&[0.5]]));
        assert_eq!(vec![1.0], cppn.calculate(&[&[4.0]]));
        assert_eq!(vec![1.0], cppn.calculate(&[&[-4.0]]));
    }

    #[test]
    fn test_find_random_unconnected_link_no_cycle() {
        let mut g: CppnGraph<CppnNode<AF>, _, _> = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let o1 = g.add_node(CppnNode::output(AF::Linear), ExternalId(2));
        let o2 = g.add_node(CppnNode::output(AF::Linear), ExternalId(3));

        let mut rng = rand::thread_rng();

        let link = g.find_random_unconnected_link_no_cycle(&mut rng);
        assert_eq!(true, link.is_some());
        let l = link.unwrap();
        assert!((i1, o1) == l || (i1, o2) == l);

        g.add_link(i1, o2, 0.0, ExternalId(1));
        let link = g.find_random_unconnected_link_no_cycle(&mut rng);
        assert_eq!(true, link.is_some());
        assert_eq!((i1, o1), link.unwrap());

        g.add_link(i1, o1, 0.0, ExternalId(2));
        let link = g.find_random_unconnected_link_no_cycle(&mut rng);
        assert_eq!(false, link.is_some());
    }

}
<commit_msg>Make kind and activation_function fields public<commit_after>use activation_function::ActivationFunction;
use acyclic_network::{NodeType, Network};
pub use acyclic_network::NodeIndex as CppnNodeIndex;
use fixedbitset::FixedBitSet;
use std::fmt::Debug;

pub trait CppnNodeType: NodeType + ActivationFunction {
    fn is_input_node(&self) -> bool;
    fn is_output_node(&self) -> bool;
}

#[derive(Clone, Copy, Debug)]
pub enum CppnNodeKind {
    Bias,
    Input,
    Output,
    Hidden,
}

/// A concrete implementation of a CppnNodeType.
#[derive(Clone, Debug)]
pub struct CppnNode {
    pub kind: CppnNodeKind,
    pub activation_function: A,
}

impl CppnNode
    where A: ActivationFunction
{
    pub fn new(kind: CppnNodeKind, activation_function: A) -> Self {
        CppnNode {
            kind: kind,
            activation_function: activation_function,
        }
    }

    pub fn input(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Input, activation_function)
    }

    pub fn output(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Output, activation_function)
    }

    pub fn hidden(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Hidden, activation_function)
    }

    pub fn bias(activation_function: A) -> Self {
        Self::new(CppnNodeKind::Bias, activation_function)
    }
}

impl ActivationFunction for CppnNode {
    fn formula_gnuplot(&self, x: String) -> String {
        match self.kind {
            CppnNodeKind::Input |
            CppnNodeKind::Output |
            CppnNodeKind::Hidden |
            CppnNodeKind::Bias => self.activation_function.formula_gnuplot(x),
        }
    }

    fn calculate(&self, input: f64) -> f64 {
        self.activation_function.calculate(input)
    }
}

impl NodeType for CppnNode {
    fn accept_incoming_links(&self) -> bool {
        match self.kind {
            CppnNodeKind::Hidden | CppnNodeKind::Output => true,
            CppnNodeKind::Input | CppnNodeKind::Bias => false,
        }
    }

    fn accept_outgoing_links(&self) -> bool {
        match self.kind {
            CppnNodeKind::Hidden | CppnNodeKind::Input | CppnNodeKind::Bias => true,
            CppnNodeKind::Output => false,
        }
    }
}

impl CppnNodeType for CppnNode {
    fn is_input_node(&self) -> bool {
        match self.kind {
            CppnNodeKind::Input => true,
            _ => false,
        }
    }
    fn is_output_node(&self) -> bool {
        match self.kind {
            CppnNodeKind::Output => true,
            _ => false,
        }
    }
}

pub type CppnGraph<N, L, EXTID>
    where N: CppnNodeType,
          L: Copy + Debug + Send + Sized + Into<f64>,
          EXTID: Copy + Debug + Send + Sized + Ord = Network<N, L, EXTID>;

/// Represents a Compositional Pattern Producing Network (CPPN)
pub struct Cppn<'a, N, L, EXTID>
    where N: CppnNodeType + 'a,
          L: Copy + Debug + Send + Sized + Into<f64> + 'a,
          EXTID: Copy + Debug + Send + Sized + Ord + 'a
{
    graph: &'a CppnGraph<N, L, EXTID>,
    inputs: Vec<CppnNodeIndex>,
    outputs: Vec<CppnNodeIndex>,

    // For each node in `graph` there exists a corresponding field in `incoming_signals` describing
    // the sum of all input signals for that node.  We could store it inline in the `CppnNode`, but
    // this would require to make the whole CppnGraph mutable.
    incoming_signals: Vec<f64>,
}

impl<'a, N, L, EXTID> Cppn<'a, N, L, EXTID>
    where N: CppnNodeType + 'a,
          L: Copy + Debug + Send + Sized + Into<f64> + 'a,
          EXTID: Copy + Debug + Send + Sized + Ord + 'a
{
    pub fn new(graph: &'a CppnGraph<N, L, EXTID>) -> Cppn<'a, N, L, EXTID> {
        let mut inputs = Vec::new();
        let mut outputs = Vec::new();

        graph.each_node_with_index(|node, index| {
            if node.node_type().is_input_node() {
                inputs.push(index);
            }
            if node.node_type().is_output_node() {
                outputs.push(index);
            }
        });

        Cppn {
            graph: graph,
            inputs: inputs,
            outputs: outputs,
            incoming_signals: graph.nodes().iter().map(|_| 0.0).collect(),
        }
    }

    fn set_signal(&mut self, node_idx: CppnNodeIndex, value: f64) {
        self.incoming_signals[node_idx.index()] = value;
    }

    fn reset_signals(&mut self) {
        for value in self.incoming_signals.iter_mut() {
            *value = 0.0;
        }
    }

    /// Forward-propagate the signals starting from `from_nodes`. We use
    /// breadth-first-search (BFS).
    fn propagate_signals(&mut self, mut nodes: Vec<CppnNodeIndex>, mut seen: FixedBitSet) {
        while let Some(node_idx) = nodes.pop() {
            let input = self.incoming_signals[node_idx.index()];
            let output = self.graph.node(node_idx).node_type().calculate(input);

            // propagate output signal to outgoing links.
            self.graph.each_active_forward_link_of_node(node_idx, |out_node_idx, weight| {
                let out_node = out_node_idx.index();
                let weight: f64 = weight.into();
                self.incoming_signals[out_node] += weight * output;
                if !seen.contains(out_node) {
                    seen.insert(out_node);
                    nodes.push(out_node_idx);
                }
            });
        }
    }

    /// Calculate all outputs
    pub fn calculate(&mut self, inputs: &[&[f64]]) -> Vec<f64> {
        assert!(self.incoming_signals.len() == self.graph.nodes().len());
        self.reset_signals();

        // assign all inputs
        let mut i = 0;
        for input_list in inputs.iter() {
            for &input in input_list.iter() {
                let input_idx = self.inputs[i];
                self.set_signal(input_idx, input);
                i += 1;
            }
        }
        assert!(i == self.inputs.len());

        let mut nodes = Vec::new(); // XXX: worst case capacity
        let mut seen = FixedBitSet::with_capacity(self.incoming_signals.len());

        // start from all nodes which have zero in_degree()
        self.graph.each_node_with_index(|node, index| {
            if node.in_degree() == 0 {
                nodes.push(index);
                seen.insert(index.index());
            }
        });

        // propagate the signals starting from the nodes with zero in degree.
        self.propagate_signals(nodes, seen);

        self.outputs
            .iter()
            .map(|&node_idx| {
                let input = self.incoming_signals[node_idx.index()];
                let output = self.graph.node(node_idx).node_type().calculate(input);
                output
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use activation_function::GeometricActivationFunction as AF;
    use super::{Cppn, CppnGraph, CppnNode};
    use acyclic_network::ExternalId;
    use rand;

    #[test]
    fn test_cycle() {
        let mut g = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let h1 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(2));
        let h2 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(3));
        assert_eq!(true, g.valid_link(i1, i1).is_err());
        assert_eq!(true, g.valid_link(h1, h1).is_err());

        assert_eq!(true, g.valid_link(h1, i1).is_err());
        assert_eq!(Ok(()), g.valid_link(i1, h1));
        assert_eq!(Ok(()), g.valid_link(i1, h2));
        assert_eq!(Ok(()), g.valid_link(h1, h2));

        g.add_link(i1, h1, 0.0, ExternalId(1));
        assert_eq!(true, g.link_would_cycle(h1, i1));
        assert_eq!(false, g.link_would_cycle(i1, h1));
        assert_eq!(false, g.link_would_cycle(i1, h2));
        assert_eq!(true, g.link_would_cycle(i1, i1));
        assert_eq!(false, g.link_would_cycle(h1, h2));
        assert_eq!(false, g.link_would_cycle(h2, h1));
        assert_eq!(false, g.link_would_cycle(h2, i1));

        g.add_link(h1, h2, 0.0, ExternalId(2));
        assert_eq!(true, g.link_would_cycle(h2, i1));
        assert_eq!(true, g.link_would_cycle(h1, i1));
        assert_eq!(true, g.link_would_cycle(h2, h1));
        assert_eq!(false, g.link_would_cycle(i1, h2));
    }

    #[test]
    fn test_simple_cppn() {
        let mut g = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let h1 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(2));
        let o1 = g.add_node(CppnNode::output(AF::Linear), ExternalId(3));
        g.add_link(i1, h1, 0.5, ExternalId(1));
        g.add_link(h1, o1, 1.0, ExternalId(2));

        let mut cppn = Cppn::new(&g);

        let f = |x| 0.5 * x * 1.0;
        assert_eq!(vec![f(0.5)], cppn.calculate(&[&[0.5]]));
        assert_eq!(vec![f(4.0)], cppn.calculate(&[&[4.0]]));
        assert_eq!(vec![f(-4.0)], cppn.calculate(&[&[-4.0]]));
    }

    #[test]
    fn test_cppn_with_output_activation_function() {
        let mut g = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let h1 = g.add_node(CppnNode::hidden(AF::Linear), ExternalId(2));
        let o1 = g.add_node(CppnNode::output(AF::Constant1), ExternalId(3));
        g.add_link(i1, h1, 0.5, ExternalId(1));
        g.add_link(h1, o1, 1.0, ExternalId(2));

        let mut cppn = Cppn::new(&g);

        assert_eq!(vec![1.0], cppn.calculate(&[&[0.5]]));
        assert_eq!(vec![1.0], cppn.calculate(&[&[4.0]]));
        assert_eq!(vec![1.0], cppn.calculate(&[&[-4.0]]));
    }

    #[test]
    fn test_find_random_unconnected_link_no_cycle() {
        let mut g: CppnGraph<CppnNode<AF>, _, _> = CppnGraph::new();
        let i1 = g.add_node(CppnNode::input(AF::Linear), ExternalId(1));
        let o1 = g.add_node(CppnNode::output(AF::Linear), ExternalId(2));
        let o2 = g.add_node(CppnNode::output(AF::Linear), ExternalId(3));

        let mut rng = rand::thread_rng();

        let link = g.find_random_unconnected_link_no_cycle(&mut rng);
        assert_eq!(true, link.is_some());
        let l = link.unwrap();
        assert!((i1, o1) == l || (i1, o2) == l);

        g.add_link(i1, o2, 0.0, ExternalId(1));
        let link = g.find_random_unconnected_link_no_cycle(&mut rng);
        assert_eq!(true, link.is_some());
        assert_eq!((i1, o1), link.unwrap());

        g.add_link(i1, o1, 0.0, ExternalId(2));
        let link = g.find_random_unconnected_link_no_cycle(&mut rng);
        assert_eq!(false, link.is_some());
    }

}
<|endoftext|>


<commit_before>//! Graph traits for associated data and graph construction.

use Graph;
use stable_graph::StableGraph;
use ::{
    EdgeType,
};
use graph::IndexType;
use graphmap::{GraphMap, NodeTrait};
use visit::{
    Data,
    NodeCount,
    NodeIndexable,
    Reversed,
};

trait_template!{
    /// Access node and edge weights (associated data).
pub trait DataMap : Data {
    @section self_ref
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight>;
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight>;
}
}

macro_rules! access0 {
    ($e:expr) => ($e.0);
}

DataMap!{delegate_impl []}
DataMap!{delegate_impl [['a, G], G, &'a mut G, deref_twice]}
DataMap!{delegate_impl [[G], G, Reversed<G>, access0]}

trait_template! {
    /// Access node and edge weights mutably.
pub trait DataMapMut : DataMap {
    @section self_mut
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight>;
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight>;
}
}

DataMapMut!{delegate_impl [['a, G], G, &'a mut G, deref_twice]}
DataMapMut!{delegate_impl [[G], G, Reversed<G>, access0]}

/// A graph that can be extended with further nodes and edges
pub trait Build : Data + NodeCount {
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId;
    /// Add a new edge. If parallel edges (duplicate) are not allowed and
    /// the edge already exists, return `None`.
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId> {
        Some(self.update_edge(a, b, weight))
    }
    /// Add or update the edge from `a` to `b`. Return the id of the affected
    /// edge.
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId;
}

/// A graph that can be created
pub trait Create : Build + Default {
    fn with_capacity(nodes: usize, edges: usize) -> Self;
}

impl<N, E, Ty, Ix> Data for Graph<N, E, Ty, Ix>
    where Ix: IndexType
{
    type NodeWeight = N;
    type EdgeWeight = E;
}

impl<N, E, Ty, Ix> DataMap for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        self.node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.edge_weight(id)
    }
}

impl<N, E, Ty, Ix> DataMapMut for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
        self.node_weight_mut(id)
    }
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
        self.edge_weight_mut(id)
    }
}

impl<N, E, Ty, Ix> DataMap for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        self.node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.edge_weight(id)
    }
}

impl<N, E, Ty, Ix> DataMapMut for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
        self.node_weight_mut(id)
    }
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
        self.edge_weight_mut(id)
    }
}

impl<N, E, Ty, Ix> Build for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        Some(self.add_edge(a, b, weight))
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.update_edge(a, b, weight)
    }
}

impl<N, E, Ty, Ix> Build for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        Some(self.add_edge(a, b, weight))
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.update_edge(a, b, weight)
    }
}

impl<N, E, Ty> Build for GraphMap<N, E, Ty>
    where Ty: EdgeType,
          N: NodeTrait,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        if self.contains_edge(a, b) {
            None
        } else {
            let r = self.add_edge(a, b, weight);
            debug_assert!(r.is_none());
            Some((a, b))
        }
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.add_edge(a, b, weight);
        (a, b)
    }
}

impl<N, E, Ty, Ix> Create for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

impl<N, E, Ty, Ix> Create for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

impl<N, E, Ty> Create for GraphMap<N, E, Ty>
    where Ty: EdgeType,
          N: NodeTrait,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

/// A graph element.
///
/// A sequence of Elements, for example an iterator, is laid out as follows:
/// Nodes are implicitly given the index of their appearance in the sequence.
/// The edges’ source and target fields refer to these indices.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Element<N, E> {
    /// A graph node.
    Node {
        weight: N,
    },
    /// A graph edge.
    Edge {
        source: usize,
        target: usize,
        weight: E,
    }
}

/// Create a graph from an iterator of elements.
pub trait FromElements : Create {
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        let mut gr = Self::with_capacity(0, 0);
        // usize -> NodeId map
        let mut map = Vec::new();
        for element in iterable {
            match element {
                Element::Node { weight } => {
                    map.push(gr.add_node(weight));
                }
                Element::Edge { source, target, weight } => {
                    gr.add_edge(map[source], map[target], weight);
                }
            }
        }
        gr
    }

}

fn from_elements_indexable<G, I>(iterable: I) -> G
    where G: Create + NodeIndexable,
          I: IntoIterator<Item=Element<G::NodeWeight, G::EdgeWeight>>,
{
    let mut gr = G::with_capacity(0, 0);
    let map = |gr: &G, i| gr.from_index(i);
    for element in iterable {
        match element {
            Element::Node { weight } => {
                gr.add_node(weight);
            }
            Element::Edge { source, target, weight } => {
                let from = map(&gr, source);
                let to = map(&gr, target);
                gr.add_edge(from, to, weight);
            }
        }
    }
    gr
}

impl<N, E, Ty, Ix> FromElements for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}

impl<N, E, Ty, Ix> FromElements for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}

impl<N, E, Ty> FromElements for GraphMap<N, E, Ty>
    where Ty: EdgeType,
          N: NodeTrait,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}

/// Iterator adaptors for iterators of `Element`.
pub trait ElementIterator<N, E> : Iterator<Item=Element<N, E>> {
    /// Create an iterator adaptor that filters graph elements.
    ///
    /// The function `f` is called with each element and if its return value
    /// is `true` the element is accepted and if `false` it is removed.
    ///
    /// `f` is called with mutable references to the node and edge weights,
    /// so that they can be mutated (but the edge endpoints can not).
    ///
    /// This filter adapts the edge source, target numbers in the
    /// stream so that they are correct after the removals.
    fn filter_elements<F>(self, f: F) -> FilterElements<Self, F>
        where Self: Sized,
              F: FnMut(Element<&mut N, &mut E>) -> bool,
    {
        FilterElements {
            iter: self,
            node_index: 0,
            map: Vec::new(),
            f: f,
        }
    }
}

impl<N, E, I> ElementIterator<N, E> for I
    where I: Iterator<Item=Element<N, E>> { }

/// An iterator that filters graph elements.
///
/// See [`.filter_elements()`][1] for more information.
///
/// [1]: trait.ElementIterator.html#method.filter_elements
pub struct FilterElements<I, F> {
    iter: I,
    node_index: usize,
    map: Vec<usize>,
    f: F,
}

impl<I, F, N, E> Iterator for FilterElements<I, F>
    where I: Iterator<Item=Element<N, E>>,
          F: FnMut(Element<&mut N, &mut E>) -> bool,
{
    type Item = Element<N, E>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let mut elt = match self.iter.next() {
                None => return None,
                Some(elt) => elt,
            };
            let keep = (self.f)(match elt {
                Element::Node { ref mut weight } => Element::Node { weight: weight },
                Element::Edge { source, target, ref mut weight }
                    => Element::Edge { source: source, target: target, weight: weight },
            });
            let is_node = if let Element::Node { .. } = elt { true } else { false };
            if !keep && is_node {
                self.map.push(self.node_index);
            }
            if is_node {
                self.node_index += 1;
            }
            if !keep {
                continue;
            }

            // map edge parts
            match elt {
                Element::Edge {
                    ref mut source,
                    ref mut target,
                    ..
                } => {
                    // Find the node indices in the map of removed ones.
                    // If a node was removed, the edge is as well.
                    // Otherwise the counts are adjusted by the number of nodes
                    // removed.
                    // Example: map: [1, 3, 4, 6]
                    // binary search for 2, result is Err(1). One node has been
                    // removed before 2.
                    match self.map.binary_search(source) {
                        Ok(_) => continue,
                        Err(i) => *source -= i,
                    }
                    match self.map.binary_search(target) {
                        Ok(_) => continue,
                        Err(i) => *target -= i,
                    }
                    println!("Letting through edge {}, {}", *source, *target);
                }
                Element::Node { .. } => {
                    println!("Letting through node {}", self.node_index - 1);
                }
            }
            return Some(elt);
        }
    }
}
<commit_msg>ElementIterator for I: ?Sized<commit_after>//! Graph traits for associated data and graph construction.

use Graph;
use stable_graph::StableGraph;
use ::{
    EdgeType,
};
use graph::IndexType;
use graphmap::{GraphMap, NodeTrait};
use visit::{
    Data,
    NodeCount,
    NodeIndexable,
    Reversed,
};

trait_template!{
    /// Access node and edge weights (associated data).
pub trait DataMap : Data {
    @section self_ref
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight>;
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight>;
}
}

macro_rules! access0 {
    ($e:expr) => ($e.0);
}

DataMap!{delegate_impl []}
DataMap!{delegate_impl [['a, G], G, &'a mut G, deref_twice]}
DataMap!{delegate_impl [[G], G, Reversed<G>, access0]}

trait_template! {
    /// Access node and edge weights mutably.
pub trait DataMapMut : DataMap {
    @section self_mut
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight>;
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight>;
}
}

DataMapMut!{delegate_impl [['a, G], G, &'a mut G, deref_twice]}
DataMapMut!{delegate_impl [[G], G, Reversed<G>, access0]}

/// A graph that can be extended with further nodes and edges
pub trait Build : Data + NodeCount {
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId;
    /// Add a new edge. If parallel edges (duplicate) are not allowed and
    /// the edge already exists, return `None`.
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId> {
        Some(self.update_edge(a, b, weight))
    }
    /// Add or update the edge from `a` to `b`. Return the id of the affected
    /// edge.
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId;
}

/// A graph that can be created
pub trait Create : Build + Default {
    fn with_capacity(nodes: usize, edges: usize) -> Self;
}

impl<N, E, Ty, Ix> Data for Graph<N, E, Ty, Ix>
    where Ix: IndexType
{
    type NodeWeight = N;
    type EdgeWeight = E;
}

impl<N, E, Ty, Ix> DataMap for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        self.node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.edge_weight(id)
    }
}

impl<N, E, Ty, Ix> DataMapMut for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
        self.node_weight_mut(id)
    }
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
        self.edge_weight_mut(id)
    }
}

impl<N, E, Ty, Ix> DataMap for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        self.node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.edge_weight(id)
    }
}

impl<N, E, Ty, Ix> DataMapMut for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
        self.node_weight_mut(id)
    }
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
        self.edge_weight_mut(id)
    }
}

impl<N, E, Ty, Ix> Build for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        Some(self.add_edge(a, b, weight))
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.update_edge(a, b, weight)
    }
}

impl<N, E, Ty, Ix> Build for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        Some(self.add_edge(a, b, weight))
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.update_edge(a, b, weight)
    }
}

impl<N, E, Ty> Build for GraphMap<N, E, Ty>
    where Ty: EdgeType,
          N: NodeTrait,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        if self.contains_edge(a, b) {
            None
        } else {
            let r = self.add_edge(a, b, weight);
            debug_assert!(r.is_none());
            Some((a, b))
        }
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.add_edge(a, b, weight);
        (a, b)
    }
}

impl<N, E, Ty, Ix> Create for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

impl<N, E, Ty, Ix> Create for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

impl<N, E, Ty> Create for GraphMap<N, E, Ty>
    where Ty: EdgeType,
          N: NodeTrait,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

/// A graph element.
///
/// A sequence of Elements, for example an iterator, is laid out as follows:
/// Nodes are implicitly given the index of their appearance in the sequence.
/// The edges’ source and target fields refer to these indices.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Element<N, E> {
    /// A graph node.
    Node {
        weight: N,
    },
    /// A graph edge.
    Edge {
        source: usize,
        target: usize,
        weight: E,
    }
}

/// Create a graph from an iterator of elements.
pub trait FromElements : Create {
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        let mut gr = Self::with_capacity(0, 0);
        // usize -> NodeId map
        let mut map = Vec::new();
        for element in iterable {
            match element {
                Element::Node { weight } => {
                    map.push(gr.add_node(weight));
                }
                Element::Edge { source, target, weight } => {
                    gr.add_edge(map[source], map[target], weight);
                }
            }
        }
        gr
    }

}

fn from_elements_indexable<G, I>(iterable: I) -> G
    where G: Create + NodeIndexable,
          I: IntoIterator<Item=Element<G::NodeWeight, G::EdgeWeight>>,
{
    let mut gr = G::with_capacity(0, 0);
    let map = |gr: &G, i| gr.from_index(i);
    for element in iterable {
        match element {
            Element::Node { weight } => {
                gr.add_node(weight);
            }
            Element::Edge { source, target, weight } => {
                let from = map(&gr, source);
                let to = map(&gr, target);
                gr.add_edge(from, to, weight);
            }
        }
    }
    gr
}

impl<N, E, Ty, Ix> FromElements for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}

impl<N, E, Ty, Ix> FromElements for StableGraph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}

impl<N, E, Ty> FromElements for GraphMap<N, E, Ty>
    where Ty: EdgeType,
          N: NodeTrait,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}

/// Iterator adaptors for iterators of `Element`.
pub trait ElementIterator<N, E> : Iterator<Item=Element<N, E>> {
    /// Create an iterator adaptor that filters graph elements.
    ///
    /// The function `f` is called with each element and if its return value
    /// is `true` the element is accepted and if `false` it is removed.
    ///
    /// `f` is called with mutable references to the node and edge weights,
    /// so that they can be mutated (but the edge endpoints can not).
    ///
    /// This filter adapts the edge source, target numbers in the
    /// stream so that they are correct after the removals.
    fn filter_elements<F>(self, f: F) -> FilterElements<Self, F>
        where Self: Sized,
              F: FnMut(Element<&mut N, &mut E>) -> bool,
    {
        FilterElements {
            iter: self,
            node_index: 0,
            map: Vec::new(),
            f: f,
        }
    }
}

impl<N, E, I: ?Sized> ElementIterator<N, E> for I
    where I: Iterator<Item=Element<N, E>> { }

/// An iterator that filters graph elements.
///
/// See [`.filter_elements()`][1] for more information.
///
/// [1]: trait.ElementIterator.html#method.filter_elements
pub struct FilterElements<I, F> {
    iter: I,
    node_index: usize,
    map: Vec<usize>,
    f: F,
}

impl<I, F, N, E> Iterator for FilterElements<I, F>
    where I: Iterator<Item=Element<N, E>>,
          F: FnMut(Element<&mut N, &mut E>) -> bool,
{
    type Item = Element<N, E>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let mut elt = match self.iter.next() {
                None => return None,
                Some(elt) => elt,
            };
            let keep = (self.f)(match elt {
                Element::Node { ref mut weight } => Element::Node { weight: weight },
                Element::Edge { source, target, ref mut weight }
                    => Element::Edge { source: source, target: target, weight: weight },
            });
            let is_node = if let Element::Node { .. } = elt { true } else { false };
            if !keep && is_node {
                self.map.push(self.node_index);
            }
            if is_node {
                self.node_index += 1;
            }
            if !keep {
                continue;
            }

            // map edge parts
            match elt {
                Element::Edge {
                    ref mut source,
                    ref mut target,
                    ..
                } => {
                    // Find the node indices in the map of removed ones.
                    // If a node was removed, the edge is as well.
                    // Otherwise the counts are adjusted by the number of nodes
                    // removed.
                    // Example: map: [1, 3, 4, 6]
                    // binary search for 2, result is Err(1). One node has been
                    // removed before 2.
                    match self.map.binary_search(source) {
                        Ok(_) => continue,
                        Err(i) => *source -= i,
                    }
                    match self.map.binary_search(target) {
                        Ok(_) => continue,
                        Err(i) => *target -= i,
                    }
                    println!("Letting through edge {}, {}", *source, *target);
                }
                Element::Node { .. } => {
                    println!("Letting through node {}", self.node_index - 1);
                }
            }
            return Some(elt);
        }
    }
}
<|endoftext|>


<commit_before>// Bitcoin secp256k1 bindings
// Written in 2015 by
//   Andrew Poelstra
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # ECDH
//! Support for shared secret computations
//!

use core::ptr;
use core::ops::{FnMut, Deref};

use key::{SecretKey, PublicKey};
use ffi::{self, CPtr};
use secp256k1_sys::types::{c_int, c_uchar, c_void};
use Error;

/// A tag used for recovering the public key from a compact signature
#[derive(Copy, Clone)]
pub struct SharedSecret {
    data: [u8; 256],
    len: usize,
}
impl_raw_debug!(SharedSecret);

// This implementes `From<N>` for all `[u8; N]` arrays from 128bits(16 byte) to 2048bits allowing known hash lengths.
// Lower than 128 bits isn't resistant to collisions any more.
impl_from_array_len!(SharedSecret, 256, (16 20 28 32 48 64 96 128 256));

impl SharedSecret {

    /// Create an empty SharedSecret
    pub(crate) fn empty() ->  SharedSecret {
        SharedSecret {
            data: [0u8; 256],
            len: 0,
        }
    }

    /// Get a pointer to the underlying data with the specified capacity.
    pub(crate) fn get_data_mut_ptr(&mut self) -> *mut u8 {
        self.data.as_mut_ptr()
    }

    /// Get the capacity of the underlying data buffer.
    pub fn capacity(&self) -> usize {
        self.data.len()
    }

    /// Get the len of the used data.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Set the length of the object.
    pub(crate) fn set_len(&mut self, len: usize) {
        debug_assert!(len <= self.data.len());
        self.len = len;
    }
}

impl PartialEq for SharedSecret {
    fn eq(&self, other: &SharedSecret) -> bool {
        self.as_ref() == other.as_ref()
    }
}

impl AsRef<[u8]> for SharedSecret {
    fn as_ref(&self) -> &[u8] {
        &self.data[..self.len]
    }
}

impl Deref for SharedSecret {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        &self.data[..self.len]
    }
}

#[cfg(feature = "std")]
unsafe extern "C" fn hash_callback<F>(output: *mut c_uchar, x: *const c_uchar, y: *const c_uchar, data: *mut c_void) -> c_int
    where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {

    use std::panic::catch_unwind;
    let res = catch_unwind(|| {
        let callback: &mut F = &mut *(data as *mut F);

        let mut x_arr = [0; 32];
        let mut y_arr = [0; 32];
        ptr::copy_nonoverlapping(x, x_arr.as_mut_ptr(), 32);
        ptr::copy_nonoverlapping(y, y_arr.as_mut_ptr(), 32);

        let secret = callback(x_arr, y_arr);
        ptr::copy_nonoverlapping(secret.as_ptr(), output as *mut u8, secret.len());

        secret.len() as c_int
    });
    if let Ok(len) = res {
        len
    } else {
        -1
    }
}

impl SharedSecret {
    /// Creates a new shared secret from a pubkey and secret key
    #[inline]
    pub fn new(point: &PublicKey, scalar: &SecretKey) -> SharedSecret {
        let mut ss = SharedSecret::empty();
        let res = unsafe {
             ffi::secp256k1_ecdh(
                ffi::secp256k1_context_no_precomp,
                ss.get_data_mut_ptr(),
                point.as_c_ptr(),
                scalar.as_c_ptr(),
                ffi::secp256k1_ecdh_hash_function_default,
                ptr::null_mut(),
            )
        };
        debug_assert_eq!(res, 1); // The default `secp256k1_ecdh_hash_function_default` should always return 1.
        ss.set_len(32); // The default hash function is SHA256, which is 32 bytes long.
        ss
    }

    /// Creates a new shared secret from a pubkey and secret key with applied custom hash function
    /// # Examples
    /// ```
    /// # use secp256k1::ecdh::SharedSecret;
    /// # use secp256k1::{Secp256k1, PublicKey, SecretKey};
    /// # fn sha2(_a: &[u8], _b: &[u8]) -> [u8; 32] {[0u8; 32]}
    /// # let secp = Secp256k1::signing_only();
    /// # let secret_key = SecretKey::from_slice(&[3u8; 32]).unwrap();
    /// # let secret_key2 = SecretKey::from_slice(&[7u8; 32]).unwrap();
    /// # let public_key = PublicKey::from_secret_key(&secp, &secret_key2);
    ///
    /// let secret = SharedSecret::new_with_hash(&public_key, &secret_key, |x,y| {
    ///     let hash: [u8; 32] = sha2(&x,&y);
    ///     hash.into()
    /// });
    ///
    /// ```
    #[cfg(feature = "std")]
    pub fn new_with_hash<F>(point: &PublicKey, scalar: &SecretKey, mut hash_function: F) -> Result<SharedSecret, Error>
        where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret
    {
        let mut ss = SharedSecret::empty();
        let hashfp: ffi::EcdhHashFn = hash_callback::<F>;

        let res =  unsafe {
            ffi::secp256k1_ecdh(
                ffi::secp256k1_context_no_precomp,
                ss.get_data_mut_ptr(),
                point.as_ptr(),
                scalar.as_ptr(),
                hashfp,
                &mut hash_function as *mut F as *mut c_void,
            )
        };
        if res == -1 {
            return Err(Error::CallbackPanicked);
        }
        debug_assert!(res >= 16); // 128 bit is the minimum for a secure hash function and the minimum we let users.
        ss.set_len(res as usize);
        Ok(ss)
    }
}

#[cfg(test)]
mod tests {
    use rand::thread_rng;
    use super::SharedSecret;
    use super::super::Secp256k1;
    use Error;

    #[test]
    fn ecdh() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let (sk2, pk2) = s.generate_keypair(&mut thread_rng());

        let sec1 = SharedSecret::new(&pk1, &sk2);
        let sec2 = SharedSecret::new(&pk2, &sk1);
        let sec_odd = SharedSecret::new(&pk1, &sk1);
        assert_eq!(sec1, sec2);
        assert!(sec_odd != sec2);
    }

    #[test]
    fn ecdh_with_hash() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let (sk2, pk2) = s.generate_keypair(&mut thread_rng());

        let sec1 = SharedSecret::new_with_hash(&pk1, &sk2, |x,_| x.into()).unwrap();
        let sec2 = SharedSecret::new_with_hash(&pk2, &sk1, |x,_| x.into()).unwrap();
        let sec_odd = SharedSecret::new_with_hash(&pk1, &sk1, |x,_| x.into()).unwrap();
        assert_eq!(sec1, sec2);
        assert_ne!(sec_odd, sec2);
    }

    #[test]
    fn ecdh_with_hash_callback() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let expect_result: [u8; 64] = [123; 64];
        let mut x_out = [0u8; 32];
        let mut y_out = [0u8; 32];
        let result = SharedSecret::new_with_hash(&pk1, &sk1, | x, y | {
            x_out = x;
            y_out = y;
            expect_result.into()
        }).unwrap();
        assert_eq!(&expect_result[..], &result[..]);
        assert_ne!(x_out, [0u8; 32]);
        assert_ne!(y_out, [0u8; 32]);
    }

    #[test]
    fn ecdh_with_hash_callback_panic() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let mut res = [0u8; 48];
        let result = SharedSecret::new_with_hash(&pk1, &sk1, | x, _ | {
            res.copy_from_slice(&x); // res.len() != x.len(). this will panic.
            res.into()
        });
        assert_eq!(result, Err(Error::CallbackPanicked));
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use rand::thread_rng;
    use test::{Bencher, black_box};

    use super::SharedSecret;
    use super::super::Secp256k1;

    #[bench]
    pub fn bench_ecdh(bh: &mut Bencher) {
        let s = Secp256k1::signing_only();
        let (sk, pk) = s.generate_keypair(&mut thread_rng());

        let s = Secp256k1::new();
        bh.iter( || {
            let res = SharedSecret::new(&pk, &sk);
            black_box(res);
        });
    }
}

<commit_msg>Add an unsafe variant of new_with_has called new_with_hash_no_panic<commit_after>// Bitcoin secp256k1 bindings
// Written in 2015 by
//   Andrew Poelstra
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # ECDH
//! Support for shared secret computations
//!

use core::ptr;
use core::ops::{FnMut, Deref};

use key::{SecretKey, PublicKey};
use ffi::{self, CPtr};
use secp256k1_sys::types::{c_int, c_uchar, c_void};
use Error;

/// A tag used for recovering the public key from a compact signature
#[derive(Copy, Clone)]
pub struct SharedSecret {
    data: [u8; 256],
    len: usize,
}
impl_raw_debug!(SharedSecret);

// This implementes `From<N>` for all `[u8; N]` arrays from 128bits(16 byte) to 2048bits allowing known hash lengths.
// Lower than 128 bits isn't resistant to collisions any more.
impl_from_array_len!(SharedSecret, 256, (16 20 28 32 48 64 96 128 256));

impl SharedSecret {

    /// Create an empty SharedSecret
    pub(crate) fn empty() ->  SharedSecret {
        SharedSecret {
            data: [0u8; 256],
            len: 0,
        }
    }

    /// Get a pointer to the underlying data with the specified capacity.
    pub(crate) fn get_data_mut_ptr(&mut self) -> *mut u8 {
        self.data.as_mut_ptr()
    }

    /// Get the capacity of the underlying data buffer.
    pub fn capacity(&self) -> usize {
        self.data.len()
    }

    /// Get the len of the used data.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Set the length of the object.
    pub(crate) fn set_len(&mut self, len: usize) {
        debug_assert!(len <= self.data.len());
        self.len = len;
    }
}

impl PartialEq for SharedSecret {
    fn eq(&self, other: &SharedSecret) -> bool {
        self.as_ref() == other.as_ref()
    }
}

impl AsRef<[u8]> for SharedSecret {
    fn as_ref(&self) -> &[u8] {
        &self.data[..self.len]
    }
}

impl Deref for SharedSecret {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        &self.data[..self.len]
    }
}

unsafe fn callback_logic<F>(output: *mut c_uchar, x: *const c_uchar, y: *const c_uchar, data: *mut c_void) -> c_int
    where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {
    let callback: &mut F = &mut *(data as *mut F);

    let mut x_arr = [0; 32];
    let mut y_arr = [0; 32];
    ptr::copy_nonoverlapping(x, x_arr.as_mut_ptr(), 32);
    ptr::copy_nonoverlapping(y, y_arr.as_mut_ptr(), 32);

    let secret = callback(x_arr, y_arr);
    ptr::copy_nonoverlapping(secret.as_ptr(), output as *mut u8, secret.len());

    secret.len() as c_int
}

#[cfg(feature = "std")]
unsafe extern "C" fn hash_callback_catch_unwind<F>(output: *mut c_uchar, x: *const c_uchar, y: *const c_uchar, data: *mut c_void) -> c_int
    where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {

    let res = ::std::panic::catch_unwind(||callback_logic::<F>(output, x, y, data));
    if let Ok(len) = res {
        len
    } else {
        -1
    }
}

unsafe extern "C" fn hash_callback_unsafe<F>(output: *mut c_uchar, x: *const c_uchar, y: *const c_uchar, data: *mut c_void) -> c_int
    where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {
    callback_logic::<F>(output, x, y, data)
}

impl SharedSecret {
    /// Creates a new shared secret from a pubkey and secret key
    #[inline]
    pub fn new(point: &PublicKey, scalar: &SecretKey) -> SharedSecret {
        let mut ss = SharedSecret::empty();
        let res = unsafe {
             ffi::secp256k1_ecdh(
                ffi::secp256k1_context_no_precomp,
                ss.get_data_mut_ptr(),
                point.as_c_ptr(),
                scalar.as_c_ptr(),
                ffi::secp256k1_ecdh_hash_function_default,
                ptr::null_mut(),
            )
        };
        debug_assert_eq!(res, 1); // The default `secp256k1_ecdh_hash_function_default` should always return 1.
        ss.set_len(32); // The default hash function is SHA256, which is 32 bytes long.
        ss
    }

    fn new_with_callback_internal<F>(point: &PublicKey, scalar: &SecretKey, mut closure: F, callback: ffi::EcdhHashFn) -> Result<SharedSecret, Error>
        where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {
        let mut ss = SharedSecret::empty();

        let res =  unsafe {
            ffi::secp256k1_ecdh(
                ffi::secp256k1_context_no_precomp,
                ss.get_data_mut_ptr(),
                point.as_ptr(),
                scalar.as_ptr(),
                callback,
                &mut closure as *mut F as *mut c_void,
            )
        };
        if res == -1 {
            return Err(Error::CallbackPanicked);
        }
        debug_assert!(res >= 16); // 128 bit is the minimum for a secure hash function and the minimum we let users.
        ss.set_len(res as usize);
        Ok(ss)

    }

    /// Creates a new shared secret from a pubkey and secret key with applied custom hash function
    /// # Examples
    /// ```
    /// # use secp256k1::ecdh::SharedSecret;
    /// # use secp256k1::{Secp256k1, PublicKey, SecretKey};
    /// # fn sha2(_a: &[u8], _b: &[u8]) -> [u8; 32] {[0u8; 32]}
    /// # let secp = Secp256k1::signing_only();
    /// # let secret_key = SecretKey::from_slice(&[3u8; 32]).unwrap();
    /// # let secret_key2 = SecretKey::from_slice(&[7u8; 32]).unwrap();
    /// # let public_key = PublicKey::from_secret_key(&secp, &secret_key2);
    ///
    /// let secret = SharedSecret::new_with_hash(&public_key, &secret_key, |x,y| {
    ///     let hash: [u8; 32] = sha2(&x,&y);
    ///     hash.into()
    /// });
    ///
    /// ```
    #[cfg(feature = "std")]
    pub fn new_with_hash<F>(point: &PublicKey, scalar: &SecretKey, hash_function: F) -> Result<SharedSecret, Error>
        where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {
            Self::new_with_callback_internal(point, scalar, hash_function, hash_callback_catch_unwind::<F>)
    }

    /// Creates a new shared secret from a pubkey and secret key with applied custom hash function
    /// Note that this function is the same as [`new_with_hash`]
    ///
    /// # Safety
    /// The function doesn't wrap the callback with [`catch_unwind`]
    /// so if the callback panics it will panic through an FFI boundray which is [`Undefined Behavior`]
    /// If possible you should use [`new_with_hash`] which does wrap the callback with [`catch_unwind`] so is safe to use.
    ///
    /// [`catch_unwind`]: https://doc.rust-lang.org/std/panic/fn.catch_unwind.html
    /// [`Undefined Behavior`]: https://doc.rust-lang.org/nomicon/ffi.html#ffi-and-panics
    /// [`new_with_hash`]: #method.new_with_hash
    /// # Examples
    /// ```
    /// # use secp256k1::ecdh::SharedSecret;
    /// # use secp256k1::{Secp256k1, PublicKey, SecretKey};
    /// # fn sha2(_a: &[u8], _b: &[u8]) -> [u8; 32] {[0u8; 32]}
    /// # let secp = Secp256k1::signing_only();
    /// # let secret_key = SecretKey::from_slice(&[3u8; 32]).unwrap();
    /// # let secret_key2 = SecretKey::from_slice(&[7u8; 32]).unwrap();
    /// # let public_key = PublicKey::from_secret_key(&secp, &secret_key2);
    //
    /// let secret = unsafe { SharedSecret::new_with_hash_no_panic(&public_key, &secret_key, |x,y| {
    ///     let hash: [u8; 32] = sha2(&x,&y);
    ///     hash.into()
    /// })};
    ///
    ///
    /// ```
    pub unsafe fn new_with_hash_no_panic<F>(point: &PublicKey, scalar: &SecretKey, hash_function: F) -> Result<SharedSecret, Error>
        where F: FnMut([u8; 32], [u8; 32]) -> SharedSecret {
        Self::new_with_callback_internal(point, scalar, hash_function, hash_callback_unsafe::<F>)
    }
}

#[cfg(test)]
mod tests {
    use rand::thread_rng;
    use super::SharedSecret;
    use super::super::Secp256k1;
    use Error;

    #[test]
    fn ecdh() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let (sk2, pk2) = s.generate_keypair(&mut thread_rng());

        let sec1 = SharedSecret::new(&pk1, &sk2);
        let sec2 = SharedSecret::new(&pk2, &sk1);
        let sec_odd = SharedSecret::new(&pk1, &sk1);
        assert_eq!(sec1, sec2);
        assert!(sec_odd != sec2);
    }

    #[test]
    fn ecdh_with_hash() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let (sk2, pk2) = s.generate_keypair(&mut thread_rng());

        let sec1 = SharedSecret::new_with_hash(&pk1, &sk2, |x,_| x.into()).unwrap();
        let sec2 = SharedSecret::new_with_hash(&pk2, &sk1, |x,_| x.into()).unwrap();
        let sec_odd = SharedSecret::new_with_hash(&pk1, &sk1, |x,_| x.into()).unwrap();
        assert_eq!(sec1, sec2);
        assert_ne!(sec_odd, sec2);
    }

    #[test]
    fn ecdh_with_hash_callback() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let expect_result: [u8; 64] = [123; 64];
        let mut x_out = [0u8; 32];
        let mut y_out = [0u8; 32];
        let result = SharedSecret::new_with_hash(&pk1, &sk1, | x, y | {
            x_out = x;
            y_out = y;
            expect_result.into()
        }).unwrap();
        let result_unsafe = unsafe {SharedSecret::new_with_hash_no_panic(&pk1, &sk1, | x, y | {
            x_out = x;
            y_out = y;
            expect_result.into()
        }).unwrap()};
        assert_eq!(&expect_result[..], &result[..]);
        assert_eq!(result, result_unsafe);
        assert_ne!(x_out, [0u8; 32]);
        assert_ne!(y_out, [0u8; 32]);
    }

    #[test]
    fn ecdh_with_hash_callback_panic() {
        let s = Secp256k1::signing_only();
        let (sk1, pk1) = s.generate_keypair(&mut thread_rng());
        let mut res = [0u8; 48];
        let result = SharedSecret::new_with_hash(&pk1, &sk1, | x, _ | {
            res.copy_from_slice(&x); // res.len() != x.len(). this will panic.
            res.into()
        });
        assert_eq!(result, Err(Error::CallbackPanicked));
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use rand::thread_rng;
    use test::{Bencher, black_box};

    use super::SharedSecret;
    use super::super::Secp256k1;

    #[bench]
    pub fn bench_ecdh(bh: &mut Bencher) {
        let s = Secp256k1::signing_only();
        let (sk, pk) = s.generate_keypair(&mut thread_rng());

        let s = Secp256k1::new();
        bh.iter( || {
            let res = SharedSecret::new(&pk, &sk);
            black_box(res);
        });
    }
}

<|endoftext|>