<commit_before>extern crate pkg_config;
extern crate regex;

use std::process::Command;
use std::collections::HashMap;
use std::env;
use regex::Regex;
use std::fs;

struct PythonVersion {
    major: u8,
    // minor == None means any minor version will do
    minor: Option<u8>
}

const CFG_KEY: &'static str = "py_sys_config";

// A list of python interpreter compile-time preprocessor defines that
// we will pick up and pass to rustc via --cfg=py_sys_config={varname};
// this allows using them conditional cfg attributes in the .rs files, so
//
// #[cfg(py_sys_config="{varname}"]
//
// is the equivalent of #ifdef {varname} name in C.
//
// see Misc/SpecialBuilds.txt in the python source for what these mean.
//
// (hrm, this is sort of re-implementing what distutils does, except
// by passing command line args instead of referring to a python.h)
static SYSCONFIG_FLAGS: [&'static str; 7] = [
    "Py_USING_UNICODE",
    "Py_UNICODE_WIDE",
    "WITH_THREAD",
    "Py_DEBUG",
    "Py_REF_DEBUG",
    "Py_TRACE_REFS",
    "COUNT_ALLOCS",
];

static SYSCONFIG_VALUES: [&'static str; 1] = [
    // cfg doesn't support flags with values, just bools - so flags
    // below are translated into bools as {varname}_{val}
    //
    // for example, Py_UNICODE_SIZE_2 or Py_UNICODE_SIZE_4
    "Py_UNICODE_SIZE" // note - not present on python 3.3+, which is always wide
];

/// Examine python's compile flags to pass to cfg by launching
/// the interpreter and printing variables of interest from
/// sysconfig.get_config_vars.
fn get_config_vars(python_path: &String) -> Result<HashMap<String, String>, String>  {
    let mut script = "import sysconfig; \
config = sysconfig.get_config_vars();".to_owned();

    for k in SYSCONFIG_FLAGS.iter().chain(SYSCONFIG_VALUES.iter()) {
        script.push_str(&format!("print(config.get('{}', {}))", k,
            if is_value(k) { "None" } else { "0" } ));
        script.push_str(";");
    }

    let mut cmd = Command::new(python_path);
    cmd.arg("-c").arg(script);

    let out = try!(cmd.output().map_err(|e| {
        format!("failed to run python interpreter `{:?}`: {}", cmd, e)
    }));

    if !out.status.success() {
        let stderr = String::from_utf8(out.stderr).unwrap();
        let mut msg = format!("python script failed with stderr:\n\n");
        msg.push_str(&stderr);
        return Err(msg);
    }

    let stdout = String::from_utf8(out.stdout).unwrap();
    let split_stdout: Vec<&str> = stdout.trim_right().split('\n').collect();
    if split_stdout.len() != SYSCONFIG_VALUES.len() + SYSCONFIG_FLAGS.len() {
        return Err(
            format!("python stdout len didn't return expected number of lines:
{}", split_stdout.len()).to_string());
    }
    let all_vars = SYSCONFIG_FLAGS.iter().chain(SYSCONFIG_VALUES.iter());
    // let var_map: HashMap<String, String> = HashMap::new();
    Ok(all_vars.zip(split_stdout.iter())
        .fold(HashMap::new(), |mut memo: HashMap<String, String>, (&k, &v)| {
            if !(v.to_owned() == "None" && is_value(k)) {
                memo.insert(k.to_owned(), v.to_owned());
            }
            memo
        }))
}

fn is_value(key: &str) -> bool {
    SYSCONFIG_VALUES.iter().find(|x| **x == key).is_some()
}

fn cfg_line_for_var(key: &str, val: &str) -> Option<String> {
    if is_value(key) {
        // is a value; suffix the key name with the value
        Some(format!("cargo:rustc-cfg={}=\"{}_{}\"\n", CFG_KEY, key, val))
    } else if val != "0" {
        // is a flag that isn't zero
        Some(format!("cargo:rustc-cfg={}=\"{}\"", CFG_KEY, key))
    } else {
        // is a flag that is zero
        None
    }
}

/// Run a python script using the 'python' located by PATH.
fn run_python_script(script: &str) -> Result<String, String> {
    let mut cmd = Command::new("python");
    cmd.arg("-c").arg(script);

    let out = try!(cmd.output().map_err(|e| {
        format!("failed to run python interpreter `{:?}`: {}", cmd, e)
    }));

    if !out.status.success() {
        let stderr = String::from_utf8(out.stderr).unwrap();
        let mut msg = format!("python script failed with stderr:\n\n");
        msg.push_str(&stderr);
        return Err(msg);
    }

    let out = String::from_utf8(out.stdout).unwrap();
    return Ok(out);
}

#[cfg(not(target_os="macos"))]
fn get_rustc_link_lib(version: &PythonVersion, enable_shared: bool) -> Result<String, String> {
    if enable_shared {
        Ok(format!("cargo:rustc-link-lib=python{}.{}", version.major,
            version.minor))
    } else {
        Ok(format!("cargo:rustc-link-lib=static=python{}.{}", version.major,
            version.minor))
    }
}

#[cfg(target_os="macos")]
fn get_macos_linkmodel() -> Result<String, String> {
    let script = "import MacOS; print MacOS.linkmodel;";
    let out = run_python_script(script).unwrap();
    Ok(out.trim_right().to_owned())
}

#[cfg(target_os="macos")]
fn get_rustc_link_lib(version: &PythonVersion, _: bool) -> Result<String, String> {
    // os x can be linked to a framework or static or dynamic, and
    // Py_ENABLE_SHARED is wrong; framework means shared library
    let dotted_version = format!("{}.{}", version.major, version.minor.unwrap());
    match get_macos_linkmodel().unwrap().as_ref() {
        "static" => Ok(format!("cargo:rustc-link-lib=static=python{}",
            dotted_version)),
        "dynamic" => Ok(format!("cargo:rustc-link-lib=python{}",
            dotted_version)),
        "framework" => Ok(format!("cargo:rustc-link-lib=python{}",
            dotted_version)),
        other => Err(format!("unknown linkmodel {}", other))
    }
}

/// Check that the interpreter's reported version matches expected_version;
/// if it does, return the specific version the interpreter reports.
fn get_interpreter_version(line: &str, expected_version: &PythonVersion)
        -> Result<PythonVersion, String> {
    let version_re = Regex::new(r"\((\d+), (\d+)\)").unwrap();
    let interpreter_version = match version_re.captures(&line) {
        Some(cap) => PythonVersion {
            major: cap.at(1).unwrap().parse().unwrap(),
            minor: Some(cap.at(2).unwrap().parse().unwrap())
        },
        None => return Err(
            format!("Unexpected response to version query {}", line))
    };

    if interpreter_version.major != expected_version.major ||
        (expected_version.minor.is_some() &&
            interpreter_version.minor != expected_version.minor)
    {
        Err(format!("'python' is not version {}.{} (is {})",
                    expected_version.major,
                    match expected_version.minor {
                        Some(v) => v.to_string(),
                        None => "*".to_owned()
                    },
                    line))
    } else {
        Ok(interpreter_version)
    }
}

/// Deduce configuration from the 'python' in the current PATH and print
/// cargo vars to stdout.
///
/// Note that if the python doesn't satisfy expected_version, this will error.
fn configure_from_path(expected_version: &PythonVersion) -> Result<String, String> {
    let script = "import sys; import sysconfig; print(sys.version_info[0:2]); \
print(sysconfig.get_config_var('LIBDIR')); \
print(sysconfig.get_config_var('Py_ENABLE_SHARED')); \
print(sys.exec_prefix);";
    let out = run_python_script(script).unwrap();
    let lines: Vec<&str> = out.split("\n").collect();
    let version: &str = lines[0];
    let libpath: &str = lines[1];
    let enable_shared: &str = lines[2];
    let exec_prefix: &str = lines[3];

    let interpreter_version = try!(get_interpreter_version(version,
        expected_version));

    println!("{}", get_rustc_link_lib(&interpreter_version,
        enable_shared == "1").unwrap());
    println!("cargo:rustc-link-search=native={}", libpath);
    return Ok(format!("{}/bin/python", exec_prefix));
}

/// Deduce configuration from the python-X.X in pkg-config and print
/// cargo vars to stdout.
fn configure_from_pkgconfig(version: &PythonVersion, pkg_name: &str)
        -> Result<String, String> {
    // this emits relevant build info to stdout, which is picked up by the
    // build chain (funny name for something with side-effects!)
    try!(pkg_config::find_library(pkg_name));

    // This seems to be a convention - unfortunately pkg-config doesn't
    // tell you the executable name, but I've noticed at least on
    // OS X homebrew the python bin dir for 3.4 doesn't actually contain
    // a 'python'.
    let exec_prefix = pkg_config::Config::get_variable(pkg_name,
        "exec_prefix").unwrap();

    // try to find the python interpreter in the exec_prefix somewhere.
    // the .pc doesn't tell us :(

    let mut attempts = vec![
        format!("/bin/python{}", version.major),
        "/bin/python".to_owned()
    ];

    // Try to seek python(major).(minor) if the user specified a minor.
    //
    // Ideally, we'd still do this even if they didn't based off the
    // specific version of the package located by pkg_config above,
    // but it's not obvious how to reliably extract that out of the .pc.
    if version.minor.is_some() {
        attempts.insert(0, format!("/bin/python{}_{}", version.major,
            version.minor.unwrap()));
    }

    for attempt in attempts.iter() {
        let possible_exec_name = format!("{}{}", exec_prefix,
            attempt);
        match fs::metadata(&possible_exec_name) {
            Ok(_) => return Ok(possible_exec_name),
            Err(_) => ()
        };
    }
    return Err("Unable to locate python interpreter".to_owned());
}

/// Determine the python version we're supposed to be building
/// from the features passed via the environment.
///
/// The environment variable can choose to omit a minor
/// version if the user doesn't care.
fn version_from_env() -> Result<PythonVersion, String> {
    let re = Regex::new(r"CARGO_FEATURE_PYTHON_(\d+)(_(\d+))?").unwrap();
    // sort env::vars so we get more explicit version specifiers first
    // so if the user passes e.g. the python-3 feature and the python-3-5
    // feature, python-3-5 takes priority.
    let mut vars = env::vars().collect::<Vec<_>>();
    vars.sort_by(|a, b| b.cmp(a));
    for (key, _) in vars {
        match re.captures(&key) {
            Some(cap) => return Ok(PythonVersion {
                major: cap.at(1).unwrap().parse().unwrap(),
                minor: match cap.at(3) {
                    Some(s) => Some(s.parse().unwrap()),
                    None => None
                }
            }),
            None => ()
        }
    }
    Err("Python version feature was not found. At least one python version \
         feature must be enabled.".to_owned())
}

fn main() {
    // 1. Setup cfg variables so we can do conditional compilation in this
    // library based on the python interpeter's compilation flags. This is
    // necessary for e.g. matching the right unicode and threading interfaces.
    //
    // By default, try to use pkgconfig - this seems to be a rust norm.
    //
    // If you want to use a different python, setting the appropriate
    // PYTHON_X.X_NO_PKG_CONFIG environment variable will cause the script
    // to pick up the python in your PATH; e.g. for python27 X.X is 2.7.
    //
    // This will work smoothly with an activated virtualenv.
    //
    // If you have troubles with your shell accepting '.' in a var name,
    // try using 'env' (sorry but this isn't our fault - it just has to
    // match the pkg-config package name, which is going to have a . in it).
    let version = version_from_env().unwrap();
    let pkg_name = match version.minor {
        Some(minor) => format!("python-{}.{}", version.major, minor),
        None => format!("python{}", version.major)
    };

    let python_interpreter_path = match configure_from_pkgconfig(&version, &pkg_name) {
        Ok(p) => p,
        // no pkgconfig - either it failed or user set the environment
        // variable "PYTHON_2.7_NO_PKG_CONFIG".
        Err(_) => configure_from_path(&version).unwrap()
    };

    let config_map = get_config_vars(&python_interpreter_path).unwrap();
    for (key, val) in &config_map {
        match cfg_line_for_var(key, val) {
            Some(line) => println!("{}", line),
            None => ()
        }
    }

    // 2. Export python interpreter compilation flags as cargo variables that
    // will be visible to dependents. All flags will be available to dependent
    // build scripts in the environment variable DEP_PYTHON27_PYTHON_FLAGS as
    // comma separated list; each item in the list looks like
    //
    // {VAL,FLAG}_{flag_name}=val;
    //
    // FLAG indicates the variable is always 0 or 1
    // VAL indicates it can take on any value
    //
    // rust-cypthon/build.rs contains an example of how to unpack this data
    // into cfg flags that replicate the ones present in this library, so
    // you can use the same cfg syntax.
    let flags: String = config_map.iter().fold("".to_owned(), |memo, (key, val)| {
        if is_value(key) {
            memo + format!("VAL_{}={},", key, val).as_ref()
        } else if val != "0" {
            memo + format!("FLAG_{}={},", key, val).as_ref()
        } else {
            memo
        }
    });
    println!("cargo:python_flags={}", &flags[..flags.len()-1]);
}
<commit_msg>clone build.rs from python27 into python3<commit_after>extern crate pkg_config;
extern crate regex;

use std::process::Command;
use std::collections::HashMap;
use std::env;
use regex::Regex;
use std::fs;

struct PythonVersion {
    major: u8,
    // minor == None means any minor version will do
    minor: Option<u8>
}

const CFG_KEY: &'static str = "py_sys_config";

// windows' python writes out lines with the windows crlf sequence;
// posix platforms and mac os should write out lines with just lf.
#[cfg(target_os="windows")]
static NEWLINE_SEQUENCE: &'static str = "\r\n";

#[cfg(not(target_os="windows"))]
static NEWLINE_SEQUENCE: &'static str = "\n";

// A list of python interpreter compile-time preprocessor defines that
// we will pick up and pass to rustc via --cfg=py_sys_config={varname};
// this allows using them conditional cfg attributes in the .rs files, so
//
// #[cfg(py_sys_config="{varname}"]
//
// is the equivalent of #ifdef {varname} name in C.
//
// see Misc/SpecialBuilds.txt in the python source for what these mean.
//
// (hrm, this is sort of re-implementing what distutils does, except
// by passing command line args instead of referring to a python.h)
#[cfg(not(target_os="windows"))]
static SYSCONFIG_FLAGS: [&'static str; 7] = [
    "Py_USING_UNICODE",
    "Py_UNICODE_WIDE",
    "WITH_THREAD",
    "Py_DEBUG",
    "Py_REF_DEBUG",
    "Py_TRACE_REFS",
    "COUNT_ALLOCS",
];

static SYSCONFIG_VALUES: [&'static str; 1] = [
    // cfg doesn't support flags with values, just bools - so flags
    // below are translated into bools as {varname}_{val}
    //
    // for example, Py_UNICODE_SIZE_2 or Py_UNICODE_SIZE_4
    "Py_UNICODE_SIZE" // note - not present on python 3.3+, which is always wide
];

/// Examine python's compile flags to pass to cfg by launching
/// the interpreter and printing variables of interest from
/// sysconfig.get_config_vars.
#[cfg(not(target_os="windows"))]
fn get_config_vars(python_path: &String) -> Result<HashMap<String, String>, String>  {
    let mut script = "import sysconfig; \
config = sysconfig.get_config_vars();".to_owned();

    for k in SYSCONFIG_FLAGS.iter().chain(SYSCONFIG_VALUES.iter()) {
        script.push_str(&format!("print(config.get('{}', {}))", k,
            if is_value(k) { "None" } else { "0" } ));
        script.push_str(";");
    }

    let mut cmd = Command::new(python_path);
    cmd.arg("-c").arg(script);

    let out = try!(cmd.output().map_err(|e| {
        format!("failed to run python interpreter `{:?}`: {}", cmd, e)
    }));

    if !out.status.success() {
        let stderr = String::from_utf8(out.stderr).unwrap();
        let mut msg = format!("python script failed with stderr:\n\n");
        msg.push_str(&stderr);
        return Err(msg);
    }

    let stdout = String::from_utf8(out.stdout).unwrap();
    let split_stdout: Vec<&str> = stdout.trim_right().split(NEWLINE_SEQUENCE).collect();
    if split_stdout.len() != SYSCONFIG_VALUES.len() + SYSCONFIG_FLAGS.len() {
        return Err(
            format!("python stdout len didn't return expected number of lines:
{}", split_stdout.len()).to_string());
    }
    let all_vars = SYSCONFIG_FLAGS.iter().chain(SYSCONFIG_VALUES.iter());
    // let var_map: HashMap<String, String> = HashMap::new();
    Ok(all_vars.zip(split_stdout.iter())
        .fold(HashMap::new(), |mut memo: HashMap<String, String>, (&k, &v)| {
            if !(v.to_owned() == "None" && is_value(k)) {
                memo.insert(k.to_owned(), v.to_owned());
            }
            memo
        }))
}

#[cfg(target_os="windows")]
fn get_config_vars(_: &String) -> Result<HashMap<String, String>, String> {
    // sysconfig is missing all the flags on windows, so we can't actually
    // query the interpreter directly for its build flags.
    //
    // For the time being, this is the flags as defined in the python source's
    // PC\pyconfig.h. This won't work correctly if someone has built their
    // python with a modified pyconfig.h - sorry if that is you, you will have
    // to comment/uncomment the lines below.
    let mut map: HashMap<String, String> = HashMap::new();
    map.insert("Py_USING_UNICODE".to_owned(), "1".to_owned());
    map.insert("Py_UNICODE_WIDE".to_owned(), "0".to_owned());
    map.insert("WITH_THREAD".to_owned(), "1".to_owned());
    map.insert("Py_UNICODE_SIZE".to_owned(), "2".to_owned());

    // This is defined #ifdef _DEBUG. The visual studio build seems to produce
    // a specially named pythonXX_d.exe and pythonXX_d.dll when you build the
    // Debug configuration, which this script doesn't currently support anyway.
    // map.insert("Py_DEBUG", "1");

    // Uncomment these manually if your python was built with these and you want
    // the cfg flags to be set in rust.
    //
    // map.insert("Py_REF_DEBUG", "1");
    // map.insert("Py_TRACE_REFS", "1");
    // map.insert("COUNT_ALLOCS", 1");
    Ok(map)
}

fn is_value(key: &str) -> bool {
    SYSCONFIG_VALUES.iter().find(|x| **x == key).is_some()
}

fn cfg_line_for_var(key: &str, val: &str) -> Option<String> {
    if is_value(key) {
        // is a value; suffix the key name with the value
        Some(format!("cargo:rustc-cfg={}=\"{}_{}\"\n", CFG_KEY, key, val))
    } else if val != "0" {
        // is a flag that isn't zero
        Some(format!("cargo:rustc-cfg={}=\"{}\"", CFG_KEY, key))
    } else {
        // is a flag that is zero
        None
    }
}

/// Run a python script using the 'python' located by PATH.
fn run_python_script(script: &str) -> Result<String, String> {
    let mut cmd = Command::new("python");
    cmd.arg("-c").arg(script);

    let out = try!(cmd.output().map_err(|e| {
        format!("failed to run python interpreter `{:?}`: {}", cmd, e)
    }));

    if !out.status.success() {
        let stderr = String::from_utf8(out.stderr).unwrap();
        let mut msg = format!("python script failed with stderr:\n\n");
        msg.push_str(&stderr);
        return Err(msg);
    }

    let out = String::from_utf8(out.stdout).unwrap();
    return Ok(out);
}

#[cfg(not(target_os="macos"))]
#[cfg(not(target_os="windows"))]
fn get_rustc_link_lib(version: &PythonVersion, enable_shared: bool) -> Result<String, String> {
    if enable_shared {
        Ok(format!("cargo:rustc-link-lib=python{}.{}", version.major,
            version.minor))
    } else {
        Ok(format!("cargo:rustc-link-lib=static=python{}.{}", version.major,
            version.minor))
    }
}

#[cfg(target_os="macos")]
fn get_macos_linkmodel() -> Result<String, String> {
    let script = "import MacOS; print MacOS.linkmodel;";
    let out = run_python_script(script).unwrap();
    Ok(out.trim_right().to_owned())
}

#[cfg(target_os="macos")]
fn get_rustc_link_lib(version: &PythonVersion, _: bool) -> Result<String, String> {
    // os x can be linked to a framework or static or dynamic, and
    // Py_ENABLE_SHARED is wrong; framework means shared library
    let dotted_version = format!("{}.{}", version.major, version.minor.unwrap());
    match get_macos_linkmodel().unwrap().as_ref() {
        "static" => Ok(format!("cargo:rustc-link-lib=static=python{}",
            dotted_version)),
        "dynamic" => Ok(format!("cargo:rustc-link-lib=python{}",
            dotted_version)),
        "framework" => Ok(format!("cargo:rustc-link-lib=python{}",
            dotted_version)),
        other => Err(format!("unknown linkmodel {}", other))
    }
}

/// Check that the interpreter's reported version matches expected_version;
/// if it does, return the specific version the interpreter reports.
fn get_interpreter_version(line: &str, expected_version: &PythonVersion)
        -> Result<PythonVersion, String> {
    let version_re = Regex::new(r"\((\d+), (\d+)\)").unwrap();
    let interpreter_version = match version_re.captures(&line) {
        Some(cap) => PythonVersion {
            major: cap.at(1).unwrap().parse().unwrap(),
            minor: Some(cap.at(2).unwrap().parse().unwrap())
        },
        None => return Err(
            format!("Unexpected response to version query {}", line))
    };

    if interpreter_version.major != expected_version.major ||
        (expected_version.minor.is_some() &&
            interpreter_version.minor != expected_version.minor)
    {
        Err(format!("'python' is not version {}.{} (is {})",
                    expected_version.major,
                    match expected_version.minor {
                        Some(v) => v.to_string(),
                        None => "*".to_owned()
                    },
                    line))
    } else {
        Ok(interpreter_version)
    }
}

#[cfg(target_os="windows")]
fn get_rustc_link_lib(version: &PythonVersion, _: bool) -> Result<String, String> {
    // Py_ENABLE_SHARED doesn't seem to be present on windows.
    Ok(format!("cargo:rustc-link-lib=python{}{}", version.major,
        match version.minor {
            Some(minor) => minor.to_string(),
            None => "".to_owned()
        }))
}

/// Deduce configuration from the 'python' in the current PATH and print
/// cargo vars to stdout.
///
/// Note that if the python doesn't satisfy expected_version, this will error.
fn configure_from_path(expected_version: &PythonVersion) -> Result<String, String> {
    let script = "import sys; import sysconfig; print(sys.version_info[0:2]); \
print(sysconfig.get_config_var('LIBDIR')); \
print(sysconfig.get_config_var('Py_ENABLE_SHARED')); \
print(sys.exec_prefix);";
    let out = run_python_script(script).unwrap();
    let lines: Vec<&str> = out.split(NEWLINE_SEQUENCE).collect();
    let version: &str = lines[0];
    let libpath: &str = lines[1];
    let enable_shared: &str = lines[2];
    let exec_prefix: &str = lines[3];

    let interpreter_version = try!(get_interpreter_version(version,
        expected_version));

    println!("{}", get_rustc_link_lib(&interpreter_version,
        enable_shared == "1").unwrap());

    if libpath != "None" {
        println!("cargo:rustc-link-search=native={}", libpath);
    }

    let rel_interpreter_path = if cfg!(target_os="windows") {
        "/python"
    } else {
        "/bin/python"
    };

    return Ok(format!("{}{}", exec_prefix, rel_interpreter_path));
}

/// Deduce configuration from the python-X.X in pkg-config and print
/// cargo vars to stdout.
fn configure_from_pkgconfig(version: &PythonVersion, pkg_name: &str)
        -> Result<String, String> {
    // this emits relevant build info to stdout, which is picked up by the
    // build chain (funny name for something with side-effects!)
    try!(pkg_config::find_library(pkg_name));

    // This seems to be a convention - unfortunately pkg-config doesn't
    // tell you the executable name, but I've noticed at least on
    // OS X homebrew the python bin dir for 3.4 doesn't actually contain
    // a 'python'.
    let exec_prefix = pkg_config::Config::get_variable(pkg_name,
        "exec_prefix").unwrap();

    // try to find the python interpreter in the exec_prefix somewhere.
    // the .pc doesn't tell us :(

    let mut attempts = vec![
        format!("/bin/python{}", version.major),
        "/bin/python".to_owned()
    ];

    // Try to seek python(major).(minor) if the user specified a minor.
    //
    // Ideally, we'd still do this even if they didn't based off the
    // specific version of the package located by pkg_config above,
    // but it's not obvious how to reliably extract that out of the .pc.
    if version.minor.is_some() {
        attempts.insert(0, format!("/bin/python{}_{}", version.major,
            version.minor.unwrap()));
    }

    for attempt in attempts.iter() {
        let possible_exec_name = format!("{}{}", exec_prefix,
            attempt);
        match fs::metadata(&possible_exec_name) {
            Ok(_) => return Ok(possible_exec_name),
            Err(_) => ()
        };
    }
    return Err("Unable to locate python interpreter".to_owned());
}

/// Determine the python version we're supposed to be building
/// from the features passed via the environment.
///
/// The environment variable can choose to omit a minor
/// version if the user doesn't care.
fn version_from_env() -> Result<PythonVersion, String> {
    let re = Regex::new(r"CARGO_FEATURE_PYTHON_(\d+)(_(\d+))?").unwrap();
    // sort env::vars so we get more explicit version specifiers first
    // so if the user passes e.g. the python-3 feature and the python-3-5
    // feature, python-3-5 takes priority.
    let mut vars = env::vars().collect::<Vec<_>>();
    vars.sort_by(|a, b| b.cmp(a));
    for (key, _) in vars {
        match re.captures(&key) {
            Some(cap) => return Ok(PythonVersion {
                major: cap.at(1).unwrap().parse().unwrap(),
                minor: match cap.at(3) {
                    Some(s) => Some(s.parse().unwrap()),
                    None => None
                }
            }),
            None => ()
        }
    }
    Err("Python version feature was not found. At least one python version \
         feature must be enabled.".to_owned())
}

fn main() {
    // 1. Setup cfg variables so we can do conditional compilation in this
    // library based on the python interpeter's compilation flags. This is
    // necessary for e.g. matching the right unicode and threading interfaces.
    //
    // By default, try to use pkgconfig - this seems to be a rust norm.
    //
    // If you want to use a different python, setting the appropriate
    // PYTHON_X.X_NO_PKG_CONFIG environment variable will cause the script
    // to pick up the python in your PATH; e.g. for python27 X.X is 2.7.
    //
    // This will work smoothly with an activated virtualenv.
    //
    // If you have troubles with your shell accepting '.' in a var name,
    // try using 'env' (sorry but this isn't our fault - it just has to
    // match the pkg-config package name, which is going to have a . in it).
    let version = version_from_env().unwrap();
    let pkg_name = match version.minor {
        Some(minor) => format!("python-{}.{}", version.major, minor),
        None => format!("python{}", version.major)
    };

    let python_interpreter_path = match configure_from_pkgconfig(&version, &pkg_name) {
        Ok(p) => p,
        // no pkgconfig - either it failed or user set the environment
        // variable "PYTHON_2.7_NO_PKG_CONFIG".
        Err(_) => configure_from_path(&version).unwrap()
    };

    let config_map = get_config_vars(&python_interpreter_path).unwrap();

    for (key, val) in &config_map {
        match cfg_line_for_var(key, val) {
            Some(line) => println!("{}", line),
            None => ()
        }
    }

    // 2. Export python interpreter compilation flags as cargo variables that
    // will be visible to dependents. All flags will be available to dependent
    // build scripts in the environment variable DEP_PYTHON27_PYTHON_FLAGS as
    // comma separated list; each item in the list looks like
    //
    // {VAL,FLAG}_{flag_name}=val;
    //
    // FLAG indicates the variable is always 0 or 1
    // VAL indicates it can take on any value
    //
    // rust-cypthon/build.rs contains an example of how to unpack this data
    // into cfg flags that replicate the ones present in this library, so
    // you can use the same cfg syntax.
    let flags: String = config_map.iter().fold("".to_owned(), |memo, (key, val)| {
        if is_value(key) {
            memo + format!("VAL_{}={},", key, val).as_ref()
        } else if val != "0" {
            memo + format!("FLAG_{}={},", key, val).as_ref()
        } else {
            memo
        }
    });
    println!("cargo:python_flags={}",
        if flags.len() > 0 { &flags[..flags.len()-1] } else { "" });
}
<|endoftext|>


<commit_before>use producer::*;
use std::rc::Rc;

pub trait BoxedConsumer<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, T);
    fn end(self: Box<Self>);
}

pub trait Consumer<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, T);
}

pub trait ConsumerRef<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, &T);
}

pub trait ConsumerRefMut<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, &mut T);
}<commit_msg>Clean up BoxedConsumer<commit_after>use producer::*;
use std::rc::Rc;

pub trait Consumer<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, T);
}

pub trait ConsumerRef<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, &T);
}

pub trait ConsumerRefMut<T> {
    fn init(&mut self, Rc<Producer>);
    fn emit(&mut self, &mut T);
}<|endoftext|>


<commit_before>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;
extern crate orbimage;
extern crate orbfont;

use std::env;
use std::path::Path;
use std::process::Command;

use orbclient::{Color, EventOption, Window, K_ESC};
use orbimage::Image;
use orbfont::Font;

use package::Package;

pub mod package;

const BAR_COLOR: Color = Color::rgb(40, 45, 57);
const BAR_HIGHLIGHT_COLOR: Color = Color::rgb(80, 86, 102);
const TEXT_COLOR: Color = Color::rgb(204, 210, 224);
const TEXT_HIGHLIGHT_COLOR: Color = Color::rgb(235, 241, 255);

fn get_packages() -> Vec<Package> {
    let read_dir = Path::new("/ui/apps/").read_dir().expect("failed to read_dir on /ui/apps/");

    let mut entries = vec![];
    for dir in read_dir {
        let dir = match dir {
            Ok(x) => x,
            Err(_) => continue,
        };
        let file_name = dir.file_name().to_string_lossy().into_owned();
        if dir.file_type().expect("failed to get file_type").is_file() {
            entries.push(file_name);
        }
    }

    entries.sort();

    let mut packages: Vec<Package> = Vec::new();
    for entry in entries.iter() {
        packages.push(Package::from_path(&format!("/ui/apps/{}", entry)));
    }

    packages
}

fn draw(window: &mut Window, packages: &Vec<Package>, start: &Image, shutdown: &Image, selected: i32){
    let w = window.width();
    let h = window.height();
    window.set(BAR_COLOR);

    let mut x = 0;
    let mut i = 0;

    {
        let y = h as isize - start.height() as isize;

        if i == selected {
            window.rect(x as i32, y as i32,
                              start.width() as u32, start.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        start.draw(window, x as i32, y as i32);

        x += start.width() as i32;
        i += 1;
    }

    for package in packages.iter() {
        let y = h as isize - package.icon.height() as isize;

        if i == selected {
            window.rect(x as i32, y as i32,
                              package.icon.width() as u32, package.icon.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        package.icon.draw(window, x as i32, y as i32);

        x += package.icon.width() as i32;
        i += 1;
    }

    {
        x = w as i32 - shutdown.width() as i32;
        let y = h as isize - shutdown.height() as isize;

        if i == selected {
            window.rect(x as i32, y as i32,
                              shutdown.width() as u32, shutdown.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        shutdown.draw(window, x as i32, y as i32);
    }

    window.sync();
}

fn draw_chooser(window: &mut Window, font: &Font, packages: &Vec<Package>, _mouse_x: i32, mouse_y: i32){
    let w = window.width();

    window.set(BAR_COLOR);

    let mut y = 0;
    for package in packages.iter() {
        let highlight = mouse_y >= y as i32 && mouse_y < y + 32;

        if highlight {
            window.rect(0, y, w, 32, BAR_HIGHLIGHT_COLOR);
        }

        package.icon.draw(window, 0, y);

        let mut c_x = 40;
        for c in package.name.chars() {
            font.render(&c.to_string(), 16.0).draw(window, c_x as i32, y + 8, if highlight { TEXT_HIGHLIGHT_COLOR } else { TEXT_COLOR });
            c_x += 8;
        }

        y += 32;
    }

    window.sync();
}

fn main() {
    let paths = env::args().skip(1);
    if paths.len() > 0 {
        for ref path in paths {
            let mut packages = get_packages();

            packages.retain(|package| -> bool {
                for accept in package.accepts.iter() {
                    if (accept.starts_with('*') && path.ends_with(&accept[1 ..])) ||
                       (accept.ends_with('*') && path.starts_with(&accept[.. accept.len() - 1])) {
                        return true;
                    }
                }
                false
            });

            if packages.len() > 1 {
                let mut window = Window::new(-1, -1, 400, packages.len() as u32 * 32, path).expect("launcher: failed to open window");
                let font = Font::find(None, None, None).expect("launcher: failed to open font");

                draw_chooser(&mut window, &font, &packages, -1, -1);
                'choosing: loop {
                    for event in window.events() {
                        match event.to_option() {
                            EventOption::Mouse(mouse_event) => {
                                draw_chooser(&mut window, &font, &packages, mouse_event.x, mouse_event.y);

                                if mouse_event.left_button {
                                    let mut y = 0;
                                    for package in packages.iter() {
                                        if mouse_event.y >= y && mouse_event.y < y + 32 {
                                            if let Err(err) = Command::new(&package.binary).arg(path).spawn() {
                                                println!("launcher: failed to launch {}: {}", package.binary, err);
                                            }
                                            break 'choosing;
                                        }
                                        y += 32;
                                    }
                                }
                            },
                            EventOption::Quit(_) => break 'choosing,
                            _ => ()
                        }
                    }
                }
            } else if let Some(package) = packages.get(0) {
                if let Err(err) = Command::new(&package.binary).arg(&path).spawn() {
                    println!("launcher: failed to launch '{}': {}", package.binary, err);
                }
            } else {
                println!("launcher: no application found for '{}'", path);
            }
        }
    } else {
        let packages = get_packages();

        let start = Image::from_path("/ui/icons/start.png").unwrap_or(Image::default());

        let shutdown = Image::from_path("/ui/icons/actions/system-log-out.png").unwrap_or(Image::default());

        let (width, height) = orbclient::get_display_size().expect("launcher: failed to get display size");
        let mut window = Window::new(0, height as i32 - 32, width, 32, "").expect("launcher: failed to open window");

        let mut selected = -1;

        draw(&mut window, &packages, &start, &shutdown, selected);
        'running: loop {
            for event in window.events() {
                match event.to_option() {
                    EventOption::Mouse(mouse_event) => {
                        let mut now_selected = -1;

                        {
                            let mut x = 0;
                            let mut i = 0;

                            {
                                let y = window.height() as i32 - start.height() as i32;
                                if mouse_event.y >= y && mouse_event.x >= x &&
                                   mouse_event.x < x + start.width() as i32 {
                                       now_selected = i;
                                }
                                x += start.width() as i32;
                                i += 1;
                            }

                            for package in packages.iter() {
                                let y = window.height() as i32 - package.icon.height() as i32;
                                if mouse_event.y >= y && mouse_event.x >= x &&
                                   mouse_event.x < x + package.icon.width() as i32 {
                                    now_selected = i;
                                }
                                x += package.icon.width() as i32;
                                i += 1;
                            }

                            {
                                x = window.width() as i32 - shutdown.width() as i32;
                                let y = window.height() as i32 - shutdown.height() as i32;
                                if mouse_event.y >= y && mouse_event.x >= x &&
                                   mouse_event.x < x + shutdown.width() as i32 {
                                       now_selected = i;
                                }
                            }
                        }

                        if now_selected != selected {
                            selected = now_selected;
                            draw(&mut window, &packages, &start, &shutdown, selected);
                        }

                        if mouse_event.left_button {
                            let mut i = 0;

                            if i == selected {
                                let start_h = packages.len() as u32 * 32;
                                let mut start_window = Window::new(0, height as i32 - 32 - start_h as i32, 400, start_h, "").unwrap();
                                let font = Font::find(None, None, None).unwrap();

                                draw_chooser(&mut start_window, &font, &packages, -1, -1);
                                'start_choosing: loop {
                                    for event in start_window.events() {
                                        match event.to_option() {
                                            EventOption::Mouse(mouse_event) => {
                                                draw_chooser(&mut start_window, &font, &packages, mouse_event.x, mouse_event.y);

                                                if mouse_event.left_button {
                                                    let mut y = 0;
                                                    for package in packages.iter() {
                                                        if mouse_event.y >= y && mouse_event.y < y + 32 {
                                                            if let Err(err) = Command::new(&package.binary).spawn() {
                                                                println!("launcher: failed to launch {}: {}", package.binary, err);
                                                            }
                                                            break 'start_choosing;
                                                        }
                                                        y += 32;
                                                    }
                                                }
                                            },
                                            EventOption::Key(key_event) => {
                                                match key_event.scancode {
                                                    K_ESC => break 'start_choosing,
                                                    _ => ()
                                                }
                                            },
                                            EventOption::Focus(focus_event) => if ! focus_event.focused {
                                                break 'start_choosing;
                                            },
                                            EventOption::Quit(_) => break 'start_choosing,
                                            _ => ()
                                        }
                                    }
                                }
                            }
                            i += 1;

                            for package in packages.iter() {
                                if i == selected {
                                    if let Err(err) = Command::new(&package.binary).spawn() {
                                        println!("{}: Failed to launch: {}", package.binary, err);
                                    }
                                }
                                i += 1;
                            }

                            if i == selected {
                                   break 'running;
                            }
                        }
                    },
                    EventOption::Quit(_) => break 'running,
                    _ => ()
                }
            }
        }
    }
}
<commit_msg>More efficient start redraw, organization of bar and chooser<commit_after>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;
extern crate orbimage;
extern crate orbfont;

use std::env;
use std::path::Path;
use std::process::Command;

use orbclient::{Color, EventOption, Window, K_ESC};
use orbimage::Image;
use orbfont::Font;

use package::Package;

pub mod package;

const BAR_COLOR: Color = Color::rgb(40, 45, 57);
const BAR_HIGHLIGHT_COLOR: Color = Color::rgb(80, 86, 102);
const TEXT_COLOR: Color = Color::rgb(204, 210, 224);
const TEXT_HIGHLIGHT_COLOR: Color = Color::rgb(235, 241, 255);

fn get_packages() -> Vec<Package> {
    let read_dir = Path::new("/ui/apps/").read_dir().expect("failed to read_dir on /ui/apps/");

    let mut entries = vec![];
    for dir in read_dir {
        let dir = match dir {
            Ok(x) => x,
            Err(_) => continue,
        };
        let file_name = dir.file_name().to_string_lossy().into_owned();
        if dir.file_type().expect("failed to get file_type").is_file() {
            entries.push(file_name);
        }
    }

    entries.sort();

    let mut packages: Vec<Package> = Vec::new();
    for entry in entries.iter() {
        packages.push(Package::from_path(&format!("/ui/apps/{}", entry)));
    }

    packages
}

fn draw(window: &mut Window, packages: &Vec<Package>, start: &Image, shutdown: &Image, selected: i32){
    let w = window.width();
    let h = window.height();
    window.set(BAR_COLOR);

    let mut x = 0;
    let mut i = 0;

    {
        let y = h as isize - start.height() as isize;

        if i == selected {
            window.rect(x as i32, y as i32,
                              start.width() as u32, start.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        start.draw(window, x as i32, y as i32);

        x += start.width() as i32;
        i += 1;
    }

    for package in packages.iter() {
        let y = h as isize - package.icon.height() as isize;

        if i == selected {
            window.rect(x as i32, y as i32,
                              package.icon.width() as u32, package.icon.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        package.icon.draw(window, x as i32, y as i32);

        x += package.icon.width() as i32;
        i += 1;
    }

    {
        x = w as i32 - shutdown.width() as i32;
        let y = h as isize - shutdown.height() as isize;

        if i == selected {
            window.rect(x as i32, y as i32,
                              shutdown.width() as u32, shutdown.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        shutdown.draw(window, x as i32, y as i32);
    }

    window.sync();
}

fn draw_chooser(window: &mut Window, font: &Font, packages: &Vec<Package>, selected: i32){
    let w = window.width();

    window.set(BAR_COLOR);

    let mut y = 0;
    for (i, package) in packages.iter().enumerate() {
        if i as i32 == selected {
            window.rect(0, y, w, 32, BAR_HIGHLIGHT_COLOR);
        }

        package.icon.draw(window, 0, y);

        let mut c_x = 40;
        for c in package.name.chars() {
            font.render(&c.to_string(), 16.0).draw(window, c_x as i32, y + 8, if i as i32 == selected { TEXT_HIGHLIGHT_COLOR } else { TEXT_COLOR });
            c_x += 8;
        }

        y += 32;
    }

    window.sync();
}

fn bar_main() {
    let packages = get_packages();

    let start = Image::from_path("/ui/icons/start.png").unwrap_or(Image::default());

    let shutdown = Image::from_path("/ui/icons/actions/system-log-out.png").unwrap_or(Image::default());

    let (width, height) = orbclient::get_display_size().expect("launcher: failed to get display size");
    let mut window = Window::new(0, height as i32 - 32, width, 32, "").expect("launcher: failed to open window");

    let mut selected = -1;

    draw(&mut window, &packages, &start, &shutdown, selected);
    'running: loop {
        for event in window.events() {
            match event.to_option() {
                EventOption::Mouse(mouse_event) => {
                    let mut now_selected = -1;

                    {
                        let mut x = 0;
                        let mut i = 0;

                        {
                            let y = window.height() as i32 - start.height() as i32;
                            if mouse_event.y >= y && mouse_event.x >= x &&
                               mouse_event.x < x + start.width() as i32 {
                                   now_selected = i;
                            }
                            x += start.width() as i32;
                            i += 1;
                        }

                        for package in packages.iter() {
                            let y = window.height() as i32 - package.icon.height() as i32;
                            if mouse_event.y >= y && mouse_event.x >= x &&
                               mouse_event.x < x + package.icon.width() as i32 {
                                now_selected = i;
                            }
                            x += package.icon.width() as i32;
                            i += 1;
                        }

                        {
                            x = window.width() as i32 - shutdown.width() as i32;
                            let y = window.height() as i32 - shutdown.height() as i32;
                            if mouse_event.y >= y && mouse_event.x >= x &&
                               mouse_event.x < x + shutdown.width() as i32 {
                                   now_selected = i;
                            }
                        }
                    }

                    if now_selected != selected {
                        selected = now_selected;
                        draw(&mut window, &packages, &start, &shutdown, selected);
                    }

                    if mouse_event.left_button {
                        let mut i = 0;

                        if i == selected {
                            let start_h = packages.len() as u32 * 32;
                            let mut start_window = Window::new(0, height as i32 - 32 - start_h as i32, 400, start_h, "").unwrap();
                            let font = Font::find(None, None, None).unwrap();

                            let mut selected = -1;

                            draw_chooser(&mut start_window, &font, &packages, selected);
                            'start_choosing: loop {
                                for event in start_window.events() {
                                    match event.to_option() {
                                        EventOption::Mouse(mouse_event) => {
                                            let mut now_selected = -1;

                                            let mut y = 0;
                                            for (i, _package) in packages.iter().enumerate() {
                                                if mouse_event.y >= y && mouse_event.y < y + 32 {
                                                    now_selected = i as i32;
                                                }
                                                y += 32;
                                            }

                                            if now_selected != selected {
                                                selected = now_selected;
                                                draw_chooser(&mut start_window, &font, &packages, selected);
                                            }

                                            if mouse_event.left_button {
                                                let mut y = 0;
                                                for package in packages.iter() {
                                                    if mouse_event.y >= y && mouse_event.y < y + 32 {
                                                        if let Err(err) = Command::new(&package.binary).spawn() {
                                                            println!("launcher: failed to launch {}: {}", package.binary, err);
                                                        }
                                                        break 'start_choosing;
                                                    }
                                                    y += 32;
                                                }
                                            }
                                        },
                                        EventOption::Key(key_event) => {
                                            match key_event.scancode {
                                                K_ESC => break 'start_choosing,
                                                _ => ()
                                            }
                                        },
                                        EventOption::Focus(focus_event) => if ! focus_event.focused {
                                            break 'start_choosing;
                                        },
                                        EventOption::Quit(_) => break 'start_choosing,
                                        _ => ()
                                    }
                                }
                            }
                        }
                        i += 1;

                        for package in packages.iter() {
                            if i == selected {
                                if let Err(err) = Command::new(&package.binary).spawn() {
                                    println!("{}: Failed to launch: {}", package.binary, err);
                                }
                            }
                            i += 1;
                        }

                        if i == selected {
                               break 'running;
                        }
                    }
                },
                EventOption::Quit(_) => break 'running,
                _ => ()
            }
        }
    }
}

fn chooser_main(paths: env::Args) {
    for ref path in paths.skip(1) {
        let mut packages = get_packages();

        packages.retain(|package| -> bool {
            for accept in package.accepts.iter() {
                if (accept.starts_with('*') && path.ends_with(&accept[1 ..])) ||
                   (accept.ends_with('*') && path.starts_with(&accept[.. accept.len() - 1])) {
                    return true;
                }
            }
            false
        });

        if packages.len() > 1 {
            let mut window = Window::new(-1, -1, 400, packages.len() as u32 * 32, path).expect("launcher: failed to open window");
            let font = Font::find(None, None, None).expect("launcher: failed to open font");

            let mut selected = -1;

            draw_chooser(&mut window, &font, &packages, selected);
            'choosing: loop {
                for event in window.events() {
                    match event.to_option() {
                        EventOption::Mouse(mouse_event) => {
                            let mut now_selected = -1;

                            let mut y = 0;
                            for (i, _package) in packages.iter().enumerate() {
                                if mouse_event.y >= y && mouse_event.y < y + 32 {
                                    now_selected = i as i32;
                                }
                                y += 32;
                            }

                            if now_selected != selected {
                                selected = now_selected;
                                draw_chooser(&mut window, &font, &packages, selected);
                            }

                            if mouse_event.left_button {
                                let mut y = 0;
                                for package in packages.iter() {
                                    if mouse_event.y >= y && mouse_event.y < y + 32 {
                                        if let Err(err) = Command::new(&package.binary).arg(path).spawn() {
                                            println!("launcher: failed to launch {}: {}", package.binary, err);
                                        }
                                        break 'choosing;
                                    }
                                    y += 32;
                                }
                            }
                        },
                        EventOption::Quit(_) => break 'choosing,
                        _ => ()
                    }
                }
            }
        } else if let Some(package) = packages.get(0) {
            if let Err(err) = Command::new(&package.binary).arg(&path).spawn() {
                println!("launcher: failed to launch '{}': {}", package.binary, err);
            }
        } else {
            println!("launcher: no application found for '{}'", path);
        }
    }
}

fn main() {
    let paths = env::args();
    if paths.len() > 1 {
        chooser_main(paths);
    } else {
        bar_main();
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::cmp::Reverse;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

use itertools::Itertools;
use regex::Regex;

use common::Solution;

#[derive(Default, Debug, Eq, PartialEq, Clone)]
struct Group {
    power: u32,
    count: u32,
    hp: u32,
    initiative: i32,
    damage: String,
    weaknesses: Vec<String>,
    immunities: Vec<String>,
    faction: char,
}

impl Group {
    pub fn is_alive(&self) -> bool {
        self.count > 0
    }

    pub fn damage_to(&self, other: &Group) -> u32 {
        let damage = self.effective_power();

        if other.weaknesses.contains(&self.damage) {
            damage * 2
        } else if other.immunities.contains(&self.damage) {
            0
        } else {
            damage
        }
    }

    pub fn effective_power(&self) -> u32 {
        self.power * self.count
    }
}

impl Ord for Group {
    fn cmp(&self, other: &Self) -> Ordering {
        self.effective_power().cmp(&other.effective_power())
            .then(other.initiative.cmp(&self.initiative))
    }
}

impl PartialOrd for Group {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Default)]
pub struct Day24 {
    units: Vec<Group>,
}

impl Day24 {
    pub fn new() -> Self {
        Default::default()
    }

    fn read_input(&mut self, input: &mut Read) {
        let matcher = Regex::new(r"(\d+) units each with (\d+) hit points (\(([^)]+)\) )?with an attack that does (\d+) (\w+) damage at initiative (\d+)").unwrap();
        let weakness_matcher = Regex::new(r"(weak|immune) to ([^;)]+)").unwrap();
        let reader = BufReader::new(input);
        let mut fac = 'X';

        for line in reader.lines() {
            let line = line.unwrap();
            match line.as_str() {
                "Immune System:" => { fac = 'D' }
                "" => {}
                "Infection:" => { fac = 'I' }
                line => {
                    let caps = matcher.captures(line).unwrap_or_else(|| {
                        panic!("{}", line);
                    });

                    let mut group = Group {
                        count: caps[1].parse().unwrap(),
                        hp: caps[2].parse().unwrap(),
                        power: caps[5].parse().unwrap(),
                        damage: caps[6].to_string(),
                        initiative: caps[7].parse().unwrap(),
                        faction: fac,
                        ..Default::default()
                    };

                    if let Some(caps) = caps.get(4) {
                        for modified in weakness_matcher.captures_iter(caps.as_str()) {
                            let target = match &modified[1] {
                                "weak" => &mut group.weaknesses,
                                "immune" => &mut group.immunities,
                                _ => panic!(),
                            };

                            for t in modified[2].split(", ") {
                                target.push(t.to_string());
                            }
                        }
                    }
                    self.units.push(group)
                }
            }
        }
    }

    fn simulate(&mut self) {
        let mut order: Vec<usize> = (0..self.units.len()).collect();
        order.sort_unstable_by_key(|&x| &self.units[x]);
        order.reverse();

        // select targets
        let mut targets: Vec<Option<usize>> = vec![None; self.units.len()];
        let mut is_targeted = vec![false; self.units.len()];

        for &i in &order {
            let unit = &self.units[i];
            if !unit.is_alive() {
                continue;
            }
            let damage = self.units.iter().map(|x| unit.damage_to(x)).collect_vec();
            let target = (0..self.units.len())
                .filter(|&x| !is_targeted[x] && self.units[x].faction != unit.faction && self.units[x].is_alive())
                .max_by(|&x, &y| {
                    damage[x].cmp(&damage[y])
                        .then(self.units[x].effective_power().cmp(&self.units[y].effective_power()))
                        .then(self.units[x].initiative.cmp(&self.units[y].initiative))
                });

            if let Some(target) = target {
                targets[i] = Some(target);
                is_targeted[target] = true;
            }
        }

        order.sort_unstable_by_key(|&x| Reverse(self.units[x].initiative));

        for attacker in order {
            if !self.units[attacker].is_alive() {
                continue;
            }

            if let Some(id) = targets[attacker] {
                let damage = self.units[attacker].damage_to(&self.units[id]);
                let defender = &mut self.units[id];

                let losses = damage / defender.hp;
                defender.count = defender.count.saturating_sub(losses);
            }
        }
    }

    fn both_alive(&self) -> bool {
        let mut seen = None;

        for unit in self.units.iter().filter(|g| g.is_alive()) {
            if let Some(seen) = seen {
                if seen != unit.faction {
                    return true;
                }
                continue;
            }

            seen = Some(unit.faction);
        }

        return false;
    }
}

impl Solution for Day24 {
    fn part1(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        while self.both_alive() {
            self.simulate();
        }

        let result: u32 = self.units.iter().map(|x| x.count).sum();
        result.to_string()
    }

    fn part2(&mut self, _input: &mut Read) -> String {
        unimplemented!()
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day24::Day24;

    const SAMPLE_INPUT: &[u8] = include_bytes!("samples/24.txt");

    #[test]
    fn sample_part1() {
        let mut instance = Day24::new();
        assert_eq!("5216", instance.part1(&mut SAMPLE_INPUT))
    }
}
<commit_msg>Add incorrect implementation of day 24 part 2.<commit_after>use std::cmp::Ordering;
use std::cmp::Reverse;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

use itertools::Itertools;
use regex::Regex;

use common::Solution;

#[derive(Default, Debug, Eq, PartialEq, Clone)]
struct Group {
    power: u32,
    count: u32,
    hp: u32,
    initiative: i32,
    damage: String,
    weaknesses: Vec<String>,
    immunities: Vec<String>,
    faction: char,
}

impl Group {
    pub fn is_alive(&self) -> bool {
        self.count > 0
    }

    pub fn damage_to(&self, other: &Group) -> u32 {
        let damage = self.effective_power();

        if other.weaknesses.contains(&self.damage) {
            damage * 2
        } else if other.immunities.contains(&self.damage) {
            0
        } else {
            damage
        }
    }

    pub fn effective_power(&self) -> u32 {
        self.power * self.count
    }
}

impl Ord for Group {
    fn cmp(&self, other: &Self) -> Ordering {
        self.effective_power().cmp(&other.effective_power())
            .then(other.initiative.cmp(&self.initiative))
    }
}

impl PartialOrd for Group {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Default)]
pub struct Day24 {
    units: Vec<Group>,
}

impl Day24 {
    pub fn new() -> Self {
        Default::default()
    }

    fn read_input(&mut self, input: &mut Read) {
        let matcher = Regex::new(r"(\d+) units each with (\d+) hit points (\(([^)]+)\) )?with an attack that does (\d+) (\w+) damage at initiative (\d+)").unwrap();
        let weakness_matcher = Regex::new(r"(weak|immune) to ([^;)]+)").unwrap();
        let reader = BufReader::new(input);
        let mut fac = 'X';

        for line in reader.lines() {
            let line = line.unwrap();
            match line.as_str() {
                "Immune System:" => { fac = 'D' }
                "" => {}
                "Infection:" => { fac = 'I' }
                line => {
                    let caps = matcher.captures(line).unwrap_or_else(|| {
                        panic!("{}", line);
                    });

                    let mut group = Group {
                        count: caps[1].parse().unwrap(),
                        hp: caps[2].parse().unwrap(),
                        power: caps[5].parse().unwrap(),
                        damage: caps[6].to_string(),
                        initiative: caps[7].parse().unwrap(),
                        faction: fac,
                        ..Default::default()
                    };

                    if let Some(caps) = caps.get(4) {
                        for modified in weakness_matcher.captures_iter(caps.as_str()) {
                            let target = match &modified[1] {
                                "weak" => &mut group.weaknesses,
                                "immune" => &mut group.immunities,
                                _ => panic!(),
                            };

                            for t in modified[2].split(", ") {
                                target.push(t.to_string());
                            }
                        }
                    }
                    self.units.push(group)
                }
            }
        }
    }

    fn simulate(&mut self) -> bool {
        let mut order: Vec<usize> = (0..self.units.len()).collect();
        order.sort_unstable_by_key(|&x| &self.units[x]);
        order.reverse();

        // select targets
        let mut targets: Vec<Option<usize>> = vec![None; self.units.len()];
        let mut is_targeted = vec![false; self.units.len()];
        let mut changes = false;

        for &i in &order {
            let unit = &self.units[i];
            if !unit.is_alive() {
                continue;
            }
            let damage = self.units.iter().map(|x| unit.damage_to(x)).collect_vec();
            let target = (0..self.units.len())
                .filter(|&x| !is_targeted[x] && self.units[x].faction != unit.faction && self.units[x].is_alive())
                .max_by(|&x, &y| {
                    damage[x].cmp(&damage[y])
                        .then(self.units[x].effective_power().cmp(&self.units[y].effective_power()))
                        .then(self.units[x].initiative.cmp(&self.units[y].initiative))
                });

            if let Some(target) = target {
                targets[i] = Some(target);
                is_targeted[target] = true;
            }
        }

        order.sort_unstable_by_key(|&x| Reverse(self.units[x].initiative));

        for attacker in order {
            if !self.units[attacker].is_alive() {
                continue;
            }

            if let Some(id) = targets[attacker] {
                let damage = self.units[attacker].damage_to(&self.units[id]);
                let defender = &mut self.units[id];

                let losses = damage / defender.hp;
                if losses > 0 {
                    defender.count = defender.count.saturating_sub(losses);
                    changes = true;
                }
            }
        }
        changes
    }

    fn both_alive(&self) -> bool {
        let mut seen = None;

        for unit in self.units.iter().filter(|g| g.is_alive()) {
            if let Some(seen) = seen {
                if seen != unit.faction {
                    return true;
                }
                continue;
            }

            seen = Some(unit.faction);
        }

        return false;
    }

    fn full_simulation(&mut self) {
        let mut changes = true;
        while self.both_alive() && changes {
            changes = self.simulate();
        }
    }

    fn faction_won(&self, faction: char) -> bool {
        if self.both_alive() {
            false
        } else {
            self.units.iter().filter(|x| x.is_alive())
                .all(|x| x.faction == faction)
        }
    }
}

impl Solution for Day24 {
    fn part1(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        self.full_simulation();

        let result: u32 = self.units.iter().map(|x| x.count).sum();
        result.to_string()
    }

    fn part2(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        let original_count = self.units.iter().map(|x| x.count).collect_vec();

        loop {
            for (unit, count) in self.units.iter_mut().zip(original_count.iter()) {
                unit.count = *count;
                if unit.faction == 'D' {
                    unit.power += 1;
                }
            }
            println!("{:?}", self.units);

            self.full_simulation();
            if self.faction_won('D') {
                let result: u32 = self.units.iter().map(|x| x.count).sum();
                return result.to_string();
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day24::Day24;

    const SAMPLE_INPUT: &[u8] = include_bytes!("samples/24.txt");

    #[test]
    fn sample_part1() {
        let mut instance = Day24::new();
        assert_eq!("5216", instance.part1(&mut SAMPLE_INPUT))
    }

    #[test]
    fn sample_part2() {
        let mut instance = Day24::new();
        assert_eq!("51", instance.part2(&mut SAMPLE_INPUT))
    }
}
<|endoftext|>


<commit_before>use std::cmp::Reverse;
use std::collections::hash_map::Entry;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::fmt::Display;
use std::io::Read;
use std::mem::swap;

use crate::common::LineIter;

type Item<const S: usize> = (u32, State<S>);
type Todo<const S: usize> = BinaryHeap<Reverse<Item<S>>>;
type Visited<const S: usize> = HashMap<State<S>, u32>;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Hash)]
enum Pod {
    A,
    B,
    C,
    D,
}

impl Pod {
    pub fn cost(self) -> u32 {
        match self {
            Pod::A => 1,
            Pod::B => 10,
            Pod::C => 100,
            Pod::D => 1000,
        }
    }

    pub fn dest(self) -> usize {
        self as usize
    }
}

impl TryFrom<char> for Pod {
    type Error = &'static str;

    fn try_from(c: char) -> Result<Self, Self::Error> {
        match c {
            'A' => Ok(Pod::A),
            'B' => Ok(Pod::B),
            'C' => Ok(Pod::C),
            'D' => Ok(Pod::D),
            _ => Err("Invalid pod"),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
struct State<const S: usize> {
    hallway: [Option<Pod>; 11],
    rooms: [[Option<Pod>; S]; 4],
}

fn room_hallway_pos(room: usize) -> usize {
    room * 2 + 2
}

fn abs_delta(a: usize, b: usize) -> usize {
    if a < b {
        b - a
    } else {
        a - b
    }
}

impl<const S: usize> State<S> {
    const VALID_HALLWAY_POS: [usize; 7] = [0, 1, 3, 5, 7, 9, 10];

    pub fn is_done(&self) -> bool {
        self == &State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); S],
                [Some(Pod::B); S],
                [Some(Pod::C); S],
                [Some(Pod::D); S],
            ],
        }
    }

    fn add_to_queue(self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        let entry = visited.entry(self.clone());

        if matches!(&entry, Entry::Occupied(entry) if *entry.get() <= cost) {
            // Already got a better one
            return;
        }

        // nightly only :'(
        // entry.insert(cost);
        *entry.or_default() = cost;

        todo.push(Reverse((cost + self.estimate(), self)))
    }

    fn estimate(&self) -> u32 {
        // A* estimate. For every entry that is not already "at rest", the cost is the cost
        // required to get it to the top of its intended room.

        // Cost to enter the hole for all pods that still need to
        let enter_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(index, room)| {
                let pod = match index {
                    0 => Pod::A,
                    1 => Pod::B,
                    2 => Pod::C,
                    3 => Pod::D,
                    _ => unreachable!(),
                };

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| entry == Some(pod))
                    .map(|(index, _)| index as u32 + 1)
                    .sum::<u32>()
                    * pod.cost()
            })
            .sum();

        // Cost for all of the hallway to move to above their intended rooms
        let hallway_estimate: u32 = self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, &pod)| {
                let pod = pod?;

                let destination_pos = room_hallway_pos(pod.dest());

                Some(abs_delta(pos, destination_pos) as u32 * pod.cost())
            })
            .sum();

        // Cost to move out of the room and above the correct rooms
        let rooms_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(room_index, room)| {
                let hallway_pos = room_hallway_pos(room_index);

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| {
                        entry.map(|pod| pod.dest() == room_index).unwrap_or(false)
                    })
                    .filter_map(|(room_pos, &pod)| {
                        let pod = pod?;

                        let destination_pos = room_hallway_pos(pod.dest());

                        let steps = 1 + room_pos + abs_delta(hallway_pos, destination_pos).max(2);

                        Some(steps as u32 * pod.cost())
                    })
                    .sum::<u32>()
            })
            .sum();

        enter_estimate + hallway_estimate + rooms_estimate
    }

    pub fn generate_next(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        self.room_to_hallway(cost, todo, visited);
        self.hallway_to_room(cost, todo, visited);
    }

    fn room_to_hallway(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (index, room) in self.rooms.iter().enumerate() {
            // Check if we even want to move anything out of this room
            if room
                .iter()
                .all(|entry| entry.map(|pod| pod.dest() == index).unwrap_or(true))
            {
                continue;
            }

            let (pos, pod) = room
                .iter()
                .enumerate()
                .find_map(|(pos, entry)| entry.map(|pod| (pos, pod)))
                .unwrap(); // Safe unwrap, we know it exists from above.

            let base_cost = 1 + pos;
            let hallway_pos = room_hallway_pos(index);

            let mut queue_new = |new_pos, new_cost| {
                let mut new_state = self.clone();
                swap(
                    &mut new_state.hallway[new_pos],
                    &mut new_state.rooms[index][pos],
                );

                new_state.add_to_queue(new_cost + cost, todo, visited)
            };

            // Check positions to the left
            for new_pos in (0..hallway_pos).rev() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + hallway_pos - new_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }

            // And to the right
            for new_pos in hallway_pos..self.hallway.len() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + new_pos - hallway_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }
        }
    }

    fn hallway_to_room(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (pos, pod) in self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, pod)| pod.map(|pod| (pos, pod)))
        {
            let room = pod.dest();
            let new_hallway_pos = room_hallway_pos(room);

            // Check if the path is free
            let in_between = if new_hallway_pos < pos {
                &self.hallway[(new_hallway_pos + 1)..pos]
            } else {
                &self.hallway[(pos + 1)..new_hallway_pos]
            };

            if in_between.iter().any(Option::is_some) {
                // Something's in the way
                continue;
            }

            // Check if we can move into the room
            if self.rooms[room]
                .iter()
                .copied()
                .flatten()
                .any(|other| other != pod)
            {
                // Scared of other pods
                continue;
            }

            let room_pos = if let Some(pos) = self.rooms[room].iter().rposition(Option::is_none) {
                pos
            } else {
                continue;
            };

            let new_cost = (abs_delta(pos, new_hallway_pos) + room_pos + 1) as u32 * pod.cost();
            let mut new_state = self.clone();
            swap(
                &mut new_state.hallway[pos],
                &mut new_state.rooms[room][room_pos],
            );
            new_state.add_to_queue(cost + new_cost, todo, visited);
        }
    }

    pub fn solve(&self) -> u32 {
        let mut todo = Todo::new();

        let mut visited = HashMap::new();
        visited.insert(self.clone(), 0);

        todo.push(Reverse((self.estimate(), self.clone())));

        while let Some(Reverse((_, state))) = todo.pop() {
            let cost = *visited.get(&state).unwrap_or(&0);

            if state.is_done() {
                return cost;
            }

            state.generate_next(cost, &mut todo, &mut visited);
        }

        panic!("No route found!")
    }
}

impl<const S: usize> Display for State<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let helper = |opt_pod| match opt_pod {
            Some(Pod::A) => 'A',
            Some(Pod::B) => 'B',
            Some(Pod::C) => 'C',
            Some(Pod::D) => 'D',
            None => '.',
        };
        writeln!(f, "#############")?;
        write!(f, "#")?;

        for entry in self.hallway {
            write!(f, "{}", helper(entry))?;
        }
        writeln!(f, "#")?;

        for i in 0..S {
            writeln!(
                f,
                "  #{}#{}#{}#{}#",
                helper(self.rooms[0][i]),
                helper(self.rooms[1][i]),
                helper(self.rooms[2][i]),
                helper(self.rooms[3][i])
            )?;
        }

        write!(f, "  #########")
    }
}

fn read_input(input: &mut dyn Read) -> State<2> {
    let mut reader = LineIter::new(input);
    let mut state = State {
        hallway: Default::default(),
        rooms: Default::default(),
    };

    let _ = reader.next();
    let _ = reader.next();

    let mut helper = |idx: usize| {
        reader
            .next()
            .unwrap()
            .chars()
            .filter_map(|c| Pod::try_from(c).ok())
            .zip(&mut state.rooms)
            .for_each(|(pod, room)| room[idx] = Some(pod))
    };

    helper(0);
    helper(1);

    state
}

pub fn part1(input: &mut dyn Read) -> String {
    let state = read_input(input);

    state.solve().to_string()
}

pub fn part2(input: &mut dyn Read) -> String {
    let state2 = read_input(input);

    let state4 = State {
        hallway: Default::default(),
        rooms: [
            [
                state2.rooms[0][0],
                Some(Pod::D),
                Some(Pod::D),
                state2.rooms[0][1],
            ],
            [
                state2.rooms[1][0],
                Some(Pod::C),
                Some(Pod::B),
                state2.rooms[1][1],
            ],
            [
                state2.rooms[2][0],
                Some(Pod::B),
                Some(Pod::A),
                state2.rooms[2][1],
            ],
            [
                state2.rooms[3][0],
                Some(Pod::A),
                Some(Pod::C),
                state2.rooms[3][1],
            ],
        ],
    };

    state4.solve().to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::test_implementation;

    const SAMPLE: &[u8] = include_bytes!("samples/23.txt");

    #[test]
    fn test_is_done() {
        let state = State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); 2],
                [Some(Pod::B); 2],
                [Some(Pod::C); 2],
                [Some(Pod::D); 2],
            ],
        };

        assert!(state.is_done());
    }

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 12521);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 44169);
    }
}
<commit_msg>Properly use TryFrom<commit_after>use std::cmp::Reverse;
use std::collections::hash_map::Entry;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::fmt::Display;
use std::io::Read;
use std::mem::swap;

use crate::common::LineIter;

type Item<const S: usize> = (u32, State<S>);
type Todo<const S: usize> = BinaryHeap<Reverse<Item<S>>>;
type Visited<const S: usize> = HashMap<State<S>, u32>;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Hash)]
enum Pod {
    A,
    B,
    C,
    D,
}

impl Pod {
    pub fn cost(self) -> u32 {
        match self {
            Pod::A => 1,
            Pod::B => 10,
            Pod::C => 100,
            Pod::D => 1000,
        }
    }

    pub fn dest(self) -> usize {
        self as usize
    }
}

impl TryFrom<usize> for Pod {
    type Error = usize;

    fn try_from(index: usize) -> Result<Self, Self::Error> {
        match index {
            0 => Ok(Pod::A),
            1 => Ok(Pod::B),
            2 => Ok(Pod::C),
            3 => Ok(Pod::D),
            _ => Err(index),
        }
    }
}

impl TryFrom<char> for Pod {
    type Error = &'static str;

    fn try_from(c: char) -> Result<Self, Self::Error> {
        match c {
            'A' => Ok(Pod::A),
            'B' => Ok(Pod::B),
            'C' => Ok(Pod::C),
            'D' => Ok(Pod::D),
            _ => Err("Invalid pod"),
        }
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
struct State<const S: usize> {
    hallway: [Option<Pod>; 11],
    rooms: [[Option<Pod>; S]; 4],
}

fn room_hallway_pos(room: usize) -> usize {
    room * 2 + 2
}

fn abs_delta(a: usize, b: usize) -> usize {
    if a < b {
        b - a
    } else {
        a - b
    }
}

impl<const S: usize> State<S> {
    const VALID_HALLWAY_POS: [usize; 7] = [0, 1, 3, 5, 7, 9, 10];

    pub fn is_done(&self) -> bool {
        self == &State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); S],
                [Some(Pod::B); S],
                [Some(Pod::C); S],
                [Some(Pod::D); S],
            ],
        }
    }

    fn add_to_queue(self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        let entry = visited.entry(self.clone());

        if matches!(&entry, Entry::Occupied(entry) if *entry.get() <= cost) {
            // Already got a better one
            return;
        }

        // nightly only :'(
        // entry.insert(cost);
        *entry.or_default() = cost;

        todo.push(Reverse((cost + self.estimate(), self)))
    }

    fn estimate(&self) -> u32 {
        // A* estimate. For every entry that is not already "at rest", the cost is the cost
        // required to get it to the top of its intended room.

        // Cost to enter the hole for all pods that still need to
        let enter_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(index, room)| {
                let pod = Pod::try_from(index).unwrap();

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| entry == Some(pod))
                    .map(|(index, _)| index as u32 + 1)
                    .sum::<u32>()
                    * pod.cost()
            })
            .sum();

        // Cost for all of the hallway to move to above their intended rooms
        let hallway_estimate: u32 = self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, &pod)| {
                let pod = pod?;

                let destination_pos = room_hallway_pos(pod.dest());

                Some(abs_delta(pos, destination_pos) as u32 * pod.cost())
            })
            .sum();

        // Cost to move out of the room and above the correct rooms
        let rooms_estimate: u32 = self
            .rooms
            .iter()
            .enumerate()
            .map(|(room_index, room)| {
                let hallway_pos = room_hallway_pos(room_index);

                room.iter()
                    .enumerate()
                    .rev()
                    .skip_while(|&(_, &entry)| {
                        entry.map(|pod| pod.dest() == room_index).unwrap_or(false)
                    })
                    .filter_map(|(room_pos, &pod)| {
                        let pod = pod?;

                        let destination_pos = room_hallway_pos(pod.dest());

                        let steps = 1 + room_pos + abs_delta(hallway_pos, destination_pos).max(2);

                        Some(steps as u32 * pod.cost())
                    })
                    .sum::<u32>()
            })
            .sum();

        enter_estimate + hallway_estimate + rooms_estimate
    }

    pub fn generate_next(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        self.room_to_hallway(cost, todo, visited);
        self.hallway_to_room(cost, todo, visited);
    }

    fn room_to_hallway(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (index, room) in self.rooms.iter().enumerate() {
            // Check if we even want to move anything out of this room
            if room
                .iter()
                .all(|entry| entry.map(|pod| pod.dest() == index).unwrap_or(true))
            {
                continue;
            }

            let (pos, pod) = room
                .iter()
                .enumerate()
                .find_map(|(pos, entry)| entry.map(|pod| (pos, pod)))
                .unwrap(); // Safe unwrap, we know it exists from above.

            let base_cost = 1 + pos;
            let hallway_pos = room_hallway_pos(index);

            let mut queue_new = |new_pos, new_cost| {
                let mut new_state = self.clone();
                swap(
                    &mut new_state.hallway[new_pos],
                    &mut new_state.rooms[index][pos],
                );

                new_state.add_to_queue(new_cost + cost, todo, visited)
            };

            // Check positions to the left
            for new_pos in (0..hallway_pos).rev() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + hallway_pos - new_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }

            // And to the right
            for new_pos in hallway_pos..self.hallway.len() {
                if self.hallway[new_pos].is_some() {
                    // Hit an occupied room
                    break;
                }

                if !Self::VALID_HALLWAY_POS.contains(&new_pos) {
                    // Not allowed to stop here
                    continue;
                }

                let new_cost = (base_cost + new_pos - hallway_pos) as u32 * pod.cost();
                queue_new(new_pos, new_cost);
            }
        }
    }

    fn hallway_to_room(&self, cost: u32, todo: &mut Todo<S>, visited: &mut Visited<S>) {
        for (pos, pod) in self
            .hallway
            .iter()
            .enumerate()
            .filter_map(|(pos, pod)| pod.map(|pod| (pos, pod)))
        {
            let room = pod.dest();
            let new_hallway_pos = room_hallway_pos(room);

            // Check if the path is free
            let in_between = if new_hallway_pos < pos {
                &self.hallway[(new_hallway_pos + 1)..pos]
            } else {
                &self.hallway[(pos + 1)..new_hallway_pos]
            };

            if in_between.iter().any(Option::is_some) {
                // Something's in the way
                continue;
            }

            // Check if we can move into the room
            if self.rooms[room]
                .iter()
                .copied()
                .flatten()
                .any(|other| other != pod)
            {
                // Scared of other pods
                continue;
            }

            let room_pos = if let Some(pos) = self.rooms[room].iter().rposition(Option::is_none) {
                pos
            } else {
                continue;
            };

            let new_cost = (abs_delta(pos, new_hallway_pos) + room_pos + 1) as u32 * pod.cost();
            let mut new_state = self.clone();
            swap(
                &mut new_state.hallway[pos],
                &mut new_state.rooms[room][room_pos],
            );
            new_state.add_to_queue(cost + new_cost, todo, visited);
        }
    }

    pub fn solve(&self) -> u32 {
        let mut todo = Todo::new();

        let mut visited = HashMap::new();
        visited.insert(self.clone(), 0);

        todo.push(Reverse((self.estimate(), self.clone())));

        while let Some(Reverse((_, state))) = todo.pop() {
            let cost = *visited.get(&state).unwrap_or(&0);

            if state.is_done() {
                return cost;
            }

            state.generate_next(cost, &mut todo, &mut visited);
        }

        panic!("No route found!")
    }
}

impl<const S: usize> Display for State<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let helper = |opt_pod| match opt_pod {
            Some(Pod::A) => 'A',
            Some(Pod::B) => 'B',
            Some(Pod::C) => 'C',
            Some(Pod::D) => 'D',
            None => '.',
        };
        writeln!(f, "#############")?;
        write!(f, "#")?;

        for entry in self.hallway {
            write!(f, "{}", helper(entry))?;
        }
        writeln!(f, "#")?;

        for i in 0..S {
            writeln!(
                f,
                "  #{}#{}#{}#{}#",
                helper(self.rooms[0][i]),
                helper(self.rooms[1][i]),
                helper(self.rooms[2][i]),
                helper(self.rooms[3][i])
            )?;
        }

        write!(f, "  #########")
    }
}

fn read_input(input: &mut dyn Read) -> State<2> {
    let mut reader = LineIter::new(input);
    let mut state = State {
        hallway: Default::default(),
        rooms: Default::default(),
    };

    let _ = reader.next();
    let _ = reader.next();

    let mut helper = |idx: usize| {
        reader
            .next()
            .unwrap()
            .chars()
            .filter_map(|c| Pod::try_from(c).ok())
            .zip(&mut state.rooms)
            .for_each(|(pod, room)| room[idx] = Some(pod))
    };

    helper(0);
    helper(1);

    state
}

pub fn part1(input: &mut dyn Read) -> String {
    let state = read_input(input);

    state.solve().to_string()
}

pub fn part2(input: &mut dyn Read) -> String {
    let state2 = read_input(input);

    let state4 = State {
        hallway: Default::default(),
        rooms: [
            [
                state2.rooms[0][0],
                Some(Pod::D),
                Some(Pod::D),
                state2.rooms[0][1],
            ],
            [
                state2.rooms[1][0],
                Some(Pod::C),
                Some(Pod::B),
                state2.rooms[1][1],
            ],
            [
                state2.rooms[2][0],
                Some(Pod::B),
                Some(Pod::A),
                state2.rooms[2][1],
            ],
            [
                state2.rooms[3][0],
                Some(Pod::A),
                Some(Pod::C),
                state2.rooms[3][1],
            ],
        ],
    };

    state4.solve().to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::test_implementation;

    const SAMPLE: &[u8] = include_bytes!("samples/23.txt");

    #[test]
    fn test_is_done() {
        let state = State {
            hallway: Default::default(),
            rooms: [
                [Some(Pod::A); 2],
                [Some(Pod::B); 2],
                [Some(Pod::C); 2],
                [Some(Pod::D); 2],
            ],
        };

        assert!(state.is_done());
    }

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 12521);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 44169);
    }
}
<|endoftext|>