<commit_before>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    CodeAction, CodeActionOrCommand, Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin,
    WorkDoneProgressEnd, WorkDoneProgressReport,
};
use std::{
    error, fmt,
    io::{BufRead, BufReader},
    path::{Path, PathBuf},
    process::{Command, Stdio},
    thread::JoinHandle,
    time::Instant,
};

mod conv;

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckOptions {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    pub task_recv: Receiver<CheckTask>,
    cmd_send: Option<Sender<CheckCommand>>,
    handle: Option<JoinHandle<()>>,
}

impl CheckWatcher {
    pub fn new(options: &CheckOptions, workspace_root: PathBuf) -> CheckWatcher {
        let options = options.clone();

        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = std::thread::spawn(move || {
            let mut check = CheckWatcherThread::new(options, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send: Some(cmd_send), handle: Some(handle) }
    }

    /// Returns a CheckWatcher that doesn't actually do anything
    pub fn dummy() -> CheckWatcher {
        CheckWatcher { task_recv: never(), cmd_send: None, handle: None }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        if let Some(cmd_send) = &self.cmd_send {
            cmd_send.send(CheckCommand::Update).unwrap();
        }
    }
}

impl std::ops::Drop for CheckWatcher {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Take the sender out of the option
            let cmd_send = self.cmd_send.take();

            // Dropping the sender finishes the thread loop
            drop(cmd_send);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { url: Url, diagnostic: Diagnostic, fixes: Vec<CodeActionOrCommand> },

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckOptions,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherThread {
    fn new(options: CheckOptions, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // Replace with a dummy watcher first so we drop the original and wait for completion
                std::mem::replace(&mut self.watcher, WatchThread::dummy());

                // Then create the actual new watcher
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result = map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root);
                if map_result.is_empty() {
                    return;
                }

                for MappedRustDiagnostic { location, diagnostic, fixes } in map_result {
                    let fixes = fixes
                        .into_iter()
                        .map(|fix| {
                            CodeAction { diagnostics: Some(vec![diagnostic.clone()]), ..fix }.into()
                        })
                        .collect();

                    task_send
                        .send(CheckTask::AddDiagnostic { url: location.uri, diagnostic, fixes })
                        .unwrap();
                }
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    fixes: Vec<CodeAction>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    handle: Option<JoinHandle<()>>,
    message_recv: Receiver<CheckEvent>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

#[derive(Debug)]
pub struct CargoError(String);

impl fmt::Display for CargoError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Cargo failed: {}", self.0)
    }
}
impl error::Error for CargoError {}

pub fn run_cargo(
    args: &[String],
    current_dir: Option<&Path>,
    on_message: &mut dyn FnMut(cargo_metadata::Message) -> bool,
) -> Result<(), CargoError> {
    let mut command = Command::new("cargo");
    if let Some(current_dir) = current_dir {
        command.current_dir(current_dir);
    }

    let mut child = command
        .args(args)
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .stdin(Stdio::null())
        .spawn()
        .expect("couldn't launch cargo");

    // We manually read a line at a time, instead of using serde's
    // stream deserializers, because the deserializer cannot recover
    // from an error, resulting in it getting stuck, because we try to
    // be resillient against failures.
    //
    // Because cargo only outputs one JSON object per line, we can
    // simply skip a line if it doesn't parse, which just ignores any
    // erroneus output.
    let stdout = BufReader::new(child.stdout.take().unwrap());
    let mut read_at_least_one_message = false;

    for line in stdout.lines() {
        let line = match line {
            Ok(line) => line,
            Err(err) => {
                log::error!("Couldn't read line from cargo: {}", err);
                continue;
            }
        };

        let message = serde_json::from_str::<cargo_metadata::Message>(&line);
        let message = match message {
            Ok(message) => message,
            Err(err) => {
                log::error!("Invalid json from cargo check, ignoring ({}): {:?} ", err, line);
                continue;
            }
        };

        read_at_least_one_message = true;

        if !on_message(message) {
            break;
        }
    }

    // It is okay to ignore the result, as it only errors if the process is already dead
    let _ = child.kill();

    let err_msg = match child.wait() {
        Ok(exit_code) if !exit_code.success() && !read_at_least_one_message => {
            // FIXME: Read the stderr to display the reason, see `read2()` reference in PR comment:
            // https://github.com/rust-analyzer/rust-analyzer/pull/3632#discussion_r395605298
            format!("the command produced no valid metadata:\n cargo {}", args.join(" "))
        }
        Err(err) => format!("io error: {:?}", err),
        Ok(_) => return Ok(()),
    };

    Err(CargoError(err_msg))
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { handle: None, message_recv: never() }
    }

    fn new(options: &CheckOptions, workspace_root: &Path) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--workspace".to_string(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.display()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let workspace_root = workspace_root.to_owned();
        let handle = if options.enable {
            Some(std::thread::spawn(move || {
                // If we trigger an error here, we will do so in the loop instead,
                // which will break out of the loop, and continue the shutdown
                let _ = message_send.send(CheckEvent::Begin);

                let res = run_cargo(&args, Some(&workspace_root), &mut |message| {
                    // Skip certain kinds of messages to only spend time on what's useful
                    match &message {
                        Message::CompilerArtifact(artifact) if artifact.fresh => return true,
                        Message::BuildScriptExecuted(_) => return true,
                        Message::Unknown => return true,
                        _ => {}
                    }

                    // if the send channel was closed, so we want to shutdown
                    message_send.send(CheckEvent::Msg(message)).is_ok()
                });

                if let Err(err) = res {
                    // FIXME: make the `message_send` to be `Sender<Result<CheckEvent, CargoError>>`
                    // to display user-caused misconfiguration errors instead of just logging them here
                    log::error!("Cargo watcher failed {:?}", err);
                }

                // We can ignore any error here, as we are already in the progress
                // of shutting down.
                let _ = message_send.send(CheckEvent::End);
            }))
        } else {
            None
        };
        WatchThread { handle, message_recv }
    }
}

impl std::ops::Drop for WatchThread {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Replace our reciever with dummy one, so we can drop and close the
            // one actually communicating with the thread
            let recv = std::mem::replace(&mut self.message_recv, never());

            // Dropping the original reciever initiates thread sub-process shutdown
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}
<commit_msg>fix: typo<commit_after>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    CodeAction, CodeActionOrCommand, Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin,
    WorkDoneProgressEnd, WorkDoneProgressReport,
};
use std::{
    error, fmt,
    io::{BufRead, BufReader},
    path::{Path, PathBuf},
    process::{Command, Stdio},
    thread::JoinHandle,
    time::Instant,
};

mod conv;

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckOptions {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    pub task_recv: Receiver<CheckTask>,
    cmd_send: Option<Sender<CheckCommand>>,
    handle: Option<JoinHandle<()>>,
}

impl CheckWatcher {
    pub fn new(options: &CheckOptions, workspace_root: PathBuf) -> CheckWatcher {
        let options = options.clone();

        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = std::thread::spawn(move || {
            let mut check = CheckWatcherThread::new(options, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send: Some(cmd_send), handle: Some(handle) }
    }

    /// Returns a CheckWatcher that doesn't actually do anything
    pub fn dummy() -> CheckWatcher {
        CheckWatcher { task_recv: never(), cmd_send: None, handle: None }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        if let Some(cmd_send) = &self.cmd_send {
            cmd_send.send(CheckCommand::Update).unwrap();
        }
    }
}

impl std::ops::Drop for CheckWatcher {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Take the sender out of the option
            let cmd_send = self.cmd_send.take();

            // Dropping the sender finishes the thread loop
            drop(cmd_send);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { url: Url, diagnostic: Diagnostic, fixes: Vec<CodeActionOrCommand> },

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckOptions,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherThread {
    fn new(options: CheckOptions, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // Replace with a dummy watcher first so we drop the original and wait for completion
                std::mem::replace(&mut self.watcher, WatchThread::dummy());

                // Then create the actual new watcher
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result = map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root);
                if map_result.is_empty() {
                    return;
                }

                for MappedRustDiagnostic { location, diagnostic, fixes } in map_result {
                    let fixes = fixes
                        .into_iter()
                        .map(|fix| {
                            CodeAction { diagnostics: Some(vec![diagnostic.clone()]), ..fix }.into()
                        })
                        .collect();

                    task_send
                        .send(CheckTask::AddDiagnostic { url: location.uri, diagnostic, fixes })
                        .unwrap();
                }
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    fixes: Vec<CodeAction>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    handle: Option<JoinHandle<()>>,
    message_recv: Receiver<CheckEvent>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

#[derive(Debug)]
pub struct CargoError(String);

impl fmt::Display for CargoError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Cargo failed: {}", self.0)
    }
}
impl error::Error for CargoError {}

pub fn run_cargo(
    args: &[String],
    current_dir: Option<&Path>,
    on_message: &mut dyn FnMut(cargo_metadata::Message) -> bool,
) -> Result<(), CargoError> {
    let mut command = Command::new("cargo");
    if let Some(current_dir) = current_dir {
        command.current_dir(current_dir);
    }

    let mut child = command
        .args(args)
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .stdin(Stdio::null())
        .spawn()
        .expect("couldn't launch cargo");

    // We manually read a line at a time, instead of using serde's
    // stream deserializers, because the deserializer cannot recover
    // from an error, resulting in it getting stuck, because we try to
    // be resillient against failures.
    //
    // Because cargo only outputs one JSON object per line, we can
    // simply skip a line if it doesn't parse, which just ignores any
    // erroneus output.
    let stdout = BufReader::new(child.stdout.take().unwrap());
    let mut read_at_least_one_message = false;

    for line in stdout.lines() {
        let line = match line {
            Ok(line) => line,
            Err(err) => {
                log::error!("Couldn't read line from cargo: {}", err);
                continue;
            }
        };

        let message = serde_json::from_str::<cargo_metadata::Message>(&line);
        let message = match message {
            Ok(message) => message,
            Err(err) => {
                log::error!("Invalid json from cargo check, ignoring ({}): {:?} ", err, line);
                continue;
            }
        };

        read_at_least_one_message = true;

        if !on_message(message) {
            break;
        }
    }

    // It is okay to ignore the result, as it only errors if the process is already dead
    let _ = child.kill();

    let err_msg = match child.wait() {
        Ok(exit_code) if !exit_code.success() && !read_at_least_one_message => {
            // FIXME: Read the stderr to display the reason, see `read2()` reference in PR comment:
            // https://github.com/rust-analyzer/rust-analyzer/pull/3632#discussion_r395605298
            format!("the command produced no valid metadata:\n cargo {}", args.join(" "))
        }
        Err(err) => format!("io error: {:?}", err),
        Ok(_) => return Ok(()),
    };

    Err(CargoError(err_msg))
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { handle: None, message_recv: never() }
    }

    fn new(options: &CheckOptions, workspace_root: &Path) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--workspace".to_string(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.display()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let workspace_root = workspace_root.to_owned();
        let handle = if options.enable {
            Some(std::thread::spawn(move || {
                // If we trigger an error here, we will do so in the loop instead,
                // which will break out of the loop, and continue the shutdown
                let _ = message_send.send(CheckEvent::Begin);

                let res = run_cargo(&args, Some(&workspace_root), &mut |message| {
                    // Skip certain kinds of messages to only spend time on what's useful
                    match &message {
                        Message::CompilerArtifact(artifact) if artifact.fresh => return true,
                        Message::BuildScriptExecuted(_) => return true,
                        Message::Unknown => return true,
                        _ => {}
                    }

                    // if the send channel was closed, we want to shutdown
                    message_send.send(CheckEvent::Msg(message)).is_ok()
                });

                if let Err(err) = res {
                    // FIXME: make the `message_send` to be `Sender<Result<CheckEvent, CargoError>>`
                    // to display user-caused misconfiguration errors instead of just logging them here
                    log::error!("Cargo watcher failed {:?}", err);
                }

                // We can ignore any error here, as we are already in the progress
                // of shutting down.
                let _ = message_send.send(CheckEvent::End);
            }))
        } else {
            None
        };
        WatchThread { handle, message_recv }
    }
}

impl std::ops::Drop for WatchThread {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Replace our reciever with dummy one, so we can drop and close the
            // one actually communicating with the thread
            let recv = std::mem::replace(&mut self.message_recv, never());

            // Dropping the original reciever initiates thread sub-process shutdown
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here
use hir_def::{
    child_by_source::ChildBySource, dyn_map::DynMap, keys, keys::Key, nameres::ModuleSource,
    ConstId, DefWithBodyId, EnumId, EnumVariantId, FunctionId, GenericDefId, ImplId, ModuleId,
    StaticId, StructId, TraitId, TypeAliasId, UnionId, VariantId,
};
use hir_expand::{name::AsName, AstId, MacroDefId, MacroDefKind};
use ra_syntax::{
    ast::{self, AstNode, NameOwner},
    match_ast, SyntaxNode,
};

use crate::{
    db::{AstDatabase, DefDatabase, HirDatabase},
    Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, ImplBlock, InFile, Local,
    MacroDef, Module, Static, Struct, StructField, Trait, TypeAlias, TypeParam, Union,
};

pub trait FromSource: Sized {
    type Ast;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self>;
}

pub trait FromSourceByContainer: Sized {
    type Ast: AstNode + 'static;
    type Id: Copy + 'static;
    const KEY: Key<Self::Ast, Self::Id>;
}

impl<T: FromSourceByContainer> FromSource for T
where
    T: From<<T as FromSourceByContainer>::Id>,
{
    type Ast = <T as FromSourceByContainer>::Ast;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        analyze_container(db, src.as_ref().map(|it| it.syntax()))[T::KEY]
            .get(&src)
            .copied()
            .map(Self::from)
    }
}

macro_rules! from_source_by_container_impls {
    ($(($hir:ident, $id:ident, $ast:path, $key:path)),* ,) => {$(
        impl FromSourceByContainer for $hir {
            type Ast = $ast;
            type Id = $id;
            const KEY: Key<Self::Ast, Self::Id> = $key;
        }
    )*}
}

from_source_by_container_impls![
    (Struct, StructId, ast::StructDef, keys::STRUCT),
    (Union, UnionId, ast::UnionDef, keys::UNION),
    (Enum, EnumId, ast::EnumDef, keys::ENUM),
    (Trait, TraitId, ast::TraitDef, keys::TRAIT),
    (Function, FunctionId, ast::FnDef, keys::FUNCTION),
    (Static, StaticId, ast::StaticDef, keys::STATIC),
    (Const, ConstId, ast::ConstDef, keys::CONST),
    (TypeAlias, TypeAliasId, ast::TypeAliasDef, keys::TYPE_ALIAS),
    (ImplBlock, ImplId, ast::ImplBlock, keys::IMPL),
];

impl FromSource for MacroDef {
    type Ast = ast::MacroCall;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        let kind = MacroDefKind::Declarative;

        let module_src = ModuleSource::from_child_node(db, src.as_ref().map(|it| it.syntax()));
        let module = Module::from_definition(db, InFile::new(src.file_id, module_src))?;
        let krate = Some(module.krate().id);

        let ast_id = Some(AstId::new(src.file_id, db.ast_id_map(src.file_id).ast_id(&src.value)));

        let id: MacroDefId = MacroDefId { krate, ast_id, kind };
        Some(MacroDef { id })
    }
}

impl FromSource for EnumVariant {
    type Ast = ast::EnumVariant;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        let parent_enum = src.value.parent_enum();
        let src_enum = InFile { file_id: src.file_id, value: parent_enum };
        let parent_enum = Enum::from_source(db, src_enum)?;
        parent_enum.id.child_by_source(db)[keys::ENUM_VARIANT]
            .get(&src)
            .copied()
            .map(EnumVariant::from)
    }
}

impl FromSource for StructField {
    type Ast = FieldSource;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        let src = src.as_ref();

        // FIXME this is buggy
        let variant_id: VariantId = match src.value {
            FieldSource::Named(field) => {
                let value = field.syntax().ancestors().find_map(ast::StructDef::cast)?;
                let src = InFile { file_id: src.file_id, value };
                let def = Struct::from_source(db, src)?;
                def.id.into()
            }
            FieldSource::Pos(field) => {
                let value = field.syntax().ancestors().find_map(ast::EnumVariant::cast)?;
                let src = InFile { file_id: src.file_id, value };
                let def = EnumVariant::from_source(db, src)?;
                EnumVariantId::from(def).into()
            }
        };

        let dyn_map = variant_id.child_by_source(db);
        match src.value {
            FieldSource::Pos(it) => dyn_map[keys::TUPLE_FIELD].get(&src.with_value(it.clone())),
            FieldSource::Named(it) => dyn_map[keys::RECORD_FIELD].get(&src.with_value(it.clone())),
        }
        .copied()
        .map(StructField::from)
    }
}

impl Local {
    pub fn from_source(db: &impl HirDatabase, src: InFile<ast::BindPat>) -> Option<Self> {
        let file_id = src.file_id;
        let parent: DefWithBody = src.value.syntax().ancestors().find_map(|it| {
            let res = match_ast! {
                match it {
                    ast::ConstDef(value) => { Const::from_source(db, InFile { value, file_id})?.into() },
                    ast::StaticDef(value) => { Static::from_source(db, InFile { value, file_id})?.into() },
                    ast::FnDef(value) => { Function::from_source(db, InFile { value, file_id})?.into() },
                    _ => return None,
                }
            };
            Some(res)
        })?;
        let (_body, source_map) = db.body_with_source_map(parent.into());
        let src = src.map(ast::Pat::from);
        let pat_id = source_map.node_pat(src.as_ref())?;
        Some(Local { parent, pat_id })
    }
}

impl TypeParam {
    pub fn from_source(db: &impl HirDatabase, src: InFile<ast::TypeParam>) -> Option<Self> {
        let file_id = src.file_id;
        let parent: GenericDefId = src.value.syntax().ancestors().find_map(|it| {
            let res = match_ast! {
                match it {
                    ast::FnDef(value) => { Function::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::StructDef(value) => { Struct::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::EnumDef(value) => { Enum::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::TraitDef(value) => { Trait::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::TypeAliasDef(value) => { TypeAlias::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::ImplBlock(value) => { ImplBlock::from_source(db, InFile { value, file_id})?.id.into() },
                    _ => return None,
                }
            };
            Some(res)
        })?;
        let &id = parent.child_by_source(db)[keys::TYPE_PARAM].get(&src)?;
        Some(TypeParam { id })
    }
}

impl Module {
    pub fn from_declaration(db: &impl DefDatabase, src: InFile<ast::Module>) -> Option<Self> {
        let parent_declaration = src.value.syntax().ancestors().skip(1).find_map(ast::Module::cast);

        let parent_module = match parent_declaration {
            Some(parent_declaration) => {
                let src_parent = InFile { file_id: src.file_id, value: parent_declaration };
                Module::from_declaration(db, src_parent)
            }
            None => {
                let source_file = db.parse(src.file_id.original_file(db)).tree();
                let src_parent =
                    InFile { file_id: src.file_id, value: ModuleSource::SourceFile(source_file) };
                Module::from_definition(db, src_parent)
            }
        }?;

        let child_name = src.value.name()?.as_name();
        let def_map = db.crate_def_map(parent_module.id.krate);
        let child_id = def_map[parent_module.id.local_id].children.get(&child_name)?;
        Some(parent_module.with_module_id(*child_id))
    }

    pub fn from_definition(db: &impl DefDatabase, src: InFile<ModuleSource>) -> Option<Self> {
        match src.value {
            ModuleSource::Module(ref module) => {
                assert!(!module.has_semi());
                return Module::from_declaration(
                    db,
                    InFile { file_id: src.file_id, value: module.clone() },
                );
            }
            ModuleSource::SourceFile(_) => (),
        };

        let original_file = src.file_id.original_file(db);

        let (krate, local_id) = db.relevant_crates(original_file).iter().find_map(|&crate_id| {
            let crate_def_map = db.crate_def_map(crate_id);
            let local_id = crate_def_map.modules_for_file(original_file).next()?;
            Some((crate_id, local_id))
        })?;
        Some(Module { id: ModuleId { krate, local_id } })
    }
}

fn analyze_container(db: &impl DefDatabase, src: InFile<&SyntaxNode>) -> DynMap {
    _analyze_container(db, src).unwrap_or_default()
}

fn _analyze_container(db: &impl DefDatabase, src: InFile<&SyntaxNode>) -> Option<DynMap> {
    // FIXME: this doesn't try to handle nested declarations
    for container in src.value.ancestors().skip(1) {
        let res = match_ast! {
            match container {
                ast::TraitDef(it) => {
                    let c = Trait::from_source(db, src.with_value(it))?;
                    c.id.child_by_source(db)
                },
                ast::ImplBlock(it) => {
                    let c = ImplBlock::from_source(db, src.with_value(it))?;
                    c.id.child_by_source(db)
                },
                ast::FnDef(it) => {
                    let f = Function::from_source(db, src.with_value(it))?;
                    DefWithBodyId::from(f.id)
                        .child_by_source(db)
                },
                _ => { continue },
            }
        };
        return Some(res);
    }

    let module_source = ModuleSource::from_child_node(db, src);
    let c = Module::from_definition(db, src.with_value(module_source))?;
    Some(c.id.child_by_source(db))
}
<commit_msg>Handle nested declrations in from_source<commit_after>//! FIXME: write short doc here
use hir_def::{
    child_by_source::ChildBySource, dyn_map::DynMap, keys, keys::Key, nameres::ModuleSource,
    ConstId, DefWithBodyId, EnumId, EnumVariantId, FunctionId, GenericDefId, ImplId, ModuleId,
    StaticId, StructId, TraitId, TypeAliasId, UnionId, VariantId,
};
use hir_expand::{name::AsName, AstId, MacroDefId, MacroDefKind};
use ra_syntax::{
    ast::{self, AstNode, NameOwner},
    match_ast, SyntaxNode,
};

use crate::{
    db::{AstDatabase, DefDatabase, HirDatabase},
    Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, ImplBlock, InFile, Local,
    MacroDef, Module, Static, Struct, StructField, Trait, TypeAlias, TypeParam, Union,
};

pub trait FromSource: Sized {
    type Ast;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self>;
}

pub trait FromSourceByContainer: Sized {
    type Ast: AstNode + 'static;
    type Id: Copy + 'static;
    const KEY: Key<Self::Ast, Self::Id>;
}

impl<T: FromSourceByContainer> FromSource for T
where
    T: From<<T as FromSourceByContainer>::Id>,
{
    type Ast = <T as FromSourceByContainer>::Ast;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        analyze_container(db, src.as_ref().map(|it| it.syntax()))[T::KEY]
            .get(&src)
            .copied()
            .map(Self::from)
    }
}

macro_rules! from_source_by_container_impls {
    ($(($hir:ident, $id:ident, $ast:path, $key:path)),* ,) => {$(
        impl FromSourceByContainer for $hir {
            type Ast = $ast;
            type Id = $id;
            const KEY: Key<Self::Ast, Self::Id> = $key;
        }
    )*}
}

from_source_by_container_impls![
    (Struct, StructId, ast::StructDef, keys::STRUCT),
    (Union, UnionId, ast::UnionDef, keys::UNION),
    (Enum, EnumId, ast::EnumDef, keys::ENUM),
    (Trait, TraitId, ast::TraitDef, keys::TRAIT),
    (Function, FunctionId, ast::FnDef, keys::FUNCTION),
    (Static, StaticId, ast::StaticDef, keys::STATIC),
    (Const, ConstId, ast::ConstDef, keys::CONST),
    (TypeAlias, TypeAliasId, ast::TypeAliasDef, keys::TYPE_ALIAS),
    (ImplBlock, ImplId, ast::ImplBlock, keys::IMPL),
];

impl FromSource for MacroDef {
    type Ast = ast::MacroCall;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        let kind = MacroDefKind::Declarative;

        let module_src = ModuleSource::from_child_node(db, src.as_ref().map(|it| it.syntax()));
        let module = Module::from_definition(db, InFile::new(src.file_id, module_src))?;
        let krate = Some(module.krate().id);

        let ast_id = Some(AstId::new(src.file_id, db.ast_id_map(src.file_id).ast_id(&src.value)));

        let id: MacroDefId = MacroDefId { krate, ast_id, kind };
        Some(MacroDef { id })
    }
}

impl FromSource for EnumVariant {
    type Ast = ast::EnumVariant;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        let parent_enum = src.value.parent_enum();
        let src_enum = InFile { file_id: src.file_id, value: parent_enum };
        let parent_enum = Enum::from_source(db, src_enum)?;
        parent_enum.id.child_by_source(db)[keys::ENUM_VARIANT]
            .get(&src)
            .copied()
            .map(EnumVariant::from)
    }
}

impl FromSource for StructField {
    type Ast = FieldSource;
    fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
        let src = src.as_ref();

        // FIXME this is buggy
        let variant_id: VariantId = match src.value {
            FieldSource::Named(field) => {
                let value = field.syntax().ancestors().find_map(ast::StructDef::cast)?;
                let src = InFile { file_id: src.file_id, value };
                let def = Struct::from_source(db, src)?;
                def.id.into()
            }
            FieldSource::Pos(field) => {
                let value = field.syntax().ancestors().find_map(ast::EnumVariant::cast)?;
                let src = InFile { file_id: src.file_id, value };
                let def = EnumVariant::from_source(db, src)?;
                EnumVariantId::from(def).into()
            }
        };

        let dyn_map = variant_id.child_by_source(db);
        match src.value {
            FieldSource::Pos(it) => dyn_map[keys::TUPLE_FIELD].get(&src.with_value(it.clone())),
            FieldSource::Named(it) => dyn_map[keys::RECORD_FIELD].get(&src.with_value(it.clone())),
        }
        .copied()
        .map(StructField::from)
    }
}

impl Local {
    pub fn from_source(db: &impl HirDatabase, src: InFile<ast::BindPat>) -> Option<Self> {
        let file_id = src.file_id;
        let parent: DefWithBody = src.value.syntax().ancestors().find_map(|it| {
            let res = match_ast! {
                match it {
                    ast::ConstDef(value) => { Const::from_source(db, InFile { value, file_id})?.into() },
                    ast::StaticDef(value) => { Static::from_source(db, InFile { value, file_id})?.into() },
                    ast::FnDef(value) => { Function::from_source(db, InFile { value, file_id})?.into() },
                    _ => return None,
                }
            };
            Some(res)
        })?;
        let (_body, source_map) = db.body_with_source_map(parent.into());
        let src = src.map(ast::Pat::from);
        let pat_id = source_map.node_pat(src.as_ref())?;
        Some(Local { parent, pat_id })
    }
}

impl TypeParam {
    pub fn from_source(db: &impl HirDatabase, src: InFile<ast::TypeParam>) -> Option<Self> {
        let file_id = src.file_id;
        let parent: GenericDefId = src.value.syntax().ancestors().find_map(|it| {
            let res = match_ast! {
                match it {
                    ast::FnDef(value) => { Function::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::StructDef(value) => { Struct::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::EnumDef(value) => { Enum::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::TraitDef(value) => { Trait::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::TypeAliasDef(value) => { TypeAlias::from_source(db, InFile { value, file_id})?.id.into() },
                    ast::ImplBlock(value) => { ImplBlock::from_source(db, InFile { value, file_id})?.id.into() },
                    _ => return None,
                }
            };
            Some(res)
        })?;
        let &id = parent.child_by_source(db)[keys::TYPE_PARAM].get(&src)?;
        Some(TypeParam { id })
    }
}

impl Module {
    pub fn from_declaration(db: &impl DefDatabase, src: InFile<ast::Module>) -> Option<Self> {
        let parent_declaration = src.value.syntax().ancestors().skip(1).find_map(ast::Module::cast);

        let parent_module = match parent_declaration {
            Some(parent_declaration) => {
                let src_parent = InFile { file_id: src.file_id, value: parent_declaration };
                Module::from_declaration(db, src_parent)
            }
            None => {
                let source_file = db.parse(src.file_id.original_file(db)).tree();
                let src_parent =
                    InFile { file_id: src.file_id, value: ModuleSource::SourceFile(source_file) };
                Module::from_definition(db, src_parent)
            }
        }?;

        let child_name = src.value.name()?.as_name();
        let def_map = db.crate_def_map(parent_module.id.krate);
        let child_id = def_map[parent_module.id.local_id].children.get(&child_name)?;
        Some(parent_module.with_module_id(*child_id))
    }

    pub fn from_definition(db: &impl DefDatabase, src: InFile<ModuleSource>) -> Option<Self> {
        match src.value {
            ModuleSource::Module(ref module) => {
                assert!(!module.has_semi());
                return Module::from_declaration(
                    db,
                    InFile { file_id: src.file_id, value: module.clone() },
                );
            }
            ModuleSource::SourceFile(_) => (),
        };

        let original_file = src.file_id.original_file(db);

        let (krate, local_id) = db.relevant_crates(original_file).iter().find_map(|&crate_id| {
            let crate_def_map = db.crate_def_map(crate_id);
            let local_id = crate_def_map.modules_for_file(original_file).next()?;
            Some((crate_id, local_id))
        })?;
        Some(Module { id: ModuleId { krate, local_id } })
    }
}

fn analyze_container(db: &impl DefDatabase, src: InFile<&SyntaxNode>) -> DynMap {
    _analyze_container(db, src).unwrap_or_default()
}

fn _analyze_container(db: &impl DefDatabase, src: InFile<&SyntaxNode>) -> Option<DynMap> {
    for container in src.value.ancestors().skip(1) {
        let res = match_ast! {
            match container {
                ast::TraitDef(it) => {
                    let def = Trait::from_source(db, src.with_value(it))?;
                    def.id.child_by_source(db)
                },
                ast::ImplBlock(it) => {
                    let def = ImplBlock::from_source(db, src.with_value(it))?;
                    def.id.child_by_source(db)
                },
                ast::FnDef(it) => {
                    let def = Function::from_source(db, src.with_value(it))?;
                    DefWithBodyId::from(def.id)
                        .child_by_source(db)
                },
                ast::StaticDef(it) => {
                    let def = Static::from_source(db, src.with_value(it))?;
                    DefWithBodyId::from(def.id)
                        .child_by_source(db)
                },
                ast::ConstDef(it) => {
                    let def = Const::from_source(db, src.with_value(it))?;
                    DefWithBodyId::from(def.id)
                        .child_by_source(db)
                },
                _ => { continue },
            }
        };
        return Some(res);
    }

    let module_source = ModuleSource::from_child_node(db, src);
    let c = Module::from_definition(db, src.with_value(module_source))?;
    Some(c.id.child_by_source(db))
}
<|endoftext|>


<commit_before>//! This module contains functions for editing syntax trees. As the trees are
//! immutable, all function here return a fresh copy of the tree, instead of
//! doing an in-place modification.
use std::{iter, ops::RangeInclusive};

use arrayvec::ArrayVec;
use rustc_hash::FxHashMap;

use crate::{
    algo,
    ast::{
        self,
        make::{self, tokens},
        AstNode, TypeBoundsOwner,
    },
    AstToken, Direction, InsertPosition, SmolStr, SyntaxElement, SyntaxKind,
    SyntaxKind::{ATTR, COMMENT, WHITESPACE},
    SyntaxNode, SyntaxToken, T,
};

impl ast::BinExpr {
    #[must_use]
    pub fn replace_op(&self, op: SyntaxKind) -> Option<ast::BinExpr> {
        let op_node: SyntaxElement = self.op_details()?.0.into();
        let to_insert: Option<SyntaxElement> = Some(tokens::op(op).into());
        let replace_range = RangeInclusive::new(op_node.clone(), op_node);
        Some(replace_children(self, replace_range, to_insert.into_iter()))
    }
}

impl ast::FnDef {
    #[must_use]
    pub fn with_body(&self, body: ast::Block) -> ast::FnDef {
        let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
        let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.body() {
            old_body.syntax().clone().into()
        } else if let Some(semi) = self.semicolon_token() {
            to_insert.push(make::tokens::single_space().into());
            semi.into()
        } else {
            to_insert.push(make::tokens::single_space().into());
            to_insert.push(body.syntax().clone().into());
            return insert_children(self, InsertPosition::Last, to_insert.into_iter());
        };
        to_insert.push(body.syntax().clone().into());
        let replace_range = RangeInclusive::new(old_body_or_semi.clone(), old_body_or_semi);
        replace_children(self, replace_range, to_insert.into_iter())
    }
}

impl ast::ItemList {
    #[must_use]
    pub fn append_items(&self, items: impl Iterator<Item = ast::ImplItem>) -> ast::ItemList {
        let mut res = self.clone();
        if !self.syntax().text().contains_char('\n') {
            res = res.make_multiline();
        }
        items.for_each(|it| res = res.append_item(it));
        res
    }

    #[must_use]
    pub fn append_item(&self, item: ast::ImplItem) -> ast::ItemList {
        let (indent, position) = match self.impl_items().last() {
            Some(it) => (
                leading_indent(it.syntax()).unwrap_or_default().to_string(),
                InsertPosition::After(it.syntax().clone().into()),
            ),
            None => match self.l_curly() {
                Some(it) => (
                    "    ".to_string() + &leading_indent(self.syntax()).unwrap_or_default(),
                    InsertPosition::After(it),
                ),
                None => return self.clone(),
            },
        };
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert: ArrayVec<[SyntaxElement; 2]> =
            [ws.ws().into(), item.syntax().clone().into()].into();
        insert_children(self, position, to_insert.into_iter())
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.syntax().children_with_tokens().find(|it| it.kind() == T!['{'])
    }

    fn make_multiline(&self) -> ast::ItemList {
        let l_curly = match self.syntax().children_with_tokens().find(|it| it.kind() == T!['{']) {
            Some(it) => it,
            None => return self.clone(),
        };
        let sibling = match l_curly.next_sibling_or_token() {
            Some(it) => it,
            None => return self.clone(),
        };
        let existing_ws = match sibling.as_token() {
            None => None,
            Some(tok) if tok.kind() != WHITESPACE => None,
            Some(ws) => {
                if ws.text().contains('\n') {
                    return self.clone();
                }
                Some(ws.clone())
            }
        };

        let indent = leading_indent(self.syntax()).unwrap_or_default();
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert = iter::once(ws.ws().into());
        match existing_ws {
            None => insert_children(self, InsertPosition::After(l_curly), to_insert),
            Some(ws) => {
                replace_children(self, RangeInclusive::new(ws.clone().into(), ws.into()), to_insert)
            }
        }
    }
}

impl ast::RecordFieldList {
    #[must_use]
    pub fn append_field(&self, field: &ast::RecordField) -> ast::RecordFieldList {
        self.insert_field(InsertPosition::Last, field)
    }

    #[must_use]
    pub fn insert_field(
        &self,
        position: InsertPosition<&'_ ast::RecordField>,
        field: &ast::RecordField,
    ) -> ast::RecordFieldList {
        let is_multiline = self.syntax().text().contains_char('\n');
        let ws;
        let space = if is_multiline {
            ws = tokens::WsBuilder::new(&format!(
                "\n{}    ",
                leading_indent(self.syntax()).unwrap_or_default()
            ));
            ws.ws()
        } else {
            tokens::single_space()
        };

        let mut to_insert: ArrayVec<[SyntaxElement; 4]> = ArrayVec::new();
        to_insert.push(space.into());
        to_insert.push(field.syntax().clone().into());
        to_insert.push(tokens::comma().into());

        macro_rules! after_l_curly {
            () => {{
                let anchor = match self.l_curly() {
                    Some(it) => it,
                    None => return self.clone(),
                };
                InsertPosition::After(anchor)
            }};
        }

        macro_rules! after_field {
            ($anchor:expr) => {
                if let Some(comma) = $anchor
                    .syntax()
                    .siblings_with_tokens(Direction::Next)
                    .find(|it| it.kind() == T![,])
                {
                    InsertPosition::After(comma)
                } else {
                    to_insert.insert(0, tokens::comma().into());
                    InsertPosition::After($anchor.syntax().clone().into())
                }
            };
        };

        let position = match position {
            InsertPosition::First => after_l_curly!(),
            InsertPosition::Last => {
                if !is_multiline {
                    // don't insert comma before curly
                    to_insert.pop();
                }
                match self.fields().last() {
                    Some(it) => after_field!(it),
                    None => after_l_curly!(),
                }
            }
            InsertPosition::Before(anchor) => {
                InsertPosition::Before(anchor.syntax().clone().into())
            }
            InsertPosition::After(anchor) => after_field!(anchor),
        };

        insert_children(self, position, to_insert.iter().cloned())
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.syntax().children_with_tokens().find(|it| it.kind() == T!['{'])
    }
}

impl ast::TypeParam {
    #[must_use]
    pub fn remove_bounds(&self) -> ast::TypeParam {
        let colon = match self.colon_token() {
            Some(it) => it,
            None => return self.clone(),
        };
        let end = match self.type_bound_list() {
            Some(it) => it.syntax().clone().into(),
            None => colon.clone().into(),
        };
        replace_children(self, RangeInclusive::new(colon.into(), end), iter::empty())
    }
}

#[must_use]
pub fn strip_attrs_and_docs<N: ast::AttrsOwner>(node: &N) -> N {
    N::cast(strip_attrs_and_docs_inner(node.syntax().clone())).unwrap()
}

fn strip_attrs_and_docs_inner(mut node: SyntaxNode) -> SyntaxNode {
    while let Some(start) =
        node.children_with_tokens().find(|it| it.kind() == ATTR || it.kind() == COMMENT)
    {
        let end = match &start.next_sibling_or_token() {
            Some(el) if el.kind() == WHITESPACE => el.clone(),
            Some(_) | None => start.clone(),
        };
        node = algo::replace_children(&node, RangeInclusive::new(start, end), &mut iter::empty());
    }
    node
}

#[must_use]
pub fn replace_descendants<N: AstNode, D: AstNode>(
    parent: &N,
    replacement_map: impl Iterator<Item = (D, D)>,
) -> N {
    let map = replacement_map
        .map(|(from, to)| (from.syntax().clone().into(), to.syntax().clone().into()))
        .collect::<FxHashMap<SyntaxElement, _>>();
    let new_syntax = algo::replace_descendants(parent.syntax(), &|n| map.get(n).cloned());
    N::cast(new_syntax).unwrap()
}

#[derive(Debug, Clone, Copy)]
pub struct IndentLevel(pub u8);

impl From<u8> for IndentLevel {
    fn from(level: u8) -> IndentLevel {
        IndentLevel(level)
    }
}

impl IndentLevel {
    pub fn from_node(node: &SyntaxNode) -> IndentLevel {
        let first_token = match node.first_token() {
            Some(it) => it,
            None => return IndentLevel(0),
        };
        for ws in prev_tokens(first_token).filter_map(ast::Whitespace::cast) {
            let text = ws.syntax().text();
            if let Some(pos) = text.rfind('\n') {
                let level = text[pos + 1..].chars().count() / 4;
                return IndentLevel(level as u8);
            }
        }
        IndentLevel(0)
    }

    pub fn increase_indent<N: AstNode>(self, node: N) -> N {
        N::cast(self._increase_indent(node.syntax().clone())).unwrap()
    }

    fn _increase_indent(self, node: SyntaxNode) -> SyntaxNode {
        let replacements: FxHashMap<SyntaxElement, SyntaxElement> = node
            .descendants_with_tokens()
            .filter_map(|el| el.into_token())
            .filter_map(ast::Whitespace::cast)
            .filter(|ws| {
                let text = ws.syntax().text();
                text.contains('\n')
            })
            .map(|ws| {
                (
                    ws.syntax().clone().into(),
                    make::tokens::whitespace(&format!(
                        "{}{:width$}",
                        ws.syntax().text(),
                        "",
                        width = self.0 as usize * 4
                    ))
                    .into(),
                )
            })
            .collect();
        algo::replace_descendants(&node, &|n| replacements.get(n).cloned())
    }

    pub fn decrease_indent<N: AstNode>(self, node: N) -> N {
        N::cast(self._decrease_indent(node.syntax().clone())).unwrap()
    }

    fn _decrease_indent(self, node: SyntaxNode) -> SyntaxNode {
        let replacements: FxHashMap<SyntaxElement, SyntaxElement> = node
            .descendants_with_tokens()
            .filter_map(|el| el.into_token())
            .filter_map(ast::Whitespace::cast)
            .filter(|ws| {
                let text = ws.syntax().text();
                text.contains('\n')
            })
            .map(|ws| {
                (
                    ws.syntax().clone().into(),
                    make::tokens::whitespace(
                        &ws.syntax()
                            .text()
                            .replace(&format!("\n{:1$}", "", self.0 as usize * 4), "\n"),
                    )
                    .into(),
                )
            })
            .collect();
        algo::replace_descendants(&node, &|n| replacements.get(n).cloned())
    }
}

// FIXME: replace usages with IndentLevel above
fn leading_indent(node: &SyntaxNode) -> Option<SmolStr> {
    for token in prev_tokens(node.first_token()?) {
        if let Some(ws) = ast::Whitespace::cast(token.clone()) {
            let ws_text = ws.text();
            if let Some(pos) = ws_text.rfind('\n') {
                return Some(ws_text[pos + 1..].into());
            }
        }
        if token.text().contains('\n') {
            break;
        }
    }
    None
}

fn prev_tokens(token: SyntaxToken) -> impl Iterator<Item = SyntaxToken> {
    iter::successors(Some(token), |token| token.prev_token())
}

#[must_use]
fn insert_children<N: AstNode>(
    parent: &N,
    position: InsertPosition<SyntaxElement>,
    mut to_insert: impl Iterator<Item = SyntaxElement>,
) -> N {
    let new_syntax = algo::insert_children(parent.syntax(), position, &mut to_insert);
    N::cast(new_syntax).unwrap()
}

#[must_use]
fn replace_children<N: AstNode>(
    parent: &N,
    to_replace: RangeInclusive<SyntaxElement>,
    mut to_insert: impl Iterator<Item = SyntaxElement>,
) -> N {
    let new_syntax = algo::replace_children(parent.syntax(), to_replace, &mut to_insert);
    N::cast(new_syntax).unwrap()
}

#[test]
fn test_increase_indent() {
    let arm_list = {
        let arm = make::match_arm(iter::once(make::placeholder_pat().into()), make::expr_unit());
        make::match_arm_list(vec![arm.clone(), arm])
    };
    assert_eq!(
        arm_list.syntax().to_string(),
        "{
    _ => (),
    _ => (),
}"
    );
    let indented = IndentLevel(2).increase_indent(arm_list);
    assert_eq!(
        indented.syntax().to_string(),
        "{
            _ => (),
            _ => (),
        }"
    );
}
<commit_msg>Simplify<commit_after>//! This module contains functions for editing syntax trees. As the trees are
//! immutable, all function here return a fresh copy of the tree, instead of
//! doing an in-place modification.
use std::{iter, ops::RangeInclusive};

use arrayvec::ArrayVec;
use rustc_hash::FxHashMap;

use crate::{
    algo,
    ast::{
        self,
        make::{self, tokens},
        AstNode, TypeBoundsOwner,
    },
    AstToken, Direction, InsertPosition, SmolStr, SyntaxElement, SyntaxKind,
    SyntaxKind::{ATTR, COMMENT, WHITESPACE},
    SyntaxNode, SyntaxToken, T,
};

impl ast::BinExpr {
    #[must_use]
    pub fn replace_op(&self, op: SyntaxKind) -> Option<ast::BinExpr> {
        let op_node: SyntaxElement = self.op_details()?.0.into();
        let to_insert: Option<SyntaxElement> = Some(tokens::op(op).into());
        let replace_range = op_node.clone()..=op_node;
        Some(replace_children(self, replace_range, to_insert.into_iter()))
    }
}

impl ast::FnDef {
    #[must_use]
    pub fn with_body(&self, body: ast::Block) -> ast::FnDef {
        let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
        let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.body() {
            old_body.syntax().clone().into()
        } else if let Some(semi) = self.semicolon_token() {
            to_insert.push(make::tokens::single_space().into());
            semi.into()
        } else {
            to_insert.push(make::tokens::single_space().into());
            to_insert.push(body.syntax().clone().into());
            return insert_children(self, InsertPosition::Last, to_insert.into_iter());
        };
        to_insert.push(body.syntax().clone().into());
        let replace_range = old_body_or_semi.clone()..=old_body_or_semi;
        replace_children(self, replace_range, to_insert.into_iter())
    }
}

impl ast::ItemList {
    #[must_use]
    pub fn append_items(&self, items: impl Iterator<Item = ast::ImplItem>) -> ast::ItemList {
        let mut res = self.clone();
        if !self.syntax().text().contains_char('\n') {
            res = res.make_multiline();
        }
        items.for_each(|it| res = res.append_item(it));
        res
    }

    #[must_use]
    pub fn append_item(&self, item: ast::ImplItem) -> ast::ItemList {
        let (indent, position) = match self.impl_items().last() {
            Some(it) => (
                leading_indent(it.syntax()).unwrap_or_default().to_string(),
                InsertPosition::After(it.syntax().clone().into()),
            ),
            None => match self.l_curly() {
                Some(it) => (
                    "    ".to_string() + &leading_indent(self.syntax()).unwrap_or_default(),
                    InsertPosition::After(it),
                ),
                None => return self.clone(),
            },
        };
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert: ArrayVec<[SyntaxElement; 2]> =
            [ws.ws().into(), item.syntax().clone().into()].into();
        insert_children(self, position, to_insert.into_iter())
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.syntax().children_with_tokens().find(|it| it.kind() == T!['{'])
    }

    fn make_multiline(&self) -> ast::ItemList {
        let l_curly = match self.syntax().children_with_tokens().find(|it| it.kind() == T!['{']) {
            Some(it) => it,
            None => return self.clone(),
        };
        let sibling = match l_curly.next_sibling_or_token() {
            Some(it) => it,
            None => return self.clone(),
        };
        let existing_ws = match sibling.as_token() {
            None => None,
            Some(tok) if tok.kind() != WHITESPACE => None,
            Some(ws) => {
                if ws.text().contains('\n') {
                    return self.clone();
                }
                Some(ws.clone())
            }
        };

        let indent = leading_indent(self.syntax()).unwrap_or_default();
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert = iter::once(ws.ws().into());
        match existing_ws {
            None => insert_children(self, InsertPosition::After(l_curly), to_insert),
            Some(ws) => replace_children(self, ws.clone().into()..=ws.into(), to_insert),
        }
    }
}

impl ast::RecordFieldList {
    #[must_use]
    pub fn append_field(&self, field: &ast::RecordField) -> ast::RecordFieldList {
        self.insert_field(InsertPosition::Last, field)
    }

    #[must_use]
    pub fn insert_field(
        &self,
        position: InsertPosition<&'_ ast::RecordField>,
        field: &ast::RecordField,
    ) -> ast::RecordFieldList {
        let is_multiline = self.syntax().text().contains_char('\n');
        let ws;
        let space = if is_multiline {
            ws = tokens::WsBuilder::new(&format!(
                "\n{}    ",
                leading_indent(self.syntax()).unwrap_or_default()
            ));
            ws.ws()
        } else {
            tokens::single_space()
        };

        let mut to_insert: ArrayVec<[SyntaxElement; 4]> = ArrayVec::new();
        to_insert.push(space.into());
        to_insert.push(field.syntax().clone().into());
        to_insert.push(tokens::comma().into());

        macro_rules! after_l_curly {
            () => {{
                let anchor = match self.l_curly() {
                    Some(it) => it,
                    None => return self.clone(),
                };
                InsertPosition::After(anchor)
            }};
        }

        macro_rules! after_field {
            ($anchor:expr) => {
                if let Some(comma) = $anchor
                    .syntax()
                    .siblings_with_tokens(Direction::Next)
                    .find(|it| it.kind() == T![,])
                {
                    InsertPosition::After(comma)
                } else {
                    to_insert.insert(0, tokens::comma().into());
                    InsertPosition::After($anchor.syntax().clone().into())
                }
            };
        };

        let position = match position {
            InsertPosition::First => after_l_curly!(),
            InsertPosition::Last => {
                if !is_multiline {
                    // don't insert comma before curly
                    to_insert.pop();
                }
                match self.fields().last() {
                    Some(it) => after_field!(it),
                    None => after_l_curly!(),
                }
            }
            InsertPosition::Before(anchor) => {
                InsertPosition::Before(anchor.syntax().clone().into())
            }
            InsertPosition::After(anchor) => after_field!(anchor),
        };

        insert_children(self, position, to_insert.iter().cloned())
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.syntax().children_with_tokens().find(|it| it.kind() == T!['{'])
    }
}

impl ast::TypeParam {
    #[must_use]
    pub fn remove_bounds(&self) -> ast::TypeParam {
        let colon = match self.colon_token() {
            Some(it) => it,
            None => return self.clone(),
        };
        let end = match self.type_bound_list() {
            Some(it) => it.syntax().clone().into(),
            None => colon.clone().into(),
        };
        replace_children(self, colon.into()..=end, iter::empty())
    }
}

#[must_use]
pub fn strip_attrs_and_docs<N: ast::AttrsOwner>(node: &N) -> N {
    N::cast(strip_attrs_and_docs_inner(node.syntax().clone())).unwrap()
}

fn strip_attrs_and_docs_inner(mut node: SyntaxNode) -> SyntaxNode {
    while let Some(start) =
        node.children_with_tokens().find(|it| it.kind() == ATTR || it.kind() == COMMENT)
    {
        let end = match &start.next_sibling_or_token() {
            Some(el) if el.kind() == WHITESPACE => el.clone(),
            Some(_) | None => start.clone(),
        };
        node = algo::replace_children(&node, start..=end, &mut iter::empty());
    }
    node
}

#[must_use]
pub fn replace_descendants<N: AstNode, D: AstNode>(
    parent: &N,
    replacement_map: impl Iterator<Item = (D, D)>,
) -> N {
    let map = replacement_map
        .map(|(from, to)| (from.syntax().clone().into(), to.syntax().clone().into()))
        .collect::<FxHashMap<SyntaxElement, _>>();
    let new_syntax = algo::replace_descendants(parent.syntax(), &|n| map.get(n).cloned());
    N::cast(new_syntax).unwrap()
}

#[derive(Debug, Clone, Copy)]
pub struct IndentLevel(pub u8);

impl From<u8> for IndentLevel {
    fn from(level: u8) -> IndentLevel {
        IndentLevel(level)
    }
}

impl IndentLevel {
    pub fn from_node(node: &SyntaxNode) -> IndentLevel {
        let first_token = match node.first_token() {
            Some(it) => it,
            None => return IndentLevel(0),
        };
        for ws in prev_tokens(first_token).filter_map(ast::Whitespace::cast) {
            let text = ws.syntax().text();
            if let Some(pos) = text.rfind('\n') {
                let level = text[pos + 1..].chars().count() / 4;
                return IndentLevel(level as u8);
            }
        }
        IndentLevel(0)
    }

    pub fn increase_indent<N: AstNode>(self, node: N) -> N {
        N::cast(self._increase_indent(node.syntax().clone())).unwrap()
    }

    fn _increase_indent(self, node: SyntaxNode) -> SyntaxNode {
        let replacements: FxHashMap<SyntaxElement, SyntaxElement> = node
            .descendants_with_tokens()
            .filter_map(|el| el.into_token())
            .filter_map(ast::Whitespace::cast)
            .filter(|ws| {
                let text = ws.syntax().text();
                text.contains('\n')
            })
            .map(|ws| {
                (
                    ws.syntax().clone().into(),
                    make::tokens::whitespace(&format!(
                        "{}{:width$}",
                        ws.syntax().text(),
                        "",
                        width = self.0 as usize * 4
                    ))
                    .into(),
                )
            })
            .collect();
        algo::replace_descendants(&node, &|n| replacements.get(n).cloned())
    }

    pub fn decrease_indent<N: AstNode>(self, node: N) -> N {
        N::cast(self._decrease_indent(node.syntax().clone())).unwrap()
    }

    fn _decrease_indent(self, node: SyntaxNode) -> SyntaxNode {
        let replacements: FxHashMap<SyntaxElement, SyntaxElement> = node
            .descendants_with_tokens()
            .filter_map(|el| el.into_token())
            .filter_map(ast::Whitespace::cast)
            .filter(|ws| {
                let text = ws.syntax().text();
                text.contains('\n')
            })
            .map(|ws| {
                (
                    ws.syntax().clone().into(),
                    make::tokens::whitespace(
                        &ws.syntax()
                            .text()
                            .replace(&format!("\n{:1$}", "", self.0 as usize * 4), "\n"),
                    )
                    .into(),
                )
            })
            .collect();
        algo::replace_descendants(&node, &|n| replacements.get(n).cloned())
    }
}

// FIXME: replace usages with IndentLevel above
fn leading_indent(node: &SyntaxNode) -> Option<SmolStr> {
    for token in prev_tokens(node.first_token()?) {
        if let Some(ws) = ast::Whitespace::cast(token.clone()) {
            let ws_text = ws.text();
            if let Some(pos) = ws_text.rfind('\n') {
                return Some(ws_text[pos + 1..].into());
            }
        }
        if token.text().contains('\n') {
            break;
        }
    }
    None
}

fn prev_tokens(token: SyntaxToken) -> impl Iterator<Item = SyntaxToken> {
    iter::successors(Some(token), |token| token.prev_token())
}

#[must_use]
fn insert_children<N: AstNode>(
    parent: &N,
    position: InsertPosition<SyntaxElement>,
    mut to_insert: impl Iterator<Item = SyntaxElement>,
) -> N {
    let new_syntax = algo::insert_children(parent.syntax(), position, &mut to_insert);
    N::cast(new_syntax).unwrap()
}

#[must_use]
fn replace_children<N: AstNode>(
    parent: &N,
    to_replace: RangeInclusive<SyntaxElement>,
    mut to_insert: impl Iterator<Item = SyntaxElement>,
) -> N {
    let new_syntax = algo::replace_children(parent.syntax(), to_replace, &mut to_insert);
    N::cast(new_syntax).unwrap()
}

#[test]
fn test_increase_indent() {
    let arm_list = {
        let arm = make::match_arm(iter::once(make::placeholder_pat().into()), make::expr_unit());
        make::match_arm_list(vec![arm.clone(), arm])
    };
    assert_eq!(
        arm_list.syntax().to_string(),
        "{
    _ => (),
    _ => (),
}"
    );
    let indented = IndentLevel(2).increase_indent(arm_list);
    assert_eq!(
        indented.syntax().to_string(),
        "{
            _ => (),
            _ => (),
        }"
    );
}
<|endoftext|>


<commit_before>//! An enum mapping type.
//!
//! It is implemented using an array type, so using it is as fast as using Rust
//! arrays.
//!
//! # Examples
//!
//! ```
//! #[macro_use]
//! extern crate enum_map;
//!
//! use enum_map::EnumMap;
//!
//! #[derive(Debug, EnumMap)]
//! enum Example {
//!     A,
//!     B,
//!     C,
//! }
//!
//! fn main() {
//!     let mut map = enum_map! {
//!         Example::A => 1,
//!         Example::B => 2,
//!         Example::C => 3,
//!     };
//!     map[Example::C] = 4;
//!
//!     assert_eq!(map[Example::A], 1);
//!
//!     for (key, &value) in &map {
//!         println!("{:?} has {} as value.", key, value);
//!     }
//! }
//! ```

#![no_std]
#![deny(missing_docs)]

#[macro_use]
extern crate array_macro;
// This allows to quiet "proc macro crates and `#[no_link]` crates have no
// effect without `#[macro_use]`" warning. Just using extern DOES have
// an effect of letting me export the macro.
#[allow(unused_imports)]
#[macro_use]
extern crate enum_map_derive;

mod enummap_impls;
mod internal;
mod iter;
mod serde;

pub use internal::Internal;
pub use iter::{Iter, IterMut};
// `*` here means re-exporting a derive procedural macro.
pub use enum_map_derive::*;

use core::slice;

/// An enum mapping.
///
/// This internally uses an array which stores a value for each possible
/// enum value. To work, it requires implementation of internal (private,
/// although public due to macro limitations) trait which allows extracting
/// information about an enum, which can be automatically generated using
/// `#[derive(EnumMap)]` from `enum_map_derive` crate.
///
/// Additionally, `bool` and `u8` automatically derives from `EnumMap`. While
/// `u8` is not technically an enum, it's convenient to consider it like one.
/// In particular, [reverse-complement in benchmark game] could be using `u8`
/// as an enum.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// use enum_map::EnumMap;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
/// }
///
/// fn main() {
///     let mut map = EnumMap::new();
///     // new initializes map with default values
///     assert_eq!(map[Example::A], 0);
///     map[Example::A] = 3;
///     assert_eq!(map[Example::A], 3);
/// }
/// ```
///
/// [reverse-complement in benchmark game]:
///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
#[derive(Debug)]
pub struct EnumMap<K: Internal<V>, V> {
    array: K::Array,
}

impl<K: Internal<V>, V: Default> EnumMap<K, V>
where
    K::Array: Default,
{
    /// Creates an enum map with default values.
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Example {
    ///     A,
    /// }
    ///
    /// fn main() {
    ///     let enum_map = EnumMap::<_, i32>::new();
    ///     assert_eq!(enum_map[Example::A], 0);
    /// }
    /// ```
    pub fn new() -> Self {
        EnumMap::default()
    }
}

impl<K: Internal<V>, V> EnumMap<K, V> {
    /// Returns number of elements in enum map.
    pub fn len(&self) -> usize {
        K::slice(&self.array).len()
    }

    /// Returns whether the enum variant set is empty.
    ///
    /// This isn't particularly useful, as there is no real reason to use
    /// enum map for enums without variants. However, it is provided for
    /// consistency with data structures providing len method (and I will
    /// admit, to avoid clippy warnings).
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Void {}
    ///
    /// #[derive(EnumMap)]
    /// enum SingleVariant {
    ///     Variant,
    /// }
    ///
    /// fn main() {
    ///     assert_eq!(EnumMap::<Void, ()>::new().is_empty(), true);
    ///     assert_eq!(EnumMap::<SingleVariant, ()>::new().is_empty(), false);
    /// }
    pub fn is_empty(&self) -> bool {
        K::slice(&self.array).is_empty()
    }

    /// Returns an iterator over enum map.
    pub fn iter(&self) -> Iter<K, V> {
        self.into_iter()
    }

    /// Returns a mutable iterator over enum map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.into_iter()
    }

    /// Swaps two indexes.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { false => 0, true => 1 };
    ///     map.swap(false, true);
    ///     assert_eq!(map[false], 1);
    ///     assert_eq!(map[true], 0);
    /// }
    /// ```
    pub fn swap(&mut self, a: K, b: K) {
        K::slice_mut(&mut self.array).swap(a.to_usize(), b.to_usize())
    }

    /// An iterator visiting all values. The iterator type is `&V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let map = enum_map! { false => 3, true => 4 };
    ///     let mut values = map.values();
    ///     assert_eq!(values.next(), Some(&3));
    ///     assert_eq!(values.next(), Some(&4));
    ///     assert_eq!(values.next(), None);
    /// }
    /// ```
    pub fn values(&self) -> slice::Iter<V> {
        K::slice(&self.array).iter()
    }

    /// An iterator visiting all values mutably. The iterator type is `&mut V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 2 };
    ///     for value in map.values_mut() {
    ///         *value += 2;
    ///     }
    ///     assert_eq!(map[false], 4);
    ///     assert_eq!(map[true], 4);
    /// }
    /// ```
    pub fn values_mut(&mut self) -> slice::IterMut<V> {
        K::slice_mut(&mut self.array).iter_mut()
    }

    /// Returns a raw pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let map = enum_map! { 5 => 42, _ => 0 };
    ///     assert_eq!(unsafe { *map.as_ptr().offset(5) }, 42);
    /// }
    /// ```
    pub fn as_ptr(&self) -> *const V {
        K::slice(&self.array).as_ptr()
    }

    /// Returns an unsafe mutable pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 0 };
    ///     unsafe {
    ///         *map.as_mut_ptr().offset(11) = 23
    ///     };
    ///     assert_eq!(map[11], 23);
    /// }
    /// ```
    pub fn as_mut_ptr(&mut self) -> *mut V {
        K::slice_mut(&mut self.array).as_mut_ptr()
    }
}

impl<F: FnMut(K) -> V, K: Internal<V>, V> From<F> for EnumMap<K, V> {
    fn from(f: F) -> Self {
        EnumMap { array: K::from_function(f) }
    }
}

/// Enum map constructor.
///
/// This macro allows to create a new enum map in a type safe way. It takes
/// a list of `,` separated pairs separated by `=>`. Left side is `|`
/// separated list of enum keys, or `_` to match all unmatched enum keys,
/// while right side is a value.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
///     D,
/// }
///
/// fn main() {
///     let enum_map = enum_map! {
///         Example::A | Example::B => 1,
///         Example::C => 2,
///         _ => 3,
///     };
///     assert_eq!(enum_map[Example::A], 1);
///     assert_eq!(enum_map[Example::B], 1);
///     assert_eq!(enum_map[Example::C], 2);
///     assert_eq!(enum_map[Example::D], 3);
/// }
/// ```
#[macro_export]
macro_rules! enum_map {
    {$($t:tt)*} => {
        $crate::EnumMap::from(|k| match k { $($t)* })
    };
}
<commit_msg>Implement as_slice and as_mut_slice<commit_after>//! An enum mapping type.
//!
//! It is implemented using an array type, so using it is as fast as using Rust
//! arrays.
//!
//! # Examples
//!
//! ```
//! #[macro_use]
//! extern crate enum_map;
//!
//! use enum_map::EnumMap;
//!
//! #[derive(Debug, EnumMap)]
//! enum Example {
//!     A,
//!     B,
//!     C,
//! }
//!
//! fn main() {
//!     let mut map = enum_map! {
//!         Example::A => 1,
//!         Example::B => 2,
//!         Example::C => 3,
//!     };
//!     map[Example::C] = 4;
//!
//!     assert_eq!(map[Example::A], 1);
//!
//!     for (key, &value) in &map {
//!         println!("{:?} has {} as value.", key, value);
//!     }
//! }
//! ```

#![no_std]
#![deny(missing_docs)]

#[macro_use]
extern crate array_macro;
// This allows to quiet "proc macro crates and `#[no_link]` crates have no
// effect without `#[macro_use]`" warning. Just using extern DOES have
// an effect of letting me export the macro.
#[allow(unused_imports)]
#[macro_use]
extern crate enum_map_derive;

mod enummap_impls;
mod internal;
mod iter;
mod serde;

pub use internal::Internal;
pub use iter::{Iter, IterMut};
// `*` here means re-exporting a derive procedural macro.
pub use enum_map_derive::*;

use core::slice;

/// An enum mapping.
///
/// This internally uses an array which stores a value for each possible
/// enum value. To work, it requires implementation of internal (private,
/// although public due to macro limitations) trait which allows extracting
/// information about an enum, which can be automatically generated using
/// `#[derive(EnumMap)]` from `enum_map_derive` crate.
///
/// Additionally, `bool` and `u8` automatically derives from `EnumMap`. While
/// `u8` is not technically an enum, it's convenient to consider it like one.
/// In particular, [reverse-complement in benchmark game] could be using `u8`
/// as an enum.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// use enum_map::EnumMap;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
/// }
///
/// fn main() {
///     let mut map = EnumMap::new();
///     // new initializes map with default values
///     assert_eq!(map[Example::A], 0);
///     map[Example::A] = 3;
///     assert_eq!(map[Example::A], 3);
/// }
/// ```
///
/// [reverse-complement in benchmark game]:
///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
#[derive(Debug)]
pub struct EnumMap<K: Internal<V>, V> {
    array: K::Array,
}

impl<K: Internal<V>, V: Default> EnumMap<K, V>
where
    K::Array: Default,
{
    /// Creates an enum map with default values.
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Example {
    ///     A,
    /// }
    ///
    /// fn main() {
    ///     let enum_map = EnumMap::<_, i32>::new();
    ///     assert_eq!(enum_map[Example::A], 0);
    /// }
    /// ```
    pub fn new() -> Self {
        EnumMap::default()
    }
}

impl<K: Internal<V>, V> EnumMap<K, V> {
    /// Returns number of elements in enum map.
    pub fn len(&self) -> usize {
        self.as_slice().len()
    }

    /// Returns whether the enum variant set is empty.
    ///
    /// This isn't particularly useful, as there is no real reason to use
    /// enum map for enums without variants. However, it is provided for
    /// consistency with data structures providing len method (and I will
    /// admit, to avoid clippy warnings).
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Void {}
    ///
    /// #[derive(EnumMap)]
    /// enum SingleVariant {
    ///     Variant,
    /// }
    ///
    /// fn main() {
    ///     assert_eq!(EnumMap::<Void, ()>::new().is_empty(), true);
    ///     assert_eq!(EnumMap::<SingleVariant, ()>::new().is_empty(), false);
    /// }
    pub fn is_empty(&self) -> bool {
        self.as_slice().is_empty()
    }

    /// Returns an iterator over enum map.
    pub fn iter(&self) -> Iter<K, V> {
        self.into_iter()
    }

    /// Returns a mutable iterator over enum map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.into_iter()
    }

    /// Swaps two indexes.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { false => 0, true => 1 };
    ///     map.swap(false, true);
    ///     assert_eq!(map[false], 1);
    ///     assert_eq!(map[true], 0);
    /// }
    /// ```
    pub fn swap(&mut self, a: K, b: K) {
        self.as_mut_slice().swap(a.to_usize(), b.to_usize())
    }

    /// An iterator visiting all values. The iterator type is `&V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let map = enum_map! { false => 3, true => 4 };
    ///     let mut values = map.values();
    ///     assert_eq!(values.next(), Some(&3));
    ///     assert_eq!(values.next(), Some(&4));
    ///     assert_eq!(values.next(), None);
    /// }
    /// ```
    pub fn values(&self) -> slice::Iter<V> {
        self.as_slice().iter()
    }

    /// An iterator visiting all values mutably. The iterator type is `&mut V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 2 };
    ///     for value in map.values_mut() {
    ///         *value += 2;
    ///     }
    ///     assert_eq!(map[false], 4);
    ///     assert_eq!(map[true], 4);
    /// }
    /// ```
    pub fn values_mut(&mut self) -> slice::IterMut<V> {
        self.as_mut_slice().iter_mut()
    }

    /// Converts an enum map to a slice representing values.
    pub fn as_slice(&self) -> &[V] {
        K::slice(&self.array)
    }

    /// Converts a mutable enum map to a mutable slice representing values.
    pub fn as_mut_slice(&mut self) -> &mut [V] {
        K::slice_mut(&mut self.array)
    }

    /// Returns a raw pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let map = enum_map! { 5 => 42, _ => 0 };
    ///     assert_eq!(unsafe { *map.as_ptr().offset(5) }, 42);
    /// }
    /// ```
    pub fn as_ptr(&self) -> *const V {
        self.as_slice().as_ptr()
    }

    /// Returns an unsafe mutable pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 0 };
    ///     unsafe {
    ///         *map.as_mut_ptr().offset(11) = 23
    ///     };
    ///     assert_eq!(map[11], 23);
    /// }
    /// ```
    pub fn as_mut_ptr(&mut self) -> *mut V {
        self.as_mut_slice().as_mut_ptr()
    }
}

impl<F: FnMut(K) -> V, K: Internal<V>, V> From<F> for EnumMap<K, V> {
    fn from(f: F) -> Self {
        EnumMap { array: K::from_function(f) }
    }
}

/// Enum map constructor.
///
/// This macro allows to create a new enum map in a type safe way. It takes
/// a list of `,` separated pairs separated by `=>`. Left side is `|`
/// separated list of enum keys, or `_` to match all unmatched enum keys,
/// while right side is a value.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
///     D,
/// }
///
/// fn main() {
///     let enum_map = enum_map! {
///         Example::A | Example::B => 1,
///         Example::C => 2,
///         _ => 3,
///     };
///     assert_eq!(enum_map[Example::A], 1);
///     assert_eq!(enum_map[Example::B], 1);
///     assert_eq!(enum_map[Example::C], 2);
///     assert_eq!(enum_map[Example::D], 3);
/// }
/// ```
#[macro_export]
macro_rules! enum_map {
    {$($t:tt)*} => {
        $crate::EnumMap::from(|k| match k { $($t)* })
    };
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::cell::RefCell;
use std::collections::VecDeque;

use communication::{Communicator, Data, Message};

// The simplest communicator remains worker-local and just queues sent messages.
pub struct Thread;
impl Communicator for Thread {
    fn index(&self) -> u64 { 0 }
    fn peers(&self) -> u64 { 1 }
    fn new_channel<T: Clone+'static, D: 'static>(&mut self) ->
            (Vec<::communication::observer::BoxedObserver<T, D>>,
             Box<::communication::Pullable<T, D>>) {
        let shared = Rc::new(RefCell::new(VecDeque::<(T, Message<D>)>::new()));
        let reverse = Rc::new(RefCell::new(Vec::new()));

        (vec![::communication::observer::BoxedObserver::new(Observer::new(shared.clone(), reverse.clone()))],
         Box::new(Pullable::new(shared, reverse)) as Box<::communication::Pullable<T, D>>)
    }
}

// an observer wrapping a Rust channel
pub struct Observer<T, D> {
    time: Option<T>,
    dest: Rc<RefCell<VecDeque<(T, Message<D>)>>>,
    reverse: Rc<RefCell<Vec<Vec<D>>>>,
}

impl<T, D> Observer<T, D> {
    pub fn new(dest: Rc<RefCell<VecDeque<(T, Message<D>)>>>, reverse: Rc<RefCell<Vec<Vec<D>>>>) -> Observer<T, D> {
        Observer { time: None, dest: dest, reverse: reverse }
    }
}

impl<T: Clone, D> ::communication::observer::Observer for Observer<T, D> {
    type Time = T;
    type Data = D;
    #[inline] fn open(&mut self, time: &Self::Time) { assert!(self.time.is_none()); self.time = Some(time.clone()); }
    #[inline] fn shut(&mut self,_time: &Self::Time) { assert!(self.time.is_some()); self.time = None; }
    #[inline] fn give(&mut self, data: &mut Message<Self::Data>) {
        assert!(self.time.is_some());
        // ALLOC : We replace with empty typed data. Not clear why typed is better than binary,
        // ALLOC : but lots of folks wouldn't expect bytes, and anyone should tolerate typed.
        if data.len() > 0 {
            let mut replacement = self.reverse.borrow_mut().pop().unwrap_or(Vec::new());
            let to_push = ::std::mem::replace(data, Message::from_typed(&mut replacement));
            self.dest.borrow_mut().push_back((self.time.clone().unwrap(), to_push));
        }
    }
}

pub struct Pullable<T, D> {
    current: Option<(T, Message<D>)>,
    source: Rc<RefCell<VecDeque<(T, Message<D>)>>>,
    reverse: Rc<RefCell<Vec<Vec<D>>>>,
}

impl<T, D> Pullable<T, D> {
    pub fn new(source: Rc<RefCell<VecDeque<(T, Message<D>)>>>, reverse: Rc<RefCell<Vec<Vec<D>>>>) -> Pullable<T, D> {
        Pullable { current: None, source: source, reverse: reverse }
    }
}

impl<T, D> ::communication::pullable::Pullable<T, D> for Pullable<T, D> {
    #[inline]
    fn pull(&mut self) -> Option<(&T, &mut Message<D>)> {

        // ALLOC : Here is one place where we drop both binary and typed data.
        // ALLOC : Ideally the lazy re-allocation elsewhere means we don't always drop.
        if let Some((_time, data)) = self.current.take() {
            let mut vector = data.into_typed();
            if vector.capacity() == Message::<D>::default_length() {
                vector.clear();
                self.reverse.borrow_mut().push(vector);
            }
        }

        self.current = self.source.borrow_mut().pop_front();

        if let Some((_, ref message)) = self.current {
            // should be ensured by test in Thread::Observer.give().
            assert!(message.len() > 0);
        }
        self.current.as_mut().map(|&mut (ref time, ref mut data)| (time, data))
    }
}
<commit_msg>disabling buffer recycling for now<commit_after>use std::rc::Rc;
use std::cell::RefCell;
use std::collections::VecDeque;

use communication::{Communicator, Data, Message};

// The simplest communicator remains worker-local and just queues sent messages.
pub struct Thread;
impl Communicator for Thread {
    fn index(&self) -> u64 { 0 }
    fn peers(&self) -> u64 { 1 }
    fn new_channel<T: Clone+'static, D: 'static>(&mut self) ->
            (Vec<::communication::observer::BoxedObserver<T, D>>,
             Box<::communication::Pullable<T, D>>) {
        let shared = Rc::new(RefCell::new(VecDeque::<(T, Message<D>)>::new()));
        let reverse = Rc::new(RefCell::new(Vec::new()));

        (vec![::communication::observer::BoxedObserver::new(Observer::new(shared.clone(), reverse.clone()))],
         Box::new(Pullable::new(shared, reverse)) as Box<::communication::Pullable<T, D>>)
    }
}

// an observer wrapping a Rust channel
pub struct Observer<T, D> {
    time: Option<T>,
    dest: Rc<RefCell<VecDeque<(T, Message<D>)>>>,
    reverse: Rc<RefCell<Vec<Vec<D>>>>,
}

impl<T, D> Observer<T, D> {
    pub fn new(dest: Rc<RefCell<VecDeque<(T, Message<D>)>>>, reverse: Rc<RefCell<Vec<Vec<D>>>>) -> Observer<T, D> {
        Observer { time: None, dest: dest, reverse: reverse }
    }
}

impl<T: Clone, D> ::communication::observer::Observer for Observer<T, D> {
    type Time = T;
    type Data = D;
    #[inline] fn open(&mut self, time: &Self::Time) { assert!(self.time.is_none()); self.time = Some(time.clone()); }
    #[inline] fn shut(&mut self,_time: &Self::Time) { assert!(self.time.is_some()); self.time = None; }
    #[inline] fn give(&mut self, data: &mut Message<Self::Data>) {
        assert!(self.time.is_some());
        // ALLOC : We replace with empty typed data. Not clear why typed is better than binary,
        // ALLOC : but lots of folks wouldn't expect bytes, and anyone should tolerate typed.
        if data.len() > 0 {
            let mut replacement = self.reverse.borrow_mut().pop().unwrap_or(Vec::new());
            let to_push = ::std::mem::replace(data, Message::from_typed(&mut replacement));
            self.dest.borrow_mut().push_back((self.time.clone().unwrap(), to_push));
        }
    }
}

pub struct Pullable<T, D> {
    current: Option<(T, Message<D>)>,
    source: Rc<RefCell<VecDeque<(T, Message<D>)>>>,
    reverse: Rc<RefCell<Vec<Vec<D>>>>,
}

impl<T, D> Pullable<T, D> {
    pub fn new(source: Rc<RefCell<VecDeque<(T, Message<D>)>>>, reverse: Rc<RefCell<Vec<Vec<D>>>>) -> Pullable<T, D> {
        Pullable { current: None, source: source, reverse: reverse }
    }
}

impl<T, D> ::communication::pullable::Pullable<T, D> for Pullable<T, D> {
    #[inline]
    fn pull(&mut self) -> Option<(&T, &mut Message<D>)> {

        // ALLOC : Here is one place where we drop both binary and typed data.
        // ALLOC : Ideally the lazy re-allocation elsewhere means we don't always drop.
        if let Some((_time, data)) = self.current.take() {
            let mut vector = data.into_typed();
            if vector.capacity() == Message::<D>::default_length() {
                vector.clear();
                // self.reverse.borrow_mut().push(vector);
            }
        }

        self.current = self.source.borrow_mut().pop_front();

        if let Some((_, ref message)) = self.current {
            // should be ensured by test in Thread::Observer.give().
            assert!(message.len() > 0);
        }
        self.current.as_mut().map(|&mut (ref time, ref mut data)| (time, data))
    }
}
<|endoftext|>