<commit_before>use std;

use std::collections::HashMap;
use std::fmt::Display;
use std::path::Path;

use ::opcodes::{AddressingMode, OpCode};
use assembler::lexer::{Lexer, LexerError};
use assembler::parser::{Parser, ParserError};
use assembler::token::{LexerToken, ParserToken};

#[derive(Debug, PartialEq)]
pub struct Label(u16);

#[derive(Debug)]
pub struct AssemblerError {
    message: String,
}

impl AssemblerError {
    fn unknown_label<S>(label: S) -> AssemblerError
        where S: Into<String> + std::fmt::Display
    {
        AssemblerError::from(format!("Unknown label: '{}'", label))
    }

    fn relative_offset_too_large<S>(context: S) -> AssemblerError
        where S: Into<String> + Display
    {
        AssemblerError::from(format!("Branch too far: {}", context))
    }
}

impl From<String> for AssemblerError {
    fn from(error: String) -> AssemblerError {
        AssemblerError { message: error }
    }
}

impl From<LexerError> for AssemblerError {
    fn from(error: LexerError) -> AssemblerError {
        AssemblerError { message: error.message }
    }
}

impl From<ParserError> for AssemblerError {
    fn from(error: ParserError) -> AssemblerError {
        AssemblerError { message: error.message }
    }
}

#[derive(Debug)]
pub struct CodeSegment {
    pub address: u16,
    pub code: Vec<u8>,
}

pub struct Assembler {
    symbol_table: HashMap<String, Label>,
}

impl Assembler {
    pub fn new() -> Assembler {
        Assembler { symbol_table: HashMap::new() }
    }

    pub fn assemble_string<S, O>(&mut self,
                                 code: S,
                                 offset: O)
                                 -> Result<Vec<CodeSegment>, AssemblerError>
        where S: Into<String>,
              O: Into<Option<u16>>
    {
        let code = code.into();
        let mut lexer = Lexer::new();
        let tokens = lexer.lex_string(code)?;
        let mut parser = Parser::new();
        let tokens = parser.parse(tokens)?;

        Ok(self.assemble(tokens, offset)?)
    }

    pub fn assemble_file<P, O>(&mut self,
                               path: P,
                               offset: O)
                               -> Result<Vec<CodeSegment>, AssemblerError>
        where P: AsRef<Path>,
              O: Into<Option<u16>>
    {
        let mut lexer = Lexer::new();
        let tokens = lexer.lex_file(path)?;
        let mut parser = Parser::new();
        let tokens = parser.parse(tokens)?;

        Ok(self.assemble(tokens, offset)?)
    }

    fn assemble<O>(&mut self,
                   tokens: Vec<ParserToken>,
                   offset: O)
                   -> Result<Vec<CodeSegment>, AssemblerError>
        where O: Into<Option<u16>>
    {
        let mut addr: u16 = offset.into().unwrap_or(0);

        // First, index the labels so we have addresses for them
        self.index_labels(&tokens, addr);

        // Now assemble the code
        let mut result = Vec::new();
        let mut last_addressing_mode = AddressingMode::Absolute;
        let mut current_segment = CodeSegment {
            address: addr,
            code: Vec::new(),
        };

        for token in tokens {
            // Push an opcode into the output and increment our address
            // offset
            if let ParserToken::OpCode(opcode) = token {
                current_segment.code.push(opcode.code);
                addr += opcode.length as u16;
                last_addressing_mode = opcode.mode;
            } else if let ParserToken::OrgDirective(org_addr) = token {
                if current_segment.code.len() > 0 {
                    result.push(current_segment);
                }
                current_segment = CodeSegment {
                    address: org_addr,
                    code: Vec::new(),
                };
                addr = org_addr;
            } else if let ParserToken::RawByte(byte) = token {
                // Push raw bytes directly into the output
                current_segment.code.push(byte);
            } else if let ParserToken::RawBytes(bytes) = token {
                // Push raw bytes directly into output
                for b in &bytes {
                    current_segment.code.push(*b);
                }
            } else if let ParserToken::LabelArg(ref label) = token {
                // Labels as arguments should be in the symbol table, look
                // it up and calculate the address direction/location
                if let Some(&Label(label_addr)) = self.symbol_table.get(label) {
                    if last_addressing_mode == AddressingMode::Absolute {
                        let low_byte = (label_addr & 0xFF) as u8;
                        let high_byte = ((label_addr >> 8) & 0xFF) as u8;

                        current_segment.code.push(low_byte);
                        current_segment.code.push(high_byte);
                    } else {
                        // Its relative.. lets generate a relative branch
                        if addr > label_addr {
                            let distance = (label_addr as i16 - addr as i16) as i8;
                            if distance < -128 || distance > 127 {
                                return Err(AssemblerError::relative_offset_too_large(format!("Attempted jump to {} at {:04X}", label, addr)));
                            }
                            current_segment.code.push(distance as u8);
                        } else {
                            let distance = label_addr - addr;
                            if distance > 127 {
                                return Err(AssemblerError::relative_offset_too_large(format!("Attempted jump to {} at {:04X}", label, addr)));
                            }
                            current_segment.code.push(distance as u8);
                        }
                    }
                } else {
                    return Err(AssemblerError::unknown_label(label.clone()));
                }
            }
        }

        result.push(current_segment);

        Ok(result)
    }

    /// Stores all labels in the code in a Symbol table for lookup later
    fn index_labels(&mut self, tokens: &[ParserToken], offset: u16) {
        let mut addr: u16 = offset;
        let mut last_addressing_mode = AddressingMode::Absolute;

        for token in tokens {
            if let &ParserToken::Label(ref label) = token {
                // Insert a label with the specified memory address
                // as its offset
                self.symbol_table.insert(label.clone(), Label(addr));
            } else if let &ParserToken::OpCode(opcode) = token {
                // Add the length of this opcode to our
                // address offset
                addr += opcode.length as u16;
                last_addressing_mode = opcode.mode;
            } else if let &ParserToken::OrgDirective(new_addr) = token {
                addr = new_addr
            }
        }
    }
}

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

    #[test]
    fn can_assemble_basic_code() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            LDA $4400
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xAD, 0x00, 0x44], &segments[0].code[..]);
    }

    #[test]
    fn can_jump_to_label_behind() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN LDA $4400
            PHA
            JMP MAIN
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xAD, 0x00, 0x44, 0x48, 0x4C, 0x00, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_jump_to_label_with_colon_behind() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN:
                LDA $4400
                PHA
                JMP MAIN
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xAD, 0x00, 0x44, 0x48, 0x4C, 0x00, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_jump_to_label_ahead() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            JMP MAIN
            PHA
            LDX #15
            MAIN LDA $4400
            RTS
        ",
                             None)
            .unwrap();

        assert_eq!(&[0x4C, 0x06, 0x00, 0x48, 0xA2, 0x0F, 0xAD, 0x00, 0x44, 0x60],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_use_variables() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN_ADDRESS = $0000
            MAIN:
            LDX #15
            JMP MAIN_ADDRESS
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA2, 0x0F, 0x4C, 0x00, 0x00], &segments[0].code[..]);
    }

    #[test]
    fn can_use_variables_assigned_to_variables() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN_ADDRESS = $0000
            MAIN_ADDRESS_INDIRECT_ONE = MAIN_ADDRESS
            MAIN_ADDRESS_INDIRECT_TWO = MAIN_ADDRESS_INDIRECT_ONE
            MAIN:
            LDX #15
            JMP MAIN_ADDRESS_INDIRECT_TWO
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA2, 0x0F, 0x4C, 0x00, 0x00], &segments[0].code[..]);
    }

    #[test]
    fn can_assemble_clearmem_implementation() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            CLRMEM  LDA #$00
                    TAY
            CLRM1   STA ($FF),Y
                    INY
                    DEX
                    BNE CLRM1
                    RTS
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA9, 0x00, 0xA8, 0x91, 0xFF, 0xC8, 0xCA, 0xD0, 0xFA, 0x60],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_assemble_clearmem_implementation_that_jumps_forward_and_is_lowercase() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            jmp     clrmem
            lda     #$00
            beq     clrm1
            nop
            nop
            clrm1   sta ($ff),y
                    iny
                    dex
                    bne clrm1
                    rts
            clrmem  lda #$00
                    tay
            jmp     clrm1
        ",
                             None)
            .unwrap();

        assert_eq!(&[0x4C, 0x10, 0x00, 0xA9, 0x00, 0xF0, 0x02, 0xEA, 0xEA, 0x91, 0xFF, 0xC8,
                     0xCA, 0xD0, 0xFA, 0x60, 0xA9, 0x00, 0xA8, 0x4C, 0x09, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_assemble_clearmem_implementation_that_jumps_forward() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            JMP     CLRMEM
            LDA     #$00
            BEQ     CLRM1
            NOP
            NOP
            BRK
            CLRM1   STA ($FF),Y
                    INY
                    DEX
                    BNE CLRM1
                    RTS
            CLRMEM  LDA #$00
                    TAY
            JMP     CLRM1
        ",
                             None)
            .unwrap();

        assert_eq!(&[0x4C, 0x11, 0x00, 0xA9, 0x00, 0xF0, 0x03, 0xEA, 0xEA, 0x00, 0x91, 0xFF,
                     0xC8, 0xCA, 0xD0, 0xFA, 0x60, 0xA9, 0x00, 0xA8, 0x4C, 0x0A, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_use_variables_for_indirect_addressing() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN_ADDRESS = $0000
            MAIN:
            LDX #15
            LDA (MAIN_ADDRESS),Y
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA2, 0x0F, 0xB1, 0x00, 0x00], &segments[0].code[..]);
    }

    #[test]
    fn can_assign_code_segments_to_different_memory_addresses() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000
            LDA #$FF
            STA $2000

            .ORG $100
            LDA #$AA
            STA $2001
        ",
                             None)
            .unwrap();

        assert_eq!(0xC000, segments[0].address);
        assert_eq!(0x0100, segments[1].address);
    }

    #[test]
    fn can_jump_between_code_segments() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000
            JMP CALLBACK

            .ORG $2000
            LDA #$AA
            STA $2001

            CALLBACK
            LDX #$0A
        ",
                             None)
            .unwrap();

        assert_eq!(0xC000, segments[0].address);
        assert_eq!(0x2000, segments[1].address);

        assert_eq!(0x05, segments[0].code[0x01]);
        assert_eq!(0x20, segments[0].code[0x02]);
    }

    #[test]
    fn can_dump_raw_bytes() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000

            .BYTE #$40, #10, #$0A
        ",
                             None)
            .unwrap();

        assert_eq!(&[64, 10, 10], &segments[0].code[..]);
    }
}<commit_msg>Add more Assembler tests for the .BYTE directive<commit_after>use std;

use std::collections::HashMap;
use std::fmt::Display;
use std::path::Path;

use ::opcodes::{AddressingMode, OpCode};
use assembler::lexer::{Lexer, LexerError};
use assembler::parser::{Parser, ParserError};
use assembler::token::{LexerToken, ParserToken};

#[derive(Debug, PartialEq)]
pub struct Label(u16);

#[derive(Debug)]
pub struct AssemblerError {
    message: String,
}

impl AssemblerError {
    fn unknown_label<S>(label: S) -> AssemblerError
        where S: Into<String> + std::fmt::Display
    {
        AssemblerError::from(format!("Unknown label: '{}'", label))
    }

    fn relative_offset_too_large<S>(context: S) -> AssemblerError
        where S: Into<String> + Display
    {
        AssemblerError::from(format!("Branch too far: {}", context))
    }
}

impl From<String> for AssemblerError {
    fn from(error: String) -> AssemblerError {
        AssemblerError { message: error }
    }
}

impl From<LexerError> for AssemblerError {
    fn from(error: LexerError) -> AssemblerError {
        AssemblerError { message: error.message }
    }
}

impl From<ParserError> for AssemblerError {
    fn from(error: ParserError) -> AssemblerError {
        AssemblerError { message: error.message }
    }
}

#[derive(Debug)]
pub struct CodeSegment {
    pub address: u16,
    pub code: Vec<u8>,
}

pub struct Assembler {
    symbol_table: HashMap<String, Label>,
}

impl Assembler {
    pub fn new() -> Assembler {
        Assembler { symbol_table: HashMap::new() }
    }

    pub fn assemble_string<S, O>(&mut self,
                                 code: S,
                                 offset: O)
                                 -> Result<Vec<CodeSegment>, AssemblerError>
        where S: Into<String>,
              O: Into<Option<u16>>
    {
        let code = code.into();
        let mut lexer = Lexer::new();
        let tokens = lexer.lex_string(code)?;
        let mut parser = Parser::new();
        let tokens = parser.parse(tokens)?;

        Ok(self.assemble(tokens, offset)?)
    }

    pub fn assemble_file<P, O>(&mut self,
                               path: P,
                               offset: O)
                               -> Result<Vec<CodeSegment>, AssemblerError>
        where P: AsRef<Path>,
              O: Into<Option<u16>>
    {
        let mut lexer = Lexer::new();
        let tokens = lexer.lex_file(path)?;
        let mut parser = Parser::new();
        let tokens = parser.parse(tokens)?;

        Ok(self.assemble(tokens, offset)?)
    }

    fn assemble<O>(&mut self,
                   tokens: Vec<ParserToken>,
                   offset: O)
                   -> Result<Vec<CodeSegment>, AssemblerError>
        where O: Into<Option<u16>>
    {
        let mut addr: u16 = offset.into().unwrap_or(0);

        // First, index the labels so we have addresses for them
        self.index_labels(&tokens, addr);

        // Now assemble the code
        let mut result = Vec::new();
        let mut last_addressing_mode = AddressingMode::Absolute;
        let mut current_segment = CodeSegment {
            address: addr,
            code: Vec::new(),
        };

        for token in tokens {
            // Push an opcode into the output and increment our address
            // offset
            if let ParserToken::OpCode(opcode) = token {
                current_segment.code.push(opcode.code);
                addr += opcode.length as u16;
                last_addressing_mode = opcode.mode;
            } else if let ParserToken::OrgDirective(org_addr) = token {
                if current_segment.code.len() > 0 {
                    result.push(current_segment);
                }
                current_segment = CodeSegment {
                    address: org_addr,
                    code: Vec::new(),
                };
                addr = org_addr;
            } else if let ParserToken::RawByte(byte) = token {
                // Push raw bytes directly into the output
                current_segment.code.push(byte);
            } else if let ParserToken::RawBytes(bytes) = token {
                // Push raw bytes directly into output
                for b in &bytes {
                    current_segment.code.push(*b);
                }
            } else if let ParserToken::LabelArg(ref label) = token {
                // Labels as arguments should be in the symbol table, look
                // it up and calculate the address direction/location
                if let Some(&Label(label_addr)) = self.symbol_table.get(label) {
                    if last_addressing_mode == AddressingMode::Absolute {
                        let low_byte = (label_addr & 0xFF) as u8;
                        let high_byte = ((label_addr >> 8) & 0xFF) as u8;

                        current_segment.code.push(low_byte);
                        current_segment.code.push(high_byte);
                    } else {
                        // Its relative.. lets generate a relative branch
                        if addr > label_addr {
                            let distance = (label_addr as i16 - addr as i16) as i8;
                            if distance < -128 || distance > 127 {
                                return Err(AssemblerError::relative_offset_too_large(format!("Attempted jump to {} at {:04X}", label, addr)));
                            }
                            current_segment.code.push(distance as u8);
                        } else {
                            let distance = label_addr - addr;
                            if distance > 127 {
                                return Err(AssemblerError::relative_offset_too_large(format!("Attempted jump to {} at {:04X}", label, addr)));
                            }
                            current_segment.code.push(distance as u8);
                        }
                    }
                } else {
                    return Err(AssemblerError::unknown_label(label.clone()));
                }
            }
        }

        result.push(current_segment);

        Ok(result)
    }

    /// Stores all labels in the code in a Symbol table for lookup later
    fn index_labels(&mut self, tokens: &[ParserToken], offset: u16) {
        let mut addr: u16 = offset;
        let mut last_addressing_mode = AddressingMode::Absolute;

        for token in tokens {
            if let &ParserToken::Label(ref label) = token {
                // Insert a label with the specified memory address
                // as its offset
                self.symbol_table.insert(label.clone(), Label(addr));
            } else if let &ParserToken::OpCode(opcode) = token {
                // Add the length of this opcode to our
                // address offset
                addr += opcode.length as u16;
                last_addressing_mode = opcode.mode;
            } else if let &ParserToken::OrgDirective(new_addr) = token {
                addr = new_addr
            }
        }
    }
}

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

    #[test]
    fn can_assemble_basic_code() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            LDA $4400
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xAD, 0x00, 0x44], &segments[0].code[..]);
    }

    #[test]
    fn can_jump_to_label_behind() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN LDA $4400
            PHA
            JMP MAIN
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xAD, 0x00, 0x44, 0x48, 0x4C, 0x00, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_jump_to_label_with_colon_behind() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN:
                LDA $4400
                PHA
                JMP MAIN
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xAD, 0x00, 0x44, 0x48, 0x4C, 0x00, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_jump_to_label_ahead() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            JMP MAIN
            PHA
            LDX #15
            MAIN LDA $4400
            RTS
        ",
                             None)
            .unwrap();

        assert_eq!(&[0x4C, 0x06, 0x00, 0x48, 0xA2, 0x0F, 0xAD, 0x00, 0x44, 0x60],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_use_variables() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN_ADDRESS = $0000
            MAIN:
            LDX #15
            JMP MAIN_ADDRESS
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA2, 0x0F, 0x4C, 0x00, 0x00], &segments[0].code[..]);
    }

    #[test]
    fn can_use_variables_assigned_to_variables() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN_ADDRESS = $0000
            MAIN_ADDRESS_INDIRECT_ONE = MAIN_ADDRESS
            MAIN_ADDRESS_INDIRECT_TWO = MAIN_ADDRESS_INDIRECT_ONE
            MAIN:
            LDX #15
            JMP MAIN_ADDRESS_INDIRECT_TWO
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA2, 0x0F, 0x4C, 0x00, 0x00], &segments[0].code[..]);
    }

    #[test]
    fn can_assemble_clearmem_implementation() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            CLRMEM  LDA #$00
                    TAY
            CLRM1   STA ($FF),Y
                    INY
                    DEX
                    BNE CLRM1
                    RTS
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA9, 0x00, 0xA8, 0x91, 0xFF, 0xC8, 0xCA, 0xD0, 0xFA, 0x60],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_assemble_clearmem_implementation_that_jumps_forward_and_is_lowercase() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            jmp     clrmem
            lda     #$00
            beq     clrm1
            nop
            nop
            clrm1   sta ($ff),y
                    iny
                    dex
                    bne clrm1
                    rts
            clrmem  lda #$00
                    tay
            jmp     clrm1
        ",
                             None)
            .unwrap();

        assert_eq!(&[0x4C, 0x10, 0x00, 0xA9, 0x00, 0xF0, 0x02, 0xEA, 0xEA, 0x91, 0xFF, 0xC8,
                     0xCA, 0xD0, 0xFA, 0x60, 0xA9, 0x00, 0xA8, 0x4C, 0x09, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_assemble_clearmem_implementation_that_jumps_forward() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            JMP     CLRMEM
            LDA     #$00
            BEQ     CLRM1
            NOP
            NOP
            BRK
            CLRM1   STA ($FF),Y
                    INY
                    DEX
                    BNE CLRM1
                    RTS
            CLRMEM  LDA #$00
                    TAY
            JMP     CLRM1
        ",
                             None)
            .unwrap();

        assert_eq!(&[0x4C, 0x11, 0x00, 0xA9, 0x00, 0xF0, 0x03, 0xEA, 0xEA, 0x00, 0x91, 0xFF,
                     0xC8, 0xCA, 0xD0, 0xFA, 0x60, 0xA9, 0x00, 0xA8, 0x4C, 0x0A, 0x00],
                   &segments[0].code[..]);
    }

    #[test]
    fn can_use_variables_for_indirect_addressing() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            MAIN_ADDRESS = $0000
            MAIN:
            LDX #15
            LDA (MAIN_ADDRESS),Y
        ",
                             None)
            .unwrap();

        assert_eq!(&[0xA2, 0x0F, 0xB1, 0x00, 0x00], &segments[0].code[..]);
    }

    #[test]
    fn can_assign_code_segments_to_different_memory_addresses() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000
            LDA #$FF
            STA $2000

            .ORG $100
            LDA #$AA
            STA $2001
        ",
                             None)
            .unwrap();

        assert_eq!(0xC000, segments[0].address);
        assert_eq!(0x0100, segments[1].address);
    }

    #[test]
    fn can_jump_between_code_segments() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000
            JMP CALLBACK

            .ORG $2000
            LDA #$AA
            STA $2001

            CALLBACK
            LDX #$0A
        ",
                             None)
            .unwrap();

        assert_eq!(0xC000, segments[0].address);
        assert_eq!(0x2000, segments[1].address);

        assert_eq!(0x05, segments[0].code[0x01]);
        assert_eq!(0x20, segments[0].code[0x02]);
    }

    #[test]
    fn can_dump_raw_bytes() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000

            .BYTE #$40, #10, #$0A
        ",
                             None)
            .unwrap();

        assert_eq!(&[64, 10, 10], &segments[0].code[..]);
    }

    #[test]
    fn can_dump_single_raw_byte() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000

            .BYTE #$FF
        ",
                             None)
            .unwrap();

        assert_eq!(&[255], &segments[0].code[..]);
    }

    #[test]
    fn can_dump_bytes_with_other_code() {
        let mut assembler = Assembler::new();
        let segments = assembler.assemble_string("
            .ORG $C000
            JMP CALLBACK
            .BYTE #$0A

            .ORG $2000
            LDA #$AA
            STA $2001
            .BYTE #$FE, #$CB

            CALLBACK
            LDX #$0A
        ",
                             None)
            .unwrap();

        assert_eq!(0x0A, segments[0].code[3]);
        assert_eq!(&[0xFE, 0xCB], &segments[1].code[5..7]);

        assert_eq!(0xC000, segments[0].address);
        assert_eq!(0x2000, segments[1].address);

        assert_eq!(0x05, segments[0].code[0x01]);
        assert_eq!(0x20, segments[0].code[0x02]);
    }
}<|endoftext|>


<commit_before>use webapi::event_target::{IEventTarget, EventTarget};
use webcore::unsafe_typed_array::UnsafeTypedArray;
use webcore::value::{
    Reference,
    Value,
};
use webcore::try_from::TryInto;
use private::TODO;

/// Use XmlHttpRequest (XHR) objects to interact with servers.
/// You can retrieve data from a URL without having to do a full page refresh.
/// This enables a Web page to update just part of a page without disrupting
/// what the user is doing. XmlHttpRequest is used heavily in Ajax programming.
///
/// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XmlHttpRequest)
// https://xhr.spec.whatwg.org/#xmlhttprequest
#[derive(Clone, Debug, PartialEq, Eq, ReferenceType)]
#[reference(instance_of = "XMLHttpRequest")]
#[reference(subclass_of(EventTarget))]
pub struct XmlHttpRequest( Reference );

/// An enum indicating the state of the `XmlHttpRequest`.
///
/// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/readyState)
// https://xhr.spec.whatwg.org/#dom-xmlhttprequest-readystate
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum XhrReadyState {
    /// Client has been created. [open()](struct.XmlHttpRequest.html#method.open) not called yet.
    Unsent,
    /// [open()](struct.XmlHttpRequest.html#method.open) has been called.
    Opened,
    /// [send()](struct.XmlHttpRequest.html#method.send) has been called, and headers and [status()](struct.XmlHttpRequest.html#method.status) are available.
    HeadersReceived,
    /// Downloading; [reponse_text()](struct.XmlHttpRequest.html#method.reponse_text) holds partial data.
    Loading,
    /// The operation is complete.
    Done,
}

/// An enum describing the type of the response to `XmlHttpRequest`
///
/// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseType)
// https://xhr.spec.whatwg.org/#dom-xmlhttprequest-responsetype
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum XhrResponseType {
    /// A JavaScript ArrayBuffer containing binary data
    ArrayBuffer,
    /// A Blob object container the binary data
    Blob,
    /// An HTML Document or XML XMLDocument
    Document,
    /// A JavaScript object parsed from JSON
    Json,
    /// Text in a String object
    Text
}

impl IEventTarget for XmlHttpRequest {}

impl XmlHttpRequest {
    /// Creates new `XmlHttpRequest`.
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest
    pub fn new() -> XmlHttpRequest {
        js!( return new XMLHttpRequest(); ).try_into().unwrap()
    }

    /// Returns the current state of the request as a [XhrReadyState](enum.XhrReadyState.html).
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/readyState)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-readystate
    pub fn ready_state(&self) -> XhrReadyState {
        use self::XhrReadyState::*;
        let state: u16 = js!( return @{self}.readyState; ).try_into().unwrap();
        match state {
            0 => Unsent,
            1 => Opened,
            2 => HeadersReceived,
            3 => Loading,
            4 => Done,
            _ => unreachable!( "Unexpected value of XMLHttpRequest::readyState: {}", state )
        }
    }

    /// Returns the type of the request as a [XhrResponseType](enum.XhrResponseType.html)
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseType)
    // https://xhr.spec.whatwg.org/#dom-xmlhttprequest-responsetype
    pub fn response_type(&self) -> XhrResponseType {
        use self::XhrResponseType::*;
        let repsonse_type: String = js! ( return @{self}.responseType; ).try_into().unwrap();
        match repsonse_type.as_str() {
            "arraybuffer" => ArrayBuffer,
            "blob" => Blob,
            "document" => Document,
            "json" => Json,
            "text" | "" => Text,
            x => unreachable!( "Unexpected value of XMLHttpRequest::responseType:: {}", x)
        }
    }

    /// Returns a string that contains the response to the request as text, or None
    /// if the request was unsuccessful or has not yet been sent.
    ///
    ///[(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseText)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-responsetext
    pub fn response_text(&self) -> Result< Option< String >, TODO > {
        let response = js!(return @{self}.responseText;);
        match response {
            Value::Null => Ok( None ),
            Value::String( resp ) => Ok( Some( resp ) ),
            _ => unreachable!(),
        }
    }

    /// Returns the object representing the response
    ///
    ///[(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/response)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-response
    pub fn raw_response(&self) -> Value {
        js!(return @{self}.response;)
    }

    /// Returns an unsigned short with the status of the response of the request.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/status)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-status
    pub fn status(&self) -> u16 {
        js!(return @{self}.status;).try_into().unwrap()
    }

    /// Open connection with given method (ie GET or POST), and the url to hit.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/open)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-open
    pub fn open(&self, method: &str, url: &str) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.open(@{method}, @{url}, true);
        };

        Ok(())
    }

    /// Returns the string containing the text of the specified header. If there
    /// are multiple response headers with the same name, then their values are
    /// returned as a single concatenated string.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/getResponseHeader)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-getresponseheader
    pub fn get_response_header(&self, header: &str) -> Option<String> {
        let header = js!( return @{self}.getResponseHeader(@{header}); );
        match header {
            Value::Null => None,
            Value::String(text) => Some(text),
            _ => unreachable!(),
        }
    }

    /// Sets the value of an HTTP request header. Must be called after `open()`,
    /// but before `send()`. If this method is called several times with the same
    /// header, the values are merged into one single request header.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/setRequestHeader)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-setrequestheader
    pub fn set_request_header(&self, header: &str, value: &str) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.setRequestHeader(@{header}, @{value});
        };

        Ok(())
    }

    /// Send request on an open connection with no data
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/send)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-send
    pub fn send(&self) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.send();
        };

        Ok(())
    }

    /// Send request on an open connection with string body
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/send)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-send
    pub fn send_with_string(&self, body: &str) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.send(@{body});
        };

        Ok(())
    }

    /// Send request on an open connection with a byte array body
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/send)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-send
    pub fn send_with_bytes(&self, body: &[u8]) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.send(@{UnsafeTypedArray(body)});
        };

        Ok(())
    }

    /// Aborts the request if it has already been sent.
    /// When a request is aborted, its [ready_state](struct.XmlHttpRequest.html#method.ready_state) is changed to [Done](enum.XhrReadyState.html#variant.Done)
    /// and the [status](struct.XmlHttpRequest.html#method.status) code is set to
    /// [Unsent](enum.XhrReadyState.html#variant.Unsent).
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/abort)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-abort
    pub fn abort(&self) {
        js! { @(no_return)
            @{self}.abort();
        };
    }
}
<commit_msg>Add `XmlHttpRequest::set_response_type` (#311)<commit_after>use webapi::event_target::{IEventTarget, EventTarget};
use webapi::dom_exception::{InvalidAccessError, InvalidStateError};
use webcore::unsafe_typed_array::UnsafeTypedArray;
use webcore::value::{
    Reference,
    Value,
};
use webcore::try_from::{TryFrom, TryInto};
use private::TODO;

/// Use XmlHttpRequest (XHR) objects to interact with servers.
/// You can retrieve data from a URL without having to do a full page refresh.
/// This enables a Web page to update just part of a page without disrupting
/// what the user is doing. XmlHttpRequest is used heavily in Ajax programming.
///
/// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XmlHttpRequest)
// https://xhr.spec.whatwg.org/#xmlhttprequest
#[derive(Clone, Debug, PartialEq, Eq, ReferenceType)]
#[reference(instance_of = "XMLHttpRequest")]
#[reference(subclass_of(EventTarget))]
pub struct XmlHttpRequest( Reference );

/// An enum indicating the state of the `XmlHttpRequest`.
///
/// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/readyState)
// https://xhr.spec.whatwg.org/#dom-xmlhttprequest-readystate
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum XhrReadyState {
    /// Client has been created. [open()](struct.XmlHttpRequest.html#method.open) not called yet.
    Unsent,
    /// [open()](struct.XmlHttpRequest.html#method.open) has been called.
    Opened,
    /// [send()](struct.XmlHttpRequest.html#method.send) has been called, and headers and [status()](struct.XmlHttpRequest.html#method.status) are available.
    HeadersReceived,
    /// Downloading; [reponse_text()](struct.XmlHttpRequest.html#method.reponse_text) holds partial data.
    Loading,
    /// The operation is complete.
    Done,
}

/// An enum describing the type of the response to `XmlHttpRequest`
///
/// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseType)
// https://xhr.spec.whatwg.org/#dom-xmlhttprequest-responsetype
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum XhrResponseType {
    /// A JavaScript ArrayBuffer containing binary data
    ArrayBuffer,
    /// A Blob object container the binary data
    Blob,
    /// An HTML Document or XML XMLDocument
    Document,
    /// A JavaScript object parsed from JSON
    Json,
    /// Text in a String object
    Text
}

impl IEventTarget for XmlHttpRequest {}

error_enum_boilerplate! {
    SetResponseTypeError,
    InvalidStateError, InvalidAccessError
}

impl XmlHttpRequest {
    /// Creates new `XmlHttpRequest`.
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest
    pub fn new() -> XmlHttpRequest {
        js!( return new XMLHttpRequest(); ).try_into().unwrap()
    }

    /// Returns the current state of the request as a [XhrReadyState](enum.XhrReadyState.html).
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/readyState)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-readystate
    pub fn ready_state(&self) -> XhrReadyState {
        use self::XhrReadyState::*;
        let state: u16 = js!( return @{self}.readyState; ).try_into().unwrap();
        match state {
            0 => Unsent,
            1 => Opened,
            2 => HeadersReceived,
            3 => Loading,
            4 => Done,
            _ => unreachable!( "Unexpected value of XMLHttpRequest::readyState: {}", state )
        }
    }

    /// Returns the type of the request as a [XhrResponseType](enum.XhrResponseType.html)
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseType)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-responsetype
    pub fn response_type(&self) -> XhrResponseType {
        use self::XhrResponseType::*;
        let repsonse_type: String = js! ( return @{self}.responseType; ).try_into().unwrap();
        match repsonse_type.as_str() {
            "arraybuffer" => ArrayBuffer,
            "blob" => Blob,
            "document" => Document,
            "json" => Json,
            "text" | "" => Text,
            x => unreachable!( "Unexpected value of XMLHttpRequest::responseType:: {}", x)
        }
    }

    /// Set the type that the XmlHttpRequest should return
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseType)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-responsetype
    pub fn set_response_type(&self, kind: XhrResponseType) -> Result<(), SetResponseTypeError> {
        use self::XhrResponseType::*;
        let response_type = match kind {
            ArrayBuffer => "arraybuffer",
            Blob => "blob",
            Document => "document",
            Json => "json",
            Text => "text"
        };

        js_try!(
            @{self}.responseType = @{response_type};
        ).unwrap()
    }

    /// Returns a string that contains the response to the request as text, or None
    /// if the request was unsuccessful or has not yet been sent.
    ///
    ///[(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/responseText)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-responsetext
    pub fn response_text(&self) -> Result< Option< String >, TODO > {
        let response = js!(return @{self}.responseText;);
        match response {
            Value::Null => Ok( None ),
            Value::String( resp ) => Ok( Some( resp ) ),
            _ => unreachable!(),
        }
    }

    /// Returns the object representing the response
    ///
    ///[(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/response)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-response
    pub fn raw_response(&self) -> Value {
        js!(return @{self}.response;)
    }

    /// Returns an unsigned short with the status of the response of the request.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/status)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-status
    pub fn status(&self) -> u16 {
        js!(return @{self}.status;).try_into().unwrap()
    }

    /// Open connection with given method (ie GET or POST), and the url to hit.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/open)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-open
    pub fn open(&self, method: &str, url: &str) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.open(@{method}, @{url}, true);
        };

        Ok(())
    }

    /// Returns the string containing the text of the specified header. If there
    /// are multiple response headers with the same name, then their values are
    /// returned as a single concatenated string.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/getResponseHeader)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-getresponseheader
    pub fn get_response_header(&self, header: &str) -> Option<String> {
        let header = js!( return @{self}.getResponseHeader(@{header}); );
        match header {
            Value::Null => None,
            Value::String(text) => Some(text),
            _ => unreachable!(),
        }
    }

    /// Sets the value of an HTTP request header. Must be called after `open()`,
    /// but before `send()`. If this method is called several times with the same
    /// header, the values are merged into one single request header.
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/setRequestHeader)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-setrequestheader
    pub fn set_request_header(&self, header: &str, value: &str) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.setRequestHeader(@{header}, @{value});
        };

        Ok(())
    }

    /// Send request on an open connection with no data
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/send)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-send
    pub fn send(&self) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.send();
        };

        Ok(())
    }

    /// Send request on an open connection with string body
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/send)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-send
    pub fn send_with_string(&self, body: &str) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.send(@{body});
        };

        Ok(())
    }

    /// Send request on an open connection with a byte array body
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/send)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-send
    pub fn send_with_bytes(&self, body: &[u8]) -> Result< (), TODO > {
        js! { @(no_return)
            @{self}.send(@{UnsafeTypedArray(body)});
        };

        Ok(())
    }

    /// Aborts the request if it has already been sent.
    /// When a request is aborted, its [ready_state](struct.XmlHttpRequest.html#method.ready_state) is changed to [Done](enum.XhrReadyState.html#variant.Done)
    /// and the [status](struct.XmlHttpRequest.html#method.status) code is set to
    /// [Unsent](enum.XhrReadyState.html#variant.Unsent).
    ///
    /// [(JavaScript docs)](https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/abort)
    // https://xhr.spec.whatwg.org/#ref-for-dom-xmlhttprequest-abort
    pub fn abort(&self) {
        js! { @(no_return)
            @{self}.abort();
        };
    }
}
<|endoftext|>


<commit_before>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to the `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freakingly cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! fsts, one can build an product fst which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an fst
//! for each library (which is assumed to never change) and an fst for each rust
//! file in the current workspace, and run a query against the union of all
//! those fsts.
use std::{
    cmp::Ordering,
    hash::{Hash, Hasher},
    sync::Arc,
    mem,
};

use fst::{self, Streamer};
use ra_syntax::{
    SyntaxNode, SyntaxNodePtr, SourceFile, SmolStr, TreeArc, AstNode,
    algo::{visit::{visitor, Visitor}, find_covering_node},
    SyntaxKind::{self, *},
    ast::{self, NameOwner},
};
use ra_db::{
    SourceRootId, FilesDatabase,
    salsa::{self, ParallelDatabase},
};
use rayon::prelude::*;

use crate::{
    FileId, Query,
    db::RootDatabase,
};

#[salsa::query_group]
pub(crate) trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let source_file = db.source_file(file_id);
    let mut symbols = source_file
        .syntax()
        .descendants()
        .filter_map(to_symbol)
        .map(move |(name, ptr)| FileSymbol { name, ptr, file_id })
        .collect::<Vec<_>>();

    for (name, text_range) in hir::source_binder::macro_symbols(db, file_id) {
        let node = find_covering_node(source_file.syntax(), text_range);
        let ptr = SyntaxNodePtr::new(node);
        symbols.push(FileSymbol { file_id, name, ptr })
    }

    Arc::new(SymbolIndex::new(symbols))
}

pub(crate) fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        db.library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect()
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.files.values().map(|&it| it))
        }

        let snap = Snap(db.snapshot());
        files
            .par_iter()
            .map_with(snap, |db, &file_id| db.0.file_symbols(file_id))
            .collect()
    };
    query.search(&buf)
}

#[derive(Default, Debug)]
pub(crate) struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map,
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(s1: &FileSymbol, s2: &FileSymbol) -> Ordering {
            unicase::Ascii::new(s1.name.as_str()).cmp(&unicase::Ascii::new(s2.name.as_str()))
        }
        symbols.par_sort_by(cmp);
        symbols.dedup_by(|s1, s2| cmp(s1, s2) == Ordering::Equal);
        let names = symbols.iter().map(|it| it.name.as_str().to_lowercase());
        let map = fst::Map::from_iter(names.into_iter().zip(0u64..)).unwrap();
        SymbolIndex { symbols, map }
    }

    pub(crate) fn len(&self) -> usize {
        self.symbols.len()
    }

    pub(crate) fn memory_size(&self) -> usize {
        self.map.as_fst().size() * self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, TreeArc<SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| {
                file.syntax()
                    .descendants()
                    .filter_map(to_symbol)
                    .map(move |(name, ptr)| FileSymbol { name, ptr, file_id })
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let file_symbols = &indices[indexed_value.index];
                let idx = indexed_value.value as usize;

                let symbol = &file_symbols.symbols[idx];
                if self.only_types && !is_type(symbol.ptr.kind()) {
                    continue;
                }
                if self.exact && symbol.name != self.query {
                    continue;
                }
                res.push(symbol.clone());
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct FileSymbol {
    pub(crate) file_id: FileId,
    pub(crate) name: SmolStr,
    pub(crate) ptr: SyntaxNodePtr,
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr)> {
    fn decl<N: NameOwner>(node: &N) -> Option<(SmolStr, SyntaxNodePtr)> {
        let name = node.name()?.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());
        Some((name, ptr))
    }
    visitor()
        .visit(decl::<ast::FnDef>)
        .visit(decl::<ast::StructDef>)
        .visit(decl::<ast::EnumDef>)
        .visit(decl::<ast::TraitDef>)
        .visit(decl::<ast::Module>)
        .visit(decl::<ast::TypeDef>)
        .visit(decl::<ast::ConstDef>)
        .visit(decl::<ast::StaticDef>)
        .accept(node)?
}
<commit_msg>print one size for symbols<commit_after>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to the `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freakingly cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! fsts, one can build an product fst which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an fst
//! for each library (which is assumed to never change) and an fst for each rust
//! file in the current workspace, and run a query against the union of all
//! those fsts.
use std::{
    cmp::Ordering,
    hash::{Hash, Hasher},
    sync::Arc,
    mem,
};

use fst::{self, Streamer};
use ra_syntax::{
    SyntaxNode, SyntaxNodePtr, SourceFile, SmolStr, TreeArc, AstNode,
    algo::{visit::{visitor, Visitor}, find_covering_node},
    SyntaxKind::{self, *},
    ast::{self, NameOwner},
};
use ra_db::{
    SourceRootId, FilesDatabase,
    salsa::{self, ParallelDatabase},
};
use rayon::prelude::*;

use crate::{
    FileId, Query,
    db::RootDatabase,
};

#[salsa::query_group]
pub(crate) trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let source_file = db.source_file(file_id);
    let mut symbols = source_file
        .syntax()
        .descendants()
        .filter_map(to_symbol)
        .map(move |(name, ptr)| FileSymbol { name, ptr, file_id })
        .collect::<Vec<_>>();

    for (name, text_range) in hir::source_binder::macro_symbols(db, file_id) {
        let node = find_covering_node(source_file.syntax(), text_range);
        let ptr = SyntaxNodePtr::new(node);
        symbols.push(FileSymbol { file_id, name, ptr })
    }

    Arc::new(SymbolIndex::new(symbols))
}

pub(crate) fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        db.library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect()
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.files.values().map(|&it| it))
        }

        let snap = Snap(db.snapshot());
        files
            .par_iter()
            .map_with(snap, |db, &file_id| db.0.file_symbols(file_id))
            .collect()
    };
    query.search(&buf)
}

#[derive(Default, Debug)]
pub(crate) struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map,
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(s1: &FileSymbol, s2: &FileSymbol) -> Ordering {
            unicase::Ascii::new(s1.name.as_str()).cmp(&unicase::Ascii::new(s2.name.as_str()))
        }
        symbols.par_sort_by(cmp);
        symbols.dedup_by(|s1, s2| cmp(s1, s2) == Ordering::Equal);
        let names = symbols.iter().map(|it| it.name.as_str().to_lowercase());
        let map = fst::Map::from_iter(names.into_iter().zip(0u64..)).unwrap();
        SymbolIndex { symbols, map }
    }

    pub(crate) fn len(&self) -> usize {
        self.symbols.len()
    }

    pub(crate) fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, TreeArc<SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| {
                file.syntax()
                    .descendants()
                    .filter_map(to_symbol)
                    .map(move |(name, ptr)| FileSymbol { name, ptr, file_id })
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let file_symbols = &indices[indexed_value.index];
                let idx = indexed_value.value as usize;

                let symbol = &file_symbols.symbols[idx];
                if self.only_types && !is_type(symbol.ptr.kind()) {
                    continue;
                }
                if self.exact && symbol.name != self.query {
                    continue;
                }
                res.push(symbol.clone());
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct FileSymbol {
    pub(crate) file_id: FileId,
    pub(crate) name: SmolStr,
    pub(crate) ptr: SyntaxNodePtr,
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr)> {
    fn decl<N: NameOwner>(node: &N) -> Option<(SmolStr, SyntaxNodePtr)> {
        let name = node.name()?.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());
        Some((name, ptr))
    }
    visitor()
        .visit(decl::<ast::FnDef>)
        .visit(decl::<ast::StructDef>)
        .visit(decl::<ast::EnumDef>)
        .visit(decl::<ast::TraitDef>)
        .visit(decl::<ast::Module>)
        .visit(decl::<ast::TypeDef>)
        .visit(decl::<ast::ConstDef>)
        .visit(decl::<ast::StaticDef>)
        .accept(node)?
}
<|endoftext|>


<commit_before>use std::{
    cmp::Ordering,
    convert::TryFrom,
    fmt,
    hash::{Hash, Hasher},
    marker::PhantomData,
    str::FromStr,
};

use crate::{crypto_algorithms::SigningKeyAlgorithm, DeviceId, KeyName};

/// A key algorithm and key name delimited by a colon.
#[repr(transparent)]
pub struct KeyId<A, K: ?Sized>(PhantomData<(A, K)>, str);

impl<A, K: ?Sized> KeyId<A, K> {
    /// Creates a new `KeyId` from an algorithm and key name.
    pub fn from_parts(algorithm: A, key_name: &K) -> Box<Self>
    where
        A: AsRef<str>,
        K: AsRef<str>,
    {
        let algorithm = algorithm.as_ref();
        let key_name = key_name.as_ref();

        let mut res = String::with_capacity(algorithm.len() + 1 + key_name.len());
        res.push_str(algorithm);
        res.push(':');
        res.push_str(key_name);

        Self::from_owned(res.into())
    }

    /// Returns key algorithm of the key ID.
    pub fn algorithm(&self) -> A
    where
        A: FromStr,
    {
        A::from_str(&self.as_str()[..self.colon_idx()]).unwrap_or_else(|_| unreachable!())
    }

    /// Returns the key name of the key ID.
    pub fn key_name<'a>(&'a self) -> &'a K
    where
        &'a K: From<&'a str>,
    {
        self.as_str()[self.colon_idx() + 1..].into()
    }

    /// Creates a string slice from this `KeyId`.
    pub fn as_str(&self) -> &str {
        &self.1
    }

    /// Creates a byte slice from this `KeyId`.
    pub fn as_bytes(&self) -> &[u8] {
        self.1.as_bytes()
    }

    fn from_borrowed(s: &str) -> &Self {
        unsafe { std::mem::transmute(s) }
    }

    fn from_owned(s: Box<str>) -> Box<Self> {
        unsafe { Box::from_raw(Box::into_raw(s) as _) }
    }

    fn into_owned(self: Box<Self>) -> Box<str> {
        unsafe { Box::from_raw(Box::into_raw(self) as _) }
    }

    fn colon_idx(&self) -> usize {
        self.as_str().find(':').unwrap()
    }
}

/// Algorithm + key name for signing keys.
pub type SigningKeyId<K> = KeyId<SigningKeyAlgorithm, K>;

/// Algorithm + key name for homeserver signing keys.
pub type ServerSigningKeyId = SigningKeyId<KeyName>;

/// Algorithm + key name for device keys.
pub type DeviceSigningKeyId = SigningKeyId<DeviceId>;

impl<A, K: ?Sized> Clone for Box<KeyId<A, K>> {
    fn clone(&self) -> Self {
        (**self).to_owned()
    }
}

impl<A, K: ?Sized> ToOwned for KeyId<A, K> {
    type Owned = Box<KeyId<A, K>>;
    fn to_owned(&self) -> Self::Owned {
        Self::from_owned(self.1.into())
    }
}

impl<A, K: ?Sized> From<&KeyId<A, K>> for Box<KeyId<A, K>> {
    fn from(id: &KeyId<A, K>) -> Self {
        id.to_owned()
    }
}

impl<A, K: ?Sized> AsRef<str> for Box<KeyId<A, K>> {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl<A, K: ?Sized> From<&KeyId<A, K>> for std::rc::Rc<KeyId<A, K>> {
    fn from(s: &KeyId<A, K>) -> std::rc::Rc<KeyId<A, K>> {
        let rc = std::rc::Rc::<str>::from(s.as_str());
        unsafe { std::rc::Rc::from_raw(std::rc::Rc::into_raw(rc) as *const KeyId<A, K>) }
    }
}

impl<A, K: ?Sized> From<&KeyId<A, K>> for std::sync::Arc<KeyId<A, K>> {
    fn from(s: &KeyId<A, K>) -> std::sync::Arc<KeyId<A, K>> {
        let arc = std::sync::Arc::<str>::from(s.as_str());
        unsafe { std::sync::Arc::from_raw(std::sync::Arc::into_raw(arc) as *const KeyId<A, K>) }
    }
}

impl<A, K: ?Sized> PartialEq<KeyId<A, K>> for Box<KeyId<A, K>> {
    fn eq(&self, other: &KeyId<A, K>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> PartialEq<&'_ KeyId<A, K>> for Box<KeyId<A, K>> {
    fn eq(&self, other: &&KeyId<A, K>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> PartialEq<Box<KeyId<A, K>>> for KeyId<A, K> {
    fn eq(&self, other: &Box<KeyId<A, K>>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> PartialEq<Box<KeyId<A, K>>> for &'_ KeyId<A, K> {
    fn eq(&self, other: &Box<KeyId<A, K>>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> AsRef<str> for KeyId<A, K> {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl<A, K: ?Sized> fmt::Display for KeyId<A, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_str().fmt(f)
    }
}

impl<A, K: ?Sized> fmt::Debug for KeyId<A, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_str().fmt(f)
    }
}

impl<A, K: ?Sized> PartialEq for KeyId<A, K> {
    fn eq(&self, other: &Self) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> Eq for KeyId<A, K> {}

impl<A, K: ?Sized> PartialOrd for KeyId<A, K> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        PartialOrd::partial_cmp(self.as_str(), other.as_str())
    }
}

impl<A, K: ?Sized> Ord for KeyId<A, K> {
    fn cmp(&self, other: &Self) -> Ordering {
        Ord::cmp(self.as_str(), other.as_str())
    }
}

impl<A, K: ?Sized> Hash for KeyId<A, K> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.as_str().hash(state);
    }
}
#[cfg(feature = "serde")]
impl<A, K: ?Sized> serde::Serialize for KeyId<A, K> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(self.as_str())
    }
}

impl<A, K: ?Sized> From<Box<KeyId<A, K>>> for String {
    fn from(id: Box<KeyId<A, K>>) -> Self {
        id.into_owned().into()
    }
}

#[cfg(feature = "serde")]
impl<'de, A, K: ?Sized> serde::Deserialize<'de> for Box<KeyId<A, K>> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;
        let s = String::deserialize(deserializer)?;
        match try_from(s) {
            Ok(o) => Ok(o),
            Err(e) => Err(D::Error::custom(e)),
        }
    }
}

fn try_from<S, A, K: ?Sized>(s: S) -> Result<Box<KeyId<A, K>>, crate::Error>
where
    S: AsRef<str> + Into<Box<str>>,
{
    ruma_identifiers_validation::key_id::validate(s.as_ref())?;
    Ok(KeyId::from_owned(s.into()))
}

impl<'a, A, K: ?Sized> TryFrom<&'a str> for &'a KeyId<A, K> {
    type Error = crate::Error;
    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
        (ruma_identifiers_validation::key_id::validate)(s)?;
        Ok(KeyId::from_borrowed(s))
    }
}

impl<A, K: ?Sized> FromStr for Box<KeyId<A, K>> {
    type Err = crate::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        try_from(s)
    }
}

impl<A, K: ?Sized> TryFrom<&str> for Box<KeyId<A, K>> {
    type Error = crate::Error;
    fn try_from(s: &str) -> Result<Self, Self::Error> {
        try_from(s)
    }
}

impl<A, K: ?Sized> TryFrom<String> for Box<KeyId<A, K>> {
    type Error = crate::Error;
    fn try_from(s: String) -> Result<Self, Self::Error> {
        try_from(s)
    }
}

#[rustfmt::skip]
partial_eq_string!(KeyId<A, K> [A, K]);
<commit_msg>identifiers: Replace qualified paths with imports<commit_after>use std::{
    cmp::Ordering,
    convert::TryFrom,
    fmt,
    hash::{Hash, Hasher},
    marker::PhantomData,
    rc::Rc,
    str::FromStr,
    sync::Arc,
};

use crate::{crypto_algorithms::SigningKeyAlgorithm, DeviceId, KeyName};

/// A key algorithm and key name delimited by a colon.
#[repr(transparent)]
pub struct KeyId<A, K: ?Sized>(PhantomData<(A, K)>, str);

impl<A, K: ?Sized> KeyId<A, K> {
    /// Creates a new `KeyId` from an algorithm and key name.
    pub fn from_parts(algorithm: A, key_name: &K) -> Box<Self>
    where
        A: AsRef<str>,
        K: AsRef<str>,
    {
        let algorithm = algorithm.as_ref();
        let key_name = key_name.as_ref();

        let mut res = String::with_capacity(algorithm.len() + 1 + key_name.len());
        res.push_str(algorithm);
        res.push(':');
        res.push_str(key_name);

        Self::from_owned(res.into())
    }

    /// Returns key algorithm of the key ID.
    pub fn algorithm(&self) -> A
    where
        A: FromStr,
    {
        A::from_str(&self.as_str()[..self.colon_idx()]).unwrap_or_else(|_| unreachable!())
    }

    /// Returns the key name of the key ID.
    pub fn key_name<'a>(&'a self) -> &'a K
    where
        &'a K: From<&'a str>,
    {
        self.as_str()[self.colon_idx() + 1..].into()
    }

    /// Creates a string slice from this `KeyId`.
    pub fn as_str(&self) -> &str {
        &self.1
    }

    /// Creates a byte slice from this `KeyId`.
    pub fn as_bytes(&self) -> &[u8] {
        self.1.as_bytes()
    }

    fn from_borrowed(s: &str) -> &Self {
        unsafe { std::mem::transmute(s) }
    }

    fn from_owned(s: Box<str>) -> Box<Self> {
        unsafe { Box::from_raw(Box::into_raw(s) as _) }
    }

    fn into_owned(self: Box<Self>) -> Box<str> {
        unsafe { Box::from_raw(Box::into_raw(self) as _) }
    }

    fn colon_idx(&self) -> usize {
        self.as_str().find(':').unwrap()
    }
}

/// Algorithm + key name for signing keys.
pub type SigningKeyId<K> = KeyId<SigningKeyAlgorithm, K>;

/// Algorithm + key name for homeserver signing keys.
pub type ServerSigningKeyId = SigningKeyId<KeyName>;

/// Algorithm + key name for device keys.
pub type DeviceSigningKeyId = SigningKeyId<DeviceId>;

impl<A, K: ?Sized> Clone for Box<KeyId<A, K>> {
    fn clone(&self) -> Self {
        (**self).to_owned()
    }
}

impl<A, K: ?Sized> ToOwned for KeyId<A, K> {
    type Owned = Box<KeyId<A, K>>;
    fn to_owned(&self) -> Self::Owned {
        Self::from_owned(self.1.into())
    }
}

impl<A, K: ?Sized> From<&KeyId<A, K>> for Box<KeyId<A, K>> {
    fn from(id: &KeyId<A, K>) -> Self {
        id.to_owned()
    }
}

impl<A, K: ?Sized> AsRef<str> for Box<KeyId<A, K>> {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl<A, K: ?Sized> From<&KeyId<A, K>> for Rc<KeyId<A, K>> {
    fn from(s: &KeyId<A, K>) -> Rc<KeyId<A, K>> {
        let rc = Rc::<str>::from(s.as_str());
        unsafe { Rc::from_raw(Rc::into_raw(rc) as *const KeyId<A, K>) }
    }
}

impl<A, K: ?Sized> From<&KeyId<A, K>> for Arc<KeyId<A, K>> {
    fn from(s: &KeyId<A, K>) -> Arc<KeyId<A, K>> {
        let arc = Arc::<str>::from(s.as_str());
        unsafe { Arc::from_raw(Arc::into_raw(arc) as *const KeyId<A, K>) }
    }
}

impl<A, K: ?Sized> PartialEq<KeyId<A, K>> for Box<KeyId<A, K>> {
    fn eq(&self, other: &KeyId<A, K>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> PartialEq<&'_ KeyId<A, K>> for Box<KeyId<A, K>> {
    fn eq(&self, other: &&KeyId<A, K>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> PartialEq<Box<KeyId<A, K>>> for KeyId<A, K> {
    fn eq(&self, other: &Box<KeyId<A, K>>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> PartialEq<Box<KeyId<A, K>>> for &'_ KeyId<A, K> {
    fn eq(&self, other: &Box<KeyId<A, K>>) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> AsRef<str> for KeyId<A, K> {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl<A, K: ?Sized> fmt::Display for KeyId<A, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_str().fmt(f)
    }
}

impl<A, K: ?Sized> fmt::Debug for KeyId<A, K> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_str().fmt(f)
    }
}

impl<A, K: ?Sized> PartialEq for KeyId<A, K> {
    fn eq(&self, other: &Self) -> bool {
        self.as_str() == other.as_str()
    }
}

impl<A, K: ?Sized> Eq for KeyId<A, K> {}

impl<A, K: ?Sized> PartialOrd for KeyId<A, K> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        PartialOrd::partial_cmp(self.as_str(), other.as_str())
    }
}

impl<A, K: ?Sized> Ord for KeyId<A, K> {
    fn cmp(&self, other: &Self) -> Ordering {
        Ord::cmp(self.as_str(), other.as_str())
    }
}

impl<A, K: ?Sized> Hash for KeyId<A, K> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.as_str().hash(state);
    }
}
#[cfg(feature = "serde")]
impl<A, K: ?Sized> serde::Serialize for KeyId<A, K> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(self.as_str())
    }
}

impl<A, K: ?Sized> From<Box<KeyId<A, K>>> for String {
    fn from(id: Box<KeyId<A, K>>) -> Self {
        id.into_owned().into()
    }
}

#[cfg(feature = "serde")]
impl<'de, A, K: ?Sized> serde::Deserialize<'de> for Box<KeyId<A, K>> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;
        let s = String::deserialize(deserializer)?;
        match try_from(s) {
            Ok(o) => Ok(o),
            Err(e) => Err(D::Error::custom(e)),
        }
    }
}

fn try_from<S, A, K: ?Sized>(s: S) -> Result<Box<KeyId<A, K>>, crate::Error>
where
    S: AsRef<str> + Into<Box<str>>,
{
    ruma_identifiers_validation::key_id::validate(s.as_ref())?;
    Ok(KeyId::from_owned(s.into()))
}

impl<'a, A, K: ?Sized> TryFrom<&'a str> for &'a KeyId<A, K> {
    type Error = crate::Error;
    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
        (ruma_identifiers_validation::key_id::validate)(s)?;
        Ok(KeyId::from_borrowed(s))
    }
}

impl<A, K: ?Sized> FromStr for Box<KeyId<A, K>> {
    type Err = crate::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        try_from(s)
    }
}

impl<A, K: ?Sized> TryFrom<&str> for Box<KeyId<A, K>> {
    type Error = crate::Error;
    fn try_from(s: &str) -> Result<Self, Self::Error> {
        try_from(s)
    }
}

impl<A, K: ?Sized> TryFrom<String> for Box<KeyId<A, K>> {
    type Error = crate::Error;
    fn try_from(s: String) -> Result<Self, Self::Error> {
        try_from(s)
    }
}

#[rustfmt::skip]
partial_eq_string!(KeyId<A, K> [A, K]);
<|endoftext|>


<commit_before>use core::ops::{Bound, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo};

// Test the Range structs and syntax.

#[test]
fn test_range() {
    let r = Range { start: 2, end: 10 };
    let mut count = 0;
    for (i, ri) in r.enumerate() {
        assert_eq!(ri, i + 2);
        assert!(ri >= 2 && ri < 10);
        count += 1;
    }
    assert_eq!(count, 8);
}

#[test]
fn test_range_from() {
    let r = RangeFrom { start: 2 };
    let mut count = 0;
    for (i, ri) in r.take(10).enumerate() {
        assert_eq!(ri, i + 2);
        assert!(ri >= 2 && ri < 12);
        count += 1;
    }
    assert_eq!(count, 10);
}

#[test]
fn test_range_to() {
    // Not much to test.
    let _ = RangeTo { end: 42 };
}

#[test]
fn test_full_range() {
    // Not much to test.
    let _ = RangeFull;
}

#[test]
fn test_range_inclusive() {
    let mut r = RangeInclusive::new(1i8, 2);
    assert_eq!(r.next(), Some(1));
    assert_eq!(r.next(), Some(2));
    assert_eq!(r.next(), None);

    r = RangeInclusive::new(127i8, 127);
    assert_eq!(r.next(), Some(127));
    assert_eq!(r.next(), None);

    r = RangeInclusive::new(-128i8, -128);
    assert_eq!(r.next_back(), Some(-128));
    assert_eq!(r.next_back(), None);

    // degenerate
    r = RangeInclusive::new(1, -1);
    assert_eq!(r.size_hint(), (0, Some(0)));
    assert_eq!(r.next(), None);
}

#[test]
fn test_range_is_empty() {
    assert!(!(0.0..10.0).is_empty());
    assert!((-0.0..0.0).is_empty());
    assert!((10.0..0.0).is_empty());

    assert!(!(f32::NEG_INFINITY..f32::INFINITY).is_empty());
    assert!((f32::EPSILON..f32::NAN).is_empty());
    assert!((f32::NAN..f32::EPSILON).is_empty());
    assert!((f32::NAN..f32::NAN).is_empty());

    assert!(!(0.0..=10.0).is_empty());
    assert!(!(-0.0..=0.0).is_empty());
    assert!((10.0..=0.0).is_empty());

    assert!(!(f32::NEG_INFINITY..=f32::INFINITY).is_empty());
    assert!((f32::EPSILON..=f32::NAN).is_empty());
    assert!((f32::NAN..=f32::EPSILON).is_empty());
    assert!((f32::NAN..=f32::NAN).is_empty());
}

#[test]
fn test_bound_cloned_unbounded() {
    assert_eq!(Bound::<&u32>::Unbounded.cloned(), Bound::Unbounded);
}

#[test]
fn test_bound_cloned_included() {
    assert_eq!(Bound::Included(&3).cloned(), Bound::Included(3));
}

#[test]
fn test_bound_cloned_excluded() {
    assert_eq!(Bound::Excluded(&3).cloned(), Bound::Excluded(3));
}

#[test]
#[allow(unused_comparisons)]
#[allow(unused_mut)]
fn test_range_syntax() {
    let mut count = 0;
    for i in 0_usize..10 {
        assert!(i >= 0 && i < 10);
        count += i;
    }
    assert_eq!(count, 45);

    let mut count = 0;
    let mut range = 0_usize..10;
    for i in range {
        assert!(i >= 0 && i < 10);
        count += i;
    }
    assert_eq!(count, 45);

    let mut count = 0;
    let mut rf = 3_usize..;
    for i in rf.take(10) {
        assert!(i >= 3 && i < 13);
        count += i;
    }
    assert_eq!(count, 75);

    let _ = 0_usize..4 + 4 - 3;

    fn foo() -> isize {
        42
    }
    let _ = 0..foo();

    let _ = { &42..&100 }; // references to literals are OK
    let _ = ..42_usize;

    // Test we can use two different types with a common supertype.
    let x = &42;
    {
        let y = 42;
        let _ = x..&y;
    }
}

#[test]
#[allow(dead_code)]
fn test_range_syntax_in_return_statement() {
    fn return_range_to() -> RangeTo<i32> {
        return ..1;
    }
    fn return_full_range() -> RangeFull {
        return ..;
    }
    // Not much to test.
}
<commit_msg>Rollup merge of #76339 - CDirkx:structural-match-range, r=Mark-Simulacrum<commit_after>use core::ops::{Bound, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};

// Test the Range structs and syntax.

#[test]
fn test_range() {
    let r = Range { start: 2, end: 10 };
    let mut count = 0;
    for (i, ri) in r.enumerate() {
        assert_eq!(ri, i + 2);
        assert!(ri >= 2 && ri < 10);
        count += 1;
    }
    assert_eq!(count, 8);
}

#[test]
fn test_range_from() {
    let r = RangeFrom { start: 2 };
    let mut count = 0;
    for (i, ri) in r.take(10).enumerate() {
        assert_eq!(ri, i + 2);
        assert!(ri >= 2 && ri < 12);
        count += 1;
    }
    assert_eq!(count, 10);
}

#[test]
fn test_range_to() {
    // Not much to test.
    let _ = RangeTo { end: 42 };
}

#[test]
fn test_full_range() {
    // Not much to test.
    let _ = RangeFull;
}

#[test]
fn test_range_inclusive() {
    let mut r = RangeInclusive::new(1i8, 2);
    assert_eq!(r.next(), Some(1));
    assert_eq!(r.next(), Some(2));
    assert_eq!(r.next(), None);

    r = RangeInclusive::new(127i8, 127);
    assert_eq!(r.next(), Some(127));
    assert_eq!(r.next(), None);

    r = RangeInclusive::new(-128i8, -128);
    assert_eq!(r.next_back(), Some(-128));
    assert_eq!(r.next_back(), None);

    // degenerate
    r = RangeInclusive::new(1, -1);
    assert_eq!(r.size_hint(), (0, Some(0)));
    assert_eq!(r.next(), None);
}

#[test]
fn test_range_to_inclusive() {
    // Not much to test.
    let _ = RangeToInclusive { end: 42 };
}

#[test]
fn test_range_is_empty() {
    assert!(!(0.0..10.0).is_empty());
    assert!((-0.0..0.0).is_empty());
    assert!((10.0..0.0).is_empty());

    assert!(!(f32::NEG_INFINITY..f32::INFINITY).is_empty());
    assert!((f32::EPSILON..f32::NAN).is_empty());
    assert!((f32::NAN..f32::EPSILON).is_empty());
    assert!((f32::NAN..f32::NAN).is_empty());

    assert!(!(0.0..=10.0).is_empty());
    assert!(!(-0.0..=0.0).is_empty());
    assert!((10.0..=0.0).is_empty());

    assert!(!(f32::NEG_INFINITY..=f32::INFINITY).is_empty());
    assert!((f32::EPSILON..=f32::NAN).is_empty());
    assert!((f32::NAN..=f32::EPSILON).is_empty());
    assert!((f32::NAN..=f32::NAN).is_empty());
}

#[test]
fn test_bound_cloned_unbounded() {
    assert_eq!(Bound::<&u32>::Unbounded.cloned(), Bound::Unbounded);
}

#[test]
fn test_bound_cloned_included() {
    assert_eq!(Bound::Included(&3).cloned(), Bound::Included(3));
}

#[test]
fn test_bound_cloned_excluded() {
    assert_eq!(Bound::Excluded(&3).cloned(), Bound::Excluded(3));
}

#[test]
#[allow(unused_comparisons)]
#[allow(unused_mut)]
fn test_range_syntax() {
    let mut count = 0;
    for i in 0_usize..10 {
        assert!(i >= 0 && i < 10);
        count += i;
    }
    assert_eq!(count, 45);

    let mut count = 0;
    let mut range = 0_usize..10;
    for i in range {
        assert!(i >= 0 && i < 10);
        count += i;
    }
    assert_eq!(count, 45);

    let mut count = 0;
    let mut rf = 3_usize..;
    for i in rf.take(10) {
        assert!(i >= 3 && i < 13);
        count += i;
    }
    assert_eq!(count, 75);

    let _ = 0_usize..4 + 4 - 3;

    fn foo() -> isize {
        42
    }
    let _ = 0..foo();

    let _ = { &42..&100 }; // references to literals are OK
    let _ = ..42_usize;

    // Test we can use two different types with a common supertype.
    let x = &42;
    {
        let y = 42;
        let _ = x..&y;
    }
}

#[test]
#[allow(dead_code)]
fn test_range_syntax_in_return_statement() {
    fn return_range_to() -> RangeTo<i32> {
        return ..1;
    }
    fn return_full_range() -> RangeFull {
        return ..;
    }
    // Not much to test.
}

#[test]
fn range_structural_match() {
    // test that all range types can be structurally matched upon

    const RANGE: Range<usize> = 0..1000;
    match RANGE {
        RANGE => {}
        _ => unreachable!(),
    }

    const RANGE_FROM: RangeFrom<usize> = 0..;
    match RANGE_FROM {
        RANGE_FROM => {}
        _ => unreachable!(),
    }

    const RANGE_FULL: RangeFull = ..;
    match RANGE_FULL {
        RANGE_FULL => {}
    }

    const RANGE_INCLUSIVE: RangeInclusive<usize> = 0..=999;
    match RANGE_INCLUSIVE {
        RANGE_INCLUSIVE => {}
        _ => unreachable!(),
    }

    const RANGE_TO: RangeTo<usize> = ..1000;
    match RANGE_TO {
        RANGE_TO => {}
        _ => unreachable!(),
    }

    const RANGE_TO_INCLUSIVE: RangeToInclusive<usize> = ..=999;
    match RANGE_TO_INCLUSIVE {
        RANGE_TO_INCLUSIVE => {}
        _ => unreachable!(),
    }
}
<|endoftext|>


<commit_before>use std::{Vec, String, ToString};
use std::collections::{BTreeMap, VecDeque};

use super::dvaddr::DVAddr;
use super::zio;

/// MRU - Most Recently Used cache
struct Mru {
    map: BTreeMap<DVAddr, Vec<u8>>,
    queue: VecDeque<DVAddr>, // Oldest DVAddrs are at the end
    size: usize, // Max mru cache size in blocks
    used: usize, // Number of used blocks in mru cache
}

impl Mru {
    pub fn new() -> Self {
        Mru {
            map: BTreeMap::new(),
            queue: VecDeque::new(),
            size: 1000,
            used: 0,
        }
    }

    fn cache_block(&mut self, dva: &DVAddr, block: Vec<u8>) -> Result<Vec<u8>, String> {
        // If necessary, make room for the block in the cache
        while self.used + (dva.asize() as usize) > self.size {
            let last_dva = match self.queue.pop_back() {
                Some(dva) => dva,
                None => return Err("No more ARC MRU items to free".to_string()),
            };
            self.map.remove(&last_dva);
            self.used -= last_dva.asize() as usize;
        }

        // Add the block to the cache
        self.used += dva.asize() as usize;
        self.map.insert(*dva, block);
        self.queue.push_front(*dva);
        Ok(self.map.get(dva).unwrap().clone())
    }
}

/// MFU - Most Frequently Used cache
struct Mfu {
    // TODO: Keep track of use counts. So mfu_map becomes (use_count: u64, Vec<u8>). Reset the use
    // count every once in a while. For instance, every 1000 reads. This will probably end up being
    // a knob for the user.
    // TODO: Keep track of minimum frequency and corresponding DVA
    map: BTreeMap<DVAddr, (u64, Vec<u8>)>,
    size: usize, // Max mfu cache size in blocks
    used: usize, // Number of used bytes in mfu cache
}

impl Mfu {
    pub fn new() -> Self {
        Mfu {
            map: BTreeMap::new(),
            size: 1000,
            used: 0,
        }
    }

    // TODO: cache_block. Remove the DVA with the lowest frequency
    /*
    fn cache_block(&mut self, dva: &DVAddr, block: Vec<u8>) -> Result<Vec<u8>, String> {
    }
    */
}

// Our implementation of the Adaptive Replacement Cache (ARC) is set up to allocate
// its buffer on the heap rather than in a private pool thing. This makes it much
// simpler to implement, but defers the fragmentation problem to the heap allocator.
// We named the type `ArCache` to avoid confusion with Rust's `Arc` reference type.
pub struct ArCache {
    mru: Mru,
    mfu: Mfu,
}

impl ArCache {
    pub fn new() -> Self {
        ArCache {
            mru: Mru::new(),
            mfu: Mfu::new(),
        }
    }

    pub fn read(&mut self, reader: &mut zio::Reader, dva: &DVAddr) -> Result<Vec<u8>, String> {
        if let Some(block) = self.mru.map.remove(dva) {
            self.mfu.map.insert(*dva, (0, block.clone()));

            // Block is cached
            return Ok(block);
        }
        if let Some(block) = self.mfu.map.get_mut(dva) {
            // Block is cached
            if block.0 > 1000 {
                block.0 = 0;
            } else {
                block.0 += 1;
            }

            return Ok(block.1.clone());
        }

        // Block isn't cached, have to read it from disk
        let block = reader.read(dva.sector() as usize, dva.asize() as usize);

        // Blocks start in MRU cache
        self.mru.cache_block(dva, block)
    }
}
<commit_msg>Finish TODO for MFU in ZFS<commit_after>use std::{Vec, String, ToString};
use std::collections::{BTreeMap, VecDeque};

use super::dvaddr::DVAddr;
use super::zio;

/// MRU - Most Recently Used cache
struct Mru {
    map: BTreeMap<DVAddr, Vec<u8>>,
    queue: VecDeque<DVAddr>, // Oldest DVAddrs are at the end
    size: usize, // Max mru cache size in blocks
    used: usize, // Number of used blocks in mru cache
}

impl Mru {
    pub fn new() -> Self {
        Mru {
            map: BTreeMap::new(),
            queue: VecDeque::new(),
            size: 1000,
            used: 0,
        }
    }

    pub fn cache_block(&mut self, dva: &DVAddr, block: Vec<u8>) -> Result<Vec<u8>, String> {
        // If necessary, make room for the block in the cache
        while self.used + (dva.asize() as usize) > self.size {
            let last_dva = match self.queue.pop_back() {
                Some(dva) => dva,
                None => return Err("No more ARC MRU items to free".to_string()),
            };
            self.map.remove(&last_dva);
            self.used -= last_dva.asize() as usize;
        }

        // Add the block to the cache
        self.used += dva.asize() as usize;
        self.map.insert(*dva, block);
        self.queue.push_front(*dva);
        Ok(self.map.get(dva).unwrap().clone())
    }
}

/// MFU - Most Frequently Used cache
struct Mfu {
    // TODO: Keep track of use counts. So mfu_map becomes (use_count: u64, Vec<u8>). Reset the use
    // count every once in a while. For instance, every 1000 reads. This will probably end up being
    // a knob for the user.
    // TODO: Keep track of minimum frequency and corresponding DVA
    map: BTreeMap<DVAddr, (u64, Vec<u8>)>,
    size: usize, // Max mfu cache size in blocks
    used: usize, // Number of used bytes in mfu cache
}

impl Mfu {
    pub fn new() -> Self {
        Mfu {
            map: BTreeMap::new(),
            size: 1000,
            used: 0,
        }
    }

    pub fn cache_block(&mut self, dva: &DVAddr, block: Vec<u8>) -> Result<Vec<u8>, String> {
        {
            let mut lowest_freq = 0;
            let mut lowest_dva: Result<DVAddr, String> = Err("No valid DVA found.".to_string());

            for (&dva_key, &(freq, _)) in self.map.iter() {
                if freq < lowest_freq {
                    lowest_freq = freq;
                    lowest_dva = Ok(dva_key);
                }
            }

            self.map.remove(&try!(lowest_dva));
        }

        // Add the block to the cache
        self.used += dva.asize() as usize;
        self.map.insert(*dva, (1, block));
        Ok(self.map.get(dva).unwrap().1.clone())
    }
}

// Our implementation of the Adaptive Replacement Cache (ARC) is set up to allocate
// its buffer on the heap rather than in a private pool thing. This makes it much
// simpler to implement, but defers the fragmentation problem to the heap allocator.
// We named the type `ArCache` to avoid confusion with Rust's `Arc` reference type.
pub struct ArCache {
    mru: Mru,
    mfu: Mfu,
}

impl ArCache {
    pub fn new() -> Self {
        ArCache {
            mru: Mru::new(),
            mfu: Mfu::new(),
        }
    }

    pub fn read(&mut self, reader: &mut zio::Reader, dva: &DVAddr) -> Result<Vec<u8>, String> {
        if let Some(block) = self.mru.map.remove(dva) {
            self.mfu.map.insert(*dva, (0, block.clone()));

            // Block is cached
            return Ok(block);
        }
        if let Some(block) = self.mfu.map.get_mut(dva) {
            // Block is cached
            if block.0 > 1000 {
                block.0 = 0;
            } else {
                block.0 += 1;
            }

            return Ok(block.1.clone());
        }

        // Block isn't cached, have to read it from disk
        let block = reader.read(dva.sector() as usize, dva.asize() as usize);

        // Blocks start in MRU cache
        self.mru.cache_block(dva, block)
    }
}
<|endoftext|>