<commit_before>// ignore-windows: File handling is not implemented yet
// compile-flags: -Zmiri-disable-isolation

use std::fs::{File, remove_file, rename};
use std::io::{Read, Write, ErrorKind, Result, Seek, SeekFrom};
use std::path::{PathBuf, Path};

fn main() {
    test_file();
    test_file_clone();
    test_seek();
    test_metadata();
    test_symlink();
    test_errors();
    test_rename();
}

/// Prepare: compute filename and make sure the file does not exist.
fn prepare(filename: &str) -> PathBuf {
    let tmp = std::env::temp_dir();
    let path = tmp.join(filename);
    // Clean the paths for robustness.
    remove_file(&path).ok();
    path
}

/// Prepare like above, and also write some initial content to the file.
fn prepare_with_content(filename: &str, content: &[u8]) -> PathBuf {
    let path = prepare(filename);
    let mut file = File::create(&path).unwrap();
    file.write(content).unwrap();
    path
}

fn test_file() {
    let bytes = b"Hello, World!\n";
    let path = prepare("miri_test_fs_file.txt");

    // Test creating, writing and closing a file (closing is tested when `file` is dropped).
    let mut file = File::create(&path).unwrap();
    // Writing 0 bytes should not change the file contents.
    file.write(&mut []).unwrap();
    assert_eq!(file.metadata().unwrap().len(), 0);

    file.write(bytes).unwrap();
    assert_eq!(file.metadata().unwrap().len(), bytes.len() as u64);
    // Test opening, reading and closing a file.
    let mut file = File::open(&path).unwrap();
    let mut contents = Vec::new();
    // Reading 0 bytes should not move the file pointer.
    file.read(&mut []).unwrap();
    // Reading until EOF should get the whole text.
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_file_clone() {
    let bytes = b"Hello, World!\n";
    let path = prepare_with_content("miri_test_fs_file_clone.txt", bytes);

    // Cloning a file should be successful.
    let file = File::open(&path).unwrap();
    let mut cloned = file.try_clone().unwrap();
    // Reading from a cloned file should get the same text.
    let mut contents = Vec::new();
    cloned.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_seek() {
    let bytes = b"Hello, entire World!\n";
    let path = prepare_with_content("miri_test_fs_seek.txt", bytes);

    let mut file = File::open(&path).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    // Test that seeking to the beginning and reading until EOF gets the text again.
    file.seek(SeekFrom::Start(0)).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());
    // Test seeking relative to the end of the file.
    file.seek(SeekFrom::End(-1)).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(&bytes[bytes.len() - 1..], contents.as_slice());
    // Test seeking relative to the current position.
    file.seek(SeekFrom::Start(5)).unwrap();
    file.seek(SeekFrom::Current(-3)).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(&bytes[2..], contents.as_slice());

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn check_metadata(bytes: &[u8], path: &Path) -> Result<()> {
    // Test that the file metadata is correct.
    let metadata = path.metadata()?;
    // `path` should point to a file.
    assert!(metadata.is_file());
    // The size of the file must be equal to the number of written bytes.
    assert_eq!(bytes.len() as u64, metadata.len());
    Ok(())
}

fn test_metadata() {
    let bytes = b"Hello, meta-World!\n";
    let path = prepare_with_content("miri_test_fs_metadata.txt", bytes);

    // Test that metadata of an absolute path is correct.
    check_metadata(bytes, &path).unwrap();
    // Test that metadata of a relative path is correct.
    std::env::set_current_dir(path.parent().unwrap()).unwrap();
    check_metadata(bytes, Path::new(path.file_name().unwrap())).unwrap();

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_symlink() {
    let bytes = b"Hello, World!\n";
    let path = prepare_with_content("miri_test_fs_link_target.txt", bytes);
    let symlink_path = prepare("miri_test_fs_symlink.txt");

    // Creating a symbolic link should succeed.
    std::os::unix::fs::symlink(&path, &symlink_path).unwrap();
    // Test that the symbolic link has the same contents as the file.
    let mut symlink_file = File::open(&symlink_path).unwrap();
    let mut contents = Vec::new();
    symlink_file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());
    // Test that metadata of a symbolic link is correct.
    check_metadata(bytes, &symlink_path).unwrap();
    // Test that the metadata of a symbolic link is correct when not following it.
    assert!(symlink_path.symlink_metadata().unwrap().file_type().is_symlink());
    // Removing symbolic link should succeed.
    remove_file(&symlink_path).unwrap();

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_errors() {
    let bytes = b"Hello, World!\n";
    let path = prepare("miri_test_fs_errors.txt");

    // The following tests also check that the `__errno_location()` shim is working properly.
    // Opening a non-existing file should fail with a "not found" error.
    assert_eq!(ErrorKind::NotFound, File::open(&path).unwrap_err().kind());
    // Removing a non-existing file should fail with a "not found" error.
    assert_eq!(ErrorKind::NotFound, remove_file(&path).unwrap_err().kind());
    // Reading the metadata of a non-existing file should fail with a "not found" error.
    assert_eq!(ErrorKind::NotFound, check_metadata(bytes, &path).unwrap_err().kind());
}

fn test_rename() {
    // Renaming a file should succeed.
    let path1 = prepare("miri_test_fs_rename_source.txt");
    let path2 = prepare("miri_test_fs_rename_destination.txt");

    let file = File::create(&path1).unwrap();
    drop(file);
    rename(&path1, &path2).unwrap();
    assert_eq!(ErrorKind::NotFound, path1.metadata().unwrap_err().kind());
    assert!(path2.metadata().unwrap().is_file());
    remove_file(&path2).unwrap();
}
<commit_msg>test a bit more<commit_after>// ignore-windows: File handling is not implemented yet
// compile-flags: -Zmiri-disable-isolation

use std::fs::{File, remove_file, rename};
use std::io::{Read, Write, ErrorKind, Result, Seek, SeekFrom};
use std::path::{PathBuf, Path};

fn main() {
    test_file();
    test_file_clone();
    test_seek();
    test_metadata();
    test_symlink();
    test_errors();
    test_rename();
}

/// Prepare: compute filename and make sure the file does not exist.
fn prepare(filename: &str) -> PathBuf {
    let tmp = std::env::temp_dir();
    let path = tmp.join(filename);
    // Clean the paths for robustness.
    remove_file(&path).ok();
    path
}

/// Prepare like above, and also write some initial content to the file.
fn prepare_with_content(filename: &str, content: &[u8]) -> PathBuf {
    let path = prepare(filename);
    let mut file = File::create(&path).unwrap();
    file.write(content).unwrap();
    path
}

fn test_file() {
    let bytes = b"Hello, World!\n";
    let path = prepare("miri_test_fs_file.txt");

    // Test creating, writing and closing a file (closing is tested when `file` is dropped).
    let mut file = File::create(&path).unwrap();
    // Writing 0 bytes should not change the file contents.
    file.write(&mut []).unwrap();
    assert_eq!(file.metadata().unwrap().len(), 0);

    file.write(bytes).unwrap();
    assert_eq!(file.metadata().unwrap().len(), bytes.len() as u64);
    // Test opening, reading and closing a file.
    let mut file = File::open(&path).unwrap();
    let mut contents = Vec::new();
    // Reading 0 bytes should not move the file pointer.
    file.read(&mut []).unwrap();
    // Reading until EOF should get the whole text.
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_file_clone() {
    let bytes = b"Hello, World!\n";
    let path = prepare_with_content("miri_test_fs_file_clone.txt", bytes);

    // Cloning a file should be successful.
    let file = File::open(&path).unwrap();
    let mut cloned = file.try_clone().unwrap();
    // Reading from a cloned file should get the same text.
    let mut contents = Vec::new();
    cloned.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_seek() {
    let bytes = b"Hello, entire World!\n";
    let path = prepare_with_content("miri_test_fs_seek.txt", bytes);

    let mut file = File::open(&path).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());
    // Test that seeking to the beginning and reading until EOF gets the text again.
    file.seek(SeekFrom::Start(0)).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());
    // Test seeking relative to the end of the file.
    file.seek(SeekFrom::End(-1)).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(&bytes[bytes.len() - 1..], contents.as_slice());
    // Test seeking relative to the current position.
    file.seek(SeekFrom::Start(5)).unwrap();
    file.seek(SeekFrom::Current(-3)).unwrap();
    let mut contents = Vec::new();
    file.read_to_end(&mut contents).unwrap();
    assert_eq!(&bytes[2..], contents.as_slice());

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn check_metadata(bytes: &[u8], path: &Path) -> Result<()> {
    // Test that the file metadata is correct.
    let metadata = path.metadata()?;
    // `path` should point to a file.
    assert!(metadata.is_file());
    // The size of the file must be equal to the number of written bytes.
    assert_eq!(bytes.len() as u64, metadata.len());
    Ok(())
}

fn test_metadata() {
    let bytes = b"Hello, meta-World!\n";
    let path = prepare_with_content("miri_test_fs_metadata.txt", bytes);

    // Test that metadata of an absolute path is correct.
    check_metadata(bytes, &path).unwrap();
    // Test that metadata of a relative path is correct.
    std::env::set_current_dir(path.parent().unwrap()).unwrap();
    check_metadata(bytes, Path::new(path.file_name().unwrap())).unwrap();

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_symlink() {
    let bytes = b"Hello, World!\n";
    let path = prepare_with_content("miri_test_fs_link_target.txt", bytes);
    let symlink_path = prepare("miri_test_fs_symlink.txt");

    // Creating a symbolic link should succeed.
    std::os::unix::fs::symlink(&path, &symlink_path).unwrap();
    // Test that the symbolic link has the same contents as the file.
    let mut symlink_file = File::open(&symlink_path).unwrap();
    let mut contents = Vec::new();
    symlink_file.read_to_end(&mut contents).unwrap();
    assert_eq!(bytes, contents.as_slice());
    // Test that metadata of a symbolic link is correct.
    check_metadata(bytes, &symlink_path).unwrap();
    // Test that the metadata of a symbolic link is correct when not following it.
    assert!(symlink_path.symlink_metadata().unwrap().file_type().is_symlink());
    // Removing symbolic link should succeed.
    remove_file(&symlink_path).unwrap();

    // Removing file should succeed.
    remove_file(&path).unwrap();
}

fn test_errors() {
    let bytes = b"Hello, World!\n";
    let path = prepare("miri_test_fs_errors.txt");

    // The following tests also check that the `__errno_location()` shim is working properly.
    // Opening a non-existing file should fail with a "not found" error.
    assert_eq!(ErrorKind::NotFound, File::open(&path).unwrap_err().kind());
    // Removing a non-existing file should fail with a "not found" error.
    assert_eq!(ErrorKind::NotFound, remove_file(&path).unwrap_err().kind());
    // Reading the metadata of a non-existing file should fail with a "not found" error.
    assert_eq!(ErrorKind::NotFound, check_metadata(bytes, &path).unwrap_err().kind());
}

fn test_rename() {
    // Renaming a file should succeed.
    let path1 = prepare("miri_test_fs_rename_source.txt");
    let path2 = prepare("miri_test_fs_rename_destination.txt");

    let file = File::create(&path1).unwrap();
    drop(file);
    rename(&path1, &path2).unwrap();
    assert_eq!(ErrorKind::NotFound, path1.metadata().unwrap_err().kind());
    assert!(path2.metadata().unwrap().is_file());
    remove_file(&path2).unwrap();
}
<|endoftext|>


<commit_before>
pub struct Task {
        id: int,
        todo: ~str,
        priority: char,
        finished: bool,
        raw_todo: ~str,
        contexts: & 'static [~str],
        projects: & 'static [~str]
}

impl Task {
        pub fn to_str(&self) -> ~str {
                return format!("{:d} {:b} ({:c}) '{:s}' \"{:s}\"",
                                self.id, self.finished, self.priority,
                                self.todo, self.raw_todo)
        }

        pub fn create(text: ~str, id: int) -> Task {
                let mut task = Task{id: id, raw_todo: text.clone(), todo: text.clone(),
                                        finished: false, priority: '^',
                                        contexts: &[], projects: &[]};

                let mut todo = text.clone();

                if (todo.slice(0, 2) == "x ") {
                        task.finished = true;
                        todo = todo.slice(2, todo.len()).to_owned();
                }

                let chars = todo.chars().to_owned_vec();
                if (chars[0] == '(' && chars[2] == ')' &&
                    chars[1].is_uppercase() && chars[3] == ' ') {

                        task.priority = chars[1];
                        todo = todo.slice(4, todo.len()).to_owned();
                }

                task.todo = todo;

                return task;
        }
}

#[test]
fn simple_todo_create_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , '^');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task");
        assert_eq!(t.raw_todo , ~"some important task");
}

#[test]
fn simple_todo_with_priority_test() {
        let t = Task::create(~"(A) some important task with priority", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , 'A');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task with priority");
        assert_eq!(t.raw_todo , ~"(A) some important task with priority");

        let ts = Task::create(~"(a) some important task with priority", 1);

        assert_eq!(ts.id , 1);
        assert_eq!(ts.priority , '^');
        assert_eq!(ts.finished , false);
        assert_eq!(ts.todo , ~"(a) some important task with priority");
        assert_eq!(ts.raw_todo , ~"(a) some important task with priority");

}

#[test]
fn todo_to_str_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.to_str(), ~"1 false (^) 'some important task' \"some important task\"")

        let x = Task::create(~"x some important task", 1);
        assert_eq!(x.to_str(), ~"1 true (^) 'some important task' \"x some important task\"")
}
<commit_msg>started with dates<commit_after>
pub struct Task {
        id: int,
        todo: ~str,
        priority: char,
        finished: bool,
        raw_todo: ~str,
        create_date ~str,
        finish_date ~str,
        contexts: & 'static [~str],
        projects: & 'static [~str]
}

impl Task {
        pub fn to_str(&self) -> ~str {
                return format!("{:d} {:b} ({:c}) '{:s}' \"{:s}\"",
                                self.id, self.finished, self.priority,
                                self.todo, self.raw_todo)
        }

        pub fn create(text: ~str, id: int) -> Task {
                let mut task = Task{id: id, raw_todo: text.clone(), todo: text.clone(),
                                        finished: false, priority: '^',
                                        contexts: &[], projects: &[]};

                let mut todo = text.clone();

                if (todo.slice(0, 2) == "x ") {
                        task.finished = true;
                        todo = todo.slice(2, todo.len()).to_owned();
                }

                let mut chars = todo.chars().to_owned_vec();
                if (chars[0] == '(' && chars[2] == ')' &&
                    chars[1].is_uppercase() && chars[3] == ' ') {

                        task.priority = chars[1];
                        todo = todo.slice(4, todo.len()).to_owned();
                }

                chars = todo.chars().to_owned_vec();

                if (chars[10] == ' ' && chars[4] == '-' && chars[7] == '-') {
                        task.create_date = todo.slice(0, 10).to_owned();
                        todo = todo.slice(11, todo.len()).to_owned();
                }

                task.todo = todo;

                return task;
        }
}

#[test]
fn simple_todo_create_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , '^');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task");
        assert_eq!(t.raw_todo , ~"some important task");
}

#[test]
fn simple_todo_with_priority_test() {
        let t = Task::create(~"(A) some important task with priority", 1);

        assert_eq!(t.id , 1);
        assert_eq!(t.priority , 'A');
        assert_eq!(t.finished , false);
        assert_eq!(t.todo , ~"some important task with priority");
        assert_eq!(t.raw_todo , ~"(A) some important task with priority");

        let ts = Task::create(~"(a) some important task with priority", 1);

        assert_eq!(ts.id , 1);
        assert_eq!(ts.priority , '^');
        assert_eq!(ts.finished , false);
        assert_eq!(ts.todo , ~"(a) some important task with priority");
        assert_eq!(ts.raw_todo , ~"(a) some important task with priority");

}

#[test]
fn todo_to_str_test() {
        let t = Task::create(~"some important task", 1);

        assert_eq!(t.to_str(), ~"1 false (^) 'some important task' \"some important task\"")

        let x = Task::create(~"x some important task", 1);
        assert_eq!(x.to_str(), ~"1 true (^) 'some important task' \"x some important task\"")
}
<|endoftext|>


<commit_before>

struct Todo {
        id: int,
        todo: ~str,
        priority: char,
        finished: bool,
        raw_todo: ~str,
        contexts: & 'static [~str],
        projects: & 'static [~str]
}

impl Todo {
        fn to_str(&self) -> ~str {
                return format!("{:d} {:b} ({:c}) '{:s}' \"{:s}\"",
                                self.id, self.finished, self.priority,
                                self.todo, self.raw_todo)
        }

        fn create(todo: ~str, id: int) -> Todo {
                let mut task = Todo{id: id, raw_todo: todo.clone(), todo: ~"",
                                        finished: false, priority: '^',
                                        contexts: &[], projects: &[]};

                if (todo[0] == "x".as_bytes()[0]) {
                        task.finished = true;
                }

                task.todo = todo.clone();

                return task;
        }
}

#[test]
fn simple_todo_create_test() {
        let t = Todo::create(~"some important task", 1);

        assert!(t.id == 1);
        assert!(t.priority == '^');
        assert!(t.finished == false);
        assert!(t.todo == ~"some important task");
        assert!(t.raw_todo == ~"some important task");
}

fn main () {

        let t = Todo::create(~"some important task", 1);
        println(t.to_str());
        let x = Todo::create(~"x some important task", 1);
        println(x.to_str());
}
<commit_msg>todo_to_str_test<commit_after>

struct Todo {
        id: int,
        todo: ~str,
        priority: char,
        finished: bool,
        raw_todo: ~str,
        contexts: & 'static [~str],
        projects: & 'static [~str]
}

impl Todo {
        fn to_str(&self) -> ~str {
                return format!("{:d} {:b} ({:c}) '{:s}' \"{:s}\"",
                                self.id, self.finished, self.priority,
                                self.todo, self.raw_todo)
        }

        fn create(todo: ~str, id: int) -> Todo {
                let mut task = Todo{id: id, raw_todo: todo.clone(), todo: ~"",
                                        finished: false, priority: '^',
                                        contexts: &[], projects: &[]};

                if (todo[0] == "x".as_bytes()[0]) {
                        task.finished = true;
                }

                task.todo = todo.clone();

                return task;
        }
}

#[test]
fn simple_todo_create_test() {
        let t = Todo::create(~"some important task", 1);

        assert!(t.id == 1);
        assert!(t.priority == '^');
        assert!(t.finished == false);
        assert!(t.todo == ~"some important task");
        assert!(t.raw_todo == ~"some important task");
}

#[test]
fn todo_to_str_test() {
        let t = Todo::create(~"some important task", 1);

        assert!(t.to_str() == ~"1 false (^) 'some important task' \"some important task\"")
}

fn main () {

        let t = Todo::create(~"some important task", 1);
        println(t.to_str());
        let x = Todo::create(~"x some important task", 1);
        println(x.to_str());
}
<|endoftext|>


<commit_before>use std::fmt;
use std::path::{Path, PathBuf};

use async_trait::async_trait;
use bollard::container::LogOutput;
use bollard::Docker;
use futures::stream::BoxStream;
use futures::{StreamExt, TryFutureExt};
use nails::execution::ExitCode;
use store::Store;
use task_executor::Executor;
use workunit_store::{in_workunit, RunningWorkunit};

use crate::local::{
  apply_chroot, create_sandbox, prepare_workdir, setup_run_sh_script, CapturedWorkdir, ChildOutput,
  KeepSandboxes,
};
use crate::{
  Context, FallibleProcessResultWithPlatform, ImmutableInputs, NamedCaches, Platform, Process,
  ProcessError,
};

pub(crate) const SANDBOX_PATH_IN_CONTAINER: &str = "/pants-sandbox";
pub(crate) const NAMED_CACHES_PATH_IN_CONTAINER: &str = "/pants-named-caches";
pub(crate) const IMMUTABLE_INPUTS_PATH_IN_CONTAINER: &str = "/pants-immutable-inputs";

/// `CommandRunner` executes processes using a local Docker client.
pub struct CommandRunner {
  docker: Docker,
  store: Store,
  executor: Executor,
  work_dir_base: PathBuf,
  named_caches: NamedCaches,
  immutable_inputs: ImmutableInputs,
  keep_sandboxes: KeepSandboxes,
  image: String,
}

impl CommandRunner {
  pub fn new(
    store: Store,
    executor: Executor,
    work_dir_base: PathBuf,
    named_caches: NamedCaches,
    immutable_inputs: ImmutableInputs,
    keep_sandboxes: KeepSandboxes,
    image: String,
  ) -> Result<Self, String> {
    let docker = Docker::connect_with_local_defaults()
      .map_err(|err| format!("Failed to connect to local Docker: {err}"))?;
    Ok(CommandRunner {
      docker,
      store,
      executor,
      work_dir_base,
      named_caches,
      immutable_inputs,
      keep_sandboxes,
      image,
    })
  }
}

impl fmt::Debug for CommandRunner {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    f.debug_struct("docker::CommandRunner")
      .finish_non_exhaustive()
  }
}

#[async_trait]
impl super::CommandRunner for CommandRunner {
  async fn run(
    &self,
    context: Context,
    _workunit: &mut RunningWorkunit,
    req: Process,
  ) -> Result<FallibleProcessResultWithPlatform, ProcessError> {
    let req_debug_repr = format!("{:#?}", req);
    in_workunit!(
      "run_local_process_via_docker",
      req.level,
      // NB: See engine::nodes::NodeKey::workunit_level for more information on why this workunit
      // renders at the Process's level.
      desc = Some(req.description.clone()),
      |_workunit| async move {
        let mut workdir = create_sandbox(
          self.executor.clone(),
          &self.work_dir_base,
          &req.description,
          self.keep_sandboxes,
        )?;

        // Start working on a mutable version of the process.
        let mut req = req;

        // Update env, replacing `{chroot}` placeholders with `/pants-sandbox`. This is the mount point
        // for the sandbox directory within the Docker container.
        //
        // DOCKER-TODO: With cached containers, the destination in the container
        // will need to be unique within that container to avoid conflicts.
        apply_chroot(SANDBOX_PATH_IN_CONTAINER, &mut req);

        // Prepare the workdir.
        // DOCKER-NOTE: The input root will be bind mounted into the container.
        let exclusive_spawn = prepare_workdir(
          workdir.path().to_owned(),
          &req,
          req.input_digests.input_files.clone(),
          self.store.clone(),
          self.executor.clone(),
          &self.named_caches,
          &self.immutable_inputs,
          Some(Path::new(NAMED_CACHES_PATH_IN_CONTAINER)),
          Some(Path::new(IMMUTABLE_INPUTS_PATH_IN_CONTAINER)),
        )
        .await?;

        // DOCKER-TODO: Add a metric for local docker execution?
        // workunit.increment_counter(Metric::LocalExecutionRequests, 1);

        let res = self
          .run_and_capture_workdir(
            req.clone(),
            context,
            self.store.clone(),
            self.executor.clone(),
            workdir.path().to_owned(),
            (
              self.immutable_inputs.workdir().to_path_buf(),
              self.named_caches.base_dir().to_path_buf(),
            ),
            exclusive_spawn,
            req
              .platform_constraint
              .unwrap_or_else(|| Platform::current().unwrap()),
          )
          .map_err(|msg| {
            // Processes that experience no infrastructure issues should result in an "Ok" return,
            // potentially with an exit code that indicates that they failed (with more information
            // on stderr). Actually failing at this level indicates a failure to start or otherwise
            // interact with the process, which would generally be an infrastructure or implementation
            // error (something missing from the sandbox, incorrect permissions, etc).
            //
            // Given that this is expected to be rare, we dump the entire process definition in the
            // error.
            ProcessError::Unclassified(format!("Failed to execute: {}\n\n{}", req_debug_repr, msg))
          })
          .await;

        if self.keep_sandboxes == KeepSandboxes::Always
          || self.keep_sandboxes == KeepSandboxes::OnFailure
            && res.as_ref().map(|r| r.exit_code).unwrap_or(1) != 0
        {
          workdir.keep(&req.description);
          setup_run_sh_script(&req.env, &req.working_directory, &req.argv, workdir.path())?;
        }

        res
      }
    )
    .await
  }
}

#[async_trait]
impl CapturedWorkdir for CommandRunner {
  type WorkdirToken = (PathBuf, PathBuf);

  async fn run_in_workdir<'a, 'b, 'c>(
    &'a self,
    workdir_path: &'b Path,
    (immutable_inputs_workdir, named_caches_workdir): Self::WorkdirToken,
    req: Process,
    _exclusive_spawn: bool,
  ) -> Result<BoxStream<'c, Result<ChildOutput, String>>, String> {
    let env = req
      .env
      .iter()
      .map(|(key, value)| format!("{}={}", key, value))
      .collect::<Vec<_>>();

    let workdir_path_as_string = workdir_path
      .to_path_buf()
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert workdir_path due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    let immutable_inputs_workdir_as_string = immutable_inputs_workdir
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert immutable_inputs_workdir due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    let named_caches_workdir_as_string = named_caches_workdir
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert named_caches_workdir due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    let working_dir = req
      .working_directory
      .map(|relpath| Path::new(SANDBOX_PATH_IN_CONTAINER).join(&relpath))
      .unwrap_or_else(|| Path::new(SANDBOX_PATH_IN_CONTAINER).to_path_buf())
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert working directory due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    // DOCKER-TODO: Set creation options so we can set platform.

    let config = bollard::container::Config {
      env: Some(env),
      entrypoint: Some(req.argv),
      working_dir: Some(working_dir),
      // DOCKER-TODO: Is this necessary on linux hosts for allowing bind mount?
      // user: Some(format!("{}", unsafe { libc::geteuid() })),
      // user: Some("0".to_string()),
      host_config: Some(bollard_stubs::models::HostConfig {
        binds: Some(vec![
          format!("{}:{}", workdir_path_as_string, SANDBOX_PATH_IN_CONTAINER),
          // DOCKER-TODO: Consider making this bind mount read-only.
          format!(
            "{}:{}",
            immutable_inputs_workdir_as_string, IMMUTABLE_INPUTS_PATH_IN_CONTAINER
          ),
          format!(
            "{}:{}",
            named_caches_workdir_as_string, NAMED_CACHES_PATH_IN_CONTAINER
          ),
        ]),
        init: Some(true),
        ..bollard_stubs::models::HostConfig::default()
      }),
      image: Some(self.image.clone()),
      attach_stdout: Some(true),
      attach_stderr: Some(true),
      ..bollard::container::Config::default()
    };

    log::trace!("creating container with config: {:?}", &config);

    let container = self
      .docker
      .create_container::<&str, String>(None, config)
      .await
      .map_err(|err| format!("Failed to create Docker container: {:?}", err))?;

    log::trace!("created container {}", &container.id);

    // DOCKER-TODO: Use a drop guard to remove the container regardless of success/failure.
    // We cannot set `.host_config.auto_remove` in `config` above because containers that exit
    // early will no longer exist and thus this code will not be able to attach and capture
    // the output (which is the entire point of this executor ...).

    self
      .docker
      .start_container::<String>(&container.id, None)
      .await
      .map_err(|err| {
        format!(
          "Failed to start Docker container `{}`: {:?}",
          &container.id, err
        )
      })?;

    log::trace!("started container {}", &container.id);

    let attach_options = bollard::container::AttachContainerOptions::<String> {
      stdout: Some(true),
      stderr: Some(true),
      logs: Some(true), // stream any output that was missed between the start_container call and now
      stream: Some(true),
      ..bollard::container::AttachContainerOptions::default()
    };

    let attach_result = self
      .docker
      .attach_container(&container.id, Some(attach_options))
      .await
      .map_err(|err| {
        format!(
          "Failed to attach to Docker container `{}`: {:?}",
          &container.id, err
        )
      })?;

    log::trace!("attached to container {}", &container.id);

    let mut output_stream = attach_result.output.boxed();

    let wait_options = bollard::container::WaitContainerOptions {
      condition: "not-running",
    };
    let mut wait_stream = self
      .docker
      .wait_container(&container.id, Some(wait_options))
      .boxed();

    let container_id = container.id.to_owned();
    let result_stream = async_stream::stream! {
      loop {
        tokio::select! {
          Some(output_msg) = output_stream.next() => {
            match output_msg {
              Ok(LogOutput::StdOut { message }) => {
                log::trace!("container {} wrote {} bytes to stdout", &container_id, message.len());
                yield Ok(ChildOutput::Stdout(message));
              }
              Ok(LogOutput::StdErr { message }) => {
                log::trace!("container {} wrote {} bytes to stderr", &container_id, message.len());
                yield Ok(ChildOutput::Stderr(message));
              }
              _ => (),
            }
          }
          Some(wait_msg) = wait_stream.next() => {
            log::trace!("wait_container stream ({}): {:?}", &container_id, wait_msg);
            match wait_msg {
              Ok(r) => {
                // DOCKER-TODO: How does Docker distinguish signal versus exit code? Improve
                // `ChildResults` to better support exit code vs signal vs error message?
                let status_code = r.status_code;
                yield Ok(ChildOutput::Exit(ExitCode(status_code as i32)));
                break;
              }
              Err(err) => {
                // DOCKER-TODO: Consider a way to pass error messages back to child status collector.
                log::error!("Docker wait failure for container {}: {:?}", &container_id, err);
                yield Err(format!("Docker wait_container failure for container {}: {:?}", &container_id, err));
                break;
              }
            }
          }
        }
      }
    }
    .boxed();

    Ok(result_stream)
  }
}
<commit_msg>remove docker containers after execution finishes (#16747)<commit_after>use std::fmt;
use std::path::{Path, PathBuf};

use async_trait::async_trait;
use bollard::container::LogOutput;
use bollard::Docker;
use futures::stream::BoxStream;
use futures::{StreamExt, TryFutureExt};
use nails::execution::ExitCode;
use store::Store;
use task_executor::Executor;
use workunit_store::{in_workunit, RunningWorkunit};

use crate::local::{
  apply_chroot, create_sandbox, prepare_workdir, setup_run_sh_script, CapturedWorkdir, ChildOutput,
  KeepSandboxes,
};
use crate::{
  Context, FallibleProcessResultWithPlatform, ImmutableInputs, NamedCaches, Platform, Process,
  ProcessError,
};

pub(crate) const SANDBOX_PATH_IN_CONTAINER: &str = "/pants-sandbox";
pub(crate) const NAMED_CACHES_PATH_IN_CONTAINER: &str = "/pants-named-caches";
pub(crate) const IMMUTABLE_INPUTS_PATH_IN_CONTAINER: &str = "/pants-immutable-inputs";

/// `CommandRunner` executes processes using a local Docker client.
pub struct CommandRunner {
  docker: Docker,
  store: Store,
  executor: Executor,
  work_dir_base: PathBuf,
  named_caches: NamedCaches,
  immutable_inputs: ImmutableInputs,
  keep_sandboxes: KeepSandboxes,
  image: String,
}

impl CommandRunner {
  pub fn new(
    store: Store,
    executor: Executor,
    work_dir_base: PathBuf,
    named_caches: NamedCaches,
    immutable_inputs: ImmutableInputs,
    keep_sandboxes: KeepSandboxes,
    image: String,
  ) -> Result<Self, String> {
    let docker = Docker::connect_with_local_defaults()
      .map_err(|err| format!("Failed to connect to local Docker: {err}"))?;
    Ok(CommandRunner {
      docker,
      store,
      executor,
      work_dir_base,
      named_caches,
      immutable_inputs,
      keep_sandboxes,
      image,
    })
  }
}

impl fmt::Debug for CommandRunner {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    f.debug_struct("docker::CommandRunner")
      .finish_non_exhaustive()
  }
}

#[async_trait]
impl super::CommandRunner for CommandRunner {
  async fn run(
    &self,
    context: Context,
    _workunit: &mut RunningWorkunit,
    req: Process,
  ) -> Result<FallibleProcessResultWithPlatform, ProcessError> {
    let req_debug_repr = format!("{:#?}", req);
    in_workunit!(
      "run_local_process_via_docker",
      req.level,
      // NB: See engine::nodes::NodeKey::workunit_level for more information on why this workunit
      // renders at the Process's level.
      desc = Some(req.description.clone()),
      |_workunit| async move {
        let mut workdir = create_sandbox(
          self.executor.clone(),
          &self.work_dir_base,
          &req.description,
          self.keep_sandboxes,
        )?;

        // Start working on a mutable version of the process.
        let mut req = req;

        // Update env, replacing `{chroot}` placeholders with `/pants-sandbox`. This is the mount point
        // for the sandbox directory within the Docker container.
        //
        // DOCKER-TODO: With cached containers, the destination in the container
        // will need to be unique within that container to avoid conflicts.
        apply_chroot(SANDBOX_PATH_IN_CONTAINER, &mut req);

        // Prepare the workdir.
        // DOCKER-NOTE: The input root will be bind mounted into the container.
        let exclusive_spawn = prepare_workdir(
          workdir.path().to_owned(),
          &req,
          req.input_digests.input_files.clone(),
          self.store.clone(),
          self.executor.clone(),
          &self.named_caches,
          &self.immutable_inputs,
          Some(Path::new(NAMED_CACHES_PATH_IN_CONTAINER)),
          Some(Path::new(IMMUTABLE_INPUTS_PATH_IN_CONTAINER)),
        )
        .await?;

        // DOCKER-TODO: Add a metric for local docker execution?
        // workunit.increment_counter(Metric::LocalExecutionRequests, 1);

        let res = self
          .run_and_capture_workdir(
            req.clone(),
            context,
            self.store.clone(),
            self.executor.clone(),
            workdir.path().to_owned(),
            (
              self.immutable_inputs.workdir().to_path_buf(),
              self.named_caches.base_dir().to_path_buf(),
            ),
            exclusive_spawn,
            req
              .platform_constraint
              .unwrap_or_else(|| Platform::current().unwrap()),
          )
          .map_err(|msg| {
            // Processes that experience no infrastructure issues should result in an "Ok" return,
            // potentially with an exit code that indicates that they failed (with more information
            // on stderr). Actually failing at this level indicates a failure to start or otherwise
            // interact with the process, which would generally be an infrastructure or implementation
            // error (something missing from the sandbox, incorrect permissions, etc).
            //
            // Given that this is expected to be rare, we dump the entire process definition in the
            // error.
            ProcessError::Unclassified(format!("Failed to execute: {}\n\n{}", req_debug_repr, msg))
          })
          .await;

        if self.keep_sandboxes == KeepSandboxes::Always
          || self.keep_sandboxes == KeepSandboxes::OnFailure
            && res.as_ref().map(|r| r.exit_code).unwrap_or(1) != 0
        {
          workdir.keep(&req.description);
          setup_run_sh_script(&req.env, &req.working_directory, &req.argv, workdir.path())?;
        }

        res
      }
    )
    .await
  }
}

#[async_trait]
impl CapturedWorkdir for CommandRunner {
  type WorkdirToken = (PathBuf, PathBuf);

  async fn run_in_workdir<'a, 'b, 'c>(
    &'a self,
    workdir_path: &'b Path,
    (immutable_inputs_workdir, named_caches_workdir): Self::WorkdirToken,
    req: Process,
    _exclusive_spawn: bool,
  ) -> Result<BoxStream<'c, Result<ChildOutput, String>>, String> {
    let env = req
      .env
      .iter()
      .map(|(key, value)| format!("{}={}", key, value))
      .collect::<Vec<_>>();

    let workdir_path_as_string = workdir_path
      .to_path_buf()
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert workdir_path due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    let immutable_inputs_workdir_as_string = immutable_inputs_workdir
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert immutable_inputs_workdir due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    let named_caches_workdir_as_string = named_caches_workdir
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert named_caches_workdir due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    let working_dir = req
      .working_directory
      .map(|relpath| Path::new(SANDBOX_PATH_IN_CONTAINER).join(&relpath))
      .unwrap_or_else(|| Path::new(SANDBOX_PATH_IN_CONTAINER).to_path_buf())
      .into_os_string()
      .into_string()
      .map_err(|s| {
        format!(
          "Unable to convert working directory due to non UTF-8 characters: {:?}",
          s
        )
      })?;

    // DOCKER-TODO: Set creation options so we can set platform.

    let config = bollard::container::Config {
      env: Some(env),
      entrypoint: Some(req.argv),
      working_dir: Some(working_dir),
      // DOCKER-TODO: Is this necessary on linux hosts for allowing bind mount?
      // user: Some(format!("{}", unsafe { libc::geteuid() })),
      // user: Some("0".to_string()),
      host_config: Some(bollard_stubs::models::HostConfig {
        binds: Some(vec![
          format!("{}:{}", workdir_path_as_string, SANDBOX_PATH_IN_CONTAINER),
          // DOCKER-TODO: Consider making this bind mount read-only.
          format!(
            "{}:{}",
            immutable_inputs_workdir_as_string, IMMUTABLE_INPUTS_PATH_IN_CONTAINER
          ),
          format!(
            "{}:{}",
            named_caches_workdir_as_string, NAMED_CACHES_PATH_IN_CONTAINER
          ),
        ]),
        init: Some(true),
        ..bollard_stubs::models::HostConfig::default()
      }),
      image: Some(self.image.clone()),
      attach_stdout: Some(true),
      attach_stderr: Some(true),
      ..bollard::container::Config::default()
    };

    log::trace!("creating container with config: {:?}", &config);

    let container = self
      .docker
      .create_container::<&str, String>(None, config)
      .await
      .map_err(|err| format!("Failed to create Docker container: {:?}", err))?;

    log::trace!("created container {}", &container.id);

    self
      .docker
      .start_container::<String>(&container.id, None)
      .await
      .map_err(|err| {
        format!(
          "Failed to start Docker container `{}`: {:?}",
          &container.id, err
        )
      })?;

    log::trace!("started container {}", &container.id);

    let attach_options = bollard::container::AttachContainerOptions::<String> {
      stdout: Some(true),
      stderr: Some(true),
      logs: Some(true), // stream any output that was missed between the start_container call and now
      stream: Some(true),
      ..bollard::container::AttachContainerOptions::default()
    };

    let attach_result = self
      .docker
      .attach_container(&container.id, Some(attach_options))
      .await
      .map_err(|err| {
        format!(
          "Failed to attach to Docker container `{}`: {:?}",
          &container.id, err
        )
      })?;

    log::trace!("attached to container {}", &container.id);

    let mut output_stream = attach_result.output.boxed();

    let wait_options = bollard::container::WaitContainerOptions {
      condition: "not-running",
    };
    let mut wait_stream = self
      .docker
      .wait_container(&container.id, Some(wait_options))
      .boxed();

    let container_id = container.id.to_owned();
    let keep_sandboxes = self.keep_sandboxes;
    let docker = self.docker.clone();
    let result_stream = async_stream::stream! {
      let was_success = loop {
        tokio::select! {
          Some(output_msg) = output_stream.next() => {
            match output_msg {
              Ok(LogOutput::StdOut { message }) => {
                log::trace!("container {} wrote {} bytes to stdout", &container_id, message.len());
                yield Ok(ChildOutput::Stdout(message));
              }
              Ok(LogOutput::StdErr { message }) => {
                log::trace!("container {} wrote {} bytes to stderr", &container_id, message.len());
                yield Ok(ChildOutput::Stderr(message));
              }
              _ => (),
            }
          }
          Some(wait_msg) = wait_stream.next() => {
            log::trace!("wait_container stream ({}): {:?}", &container_id, wait_msg);
            match wait_msg {
              Ok(r) => {
                // DOCKER-TODO: How does Docker distinguish signal versus exit code? Improve
                // `ChildResults` to better support exit code vs signal vs error message?
                let status_code = r.status_code;
                yield Ok(ChildOutput::Exit(ExitCode(status_code as i32)));
                break status_code == 0;
              }
              Err(err) => {
                // DOCKER-TODO: Consider a way to pass error messages back to child status collector.
                log::error!("Docker wait failure for container {}: {:?}", &container_id, err);
                yield Err(format!("Docker wait_container failure for container {}: {:?}", &container_id, err));
                break false;
              }
            }
          }
        }
      };

      let do_remove_container = match keep_sandboxes {
        KeepSandboxes::Always => false,
        KeepSandboxes::Never => true,
        KeepSandboxes::OnFailure => !was_success,
      };

      if do_remove_container {
        let remove_options = bollard::container::RemoveContainerOptions {
          force: true,
          ..bollard::container::RemoveContainerOptions::default()
        };

        let remove_result = docker
          .remove_container(&container_id, Some(remove_options))
          .await
          .map_err(|err| format!("Failed to remove container `{}`: {:?}", &container_id, err));

        if let Err(err) = remove_result {
          log::warn!("{}", err);
        }
      }
    }
    .boxed();

    Ok(result_stream)
  }
}
<|endoftext|>


<commit_before>use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::collections::HashSet;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

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

use common::Solution;

type Coordinate = (usize, usize);

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Copy, Clone, Default)]
struct Node {
    pos: Coordinate,
    capacity: usize,
    used: usize,
}

impl Node {
    fn fits(self, other: Node) -> bool {
        self.capacity - self.used >= other.used
    }
}

fn get_grid<T>(list: &[T], size: usize, (x, y): Coordinate) -> &T {
    &list[x + size * y]
}

fn get_mut_grid<T>(list: &mut [T], size: usize, (x, y): Coordinate) -> &mut T {
    &mut list[x + size * y]
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
struct State {
    pos: Coordinate,
    empty_pos: Coordinate,
    grid: Vec<usize>,
}

impl State {
    fn estimate(&self) -> usize {
        self.pos.0 + self.pos.1
            + self.pos.0.max(self.empty_pos.0) + self.pos.1.max(self.empty_pos.1)
            - self.pos.0.min(self.empty_pos.0) - self.pos.1.min(self.empty_pos.1)
            - 1
    }
}

#[derive(Default)]
pub struct Day22 {
    nodes: Vec<Node>,
}

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

    fn read_input(&mut self, input: &mut Read) {
        let reader = BufReader::new(input);
        let matcher = Regex::new(r"/dev/grid/node-x(\d+)-y(\d+)(?:\s+)(\d+)T(?:\s+)+(\d+)T").unwrap();

        for line in reader.lines() {
            let line = line.unwrap();
            let node = match matcher.captures(&line) {
                Some(captures) => Node {
                    pos: (captures[1].parse().unwrap(), captures[2].parse().unwrap()),
                    capacity: captures[3].parse().unwrap(),
                    used: captures[4].parse().unwrap(),
                },

                None => continue
            };

            self.nodes.push(node)
        }

        self.nodes.sort_unstable();

        assert_ne!(0, self.nodes.len())
    }

    fn build_grid(&self) -> (usize, Vec<usize>, Vec<usize>) {
        let size = self.nodes.iter().map(|node| node.pos.0).max().unwrap_or(0) + 1;

        let mut used = vec![0usize; self.nodes.len()];
        let mut capacity = used.clone();

        for node in &self.nodes {
            let pos = node.pos;
            *get_mut_grid(&mut used, size, pos) = node.used;
            *get_mut_grid(&mut capacity, size, pos) = node.capacity;
        }

        (size, used, capacity)
    }
}

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

        let fitting = iproduct!(&self.nodes, &self.nodes)
            .filter(|&(_, b)| b.used > 0)
            .filter(|&(a, b)| a != b)
            .filter(|&(a, b)| a.fits(*b))
            .count();

        fitting.to_string()
    }

    fn part2(&mut self, input: &mut Read) -> String {
        self.read_input(input);

        let (size, used, capacity) = self.build_grid();

        let (empty_pos, _) = used.iter()
            .find_position(|&&used| used == 0)
            .unwrap();

        let height = used.len() / size;

        let state = State {
            pos: (size - 1, 0),
            empty_pos: (empty_pos % size, empty_pos / size),
            grid: used,
        };

        let mut visited = HashSet::new();
        visited.insert((state.pos, state.empty_pos));

        let mut todo = BinaryHeap::new();
        todo.push(Reverse((state.estimate(), 0usize, state)));

        while let Some(Reverse((_, dist, state))) = todo.pop() {
            if state.pos == (0, 0) {
                return dist.to_string();
            }
            let (xe, ye) = state.empty_pos;
            let empty_capacity = *get_grid(&capacity, size, state.empty_pos);

            let valid_x = [if xe > 0 { Some(xe - 1) } else { None }, Some(xe), Some(xe), if xe < size - 1 { Some(xe + 1) } else { None }];
            let valid_y = [Some(ye), if ye > 0 { Some(ye - 1) } else { None }, if ye < height - 1 { Some(ye + 1) } else { None }, Some(ye)];
            for (x, y) in valid_x.iter().zip(valid_y.iter()) {
                if x.is_none() || y.is_none() {
                    continue
                }
                let switch = (x.unwrap(), y.unwrap());
                let contents = *get_grid(&state.grid, size, switch);

                if contents >  empty_capacity {
                    // Not enough capacity
                    continue;
                }

                let mut new_state = State {
                    grid: state.grid.clone(),
                    empty_pos: switch,
                    pos: if switch == state.pos { state.empty_pos } else { state.pos },
                };

                if !visited.contains(&(new_state.pos, new_state.empty_pos)) {
                    visited.insert((new_state.pos, new_state.empty_pos));
                    *get_mut_grid(&mut new_state.grid, size, switch) = 0;
                    *get_mut_grid(&mut new_state.grid, size, state.empty_pos) = contents;

                    todo.push(Reverse((dist + 1 + new_state.estimate(), dist + 1, new_state)));
                }
            }
        }

        unreachable!("Did not arrive at an end state")
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day22::Day22;

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

    #[test]
    fn sample_part2() {
        let mut instance = Day22::new();
        let result = instance.part2(&mut SAMPLE_INPUT);
        assert_eq!("7", result);
    }
}
<commit_msg>Small simplification.<commit_after>use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::collections::HashSet;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

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

use common::Solution;

type Coordinate = (usize, usize);

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Copy, Clone, Default)]
struct Node {
    pos: Coordinate,
    capacity: usize,
    used: usize,
}

impl Node {
    fn fits(self, other: Node) -> bool {
        self.capacity - self.used >= other.used
    }
}

fn get_grid<T>(list: &[T], size: usize, (x, y): Coordinate) -> &T {
    &list[x + size * y]
}

fn get_mut_grid<T>(list: &mut [T], size: usize, (x, y): Coordinate) -> &mut T {
    &mut list[x + size * y]
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
struct State {
    pos: Coordinate,
    empty_pos: Coordinate,
    grid: Vec<usize>,
}

impl State {
    fn estimate(&self) -> usize {
        self.pos.0 + self.pos.1
            + self.pos.0.max(self.empty_pos.0) + self.pos.1.max(self.empty_pos.1)
            - self.pos.0.min(self.empty_pos.0) - self.pos.1.min(self.empty_pos.1)
            - 1
    }
}

#[derive(Default)]
pub struct Day22 {
    nodes: Vec<Node>,
}

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

    fn read_input(&mut self, input: &mut Read) {
        let reader = BufReader::new(input);
        let matcher = Regex::new(r"/dev/grid/node-x(\d+)-y(\d+)(?:\s+)(\d+)T(?:\s+)+(\d+)T").unwrap();

        for line in reader.lines() {
            let line = line.unwrap();
            let node = match matcher.captures(&line) {
                Some(captures) => Node {
                    pos: (captures[1].parse().unwrap(), captures[2].parse().unwrap()),
                    capacity: captures[3].parse().unwrap(),
                    used: captures[4].parse().unwrap(),
                },

                None => continue
            };

            self.nodes.push(node)
        }

        self.nodes.sort_unstable();

        assert_ne!(0, self.nodes.len())
    }

    fn build_grid(&self) -> (usize, Vec<usize>, Vec<usize>) {
        let size = self.nodes.iter().map(|node| node.pos.0).max().unwrap_or(0) + 1;

        let mut used = vec![0usize; self.nodes.len()];
        let mut capacity = used.clone();

        for node in &self.nodes {
            let pos = node.pos;
            *get_mut_grid(&mut used, size, pos) = node.used;
            *get_mut_grid(&mut capacity, size, pos) = node.capacity;
        }

        (size, used, capacity)
    }
}

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

        let fitting = iproduct!(&self.nodes, &self.nodes)
            .filter(|&(_, b)| b.used > 0)
            .filter(|&(a, b)| a != b)
            .filter(|&(a, b)| a.fits(*b))
            .count();

        fitting.to_string()
    }

    fn part2(&mut self, input: &mut Read) -> String {
        self.read_input(input);

        let (size, used, capacity) = self.build_grid();

        let (empty_pos, _) = used.iter()
            .find_position(|&&used| used == 0)
            .unwrap();

        let height = used.len() / size;

        let state = State {
            pos: (size - 1, 0),
            empty_pos: (empty_pos % size, empty_pos / size),
            grid: used,
        };

        let mut visited = HashSet::new();
        visited.insert((state.pos, state.empty_pos));

        let mut todo = BinaryHeap::new();
        todo.push(Reverse((state.estimate(), 0usize, state)));

        while let Some(Reverse((_, dist, state))) = todo.pop() {
            if state.pos == (0, 0) {
                return dist.to_string();
            }
            let (xe, ye) = state.empty_pos;
            let empty_capacity = *get_grid(&capacity, size, state.empty_pos);

            let valid_x = [xe.checked_sub(1), Some(xe), Some(xe), if xe < size - 1 { Some(xe + 1) } else { None }];
            let valid_y = [Some(ye), ye.checked_sub(1), if ye < height - 1 { Some(ye + 1) } else { None }, Some(ye)];
            for (x, y) in valid_x.iter().zip(valid_y.iter()) {
                if x.is_none() || y.is_none() {
                    continue;
                }
                let switch = (x.unwrap(), y.unwrap());
                let contents = *get_grid(&state.grid, size, switch);

                if contents > empty_capacity {
                    // Not enough capacity
                    continue;
                }

                let mut new_state = State {
                    grid: state.grid.clone(),
                    empty_pos: switch,
                    pos: if switch == state.pos { state.empty_pos } else { state.pos },
                };

                if !visited.contains(&(new_state.pos, new_state.empty_pos)) {
                    visited.insert((new_state.pos, new_state.empty_pos));
                    *get_mut_grid(&mut new_state.grid, size, switch) = 0;
                    *get_mut_grid(&mut new_state.grid, size, state.empty_pos) = contents;

                    todo.push(Reverse((dist + 1 + new_state.estimate(), dist + 1, new_state)));
                }
            }
        }

        unreachable!("Did not arrive at an end state")
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day22::Day22;

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

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


<commit_before>#![feature(globs)]
extern crate geoip2;
extern crate maxminddb;
extern crate serialize;

use std::io::net::ip::IpAddr;
use str::FromStr;
use serialize::Decodable;

use geoip2::Country;

fn main() {
    let r = maxminddb::Reader::open("GeoLite2-City.mmdb").unwrap();
    let ip: IpAddr = FromStr::from_str("128.101.101.101").unwrap();
    let dr = r.lookup(ip);
    //print!("{}", dr)

    let mut decoder = maxminddb::Decoder::new(dr.unwrap());
    let decoded_object: Country = match Decodable::decode(&mut decoder) {
        Ok(v) => v,
        Err(e) => panic!("Decoding error: {}", e)
    }; // create the final object
    print!("{}\n", decoded_object);

}
<commit_msg>Fix example<commit_after>#![feature(globs)]
extern crate geoip2;
extern crate maxminddb;
extern crate serialize;

use std::io::net::ip::IpAddr;
use std::str::FromStr;
use serialize::Decodable;

use geoip2::Country;

fn main() {
    let r = maxminddb::Reader::open("GeoLite2-City.mmdb").unwrap();
    let ip: IpAddr = FromStr::from_str("128.101.101.101").unwrap();
    let dr = r.lookup(ip);
    //print!("{}", dr)

    let mut decoder = maxminddb::Decoder::new(dr.unwrap());
    let decoded_object: Country = match Decodable::decode(&mut decoder) {
        Ok(v) => v,
        Err(e) => panic!("Decoding error: {}", e)
    }; // create the final object
    print!("{}\n", decoded_object);

}
<|endoftext|>


<commit_before>use ::ast::*;
use ::ast::ExprNode::*;
use super::*;
use super::parser_combinators::{Parser,parser};
use std::char;

#[test]
fn test_basic_ident() {
    assert_eq!(
        parser(expr).parse("ident"),
        Ok((Name(NameNode { name: "ident".to_string() }), ""))
        );
    assert_eq!(
        parser(expr).parse("a"),
        Ok((Name(NameNode { name: "a".to_string() }), ""))
        );
    assert_eq!(
        parser(expr).parse("ident_With\\special!Chars:~-+"),
        Ok((Name(NameNode { name: "ident_With\\special!Chars:~-+".to_string() }), ""))
        );
}

#[test]
fn test_basic_sexpr() {
    assert_eq!(
        parser(expr).parse("(ident arg1 arg2)"),
        Ok((SExpr(SExprNode {
            operator: NameNode { name: "ident".to_string() },
            operands: vec![
                Name(NameNode { name: "arg1".to_string() }),
                Name(NameNode { name: "arg2".to_string() })
            ]
        }), ""))
        );
}

#[test]
fn test_parse_sint_pos() {
    assert_eq!(
        parser(number).parse("1234"),
        Ok((NumNode::IntConst(IntNode { value: 1234isize }), ""))
        );
    assert_eq!(
        parser(number).parse("#d1234"),
        Ok((NumNode::IntConst(IntNode { value: 1234isize }), ""))
        );
    assert_eq!(
        parser(number).parse("#D1234"),
        Ok((NumNode::IntConst(IntNode { value: 1234isize }), ""))
        );
}

#[test]
fn test_parse_sint_neg() {
    assert_eq!(
        parser(number).parse("-1234"),
        Ok((NumNode::IntConst(IntNode { value: -1234isize }), ""))
        );
}

#[test]
fn test_parse_sint_hex() {
    assert_eq!(
        parser(number).parse("#x0ff"),
        Ok((NumNode::IntConst(IntNode { value: 0x0ffisize }), ""))
        );
    assert_eq!(
        parser(number).parse("#X0FF"),
        Ok((NumNode::IntConst(IntNode { value: 0x0ffisize }), ""))
        );
}
/* // Currently unsupported
#[test]
fn test_parse_sint_bin_upper() {
    assert_eq!(
        parser(number).parse("0B01"),
        Ok((NumNode::IntConst(IntNode { value: 0b01isize }), ""))
        );
    assert_eq!(
        parser(number).parse("0b01"),
        Ok((NumNode::IntConst(IntNode { value: 0b01isize }), ""))
        );
}*/

#[test]
fn test_parse_uint() {
    assert_eq!(
        parser(number).parse("1234u"),
        Ok((NumNode::UIntConst(UIntNode { value: 1234usize }), ""))
        );
    assert_eq!(
        parser(number).parse("4321U"),
        Ok((NumNode::UIntConst(UIntNode { value: 4321usize }), ""))
        );
}

#[test]
fn test_parse_uint_hex() {
    assert_eq!(
        parser(number).parse("#x0ffu"),
        Ok((NumNode::UIntConst(UIntNode { value: 0x0ffusize }), ""))
        );
    assert_eq!(
        parser(number).parse("#X0FFu"),
        Ok((NumNode::UIntConst(UIntNode { value: 0x0ffusize }), ""))
        );
}

#[test]
fn test_parse_float() {
    assert_eq!(
        parser(number).parse("1.0"),
        Ok((NumNode::FloatConst(FloatNode { value: 1.0f64 }), ""))
        );/* // Unsupported
    assert_eq!(
        parser(number).parse("1f").unwrap(),
        (NumNode::FloatConst(FloatNode { value: 1.0f64 }), "")
        );
    assert_eq!(
        parser(number).parse("22.2222").unwrap(),
        (NumNode::FloatConst(FloatNode { value: 22.2222f64 }), "")
        );
    assert_eq!(
        parser(number).parse("22.2222f").unwrap(),
        (NumNode::FloatConst(FloatNode { value: 22.2222f64 }), "")
        );*/
}

#[test]
fn test_parse_bool() {
    assert_eq!(
        parser(bool_const).parse("#t"),
        Ok((BoolNode { value: true}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("#T"),
        Ok((BoolNode { value: true}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("true"),
        Ok((BoolNode { value: true}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("#f"),
        Ok((BoolNode { value: false}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("#F"),
        Ok((BoolNode { value: false}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("false"),
        Ok((BoolNode { value: false}, ""))
        );
}

#[test]
fn test_parse_char() {
    assert_eq!(
        parser(character).parse("#\\c"),
        Ok((CharNode { value: 'c'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\A"),
        Ok((CharNode { value: 'A'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\tab"),
        Ok((CharNode { value: '\t'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\newline"),
        Ok((CharNode { value: '\n'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\nul"),
        Ok((CharNode { value: char::from_u32(0x0000).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\backspace"),
        Ok((CharNode { value: char::from_u32(0x0008).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\vtab"),
        Ok((CharNode { value: char::from_u32(0x000B).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\page"),
        Ok((CharNode { value: char::from_u32(0x000C).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\return"),
        Ok((CharNode { value: char::from_u32(0x000D).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\esc"),
        Ok((CharNode { value: char::from_u32(0x001B).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\delete"),
        Ok((CharNode { value: char::from_u32(0x007F).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\alarm"),
        Ok((CharNode { value: char::from_u32(0x0007).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\linefeed"),
        Ok((CharNode { value: '\n'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\space"),
        Ok((CharNode { value: ' '}, ""))
        );
}

<commit_msg>Add test for hex char constant lexing<commit_after>use ::ast::*;
use ::ast::ExprNode::*;
use super::*;
use super::parser_combinators::{Parser,parser};
use std::char;

#[test]
fn test_basic_ident() {
    assert_eq!(
        parser(expr).parse("ident"),
        Ok((Name(NameNode { name: "ident".to_string() }), ""))
        );
    assert_eq!(
        parser(expr).parse("a"),
        Ok((Name(NameNode { name: "a".to_string() }), ""))
        );
    assert_eq!(
        parser(expr).parse("ident_With\\special!Chars:~-+"),
        Ok((Name(NameNode { name: "ident_With\\special!Chars:~-+".to_string() }), ""))
        );
}

#[test]
fn test_basic_sexpr() {
    assert_eq!(
        parser(expr).parse("(ident arg1 arg2)"),
        Ok((SExpr(SExprNode {
            operator: NameNode { name: "ident".to_string() },
            operands: vec![
                Name(NameNode { name: "arg1".to_string() }),
                Name(NameNode { name: "arg2".to_string() })
            ]
        }), ""))
        );
}

#[test]
fn test_parse_sint_pos() {
    assert_eq!(
        parser(number).parse("1234"),
        Ok((NumNode::IntConst(IntNode { value: 1234isize }), ""))
        );
    assert_eq!(
        parser(number).parse("#d1234"),
        Ok((NumNode::IntConst(IntNode { value: 1234isize }), ""))
        );
    assert_eq!(
        parser(number).parse("#D1234"),
        Ok((NumNode::IntConst(IntNode { value: 1234isize }), ""))
        );
}

#[test]
fn test_parse_sint_neg() {
    assert_eq!(
        parser(number).parse("-1234"),
        Ok((NumNode::IntConst(IntNode { value: -1234isize }), ""))
        );
}

#[test]
fn test_parse_sint_hex() {
    assert_eq!(
        parser(number).parse("#x0ff"),
        Ok((NumNode::IntConst(IntNode { value: 0x0ffisize }), ""))
        );
    assert_eq!(
        parser(number).parse("#X0FF"),
        Ok((NumNode::IntConst(IntNode { value: 0x0ffisize }), ""))
        );
}
/* // Currently unsupported
#[test]
fn test_parse_sint_bin_upper() {
    assert_eq!(
        parser(number).parse("0B01"),
        Ok((NumNode::IntConst(IntNode { value: 0b01isize }), ""))
        );
    assert_eq!(
        parser(number).parse("0b01"),
        Ok((NumNode::IntConst(IntNode { value: 0b01isize }), ""))
        );
}*/

#[test]
fn test_parse_uint() {
    assert_eq!(
        parser(number).parse("1234u"),
        Ok((NumNode::UIntConst(UIntNode { value: 1234usize }), ""))
        );
    assert_eq!(
        parser(number).parse("4321U"),
        Ok((NumNode::UIntConst(UIntNode { value: 4321usize }), ""))
        );
}

#[test]
fn test_parse_uint_hex() {
    assert_eq!(
        parser(number).parse("#x0ffu"),
        Ok((NumNode::UIntConst(UIntNode { value: 0x0ffusize }), ""))
        );
    assert_eq!(
        parser(number).parse("#X0FFu"),
        Ok((NumNode::UIntConst(UIntNode { value: 0x0ffusize }), ""))
        );
}

#[test]
fn test_parse_float() {
    assert_eq!(
        parser(number).parse("1.0"),
        Ok((NumNode::FloatConst(FloatNode { value: 1.0f64 }), ""))
        );/* // Unsupported
    assert_eq!(
        parser(number).parse("1f").unwrap(),
        (NumNode::FloatConst(FloatNode { value: 1.0f64 }), "")
        );
    assert_eq!(
        parser(number).parse("22.2222").unwrap(),
        (NumNode::FloatConst(FloatNode { value: 22.2222f64 }), "")
        );
    assert_eq!(
        parser(number).parse("22.2222f").unwrap(),
        (NumNode::FloatConst(FloatNode { value: 22.2222f64 }), "")
        );*/
}

#[test]
fn test_parse_bool() {
    assert_eq!(
        parser(bool_const).parse("#t"),
        Ok((BoolNode { value: true}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("#T"),
        Ok((BoolNode { value: true}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("true"),
        Ok((BoolNode { value: true}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("#f"),
        Ok((BoolNode { value: false}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("#F"),
        Ok((BoolNode { value: false}, ""))
        );
    assert_eq!(
        parser(bool_const).parse("false"),
        Ok((BoolNode { value: false}, ""))
        );
}

#[test]
fn test_parse_char() {
    assert_eq!(
        parser(character).parse("#\\c"),
        Ok((CharNode { value: 'c'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\A"),
        Ok((CharNode { value: 'A'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\tab"),
        Ok((CharNode { value: '\t'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\newline"),
        Ok((CharNode { value: '\n'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\nul"),
        Ok((CharNode { value: char::from_u32(0x0000).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\backspace"),
        Ok((CharNode { value: char::from_u32(0x0008).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\vtab"),
        Ok((CharNode { value: char::from_u32(0x000B).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\page"),
        Ok((CharNode { value: char::from_u32(0x000C).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\return"),
        Ok((CharNode { value: char::from_u32(0x000D).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\esc"),
        Ok((CharNode { value: char::from_u32(0x001B).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\delete"),
        Ok((CharNode { value: char::from_u32(0x007F).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\alarm"),
        Ok((CharNode { value: char::from_u32(0x0007).unwrap()}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\linefeed"),
        Ok((CharNode { value: '\n'}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\space"),
        Ok((CharNode { value: ' '}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\x0020"),
        Ok((CharNode { value: ' '}, ""))
        );
    assert_eq!(
        parser(character).parse("#\\x001B"),
        Ok((CharNode { value: char::from_u32(0x001B).unwrap()}, ""))
        );
}

<|endoftext|>


<commit_before>use std::io;
use std::str::FromStr;
use std::collections::HashMap;

enum VMState {
	Definition,
	Execution
}

pub struct VM {
	state: VMState,
	op_stack: Vec<Operator>,
	data_stack: Vec<Operator>,
	words: HashMap<String, Vec<Operator>>,
	current_def_word: Option<String>,
}

impl VM{
	pub fn pop_data <'a> (&'a mut self) -> Operator {
		return self.data_stack.pop().expect("data stack is empty!");
	}

	pub fn push_data <'a> (&'a mut self, op: Operator) {
		self.data_stack.push(op);
	}

	pub fn pop_op <'a> (&'a mut self) -> Operator {
		return self.op_stack.pop().expect("operand stack is empty!");
	}

	pub fn push_op <'a> (&'a mut self, op: Operator) {
		match op {
			Operator::Define => {
				self.state = VMState::Definition;
				return;
			},
			Operator::EndDefine => {
				self.state = VMState::Execution;
				self.current_def_word = None;
				return;
			},
			_ => {}
		};
		let should_def = match self.state {
			VMState::Definition => true,
			VMState::Execution => false
		};
		if should_def {
			let current_def_word = self.push_def_word(op);
			self.current_def_word = Some(current_def_word);
		} else {
			self.op_stack.push(op);
		}
	}

	fn push_def_word <'a> (&'a mut self, op: Operator) -> String{
		let wrapped = self.current_def_word.as_ref();
		let new_key = match wrapped {
			Some(w) => {
				self.words.get_mut(w).expect("error").push(op);
				w.clone()
			},
			None => {
				let key = match op {
					Operator::Other(key) => key,
					_ => panic!("key is broken")
				};
				self.words.insert(key.clone(), Vec::new());
				key
			}
		};

		return new_key;
	}

	pub fn eval <'a> (&'a mut self) {
		let op = self.pop_op();
		match op {
			Operator::Add => self.add(),
			Operator::Subtract => self.subtract(),
			Operator::Multiply => self.multiply(),
			Operator::Divide => self.divide(),
			Operator::Modulus => self.modulus(),
			Operator::And => self.and(),
			Operator::Or => self.or(),
			Operator::Xor => self.xor(),
			Operator::Lshift => self.lshift(),
			Operator::Rshift => self.rshift(),
			Operator::Abs => self.abs(),
			Operator::Max => self.max(),
			Operator::Min => self.min(),
			Operator::Invert => self.invert(),
			Operator::Negate => self.negate(),
			Operator::Dup => self.dup(),
			Operator::Drop => self.drop(),
			Operator::Swap => self.swap(),
			Operator::Over => self.over(),
			Operator::Rot => self.rot(),
			Operator::Value(i) => self.push_data(Operator::Value(i)),
			Operator::Print => self.print(),
			Operator::Other(s) => self.custom(s),
			x => panic!("cannot eval: {}", x.stringify())
		}
	}

	fn custom <'a> (&'a mut self, word: String) {

	}

	fn pop_two_ints <'a> (&'a mut self) -> (i32, i32){
		let a = self.pop_one_int();
		let b = self.pop_one_int();
		(a, b)
	}

	fn pop_one_int <'a> (&'a mut self) -> i32 {
		let a = self.pop_data();
		match a {
			Operator::Value(j) => j,
			_ => panic!("can only add ints"),
		}
	}

	// Math Functions
	pub fn add <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x + y));
	}

	pub fn subtract <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x - y));
	}

	pub fn multiply <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x * y));
	}

	pub fn divide <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(y / x));
	}

	pub fn modulus <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x % y));
	}

	pub fn and <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x & y));
	}

	pub fn or <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x | y));
	}

	pub fn xor <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x ^ y));
	}

	pub fn lshift <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x << y));
	}

	pub fn rshift <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x >> y));
	}

	pub fn abs <'a> (&'a mut self) {
		let x = self.pop_one_int();
		self.push_data(Operator::Value(if x < 0 {-x} else {x}));
	}

	pub fn max <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(if y > x {y} else {x}));
	}

	pub fn min <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(if y < x {y} else {x}));
	}

	pub fn invert <'a> (&'a mut self) {
		let x = self.pop_one_int();
		self.push_data(Operator::Value(!x));
	}

	pub fn negate <'a> (&'a mut self) {
		let x = self.pop_one_int();
		self.push_data(Operator::Value(-x));
	}

	//Stack Functions
	pub fn dup <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = x.clone();
		self.push_data(x);
		self.push_data(y);
	}

	pub fn drop <'a> (&'a mut self) {
		self.pop_data();
	}

	pub fn swap <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = self.pop_data();
		self.push_data(x);
		self.push_data(y);
	}

	pub fn over <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = self.pop_data();
		let z = x.clone();
		self.push_data(x);
		self.push_data(y);
		self.push_data(z);
	}

	pub fn rot <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = self.pop_data();
		let z = self.pop_data();
		self.push_data(x);
		self.push_data(z);
		self.push_data(y);
	}

	//Debug Functions
	pub fn print <'a> (&'a mut self) {
		let x = self.pop_data();
		println!("{}",match x{
			Operator::Value(i) => i,
			_ => panic!("tried to print something other than a value"),
		});
	}
}

#[derive(Clone)]
pub enum Operator {
	//stack
	Dup,
	Drop,
	Swap,
	Over,
	Rot,

	//debug
	Print,

	//math
	Add,
	Subtract,
	Multiply,
	Divide,
	Modulus,
	And,
	Or,
	Xor,
	Lshift,
	Rshift,
	Abs,
	Max,
	Min,
	Invert,
	Negate,

	//custom
	Define,
	EndDefine,

	//values and custom calls
	Value(i32),
	Other(String),
}

impl Operator {
	fn stringify(self) -> String {
		match self{
			Operator::Add => format!("+"),
			Operator::Subtract => format!("-"),
			Operator::Multiply => format!("*"),
			Operator::Divide => format!("/"),
			Operator::Modulus => format!("mod"),
			Operator::And => format!("and"),
			Operator::Or => format!("or"),
			Operator::Xor => format!("xor"),
			Operator::Lshift => format!("<<"),
			Operator::Rshift => format!(">>"),
			Operator::Abs => format!("abs"),
			Operator::Max => format!("max"),
			Operator::Min => format!("min"),
			Operator::Invert => format!("invert"),
			Operator::Negate => format!("negate"),
			Operator::Dup => format!("dup"),
			Operator::Drop => format!("drop"),
			Operator::Swap => format!("swap"),
			Operator::Over => format!("over"),
			Operator::Rot => format!("rot"),
			Operator::Value(i) => format!("{}", i),
			Operator::Print => format!("."),
			Operator::Define => format!(":"),
			Operator::EndDefine => format!(";"),
			Operator::Other(_) => format!("other"),
		}
	}
}

fn parse(s: &str) -> Operator {
	match s {
		//stack operations:
		"dup" => Operator::Dup,
		"drop" => Operator::Drop,
		"swap" => Operator::Swap,
		"over" => Operator::Over,
		"rot" => Operator::Rot,

		//debug operations:
		"." => Operator::Print,

		//math operations:
		"+" => Operator::Add,
		"-" => Operator::Subtract,
		"*" => Operator::Multiply,
		"/" => Operator::Divide,
		"mod" => Operator::Modulus,
		"and" => Operator::And,
		"or" => Operator::Or,
		"xor" => Operator::Xor,
		"lshift" => Operator::Lshift,
		"<<" => Operator::Lshift,
		"rshift" => Operator::Rshift,
		">>" => Operator::Rshift,
		"abs" => Operator::Abs,
		"max" => Operator::Max,
		"min" => Operator::Min,
		"invert" => Operator::Invert,
		"negate" => Operator::Negate,

		//custom definitions
		":" => Operator::Define,
		";" => Operator::EndDefine,

		//value or custom operator
		x => {
			let y: Option<i32> = FromStr::from_str(x).ok();
			match y{
				Some(i) => Operator::Value(i),
				None => Operator::Other(format!("other")),
			}
		},
	}
}

fn main() {
	let x = io::stdin();
	let mut vm = VM {state: VMState::Execution,
					 op_stack: vec![],
					 data_stack: vec![],
					 words: HashMap::new(),
				     current_def_word: None};

	loop{
		let mut line: String = "".to_string();
		let _ = x.read_line(&mut line);

		if line.trim() == "" {break;}

		for o in line.rsplit(" "){
			let b = parse(o.trim());
			vm.push_op(b)
		}

		while vm.op_stack.len() > 0 {
			vm.eval();
		}
	}
}
<commit_msg>custom definitions working<commit_after>use std::io;
use std::str::FromStr;
use std::collections::HashMap;

enum VMState {
	Definition,
	Execution
}

pub struct VM {
	state: VMState,
	op_stack: Vec<Operator>,
	data_stack: Vec<Operator>,
	words: HashMap<String, Vec<Operator>>,
	current_def_word: Option<String>,
}

impl VM{
	pub fn pop_data <'a> (&'a mut self) -> Operator {
		let data = self.data_stack.pop().expect("data stack is empty!");
		return data;
	}

	pub fn push_data <'a> (&'a mut self, op: Operator) {
		self.data_stack.push(op);
	}

	pub fn pop_op <'a> (&'a mut self) -> Operator {
		let op = self.op_stack.pop().expect("operand stack is empty!");
		return op;
	}

	pub fn push_op <'a> (&'a mut self, op: Operator) {
		match op {
			Operator::Define => {
				self.state = VMState::Definition;
				return;
			},
			Operator::EndDefine => {
				self.state = VMState::Execution;
				self.current_def_word = None;
				return;
			},
			_ => {}
		};
		let should_def = match self.state {
			VMState::Definition => true,
			VMState::Execution => false
		};
		if should_def {
			let current_def_word = self.push_def_word(op);
			self.current_def_word = Some(current_def_word);
		} else {
			self.op_stack.insert(0, op);
		}
	}

	fn push_def_word <'a> (&'a mut self, op: Operator) -> String{
		let wrapped = self.current_def_word.as_ref();
		let new_key = match wrapped {
			Some(w) => {
				self.words.get_mut(w).expect("error").insert(0, op);
				w.clone()
			},
			None => {
				let key = match op {
					Operator::Other(key) => key,
					_ => panic!("key is broken")
				};
				self.words.insert(key.clone(), Vec::new());
				key
			}
		};

		return new_key;
	}

	pub fn eval <'a> (&'a mut self) {
		let op = self.pop_op();
		match op {
			Operator::Add => self.add(),
			Operator::Subtract => self.subtract(),
			Operator::Multiply => self.multiply(),
			Operator::Divide => self.divide(),
			Operator::Modulus => self.modulus(),
			Operator::And => self.and(),
			Operator::Or => self.or(),
			Operator::Xor => self.xor(),
			Operator::Lshift => self.lshift(),
			Operator::Rshift => self.rshift(),
			Operator::Abs => self.abs(),
			Operator::Max => self.max(),
			Operator::Min => self.min(),
			Operator::Invert => self.invert(),
			Operator::Negate => self.negate(),
			Operator::Dup => self.dup(),
			Operator::Drop => self.drop(),
			Operator::Swap => self.swap(),
			Operator::Over => self.over(),
			Operator::Rot => self.rot(),
			Operator::Value(i) => {
				self.push_data(Operator::Value(i));
			},
			Operator::Print => self.print(),
			Operator::Other(s) => self.custom(s),
			x => panic!("cannot eval: {}", x.stringify())
		}
	}

	fn custom <'a> (&'a mut self, word: String) {
		let words = self.words_from_word(word);
		for op in words{
			self.op_stack.push(op.clone());
		}
	}

	fn words_from_word <'a> (&'a mut self, word: String) -> Vec<Operator>{
		let mut ret : Vec<Operator> = Vec::new();
		let ops = self.words.get(&word).expect("word not defined");
		for some_op in ops{
			ret.push(some_op.clone());
		}
		return ret;
	}

	fn pop_two_ints <'a> (&'a mut self) -> (i32, i32){
		let a = self.pop_one_int();
		let b = self.pop_one_int();
		(a, b)
	}

	fn pop_one_int <'a> (&'a mut self) -> i32 {
		let a = self.pop_data();
		match a {
			Operator::Value(j) => j,
			_ => panic!("can only add ints"),
		}
	}

	// Math Functions
	pub fn add <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x + y));
	}

	pub fn subtract <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x - y));
	}

	pub fn multiply <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x * y));
	}

	pub fn divide <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(y / x));
	}

	pub fn modulus <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x % y));
	}

	pub fn and <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x & y));
	}

	pub fn or <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x | y));
	}

	pub fn xor <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x ^ y));
	}

	pub fn lshift <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x << y));
	}

	pub fn rshift <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(x >> y));
	}

	pub fn abs <'a> (&'a mut self) {
		let x = self.pop_one_int();
		self.push_data(Operator::Value(if x < 0 {-x} else {x}));
	}

	pub fn max <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(if y > x {y} else {x}));
	}

	pub fn min <'a> (&'a mut self) {
		let (x, y) = self.pop_two_ints();
		self.push_data(Operator::Value(if y < x {y} else {x}));
	}

	pub fn invert <'a> (&'a mut self) {
		let x = self.pop_one_int();
		self.push_data(Operator::Value(!x));
	}

	pub fn negate <'a> (&'a mut self) {
		let x = self.pop_one_int();
		self.push_data(Operator::Value(-x));
	}

	//Stack Functions
	pub fn dup <'a> (&'a mut self) {
		let x = self.pop_data();
		self.push_data(x.clone());
		self.push_data(x.clone());
	}

	pub fn drop <'a> (&'a mut self) {
		self.pop_data();
	}

	pub fn swap <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = self.pop_data();
		self.push_data(x);
		self.push_data(y);
	}

	pub fn over <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = self.pop_data();
		let z = x.clone();
		self.push_data(x);
		self.push_data(y);
		self.push_data(z);
	}

	pub fn rot <'a> (&'a mut self) {
		let x = self.pop_data();
		let y = self.pop_data();
		let z = self.pop_data();
		self.push_data(x);
		self.push_data(z);
		self.push_data(y);
	}

	//Debug Functions
	pub fn print <'a> (&'a mut self) {
		let x = self.pop_data();
		println!("{}",match x{
			Operator::Value(i) => i,
			_ => panic!("tried to print something other than a value"),
		});
	}
}

#[derive(Clone)]
pub enum Operator {
	//stack
	Dup,
	Drop,
	Swap,
	Over,
	Rot,

	//debug
	Print,

	//math
	Add,
	Subtract,
	Multiply,
	Divide,
	Modulus,
	And,
	Or,
	Xor,
	Lshift,
	Rshift,
	Abs,
	Max,
	Min,
	Invert,
	Negate,

	//custom
	Define,
	EndDefine,

	//values and custom calls
	Value(i32),
	Other(String),
}

impl Operator {
	fn stringify(self) -> String {
		match self{
			Operator::Add => format!("+"),
			Operator::Subtract => format!("-"),
			Operator::Multiply => format!("*"),
			Operator::Divide => format!("/"),
			Operator::Modulus => format!("mod"),
			Operator::And => format!("and"),
			Operator::Or => format!("or"),
			Operator::Xor => format!("xor"),
			Operator::Lshift => format!("<<"),
			Operator::Rshift => format!(">>"),
			Operator::Abs => format!("abs"),
			Operator::Max => format!("max"),
			Operator::Min => format!("min"),
			Operator::Invert => format!("invert"),
			Operator::Negate => format!("negate"),
			Operator::Dup => format!("dup"),
			Operator::Drop => format!("drop"),
			Operator::Swap => format!("swap"),
			Operator::Over => format!("over"),
			Operator::Rot => format!("rot"),
			Operator::Value(i) => format!("{}", i),
			Operator::Print => format!("."),
			Operator::Define => format!(":"),
			Operator::EndDefine => format!(";"),
			Operator::Other(x) => format!("other {}", x),
		}
	}
}

fn parse(s: &str) -> Operator {
	match s {
		//stack operations:
		"dup" => Operator::Dup,
		"drop" => Operator::Drop,
		"swap" => Operator::Swap,
		"over" => Operator::Over,
		"rot" => Operator::Rot,

		//debug operations:
		"." => Operator::Print,

		//math operations:
		"+" => Operator::Add,
		"-" => Operator::Subtract,
		"*" => Operator::Multiply,
		"/" => Operator::Divide,
		"mod" => Operator::Modulus,
		"and" => Operator::And,
		"or" => Operator::Or,
		"xor" => Operator::Xor,
		"lshift" => Operator::Lshift,
		"<<" => Operator::Lshift,
		"rshift" => Operator::Rshift,
		">>" => Operator::Rshift,
		"abs" => Operator::Abs,
		"max" => Operator::Max,
		"min" => Operator::Min,
		"invert" => Operator::Invert,
		"negate" => Operator::Negate,

		//custom definitions
		":" => Operator::Define,
		";" => Operator::EndDefine,

		//value or custom operator
		x => {
			let y: Option<i32> = FromStr::from_str(x).ok();
			match y{
				Some(i) => Operator::Value(i),
				None => Operator::Other(format!("{}", x)),
			}
		},
	}
}

fn main() {
	let x = io::stdin();
	let mut vm = VM {state: VMState::Execution,
					 op_stack: vec![],
					 data_stack: vec![],
					 words: HashMap::new(),
				     current_def_word: None};

	loop{
		let mut line: String = "".to_string();
		let _ = x.read_line(&mut line);

		if line.trim() == "" {break;}

		for o in line.split(" "){
			let b = parse(o.trim());
			vm.push_op(b)
		}

		while vm.op_stack.len() > 0 {
			vm.eval();
		}
	}
}
<|endoftext|>