<commit_before>use std::io::{Read, Write};

use super::{CountedList, CountedListWriter, CountedWriter, Deserialize, Error, Serialize, VarInt32, VarUint32, VarUint7};

const R_WEBASSEMBLY_FUNCTION_INDEX_LEB: u8 = 0;
const R_WEBASSEMBLY_TABLE_INDEX_SLEB: u8 = 1;
const R_WEBASSEMBLY_TABLE_INDEX_I32: u8 = 2;
const R_WEBASSEMBLY_MEMORY_ADDR_LEB: u8 = 3;
const R_WEBASSEMBLY_MEMORY_ADDR_SLEB: u8 = 4;
const R_WEBASSEMBLY_MEMORY_ADDR_I32: u8 = 5;
const R_WEBASSEMBLY_TYPE_INDEX_LEB: u8 = 6;
const R_WEBASSEMBLY_GLOBAL_INDEX_LEB: u8 = 7;

/// Relocation information.
#[derive(Clone, Debug)]
pub struct RelocSection {
	// Name of this section.
	name: String,

	// ID of the section containing the relocations described in this section.
	section_id: u32,

	// Name of the section containing the relocations described in this section. Only set if section_id is 0.
	relocation_section_name: Option<String>,

	// Relocation entries.
	entries: Vec<RelocationEntry>,
}

impl RelocSection {
	pub fn deserialize<R: Read>(
		name: String,
		rdr: &mut R,
	) -> Result<Self, Error> {
		let section_id = VarUint32::deserialize(rdr)?.into();
		let relocation_section_name =
			if section_id == 0 {
				Some(String::deserialize(rdr)?)
			}
			else {
				None
			};

		let entries = CountedList::deserialize(rdr)?.into_inner();

		Ok(RelocSection {
			name,
			section_id,
			relocation_section_name,
			entries,
		})
	}
}

impl Serialize for RelocSection {
	type Error = Error;

	fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
		let mut counted_writer = CountedWriter::new(wtr);
		self.name.serialize(&mut counted_writer)?;

		VarUint32::from(self.section_id).serialize(&mut counted_writer)?;

		if let Some(relocation_section_name) = self.relocation_section_name {
			relocation_section_name.serialize(&mut counted_writer)?;
		}

		let counted_list = CountedListWriter(self.entries.len(), self.entries.into_iter());
		counted_list.serialize(&mut counted_writer)?;

		counted_writer.done()?;

		Ok(())
	}
}

/// Relocation entry.
#[derive(Clone, Debug)]
pub enum RelocationEntry {
	// Function index.
	FunctionIndexLeb {
		offset: u32,
		index: u32,
	},

	// Function table index.
	TableIndexSleb {
		offset: u32,
		index: u32,
	},

	// Function table index.
	TableIndexI32 {
		offset: u32,
		index: u32,
	},

	// Linear memory index.
	MemoryAddressLeb {
		offset: u32,
		index: u32,
		addend: i32,
	},

	// Linear memory index.
	MemoryAddressSleb {
		offset: u32,
		index: u32,
		addend: i32,
	},

	// Linear memory index.
	MemoryAddressI32 {
		offset: u32,
		index: u32,
		addend: i32,
	},

	// Type table index.
	TypeIndexLeb {
		offset: u32,
		index: u32,
	},

	// Global index.
	GlobalIndexLeb {
		offset: u32,
		index: u32,
	},
}

impl Deserialize for RelocationEntry {
	type Error = Error;

	fn deserialize<R: Read>(rdr: &mut R) -> Result<Self, Self::Error> {
		match VarUint7::deserialize(rdr)?.into() {
			R_WEBASSEMBLY_FUNCTION_INDEX_LEB => Ok(RelocationEntry::FunctionIndexLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_TABLE_INDEX_SLEB => Ok(RelocationEntry::TableIndexSleb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_TABLE_INDEX_I32 => Ok(RelocationEntry::TableIndexI32 {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_MEMORY_ADDR_LEB => Ok(RelocationEntry::MemoryAddressLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
				addend: VarInt32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_MEMORY_ADDR_SLEB => Ok(RelocationEntry::MemoryAddressSleb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
				addend: VarInt32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_MEMORY_ADDR_I32 => Ok(RelocationEntry::MemoryAddressI32 {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
				addend: VarInt32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_TYPE_INDEX_LEB => Ok(RelocationEntry::TypeIndexLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			R_WEBASSEMBLY_GLOBAL_INDEX_LEB => Ok(RelocationEntry::GlobalIndexLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			entry_type => Err(Error::UnknownValueType(entry_type as i8)),
		}
	}
}

impl Serialize for RelocationEntry {
	type Error = Error;

	fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
		match self {
			RelocationEntry::FunctionIndexLeb { offset, index } => {
				VarUint7::from(R_WEBASSEMBLY_FUNCTION_INDEX_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::TableIndexSleb { offset, index } => {
				VarUint7::from(R_WEBASSEMBLY_TABLE_INDEX_SLEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::TableIndexI32 { offset, index } => {
				VarUint7::from(R_WEBASSEMBLY_TABLE_INDEX_I32).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::MemoryAddressLeb { offset, index, addend } => {
				VarUint7::from(R_WEBASSEMBLY_MEMORY_ADDR_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
				VarInt32::from(addend).serialize(wtr)?;
			},

			RelocationEntry::MemoryAddressSleb { offset, index, addend } => {
				VarUint7::from(R_WEBASSEMBLY_MEMORY_ADDR_SLEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
				VarInt32::from(addend).serialize(wtr)?;
			},

			RelocationEntry::MemoryAddressI32 { offset, index, addend } => {
				VarUint7::from(R_WEBASSEMBLY_MEMORY_ADDR_I32).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
				VarInt32::from(addend).serialize(wtr)?;
			},

			RelocationEntry::TypeIndexLeb { offset, index } => {
				VarUint7::from(R_WEBASSEMBLY_TYPE_INDEX_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::GlobalIndexLeb { offset, index } => {
				VarUint7::from(R_WEBASSEMBLY_GLOBAL_INDEX_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},
		}

		Ok(())
	}
}
<commit_msg>Fixed reloc section doc comments and removed prefix from entry type constants.<commit_after>use std::io::{Read, Write};

use super::{CountedList, CountedListWriter, CountedWriter, Deserialize, Error, Serialize, VarInt32, VarUint32, VarUint7};

const FUNCTION_INDEX_LEB: u8 = 0;
const TABLE_INDEX_SLEB: u8 = 1;
const TABLE_INDEX_I32: u8 = 2;
const MEMORY_ADDR_LEB: u8 = 3;
const MEMORY_ADDR_SLEB: u8 = 4;
const MEMORY_ADDR_I32: u8 = 5;
const TYPE_INDEX_LEB: u8 = 6;
const GLOBAL_INDEX_LEB: u8 = 7;

/// Relocation information.
#[derive(Clone, Debug)]
pub struct RelocSection {
	/// Name of this section.
	name: String,

	/// ID of the section containing the relocations described in this section.
	section_id: u32,

	/// Name of the section containing the relocations described in this section. Only set if section_id is 0.
	relocation_section_name: Option<String>,

	/// Relocation entries.
	entries: Vec<RelocationEntry>,
}

impl RelocSection {
	pub fn deserialize<R: Read>(
		name: String,
		rdr: &mut R,
	) -> Result<Self, Error> {
		let section_id = VarUint32::deserialize(rdr)?.into();

		let relocation_section_name =
			if section_id == 0 {
				Some(String::deserialize(rdr)?)
			}
			else {
				None
			};

		let entries = CountedList::deserialize(rdr)?.into_inner();

		Ok(RelocSection {
			name,
			section_id,
			relocation_section_name,
			entries,
		})
	}
}

impl Serialize for RelocSection {
	type Error = Error;

	fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
		let mut counted_writer = CountedWriter::new(wtr);

		self.name.serialize(&mut counted_writer)?;

		VarUint32::from(self.section_id).serialize(&mut counted_writer)?;

		if let Some(relocation_section_name) = self.relocation_section_name {
			relocation_section_name.serialize(&mut counted_writer)?;
		}

		let counted_list = CountedListWriter(self.entries.len(), self.entries.into_iter());
		counted_list.serialize(&mut counted_writer)?;

		counted_writer.done()?;

		Ok(())
	}
}

/// Relocation entry.
#[derive(Clone, Debug)]
pub enum RelocationEntry {
	/// Function index.
	FunctionIndexLeb {
		/// Offset of the value to rewrite.
		offset: u32,

		/// Index of the function symbol in the symbol table.
		index: u32,
	},

	/// Function table index.
	TableIndexSleb {
		offset: u32,
		index: u32,
	},

	/// Function table index.
	TableIndexI32 {
		offset: u32,
		index: u32,
	},

	/// Linear memory index.
	MemoryAddressLeb {
		offset: u32,
		index: u32,
		addend: i32,
	},

	/// Linear memory index.
	MemoryAddressSleb {
		offset: u32,
		index: u32,
		addend: i32,
	},

	/// Linear memory index.
	MemoryAddressI32 {
		offset: u32,
		index: u32,
		addend: i32,
	},

	/// Type table index.
	TypeIndexLeb {
		offset: u32,
		index: u32,
	},

	/// Global index.
	GlobalIndexLeb {
		offset: u32,
		index: u32,
	},
}

impl Deserialize for RelocationEntry {
	type Error = Error;

	fn deserialize<R: Read>(rdr: &mut R) -> Result<Self, Self::Error> {
		match VarUint7::deserialize(rdr)?.into() {
			FUNCTION_INDEX_LEB => Ok(RelocationEntry::FunctionIndexLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			TABLE_INDEX_SLEB => Ok(RelocationEntry::TableIndexSleb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			TABLE_INDEX_I32 => Ok(RelocationEntry::TableIndexI32 {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			MEMORY_ADDR_LEB => Ok(RelocationEntry::MemoryAddressLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
				addend: VarInt32::deserialize(rdr)?.into(),
			}),

			MEMORY_ADDR_SLEB => Ok(RelocationEntry::MemoryAddressSleb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
				addend: VarInt32::deserialize(rdr)?.into(),
			}),

			MEMORY_ADDR_I32 => Ok(RelocationEntry::MemoryAddressI32 {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
				addend: VarInt32::deserialize(rdr)?.into(),
			}),

			TYPE_INDEX_LEB => Ok(RelocationEntry::TypeIndexLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			GLOBAL_INDEX_LEB => Ok(RelocationEntry::GlobalIndexLeb {
				offset: VarUint32::deserialize(rdr)?.into(),
				index: VarUint32::deserialize(rdr)?.into(),
			}),

			entry_type => Err(Error::UnknownValueType(entry_type as i8)),
		}
	}
}

impl Serialize for RelocationEntry {
	type Error = Error;

	fn serialize<W: Write>(self, wtr: &mut W) -> Result<(), Error> {
		match self {
			RelocationEntry::FunctionIndexLeb { offset, index } => {
				VarUint7::from(FUNCTION_INDEX_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::TableIndexSleb { offset, index } => {
				VarUint7::from(TABLE_INDEX_SLEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::TableIndexI32 { offset, index } => {
				VarUint7::from(TABLE_INDEX_I32).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::MemoryAddressLeb { offset, index, addend } => {
				VarUint7::from(MEMORY_ADDR_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
				VarInt32::from(addend).serialize(wtr)?;
			},

			RelocationEntry::MemoryAddressSleb { offset, index, addend } => {
				VarUint7::from(MEMORY_ADDR_SLEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
				VarInt32::from(addend).serialize(wtr)?;
			},

			RelocationEntry::MemoryAddressI32 { offset, index, addend } => {
				VarUint7::from(MEMORY_ADDR_I32).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
				VarInt32::from(addend).serialize(wtr)?;
			},

			RelocationEntry::TypeIndexLeb { offset, index } => {
				VarUint7::from(TYPE_INDEX_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},

			RelocationEntry::GlobalIndexLeb { offset, index } => {
				VarUint7::from(GLOBAL_INDEX_LEB).serialize(wtr)?;
				VarUint32::from(offset).serialize(wtr)?;
				VarUint32::from(index).serialize(wtr)?;
			},
		}

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


<commit_before>use anyhow::{bail, Context};
//use indicatif::{ProgressBar, ProgressStyle};
use serde::Deserialize;
use std::{
    collections::BTreeMap,
    fs::File,
    io::{BufRead, BufReader, Write},
    process::{Command, Output, Stdio},
    str,
    time::SystemTime,
};
use structopt::StructOpt;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
use vcpkg::{find_vcpkg_root, Config};

// settings for a specific Rust target
#[derive(Debug, Deserialize)]
struct Target {
    triplet: Option<String>,
    // this install key for a specific target overrides the main entry
    // so a the target can opt out of installing packages
    // #[serde(default = "Vec::new")]
    install: Option<Vec<String>>,
}

#[derive(Debug, Deserialize)]
struct Vcpkg {
    vcpkg_root: Option<String>,
    #[serde(default = "BTreeMap::new")]
    target: BTreeMap<String, Target>,
    branch: Option<String>,
    rev: Option<String>,
    git: Option<String>,
    tag: Option<String>,
    #[serde(default = "Vec::new")]
    install: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct Metadata {
    vcpkg: Vcpkg,
}
#[derive(Debug, PartialEq, StructOpt)]
/// Install vcpkg and build packages
///
/// This program clones vcpkg from the specified source and
/// compiles it. It then builds packages required by crates
/// that are depended on by the top level crate being built.
#[structopt(rename_all = "kebab-case")]
struct Opt {
    ///
    #[structopt(short, long)]
    verbose: bool,

    // #[structopt(long)]
    // manifest_path: Option<String>,
    #[structopt(subcommand)]
    sub: Subcommands,
}

#[derive(Debug, PartialEq, StructOpt)]
enum Subcommands {
    /// Build packages
    ///
    /// This command will clone or update a vcpkg tree to the version specified
    /// in Cargo.toml and build the required packages.
    Build {
        #[structopt(long)]
        /// Path to Cargo.toml
        manifest_path: Option<String>,

        #[structopt(long)]
        /// Build for the target triple
        target: Option<String>,
    },
    // `external_subcommand` tells structopt to put
    // all the extra arguments into this Vec
    // #[structopt(external_subcommand)]
    // Cargo(Vec<String>),
}

fn main() {
    // cargo passes us the "vcpkg" arg when it calls us. Drop it before
    // parsing the arg list so t doesn't end up the usage description
    let mut args = std::env::args().collect::<Vec<_>>();
    if args.len() >= 2 && args[1] == "vcpkg" {
        args.remove(1);
    }
    let args = Opt::from_iter(args);

    match args.sub {
        Subcommands::Build { .. } => {
            build(args).unwrap_or_else(|e| {
                eprintln!("cargo-vcpkg: {}", e);
                std::process::exit(1);
            });
        }
    }
}

fn build(opt: Opt) -> Result<(), anyhow::Error> {
    let start_time = SystemTime::now();

    let target_triple = target_triple();

    let verbose = opt.verbose;

    let mut args = std::env::args().skip_while(|val| !val.starts_with("--manifest-path"));
    let mut cmd = cargo_metadata::MetadataCommand::new();

    match args.next() {
        Some(p) if p == "--manifest-path" => {
            cmd.manifest_path(args.next().unwrap());
        }
        Some(p) => {
            cmd.manifest_path(p.trim_start_matches("--manifest-path="));
        }
        None => {}
    }
    let metadata = cmd.exec()?;

    let resolve = metadata.resolve.as_ref().unwrap();

    let root_crate = resolve
        .root
        .as_ref()
        .context("cannot run on a virtual manifest, this command requires running against an actual package in this workspace.")?;

    let mut git_url = None;
    let mut vcpkg_ports = Vec::new();
    let mut rev_tag_branch: Option<String> = None;
    let mut vcpkg_triplet = None;
    for p in &metadata.packages {
        if let Ok(v) = serde_json::from_value::<Metadata>(p.metadata.clone()) {
            let v = v.vcpkg;
            let is_root_crate = p.id == *root_crate;

            // only use git url and rev from the root crate
            if v.git.is_some() && is_root_crate {
                git_url = v.git;

                // TODO: check the target and use it's package set if required
                // TODO: get the correct target
                // TODO: make sure to pull if it's a branch
                rev_tag_branch = match (&v.branch, &v.tag, &v.rev) {
                    (Some(b), None, None) => Some(b.into()),
                    (None, Some(t), None) => Some(t.into()),
                    (None, None, Some(r)) => Some(r.into()),
                    _ => {
                        bail!("must specify one of branch,rev,tag for git source");
                    }
                };
            }

            // if there is specific configuration for the target and it has
            // an install key, use that rather than the general install key
            match v.target.get(&target_triple) {
                Some(target) => {
                    if target.install.is_some() {
                        vcpkg_ports.extend_from_slice(&target.install.as_ref().unwrap().as_slice());
                    }
                    if is_root_crate {
                        vcpkg_triplet = target.triplet.clone();
                    }
                }
                _ => {
                    // not found or install is empty
                    vcpkg_ports.extend_from_slice(&v.install.as_slice());
                }
            }
        }
    }

    // should we modify the existing?
    // let mut allow_updates = true;

    // find the vcpkg root
    let vcpkg_root = find_vcpkg_root(&Config::default()).unwrap_or_else(|_| {
        let target_directory = metadata.target_directory.clone();
        let mut vcpkg_root = target_directory.clone();
        vcpkg_root.push("vcpkg");
        vcpkg_root.to_path_buf();
        vcpkg_root
    });
    if verbose {
        println!("vcpkg root is {}", vcpkg_root.display());
    }
    // if it does not exist, clone vcpkg from git
    let mut vcpkg_root_file = vcpkg_root.clone();
    vcpkg_root_file.push(".vcpkg-root");
    if !vcpkg_root_file.exists() {
        let git_url = git_url.context(format!(
            "could not find a vcpkg installation and crate \
        {} does not specify a git repository to clone from.",
            root_crate
        ))?;
        print_tag("Cloning", &git_url);
        let mut cmd = Command::new("git");
        cmd.arg("clone");
        cmd.arg(git_url);
        cmd.arg(&vcpkg_root);
        let _output = run_command(cmd, verbose).context("failed to run git clone")?;

    //eprintln!("git clone done = {:?}", output.status);
    } else {
        print_tag("Fetching", "vcpkg");
        let mut cmd = Command::new("git");
        cmd.arg("fetch");
        cmd.arg("--verbose");
        cmd.arg("--all");
        let output = run_command(cmd, verbose).context("failed to run git fetch")?;

        if !output.status.success() {
            bail!("fetch failed");
        }
    }

    // create a cargo-vcpkg.toml in the vcpkg tree
    let mut cargo_vcpkg_config_file = vcpkg_root.clone();
    cargo_vcpkg_config_file.push("downloads");
    std::fs::create_dir_all(&cargo_vcpkg_config_file)
        .context("could not create downloads directory in vcpkg tree")?;
    cargo_vcpkg_config_file.push("cargo-vcpkg.toml");
    if !cargo_vcpkg_config_file.exists() {
        let mut file =
            File::create(cargo_vcpkg_config_file).context("could not create cargo-vcpkg.toml")?;
        file.write_all(b"# This file was created automatically by cargo-vcpkg\n")?;
    }

    // otherwise, check that the rev is where we want it to be
    // there needs to be some serious thought here because if we are on a branch
    // does this mean we should fetch?

    // check out the required rev
    let rev_tag_branch = rev_tag_branch.unwrap();
    print_tag("Checkout", &format!("rev/tag/branch {}", rev_tag_branch));
    let mut cmd = Command::new("git");
    cmd.arg("checkout");
    cmd.arg(rev_tag_branch);
    cmd.current_dir(&vcpkg_root);
    run_command(cmd, verbose).context("failed to execute process")?;

    // try and run 'vcpkg update' and if it fails or gives the version warning, rebuild it
    let require_bootstrap = match vcpkg_command(&vcpkg_root, &vcpkg_triplet)
        .arg("update")
        .output()
    {
        Ok(output) => {
            if verbose {
                println!("-- stdout --\n{}", String::from_utf8_lossy(&output.stdout));
                println!("-- stderr --\n{}", String::from_utf8_lossy(&output.stderr));
                println!("{:?}", output.status);
            }
            if output.status.success()
                && !str::from_utf8(&output.stdout)
                    .unwrap()
                    .contains("Warning: Different source is available for vcpkg")
            {
                false
            } else {
                true
            }
        }
        Err(_) => true,
    };

    if require_bootstrap {
        print_tag("Compiling", "vcpkg");
        let mut cmd = if cfg!(windows) {
            let mut cmd = Command::new("cmd");
            cmd.arg("/C");
            cmd.arg("bootstrap-vcpkg.bat");
            cmd
        } else {
            let mut cmd = Command::new("sh");
            cmd.arg("-c");
            cmd.arg("./bootstrap-vcpkg.sh");
            cmd
        };
        cmd.arg("-disableMetrics");
        cmd.current_dir(&vcpkg_root);
        run_command(cmd, verbose).context("failed to run vcpkg bootstrap")?;
    }

    // TODO: upgrade anything that is installed
    print_tag("Installing", &vcpkg_ports.join(" "));
    let mut v = vcpkg_command(&vcpkg_root, &vcpkg_triplet);
    v.arg("install");
    v.arg("--recurse");
    v.args(vcpkg_ports.as_slice());
    v.stdout(Stdio::piped());

    let mut output = v.spawn()?;

    let reader = BufReader::new(output.stdout.take().context("could not get stdout")?);

    // let style = ProgressStyle::default_bar()
    //     .progress_chars("=> ")
    //     .template("    Building [{bar:40.cyan/blue}] {pos:>7}/{len:7} {msg}");
    // let bar = ProgressBar::new(10).with_style(style);
    for line in reader.lines().flat_map(Result::ok) {
        parse_build_line(&line).map(|(pkg, triplet, _num, _tot)| {
            print_tag("Compiling", &format!("{} (triplet {})", pkg, triplet))
        });

        if verbose {
            println!("{}", line);
        }
        // bar.set_length(tot);
        // bar.set_position(num);
    }
    // grab anything that is left
    let output = output.wait_with_output()?;

    if !output.status.success() && !verbose {
        println!("-- stdout --\n{}", String::from_utf8_lossy(&output.stdout));
        println!("-- stderr --\n{}", String::from_utf8_lossy(&output.stderr));
        bail!("failed");
    }

    let duration = SystemTime::now().duration_since(start_time).unwrap();
    print_tag("Finished", &format!("in {:0.2}s", duration.as_secs_f32()));
    Ok(())
}

fn target_triple() -> String {
    let mut args = std::env::args().skip_while(|val| !val.starts_with("--target"));
    match args.next() {
        Some(p) if p == "--target" => args.next().unwrap(),
        Some(p) => p.trim_start_matches("--target=").into(),
        None => std::env!("TARGET").into(),
    }
}

fn vcpkg_command(vcpkg_root: &std::path::Path, vcpkg_triplet: &Option<String>) -> Command {
    let mut x = vcpkg_root.to_path_buf();
    if cfg!(windows) {
        x.push("vcpkg.exe");
    } else {
        x.push("vcpkg")
    }
    let mut command = Command::new(x);
    command.current_dir(&vcpkg_root);
    if let Some(triplet) = &vcpkg_triplet {
        command.arg("--triplet");
        command.arg(triplet);
    }
    command
}

fn run_command(mut cmd: Command, verbose: bool) -> Result<Output, anyhow::Error> {
    if verbose {
        cmd.stdout(Stdio::inherit());
        cmd.stderr(Stdio::inherit());
    }
    let output = cmd.output()?;

    if !output.status.success() && !verbose {
        println!("-- stdout --\n{}", String::from_utf8_lossy(&output.stdout));
        println!("-- stderr --\n{}", String::from_utf8_lossy(&output.stderr));
        bail!("failed");
    }

    Ok(output)
}

fn print_tag(tag: &str, detail: &str) {
    let mut stdout = StandardStream::stdout(ColorChoice::Auto);
    stdout
        .set_color(ColorSpec::new().set_fg(Some(Color::Green)).set_bold(true))
        .unwrap();

    print!("{:>12} ", tag);
    stdout.reset().unwrap();
    println!("{}", detail);
}

fn parse_build_line(line: &str) -> Option<(String, String, u64, u64)> {
    let line = Some(line)
        .filter(|line| line.starts_with("Starting package "))
        .map(|line| line.trim_start_matches("Starting package ").to_string())?;

    let progress_and_pkg_trp = line.splitn(2, ":").collect::<Vec<_>>();
    if progress_and_pkg_trp.len() != 2 {
        return None;
    }

    let pkg_with_triplet = progress_and_pkg_trp[1].trim();

    let (pkg, triplet) = match pkg_with_triplet
        .rsplitn(2, ":")
        .collect::<Vec<_>>()
        .as_slice()
    {
        [t, p] => (p.to_string(), t.to_string()),
        _ => return None,
    };

    let (cnt, tot) = match &progress_and_pkg_trp[0]
        .splitn(2, "/")
        .filter_map(|s| s.parse::<u64>().ok())
        .collect::<Vec<_>>()
        .as_slice()
    {
        [cnt, tot] => (*cnt, *tot),
        _ => (0, 0),
    };

    Some((pkg, triplet, cnt, tot))
}
//    Building [==============================================> ] 58/59
<commit_msg>flesh out some help comments<commit_after>use anyhow::{bail, Context};
//use indicatif::{ProgressBar, ProgressStyle};
use serde::Deserialize;
use std::{
    collections::BTreeMap,
    fs::File,
    io::{BufRead, BufReader, Write},
    process::{Command, Output, Stdio},
    str,
    time::SystemTime,
};
use structopt::StructOpt;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
use vcpkg::{find_vcpkg_root, Config};

// settings for a specific Rust target
#[derive(Debug, Deserialize)]
struct Target {
    triplet: Option<String>,
    // this install key for a specific target overrides the main entry
    // so a the target can opt out of installing packages
    // #[serde(default = "Vec::new")]
    install: Option<Vec<String>>,
}

#[derive(Debug, Deserialize)]
struct Vcpkg {
    vcpkg_root: Option<String>,
    #[serde(default = "BTreeMap::new")]
    target: BTreeMap<String, Target>,
    branch: Option<String>,
    rev: Option<String>,
    git: Option<String>,
    tag: Option<String>,
    #[serde(default = "Vec::new")]
    install: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct Metadata {
    vcpkg: Vcpkg,
}
#[derive(Debug, PartialEq, StructOpt)]
/// Install vcpkg and build packages
///
/// This program clones vcpkg from the specified source and
/// compiles it. It then builds packages required by crates
/// that are depended on by the top level crate being built.
#[structopt(rename_all = "kebab-case")]
struct Opt {
    /// Display more information while building
    ///
    /// This will display the output of git and
    /// vcpkg operations.
    #[structopt(short, long)]
    verbose: bool,

    // #[structopt(long)]
    // manifest_path: Option<String>,
    #[structopt(subcommand)]
    sub: Subcommands,
}

#[derive(Debug, PartialEq, StructOpt)]
enum Subcommands {
    /// Build packages, checking out the correct version
    ///
    /// This command will clone or update a vcpkg tree to the version specified
    /// in Cargo.toml and build the required packages.
    Build {
        #[structopt(long)]
        /// Path to Cargo.toml
        manifest_path: Option<String>,

        #[structopt(long)]
        /// Build for the target triple
        target: Option<String>,
    },
    // `external_subcommand` tells structopt to put
    // all the extra arguments into this Vec
    // #[structopt(external_subcommand)]
    // Cargo(Vec<String>),
}

fn main() {
    // cargo passes us the "vcpkg" arg when it calls us. Drop it before
    // parsing the arg list so t doesn't end up the usage description
    let mut args = std::env::args().collect::<Vec<_>>();
    if args.len() >= 2 && args[1] == "vcpkg" {
        args.remove(1);
    }
    let args = Opt::from_iter(args);

    match args.sub {
        Subcommands::Build { .. } => {
            build(args).unwrap_or_else(|e| {
                eprintln!("cargo-vcpkg: {}", e);
                std::process::exit(1);
            });
        }
    }
}

fn build(opt: Opt) -> Result<(), anyhow::Error> {
    let start_time = SystemTime::now();

    let target_triple = target_triple();

    let verbose = opt.verbose;

    let mut args = std::env::args().skip_while(|val| !val.starts_with("--manifest-path"));
    let mut cmd = cargo_metadata::MetadataCommand::new();

    match args.next() {
        Some(p) if p == "--manifest-path" => {
            cmd.manifest_path(args.next().unwrap());
        }
        Some(p) => {
            cmd.manifest_path(p.trim_start_matches("--manifest-path="));
        }
        None => {}
    }
    let metadata = cmd.exec()?;

    let resolve = metadata.resolve.as_ref().unwrap();

    let root_crate = resolve
        .root
        .as_ref()
        .context("cannot run on a virtual manifest, this command requires running against an actual package in this workspace.")?;

    let mut git_url = None;
    let mut vcpkg_ports = Vec::new();
    let mut rev_tag_branch: Option<String> = None;
    let mut vcpkg_triplet = None;
    for p in &metadata.packages {
        if let Ok(v) = serde_json::from_value::<Metadata>(p.metadata.clone()) {
            let v = v.vcpkg;
            let is_root_crate = p.id == *root_crate;

            // only use git url and rev from the root crate
            if v.git.is_some() && is_root_crate {
                git_url = v.git;

                // TODO: check the target and use it's package set if required
                // TODO: get the correct target
                // TODO: make sure to pull if it's a branch
                rev_tag_branch = match (&v.branch, &v.tag, &v.rev) {
                    (Some(b), None, None) => Some(b.into()),
                    (None, Some(t), None) => Some(t.into()),
                    (None, None, Some(r)) => Some(r.into()),
                    _ => {
                        bail!("must specify one of branch,rev,tag for git source");
                    }
                };
            }

            // if there is specific configuration for the target and it has
            // an install key, use that rather than the general install key
            match v.target.get(&target_triple) {
                Some(target) => {
                    if target.install.is_some() {
                        vcpkg_ports.extend_from_slice(&target.install.as_ref().unwrap().as_slice());
                    }
                    if is_root_crate {
                        vcpkg_triplet = target.triplet.clone();
                    }
                }
                _ => {
                    // not found or install is empty
                    vcpkg_ports.extend_from_slice(&v.install.as_slice());
                }
            }
        }
    }

    // should we modify the existing?
    // let mut allow_updates = true;

    // find the vcpkg root
    let vcpkg_root = find_vcpkg_root(&Config::default()).unwrap_or_else(|_| {
        let target_directory = metadata.target_directory.clone();
        let mut vcpkg_root = target_directory.clone();
        vcpkg_root.push("vcpkg");
        vcpkg_root.to_path_buf();
        vcpkg_root
    });
    if verbose {
        println!("vcpkg root is {}", vcpkg_root.display());
    }
    // if it does not exist, clone vcpkg from git
    let mut vcpkg_root_file = vcpkg_root.clone();
    vcpkg_root_file.push(".vcpkg-root");
    if !vcpkg_root_file.exists() {
        let git_url = git_url.context(format!(
            "could not find a vcpkg installation and crate \
        {} does not specify a git repository to clone from.",
            root_crate
        ))?;
        print_tag("Cloning", &git_url);
        let mut cmd = Command::new("git");
        cmd.arg("clone");
        cmd.arg(git_url);
        cmd.arg(&vcpkg_root);
        let _output = run_command(cmd, verbose).context("failed to run git clone")?;

    //eprintln!("git clone done = {:?}", output.status);
    } else {
        print_tag("Fetching", "vcpkg");
        let mut cmd = Command::new("git");
        cmd.arg("fetch");
        cmd.arg("--verbose");
        cmd.arg("--all");
        let output = run_command(cmd, verbose).context("failed to run git fetch")?;

        if !output.status.success() {
            bail!("fetch failed");
        }
    }

    // create a cargo-vcpkg.toml in the vcpkg tree
    let mut cargo_vcpkg_config_file = vcpkg_root.clone();
    cargo_vcpkg_config_file.push("downloads");
    std::fs::create_dir_all(&cargo_vcpkg_config_file)
        .context("could not create downloads directory in vcpkg tree")?;
    cargo_vcpkg_config_file.push("cargo-vcpkg.toml");
    if !cargo_vcpkg_config_file.exists() {
        let mut file =
            File::create(cargo_vcpkg_config_file).context("could not create cargo-vcpkg.toml")?;
        file.write_all(b"# This file was created automatically by cargo-vcpkg\n")?;
    }

    // otherwise, check that the rev is where we want it to be
    // there needs to be some serious thought here because if we are on a branch
    // does this mean we should fetch?

    // check out the required rev
    let rev_tag_branch = rev_tag_branch.unwrap();
    print_tag("Checkout", &format!("rev/tag/branch {}", rev_tag_branch));
    let mut cmd = Command::new("git");
    cmd.arg("checkout");
    cmd.arg(rev_tag_branch);
    cmd.current_dir(&vcpkg_root);
    run_command(cmd, verbose).context("failed to execute process")?;

    // try and run 'vcpkg update' and if it fails or gives the version warning, rebuild it
    let require_bootstrap = match vcpkg_command(&vcpkg_root, &vcpkg_triplet)
        .arg("update")
        .output()
    {
        Ok(output) => {
            if verbose {
                println!("-- stdout --\n{}", String::from_utf8_lossy(&output.stdout));
                println!("-- stderr --\n{}", String::from_utf8_lossy(&output.stderr));
                println!("{:?}", output.status);
            }
            if output.status.success()
                && !str::from_utf8(&output.stdout)
                    .unwrap()
                    .contains("Warning: Different source is available for vcpkg")
            {
                false
            } else {
                true
            }
        }
        Err(_) => true,
    };

    if require_bootstrap {
        print_tag("Compiling", "vcpkg");
        let mut cmd = if cfg!(windows) {
            let mut cmd = Command::new("cmd");
            cmd.arg("/C");
            cmd.arg("bootstrap-vcpkg.bat");
            cmd
        } else {
            let mut cmd = Command::new("sh");
            cmd.arg("-c");
            cmd.arg("./bootstrap-vcpkg.sh");
            cmd
        };
        cmd.arg("-disableMetrics");
        cmd.current_dir(&vcpkg_root);
        run_command(cmd, verbose).context("failed to run vcpkg bootstrap")?;
    }

    // TODO: upgrade anything that is installed
    print_tag("Installing", &vcpkg_ports.join(" "));
    let mut v = vcpkg_command(&vcpkg_root, &vcpkg_triplet);
    v.arg("install");
    v.arg("--recurse");
    v.args(vcpkg_ports.as_slice());
    v.stdout(Stdio::piped());

    let mut output = v.spawn()?;

    let reader = BufReader::new(output.stdout.take().context("could not get stdout")?);

    // let style = ProgressStyle::default_bar()
    //     .progress_chars("=> ")
    //     .template("    Building [{bar:40.cyan/blue}] {pos:>7}/{len:7} {msg}");
    // let bar = ProgressBar::new(10).with_style(style);
    for line in reader.lines().flat_map(Result::ok) {
        parse_build_line(&line).map(|(pkg, triplet, _num, _tot)| {
            print_tag("Compiling", &format!("{} (triplet {})", pkg, triplet))
        });

        if verbose {
            println!("{}", line);
        }
        // bar.set_length(tot);
        // bar.set_position(num);
    }
    // grab anything that is left
    let output = output.wait_with_output()?;

    if !output.status.success() && !verbose {
        println!("-- stdout --\n{}", String::from_utf8_lossy(&output.stdout));
        println!("-- stderr --\n{}", String::from_utf8_lossy(&output.stderr));
        bail!("failed");
    }

    let duration = SystemTime::now().duration_since(start_time).unwrap();
    print_tag("Finished", &format!("in {:0.2}s", duration.as_secs_f32()));
    Ok(())
}

fn target_triple() -> String {
    let mut args = std::env::args().skip_while(|val| !val.starts_with("--target"));
    match args.next() {
        Some(p) if p == "--target" => args.next().unwrap(),
        Some(p) => p.trim_start_matches("--target=").into(),
        None => std::env!("TARGET").into(),
    }
}

fn vcpkg_command(vcpkg_root: &std::path::Path, vcpkg_triplet: &Option<String>) -> Command {
    let mut x = vcpkg_root.to_path_buf();
    if cfg!(windows) {
        x.push("vcpkg.exe");
    } else {
        x.push("vcpkg")
    }
    let mut command = Command::new(x);
    command.current_dir(&vcpkg_root);
    if let Some(triplet) = &vcpkg_triplet {
        command.arg("--triplet");
        command.arg(triplet);
    }
    command
}

fn run_command(mut cmd: Command, verbose: bool) -> Result<Output, anyhow::Error> {
    if verbose {
        cmd.stdout(Stdio::inherit());
        cmd.stderr(Stdio::inherit());
    }
    let output = cmd.output()?;

    if !output.status.success() && !verbose {
        println!("-- stdout --\n{}", String::from_utf8_lossy(&output.stdout));
        println!("-- stderr --\n{}", String::from_utf8_lossy(&output.stderr));
        bail!("failed");
    }

    Ok(output)
}

fn print_tag(tag: &str, detail: &str) {
    let mut stdout = StandardStream::stdout(ColorChoice::Auto);
    stdout
        .set_color(ColorSpec::new().set_fg(Some(Color::Green)).set_bold(true))
        .unwrap();

    print!("{:>12} ", tag);
    stdout.reset().unwrap();
    println!("{}", detail);
}

fn parse_build_line(line: &str) -> Option<(String, String, u64, u64)> {
    let line = Some(line)
        .filter(|line| line.starts_with("Starting package "))
        .map(|line| line.trim_start_matches("Starting package ").to_string())?;

    let progress_and_pkg_trp = line.splitn(2, ":").collect::<Vec<_>>();
    if progress_and_pkg_trp.len() != 2 {
        return None;
    }

    let pkg_with_triplet = progress_and_pkg_trp[1].trim();

    let (pkg, triplet) = match pkg_with_triplet
        .rsplitn(2, ":")
        .collect::<Vec<_>>()
        .as_slice()
    {
        [t, p] => (p.to_string(), t.to_string()),
        _ => return None,
    };

    let (cnt, tot) = match &progress_and_pkg_trp[0]
        .splitn(2, "/")
        .filter_map(|s| s.parse::<u64>().ok())
        .collect::<Vec<_>>()
        .as_slice()
    {
        [cnt, tot] => (*cnt, *tot),
        _ => (0, 0),
    };

    Some((pkg, triplet, cnt, tot))
}
//    Building [==============================================> ] 58/59
<|endoftext|>


<commit_before>use std::io;
use std::net::ToSocketAddrs;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use futures::{ready, Sink, Stream, StreamExt};
use pin_project::{pin_project, project};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;
use tokio_rustls::rustls::ClientConfig;
use tokio_rustls::webpki::DNSNameRef;
use tokio_rustls::{client::TlsStream, TlsConnector};
use tokio_util::codec::{Decoder, Framed};

use crate::codec::{ImapCodec, ResponseData};
use imap_proto::builders::command::Command;
use imap_proto::{Request, RequestId, State};

pub type TlsClient = Client<TlsStream<TcpStream>>;

pub struct Client<T> {
    transport: Framed<T, ImapCodec>,
    state: State,
    request_ids: IdGenerator,
}

impl TlsClient {
    pub async fn connect(server: &str) -> io::Result<(ResponseData, Self)> {
        let addr = (server, 993).to_socket_addrs()?.next().ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::Other,
                format!("no IP addresses found for {}", server),
            )
        })?;

        let mut tls_config = ClientConfig::new();
        tls_config
            .root_store
            .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
        let connector: TlsConnector = Arc::new(tls_config).into();

        let stream = TcpStream::connect(&addr).await?;
        let stream = connector
            .connect(DNSNameRef::try_from_ascii_str(server).unwrap(), stream)
            .await?;
        let mut transport = ImapCodec::default().framed(stream);

        let greeting = match transport.next().await {
            Some(greeting) => Ok(greeting),
            None => Err(io::Error::new(io::ErrorKind::Other, "no greeting found")),
        }?;
        let client = Client {
            transport,
            state: State::NotAuthenticated,
            request_ids: IdGenerator::new(),
        };

        greeting.map(|greeting| (greeting, client))
    }

    pub fn call<C: Into<Command>>(&mut self, cmd: C) -> ResponseStream<TlsStream<TcpStream>> {
        let request_id = self.request_ids.next().unwrap(); // safe: never returns Err,
        ResponseStream {
            client: self,
            request_id,
            cmd: cmd.into(),
            state: ResponseStreamState::Start,
        }
    }
}

#[pin_project]
pub struct ResponseStream<'a, T> {
    #[pin]
    client: &'a mut Client<T>,
    request_id: RequestId,
    cmd: Command,
    state: ResponseStreamState,
}

impl<'a, T> Stream for ResponseStream<'a, T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    type Item = Result<ResponseData, io::Error>;

    #[project]
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let mut me = self.project();
        loop {
            match me.state {
                ResponseStreamState::Start => {
                    ready!(Pin::new(&mut me.client.transport).poll_ready(cx))?;
                    let pinned = Pin::new(&mut me.client.transport);
                    pinned.start_send(&Request(me.request_id.as_bytes(), &me.cmd.args))?;
                    *me.state = ResponseStreamState::Sending;
                }
                ResponseStreamState::Sending => {
                    let pinned = Pin::new(&mut me.client.transport);
                    ready!(pinned.poll_flush(cx))?;
                    *me.state = ResponseStreamState::Receiving;
                }
                ResponseStreamState::Receiving => {
                    match ready!(Pin::new(&mut me.client.transport).poll_next(cx)) {
                        Some(Ok(rsp)) => {
                            match rsp.request_id() {
                                Some(req_id) if req_id == me.request_id => {}
                                Some(_) | None => return Poll::Ready(Some(Ok(rsp))),
                            }

                            if let Some(next_state) = me.cmd.next_state.as_ref() {
                                me.client.state = *next_state;
                            }
                            *me.state = ResponseStreamState::Done;
                            return Poll::Ready(Some(Ok(rsp)));
                        }
                        Some(Err(e)) => return Poll::Ready(Some(Err(e))),
                        None => {
                            return Poll::Ready(Some(Err(io::Error::new(
                                io::ErrorKind::Other,
                                "stream ended before command completion",
                            ))))
                        }
                    }
                }
                ResponseStreamState::Done => {
                    return Poll::Ready(None);
                }
            }
        }
    }
}

enum ResponseStreamState {
    Start,
    Sending,
    Receiving,
    Done,
}

pub struct IdGenerator {
    next: u64,
}

impl IdGenerator {
    pub fn new() -> Self {
        Self { next: 0 }
    }
}

impl Default for IdGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl Iterator for IdGenerator {
    type Item = RequestId;
    fn next(&mut self) -> Option<Self::Item> {
        self.next += 1;
        Some(RequestId(format!("A{:04}", self.next % 10_000)))
    }
}
<commit_msg>Use new-style pin projections<commit_after>use std::io;
use std::net::ToSocketAddrs;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use futures::{ready, Sink, Stream, StreamExt};
use pin_project::pin_project;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;
use tokio_rustls::rustls::ClientConfig;
use tokio_rustls::webpki::DNSNameRef;
use tokio_rustls::{client::TlsStream, TlsConnector};
use tokio_util::codec::{Decoder, Framed};

use crate::codec::{ImapCodec, ResponseData};
use imap_proto::builders::command::Command;
use imap_proto::{Request, RequestId, State};

pub type TlsClient = Client<TlsStream<TcpStream>>;

pub struct Client<T> {
    transport: Framed<T, ImapCodec>,
    state: State,
    request_ids: IdGenerator,
}

impl TlsClient {
    pub async fn connect(server: &str) -> io::Result<(ResponseData, Self)> {
        let addr = (server, 993).to_socket_addrs()?.next().ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::Other,
                format!("no IP addresses found for {}", server),
            )
        })?;

        let mut tls_config = ClientConfig::new();
        tls_config
            .root_store
            .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
        let connector: TlsConnector = Arc::new(tls_config).into();

        let stream = TcpStream::connect(&addr).await?;
        let stream = connector
            .connect(DNSNameRef::try_from_ascii_str(server).unwrap(), stream)
            .await?;
        let mut transport = ImapCodec::default().framed(stream);

        let greeting = match transport.next().await {
            Some(greeting) => Ok(greeting),
            None => Err(io::Error::new(io::ErrorKind::Other, "no greeting found")),
        }?;
        let client = Client {
            transport,
            state: State::NotAuthenticated,
            request_ids: IdGenerator::new(),
        };

        greeting.map(|greeting| (greeting, client))
    }

    pub fn call<C: Into<Command>>(&mut self, cmd: C) -> ResponseStream<TlsStream<TcpStream>> {
        let request_id = self.request_ids.next().unwrap(); // safe: never returns Err,
        ResponseStream {
            client: self,
            request_id,
            cmd: cmd.into(),
            state: ResponseStreamState::Start,
        }
    }
}

#[pin_project]
pub struct ResponseStream<'a, T> {
    #[pin]
    client: &'a mut Client<T>,
    request_id: RequestId,
    cmd: Command,
    state: ResponseStreamState,
}

impl<'a, T> Stream for ResponseStream<'a, T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    type Item = Result<ResponseData, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let mut me = self.project();
        loop {
            match me.state {
                ResponseStreamState::Start => {
                    ready!(Pin::new(&mut me.client.transport).poll_ready(cx))?;
                    let pinned = Pin::new(&mut me.client.transport);
                    pinned.start_send(&Request(me.request_id.as_bytes(), &me.cmd.args))?;
                    *me.state = ResponseStreamState::Sending;
                }
                ResponseStreamState::Sending => {
                    let pinned = Pin::new(&mut me.client.transport);
                    ready!(pinned.poll_flush(cx))?;
                    *me.state = ResponseStreamState::Receiving;
                }
                ResponseStreamState::Receiving => {
                    match ready!(Pin::new(&mut me.client.transport).poll_next(cx)) {
                        Some(Ok(rsp)) => {
                            match rsp.request_id() {
                                Some(req_id) if req_id == me.request_id => {}
                                Some(_) | None => return Poll::Ready(Some(Ok(rsp))),
                            }

                            if let Some(next_state) = me.cmd.next_state.as_ref() {
                                me.client.state = *next_state;
                            }
                            *me.state = ResponseStreamState::Done;
                            return Poll::Ready(Some(Ok(rsp)));
                        }
                        Some(Err(e)) => return Poll::Ready(Some(Err(e))),
                        None => {
                            return Poll::Ready(Some(Err(io::Error::new(
                                io::ErrorKind::Other,
                                "stream ended before command completion",
                            ))))
                        }
                    }
                }
                ResponseStreamState::Done => {
                    return Poll::Ready(None);
                }
            }
        }
    }
}

enum ResponseStreamState {
    Start,
    Sending,
    Receiving,
    Done,
}

pub struct IdGenerator {
    next: u64,
}

impl IdGenerator {
    pub fn new() -> Self {
        Self { next: 0 }
    }
}

impl Default for IdGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl Iterator for IdGenerator {
    type Item = RequestId;
    fn next(&mut self) -> Option<Self::Item> {
        self.next += 1;
        Some(RequestId(format!("A{:04}", self.next % 10_000)))
    }
}
<|endoftext|>


<commit_before>//! Descriptor sets and layouts.
//!
//! A [`Descriptor`] is an object that describes the connection between a resource, such as
//! an `Image` or `Buffer`, and a variable in a shader. Descriptors are organized into
//! `DescriptorSet`s, each of which contains multiple descriptors that are bound and unbound to
//! shaders as a single unit. The contents of each descriptor in a set is defined by a
//! `DescriptorSetLayout` which is in turn built of [`DescriptorSetLayoutBinding`]s. A `DescriptorSet`
//! is then allocated from a [`DescriptorPool`] using the `DescriptorSetLayout`, and specific [`Descriptor`]s are
//! then bound to each binding point in the set using a [`DescriptorSetWrite`] and/or [`DescriptorSetCopy`].
//! Each descriptor set may contain descriptors to multiple different sorts of resources, and a shader may
//! use multiple descriptor sets at a time.
//!
//! [`Descriptor`]: enum.Descriptor.html
//! [`DescriptorSetLayoutBinding`]: struct.DescriptorSetLayoutBinding.html
//! [`DescriptorPool`]: trait.DescriptorPool.html
//! [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
//! [`DescriptorSetCopy`]: struct.DescriptorSetWrite.html

use smallvec::SmallVec;
use std::{borrow::Borrow, fmt, iter, ops::Range};

use crate::{
    buffer::Offset,
    image::Layout,
    pso::ShaderStageFlags,
    Backend,
};

///
pub type DescriptorSetIndex = u16;
///
pub type DescriptorBinding = u32;
///
pub type DescriptorArrayIndex = usize;

/// DOC TODO: Grasping and remembering the differences between these
///       types is a tough task. We might be able to come up with better names?
///       Or even use tuples to describe functionality instead of coming up with fancy names.
#[repr(C)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum DescriptorType {
    /// Controls filtering parameters for sampling from images.
    Sampler = 0,
    ///
    CombinedImageSampler = 1,
    /// Allows sampling (filtered loading) from associated image memory.
    /// Usually combined with a `Sampler`.
    SampledImage = 2,
    /// Allows atomic operations, (non-filtered) loads and stores on image memory.
    StorageImage = 3,
    /// Read-only, formatted buffer.
    UniformTexelBuffer = 4,
    /// Read-Write, formatted buffer.
    StorageTexelBuffer = 5,
    /// Read-only, structured buffer.
    UniformBuffer = 6,
    /// Read-Write, structured buffer.
    StorageBuffer = 7,
    /// A uniform buffer that can be bound with an offset into its memory with minimal performance impact,
    /// usually used to store pieces of "uniform" data that change per draw call rather than
    /// per render pass.
    UniformBufferDynamic = 8,
    ///
    StorageBufferDynamic = 9,
    /// Allows unfiltered loads of pixel local data in the fragment shader.
    InputAttachment = 10,
}

/// Information about the contents of and in which stages descriptors may be bound to a descriptor
/// set at a certain binding point. Multiple `DescriptorSetLayoutBinding`s are assembled into
/// a `DescriptorSetLayout`, which is then allocated into a `DescriptorSet` using a
/// [`DescriptorPool`].
///
/// A descriptor set consists of multiple binding points.
/// Each binding point contains one or multiple descriptors of a certain type.
/// The binding point is only valid for the pipelines stages specified.
///
/// The binding _must_ match with the corresponding shader interface.
///
/// [`DescriptorPool`]: trait.DescriptorPool.html
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct DescriptorSetLayoutBinding {
    /// Descriptor bindings range.
    pub binding: DescriptorBinding,
    /// Type of the bound descriptors.
    pub ty: DescriptorType,
    /// Number of descriptors in the array.
    ///
    /// *Note*: If count is zero, the binding point is reserved
    /// and can't be accessed from any shader stages.
    pub count: DescriptorArrayIndex,
    /// Valid shader stages.
    pub stage_flags: ShaderStageFlags,
    /// Use the associated list of immutable samplers.
    pub immutable_samplers: bool,
}

/// Set of descriptors of a specific type.
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct DescriptorRangeDesc {
    /// Type of the stored descriptors.
    pub ty: DescriptorType,
    /// Amount of space.
    pub count: usize,
}

/// An error allocating descriptor sets from a pool.
#[derive(Clone, Debug, PartialEq)]
pub enum AllocationError {
    /// Memory allocation on the host side failed.
    /// This could be caused by a lack of memory or pool fragmentation.
    Host,
    /// Memory allocation on the host side failed.
    /// This could be caused by a lack of memory or pool fragmentation.
    Device,
    /// Memory allocation failed as there is not enough in the pool.
    /// This could be caused by too many descriptor sets being created.
    OutOfPoolMemory,
    /// Memory allocation failed due to pool fragmentation.
    FragmentedPool,
    /// Descriptor set allocation failed as the layout is incompatible with the pool.
    IncompatibleLayout,
}

impl std::fmt::Display for AllocationError {
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AllocationError::Host => write!(fmt, "Failed to allocate descriptor set: Out of host memory"),
            AllocationError::Device => write!(fmt, "Failed to allocate descriptor set: Out of device memory"),
            AllocationError::OutOfPoolMemory => write!(fmt, "Failed to allocate descriptor set: Out of pool memory"),
            AllocationError::FragmentedPool => write!(fmt, "Failed to allocate descriptor set: Pool is fragmented"),
            AllocationError::IncompatibleLayout => write!(fmt, "Failed to allocate descriptor set: Incompatible layout"),
        }
    }
}

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

/// A descriptor pool is a collection of memory from which descriptor sets are allocated.
pub trait DescriptorPool<B: Backend>: Send + Sync + fmt::Debug {
    /// Allocate a descriptor set from the pool.
    ///
    /// The descriptor set will be allocated from the pool according to the corresponding set layout. However,
    /// specific descriptors must still be written to the set before use using a [`DescriptorSetWrite`] or
    /// [`DescriptorSetCopy`].
    ///
    /// Descriptors will become invalid once the pool is reset. Usage of invalidated descriptor sets results
    /// in undefined behavior.
    ///
    /// [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
    /// [`DescriptorSetCopy`]: struct.DescriptorSetCopy.html
    unsafe fn allocate_set(
        &mut self,
        layout: &B::DescriptorSetLayout,
    ) -> Result<B::DescriptorSet, AllocationError> {
        let mut sets = SmallVec::new();
        self.allocate_sets(iter::once(layout), &mut sets)
            .map(|_| sets.remove(0))
    }

    /// Allocate one or multiple descriptor sets from the pool.
    ///
    /// The descriptor set will be allocated from the pool according to the corresponding set layout. However,
    /// specific descriptors must still be written to the set before use using a [`DescriptorSetWrite`] or
    /// [`DescriptorSetCopy`].
    ///
    /// Each descriptor set will be allocated from the pool according to the corresponding set layout.
    /// Descriptors will become invalid once the pool is reset. Usage of invalidated descriptor sets results
    /// in undefined behavior.
    ///
    /// [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
    /// [`DescriptorSetCopy`]: struct.DescriptorSetCopy.html
    unsafe fn allocate_sets<I>(
        &mut self,
        layouts: I,
        sets: &mut SmallVec<[B::DescriptorSet; 1]>,
    ) -> Result<(), AllocationError>
    where
        I: IntoIterator,
        I::Item: Borrow<B::DescriptorSetLayout>,
    {
        let base = sets.len();
        for layout in layouts {
            match self.allocate_set(layout.borrow()) {
                Ok(set) => sets.push(set),
                Err(e) => {
                    while sets.len() != base {
                        self.free_sets(sets.pop());
                    }
                    return Err(e);
                }
            }
        }
        Ok(())
    }

    /// Free the given descriptor sets provided as an iterator.
    unsafe fn free_sets<I>(&mut self, descriptor_sets: I)
    where
        I: IntoIterator<Item = B::DescriptorSet>;

    /// Resets a descriptor pool, releasing all resources from all the descriptor sets
    /// allocated from it and freeing the descriptor sets. Invalidates all descriptor
    /// sets allocated from the pool; trying to use one after the pool has been reset
    /// is undefined behavior.
    unsafe fn reset(&mut self);
}

/// Writes the actual descriptors to be bound into a descriptor set. Should be provided
/// to the `write_descriptor_sets` method of a `Device`.
#[allow(missing_docs)]
#[derive(Debug)]
pub struct DescriptorSetWrite<'a, B: Backend, WI>
where
    WI: IntoIterator,
    WI::Item: Borrow<Descriptor<'a, B>>,
{
    pub set: &'a B::DescriptorSet,
    /// *Note*: when there is more descriptors provided than
    /// array elements left in the specified binding starting
    /// at specified, offset, the updates are spilled onto
    /// the next binding (starting with offset 0), and so on.
    pub binding: DescriptorBinding,
    pub array_offset: DescriptorArrayIndex,
    pub descriptors: WI,
}

/// A handle to a specific shader resource that can be bound for use in a `DescriptorSet`.
/// Usually provided in a [`DescriptorSetWrite`]
///
/// [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
#[allow(missing_docs)]
#[derive(Clone, Debug)]
pub enum Descriptor<'a, B: Backend> {
    Sampler(&'a B::Sampler),
    Image(&'a B::ImageView, Layout),
    CombinedImageSampler(&'a B::ImageView, Layout, &'a B::Sampler),
    Buffer(&'a B::Buffer, Range<Option<Offset>>),
    UniformTexelBuffer(&'a B::BufferView),
    StorageTexelBuffer(&'a B::BufferView),
}

/// Copies a range of descriptors to be bound from one descriptor set to another Should be
/// provided to the `copy_descriptor_sets` method of a `Device`.
#[allow(missing_docs)]
#[derive(Clone, Copy, Debug)]
pub struct DescriptorSetCopy<'a, B: Backend> {
    pub src_set: &'a B::DescriptorSet,
    pub src_binding: DescriptorBinding,
    pub src_array_offset: DescriptorArrayIndex,
    pub dst_set: &'a B::DescriptorSet,
    pub dst_binding: DescriptorBinding,
    pub dst_array_offset: DescriptorArrayIndex,
    pub count: usize,
}

bitflags! {
    /// Descriptor pool creation flags.
    pub struct DescriptorPoolCreateFlags: u32 {
        /// Specifies that descriptor sets are allowed to be freed from the pool
        /// individually.
        const FREE_DESCRIPTOR_SET = 0x1;
    }
}
<commit_msg>Add RichDescriptorType, subtypes, and conversion<commit_after>//! Descriptor sets and layouts.
//!
//! A [`Descriptor`] is an object that describes the connection between a resource, such as
//! an `Image` or `Buffer`, and a variable in a shader. Descriptors are organized into
//! `DescriptorSet`s, each of which contains multiple descriptors that are bound and unbound to
//! shaders as a single unit. The contents of each descriptor in a set is defined by a
//! `DescriptorSetLayout` which is in turn built of [`DescriptorSetLayoutBinding`]s. A `DescriptorSet`
//! is then allocated from a [`DescriptorPool`] using the `DescriptorSetLayout`, and specific [`Descriptor`]s are
//! then bound to each binding point in the set using a [`DescriptorSetWrite`] and/or [`DescriptorSetCopy`].
//! Each descriptor set may contain descriptors to multiple different sorts of resources, and a shader may
//! use multiple descriptor sets at a time.
//!
//! [`Descriptor`]: enum.Descriptor.html
//! [`DescriptorSetLayoutBinding`]: struct.DescriptorSetLayoutBinding.html
//! [`DescriptorPool`]: trait.DescriptorPool.html
//! [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
//! [`DescriptorSetCopy`]: struct.DescriptorSetWrite.html

use smallvec::SmallVec;
use std::{borrow::Borrow, fmt, iter, ops::Range};

use crate::{
    buffer::Offset,
    image::Layout,
    pso::ShaderStageFlags,
    Backend,
};

///
pub type DescriptorSetIndex = u16;
///
pub type DescriptorBinding = u32;
///
pub type DescriptorArrayIndex = usize;

/// DOC TODO: Grasping and remembering the differences between these
///       types is a tough task. We might be able to come up with better names?
///       Or even use tuples to describe functionality instead of coming up with fancy names.
#[repr(C)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum DescriptorType {
    /// Controls filtering parameters for sampling from images.
    Sampler = 0,
    ///
    CombinedImageSampler = 1,
    /// Allows sampling (filtered loading) from associated image memory.
    /// Usually combined with a `Sampler`.
    SampledImage = 2,
    /// Allows atomic operations, (non-filtered) loads and stores on image memory.
    StorageImage = 3,
    /// Read-only, formatted buffer.
    UniformTexelBuffer = 4,
    /// Read-Write, formatted buffer.
    StorageTexelBuffer = 5,
    /// Read-only, structured buffer.
    UniformBuffer = 6,
    /// Read-Write, structured buffer.
    StorageBuffer = 7,
    /// A uniform buffer that can be bound with an offset into its memory with minimal performance impact,
    /// usually used to store pieces of "uniform" data that change per draw call rather than
    /// per render pass.
    UniformBufferDynamic = 8,
    ///
    StorageBufferDynamic = 9,
    /// Allows unfiltered loads of pixel local data in the fragment shader.
    InputAttachment = 10,
}

impl std::convert::From<RichDescriptorType> for DescriptorType {
    fn from(ty: RichDescriptorType) -> Self {
        match ty {
            RichDescriptorType::Sampler => DescriptorType::Sampler,
            RichDescriptorType::CombinedImageSampler => DescriptorType::CombinedImageSampler,
            RichDescriptorType::Image { ty } => match ty {
                ImageDescriptorType::Storage => DescriptorType::StorageImage,
                ImageDescriptorType::Sampled => DescriptorType::SampledImage,
            }
            RichDescriptorType::Buffer { access, format } => match access {
                BufferDescriptorAccess::Storage => match format {
                    BufferDescriptorFormat::Dynamic => DescriptorType::StorageBufferDynamic,
                    BufferDescriptorFormat::Structured => DescriptorType::StorageBuffer,
                    BufferDescriptorFormat::Texel => DescriptorType::StorageTexelBuffer,
                },
                BufferDescriptorAccess::Uniform => match format {
                    BufferDescriptorFormat::Dynamic => DescriptorType::UniformBufferDynamic,
                    BufferDescriptorFormat::Structured => DescriptorType::UniformBuffer,
                    BufferDescriptorFormat::Texel => DescriptorType::UniformTexelBuffer,
                },
            }
            RichDescriptorType::InputAttachment => DescriptorType::InputAttachment,
        }
    }
}

/// Access type of a buffer.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BufferDescriptorAccess {
    /// Storage buffers allow load, store, and atomic operations.
    Storage,
    /// Uniform buffers allow only load operations.
    Uniform,
}

/// Format of a buffer.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BufferDescriptorFormat {
    // TODO: seems to have no effect in hal, ask about removing or documenting that it's a no-op
    /// TODO
    Dynamic,
    /// The buffer is interpreted as a structure defined in a shader.
    Structured,
    /// The buffer is interpreted as a 1-D array of texels, which undergo format
    /// conversion when loaded in a shader.
    Texel,
}

/// Specific type of an image descriptor.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum ImageDescriptorType {
    /// A storage image allows load, store and atomic operations.
    Storage,

    /// A sampled image allows sampling operations.
    Sampled,
}

/// The type of a descriptor. TODO: more specific doc
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum RichDescriptorType {
    /// A descriptor associated with sampler.
    Sampler,
    /// A single descriptor associated with both a sampler and an image to be
    /// used together.
    CombinedImageSampler,
    /// A descriptor associated with an image.
    Image {
        /// The specific type of this image descriptor.
        ty: ImageDescriptorType,
    },
    /// A descriptor associated with a buffer.
    Buffer {
        /// The access type of this buffer descriptor.
        access: BufferDescriptorAccess,
        /// The format of this buffer descriptor.
        format: BufferDescriptorFormat,
    },
    /// A descriptor associated with an input attachment.
    InputAttachment,
}

/// Information about the contents of and in which stages descriptors may be bound to a descriptor
/// set at a certain binding point. Multiple `DescriptorSetLayoutBinding`s are assembled into
/// a `DescriptorSetLayout`, which is then allocated into a `DescriptorSet` using a
/// [`DescriptorPool`].
///
/// A descriptor set consists of multiple binding points.
/// Each binding point contains one or multiple descriptors of a certain type.
/// The binding point is only valid for the pipelines stages specified.
///
/// The binding _must_ match with the corresponding shader interface.
///
/// [`DescriptorPool`]: trait.DescriptorPool.html
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct DescriptorSetLayoutBinding {
    /// Descriptor bindings range.
    pub binding: DescriptorBinding,
    /// Type of the bound descriptors.
    pub ty: DescriptorType,
    /// Number of descriptors in the array.
    ///
    /// *Note*: If count is zero, the binding point is reserved
    /// and can't be accessed from any shader stages.
    pub count: DescriptorArrayIndex,
    /// Valid shader stages.
    pub stage_flags: ShaderStageFlags,
    /// Use the associated list of immutable samplers.
    pub immutable_samplers: bool,
}

/// Set of descriptors of a specific type.
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct DescriptorRangeDesc {
    /// Type of the stored descriptors.
    pub ty: DescriptorType,
    /// Amount of space.
    pub count: usize,
}

/// An error allocating descriptor sets from a pool.
#[derive(Clone, Debug, PartialEq)]
pub enum AllocationError {
    /// Memory allocation on the host side failed.
    /// This could be caused by a lack of memory or pool fragmentation.
    Host,
    /// Memory allocation on the host side failed.
    /// This could be caused by a lack of memory or pool fragmentation.
    Device,
    /// Memory allocation failed as there is not enough in the pool.
    /// This could be caused by too many descriptor sets being created.
    OutOfPoolMemory,
    /// Memory allocation failed due to pool fragmentation.
    FragmentedPool,
    /// Descriptor set allocation failed as the layout is incompatible with the pool.
    IncompatibleLayout,
}

impl std::fmt::Display for AllocationError {
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AllocationError::Host => write!(fmt, "Failed to allocate descriptor set: Out of host memory"),
            AllocationError::Device => write!(fmt, "Failed to allocate descriptor set: Out of device memory"),
            AllocationError::OutOfPoolMemory => write!(fmt, "Failed to allocate descriptor set: Out of pool memory"),
            AllocationError::FragmentedPool => write!(fmt, "Failed to allocate descriptor set: Pool is fragmented"),
            AllocationError::IncompatibleLayout => write!(fmt, "Failed to allocate descriptor set: Incompatible layout"),
        }
    }
}

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

/// A descriptor pool is a collection of memory from which descriptor sets are allocated.
pub trait DescriptorPool<B: Backend>: Send + Sync + fmt::Debug {
    /// Allocate a descriptor set from the pool.
    ///
    /// The descriptor set will be allocated from the pool according to the corresponding set layout. However,
    /// specific descriptors must still be written to the set before use using a [`DescriptorSetWrite`] or
    /// [`DescriptorSetCopy`].
    ///
    /// Descriptors will become invalid once the pool is reset. Usage of invalidated descriptor sets results
    /// in undefined behavior.
    ///
    /// [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
    /// [`DescriptorSetCopy`]: struct.DescriptorSetCopy.html
    unsafe fn allocate_set(
        &mut self,
        layout: &B::DescriptorSetLayout,
    ) -> Result<B::DescriptorSet, AllocationError> {
        let mut sets = SmallVec::new();
        self.allocate_sets(iter::once(layout), &mut sets)
            .map(|_| sets.remove(0))
    }

    /// Allocate one or multiple descriptor sets from the pool.
    ///
    /// The descriptor set will be allocated from the pool according to the corresponding set layout. However,
    /// specific descriptors must still be written to the set before use using a [`DescriptorSetWrite`] or
    /// [`DescriptorSetCopy`].
    ///
    /// Each descriptor set will be allocated from the pool according to the corresponding set layout.
    /// Descriptors will become invalid once the pool is reset. Usage of invalidated descriptor sets results
    /// in undefined behavior.
    ///
    /// [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
    /// [`DescriptorSetCopy`]: struct.DescriptorSetCopy.html
    unsafe fn allocate_sets<I>(
        &mut self,
        layouts: I,
        sets: &mut SmallVec<[B::DescriptorSet; 1]>,
    ) -> Result<(), AllocationError>
    where
        I: IntoIterator,
        I::Item: Borrow<B::DescriptorSetLayout>,
    {
        let base = sets.len();
        for layout in layouts {
            match self.allocate_set(layout.borrow()) {
                Ok(set) => sets.push(set),
                Err(e) => {
                    while sets.len() != base {
                        self.free_sets(sets.pop());
                    }
                    return Err(e);
                }
            }
        }
        Ok(())
    }

    /// Free the given descriptor sets provided as an iterator.
    unsafe fn free_sets<I>(&mut self, descriptor_sets: I)
    where
        I: IntoIterator<Item = B::DescriptorSet>;

    /// Resets a descriptor pool, releasing all resources from all the descriptor sets
    /// allocated from it and freeing the descriptor sets. Invalidates all descriptor
    /// sets allocated from the pool; trying to use one after the pool has been reset
    /// is undefined behavior.
    unsafe fn reset(&mut self);
}

/// Writes the actual descriptors to be bound into a descriptor set. Should be provided
/// to the `write_descriptor_sets` method of a `Device`.
#[allow(missing_docs)]
#[derive(Debug)]
pub struct DescriptorSetWrite<'a, B: Backend, WI>
where
    WI: IntoIterator,
    WI::Item: Borrow<Descriptor<'a, B>>,
{
    pub set: &'a B::DescriptorSet,
    /// *Note*: when there is more descriptors provided than
    /// array elements left in the specified binding starting
    /// at specified, offset, the updates are spilled onto
    /// the next binding (starting with offset 0), and so on.
    pub binding: DescriptorBinding,
    pub array_offset: DescriptorArrayIndex,
    pub descriptors: WI,
}

/// A handle to a specific shader resource that can be bound for use in a `DescriptorSet`.
/// Usually provided in a [`DescriptorSetWrite`]
///
/// [`DescriptorSetWrite`]: struct.DescriptorSetWrite.html
#[allow(missing_docs)]
#[derive(Clone, Debug)]
pub enum Descriptor<'a, B: Backend> {
    Sampler(&'a B::Sampler),
    Image(&'a B::ImageView, Layout),
    CombinedImageSampler(&'a B::ImageView, Layout, &'a B::Sampler),
    Buffer(&'a B::Buffer, Range<Option<Offset>>),
    UniformTexelBuffer(&'a B::BufferView),
    StorageTexelBuffer(&'a B::BufferView),
}

/// Copies a range of descriptors to be bound from one descriptor set to another Should be
/// provided to the `copy_descriptor_sets` method of a `Device`.
#[allow(missing_docs)]
#[derive(Clone, Copy, Debug)]
pub struct DescriptorSetCopy<'a, B: Backend> {
    pub src_set: &'a B::DescriptorSet,
    pub src_binding: DescriptorBinding,
    pub src_array_offset: DescriptorArrayIndex,
    pub dst_set: &'a B::DescriptorSet,
    pub dst_binding: DescriptorBinding,
    pub dst_array_offset: DescriptorArrayIndex,
    pub count: usize,
}

bitflags! {
    /// Descriptor pool creation flags.
    pub struct DescriptorPoolCreateFlags: u32 {
        /// Specifies that descriptor sets are allowed to be freed from the pool
        /// individually.
        const FREE_DESCRIPTOR_SET = 0x1;
    }
}
<|endoftext|>


<commit_before>use helpers::{HelperDef};
use registry::{Registry};
use context::{Context};
use render::{Renderable, RenderContext, RenderError, Helper};

#[derive(Clone, Copy)]
pub struct IncludeHelper;

#[derive(Clone, Copy)]
pub struct BlockHelper;

#[derive(Clone, Copy)]
pub struct PartialHelper;

impl HelperDef for IncludeHelper {
    fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
        let template = match h.params().get(0) {
            Some(ref t) => r.get_template(t),
            None => return Err(RenderError::new("Param not found for helper")),
        };

        let context_param = h.params().get(1);
        let old_path = match context_param {
            Some(p) => {
                let old_path = rc.get_path().clone();
                rc.promote_local_vars();
                let new_path = format!("{}/{}", old_path, p);
                rc.set_path(new_path);
                Some(old_path)
            },
            None => None,
        };

        let result = match template {
            Some(t) => {
                if h.hash().is_empty() {
                    t.render(c, r, rc)
                } else {
                    let new_ctx = c.extend(h.hash());
                    t.render(&new_ctx, r, rc)
                }
            },
            None => Err(RenderError::new("Template not found.")),
        };

        if let Some(path) = old_path {
            rc.set_path(path);
            rc.demote_local_vars();
        }

        result
    }
}

impl HelperDef for BlockHelper {
    fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
        let param = h.params().get(0);

        if param.is_none() {
            return Err(RenderError::new("Param not found for helper"));
        }

        let partial_template = rc.get_partial(param.unwrap());

        match partial_template {
            Some(partial_template) => {
                partial_template.render(c, r, rc)
            },
            None => {
                h.template().unwrap().render(c, r, rc)
            }
        }
    }
}

impl HelperDef for PartialHelper {
    fn call(&self, _: &Context, h: &Helper, _: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
        let param = h.params().get(0);

        if param.is_none() {
            return Err(RenderError::new("Param not found for helper"));
        }

        rc.set_partial(param.unwrap().clone(), h.template().unwrap().clone());
        Ok(())
    }
}

pub static INCLUDE_HELPER: IncludeHelper = IncludeHelper;
pub static BLOCK_HELPER: BlockHelper = BlockHelper;
pub static PARTIAL_HELPER: PartialHelper = PartialHelper;

#[cfg(test)]
mod test {
    use template::{Template};
    use registry::{Registry};
    use std::collections::BTreeMap;

    #[test]
    fn test() {
        let t0 = Template::compile("
{{#block title}}default{{/block}}
".to_string()).ok().unwrap();
        let t1 = Template::compile("{{#partial title}}{{this}}{{/partial}}{{> t0}}".to_string()).ok().unwrap();
        let t2 = Template::compile("{{> t0}}
{{this}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);
        handlebars.register_template("t1", t1);
        handlebars.register_template("t2", t2);

        let r0 = handlebars.render("t1", &true);
        assert_eq!(r0.ok().unwrap(), "
true
".to_string());

        let r1 = handlebars.render("t2", &true);
        assert_eq!(r1.ok().unwrap(), "
default

true
".to_string());
    }

    #[test]
    fn test_context() {
        let t0 = Template::compile("
{{> (body) data}}
".to_string()).ok().unwrap();
        let t1 = Template::compile("
{{this}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);
        handlebars.register_template("t1", t1);

        let mut map: BTreeMap<String, String> = BTreeMap::new();
        map.insert("body".into(), "t1".into());
        map.insert("data".into(), "hello".into());

        let r0 = handlebars.render("t0", &map);
        assert_eq!(r0.ok().unwrap(), "

hello

".to_string());
    }

    #[test]
    fn test_partial_hash_context() {
        let t0 = Template::compile("
{{> t1 hello=world}}
".to_string()).ok().unwrap();
        let t1 = Template::compile("
{{data}}

{{hello}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);
        handlebars.register_template("t1", t1);

        let mut map: BTreeMap<String, String> = BTreeMap::new();
        map.insert("data".into(), "hello".into());

        let r0 = handlebars.render("t0", &map);
        assert_eq!(r0.ok().unwrap(), "

hello

world

".to_string());
    }

    #[test]
    fn test_inline_partial() {
        let t0 = Template::compile("{{#partial title}}hello {{name}}{{/partial}}
include partial: {{#block title}}{{/block}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);

        let mut map: BTreeMap<String, String> = BTreeMap::new();
        map.insert("name".into(), "world".into());

        let r0 = handlebars.render("t0", &map);
        assert_eq!(r0.ok().unwrap(), "
include partial: hello world
".to_string());
    }
}
<commit_msg>(feat) check current_template to avoid stack overflow [#44]<commit_after>use helpers::{HelperDef};
use registry::{Registry};
use context::{Context};
use render::{Renderable, RenderContext, RenderError, Helper};

#[derive(Clone, Copy)]
pub struct IncludeHelper;

#[derive(Clone, Copy)]
pub struct BlockHelper;

#[derive(Clone, Copy)]
pub struct PartialHelper;

impl HelperDef for IncludeHelper {
    fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
        let template = match h.params().get(0) {
            Some(ref t) => {
                if rc.current_template == Some((*t).to_string()) {
                    return Err(RenderError::new("Cannot include self in >"));
                } else {
                    r.get_template(t)
                }
            }
            None => return Err(RenderError::new("Param not found for helper")),
        };

        let context_param = h.params().get(1);
        let old_path = match context_param {
            Some(p) => {
                let old_path = rc.get_path().clone();
                rc.promote_local_vars();
                let new_path = format!("{}/{}", old_path, p);
                rc.set_path(new_path);
                Some(old_path)
            },
            None => None,
        };

        let result = match template {
            Some(t) => {
                if h.hash().is_empty() {
                    t.render(c, r, rc)
                } else {
                    let new_ctx = c.extend(h.hash());
                    t.render(&new_ctx, r, rc)
                }
            },
            None => Err(RenderError::new("Template not found.")),
        };

        if let Some(path) = old_path {
            rc.set_path(path);
            rc.demote_local_vars();
        }

        result
    }
}

impl HelperDef for BlockHelper {
    fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
        let param = h.params().get(0);

        if param.is_none() {
            return Err(RenderError::new("Param not found for helper"));
        }

        let partial_template = rc.get_partial(param.unwrap());

        match partial_template {
            Some(partial_template) => {
                partial_template.render(c, r, rc)
            },
            None => {
                h.template().unwrap().render(c, r, rc)
            }
        }
    }
}

impl HelperDef for PartialHelper {
    fn call(&self, _: &Context, h: &Helper, _: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
        let param = h.params().get(0);

        if param.is_none() {
            return Err(RenderError::new("Param not found for helper"));
        }

        rc.set_partial(param.unwrap().clone(), h.template().unwrap().clone());
        Ok(())
    }
}

pub static INCLUDE_HELPER: IncludeHelper = IncludeHelper;
pub static BLOCK_HELPER: BlockHelper = BlockHelper;
pub static PARTIAL_HELPER: PartialHelper = PartialHelper;

#[cfg(test)]
mod test {
    use template::{Template};
    use registry::{Registry};
    use std::collections::BTreeMap;

    #[test]
    fn test() {
        let t0 = Template::compile("
{{#block title}}default{{/block}}
".to_string()).ok().unwrap();
        let t1 = Template::compile("{{#partial title}}{{this}}{{/partial}}{{> t0}}".to_string()).ok().unwrap();
        let t2 = Template::compile("{{> t0}}
{{this}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);
        handlebars.register_template("t1", t1);
        handlebars.register_template("t2", t2);

        let r0 = handlebars.render("t1", &true);
        assert_eq!(r0.ok().unwrap(), "
true
".to_string());

        let r1 = handlebars.render("t2", &true);
        assert_eq!(r1.ok().unwrap(), "
default

true
".to_string());
    }

    #[test]
    fn test_context() {
        let t0 = Template::compile("
{{> (body) data}}
".to_string()).ok().unwrap();
        let t1 = Template::compile("
{{this}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);
        handlebars.register_template("t1", t1);

        let mut map: BTreeMap<String, String> = BTreeMap::new();
        map.insert("body".into(), "t1".into());
        map.insert("data".into(), "hello".into());

        let r0 = handlebars.render("t0", &map);
        assert_eq!(r0.ok().unwrap(), "

hello

".to_string());
    }

    #[test]
    fn test_partial_hash_context() {
        let t0 = Template::compile("
{{> t1 hello=world}}
".to_string()).ok().unwrap();
        let t1 = Template::compile("
{{data}}

{{hello}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);
        handlebars.register_template("t1", t1);

        let mut map: BTreeMap<String, String> = BTreeMap::new();
        map.insert("data".into(), "hello".into());

        let r0 = handlebars.render("t0", &map);
        assert_eq!(r0.ok().unwrap(), "

hello

world

".to_string());
    }

    #[test]
    fn test_inline_partial() {
        let t0 = Template::compile("{{#partial title}}hello {{name}}{{/partial}}
include partial: {{#block title}}{{/block}}
".to_string()).ok().unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);

        let mut map: BTreeMap<String, String> = BTreeMap::new();
        map.insert("name".into(), "world".into());

        let r0 = handlebars.render("t0", &map);
        assert_eq!(r0.ok().unwrap(), "
include partial: hello world
".to_string());
    }

    #[test]
    fn test_include_self() {
        let t0 = Template::compile("
{{> t0}}
".to_string()).ok().unwrap();
        let mut handlebars = Registry::new();
        handlebars.register_template("t0", t0);

        let map: BTreeMap<String, String> = BTreeMap::new();

        let r0 = handlebars.render("t0", &map);
        assert!(r0.is_err());
    }
}
<|endoftext|>


<commit_before>use criterion::{criterion_group, criterion_main, Criterion};
use ipnetwork::{Ipv4Network, Ipv6Network};
use std::net::{Ipv4Addr, Ipv6Addr};

fn parse_ipv4_benchmark(c: &mut Criterion) {
    c.bench_function("parse ipv4", |b| {
        b.iter(|| "127.1.0.0/24".parse::<Ipv4Network>().unwrap())
    });
}

fn parse_ipv6_benchmark(c: &mut Criterion) {
    c.bench_function("parse ipv6", |b| {
        b.iter(|| "FF01:0:0:17:0:0:0:2/64".parse::<Ipv6Network>().unwrap())
    });
}

fn contains_ipv4_benchmark(c: &mut Criterion) {
    c.bench_function("contains ipv4", |b| {
        b.iter(|| {
            let cidr = "74.125.227.0/25".parse::<Ipv4Network>().unwrap();
            cidr.contains(Ipv4Addr::new(74, 125, 227, 4))
        })
    });
}

fn contains_ipv6_benchmark(c: &mut Criterion) {
    c.bench_function("contains ipv6", |b| {
        b.iter(|| {
            let cidr = "FF01:0:0:17:0:0:0:2/65".parse::<Ipv6Network>().unwrap();
            cidr.contains(Ipv6Addr::new(0xff01, 0, 0, 0x17, 0x7fff, 0, 0, 0x2))
        })
    });
}

criterion_group!(
    benches,
    parse_ipv4_benchmark,
    parse_ipv6_benchmark,
    contains_ipv4_benchmark,
    contains_ipv6_benchmark
);
criterion_main!(benches);
<commit_msg>Add a benchmark for parsing IPv4 netmask<commit_after>use criterion::{criterion_group, criterion_main, Criterion};
use ipnetwork::{Ipv4Network, Ipv6Network};
use std::net::{Ipv4Addr, Ipv6Addr};

fn parse_ipv4_prefix_benchmark(c: &mut Criterion) {
    c.bench_function("parse ipv4 prefix", |b| {
        b.iter(|| "127.1.0.0/24".parse::<Ipv4Network>().unwrap())
    });
}

fn parse_ipv6_benchmark(c: &mut Criterion) {
    c.bench_function("parse ipv6", |b| {
        b.iter(|| "FF01:0:0:17:0:0:0:2/64".parse::<Ipv6Network>().unwrap())
    });
}

fn parse_ipv4_netmask_benchmark(c: &mut Criterion) {
    c.bench_function("parse ipv4 netmask", |b| {
        b.iter(|| "127.1.0.0/255.255.255.0".parse::<Ipv4Network>().unwrap())
    });
}

fn contains_ipv4_benchmark(c: &mut Criterion) {
    c.bench_function("contains ipv4", |b| {
        b.iter(|| {
            let cidr = "74.125.227.0/25".parse::<Ipv4Network>().unwrap();
            cidr.contains(Ipv4Addr::new(74, 125, 227, 4))
        })
    });
}

fn contains_ipv6_benchmark(c: &mut Criterion) {
    c.bench_function("contains ipv6", |b| {
        b.iter(|| {
            let cidr = "FF01:0:0:17:0:0:0:2/65".parse::<Ipv6Network>().unwrap();
            cidr.contains(Ipv6Addr::new(0xff01, 0, 0, 0x17, 0x7fff, 0, 0, 0x2))
        })
    });
}

criterion_group!(
    benches,
    parse_ipv4_prefix_benchmark,
    parse_ipv6_benchmark,
    parse_ipv4_netmask_benchmark,
    contains_ipv4_benchmark,
    contains_ipv6_benchmark
);
criterion_main!(benches);
<|endoftext|>