<commit_before>// (C)opyleft 2013-2015 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![crate_name="bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]
#![feature(collections)]

extern crate rand;

use std::cmp;
use std::f64;
use std::hash::{Hash, Hasher, SipHasher};
use std::collections::BitVec;

#[cfg(test)]
use rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: BitVec,
    bitmap_bits: u64,
    k_num: u32,
    sips: [SipHasher; 2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: usize, items_count: usize) -> Bloom {
        assert!(bitmap_size > 0 && items_count > 0);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = BitVec::from_elem(bitmap_bits as usize, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: usize, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: usize, fp_p: f64) -> usize {
        assert!(items_count > 0);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * f64::ln(fp_p) / (-8.0 * log2_2)).ceil() as usize
    }

/// Record the presence of an item.
    pub fn set<T>(& mut self, item: T) where T: Hash {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in (0..self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as usize;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T>(&self, item: T) -> bool where T: Hash {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in (0..self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T>(&mut self, item: T)
                               -> bool where T: Hash {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in (0..self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

/// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

/// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> u32 {
        self.k_num
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: usize) -> u32 {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * f64::ln(2.0f64)).ceil() as u32;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T>(&self, hashes: & mut [u64; 2],
                  item: &T, k_i: u32) -> u64 where T: Hash {
        if k_i < 2 {
            let sip = &mut self.sips[k_i as usize].clone();
            item.hash(sip);
            let hash = sip.finish();
            hashes[k_i as usize] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::thread_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::thread_rng().gen_iter::<u8>().take(16).collect();
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::thread_rng().gen_iter::<u8>().take(16).collect();
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<commit_msg>collections -> bitvec for bit vectors<commit_after>// (C)opyleft 2013-2015 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![crate_name="bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]
#![feature(bitvec)]

extern crate rand;

use std::cmp;
use std::f64;
use std::hash::{Hash, Hasher, SipHasher};
use std::collections::BitVec;

#[cfg(test)]
use rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: BitVec,
    bitmap_bits: u64,
    k_num: u32,
    sips: [SipHasher; 2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: usize, items_count: usize) -> Bloom {
        assert!(bitmap_size > 0 && items_count > 0);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = BitVec::from_elem(bitmap_bits as usize, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: usize, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: usize, fp_p: f64) -> usize {
        assert!(items_count > 0);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * f64::ln(fp_p) / (-8.0 * log2_2)).ceil() as usize
    }

/// Record the presence of an item.
    pub fn set<T>(& mut self, item: T) where T: Hash {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in (0..self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as usize;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T>(&self, item: T) -> bool where T: Hash {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in (0..self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T>(&mut self, item: T)
                               -> bool where T: Hash {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in (0..self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as usize;
            if self.bitmap.get(bit_offset).unwrap() == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

/// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

/// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> u32 {
        self.k_num
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: usize) -> u32 {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * f64::ln(2.0f64)).ceil() as u32;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T>(&self, hashes: & mut [u64; 2],
                  item: &T, k_i: u32) -> u64 where T: Hash {
        if k_i < 2 {
            let sip = &mut self.sips[k_i as usize].clone();
            item.hash(sip);
            let hash = sip.finish();
            hashes[k_i as usize] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::thread_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::thread_rng().gen_iter::<u8>().take(16).collect();
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::thread_rng().gen_iter::<u8>().take(16).collect();
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<|endoftext|>


<commit_before>
#![allow(unused)]
include!(concat!(env!("OUT_DIR"), "/../entry_capnp.rs"));

use std;
use std::io;

use capnp;

pub enum Ownership {
    Unknown,
    Posix(String, String),
}

#[derive(Clone, Debug)]
pub struct FileEntry {
    pub len: u64,
    pub paths: Vec<String>,
    pub atime: u64,
    pub mtime: u64,
    pub ctime: u64,
    pub btime: u64,
    pub normal_file: bool,
    pub content_follows: bool,
}

pub fn read_entry<'a, R: io::Read>(mut from: &mut R) -> capnp::Result<Option<FileEntry>> {
    let message = capnp::serialize::read_message(&mut from, capnp::message::ReaderOptions::new())?;

    let entry = message.get_root::<entry::Reader>()?;

    if 0x0100C1C1 != entry.get_magic() {
        return Err(capnp::Error::failed("invalid magic after decoding; invalid stream?".to_string()));
    }

    let entry_paths = entry.get_paths()?;
    let entry_paths_len = entry_paths.len();

    let mut paths = Vec::with_capacity(entry_paths_len as usize);

    for i in 0..entry_paths_len {
        paths.push(entry_paths.get(i)?.to_string());
    }

    Ok(Some(FileEntry {
        len: entry.get_len(),
        atime: entry.get_atime(),
        mtime: entry.get_mtime(),
        ctime: entry.get_ctime(),
        btime: entry.get_btime(),
        paths,
        normal_file: match entry.get_type().which()? {
            entry::type_::Which::Normal(()) => true,
            _ => false,
        },
        content_follows: match entry.get_content().which()? {
            entry::content::Which::Follows(()) => true,
            _ => false,
        }
    }))
}<commit_msg>peeky reader is a cheeky reader<commit_after>
#![allow(unused)]
include!(concat!(env!("OUT_DIR"), "/../entry_capnp.rs"));

use std;
use std::io;

use capnp;

pub enum Ownership {
    Unknown,
    Posix(String, String),
}

#[derive(Clone, Debug)]
pub struct FileEntry {
    pub len: u64,
    pub paths: Vec<String>,
    pub atime: u64,
    pub mtime: u64,
    pub ctime: u64,
    pub btime: u64,
    pub normal_file: bool,
    pub content_follows: bool,
}

struct PeekyRead<'a, R: io::Read + 'a> {
    inner: &'a mut R,
    peeked: Option<u8>,
}

impl<'a, R: io::Read + 'a> io::Read for PeekyRead<'a, R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.peeked {
            Some(c) => {
                buf[0] = c;
                self.peeked = None;
                Ok(1)
            }
            None => {
                self.inner.read(buf)
            }
        }
    }
}

impl<'a, R: io::Read> PeekyRead<'a, R> {
    fn new(inner: &mut R) -> PeekyRead<R> {
        PeekyRead {
            inner,
            peeked: None,
        }
    }

    fn check_eof(&mut self) -> io::Result<bool> {
        if self.peeked.is_some() {
            // we have something more to return; read won't return 0 bytes
            return Ok(false);
        }

        let mut buf = [0; 1];
        Ok(match self.inner.read(&mut buf)? {
            0 => true,
            1 => {
                self.peeked = Some(buf[0]);
                false
            },
            _ => unreachable!(),
        })
    }
}

pub fn read_entry<'a, R: io::Read>(mut from: &mut R) -> capnp::Result<Option<FileEntry>> {
    let mut from = PeekyRead::new(from);

    if from.check_eof()? {
        return Ok(None);
    }

    let message = capnp::serialize::read_message(&mut from, capnp::message::ReaderOptions::new())?;

    let entry = message.get_root::<entry::Reader>()?;

    if 0x0100C1C1 != entry.get_magic() {
        return Err(capnp::Error::failed("invalid magic after decoding; invalid stream?".to_string()));
    }

    let entry_paths = entry.get_paths()?;
    let entry_paths_len = entry_paths.len();

    let mut paths = Vec::with_capacity(entry_paths_len as usize);

    for i in 0..entry_paths_len {
        paths.push(entry_paths.get(i)?.to_string());
    }

    Ok(Some(FileEntry {
        len: entry.get_len(),
        atime: entry.get_atime(),
        mtime: entry.get_mtime(),
        ctime: entry.get_ctime(),
        btime: entry.get_btime(),
        paths,
        normal_file: match entry.get_type().which()? {
            entry::type_::Which::Normal(()) => true,
            _ => false,
        },
        content_follows: match entry.get_content().which()? {
            entry::content::Which::Follows(()) => true,
            _ => false,
        }
    }))
}<|endoftext|>


<commit_before>use lsp_types::{
    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
    DocumentFilter, DocumentHighlight, DocumentHighlightParams, FoldingRange, FoldingRangeParams,
    GotoDefinitionParams, GotoDefinitionResponse, InitializeParams, InitializeResult,
    InitializedParams, Location, MessageType, ReferenceParams, Registration, SemanticTokens,
    SemanticTokensFullOptions, SemanticTokensOptions, SemanticTokensParams,
    SemanticTokensRegistrationOptions, SemanticTokensResult, SemanticTokensServerCapabilities,
    ServerCapabilities, ServerInfo, StaticRegistrationOptions,
    TextDocumentChangeRegistrationOptions, TextDocumentContentChangeEvent,
    TextDocumentRegistrationOptions, Url, WorkDoneProgressOptions,
};
use lspower::{jsonrpc, Client, LanguageServer};
use tokio::sync::Mutex;

use crate::{
    compiler::{ast_to_hir::AstToHir, cst_to_ast::CstToAst, string_to_cst::StringToCst},
    input::{Input, InputDb},
    language_server::hints::HintsDb,
    Database,
};

use self::{
    definition::find_definition,
    folding_range::FoldingRangeDb,
    hints::HintsNotification,
    references::{find_document_highlights, find_references},
    semantic_tokens::SemanticTokenDb,
    utils::LspPositionConversion,
};

pub mod definition;
pub mod folding_range;
pub mod hints;
pub mod references;
pub mod semantic_tokens;
pub mod utils;

pub struct CandyLanguageServer {
    pub client: Client,
    pub db: Mutex<Database>,
}
impl CandyLanguageServer {
    pub fn from_client(client: Client) -> Self {
        Self {
            client,
            db: Mutex::new(Database::default()),
        }
    }
}

#[lspower::async_trait]
impl LanguageServer for CandyLanguageServer {
    async fn initialize(&self, _: InitializeParams) -> jsonrpc::Result<InitializeResult> {
        log::info!("LSP: initialize");
        self.client
            .log_message(MessageType::INFO, "Initializing!")
            .await;
        Ok(InitializeResult {
            // We only support dynamic registration for now.
            capabilities: ServerCapabilities::default(),
            server_info: Some(ServerInfo {
                name: "🍭 Candy Language Server".to_owned(),
                version: None,
            }),
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        log::info!("LSP: initialized");
        let candy_files = DocumentFilter {
            language: Some("candy".to_owned()),
            scheme: Some("file".to_owned()),
            pattern: None,
        };
        let text_document_registration_options = TextDocumentRegistrationOptions {
            document_selector: Some(vec![candy_files.clone()]),
        };
        self.client
            .register_capability(vec![
                Registration {
                    id: "0".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "1".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "2".to_owned(),
                    method: "textDocument/didChange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(TextDocumentChangeRegistrationOptions {
                            document_selector: Some(vec![candy_files]),
                            sync_kind: 2, // incremental
                        })
                        .unwrap(),
                    ),
                },
                Registration {
                    id: "3".to_owned(),
                    method: "textDocument/definition".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "4".to_owned(),
                    method: "textDocument/references".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "5".to_owned(),
                    method: "textDocument/documentHighlight".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "6".to_owned(),
                    method: "textDocument/foldingRange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "7".to_owned(),
                    method: "textDocument/semanticTokens".to_owned(),
                    register_options: Some(
                        serde_json::to_value(
                            SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
                                SemanticTokensRegistrationOptions {
                                    text_document_registration_options,
                                    semantic_tokens_options: SemanticTokensOptions {
                                        work_done_progress_options: WorkDoneProgressOptions {
                                            work_done_progress: None,
                                        },
                                        legend: semantic_tokens::LEGEND.clone(),
                                        range: Some(false),
                                        full: Some(SemanticTokensFullOptions::Bool(true)),
                                    },
                                    static_registration_options: StaticRegistrationOptions {
                                        id: None,
                                    },
                                },
                            ),
                        )
                        .unwrap(),
                    ),
                },
            ])
            .await
            .expect("Dynamic capability registration failed.");
        self.client
            .log_message(MessageType::INFO, "server initialized!")
            .await;
    }

    async fn shutdown(&self) -> jsonrpc::Result<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let input = params.text_document.uri.into();
        {
            let mut db = self.db.lock().await;
            db.did_open_input(&input, params.text_document.text);
        }
        self.analyze_file(input).await;
    }
    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let input: Input = params.text_document.uri.into();
        {
            let mut db = self.db.lock().await;
            let text = apply_text_changes(&db, input.clone(), params.content_changes);
            db.did_change_input(&input, text);
        }
        self.analyze_file(input).await;
    }
    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        let input = params.text_document.uri.into();
        let mut db = self.db.lock().await;
        db.did_close_input(&input);
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> jsonrpc::Result<Option<GotoDefinitionResponse>> {
        let db = self.db.lock().await;
        Ok(find_definition(&db, params))
    }

    async fn references(&self, params: ReferenceParams) -> jsonrpc::Result<Option<Vec<Location>>> {
        let db = self.db.lock().await;
        Ok(find_references(&db, params))
    }
    async fn document_highlight(
        &self,
        params: DocumentHighlightParams,
    ) -> jsonrpc::Result<Option<Vec<DocumentHighlight>>> {
        let db = self.db.lock().await;
        Ok(find_document_highlights(&db, params))
    }

    async fn folding_range(
        &self,
        params: FoldingRangeParams,
    ) -> jsonrpc::Result<Option<Vec<FoldingRange>>> {
        let db = self.db.lock().await;
        let ranges = db.folding_ranges(params.text_document.uri.into());
        Ok(Some(ranges))
    }

    async fn semantic_tokens_full(
        &self,
        params: SemanticTokensParams,
    ) -> jsonrpc::Result<Option<SemanticTokensResult>> {
        let db = self.db.lock().await;
        let tokens = db.semantic_tokens(params.text_document.uri.into());
        Ok(Some(SemanticTokensResult::Tokens(SemanticTokens {
            result_id: None,
            data: tokens,
        })))
    }
}

impl CandyLanguageServer {
    async fn analyze_file(&self, input: Input) {
        log::debug!("Analyzing file. Locking guard.");
        let db = self.db.lock().await;
        log::debug!("Locked.");

        let (_, cst_errors) = db.cst_raw(input.clone()).unwrap();
        let (_, _, ast_errors) = db.ast_raw(input.clone()).unwrap();
        let (_, _, hir_errors) = db.hir_raw(input.clone()).unwrap();

        log::debug!("Mapping errors to diagnostics.");
        let diagnostics = cst_errors
            .into_iter()
            .chain(ast_errors.into_iter())
            .chain(hir_errors.into_iter())
            .map(|it| it.to_diagnostic(&db, input.clone()))
            .collect();
        log::debug!("Publishing diagnostics.");
        self.client
            .publish_diagnostics(input.clone().into(), diagnostics, None)
            .await;
        log::debug!("Published.");
        let hints = db.hints(input.clone());
        self.client
            .send_custom_notification::<HintsNotification>(HintsNotification {
                uri: Url::from(input).to_string(),
                hints,
            })
            .await;
    }
}

fn apply_text_changes(
    db: &Database,
    input: Input,
    changes: Vec<TextDocumentContentChangeEvent>,
) -> String {
    let mut text = db.get_input(input.clone()).unwrap().as_ref().to_owned();
    for change in changes {
        match change.range {
            Some(range) => {
                let start =
                    db.offset_from_lsp(input.clone(), range.start.line, range.start.character);
                let end = db.offset_from_lsp(input.clone(), range.end.line, range.end.character);
                text = format!("{}{}{}", &text[..start], &change.text, &text[end..]);
            }
            None => text = change.text,
        }
    }
    text
}
<commit_msg>feat(compiler): support untitled (non-saved) files in LS<commit_after>use lsp_types::{
    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
    DocumentFilter, DocumentHighlight, DocumentHighlightParams, FoldingRange, FoldingRangeParams,
    GotoDefinitionParams, GotoDefinitionResponse, InitializeParams, InitializeResult,
    InitializedParams, Location, MessageType, ReferenceParams, Registration, SemanticTokens,
    SemanticTokensFullOptions, SemanticTokensOptions, SemanticTokensParams,
    SemanticTokensRegistrationOptions, SemanticTokensResult, SemanticTokensServerCapabilities,
    ServerCapabilities, ServerInfo, StaticRegistrationOptions,
    TextDocumentChangeRegistrationOptions, TextDocumentContentChangeEvent,
    TextDocumentRegistrationOptions, Url, WorkDoneProgressOptions,
};
use lspower::{jsonrpc, Client, LanguageServer};
use tokio::sync::Mutex;

use crate::{
    compiler::{ast_to_hir::AstToHir, cst_to_ast::CstToAst, string_to_cst::StringToCst},
    input::{Input, InputDb},
    language_server::hints::HintsDb,
    Database,
};

use self::{
    definition::find_definition,
    folding_range::FoldingRangeDb,
    hints::HintsNotification,
    references::{find_document_highlights, find_references},
    semantic_tokens::SemanticTokenDb,
    utils::LspPositionConversion,
};

pub mod definition;
pub mod folding_range;
pub mod hints;
pub mod references;
pub mod semantic_tokens;
pub mod utils;

pub struct CandyLanguageServer {
    pub client: Client,
    pub db: Mutex<Database>,
}
impl CandyLanguageServer {
    pub fn from_client(client: Client) -> Self {
        Self {
            client,
            db: Mutex::new(Database::default()),
        }
    }
}

#[lspower::async_trait]
impl LanguageServer for CandyLanguageServer {
    async fn initialize(&self, _: InitializeParams) -> jsonrpc::Result<InitializeResult> {
        log::info!("LSP: initialize");
        self.client
            .log_message(MessageType::INFO, "Initializing!")
            .await;
        Ok(InitializeResult {
            // We only support dynamic registration for now.
            capabilities: ServerCapabilities::default(),
            server_info: Some(ServerInfo {
                name: "🍭 Candy Language Server".to_owned(),
                version: None,
            }),
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        log::info!("LSP: initialized");
        let candy_files = vec![
            DocumentFilter {
                language: Some("candy".to_owned()),
                scheme: Some("file".to_owned()),
                pattern: None,
            },
            DocumentFilter {
                language: Some("candy".to_owned()),
                scheme: Some("untitled".to_owned()),
                pattern: None,
            },
        ];
        let text_document_registration_options = TextDocumentRegistrationOptions {
            document_selector: Some(candy_files.clone()),
        };
        self.client
            .register_capability(vec![
                Registration {
                    id: "0".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "1".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "2".to_owned(),
                    method: "textDocument/didChange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(TextDocumentChangeRegistrationOptions {
                            document_selector: Some(candy_files),
                            sync_kind: 2, // incremental
                        })
                        .unwrap(),
                    ),
                },
                Registration {
                    id: "3".to_owned(),
                    method: "textDocument/definition".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "4".to_owned(),
                    method: "textDocument/references".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "5".to_owned(),
                    method: "textDocument/documentHighlight".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "6".to_owned(),
                    method: "textDocument/foldingRange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "7".to_owned(),
                    method: "textDocument/semanticTokens".to_owned(),
                    register_options: Some(
                        serde_json::to_value(
                            SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
                                SemanticTokensRegistrationOptions {
                                    text_document_registration_options,
                                    semantic_tokens_options: SemanticTokensOptions {
                                        work_done_progress_options: WorkDoneProgressOptions {
                                            work_done_progress: None,
                                        },
                                        legend: semantic_tokens::LEGEND.clone(),
                                        range: Some(false),
                                        full: Some(SemanticTokensFullOptions::Bool(true)),
                                    },
                                    static_registration_options: StaticRegistrationOptions {
                                        id: None,
                                    },
                                },
                            ),
                        )
                        .unwrap(),
                    ),
                },
            ])
            .await
            .expect("Dynamic capability registration failed.");
        self.client
            .log_message(MessageType::INFO, "server initialized!")
            .await;
    }

    async fn shutdown(&self) -> jsonrpc::Result<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let input = params.text_document.uri.into();
        {
            let mut db = self.db.lock().await;
            db.did_open_input(&input, params.text_document.text);
        }
        self.analyze_file(input).await;
    }
    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let input: Input = params.text_document.uri.into();
        {
            let mut db = self.db.lock().await;
            let text = apply_text_changes(&db, input.clone(), params.content_changes);
            db.did_change_input(&input, text);
        }
        self.analyze_file(input).await;
    }
    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        let input = params.text_document.uri.into();
        let mut db = self.db.lock().await;
        db.did_close_input(&input);
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> jsonrpc::Result<Option<GotoDefinitionResponse>> {
        let db = self.db.lock().await;
        Ok(find_definition(&db, params))
    }

    async fn references(&self, params: ReferenceParams) -> jsonrpc::Result<Option<Vec<Location>>> {
        let db = self.db.lock().await;
        Ok(find_references(&db, params))
    }
    async fn document_highlight(
        &self,
        params: DocumentHighlightParams,
    ) -> jsonrpc::Result<Option<Vec<DocumentHighlight>>> {
        let db = self.db.lock().await;
        Ok(find_document_highlights(&db, params))
    }

    async fn folding_range(
        &self,
        params: FoldingRangeParams,
    ) -> jsonrpc::Result<Option<Vec<FoldingRange>>> {
        let db = self.db.lock().await;
        let ranges = db.folding_ranges(params.text_document.uri.into());
        Ok(Some(ranges))
    }

    async fn semantic_tokens_full(
        &self,
        params: SemanticTokensParams,
    ) -> jsonrpc::Result<Option<SemanticTokensResult>> {
        let db = self.db.lock().await;
        let tokens = db.semantic_tokens(params.text_document.uri.into());
        Ok(Some(SemanticTokensResult::Tokens(SemanticTokens {
            result_id: None,
            data: tokens,
        })))
    }
}

impl CandyLanguageServer {
    async fn analyze_file(&self, input: Input) {
        log::debug!("Analyzing file. Locking guard.");
        let db = self.db.lock().await;
        log::debug!("Locked.");

        let (_, cst_errors) = db.cst_raw(input.clone()).unwrap();
        let (_, _, ast_errors) = db.ast_raw(input.clone()).unwrap();
        let (_, _, hir_errors) = db.hir_raw(input.clone()).unwrap();

        log::debug!("Mapping errors to diagnostics.");
        let diagnostics = cst_errors
            .into_iter()
            .chain(ast_errors.into_iter())
            .chain(hir_errors.into_iter())
            .map(|it| it.to_diagnostic(&db, input.clone()))
            .collect();
        log::debug!("Publishing diagnostics.");
        self.client
            .publish_diagnostics(input.clone().into(), diagnostics, None)
            .await;
        log::debug!("Published.");
        let hints = db.hints(input.clone());
        self.client
            .send_custom_notification::<HintsNotification>(HintsNotification {
                uri: Url::from(input).to_string(),
                hints,
            })
            .await;
    }
}

fn apply_text_changes(
    db: &Database,
    input: Input,
    changes: Vec<TextDocumentContentChangeEvent>,
) -> String {
    let mut text = db.get_input(input.clone()).unwrap().as_ref().to_owned();
    for change in changes {
        match change.range {
            Some(range) => {
                let start =
                    db.offset_from_lsp(input.clone(), range.start.line, range.start.character);
                let end = db.offset_from_lsp(input.clone(), range.end.line, range.end.character);
                text = format!("{}{}{}", &text[..start], &change.text, &text[end..]);
            }
            None => text = change.text,
        }
    }
    text
}
<|endoftext|>


<commit_before>//! Implements castling rights logic.

use basetypes::*;

/// Holds information about which player is allowed to castle on which
/// side.
///
/// The castling rights are held in a `usize` value. The lowest 4 bits
/// of the value contain the whole needed information. It is laid out
/// in the following way:
///
/// ```text
///  usize             3           0
///  +---------------+---+---+---+---+
///  |               |   |   |   |   |
///  |    Unused     |Castling flags |
///  |               |   |   |   |   |
///  +---------------+---+---+---+---+
///
///  bit 0 -- if set, white can castle on queen-side;
///  bit 1 -- if set, white can castle on king-side;
///  bit 2 -- if set, black can castle on queen-side;
///  bit 3 -- if set, black can castle on king-side.
/// ```
#[derive(Debug)]
#[derive(Clone, Copy)]
pub struct CastlingRights(usize);

impl CastlingRights {
    /// Creates a new instance. No player can castle on any side.
    #[inline(always)]
    pub fn new() -> CastlingRights {
        CastlingRights(0)
    }

    /// Returns the contained raw value.
    #[inline(always)]
    pub fn value(&self) -> usize {
        self.0
    }

    /// Grants a given player the right to castle on a given side.
    ///
    /// This method returns `true` if the player had the right to
    /// castle on the given side before this method was called, and
    /// `false` otherwise.
    pub fn grant(&mut self, player: Color, side: CastlingSide) -> bool {
        if player > 1 || side > 1 {
            panic!("invalid arguments");
        }
        let before = self.0;
        let mask = 1 << (player << 1) << side;
        self.0 |= mask;
        !before & mask == 0
    }

    /// Updates the castling rights after played move.
    ///
    /// `orig_square` and `dest_square` describe the played move.
    #[inline]
    pub fn update(&mut self, orig_square: Square, dest_square: Square) {
        // On each move, the value of `CASTLING_RELATION` for the
        // origin and destination squares should be &-ed with the
        // castling rights value, to derive the updated castling
        // rights.
        const CASTLING_RELATION: [usize; 64] = [
            !CASTLE_WHITE_QUEENSIDE, !0, !0, !0,
            !(CASTLE_WHITE_QUEENSIDE | CASTLE_WHITE_KINGSIDE), !0, !0, !CASTLE_WHITE_KINGSIDE,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !CASTLE_BLACK_QUEENSIDE, !0, !0, !0,
            !(CASTLE_BLACK_QUEENSIDE | CASTLE_BLACK_KINGSIDE), !0, !0, !CASTLE_BLACK_KINGSIDE,
        ];
        self.0 &= CASTLING_RELATION[orig_square] & CASTLING_RELATION[dest_square];
    }

    /// Returns if a given player can castle on a given side.
    #[inline]
    pub fn can_castle(&self, player: Color, side: CastlingSide) -> bool {
        assert!(player <= 1);
        assert!(side <= 1);
        (1 << (player << 1) << side) & self.0 != 0
    }

    /// Returns a bitboard with potential castling obstacles.
    ///
    /// Returns a bitboard with the set of squares that should be
    /// vacant in order for the specified (`player`, `side`) castling
    /// move to be possible. If `player` can never castle on `side`,
    /// because the king or the rook had been moved, this method
    /// returns universal set (`0xffffffffffffffff`).
    #[inline]
    pub fn obstacles(&self, player: Color, side: CastlingSide) -> u64 {
        const OBSTACLES: [[u64; 2]; 2] = [[1 << B1 | 1 << C1 | 1 << D1, 1 << F1 | 1 << G1],
                                          [1 << B8 | 1 << C8 | 1 << D8, 1 << F8 | 1 << G8]];
        if self.can_castle(player, side) {
            OBSTACLES[player][side]
        } else {
            // Castling is not allowed, therefore every piece on every
            // square on the board can be considered an obstacle.
            !0
        }
    }

    /// Returns a value from 0 to 3 representing the castling rights
    /// for `player`.
    #[inline(always)]
    pub fn get_for(&self, player: Color) -> usize {
        assert!(player <= 1);
        if player == WHITE {
            self.0 & 0b0011
        } else {
            self.0 >> 2
        }
    }

    /// Sets the castling rights for `player` with a value from 0 to
    /// 3.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it is performance-critical, and
    /// so it does not verify that the passed arguments are
    /// valid. Users of this method should make sure that `rights <=
    /// 3`.
    #[inline(always)]
    pub unsafe fn set_for(&mut self, player: Color, rights: usize) {
        assert!(player <= 1);
        assert!(rights <= 3);
        self.0 = if player == WHITE {
            self.0 & 0b1100 | rights
        } else {
            self.0 & 0b0011 | rights << 2
        }
    }
}

/// `QUEENSIDE` of `KINGSIDE`.
pub type CastlingSide = usize;

/// Queen-side castling.
pub const QUEENSIDE: CastlingSide = 0;

/// King-side castling.
pub const KINGSIDE: CastlingSide = 1;

// White can castle on the queen-side.
const CASTLE_WHITE_QUEENSIDE: usize = 1 << 0;

// White can castle on the king-side.
const CASTLE_WHITE_KINGSIDE: usize = 1 << 1;

// Black can castle on the queen-side.
const CASTLE_BLACK_QUEENSIDE: usize = 1 << 2;

// Black can castle on the king-side.
const CASTLE_BLACK_KINGSIDE: usize = 1 << 3;

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

    #[test]
    fn test_castling_rights() {
        use basetypes::*;
        use bitsets::*;

        let mut c = CastlingRights::new();
        unsafe {
            c.set_for(WHITE, 0b10);
            c.set_for(BLACK, 0b11);
        }
        assert_eq!(c.can_castle(WHITE, QUEENSIDE), false);
        assert_eq!(c.can_castle(WHITE, KINGSIDE), true);
        assert_eq!(c.can_castle(BLACK, QUEENSIDE), true);
        assert_eq!(c.can_castle(BLACK, KINGSIDE), true);
        c.update(H8, H7);
        assert_eq!(c.can_castle(WHITE, QUEENSIDE), false);
        assert_eq!(c.can_castle(WHITE, KINGSIDE), true);
        assert_eq!(c.can_castle(BLACK, QUEENSIDE), true);
        assert_eq!(c.can_castle(BLACK, KINGSIDE), false);
        assert_eq!(c.get_for(BLACK), 0b01);
        assert_eq!(c.get_for(WHITE), 0b10);
        assert_eq!(c.value(), 0b0110);
        let granted = c.grant(BLACK, KINGSIDE);
        assert_eq!(granted, false);
        let granted = c.grant(BLACK, KINGSIDE);
        assert_eq!(granted, true);
        assert_eq!(c.value(), 0b1110);
        assert_eq!(c.obstacles(WHITE, QUEENSIDE), UNIVERSAL_SET);
        assert_eq!(c.obstacles(BLACK, QUEENSIDE), 1 << B8 | 1 << C8 | 1 << D8);
    }
}
<commit_msg>update() is now unsafe<commit_after>//! Implements castling rights logic.

use basetypes::*;

/// Holds information about which player is allowed to castle on which
/// side.
///
/// The castling rights are held in a `usize` value. The lowest 4 bits
/// of the value contain the whole needed information. It is laid out
/// in the following way:
///
/// ```text
///  usize             3           0
///  +---------------+---+---+---+---+
///  |               |   |   |   |   |
///  |    Unused     |Castling flags |
///  |               |   |   |   |   |
///  +---------------+---+---+---+---+
///
///  bit 0 -- if set, white can castle on queen-side;
///  bit 1 -- if set, white can castle on king-side;
///  bit 2 -- if set, black can castle on queen-side;
///  bit 3 -- if set, black can castle on king-side.
/// ```
#[derive(Debug)]
#[derive(Clone, Copy)]
pub struct CastlingRights(usize);

impl CastlingRights {
    /// Creates a new instance. No player can castle on any side.
    #[inline(always)]
    pub fn new() -> CastlingRights {
        CastlingRights(0)
    }

    /// Returns the contained raw value.
    #[inline(always)]
    pub fn value(&self) -> usize {
        self.0
    }

    /// Grants a given player the right to castle on a given side.
    ///
    /// This method returns `true` if the player had the right to
    /// castle on the given side before this method was called, and
    /// `false` otherwise.
    pub fn grant(&mut self, player: Color, side: CastlingSide) -> bool {
        if player > 1 || side > 1 {
            panic!("invalid arguments");
        }
        let before = self.0;
        let mask = 1 << (player << 1) << side;
        self.0 |= mask;
        !before & mask == 0
    }

    /// Updates the castling rights after played move.
    ///
    /// `orig_square` and `dest_square` describe the played move.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it is performance-critical, and
    /// so it does not verify that the passed arguments are
    /// valid. Users of this method should make sure that:
    ///
    /// * `orig_square <= 63`.
    /// * `dest_square <= 63`.
    #[inline]
    pub unsafe fn update(&mut self, orig_square: Square, dest_square: Square) {
        assert!(orig_square <= 63);
        assert!(dest_square <= 63);
        // On each move, the value of `CASTLING_RELATION` for the
        // origin and destination squares should be &-ed with the
        // castling rights value, to derive the updated castling
        // rights.
        const CASTLING_RELATION: [usize; 64] = [
            !CASTLE_WHITE_QUEENSIDE, !0, !0, !0,
            !(CASTLE_WHITE_QUEENSIDE | CASTLE_WHITE_KINGSIDE), !0, !0, !CASTLE_WHITE_KINGSIDE,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !0, !0, !0, !0, !0, !0, !0, !0,
            !CASTLE_BLACK_QUEENSIDE, !0, !0, !0,
            !(CASTLE_BLACK_QUEENSIDE | CASTLE_BLACK_KINGSIDE), !0, !0, !CASTLE_BLACK_KINGSIDE,
        ];
        self.0 &= *CASTLING_RELATION.get_unchecked(orig_square) &
                  *CASTLING_RELATION.get_unchecked(dest_square);
    }

    /// Returns if a given player can castle on a given side.
    #[inline]
    pub fn can_castle(&self, player: Color, side: CastlingSide) -> bool {
        assert!(player <= 1);
        assert!(side <= 1);
        (1 << (player << 1) << side) & self.0 != 0
    }

    /// Returns a bitboard with potential castling obstacles.
    ///
    /// Returns a bitboard with the set of squares that should be
    /// vacant in order for the specified (`player`, `side`) castling
    /// move to be possible. If `player` can never castle on `side`,
    /// because the king or the rook had been moved, this method
    /// returns universal set (`0xffffffffffffffff`).
    #[inline]
    pub fn obstacles(&self, player: Color, side: CastlingSide) -> u64 {
        const OBSTACLES: [[u64; 2]; 2] = [[1 << B1 | 1 << C1 | 1 << D1, 1 << F1 | 1 << G1],
                                          [1 << B8 | 1 << C8 | 1 << D8, 1 << F8 | 1 << G8]];
        if self.can_castle(player, side) {
            OBSTACLES[player][side]
        } else {
            // Castling is not allowed, therefore every piece on every
            // square on the board can be considered an obstacle.
            !0
        }
    }

    /// Returns a value from 0 to 3 representing the castling rights
    /// for `player`.
    #[inline(always)]
    pub fn get_for(&self, player: Color) -> usize {
        assert!(player <= 1);
        if player == WHITE {
            self.0 & 0b0011
        } else {
            self.0 >> 2
        }
    }

    /// Sets the castling rights for `player` with a value from 0 to
    /// 3.
    ///
    /// # Safety
    ///
    /// This method is unsafe because it is performance-critical, and
    /// so it does not verify that the passed arguments are
    /// valid. Users of this method should make sure that `rights <=
    /// 3`.
    #[inline(always)]
    pub unsafe fn set_for(&mut self, player: Color, rights: usize) {
        assert!(player <= 1);
        assert!(rights <= 3);
        self.0 = if player == WHITE {
            self.0 & 0b1100 | rights
        } else {
            self.0 & 0b0011 | rights << 2
        }
    }
}

/// `QUEENSIDE` of `KINGSIDE`.
pub type CastlingSide = usize;

/// Queen-side castling.
pub const QUEENSIDE: CastlingSide = 0;

/// King-side castling.
pub const KINGSIDE: CastlingSide = 1;

// White can castle on the queen-side.
const CASTLE_WHITE_QUEENSIDE: usize = 1 << 0;

// White can castle on the king-side.
const CASTLE_WHITE_KINGSIDE: usize = 1 << 1;

// Black can castle on the queen-side.
const CASTLE_BLACK_QUEENSIDE: usize = 1 << 2;

// Black can castle on the king-side.
const CASTLE_BLACK_KINGSIDE: usize = 1 << 3;

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

    #[test]
    fn test_castling_rights() {
        use basetypes::*;
        use bitsets::*;

        let mut c = CastlingRights::new();
        unsafe {
            c.set_for(WHITE, 0b10);
            c.set_for(BLACK, 0b11);
        }
        assert_eq!(c.can_castle(WHITE, QUEENSIDE), false);
        assert_eq!(c.can_castle(WHITE, KINGSIDE), true);
        assert_eq!(c.can_castle(BLACK, QUEENSIDE), true);
        assert_eq!(c.can_castle(BLACK, KINGSIDE), true);
        unsafe {
            c.update(H8, H7);
        }
        assert_eq!(c.can_castle(WHITE, QUEENSIDE), false);
        assert_eq!(c.can_castle(WHITE, KINGSIDE), true);
        assert_eq!(c.can_castle(BLACK, QUEENSIDE), true);
        assert_eq!(c.can_castle(BLACK, KINGSIDE), false);
        assert_eq!(c.get_for(BLACK), 0b01);
        assert_eq!(c.get_for(WHITE), 0b10);
        assert_eq!(c.value(), 0b0110);
        let granted = c.grant(BLACK, KINGSIDE);
        assert_eq!(granted, false);
        let granted = c.grant(BLACK, KINGSIDE);
        assert_eq!(granted, true);
        assert_eq!(c.value(), 0b1110);
        assert_eq!(c.obstacles(WHITE, QUEENSIDE), UNIVERSAL_SET);
        assert_eq!(c.obstacles(BLACK, QUEENSIDE), 1 << B8 | 1 << C8 | 1 << D8);
    }
}
<|endoftext|>


<commit_before>use utils::bigint::M256;
use utils::gas::Gas;
use super::commit::{AccountState, BlockhashState};
use super::errors::{RequireError, MachineError, CommitError};
use super::{Stack, Context, BlockHeader, Patch, PC, Storage, Memory, AccountCommitment};

pub mod cost;
pub mod run;

/// A VM state without PC.
pub struct State<M, S> {
    pub memory: M,
    pub stack: Stack,

    pub context: Context,
    pub block: BlockHeader,
    pub patch: Patch,

    pub out: Vec<u8>,

    pub memory_gas: Gas,
    pub used_gas: Gas,
    pub refunded_gas: Gas,

    pub account_state: AccountState<S>,
    pub blockhash_state: BlockhashState,
}

/// A VM state with PC.
pub struct Machine<M, S> {
    state: State<M, S>,
    pc: PC,
    status: Status,
}

#[derive(Debug, Clone)]
pub enum Status {
    Running,
    ExitedOk,
    ExitedError(MachineError),
    InvokeCall(Context, (M256, M256)),
}

impl<M: Memory + Default, S: Storage + Default + Clone> Machine<M, S> {
    #[allow(unused_variables)]
    pub fn new(context: Context, block: BlockHeader, patch: Patch) -> Self {
        unimplemented!()
    }

    pub fn derive(&self, context: Context) -> Self {
        Machine {
            pc: PC::new(context.code.as_slice()),
            status: Status::Running,
            state: State {
                memory: M::default(),
                stack: Stack::default(),

                context: context,
                block: self.state.block.clone(),
                patch: self.state.patch.clone(),

                out: Vec::new(),

                memory_gas: Gas::zero(),
                used_gas: Gas::zero(),
                refunded_gas: Gas::zero(),

                account_state: self.state.account_state.clone(),
                blockhash_state: self.state.blockhash_state.clone()
            },
        }
    }

    pub fn commit_account(&mut self, commitment: AccountCommitment<S>) -> Result<(), CommitError> {
        self.state.account_state.commit(commitment)
    }

    pub fn commit_blockhash(&mut self, number: M256, hash: M256) -> Result<(), CommitError> {
        self.state.blockhash_state.commit(number, hash)
    }

    #[allow(unused_variables)]
    pub fn apply_sub(&mut self, sub: Machine<M, S>) {
        unimplemented!()
    }

    #[allow(unused_variables)]
    fn invoke_sub(&mut self, context: Context, from: M256, len: M256) {
        unimplemented!()
    }

    pub fn step(&mut self) -> Result<(), RequireError> {
        unimplemented!()
    }

    pub fn status(&self) -> &Status {
        &self.status
    }
}
<commit_msg>Problem: new is unimplemented Solution: implement it<commit_after>use utils::bigint::M256;
use utils::gas::Gas;
use super::commit::{AccountState, BlockhashState};
use super::errors::{RequireError, MachineError, CommitError};
use super::{Stack, Context, BlockHeader, Patch, PC, Storage, Memory, AccountCommitment};

pub mod cost;
pub mod run;

/// A VM state without PC.
pub struct State<M, S> {
    pub memory: M,
    pub stack: Stack,

    pub context: Context,
    pub block: BlockHeader,
    pub patch: Patch,

    pub out: Vec<u8>,

    pub memory_gas: Gas,
    pub used_gas: Gas,
    pub refunded_gas: Gas,

    pub account_state: AccountState<S>,
    pub blockhash_state: BlockhashState,
}

/// A VM state with PC.
pub struct Machine<M, S> {
    state: State<M, S>,
    pc: PC,
    status: Status,
}

#[derive(Debug, Clone)]
pub enum Status {
    Running,
    ExitedOk,
    ExitedError(MachineError),
    InvokeCall(Context, (M256, M256)),
}

impl<M: Memory + Default, S: Storage + Default + Clone> Machine<M, S> {
    pub fn new(context: Context, block: BlockHeader, patch: Patch) -> Self {
        Machine {
            pc: PC::new(context.code.as_slice()),
            status: Status::Running,
            state: State {
                memory: M::default(),
                stack: Stack::default(),

                context: context,
                block: block,
                patch: patch,

                out: Vec::new(),

                memory_gas: Gas::zero(),
                used_gas: Gas::zero(),
                refunded_gas: Gas::zero(),

                account_state: AccountState::default(),
                blockhash_state: BlockhashState::default(),
            },
        }
    }

    pub fn derive(&self, context: Context) -> Self {
        Machine {
            pc: PC::new(context.code.as_slice()),
            status: Status::Running,
            state: State {
                memory: M::default(),
                stack: Stack::default(),

                context: context,
                block: self.state.block.clone(),
                patch: self.state.patch.clone(),

                out: Vec::new(),

                memory_gas: Gas::zero(),
                used_gas: Gas::zero(),
                refunded_gas: Gas::zero(),

                account_state: self.state.account_state.clone(),
                blockhash_state: self.state.blockhash_state.clone()
            },
        }
    }

    pub fn commit_account(&mut self, commitment: AccountCommitment<S>) -> Result<(), CommitError> {
        self.state.account_state.commit(commitment)
    }

    pub fn commit_blockhash(&mut self, number: M256, hash: M256) -> Result<(), CommitError> {
        self.state.blockhash_state.commit(number, hash)
    }

    #[allow(unused_variables)]
    pub fn apply_sub(&mut self, sub: Machine<M, S>) {
        unimplemented!()
    }

    #[allow(unused_variables)]
    fn invoke_sub(&mut self, context: Context, from: M256, len: M256) {
        unimplemented!()
    }

    pub fn step(&mut self) -> Result<(), RequireError> {
        unimplemented!()
    }

    pub fn status(&self) -> &Status {
        &self.status
    }
}
<|endoftext|>


<commit_before>use crate::{
    channel_closer::ChannelCloser,
    consumer_canceler::ConsumerCanceler,
    consumer_status::{ConsumerState, ConsumerStatus},
    executor::Executor,
    internal_rpc::InternalRPCHandle,
    message::{Delivery, DeliveryResult},
    types::ShortString,
    BasicProperties, Channel, Error, Result,
};
use flume::{Receiver, Sender};
use futures_lite::Stream;
use parking_lot::Mutex;
use std::{
    fmt,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
};
use tracing::trace;

pub trait ConsumerDelegate: Send + Sync {
    fn on_new_delivery(&self, delivery: DeliveryResult)
        -> Pin<Box<dyn Future<Output = ()> + Send>>;
    fn drop_prefetched_messages(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(async move {})
    }
}

impl<
        F: Future<Output = ()> + Send + 'static,
        DeliveryHandler: Fn(DeliveryResult) -> F + Send + Sync + 'static,
    > ConsumerDelegate for DeliveryHandler
{
    fn on_new_delivery(
        &self,
        delivery: DeliveryResult,
    ) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(self(delivery))
    }
}

/// Continuously consumes message from a Queue.
///
/// A consumer represents a stream of messages created from
/// the [basic.consume](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.consume) AMQP command.
/// It continuously receives messages from the queue, as opposed to the
/// [basic.get](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.get) command, which
/// retrieves only a single message.
///
/// A consumer is obtained by calling [`Channel::basic_consume`] with the queue name.
///
/// New messages from this consumer can be accessed by obtaining the iterator from the consumer.
/// This iterator returns new messages and the associated channel in the form of a
/// [`DeliveryResult`] for as long as the consumer is subscribed to the queue.
///
/// It is also possible to set a delegate to be spawned via [`set_delegate`].
///
/// ## Message acknowledgment
///
/// There are two ways for acknowledging a message:
///
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `true` while obtaining the consumer from
///   [`Channel::basic_consume`], the server implicitely acknowledges each message after it has been
///   sent.
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `false`, a message has to be explicitely
///   acknowledged or rejected with [`Channel::basic_ack`],
///   [`Channel::basic_reject`] or [`Channel::basic_nack`]. See the documentation at [`Delivery`]
///   for further information.
///
/// Also see the RabbitMQ documentation about
/// [Acknowledgement Modes](https://www.rabbitmq.com/consumers.html#acknowledgement-modes).
///
/// ## Consumer Prefetch
///
/// To limit the maximum number of unacknowledged messages arriving, you can call [`Channel::basic_qos`]
/// before creating the consumer.
///
/// Also see the RabbitMQ documentation about
/// [Consumer Prefetch](https://www.rabbitmq.com/consumer-prefetch.html).
///
/// ## Cancel subscription
///
/// To stop receiving messages, call [`Channel::basic_cancel`] with the consumer tag of this
/// consumer.
///
///
/// ## Example
/// ```rust,no_run
/// use lapin::{options::*, types::FieldTable, Connection, ConnectionProperties, Result};
/// use futures_lite::stream::StreamExt;
/// use std::future::Future;
///
/// let addr = std::env::var("AMQP_ADDR").unwrap_or_else(|_| "amqp://127.0.0.1:5672/%2f".into());
///
/// let res: Result<()> = async_global_executor::block_on(async {
///     let conn = Connection::connect(
///         &addr,
///         ConnectionProperties::default(),
///     )
///     .await?;
///     let channel = conn.create_channel().await?;
///     let mut consumer = channel
///         .basic_consume(
///             "hello",
///             "my_consumer",
///             BasicConsumeOptions::default(),
///             FieldTable::default(),
///         )
///         .await?;
///
///     while let Some(delivery) = consumer.next().await {
///         let (channel, delivery) = delivery.expect("error in consumer");
///         channel
///             .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
///             .await?;
///     }
///     Ok(())
/// });
/// ```
///
/// [`Channel::basic_consume`]: ./struct.Channel.html#method.basic_consume
/// [`Channel::basic_qos`]: ./struct.Channel.html#method.basic_qos
/// [`Channel::basic_ack`]: ./struct.Channel.html#method.basic_ack
/// [`Channel::basic_reject`]: ./struct.Channel.html#method.basic_reject
/// [`Channel::basic_nack`]: ./struct.Channel.html#method.basic_nack
/// [`Channel::basic_cancel`]: ./struct.Channel.html#method.basic_cancel
/// [`DeliveryResult`]: ./message/type.DeliveryResult.html
/// [`BasicConsumeOptions::no_ack`]: ./options/struct.BasicConsumeOptions.html#structfield.no_ack
/// [`set_delegate`]: #method.set_delegate
#[derive(Clone)]
pub struct Consumer {
    inner: Arc<Mutex<ConsumerInner>>,
    status: ConsumerStatus,
    channel_closer: Option<Arc<ChannelCloser>>,
    consumer_canceler: Option<Arc<ConsumerCanceler>>,
}

impl Consumer {
    pub(crate) fn new(
        consumer_tag: ShortString,
        executor: Arc<dyn Executor>,
        channel_closer: Option<Arc<ChannelCloser>>,
    ) -> Self {
        let status = ConsumerStatus::default();
        Self {
            inner: Arc::new(Mutex::new(ConsumerInner::new(
                status.clone(),
                consumer_tag,
                executor,
            ))),
            status,
            channel_closer,
            consumer_canceler: None,
        }
    }

    pub(crate) fn external(&self, channel_id: u16, internal_rpc_handle: InternalRPCHandle) -> Self {
        Self {
            inner: self.inner.clone(),
            status: self.status.clone(),
            channel_closer: None,
            consumer_canceler: Some(Arc::new(ConsumerCanceler::new(
                channel_id,
                self.tag().to_string(),
                self.status.clone(),
                internal_rpc_handle,
            ))),
        }
    }

    /// Gets the consumer tag.
    ///
    /// If no consumer tag was specified when obtaining the consumer from the channel,
    /// this contains the server generated consumer tag.
    pub fn tag(&self) -> ShortString {
        self.inner.lock().tag.clone()
    }

    /// Gets the current state of the Consumer.
    pub fn state(&self) -> ConsumerState {
        self.status.state()
    }

    /// Automatically spawns the delegate on the executor for each message.
    ///
    /// Enables parallel handling of the messages.
    pub fn set_delegate<D: ConsumerDelegate + 'static>(&self, delegate: D) {
        let mut inner = self.inner.lock();
        let mut status = self.status.lock();
        while let Some(delivery) = inner.next_delivery() {
            inner.executor.spawn(delegate.on_new_delivery(delivery));
        }
        inner.delegate = Some(Arc::new(Box::new(delegate)));
        status.set_delegate();
    }

    pub(crate) fn start_new_delivery(&mut self, delivery: Delivery) {
        self.inner.lock().current_message = Some(delivery)
    }

    pub(crate) fn set_delivery_properties(&mut self, properties: BasicProperties) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.properties = properties;
        }
    }

    pub(crate) fn receive_delivery_content(&mut self, payload: Vec<u8>) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.receive_content(payload);
        }
    }

    pub(crate) fn new_delivery_complete(&mut self, channel: Channel) {
        let mut inner = self.inner.lock();
        if let Some(delivery) = inner.current_message.take() {
            inner.new_delivery(channel, delivery);
        }
    }

    pub(crate) fn drop_prefetched_messages(&self) {
        self.inner.lock().drop_prefetched_messages();
    }

    pub(crate) fn cancel(&self) {
        self.inner.lock().cancel();
    }

    pub(crate) fn set_error(&self, error: Error) {
        self.inner.lock().set_error(error);
    }
}

struct ConsumerInner {
    status: ConsumerStatus,
    current_message: Option<Delivery>,
    deliveries_in: Sender<DeliveryResult>,
    deliveries_out: Receiver<DeliveryResult>,
    task: Option<Waker>,
    tag: ShortString,
    delegate: Option<Arc<Box<dyn ConsumerDelegate>>>,
    executor: Arc<dyn Executor>,
}

pub struct ConsumerIterator {
    receiver: Receiver<DeliveryResult>,
}

impl Iterator for ConsumerIterator {
    type Item = Result<(Channel, Delivery)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok().and_then(Result::transpose)
    }
}

impl IntoIterator for Consumer {
    type Item = Result<(Channel, Delivery)>;
    type IntoIter = ConsumerIterator;

    fn into_iter(self) -> Self::IntoIter {
        ConsumerIterator {
            receiver: self.inner.lock().deliveries_out.clone(),
        }
    }
}

impl fmt::Debug for Consumer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Consumer");
        if let Some(inner) = self.inner.try_lock() {
            debug
                .field("tag", &inner.tag)
                .field("executor", &inner.executor)
                .field("task", &inner.task);
        }
        if let Some(status) = self.status.try_lock() {
            debug.field("state", &status.state());
        }
        debug.finish()
    }
}

impl ConsumerInner {
    fn new(status: ConsumerStatus, consumer_tag: ShortString, executor: Arc<dyn Executor>) -> Self {
        let (sender, receiver) = flume::unbounded();
        Self {
            status,
            current_message: None,
            deliveries_in: sender,
            deliveries_out: receiver,
            task: None,
            tag: consumer_tag,
            delegate: None,
            executor,
        }
    }

    fn next_delivery(&mut self) -> Option<DeliveryResult> {
        self.deliveries_out.try_recv().ok()
    }

    fn new_delivery(&mut self, channel: Channel, delivery: Delivery) {
        trace!(consumer_tag=%self.tag, "new_delivery");
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor
                .spawn(delegate.on_new_delivery(Ok(Some((channel, delivery)))));
        } else {
            self.deliveries_in
                .send(Ok(Some((channel, delivery))))
                .expect("failed to send delivery to consumer");
        }
        if let Some(task) = self.task.as_ref() {
            task.wake_by_ref();
        }
    }

    fn drop_prefetched_messages(&mut self) {
        trace!(consumer_tag=%self.tag, "drop_prefetched_messages");
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.drop_prefetched_messages());
        }
        while self.next_delivery().is_some() {}
    }

    fn cancel(&mut self) {
        trace!(consumer_tag=%self.tag, "cancel");
        let mut status = self.status.lock();
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Ok(None)));
        } else {
            self.deliveries_in
                .send(Ok(None))
                .expect("failed to send cancel to consumer");
        }
        if let Some(task) = self.task.take() {
            task.wake();
        }
        status.cancel();
    }

    fn set_error(&mut self, error: Error) {
        trace!(consumer_tag=%self.tag, "set_error");
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Err(error)));
        } else {
            self.deliveries_in
                .send(Err(error))
                .expect("failed to send error to consumer");
        }
        self.cancel();
    }
}

impl Stream for Consumer {
    type Item = Result<(Channel, Delivery)>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        trace!("consumer poll_next");
        let mut inner = self.inner.lock();
        trace!(
            consumer_tag=%inner.tag,
            "consumer poll; acquired inner lock"
        );
        inner.task = Some(cx.waker().clone());
        if let Some(delivery) = inner.next_delivery() {
            match delivery {
                Ok(Some((channel, delivery))) => {
                    trace!(
                        channel=%channel.id(),
                        consumer_tag=%inner.tag,
                        delivery_tag=?delivery.delivery_tag,
                        "delivery"
                    );
                    Poll::Ready(Some(Ok((channel, delivery))))
                }
                Ok(None) => {
                    trace!(consumer_tag=%inner.tag, "consumer canceled");
                    Poll::Ready(None)
                }
                Err(error) => Poll::Ready(Some(Err(error))),
            }
        } else {
            trace!(consumer_tag=%inner.tag, "delivery; status=NotReady");
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod futures_tests {
    use super::*;
    use crate::executor::DefaultExecutor;

    use std::sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    };
    use std::task::{Context, Poll};

    use futures_lite::stream::StreamExt;
    use waker_fn::waker_fn;

    #[test]
    fn stream_on_cancel() {
        let awoken_count = Arc::new(AtomicUsize::new(0));
        let waker = {
            let awoken_count = awoken_count.clone();
            waker_fn(move || {
                awoken_count.fetch_add(1, Ordering::SeqCst);
            })
        };
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            DefaultExecutor::default().unwrap(),
            None,
        );
        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 0);
            assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Pending);
        }

        consumer.cancel();

        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 1);
            assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Ready(None));
        }
    }

    #[test]
    fn stream_on_error() {
        let awoken_count = Arc::new(AtomicUsize::new(0));
        let waker = {
            let awoken_count = awoken_count.clone();
            waker_fn(move || {
                awoken_count.fetch_add(1, Ordering::SeqCst);
            })
        };
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            DefaultExecutor::default().unwrap(),
            None,
        );
        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 0);
            assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Pending);
        }

        consumer.set_error(Error::ChannelsLimitReached);

        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 1);
            assert_eq!(
                Pin::new(&mut next).poll(&mut cx),
                Poll::Ready(Some(Err(Error::ChannelsLimitReached)))
            );
        }
    }
}
<commit_msg>consumer: use Wakers<commit_after>use crate::{
    channel_closer::ChannelCloser,
    consumer_canceler::ConsumerCanceler,
    consumer_status::{ConsumerState, ConsumerStatus},
    executor::Executor,
    internal_rpc::InternalRPCHandle,
    message::{Delivery, DeliveryResult},
    types::ShortString,
    wakers::Wakers,
    BasicProperties, Channel, Error, Result,
};
use flume::{Receiver, Sender};
use futures_lite::Stream;
use parking_lot::Mutex;
use std::{
    fmt,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tracing::trace;

pub trait ConsumerDelegate: Send + Sync {
    fn on_new_delivery(&self, delivery: DeliveryResult)
        -> Pin<Box<dyn Future<Output = ()> + Send>>;
    fn drop_prefetched_messages(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(async move {})
    }
}

impl<
        F: Future<Output = ()> + Send + 'static,
        DeliveryHandler: Fn(DeliveryResult) -> F + Send + Sync + 'static,
    > ConsumerDelegate for DeliveryHandler
{
    fn on_new_delivery(
        &self,
        delivery: DeliveryResult,
    ) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(self(delivery))
    }
}

/// Continuously consumes message from a Queue.
///
/// A consumer represents a stream of messages created from
/// the [basic.consume](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.consume) AMQP command.
/// It continuously receives messages from the queue, as opposed to the
/// [basic.get](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.get) command, which
/// retrieves only a single message.
///
/// A consumer is obtained by calling [`Channel::basic_consume`] with the queue name.
///
/// New messages from this consumer can be accessed by obtaining the iterator from the consumer.
/// This iterator returns new messages and the associated channel in the form of a
/// [`DeliveryResult`] for as long as the consumer is subscribed to the queue.
///
/// It is also possible to set a delegate to be spawned via [`set_delegate`].
///
/// ## Message acknowledgment
///
/// There are two ways for acknowledging a message:
///
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `true` while obtaining the consumer from
///   [`Channel::basic_consume`], the server implicitely acknowledges each message after it has been
///   sent.
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `false`, a message has to be explicitely
///   acknowledged or rejected with [`Channel::basic_ack`],
///   [`Channel::basic_reject`] or [`Channel::basic_nack`]. See the documentation at [`Delivery`]
///   for further information.
///
/// Also see the RabbitMQ documentation about
/// [Acknowledgement Modes](https://www.rabbitmq.com/consumers.html#acknowledgement-modes).
///
/// ## Consumer Prefetch
///
/// To limit the maximum number of unacknowledged messages arriving, you can call [`Channel::basic_qos`]
/// before creating the consumer.
///
/// Also see the RabbitMQ documentation about
/// [Consumer Prefetch](https://www.rabbitmq.com/consumer-prefetch.html).
///
/// ## Cancel subscription
///
/// To stop receiving messages, call [`Channel::basic_cancel`] with the consumer tag of this
/// consumer.
///
///
/// ## Example
/// ```rust,no_run
/// use lapin::{options::*, types::FieldTable, Connection, ConnectionProperties, Result};
/// use futures_lite::stream::StreamExt;
/// use std::future::Future;
///
/// let addr = std::env::var("AMQP_ADDR").unwrap_or_else(|_| "amqp://127.0.0.1:5672/%2f".into());
///
/// let res: Result<()> = async_global_executor::block_on(async {
///     let conn = Connection::connect(
///         &addr,
///         ConnectionProperties::default(),
///     )
///     .await?;
///     let channel = conn.create_channel().await?;
///     let mut consumer = channel
///         .basic_consume(
///             "hello",
///             "my_consumer",
///             BasicConsumeOptions::default(),
///             FieldTable::default(),
///         )
///         .await?;
///
///     while let Some(delivery) = consumer.next().await {
///         let (channel, delivery) = delivery.expect("error in consumer");
///         channel
///             .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
///             .await?;
///     }
///     Ok(())
/// });
/// ```
///
/// [`Channel::basic_consume`]: ./struct.Channel.html#method.basic_consume
/// [`Channel::basic_qos`]: ./struct.Channel.html#method.basic_qos
/// [`Channel::basic_ack`]: ./struct.Channel.html#method.basic_ack
/// [`Channel::basic_reject`]: ./struct.Channel.html#method.basic_reject
/// [`Channel::basic_nack`]: ./struct.Channel.html#method.basic_nack
/// [`Channel::basic_cancel`]: ./struct.Channel.html#method.basic_cancel
/// [`DeliveryResult`]: ./message/type.DeliveryResult.html
/// [`BasicConsumeOptions::no_ack`]: ./options/struct.BasicConsumeOptions.html#structfield.no_ack
/// [`set_delegate`]: #method.set_delegate
#[derive(Clone)]
pub struct Consumer {
    inner: Arc<Mutex<ConsumerInner>>,
    status: ConsumerStatus,
    channel_closer: Option<Arc<ChannelCloser>>,
    consumer_canceler: Option<Arc<ConsumerCanceler>>,
}

impl Consumer {
    pub(crate) fn new(
        consumer_tag: ShortString,
        executor: Arc<dyn Executor>,
        channel_closer: Option<Arc<ChannelCloser>>,
    ) -> Self {
        let status = ConsumerStatus::default();
        Self {
            inner: Arc::new(Mutex::new(ConsumerInner::new(
                status.clone(),
                consumer_tag,
                executor,
            ))),
            status,
            channel_closer,
            consumer_canceler: None,
        }
    }

    pub(crate) fn external(&self, channel_id: u16, internal_rpc_handle: InternalRPCHandle) -> Self {
        Self {
            inner: self.inner.clone(),
            status: self.status.clone(),
            channel_closer: None,
            consumer_canceler: Some(Arc::new(ConsumerCanceler::new(
                channel_id,
                self.tag().to_string(),
                self.status.clone(),
                internal_rpc_handle,
            ))),
        }
    }

    /// Gets the consumer tag.
    ///
    /// If no consumer tag was specified when obtaining the consumer from the channel,
    /// this contains the server generated consumer tag.
    pub fn tag(&self) -> ShortString {
        self.inner.lock().tag.clone()
    }

    /// Gets the current state of the Consumer.
    pub fn state(&self) -> ConsumerState {
        self.status.state()
    }

    /// Automatically spawns the delegate on the executor for each message.
    ///
    /// Enables parallel handling of the messages.
    pub fn set_delegate<D: ConsumerDelegate + 'static>(&self, delegate: D) {
        let mut inner = self.inner.lock();
        let mut status = self.status.lock();
        while let Some(delivery) = inner.next_delivery() {
            inner.executor.spawn(delegate.on_new_delivery(delivery));
        }
        inner.delegate = Some(Arc::new(Box::new(delegate)));
        status.set_delegate();
    }

    pub(crate) fn start_new_delivery(&mut self, delivery: Delivery) {
        self.inner.lock().current_message = Some(delivery)
    }

    pub(crate) fn set_delivery_properties(&mut self, properties: BasicProperties) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.properties = properties;
        }
    }

    pub(crate) fn receive_delivery_content(&mut self, payload: Vec<u8>) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.receive_content(payload);
        }
    }

    pub(crate) fn new_delivery_complete(&mut self, channel: Channel) {
        let mut inner = self.inner.lock();
        if let Some(delivery) = inner.current_message.take() {
            inner.new_delivery(channel, delivery);
        }
    }

    pub(crate) fn drop_prefetched_messages(&self) {
        self.inner.lock().drop_prefetched_messages();
    }

    pub(crate) fn cancel(&self) {
        self.inner.lock().cancel();
    }

    pub(crate) fn set_error(&self, error: Error) {
        self.inner.lock().set_error(error);
    }
}

struct ConsumerInner {
    status: ConsumerStatus,
    current_message: Option<Delivery>,
    deliveries_in: Sender<DeliveryResult>,
    deliveries_out: Receiver<DeliveryResult>,
    wakers: Wakers,
    tag: ShortString,
    delegate: Option<Arc<Box<dyn ConsumerDelegate>>>,
    executor: Arc<dyn Executor>,
}

pub struct ConsumerIterator {
    receiver: Receiver<DeliveryResult>,
}

impl Iterator for ConsumerIterator {
    type Item = Result<(Channel, Delivery)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok().and_then(Result::transpose)
    }
}

impl IntoIterator for Consumer {
    type Item = Result<(Channel, Delivery)>;
    type IntoIter = ConsumerIterator;

    fn into_iter(self) -> Self::IntoIter {
        ConsumerIterator {
            receiver: self.inner.lock().deliveries_out.clone(),
        }
    }
}

impl fmt::Debug for Consumer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Consumer");
        if let Some(inner) = self.inner.try_lock() {
            debug
                .field("tag", &inner.tag)
                .field("executor", &inner.executor);
        }
        if let Some(status) = self.status.try_lock() {
            debug.field("state", &status.state());
        }
        debug.finish()
    }
}

impl ConsumerInner {
    fn new(status: ConsumerStatus, consumer_tag: ShortString, executor: Arc<dyn Executor>) -> Self {
        let (sender, receiver) = flume::unbounded();
        Self {
            status,
            current_message: None,
            deliveries_in: sender,
            deliveries_out: receiver,
            wakers: Wakers::default(),
            tag: consumer_tag,
            delegate: None,
            executor,
        }
    }

    fn next_delivery(&mut self) -> Option<DeliveryResult> {
        self.deliveries_out.try_recv().ok()
    }

    fn new_delivery(&mut self, channel: Channel, delivery: Delivery) {
        trace!(consumer_tag=%self.tag, "new_delivery");
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor
                .spawn(delegate.on_new_delivery(Ok(Some((channel, delivery)))));
        } else {
            self.deliveries_in
                .send(Ok(Some((channel, delivery))))
                .expect("failed to send delivery to consumer");
        }
        self.wakers.wake();
    }

    fn drop_prefetched_messages(&mut self) {
        trace!(consumer_tag=%self.tag, "drop_prefetched_messages");
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.drop_prefetched_messages());
        }
        while self.next_delivery().is_some() {}
    }

    fn cancel(&mut self) {
        trace!(consumer_tag=%self.tag, "cancel");
        let mut status = self.status.lock();
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Ok(None)));
        } else {
            self.deliveries_in
                .send(Ok(None))
                .expect("failed to send cancel to consumer");
        }
        self.wakers.wake();
        status.cancel();
    }

    fn set_error(&mut self, error: Error) {
        trace!(consumer_tag=%self.tag, "set_error");
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Err(error)));
        } else {
            self.deliveries_in
                .send(Err(error))
                .expect("failed to send error to consumer");
        }
        self.cancel();
    }
}

impl Stream for Consumer {
    type Item = Result<(Channel, Delivery)>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        trace!("consumer poll_next");
        let mut inner = self.inner.lock();
        trace!(
            consumer_tag=%inner.tag,
            "consumer poll; acquired inner lock"
        );
        inner.wakers.register(cx.waker().clone());
        if let Some(delivery) = inner.next_delivery() {
            match delivery {
                Ok(Some((channel, delivery))) => {
                    trace!(
                        channel=%channel.id(),
                        consumer_tag=%inner.tag,
                        delivery_tag=?delivery.delivery_tag,
                        "delivery"
                    );
                    Poll::Ready(Some(Ok((channel, delivery))))
                }
                Ok(None) => {
                    trace!(consumer_tag=%inner.tag, "consumer canceled");
                    Poll::Ready(None)
                }
                Err(error) => Poll::Ready(Some(Err(error))),
            }
        } else {
            trace!(consumer_tag=%inner.tag, "delivery; status=NotReady");
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod futures_tests {
    use super::*;
    use crate::executor::DefaultExecutor;

    use std::sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    };
    use std::task::{Context, Poll};

    use futures_lite::stream::StreamExt;
    use waker_fn::waker_fn;

    #[test]
    fn stream_on_cancel() {
        let awoken_count = Arc::new(AtomicUsize::new(0));
        let waker = {
            let awoken_count = awoken_count.clone();
            waker_fn(move || {
                awoken_count.fetch_add(1, Ordering::SeqCst);
            })
        };
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            DefaultExecutor::default().unwrap(),
            None,
        );
        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 0);
            assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Pending);
        }

        consumer.cancel();

        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 1);
            assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Ready(None));
        }
    }

    #[test]
    fn stream_on_error() {
        let awoken_count = Arc::new(AtomicUsize::new(0));
        let waker = {
            let awoken_count = awoken_count.clone();
            waker_fn(move || {
                awoken_count.fetch_add(1, Ordering::SeqCst);
            })
        };
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            DefaultExecutor::default().unwrap(),
            None,
        );
        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 0);
            assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Pending);
        }

        consumer.set_error(Error::ChannelsLimitReached);

        {
            let mut next = consumer.next();

            assert_eq!(awoken_count.load(Ordering::SeqCst), 1);
            assert_eq!(
                Pin::new(&mut next).poll(&mut cx),
                Poll::Ready(Some(Err(Error::ChannelsLimitReached)))
            );
        }
    }
}
<|endoftext|>


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

use crate::parser::{Expr, Macro, Node};
use crate::Config;

pub struct Heritage<'a> {
    pub root: &'a Context<'a>,
    pub blocks: BlockAncestry<'a>,
}

impl<'a> Heritage<'a> {
    pub fn new<'n>(
        mut ctx: &'n Context<'n>,
        contexts: &'n HashMap<&'n PathBuf, Context<'n>>,
    ) -> Heritage<'n> {
        let mut blocks: BlockAncestry<'n> = ctx
            .blocks
            .iter()
            .map(|(name, def)| (*name, vec![(ctx, *def)]))
            .collect();

        while let Some(ref path) = ctx.extends {
            ctx = &contexts[&path];
            for (name, def) in &ctx.blocks {
                blocks
                    .entry(name)
                    .or_insert_with(|| vec![])
                    .push((ctx, def));
            }
        }

        Heritage { root: ctx, blocks }
    }
}

type BlockAncestry<'a> = HashMap<&'a str, Vec<(&'a Context<'a>, &'a Node<'a>)>>;

pub struct Context<'a> {
    pub nodes: &'a [Node<'a>],
    pub extends: Option<PathBuf>,
    pub blocks: HashMap<&'a str, &'a Node<'a>>,
    pub macros: HashMap<&'a str, &'a Macro<'a>>,
    pub imports: HashMap<&'a str, PathBuf>,
}

impl<'a> Context<'a> {
    pub fn new<'n>(config: &Config, path: &Path, nodes: &'n [Node<'n>]) -> Context<'n> {
        let mut extends = None;
        let mut blocks = Vec::new();
        let mut macros = HashMap::new();
        let mut imports = HashMap::new();

        for n in nodes {
            match n {
                Node::Extends(Expr::StrLit(extends_path)) => match extends {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => {
                        extends = Some(config.find_template(extends_path, Some(path)));
                    }
                },
                def @ Node::BlockDef(_, _, _, _) => {
                    blocks.push(def);
                }
                Node::Macro(name, m) => {
                    macros.insert(*name, m);
                }
                Node::Import(_, import_path, scope) => {
                    let path = config.find_template(import_path, Some(path));
                    imports.insert(*scope, path);
                }
                _ => {}
            }
        }

        let mut check_nested = 0;
        let mut nested_blocks = Vec::new();
        while check_nested < blocks.len() {
            if let Node::BlockDef(_, _, ref nodes, _) = blocks[check_nested] {
                for n in nodes {
                    if let def @ Node::BlockDef(_, _, _, _) = n {
                        nested_blocks.push(def);
                    }
                }
            } else {
                panic!("non block found in list of blocks");
            }
            blocks.append(&mut nested_blocks);
            check_nested += 1;
        }

        let blocks: HashMap<_, _> = blocks
            .iter()
            .map(|def| {
                if let Node::BlockDef(_, name, _, _) = def {
                    (*name, *def)
                } else {
                    unreachable!()
                }
            })
            .collect();

        Context {
            nodes,
            extends,
            blocks,
            macros,
            imports,
        }
    }
}
<commit_msg>Fix clippy suggestion<commit_after>use std::collections::HashMap;
use std::path::{Path, PathBuf};

use crate::parser::{Expr, Macro, Node};
use crate::Config;

pub struct Heritage<'a> {
    pub root: &'a Context<'a>,
    pub blocks: BlockAncestry<'a>,
}

impl<'a> Heritage<'a> {
    pub fn new<'n>(
        mut ctx: &'n Context<'n>,
        contexts: &'n HashMap<&'n PathBuf, Context<'n>>,
    ) -> Heritage<'n> {
        let mut blocks: BlockAncestry<'n> = ctx
            .blocks
            .iter()
            .map(|(name, def)| (*name, vec![(ctx, *def)]))
            .collect();

        while let Some(ref path) = ctx.extends {
            ctx = &contexts[&path];
            for (name, def) in &ctx.blocks {
                blocks
                    .entry(name)
                    .or_insert_with(Vec::new)
                    .push((ctx, def));
            }
        }

        Heritage { root: ctx, blocks }
    }
}

type BlockAncestry<'a> = HashMap<&'a str, Vec<(&'a Context<'a>, &'a Node<'a>)>>;

pub struct Context<'a> {
    pub nodes: &'a [Node<'a>],
    pub extends: Option<PathBuf>,
    pub blocks: HashMap<&'a str, &'a Node<'a>>,
    pub macros: HashMap<&'a str, &'a Macro<'a>>,
    pub imports: HashMap<&'a str, PathBuf>,
}

impl<'a> Context<'a> {
    pub fn new<'n>(config: &Config, path: &Path, nodes: &'n [Node<'n>]) -> Context<'n> {
        let mut extends = None;
        let mut blocks = Vec::new();
        let mut macros = HashMap::new();
        let mut imports = HashMap::new();

        for n in nodes {
            match n {
                Node::Extends(Expr::StrLit(extends_path)) => match extends {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => {
                        extends = Some(config.find_template(extends_path, Some(path)));
                    }
                },
                def @ Node::BlockDef(_, _, _, _) => {
                    blocks.push(def);
                }
                Node::Macro(name, m) => {
                    macros.insert(*name, m);
                }
                Node::Import(_, import_path, scope) => {
                    let path = config.find_template(import_path, Some(path));
                    imports.insert(*scope, path);
                }
                _ => {}
            }
        }

        let mut check_nested = 0;
        let mut nested_blocks = Vec::new();
        while check_nested < blocks.len() {
            if let Node::BlockDef(_, _, ref nodes, _) = blocks[check_nested] {
                for n in nodes {
                    if let def @ Node::BlockDef(_, _, _, _) = n {
                        nested_blocks.push(def);
                    }
                }
            } else {
                panic!("non block found in list of blocks");
            }
            blocks.append(&mut nested_blocks);
            check_nested += 1;
        }

        let blocks: HashMap<_, _> = blocks
            .iter()
            .map(|def| {
                if let Node::BlockDef(_, name, _, _) = def {
                    (*name, *def)
                } else {
                    unreachable!()
                }
            })
            .collect();

        Context {
            nodes,
            extends,
            blocks,
            macros,
            imports,
        }
    }
}
<|endoftext|>