<commit_before>// 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 http://mozilla.org/MPL/2.0/.

// Discover or identify devices that may belong to Stratis.

use std::{
    collections::HashMap,
    fs::OpenOptions,
    path::{Path, PathBuf},
};

use libudev;

use devicemapper::Device;

use crate::engine::{
    strat_engine::{
        backstore::metadata::device_identifiers,
        udev::{block_enumerator, decide_ownership, is_multipath_member, UdevOwnership},
    },
    types::PoolUuid,
};

// A wrapper for obtaining the device number as a devicemapper Device
// which interprets absence of the value as an error, which it is in this
// context.
fn device_to_devno_wrapper(device: &libudev::Device) -> Result<Device, String> {
    device
        .devnum()
        .ok_or_else(|| "udev entry did not contain a device number".into())
        .map(Device::from)
}

// A wrapper around the metadata module's device_identifers method
// which also handles failure to open a device for reading.
// Returns an error if the device could not be opened for reading.
// Returns Ok(Err(...)) if there was an error while reading the
// Stratis identifiers from the device.
// Returns Ok(Ok(None)) if the identifers did not appear to be on
// the device.
fn device_identifiers_wrapper(devnode: &Path) -> Result<Result<Option<PoolUuid>, String>, String> {
    OpenOptions::new()
        .read(true)
        .open(devnode)
        .as_mut()
        .map_err(|err| {
            format!(
                "device {} could not be opened for reading: {}",
                devnode.display(),
                err
            )
        })
        .map(|f| {
            device_identifiers(f)
                .map_err(|err| {
                    format!(
                        "encountered an error while reading Stratis header for device {}: {}",
                        devnode.display(),
                        err
                    )
                })
                .map(|maybe_ids| maybe_ids.map(|(pool_uuid, _)| pool_uuid))
        })
}

// Use udev to identify all block devices and return the subset of those
// that have Stratis signatures.
fn find_all_block_devices_with_stratis_signatures(
) -> libudev::Result<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    let context = libudev::Context::new()?;
    let mut enumerator = block_enumerator(&context)?;

    let pool_map = enumerator.scan_devices()?
        .filter(|dev| {
            let initialized = dev.is_initialized();
            if !initialized {
                debug!("Found a udev entry for a device identified as a block device, but udev also identified it as uninitialized, omitting the device from the set of devices to process, for safety");
            };
            initialized
        })
        .filter(|dev| {
            decide_ownership(dev)
                .map_err(|err| {
                    warn!("Could not determine ownership of a udev block device because of an error processing udev information, omitting the device from the set of devices to process, for safety: {}",
                          err);
                })
                .map(|decision| match decision {
                    UdevOwnership::Stratis | UdevOwnership::Unowned => true,
                    _ => false,
                })
                .unwrap_or(false)
        })
        .filter_map(|dev| match dev.devnode() {
            Some(devnode) => {
                match (device_to_devno_wrapper(&dev), device_identifiers_wrapper(devnode)) {
                    (Err(err), _) | (_, Err(err)) => {
                        warn!("udev identified device {} as a block device but {}, omitting the device from the set of devices to process",
                              devnode.display(),
                              err);
                        None
                    }
                    // FIXME: Refine error return in StaticHeader::setup(),
                    // so it can be used to distinguish between signficant
                    // and insignficant errors and then use that ability to
                    // distinguish here between different levels of
                    // severity.
                    (_, Ok(Err(err))) => {
                        debug!("udev identified device {} as a block device but {}, omitting the device from the set of devices to process",
                               devnode.display(),
                               err);
                        None
                    }
                    (_, Ok(Ok(None))) => None,
                    (Ok(devno), Ok(Ok(Some(pool_uuid)))) => Some((pool_uuid, devno, devnode.to_path_buf())),
                }
            }
            None => {
                warn!("udev identified a device as a block device, but the udev entry for the device had no device node, omitting the device from the set of devices to process");
                None
            }
        })
        .fold(HashMap::new(), |mut acc, (pool_uuid, device, devnode)| {
            acc.entry(pool_uuid).or_insert_with(HashMap::new).insert(device, devnode);
            acc
        });

    Ok(pool_map)
}

// Find all devices identified by udev as Stratis devices.
fn find_all_stratis_devices() -> libudev::Result<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    let context = libudev::Context::new()?;
    let mut enumerator = block_enumerator(&context)?;
    enumerator.match_property("ID_FS_TYPE", "stratis")?;

    let pool_map = enumerator.scan_devices()?
        .filter(|dev| {
            let initialized = dev.is_initialized();
            if !initialized {
                warn!("Found a udev entry for a device identified as a Stratis device, but udev also identified it as uninitialized, omitting the device from the set of devices to process, for safety");
            };
            initialized
        })
        .filter(|dev| !is_multipath_member(dev)
                .map_err(|err| {
                    warn!("Could not certainly determine whether a device was a multipath member because of an error processing udev information, omitting the device from the set of devices to process, for safety: {}",
                          err);
                })
                .unwrap_or(true))
        .filter_map(|dev| match dev.devnode() {
            Some(devnode) => {
                match (device_to_devno_wrapper(&dev), device_identifiers_wrapper(devnode)) {
                    (Err(err), _) | (_, Err(err)) | (_, Ok(Err(err)))=> {
                        warn!("udev identified device {} as a Stratis device but {}, omitting the device from the set of devices to process",
                              devnode.display(),
                              err);
                        None
                    }
                    (_, Ok(Ok(None))) => {
                            warn!("udev identified device {} as a Stratis device but there appeared to be no Stratis metadata on the device, omitting the device from the set of devices to process",
                                  devnode.display());
                            None
                    }
                    (Ok(devno), Ok(Ok(Some(pool_uuid)))) => Some((pool_uuid, devno, devnode.to_path_buf())),
                }
            }
            None => {
                warn!("udev identified a device as a Stratis device, but the udev entry for the device had no device node, omitting the the device from the set of devices to process");
                None
            }
        })
        .fold(HashMap::new(), |mut acc, (pool_uuid, device, devnode)| {
            acc.entry(pool_uuid).or_insert_with(HashMap::new).insert(device, devnode);
            acc
        });
    Ok(pool_map)
}

/// Retrieve all block devices that should be made use of by the
/// Stratis engine. This excludes Stratis block devices that appear to be
/// multipath members.
///
/// Includes a fallback path, which is used if no Stratis block devices are
/// found using the obvious udev property- and enumerator-based approach.
/// This fallback path is more expensive, because it must search all block
/// devices via udev rather than just all Stratis block devices.
///
/// Omits any device that appears problematic in some way.
///
/// Return an error only on a failure to construct or scan with a udev
/// enumerator.
///
/// Returns a map of pool uuids to a map of devices to devnodes for each pool.
pub fn find_all() -> libudev::Result<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    info!("Beginning initial search for Stratis block devices");
    let pool_map = find_all_stratis_devices()?;

    if pool_map.is_empty() {
        // No Stratis devices have been found, possible reasons are:
        // 1. There are none
        // 2. There are some but libblkid version is less than 2.32, so
        // Stratis devices are not recognized by udev.
        // 3. There are many incomplete udev entries, because this code is
        // being run before the udev database is populated.
        //
        // Try again to find Stratis block devices, but this time enumerate
        // all block devices, not just all the ones that can be identified
        // as Stratis blockdevs by udev, and process only those that appear
        // unclaimed or appear to be claimed by Stratis (and not
        // multipath members).

        info!("Could not identify any Stratis devices by a udev search for devices with ID_FS_TYPE=\"stratis\"; using fallback search mechanism");

        find_all_block_devices_with_stratis_signatures()
    } else {
        Ok(pool_map)
    }
}
<commit_msg>Have device_identifiers_wrapper return both device identifiers<commit_after>// 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 http://mozilla.org/MPL/2.0/.

// Discover or identify devices that may belong to Stratis.

use std::{
    collections::HashMap,
    fs::OpenOptions,
    path::{Path, PathBuf},
};

use libudev;

use devicemapper::Device;

use crate::engine::{
    strat_engine::{
        backstore::metadata::device_identifiers,
        udev::{block_enumerator, decide_ownership, is_multipath_member, UdevOwnership},
    },
    types::{DevUuid, PoolUuid},
};

// A wrapper for obtaining the device number as a devicemapper Device
// which interprets absence of the value as an error, which it is in this
// context.
fn device_to_devno_wrapper(device: &libudev::Device) -> Result<Device, String> {
    device
        .devnum()
        .ok_or_else(|| "udev entry did not contain a device number".into())
        .map(Device::from)
}

// A wrapper around the metadata module's device_identifers method
// which also handles failure to open a device for reading.
// Returns an error if the device could not be opened for reading.
// Returns Ok(Err(...)) if there was an error while reading the
// Stratis identifiers from the device.
// Returns Ok(Ok(None)) if the identifers did not appear to be on
// the device.
fn device_identifiers_wrapper(
    devnode: &Path,
) -> Result<Result<Option<(PoolUuid, DevUuid)>, String>, String> {
    OpenOptions::new()
        .read(true)
        .open(devnode)
        .as_mut()
        .map_err(|err| {
            format!(
                "device {} could not be opened for reading: {}",
                devnode.display(),
                err
            )
        })
        .map(|f| {
            device_identifiers(f).map_err(|err| {
                format!(
                    "encountered an error while reading Stratis header for device {}: {}",
                    devnode.display(),
                    err
                )
            })
        })
}

// Use udev to identify all block devices and return the subset of those
// that have Stratis signatures.
fn find_all_block_devices_with_stratis_signatures(
) -> libudev::Result<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    let context = libudev::Context::new()?;
    let mut enumerator = block_enumerator(&context)?;

    let pool_map = enumerator.scan_devices()?
        .filter(|dev| {
            let initialized = dev.is_initialized();
            if !initialized {
                debug!("Found a udev entry for a device identified as a block device, but udev also identified it as uninitialized, omitting the device from the set of devices to process, for safety");
            };
            initialized
        })
        .filter(|dev| {
            decide_ownership(dev)
                .map_err(|err| {
                    warn!("Could not determine ownership of a udev block device because of an error processing udev information, omitting the device from the set of devices to process, for safety: {}",
                          err);
                })
                .map(|decision| match decision {
                    UdevOwnership::Stratis | UdevOwnership::Unowned => true,
                    _ => false,
                })
                .unwrap_or(false)
        })
        .filter_map(|dev| match dev.devnode() {
            Some(devnode) => {
                match (device_to_devno_wrapper(&dev), device_identifiers_wrapper(devnode)) {
                    (Err(err), _) | (_, Err(err)) => {
                        warn!("udev identified device {} as a block device but {}, omitting the device from the set of devices to process",
                              devnode.display(),
                              err);
                        None
                    }
                    // FIXME: Refine error return in StaticHeader::setup(),
                    // so it can be used to distinguish between signficant
                    // and insignficant errors and then use that ability to
                    // distinguish here between different levels of
                    // severity.
                    (_, Ok(Err(err))) => {
                        debug!("udev identified device {} as a block device but {}, omitting the device from the set of devices to process",
                               devnode.display(),
                               err);
                        None
                    }
                    (_, Ok(Ok(None))) => None,
                    (Ok(devno), Ok(Ok(Some((pool_uuid, _))))) => Some((pool_uuid, devno, devnode.to_path_buf())),
                }
            }
            None => {
                warn!("udev identified a device as a block device, but the udev entry for the device had no device node, omitting the device from the set of devices to process");
                None
            }
        })
        .fold(HashMap::new(), |mut acc, (pool_uuid, device, devnode)| {
            acc.entry(pool_uuid).or_insert_with(HashMap::new).insert(device, devnode);
            acc
        });

    Ok(pool_map)
}

// Find all devices identified by udev as Stratis devices.
fn find_all_stratis_devices() -> libudev::Result<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    let context = libudev::Context::new()?;
    let mut enumerator = block_enumerator(&context)?;
    enumerator.match_property("ID_FS_TYPE", "stratis")?;

    let pool_map = enumerator.scan_devices()?
        .filter(|dev| {
            let initialized = dev.is_initialized();
            if !initialized {
                warn!("Found a udev entry for a device identified as a Stratis device, but udev also identified it as uninitialized, omitting the device from the set of devices to process, for safety");
            };
            initialized
        })
        .filter(|dev| !is_multipath_member(dev)
                .map_err(|err| {
                    warn!("Could not certainly determine whether a device was a multipath member because of an error processing udev information, omitting the device from the set of devices to process, for safety: {}",
                          err);
                })
                .unwrap_or(true))
        .filter_map(|dev| match dev.devnode() {
            Some(devnode) => {
                match (device_to_devno_wrapper(&dev), device_identifiers_wrapper(devnode)) {
                    (Err(err), _) | (_, Err(err)) | (_, Ok(Err(err)))=> {
                        warn!("udev identified device {} as a Stratis device but {}, omitting the device from the set of devices to process",
                              devnode.display(),
                              err);
                        None
                    }
                    (_, Ok(Ok(None))) => {
                            warn!("udev identified device {} as a Stratis device but there appeared to be no Stratis metadata on the device, omitting the device from the set of devices to process",
                                  devnode.display());
                            None
                    }
                    (Ok(devno), Ok(Ok(Some((pool_uuid, _))))) => Some((pool_uuid, devno, devnode.to_path_buf())),
                }
            }
            None => {
                warn!("udev identified a device as a Stratis device, but the udev entry for the device had no device node, omitting the the device from the set of devices to process");
                None
            }
        })
        .fold(HashMap::new(), |mut acc, (pool_uuid, device, devnode)| {
            acc.entry(pool_uuid).or_insert_with(HashMap::new).insert(device, devnode);
            acc
        });
    Ok(pool_map)
}

/// Retrieve all block devices that should be made use of by the
/// Stratis engine. This excludes Stratis block devices that appear to be
/// multipath members.
///
/// Includes a fallback path, which is used if no Stratis block devices are
/// found using the obvious udev property- and enumerator-based approach.
/// This fallback path is more expensive, because it must search all block
/// devices via udev rather than just all Stratis block devices.
///
/// Omits any device that appears problematic in some way.
///
/// Return an error only on a failure to construct or scan with a udev
/// enumerator.
///
/// Returns a map of pool uuids to a map of devices to devnodes for each pool.
pub fn find_all() -> libudev::Result<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    info!("Beginning initial search for Stratis block devices");
    let pool_map = find_all_stratis_devices()?;

    if pool_map.is_empty() {
        // No Stratis devices have been found, possible reasons are:
        // 1. There are none
        // 2. There are some but libblkid version is less than 2.32, so
        // Stratis devices are not recognized by udev.
        // 3. There are many incomplete udev entries, because this code is
        // being run before the udev database is populated.
        //
        // Try again to find Stratis block devices, but this time enumerate
        // all block devices, not just all the ones that can be identified
        // as Stratis blockdevs by udev, and process only those that appear
        // unclaimed or appear to be claimed by Stratis (and not
        // multipath members).

        info!("Could not identify any Stratis devices by a udev search for devices with ID_FS_TYPE=\"stratis\"; using fallback search mechanism");

        find_all_block_devices_with_stratis_signatures()
    } else {
        Ok(pool_map)
    }
}
<|endoftext|>


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

extern crate word_count;

fn check_word_count(s: &str, pairs: Vec<(&str, u32)>) {
    // The reason for the awkward code in here is to ensure that the failure
    // message for assert_eq! is as informative as possible. A simpler
    // solution would simply check the length of the map, and then
    // check for the presence and value of each key in the given pairs vector.
    let mut m: HashMap<String, u32> = word_count::word_count(s);
    for &(k, v) in pairs.iter() {
        assert_eq!((k, m.remove(&k.to_string()).unwrap()), (k, v));
    }
    // may fail with a message that clearly shows all extra pairs in the map
    assert_eq!(m.iter().collect::<Vec<(&String,&u32)>>(), vec!());
}

#[test]
fn test_no_words() {
    check_word_count("", vec![]);
    check_word_count("-!@#$ $#() @@@", vec![]);
}

#[test]
#[ignore]
fn test_count_one_word() {
    check_word_count("word", vec![("word", 1)]);
}

#[test]
#[ignore]
fn test_count_one_of_each() {
    check_word_count(
        "one of each",
        vec![("one", 1),
             ("of", 1),
             ("each", 1)]);
}

#[test]
#[ignore]
fn test_count_multiple_occurrences() {
    check_word_count(
        "one fish two fish red fish blue fish",
        vec![("one", 1),
             ("fish", 4),
             ("two", 1),
             ("red", 1),
             ("blue", 1)]);
}

#[test]
#[ignore]
fn test_ignore_punctuation() {
    check_word_count(
        "car : carpet as java : javascript!!&@$%^&",
        vec![("car", 1),
             ("carpet", 1),
             ("as", 1),
             ("java", 1),
             ("javascript", 1)]);
}

#[test]
#[ignore]
fn test_include_numbers() {
    check_word_count(
        "testing, 1, 2 testing",
        vec![("testing", 2),
             ("1", 1),
             ("2", 1)]);
}

#[test]
#[ignore]
fn test_normalize_case() {
    check_word_count(
        "go Go GO",
        vec![("go", 3)]);
}

#[test]
#[ignore]
fn test_prefix_punctuation() {
    check_word_count(
        "!%%#testing, 1, 2 testing",
        vec![("testing", 2),
             ("1", 1),
             ("2", 1)]);
}

#[test]
#[ignore]
fn test_symbols_are_separators() {
    check_word_count(
        "hey,my_spacebar_is_broken.",
        vec![("hey", 1),
             ("my", 1),
             ("spacebar", 1),
             ("is", 1),
             ("broken", 1)]);
}
<commit_msg>Update test to prevent panic with missing value<commit_after>use std::collections::HashMap;

extern crate word_count;

fn check_word_count(s: &str, pairs: Vec<(&str, u32)>) {
    // The reason for the awkward code in here is to ensure that the failure
    // message for assert_eq! is as informative as possible. A simpler
    // solution would simply check the length of the map, and then
    // check for the presence and value of each key in the given pairs vector.
    let mut m: HashMap<String, u32> = word_count::word_count(s);
    for &(k, v) in pairs.iter() {
        assert_eq!((k, m.remove(&k.to_string()).unwrap_or(0)), (k, v));
    }
    // may fail with a message that clearly shows all extra pairs in the map
    assert_eq!(m.iter().collect::<Vec<(&String,&u32)>>(), vec!());
}

#[test]
fn test_no_words() {
    check_word_count("", vec![]);
    check_word_count("-!@#$ $#() @@@", vec![]);
}

#[test]
#[ignore]
fn test_count_one_word() {
    check_word_count("word", vec![("word", 1)]);
}

#[test]
#[ignore]
fn test_count_one_of_each() {
    check_word_count(
        "one of each",
        vec![("one", 1),
             ("of", 1),
             ("each", 1)]);
}

#[test]
#[ignore]
fn test_count_multiple_occurrences() {
    check_word_count(
        "one fish two fish red fish blue fish",
        vec![("one", 1),
             ("fish", 4),
             ("two", 1),
             ("red", 1),
             ("blue", 1)]);
}

#[test]
#[ignore]
fn test_ignore_punctuation() {
    check_word_count(
        "car : carpet as java : javascript!!&@$%^&",
        vec![("car", 1),
             ("carpet", 1),
             ("as", 1),
             ("java", 1),
             ("javascript", 1)]);
}

#[test]
#[ignore]
fn test_include_numbers() {
    check_word_count(
        "testing, 1, 2 testing",
        vec![("testing", 2),
             ("1", 1),
             ("2", 1)]);
}

#[test]
#[ignore]
fn test_normalize_case() {
    check_word_count(
        "go Go GO",
        vec![("go", 3)]);
}

#[test]
#[ignore]
fn test_prefix_punctuation() {
    check_word_count(
        "!%%#testing, 1, 2 testing",
        vec![("testing", 2),
             ("1", 1),
             ("2", 1)]);
}

#[test]
#[ignore]
fn test_symbols_are_separators() {
    check_word_count(
        "hey,my_spacebar_is_broken.",
        vec![("hey", 1),
             ("my", 1),
             ("spacebar", 1),
             ("is", 1),
             ("broken", 1)]);
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::ast::{Ast, AstBuilder, AstKey, ObjectEntry, Primitive, RequestParameters};
use crate::build_ast::{
    build_fragment, build_operation, build_request, build_request_params,
    build_request_params_ast_key,
};
use crate::constants::CODEGEN_CONSTANTS;
use crate::indentation::print_indentation;
use crate::utils::escape;

use graphql_ir::{FragmentDefinition, OperationDefinition};
use schema::Schema;

use fnv::{FnvBuildHasher, FnvHashSet};
use indexmap::IndexMap;
use interner::StringKey;
use std::fmt::{Result as FmtResult, Write};

pub struct Printer {
    builder: AstBuilder,
}

type VariableDefinitions = IndexMap<AstKey, String, FnvBuildHasher>;
struct DedupedJSONPrinter<'b> {
    variable_definitions: VariableDefinitions,
    duplicates: FnvHashSet<AstKey>,
    builder: &'b AstBuilder,
    root_key: AstKey,
}

impl<'b> DedupedJSONPrinter<'b> {
    fn new(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        let mut visited = Default::default();
        let mut duplicates = Default::default();
        Self::collect_value_duplicates(&mut visited, &mut duplicates, builder, root_key);
        Self {
            variable_definitions: Default::default(),
            duplicates,
            builder,
            root_key,
        }
    }

    fn new_without_dedupe(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        Self {
            variable_definitions: Default::default(),
            duplicates: Default::default(),
            builder,
            root_key,
        }
    }

    /// We don't dedupe in an already deduped AST unless the duplicate
    /// also appears on a subtree that's not a duplicate
    /// Input:
    /// [
    ///     [{a: 1}, {b:2}],
    ///     [{a: 1}, {b:2}],
    ///     {b: 2}
    /// ]
    /// Output:
    /// v0 = {b: 2};
    /// v1 = [{a: 1}, v0];
    fn collect_value_duplicates(
        visited: &mut FnvHashSet<AstKey>,
        duplicates: &mut FnvHashSet<AstKey>,
        builder: &AstBuilder,
        key: AstKey,
    ) {
        match builder.lookup(key) {
            Ast::Array(array) => {
                if array.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for val in array {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
            Ast::Object(object) => {
                if object.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for (_, val) in object {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
        }
    }

    fn print(mut self) -> String {
        let mut result = String::new();
        self.print_ast(&mut result, self.root_key, 0, false);
        if self.variable_definitions.is_empty() {
            result
        } else {
            let mut with_variables = String::new();
            with_variables.push_str("(function(){\nvar ");
            let last = self.variable_definitions.len() - 1;
            for (i, (_, value)) in self.variable_definitions.drain(..).enumerate() {
                writeln!(
                    &mut with_variables,
                    "v{} = {}{}",
                    i,
                    value,
                    if i == last { ";" } else { "," }
                )
                .unwrap();
            }
            write!(&mut with_variables, "return {};\n}})()", result).unwrap();
            with_variables
        }
    }

    fn print_ast(&mut self, f: &mut String, key: AstKey, indent: usize, is_dedupe_var: bool) {
        // Only use variable references at depth beyond the top level.
        if indent > 0 && self.duplicates.contains(&key) {
            let v = if self.variable_definitions.contains_key(&key) {
                self.variable_definitions.get_full(&key).unwrap().0
            } else {
                let mut variable = String::new();
                self.print_ast(&mut variable, key, 0, true);
                let v = self.variable_definitions.len();
                self.variable_definitions.insert(key, variable);
                v
            };
            return write!(f, "(v{}/*: any*/)", v).unwrap();
        }

        let ast = self.builder.lookup(key);
        match ast {
            Ast::Object(object) => {
                if object.is_empty() {
                    f.push_str("{}");
                } else {
                    let next_indent = indent + 1;
                    f.push('{');
                    for (key, value) in object {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        write!(f, "\"{}\": ", key.lookup()).unwrap();
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push('}');
                }
            }
            Ast::Array(array) => {
                if array.is_empty() {
                    if is_dedupe_var {
                        // Empty arrays can only have one inferred flow type and then conflict if
                        // used in different places, this is unsound if we would write to them but
                        // this whole module is based on the idea of a read only JSON tree.
                        f.push_str("([]/*: any*/)");
                    } else {
                        f.push_str("[]");
                    }
                } else {
                    f.push('[');
                    let next_indent = indent + 1;
                    for value in array {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push(']');
                }
            }
        }
    }

    fn print_primitive(
        &mut self,
        f: &mut String,
        primitive: &Primitive,
        indent: usize,
        is_dedupe_var: bool,
    ) -> FmtResult {
        match primitive {
            Primitive::Null => write!(f, "null"),
            Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
            Primitive::RawString(str) => {
                f.push('\"');
                escape(str, f);
                f.push('\"');
                Ok(())
            }
            Primitive::String(key) => write!(f, "\"{}\"", key),
            Primitive::Float(value) => write!(f, "{}", value.as_float()),
            Primitive::Int(value) => write!(f, "{}", value),
            Primitive::Key(key) => {
                self.print_ast(f, *key, indent, is_dedupe_var);
                Ok(())
            }
            Primitive::StorageKey(field_name, key) => {
                print_static_storage_key(f, &self.builder, *field_name, *key)
            }
            Primitive::ModuleDependency(key) => write!(f, "require('{}.graphql')", key),
        }
    }
}

impl Default for Printer {
    fn default() -> Self {
        Self {
            builder: AstBuilder::default(),
        }
    }
}

pub fn print_operation(schema: &Schema, operation: &OperationDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_operation(schema, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_fragment(schema: &Schema, fragment: &FragmentDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_fragment(schema, &mut builder, fragment);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request(
    schema: &Schema,
    operation: &OperationDefinition,
    fragment: &FragmentDefinition,
    request_parameters: RequestParameters,
) -> String {
    let mut builder = AstBuilder::default();
    let request_parameters =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let key = build_request(
        schema,
        &mut builder,
        operation,
        fragment,
        request_parameters,
    );
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request_params(
    schema: &Schema,
    operation: &OperationDefinition,
    query_id: Option<String>,
) -> String {
    let mut request_parameters = build_request_params(operation);
    request_parameters.id = query_id;

    let mut builder = AstBuilder::default();
    let request_parameters_ast_key =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, request_parameters_ast_key);

    printer.print()
}

impl Printer {
    pub fn print_request_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
        fragment: &FragmentDefinition,
        request_parameters: RequestParameters,
    ) -> String {
        let request_parameters =
            build_request_params_ast_key(schema, request_parameters, &mut self.builder, operation);
        let key = build_request(
            schema,
            &mut self.builder,
            operation,
            fragment,
            request_parameters,
        );
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_operation_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
    ) -> String {
        let key = build_operation(schema, &mut self.builder, operation);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_fragment_deduped(
        &mut self,
        schema: &Schema,
        fragment: &FragmentDefinition,
    ) -> String {
        let key = build_fragment(schema, &mut self.builder, fragment);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }
}

/// Pre-computes storage key if possible and advantageous. Storage keys are
/// generated for fields with supplied arguments that are all statically known
/// (ie. literals, no variables) at build time.
fn print_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args_key: AstKey,
) -> FmtResult {
    // TODO (T64585375): JS compiler has an option to force a storageKey.
    let args = builder.lookup(args_key).assert_array();
    write_static_storage_key(f, builder, field_name, args)
}

fn write_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args: &[Primitive],
) -> FmtResult {
    write!(f, "\"{}(", field_name)?;
    for arg_key in args {
        let arg = builder.lookup(arg_key.assert_key()).assert_object();
        let (_, name) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
            .expect("Expected `name` to exist");
        let name = name.assert_string();
        write!(f, "{}:", name)?;
        write_argument_value(f, builder, arg)?;
        f.push(',');
    }
    f.pop(); // args won't be empty
    f.push_str(")\"");
    Ok(())
}

fn write_argument_value(f: &mut String, builder: &AstBuilder, arg: &[ObjectEntry]) -> FmtResult {
    let (_, key) = arg
        .iter()
        .find(|(key, _)| *key == CODEGEN_CONSTANTS.kind)
        .expect("Expected `kind` to exist");
    let key = key.assert_string();
    // match doesn't allow `CODEGEN_CONSTANTS.<>` on the match arm, falling back to if statements
    if key == CODEGEN_CONSTANTS.literal {
        let (_, literal) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.value)
            .expect("Expected `name` to exist");
        write_constant_value(f, builder, literal)?;
    } else if key == CODEGEN_CONSTANTS.list_value {
        let (_, items) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.items)
            .expect("Expected `items` to exist");
        let array = builder.lookup(items.assert_key()).assert_array();

        f.push('[');
        for key_or_null in array {
            match key_or_null {
                Primitive::Null => f.push_str("null,"),
                Primitive::Key(key) => {
                    let object = builder.lookup(*key).assert_object();
                    write_argument_value(f, builder, object)?;
                    f.push(',');
                }
                _ => panic!("Expected an object key or null"),
            }
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push(']');
    } else {
        // We filtered out Variables, here it should only be ObjectValue
        let (_, value) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.fields)
            .expect("Expected `fields` to exist");
        let array = builder.lookup(value.assert_key()).assert_array();

        f.push('{');
        for key in array {
            let (_, name) = arg
                .iter()
                .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
                .expect("Expected `name` to exist");
            let name = name.assert_string();
            write!(f, "\\\"{}\\\":", name)?;
            let object = builder.lookup(key.assert_key()).assert_object();
            write_argument_value(f, builder, object)?;
            f.push(',');
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push('}');
    }
    Ok(())
}

fn write_constant_value(f: &mut String, builder: &AstBuilder, value: &Primitive) -> FmtResult {
    match value {
        Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
        Primitive::String(key) => write!(f, "\\\"{}\\\"", key),
        Primitive::Float(value) => write!(f, "{}", value.as_float()),
        Primitive::Int(value) => write!(f, "{}", value),
        Primitive::Key(key) => {
            let ast = builder.lookup(*key);
            match ast {
                Ast::Array(arr) => {
                    f.push('[');
                    for value in arr {
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !arr.is_empty() {
                        f.pop();
                    }
                    f.push(']');
                    Ok(())
                }
                Ast::Object(obj) => {
                    f.push('{');
                    for (name, value) in obj {
                        write!(f, "\\\"{}\\\":", name)?;
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !obj.is_empty() {
                        f.pop();
                    }
                    f.push('}');
                    Ok(())
                }
            }
        }
        Primitive::Null => {
            f.push_str("null");
            Ok(())
        }
        Primitive::StorageKey(_, _) => panic!("Unexpected StorageKey"),
        Primitive::RawString(_) => panic!("Unexpected RawString"),
        Primitive::ModuleDependency(_) => panic!("Unexpected ModuleDependency"),
    }
}
<commit_msg>Skip null in ListValues in the storageKey<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::ast::{Ast, AstBuilder, AstKey, ObjectEntry, Primitive, RequestParameters};
use crate::build_ast::{
    build_fragment, build_operation, build_request, build_request_params,
    build_request_params_ast_key,
};
use crate::constants::CODEGEN_CONSTANTS;
use crate::indentation::print_indentation;
use crate::utils::escape;

use graphql_ir::{FragmentDefinition, OperationDefinition};
use schema::Schema;

use fnv::{FnvBuildHasher, FnvHashSet};
use indexmap::IndexMap;
use interner::StringKey;
use std::fmt::{Result as FmtResult, Write};

pub struct Printer {
    builder: AstBuilder,
}

type VariableDefinitions = IndexMap<AstKey, String, FnvBuildHasher>;
struct DedupedJSONPrinter<'b> {
    variable_definitions: VariableDefinitions,
    duplicates: FnvHashSet<AstKey>,
    builder: &'b AstBuilder,
    root_key: AstKey,
}

impl<'b> DedupedJSONPrinter<'b> {
    fn new(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        let mut visited = Default::default();
        let mut duplicates = Default::default();
        Self::collect_value_duplicates(&mut visited, &mut duplicates, builder, root_key);
        Self {
            variable_definitions: Default::default(),
            duplicates,
            builder,
            root_key,
        }
    }

    fn new_without_dedupe(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        Self {
            variable_definitions: Default::default(),
            duplicates: Default::default(),
            builder,
            root_key,
        }
    }

    /// We don't dedupe in an already deduped AST unless the duplicate
    /// also appears on a subtree that's not a duplicate
    /// Input:
    /// [
    ///     [{a: 1}, {b:2}],
    ///     [{a: 1}, {b:2}],
    ///     {b: 2}
    /// ]
    /// Output:
    /// v0 = {b: 2};
    /// v1 = [{a: 1}, v0];
    fn collect_value_duplicates(
        visited: &mut FnvHashSet<AstKey>,
        duplicates: &mut FnvHashSet<AstKey>,
        builder: &AstBuilder,
        key: AstKey,
    ) {
        match builder.lookup(key) {
            Ast::Array(array) => {
                if array.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for val in array {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
            Ast::Object(object) => {
                if object.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for (_, val) in object {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
        }
    }

    fn print(mut self) -> String {
        let mut result = String::new();
        self.print_ast(&mut result, self.root_key, 0, false);
        if self.variable_definitions.is_empty() {
            result
        } else {
            let mut with_variables = String::new();
            with_variables.push_str("(function(){\nvar ");
            let last = self.variable_definitions.len() - 1;
            for (i, (_, value)) in self.variable_definitions.drain(..).enumerate() {
                writeln!(
                    &mut with_variables,
                    "v{} = {}{}",
                    i,
                    value,
                    if i == last { ";" } else { "," }
                )
                .unwrap();
            }
            write!(&mut with_variables, "return {};\n}})()", result).unwrap();
            with_variables
        }
    }

    fn print_ast(&mut self, f: &mut String, key: AstKey, indent: usize, is_dedupe_var: bool) {
        // Only use variable references at depth beyond the top level.
        if indent > 0 && self.duplicates.contains(&key) {
            let v = if self.variable_definitions.contains_key(&key) {
                self.variable_definitions.get_full(&key).unwrap().0
            } else {
                let mut variable = String::new();
                self.print_ast(&mut variable, key, 0, true);
                let v = self.variable_definitions.len();
                self.variable_definitions.insert(key, variable);
                v
            };
            return write!(f, "(v{}/*: any*/)", v).unwrap();
        }

        let ast = self.builder.lookup(key);
        match ast {
            Ast::Object(object) => {
                if object.is_empty() {
                    f.push_str("{}");
                } else {
                    let next_indent = indent + 1;
                    f.push('{');
                    for (key, value) in object {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        write!(f, "\"{}\": ", key.lookup()).unwrap();
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push('}');
                }
            }
            Ast::Array(array) => {
                if array.is_empty() {
                    if is_dedupe_var {
                        // Empty arrays can only have one inferred flow type and then conflict if
                        // used in different places, this is unsound if we would write to them but
                        // this whole module is based on the idea of a read only JSON tree.
                        f.push_str("([]/*: any*/)");
                    } else {
                        f.push_str("[]");
                    }
                } else {
                    f.push('[');
                    let next_indent = indent + 1;
                    for value in array {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push(']');
                }
            }
        }
    }

    fn print_primitive(
        &mut self,
        f: &mut String,
        primitive: &Primitive,
        indent: usize,
        is_dedupe_var: bool,
    ) -> FmtResult {
        match primitive {
            Primitive::Null => write!(f, "null"),
            Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
            Primitive::RawString(str) => {
                f.push('\"');
                escape(str, f);
                f.push('\"');
                Ok(())
            }
            Primitive::String(key) => write!(f, "\"{}\"", key),
            Primitive::Float(value) => write!(f, "{}", value.as_float()),
            Primitive::Int(value) => write!(f, "{}", value),
            Primitive::Key(key) => {
                self.print_ast(f, *key, indent, is_dedupe_var);
                Ok(())
            }
            Primitive::StorageKey(field_name, key) => {
                print_static_storage_key(f, &self.builder, *field_name, *key)
            }
            Primitive::ModuleDependency(key) => write!(f, "require('{}.graphql')", key),
        }
    }
}

impl Default for Printer {
    fn default() -> Self {
        Self {
            builder: AstBuilder::default(),
        }
    }
}

pub fn print_operation(schema: &Schema, operation: &OperationDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_operation(schema, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_fragment(schema: &Schema, fragment: &FragmentDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_fragment(schema, &mut builder, fragment);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request(
    schema: &Schema,
    operation: &OperationDefinition,
    fragment: &FragmentDefinition,
    request_parameters: RequestParameters,
) -> String {
    let mut builder = AstBuilder::default();
    let request_parameters =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let key = build_request(
        schema,
        &mut builder,
        operation,
        fragment,
        request_parameters,
    );
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request_params(
    schema: &Schema,
    operation: &OperationDefinition,
    query_id: Option<String>,
) -> String {
    let mut request_parameters = build_request_params(operation);
    request_parameters.id = query_id;

    let mut builder = AstBuilder::default();
    let request_parameters_ast_key =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, request_parameters_ast_key);

    printer.print()
}

impl Printer {
    pub fn print_request_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
        fragment: &FragmentDefinition,
        request_parameters: RequestParameters,
    ) -> String {
        let request_parameters =
            build_request_params_ast_key(schema, request_parameters, &mut self.builder, operation);
        let key = build_request(
            schema,
            &mut self.builder,
            operation,
            fragment,
            request_parameters,
        );
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_operation_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
    ) -> String {
        let key = build_operation(schema, &mut self.builder, operation);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_fragment_deduped(
        &mut self,
        schema: &Schema,
        fragment: &FragmentDefinition,
    ) -> String {
        let key = build_fragment(schema, &mut self.builder, fragment);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }
}

/// Pre-computes storage key if possible and advantageous. Storage keys are
/// generated for fields with supplied arguments that are all statically known
/// (ie. literals, no variables) at build time.
fn print_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args_key: AstKey,
) -> FmtResult {
    // TODO (T64585375): JS compiler has an option to force a storageKey.
    let args = builder.lookup(args_key).assert_array();
    write_static_storage_key(f, builder, field_name, args)
}

fn write_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args: &[Primitive],
) -> FmtResult {
    write!(f, "\"{}(", field_name)?;
    for arg_key in args {
        let arg = builder.lookup(arg_key.assert_key()).assert_object();
        let (_, name) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
            .expect("Expected `name` to exist");
        let name = name.assert_string();
        write!(f, "{}:", name)?;
        write_argument_value(f, builder, arg)?;
        f.push(',');
    }
    f.pop(); // args won't be empty
    f.push_str(")\"");
    Ok(())
}

fn write_argument_value(f: &mut String, builder: &AstBuilder, arg: &[ObjectEntry]) -> FmtResult {
    let (_, key) = arg
        .iter()
        .find(|(key, _)| *key == CODEGEN_CONSTANTS.kind)
        .expect("Expected `kind` to exist");
    let key = key.assert_string();
    // match doesn't allow `CODEGEN_CONSTANTS.<>` on the match arm, falling back to if statements
    if key == CODEGEN_CONSTANTS.literal {
        let (_, literal) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.value)
            .expect("Expected `name` to exist");
        write_constant_value(f, builder, literal)?;
    } else if key == CODEGEN_CONSTANTS.list_value {
        let (_, items) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.items)
            .expect("Expected `items` to exist");
        let array = builder.lookup(items.assert_key()).assert_array();

        f.push('[');
        let mut after_first = false;
        for key_or_null in array {
            match key_or_null {
                Primitive::Null => {}
                Primitive::Key(key) => {
                    if after_first {
                        f.push(',');
                    } else {
                        after_first = true;
                    }
                    let object = builder.lookup(*key).assert_object();
                    write_argument_value(f, builder, object)?;
                }
                _ => panic!("Expected an object key or null"),
            }
        }
        f.push(']');
    } else {
        // We filtered out Variables, here it should only be ObjectValue
        let (_, value) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.fields)
            .expect("Expected `fields` to exist");
        let array = builder.lookup(value.assert_key()).assert_array();

        f.push('{');
        for key in array {
            let (_, name) = arg
                .iter()
                .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
                .expect("Expected `name` to exist");
            let name = name.assert_string();
            write!(f, "\\\"{}\\\":", name)?;
            let object = builder.lookup(key.assert_key()).assert_object();
            write_argument_value(f, builder, object)?;
            f.push(',');
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push('}');
    }
    Ok(())
}

fn write_constant_value(f: &mut String, builder: &AstBuilder, value: &Primitive) -> FmtResult {
    match value {
        Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
        Primitive::String(key) => write!(f, "\\\"{}\\\"", key),
        Primitive::Float(value) => write!(f, "{}", value.as_float()),
        Primitive::Int(value) => write!(f, "{}", value),
        Primitive::Key(key) => {
            let ast = builder.lookup(*key);
            match ast {
                Ast::Array(arr) => {
                    f.push('[');
                    for value in arr {
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !arr.is_empty() {
                        f.pop();
                    }
                    f.push(']');
                    Ok(())
                }
                Ast::Object(obj) => {
                    f.push('{');
                    for (name, value) in obj {
                        write!(f, "\\\"{}\\\":", name)?;
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !obj.is_empty() {
                        f.pop();
                    }
                    f.push('}');
                    Ok(())
                }
            }
        }
        Primitive::Null => {
            f.push_str("null");
            Ok(())
        }
        Primitive::StorageKey(_, _) => panic!("Unexpected StorageKey"),
        Primitive::RawString(_) => panic!("Unexpected RawString"),
        Primitive::ModuleDependency(_) => panic!("Unexpected ModuleDependency"),
    }
}
<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use amq_protocol::frame::{AMQPFrame, gen_frame, parse_frame};

use nom::Offset;
use cookie_factory::GenError;
use bytes::{BufMut, BytesMut};
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use futures::{Async,AsyncSink,Poll,Sink,StartSend,Stream,Future,future};
use tokio_codec::{Decoder,Encoder,Framed};
use tokio_io::{AsyncRead,AsyncWrite};
use channel::BasicProperties;
use client::ConnectionOptions;

/// During my testing, it appeared to be the "best" value.
/// To fine-tune it, use a queue with a very large amount of messages, consume it and compare the
/// delivery rate with the ack one.
/// Decreasing this number will increase the ack rate slightly, but decrease the delivery rate by a
/// lot.
/// Increasing this number will increate the delivery rate slightly, but the ack rate will be very
/// close to 0.
const POLL_RECV_LIMIT: u32 = 128;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<AMQPFrame>, io::Error> {
        let (consumed, f) = match parse_frame(buf) {
          Err(e) => {
            if e.is_incomplete() {
              return Ok(None);
            } else {
              return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)));
            }
          },
          Ok((i, frame)) => {
            (buf.offset(i), frame)
          }
        };

        trace!("amqp decoder; frame={:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = io::Error;

    fn encode(&mut self, frame: AMQPFrame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      trace!("encoder; frame={:?}", frame);
      let offset = buf.len();
      loop {
        // If the buffer starts running out of capacity (the threshold is 1/4 of a frame), we
        // reserve more bytes upfront to avoid putting too much strain on the allocator.
        if buf.remaining_mut() < frame_max / 4 {
          trace!("encoder; reserve={}", frame_max * 2);
          buf.reserve(frame_max * 2);
        }

        let gen_res = gen_frame((buf, offset), &frame).map(|tup| tup.1);

        match gen_res {
          Ok(sz) => {
            trace!("encoder; frame_size={}", sz - offset);
            return Ok(());
          },
          Err(GenError::BufferTooSmall(sz)) => {
            // BufferTooSmall error variant returns the index the next write would have
            // occured if there was enough space in the buffer. Thus we subtract the
            // buffer's length to know how much bytes we sould make available.
            let length = buf.len();
            trace!("encoder; sz={} length={} extend={}", sz, length, sz - length);
            buf.extend(repeat(0).take(sz - length));
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream:  Framed<T,AMQPCodec>,
  pub conn:  Connection,
  heartbeat: Option<AMQPFrame>,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: ConnectionOptions) -> impl Future<Item = AMQPTransport<T>, Error = io::Error> + Send + 'static {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);

    future::result(conn.connect()).map_err(|e| {
      let err = format!("Failed to connect: {:?}", e);
      Error::new(ErrorKind::ConnectionAborted, err)
    }).and_then(|_| {
        let codec = AMQPCodec {
          frame_max: conn.configuration.frame_max,
        };
        let t = AMQPTransport {
          upstream:  codec.framed(stream),
          conn:      conn,
          heartbeat: Some(AMQPFrame::Heartbeat(0)),
        };

        AMQPTransportConnector {
          transport: Some(t),
        }
    })
  }

  /// Send a frame to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frame to a queue, to actually send the frame you have to
  /// call either `poll` or `poll_send`.
  pub fn send_frame(&mut self, frame: AMQPFrame) {
    self.conn.frame_queue.push_back(frame);
  }

  /// Send content frames to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frames to a queue, to actually send the frames you have to
  /// call either `poll` or `poll_send`.
  pub fn send_content_frames(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) {
    self.conn.send_content_frames(channel_id, 60, payload, properties);
  }

  /// Preemptively send an heartbeat frame
  pub fn send_heartbeat(&mut self) -> Poll<(), io::Error> {
    if let Some(frame) = self.heartbeat.take() {
      self.conn.frame_queue.push_front(frame);
    }
    self.poll_send().map(|r| r.map(|r| {
      // poll_send succeeded, reinitialize self.heartbeat so that we sned a new frame on the next
      // send_heartbeat call
      self.heartbeat = Some(AMQPFrame::Heartbeat(0));
      r
    }))
  }

  /// Poll the network to receive & handle incoming frames.
  ///
  /// # Return value
  ///
  /// This function will always return `Ok(Async::NotReady)` except in two cases:
  ///
  /// * In case of error, it will return `Err(e)`
  /// * If the socket was closed, it will return `Ok(Async::Ready(()))`
  fn poll_recv(&mut self) -> Poll<(), io::Error> {
    for _ in 0..POLL_RECV_LIMIT {
      match self.upstream.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("transport poll_recv; frame={:?}", frame);
          if let Err(e) = self.conn.handle_frame(frame) {
            let err = format!("failed to handle frame: {:?}", e);
            return Err(io::Error::new(io::ErrorKind::Other, err));
          }
        },
        Ok(Async::Ready(None)) => {
          trace!("transport poll_recv; status=Ready(None)");
          return Ok(Async::Ready(()));
        },
        Ok(Async::NotReady) => {
          trace!("transport poll_recv; status=NotReady");
          return Ok(Async::NotReady);
        },
        Err(e) => {
          error!("transport poll_recv; status=Err({:?})", e);
          return Err(From::from(e));
        },
      };
    }
    Ok(Async::NotReady)
  }

  /// Poll the network to send outcoming frames.
  fn poll_send(&mut self) -> Poll<(), io::Error> {
    let mut ret = Ok(Async::Ready(()));
    while let Some(frame) = self.conn.next_frame() {
      trace!("transport poll_send; frame={:?}", frame);
      match self.start_send(frame)? {
        AsyncSink::Ready => {
          trace!("transport poll_send; status=Ready");
        },
        AsyncSink::NotReady(frame) => {
          trace!("transport poll_send; status=NotReady");
          self.conn.frame_queue.push_front(frame);
          ret = Ok(Async::NotReady);
          break;
        }
      }
    }
    if self.poll_complete()?.is_ready() {
      ret
    } else {
      Ok(Async::NotReady)
    }
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("transport poll");
      if let Async::Ready(()) = self.poll_recv()? {
        trace!("poll transport; status=Ready");
        return Err(io::Error::new(io::ErrorKind::ConnectionAborted, "The connection was closed by the remote peer"));
      }
      self.poll_send().map(|r| r.map(Some))
    }
}

impl <T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Send,
          T: 'static {
    type SinkItem = AMQPFrame;
    type SinkError = io::Error;

    fn start_send(&mut self, frame: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        trace!("transport start_send; frame={:?}", frame);
        self.upstream.start_send(frame)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        trace!("transport poll_complete");
        self.upstream.poll_complete()
    }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    trace!("connector poll; has_transport={:?}", !self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    transport.poll()?;

    trace!("connector poll; state=ConnectionState::{:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      return Ok(Async::Ready(transport))
    }

    self.transport = Some(transport);
    Ok(Async::NotReady)
  }
}

#[macro_export]
macro_rules! lock_transport (
    ($t: expr) => ({
        match $t.lock() {
            Ok(t) => t,
            Err(_) => if $t.is_poisoned() {
                return Err(io::Error::new(io::ErrorKind::Other, "Transport mutex is poisoned"))
            } else {
                task::current().notify();
                return Ok(Async::NotReady)
            }
        }
    });
);

#[cfg(test)]
mod tests {
  extern crate env_logger;

  use super::*;

  #[test]
  fn encode_multiple_frames() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(8, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(16, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(24, buffer.len());
  }

  #[test]
  fn encode_nested_frame() {
    use amq_protocol::frame::AMQPContentHeader;

    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let frame = AMQPFrame::Header(0, 10, AMQPContentHeader {
      class_id: 10,
      weight: 0,
      body_size: 64,
      properties: BasicProperties::default()
    });
    let r = codec.encode(frame, &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(22, buffer.len());
  }

  #[test]
  fn encode_initial_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(true, buffer.capacity() >= frame_max);
    assert_eq!(8, buffer.len());
  }

  #[test]
  fn encode_anticipation_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());
    assert_eq!(8, buffer.len());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it shouldn't trigger buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it should trigger a buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 4, buffer.capacity());
  }
}
<commit_msg>Revert "futures: fix sink usage"<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use amq_protocol::frame::{AMQPFrame, gen_frame, parse_frame};

use nom::Offset;
use cookie_factory::GenError;
use bytes::{BufMut, BytesMut};
use std::cmp;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use futures::{Async,AsyncSink,Poll,Sink,StartSend,Stream,Future,future};
use tokio_codec::{Decoder,Encoder,Framed};
use tokio_io::{AsyncRead,AsyncWrite};
use channel::BasicProperties;
use client::ConnectionOptions;

/// During my testing, it appeared to be the "best" value.
/// To fine-tune it, use a queue with a very large amount of messages, consume it and compare the
/// delivery rate with the ack one.
/// Decreasing this number will increase the ack rate slightly, but decrease the delivery rate by a
/// lot.
/// Increasing this number will increate the delivery rate slightly, but the ack rate will be very
/// close to 0.
const POLL_RECV_LIMIT: u32 = 128;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<AMQPFrame>, io::Error> {
        let (consumed, f) = match parse_frame(buf) {
          Err(e) => {
            if e.is_incomplete() {
              return Ok(None);
            } else {
              return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)));
            }
          },
          Ok((i, frame)) => {
            (buf.offset(i), frame)
          }
        };

        trace!("amqp decoder; frame={:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = io::Error;

    fn encode(&mut self, frame: AMQPFrame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      trace!("encoder; frame={:?}", frame);
      let offset = buf.len();
      loop {
        // If the buffer starts running out of capacity (the threshold is 1/4 of a frame), we
        // reserve more bytes upfront to avoid putting too much strain on the allocator.
        if buf.remaining_mut() < frame_max / 4 {
          trace!("encoder; reserve={}", frame_max * 2);
          buf.reserve(frame_max * 2);
        }

        let gen_res = gen_frame((buf, offset), &frame).map(|tup| tup.1);

        match gen_res {
          Ok(sz) => {
            trace!("encoder; frame_size={}", sz - offset);
            return Ok(());
          },
          Err(GenError::BufferTooSmall(sz)) => {
            // BufferTooSmall error variant returns the index the next write would have
            // occured if there was enough space in the buffer. Thus we subtract the
            // buffer's length to know how much bytes we sould make available.
            let length = buf.len();
            trace!("encoder; sz={} length={} extend={}", sz, length, sz - length);
            buf.extend(repeat(0).take(sz - length));
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream:  Framed<T,AMQPCodec>,
  pub conn:  Connection,
  heartbeat: Option<AMQPFrame>,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: ConnectionOptions) -> impl Future<Item = AMQPTransport<T>, Error = io::Error> + Send + 'static {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);

    future::result(conn.connect()).map_err(|e| {
      let err = format!("Failed to connect: {:?}", e);
      Error::new(ErrorKind::ConnectionAborted, err)
    }).and_then(|_| {
        let codec = AMQPCodec {
          frame_max: conn.configuration.frame_max,
        };
        let t = AMQPTransport {
          upstream:  codec.framed(stream),
          conn:      conn,
          heartbeat: Some(AMQPFrame::Heartbeat(0)),
        };

        AMQPTransportConnector {
          transport: Some(t),
        }
    })
  }

  /// Send a frame to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frame to a queue, to actually send the frame you have to
  /// call either `poll` or `poll_send`.
  pub fn send_frame(&mut self, frame: AMQPFrame) {
    self.conn.frame_queue.push_back(frame);
  }

  /// Send content frames to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frames to a queue, to actually send the frames you have to
  /// call either `poll` or `poll_send`.
  pub fn send_content_frames(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) {
    self.conn.send_content_frames(channel_id, 60, payload, properties);
  }

  /// Preemptively send an heartbeat frame
  pub fn send_heartbeat(&mut self) -> Poll<(), io::Error> {
    if let Some(frame) = self.heartbeat.take() {
      self.conn.frame_queue.push_front(frame);
    }
    self.poll_send().map(|r| r.map(|r| {
      // poll_send succeeded, reinitialize self.heartbeat so that we sned a new frame on the next
      // send_heartbeat call
      self.heartbeat = Some(AMQPFrame::Heartbeat(0));
      r
    }))
  }

  /// Poll the network to receive & handle incoming frames.
  ///
  /// # Return value
  ///
  /// This function will always return `Ok(Async::NotReady)` except in two cases:
  ///
  /// * In case of error, it will return `Err(e)`
  /// * If the socket was closed, it will return `Ok(Async::Ready(()))`
  fn poll_recv(&mut self) -> Poll<(), io::Error> {
    for _ in 0..POLL_RECV_LIMIT {
      match self.upstream.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("transport poll_recv; frame={:?}", frame);
          if let Err(e) = self.conn.handle_frame(frame) {
            let err = format!("failed to handle frame: {:?}", e);
            return Err(io::Error::new(io::ErrorKind::Other, err));
          }
        },
        Ok(Async::Ready(None)) => {
          trace!("transport poll_recv; status=Ready(None)");
          return Ok(Async::Ready(()));
        },
        Ok(Async::NotReady) => {
          trace!("transport poll_recv; status=NotReady");
          return Ok(Async::NotReady);
        },
        Err(e) => {
          error!("transport poll_recv; status=Err({:?})", e);
          return Err(From::from(e));
        },
      };
    }
    Ok(Async::NotReady)
  }

  /// Poll the network to send outcoming frames.
  fn poll_send(&mut self) -> Poll<(), io::Error> {
    while let Some(frame) = self.conn.next_frame() {
      trace!("transport poll_send; frame={:?}", frame);
      match self.start_send(frame)? {
        AsyncSink::Ready => {
          trace!("transport poll_send; status=Ready");
        },
        AsyncSink::NotReady(frame) => {
          trace!("transport poll_send; status=NotReady");
          self.conn.frame_queue.push_front(frame);
          return Ok(Async::NotReady);
        }
      }
    }
    self.poll_complete()
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("transport poll");
      if let Async::Ready(()) = self.poll_recv()? {
        trace!("poll transport; status=Ready");
        return Err(io::Error::new(io::ErrorKind::ConnectionAborted, "The connection was closed by the remote peer"));
      }
      self.poll_send().map(|r| r.map(Some))
    }
}

impl <T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Send,
          T: 'static {
    type SinkItem = AMQPFrame;
    type SinkError = io::Error;

    fn start_send(&mut self, frame: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        trace!("transport start_send; frame={:?}", frame);
        self.upstream.start_send(frame)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        trace!("transport poll_complete");
        self.upstream.poll_complete()
    }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    trace!("connector poll; has_transport={:?}", !self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    transport.poll()?;

    trace!("connector poll; state=ConnectionState::{:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      return Ok(Async::Ready(transport))
    }

    self.transport = Some(transport);
    Ok(Async::NotReady)
  }
}

#[macro_export]
macro_rules! lock_transport (
    ($t: expr) => ({
        match $t.lock() {
            Ok(t) => t,
            Err(_) => if $t.is_poisoned() {
                return Err(io::Error::new(io::ErrorKind::Other, "Transport mutex is poisoned"))
            } else {
                task::current().notify();
                return Ok(Async::NotReady)
            }
        }
    });
);

#[cfg(test)]
mod tests {
  extern crate env_logger;

  use super::*;

  #[test]
  fn encode_multiple_frames() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(8, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(16, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(24, buffer.len());
  }

  #[test]
  fn encode_nested_frame() {
    use amq_protocol::frame::AMQPContentHeader;

    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let frame = AMQPFrame::Header(0, 10, AMQPContentHeader {
      class_id: 10,
      weight: 0,
      body_size: 64,
      properties: BasicProperties::default()
    });
    let r = codec.encode(frame, &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(22, buffer.len());
  }

  #[test]
  fn encode_initial_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(true, buffer.capacity() >= frame_max);
    assert_eq!(8, buffer.len());
  }

  #[test]
  fn encode_anticipation_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());
    assert_eq!(8, buffer.len());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it shouldn't trigger buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it should trigger a buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 4, buffer.capacity());
  }
}
<|endoftext|>