<commit_before>//! An Implementation of RFC 1951
//!
//! The DEFLATE compression algorithm
//!
//! # Related Links
//! *http://tools.ietf.org/html/rfc1951 - DEFLATE Compressed Data Format Specification

use std::cmp;
use std::iter::repeat;
use std::num::Wrapping as w;
use std::io::{self, Read};
use byteorder::ReadBytesExt;

static LITERALLENGTHCODES: u16 = 286;
static DISTANCECODES: u16 = 30;
static CODEORDER: [u8; 19] = [
    16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15
];
static LENGTHS: [u16; 29] = [
    3,  4,  5,   6,   7,   8,   9,  10,  11, 13,
    15, 17, 19,  23,  27,  31,  35,  43,  51, 59,
    67, 83, 99, 115, 131, 163, 195, 227, 258
];
static EXTRA_LENGTHS: [u8; 29] = [
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1,
    1, 1, 2, 2, 2, 2, 3, 3, 3, 3,
    4, 4, 4, 4, 5, 5, 5, 5, 0
];
static DISTANCES: [u16; 30] = [
    1,    2,      3,    4,    5,    7,    9,    13,    17,    25,
    33,   49,     65,   97,  129,  193,  257,   385,   513,   769,
    1025,  1537,  2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577
];
static EXTRA_DISTANCES: [u8; 30] = [
    0, 0,  0,  0,  1,  1,  2,  2,  3,  3,
    4, 4,  5,  5,  6,  6,  7,  7,  8,  8,
    9, 9, 10, 10, 11, 11, 12, 12, 13, 13
];

static TABLESIZE: u8 = 9;

#[derive(PartialEq, Clone)]
enum TableElement {
    Symbol(u16, u8),
    Table(u16, Vec<TableElement>),
    Nothing
}

impl TableElement {
    pub fn put(&mut self, index: u16, elem: TableElement) {
        match *self {
            TableElement::Table(_, ref mut a) => a[index as usize] = elem,
            _		                          => panic!("requires Table()"),
        }
    }
}

enum BlockType {
    Stored,
    Compressed
}

/// A DEFLATE compressed stream decoder.
pub struct Inflater<R> {
    h:
    HuffReader<R>,

    buf: Vec<u8>,
    pos: u64,

    finished: bool,
    btype: BlockType,
    block_length: u32,

    ctable: Vec<TableElement>,
    lltable: Vec<TableElement>,
    dtable: Vec<TableElement>,
}

impl<R: Read> Inflater<R> {
    /// Create a new decoder that decodes from a Reader
    pub fn new(r: R) -> Inflater<R> {
        Inflater {
            h: HuffReader::new(r),

            buf: Vec::new(),
            pos: 0,

            finished: false,
            block_length: 0,
            btype: BlockType::Stored,

            ctable: Vec::new(),
            lltable: Vec::new(),
            dtable: Vec::new(),
        }
    }

    /// Indicate whether the end of the stream has been reached.
    pub fn eof(&self) -> bool {
        self.finished && (self.pos as usize == self.buf.len())
    }

    /// Return a mutable reference to the wrapped Reader
    pub fn inner(&mut self) -> &mut R {
        &mut self.h.r
    }

    fn read_block_type(&mut self) -> io::Result<()> {
        let is_final = try!(self.h.receive(1));
        self.finished = is_final == 1;

        let bits = try!(self.h.receive(2));
        match bits {
            0b00 => {
                let _ = try!(self.read_stored_block_length());
                self.btype = BlockType::Stored;
            }
            0b01 => {
                self.create_fixed_tables();
                self.btype = BlockType::Compressed;
            }
            0b10 => {
                let _ = try!(self.read_dynamic_tables());
                self.btype = BlockType::Compressed;
            }
            _ => panic!("reserved block type")
        }

        Ok(())
    }

    fn read_dynamic_tables(&mut self) -> io::Result<()> {
        let totalcodes = LITERALLENGTHCODES + DISTANCECODES;

        let hlit  = try!(self.h.receive(5)) + 257;
        let hdist = try!(self.h.receive(5)) + 1;
        let hclen = try!(self.h.receive(4)) + 4;

        let mut code_lengths: Vec<u8> = repeat(0u8).take(CODEORDER.len()).collect();

        for i in (0..hclen as usize) {
            let length = try!(self.h.receive(3));
            code_lengths[CODEORDER[i] as usize] = length as u8;
        }

        self.ctable = table_from_lengths(&code_lengths);
        let mut all_lengths: Vec<u8> = repeat(0u8).take(totalcodes as usize).collect();

        let mut i = 0;
        while i < hlit + hdist {
            let s = try!(self.h.decode_symbol(&self.ctable));

            match s {
                0 ... 15 => {
                    all_lengths[i as usize] = s as u8;
                    i += 1;
                }

                16 => {
                    let repeat = 3 + try!(self.h.receive(2));

                    for _ in (0..repeat) {
                        all_lengths[i as usize] = all_lengths[i as usize - 1];
                        i += 1;
                    }
                }

                17 => i += 3 + try!(self.h.receive(3)),

                18 => i += 11 + try!(self.h.receive(7)),

                _ => panic!("out of range code length code symbol")
            }
        }

        let ll_lengths = &all_lengths[..hlit as usize];
        let d_lengths  = &all_lengths[hlit as usize..];

        self.lltable = table_from_lengths(ll_lengths);
        self.dtable = table_from_lengths(d_lengths);

        Ok(())
    }

    fn create_fixed_tables(&mut self) {
        let lengths: Vec<u8> = (0u32..288).map(|i|
            if i < 144 { 8u8 }
            else if i < 256 { 9u8 }
            else if i < 280 { 7u8 }
            else { 8u8 }
        ).collect();
        self.lltable = table_from_lengths(&lengths);

        let lengths: Vec<u8> = repeat(5u8).take(DISTANCECODES as usize).collect();
        self.dtable = table_from_lengths(&lengths);
    }

    fn read_stored_block_length(&mut self) -> io::Result<()> {
        self.h.byte_align();

        let len   = try!(self.h.receive(16));
        let _nlen = try!(self.h.receive(16));

        self.block_length = len as u32;

        Ok(())
    }

    fn read_stored_block(&mut self) -> io::Result<()> {
        while self.block_length > 0 {
            let a = try!(self.h.receive(8));

            self.buf.push(a as u8);
            self.h.consume(8);
            self.block_length -= 1;
        }

        Ok(())
    }

    fn read_compressed_block(&mut self) -> io::Result<()> {
        loop {
            let s = try!(self.h.decode_symbol(&self.lltable));

            match s {
                literal @ 0 ... 255 => self.buf.push(literal as u8),

                256 => break,

                length @ 257 ... 285 => {
                    let length = length - 257;

                    let bits = EXTRA_LENGTHS[length as usize];
                    let extra = try!(self.h.receive(bits));

                    let length = LENGTHS[length as usize] + extra;

                    let distance = try!(self.h.decode_symbol(&self.dtable));

                    let bits = EXTRA_DISTANCES[distance as usize];
                    let extra = try!(self.h.receive(bits));

                    let distance = DISTANCES[distance as usize] + extra;

                    let len = self.buf.len();
                    for i in (0..length) {
                        let s = self.buf[len - distance as usize + i as usize];
                        self.buf.push(s);
                    }
                }

                _ => panic!("out of range symbol")
            }
        }

        Ok(())
    }
}

impl<R: Read> Read for Inflater<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.pos as usize == self.buf.len() {
            if self.finished {
                return Ok(0);
            }

            let _ = try!(self.read_block_type());
            let _ = match self.btype {
                BlockType::Stored => try!(self.read_stored_block()),
                BlockType::Compressed => try!(self.read_compressed_block())
            };
        }

        let n = cmp::min(buf.len(), self.buf.len() - self.pos as usize);
        for i in (0usize..n) {
            buf[i] = self.buf[self.pos as usize + i];
        }

        self.pos += n as u64;
        Ok(n)
    }
}

fn reverse(a: u16) -> u16 {
    let b = (((!0x5555) & a) >> 1) | ((0x5555 & a) << 1);
    let c = (((!0x3333) & b) >> 2) | ((0x3333 & b) << 2);
    let d = (((!0x0F0F) & c) >> 4) | ((0x0F0F & c) << 4);

    (((!0x00FF) & d) >> 8) | ((0x00FF & d) << 8)
}

fn table_from_lengths(lengths: &[u8]) -> Vec<TableElement> {
    let mut max_len = 0;
    let mut code = 0u16;
    let mut next_code: Vec<u16> = repeat(0u16).take(16).collect();
    let mut bl_count: Vec<w<u8>> = repeat(w(0u8)).take(16).collect();

    for &len in lengths.iter() {
        bl_count[len as usize] = bl_count[len as usize] + w(1);
        if len > max_len {
            max_len = len;
        }
    }

    let max_overflow = (w(max_len) - w(TABLESIZE)).0;
    bl_count[0] = w(0);

    for bits in (1usize..16) {
        code = (code + bl_count[bits - 1].0 as u16) << 1;
        next_code[bits] = code;
    }

    let mut lut: Vec<TableElement> = repeat(TableElement::Nothing).take(1 << TABLESIZE as usize).collect();

    for (i, &len) in lengths.iter().enumerate() {
        if len == 0 {
            continue
        }

        let code = next_code[len as usize];
        let code = reverse(code) >> (16 - len) as usize;

        if len <= TABLESIZE {
            let r = TABLESIZE - len;

            for j in (0u16..1 << r as usize) {
                let index = (j << len as usize) + code;
                lut[index as usize] = TableElement::Symbol(i as u16, len);
            }
        } else {
            let index = code & ((1 << TABLESIZE as usize) - 1);

            if lut[index as usize] == TableElement::Nothing {
                let mask  = (1 << max_overflow) - 1;
                let array: Vec<TableElement> = repeat(TableElement::Nothing).take(1 << max_overflow).collect();

                lut[index as usize] = TableElement::Table(mask, array);
            }

            let code = code >> TABLESIZE as usize;
            let r = max_len - len;

            for j in (0u16..1 << r as usize) {
                let k = (j << (len - TABLESIZE) as usize) + code;
                let s = TableElement::Symbol(i as u16, len - TABLESIZE);

                lut[index as usize].put(k, s);
            }
        }

        next_code[len as usize] += 1;
    }

    lut
}

struct HuffReader<R> {
    pub r: R,

    bits: u32,
    num_bits: u8,
}

impl<R: Read> HuffReader<R> {
    pub fn new(r: R) -> HuffReader<R> {
        HuffReader {r: r, bits: 0, num_bits: 0}
    }

    pub fn guarantee(&mut self, n: u8) -> io::Result<()> {
        while self.num_bits < n {
            let byte = try!(self.r.read_u8());

            self.bits |= (byte as u32) << self.num_bits as usize;
            self.num_bits += 8;
        }

        Ok(())
    }

    pub fn byte_align(&mut self) {
        let n = self.bits & 0b111;

        self.bits >>= n as usize;
        self.num_bits -= n as u8;
    }

    pub fn consume(&mut self, n: u8) {
        self.bits >>= n as usize;
        self.num_bits -= n;
    }

    pub fn receive(&mut self, n: u8) -> io::Result<u16> {
        let _ = try!(self.guarantee(n));

        let val = self.bits & ((1 << n as usize) - 1);
        self.consume(n);

        Ok(val as u16)
    }

    pub fn decode_symbol(&mut self, table: &[TableElement]) -> io::Result<u16> {
        let _ = try!(self.guarantee(1));

        loop {
            let index = self.bits & ((1 << TABLESIZE as usize) - 1);

            let (val, size) = match table[index as usize] {
                TableElement::Symbol(val, size) => (val, size),

                TableElement::Table(mask, ref a) => {
                    let index = (self.bits >> TABLESIZE as usize) & mask as u32;

                    match a[index as usize] {
                        TableElement::Symbol(val, size) => (val, size + TABLESIZE),
                        _ 		  => panic!("bad huffman code")
                    }
                }

                TableElement::Nothing => panic!("bad huffman code")
            };

            if size <= self.num_bits {
                self.consume(size);
                return	Ok(val)
            }

            let _ = try!(self.guarantee(size));
        }
    }
}
<commit_msg>Fix inflate with uncompressed blocks<commit_after>//! An Implementation of RFC 1951
//!
//! The DEFLATE compression algorithm
//!
//! # Related Links
//! *http://tools.ietf.org/html/rfc1951 - DEFLATE Compressed Data Format Specification

use std::cmp;
use std::iter::repeat;
use std::num::Wrapping as w;
use std::io::{self, Read};
use byteorder::ReadBytesExt;

static LITERALLENGTHCODES: u16 = 286;
static DISTANCECODES: u16 = 30;
static CODEORDER: [u8; 19] = [
    16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15
];
static LENGTHS: [u16; 29] = [
    3,  4,  5,   6,   7,   8,   9,  10,  11, 13,
    15, 17, 19,  23,  27,  31,  35,  43,  51, 59,
    67, 83, 99, 115, 131, 163, 195, 227, 258
];
static EXTRA_LENGTHS: [u8; 29] = [
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1,
    1, 1, 2, 2, 2, 2, 3, 3, 3, 3,
    4, 4, 4, 4, 5, 5, 5, 5, 0
];
static DISTANCES: [u16; 30] = [
    1,    2,      3,    4,    5,    7,    9,    13,    17,    25,
    33,   49,     65,   97,  129,  193,  257,   385,   513,   769,
    1025,  1537,  2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577
];
static EXTRA_DISTANCES: [u8; 30] = [
    0, 0,  0,  0,  1,  1,  2,  2,  3,  3,
    4, 4,  5,  5,  6,  6,  7,  7,  8,  8,
    9, 9, 10, 10, 11, 11, 12, 12, 13, 13
];

static TABLESIZE: u8 = 9;

#[derive(PartialEq, Clone)]
enum TableElement {
    Symbol(u16, u8),
    Table(u16, Vec<TableElement>),
    Nothing
}

impl TableElement {
    pub fn put(&mut self, index: u16, elem: TableElement) {
        match *self {
            TableElement::Table(_, ref mut a) => a[index as usize] = elem,
            _		                          => panic!("requires Table()"),
        }
    }
}

enum BlockType {
    Stored,
    Compressed
}

/// A DEFLATE compressed stream decoder.
pub struct Inflater<R> {
    h:
    HuffReader<R>,

    buf: Vec<u8>,
    pos: u64,

    finished: bool,
    btype: BlockType,
    block_length: u32,

    ctable: Vec<TableElement>,
    lltable: Vec<TableElement>,
    dtable: Vec<TableElement>,
}

impl<R: Read> Inflater<R> {
    /// Create a new decoder that decodes from a Reader
    pub fn new(r: R) -> Inflater<R> {
        Inflater {
            h: HuffReader::new(r),

            buf: Vec::new(),
            pos: 0,

            finished: false,
            block_length: 0,
            btype: BlockType::Stored,

            ctable: Vec::new(),
            lltable: Vec::new(),
            dtable: Vec::new(),
        }
    }

    /// Indicate whether the end of the stream has been reached.
    pub fn eof(&self) -> bool {
        self.finished && (self.pos as usize == self.buf.len())
    }

    /// Return a mutable reference to the wrapped Reader
    pub fn inner(&mut self) -> &mut R {
        &mut self.h.r
    }

    fn read_block_type(&mut self) -> io::Result<()> {
        let is_final = try!(self.h.receive(1));
        self.finished = is_final == 1;

        let bits = try!(self.h.receive(2));
        match bits {
            0b00 => {
                let _ = try!(self.read_stored_block_length());
                self.btype = BlockType::Stored;
            }
            0b01 => {
                self.create_fixed_tables();
                self.btype = BlockType::Compressed;
            }
            0b10 => {
                let _ = try!(self.read_dynamic_tables());
                self.btype = BlockType::Compressed;
            }
            _ => panic!("reserved block type")
        }

        Ok(())
    }

    fn read_dynamic_tables(&mut self) -> io::Result<()> {
        let totalcodes = LITERALLENGTHCODES + DISTANCECODES;

        let hlit  = try!(self.h.receive(5)) + 257;
        let hdist = try!(self.h.receive(5)) + 1;
        let hclen = try!(self.h.receive(4)) + 4;

        let mut code_lengths: Vec<u8> = repeat(0u8).take(CODEORDER.len()).collect();

        for i in (0..hclen as usize) {
            let length = try!(self.h.receive(3));
            code_lengths[CODEORDER[i] as usize] = length as u8;
        }

        self.ctable = table_from_lengths(&code_lengths);
        let mut all_lengths: Vec<u8> = repeat(0u8).take(totalcodes as usize).collect();

        let mut i = 0;
        while i < hlit + hdist {
            let s = try!(self.h.decode_symbol(&self.ctable));

            match s {
                0 ... 15 => {
                    all_lengths[i as usize] = s as u8;
                    i += 1;
                }

                16 => {
                    let repeat = 3 + try!(self.h.receive(2));

                    for _ in (0..repeat) {
                        all_lengths[i as usize] = all_lengths[i as usize - 1];
                        i += 1;
                    }
                }

                17 => i += 3 + try!(self.h.receive(3)),

                18 => i += 11 + try!(self.h.receive(7)),

                _ => panic!("out of range code length code symbol")
            }
        }

        let ll_lengths = &all_lengths[..hlit as usize];
        let d_lengths  = &all_lengths[hlit as usize..];

        self.lltable = table_from_lengths(ll_lengths);
        self.dtable = table_from_lengths(d_lengths);

        Ok(())
    }

    fn create_fixed_tables(&mut self) {
        let lengths: Vec<u8> = (0u32..288).map(|i|
            if i < 144 { 8u8 }
            else if i < 256 { 9u8 }
            else if i < 280 { 7u8 }
            else { 8u8 }
        ).collect();
        self.lltable = table_from_lengths(&lengths);

        let lengths: Vec<u8> = repeat(5u8).take(DISTANCECODES as usize).collect();
        self.dtable = table_from_lengths(&lengths);
    }

    fn read_stored_block_length(&mut self) -> io::Result<()> {
        self.h.byte_align();

        let len   = try!(self.h.receive(16));
        let _nlen = try!(self.h.receive(16));

        self.block_length = len as u32;

        Ok(())
    }

    fn read_stored_block(&mut self) -> io::Result<()> {
        while self.block_length > 0 {
            let a = try!(self.h.receive(8));
            self.buf.push(a as u8);
            self.block_length -= 1;
        }

        Ok(())
    }

    fn read_compressed_block(&mut self) -> io::Result<()> {
        loop {
            let s = try!(self.h.decode_symbol(&self.lltable));

            match s {
                literal @ 0 ... 255 => self.buf.push(literal as u8),

                256 => break,

                length @ 257 ... 285 => {
                    let length = length - 257;

                    let bits = EXTRA_LENGTHS[length as usize];
                    let extra = try!(self.h.receive(bits));

                    let length = LENGTHS[length as usize] + extra;

                    let distance = try!(self.h.decode_symbol(&self.dtable));

                    let bits = EXTRA_DISTANCES[distance as usize];
                    let extra = try!(self.h.receive(bits));

                    let distance = DISTANCES[distance as usize] + extra;

                    let len = self.buf.len();
                    for i in (0..length) {
                        let s = self.buf[len - distance as usize + i as usize];
                        self.buf.push(s);
                    }
                }

                _ => panic!("out of range symbol")
            }
        }

        Ok(())
    }
}

impl<R: Read> Read for Inflater<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.pos as usize == self.buf.len() {
            if self.finished {
                return Ok(0);
            }

            let _ = try!(self.read_block_type());
            let _ = match self.btype {
                BlockType::Stored => try!(self.read_stored_block()),
                BlockType::Compressed => try!(self.read_compressed_block())
            };
        }

        let n = cmp::min(buf.len(), self.buf.len() - self.pos as usize);
        for i in (0usize..n) {
            buf[i] = self.buf[self.pos as usize + i];
        }

        self.pos += n as u64;
        Ok(n)
    }
}

fn reverse(a: u16) -> u16 {
    let b = (((!0x5555) & a) >> 1) | ((0x5555 & a) << 1);
    let c = (((!0x3333) & b) >> 2) | ((0x3333 & b) << 2);
    let d = (((!0x0F0F) & c) >> 4) | ((0x0F0F & c) << 4);

    (((!0x00FF) & d) >> 8) | ((0x00FF & d) << 8)
}

fn table_from_lengths(lengths: &[u8]) -> Vec<TableElement> {
    let mut max_len = 0;
    let mut code = 0u16;
    let mut next_code: Vec<u16> = repeat(0u16).take(16).collect();
    let mut bl_count: Vec<w<u8>> = repeat(w(0u8)).take(16).collect();

    for &len in lengths.iter() {
        bl_count[len as usize] = bl_count[len as usize] + w(1);
        if len > max_len {
            max_len = len;
        }
    }

    let max_overflow = (w(max_len) - w(TABLESIZE)).0;
    bl_count[0] = w(0);

    for bits in (1usize..16) {
        code = (code + bl_count[bits - 1].0 as u16) << 1;
        next_code[bits] = code;
    }

    let mut lut: Vec<TableElement> = repeat(TableElement::Nothing).take(1 << TABLESIZE as usize).collect();

    for (i, &len) in lengths.iter().enumerate() {
        if len == 0 {
            continue
        }

        let code = next_code[len as usize];
        let code = reverse(code) >> (16 - len) as usize;

        if len <= TABLESIZE {
            let r = TABLESIZE - len;

            for j in (0u16..1 << r as usize) {
                let index = (j << len as usize) + code;
                lut[index as usize] = TableElement::Symbol(i as u16, len);
            }
        } else {
            let index = code & ((1 << TABLESIZE as usize) - 1);

            if lut[index as usize] == TableElement::Nothing {
                let mask  = (1 << max_overflow) - 1;
                let array: Vec<TableElement> = repeat(TableElement::Nothing).take(1 << max_overflow).collect();

                lut[index as usize] = TableElement::Table(mask, array);
            }

            let code = code >> TABLESIZE as usize;
            let r = max_len - len;

            for j in (0u16..1 << r as usize) {
                let k = (j << (len - TABLESIZE) as usize) + code;
                let s = TableElement::Symbol(i as u16, len - TABLESIZE);

                lut[index as usize].put(k, s);
            }
        }

        next_code[len as usize] += 1;
    }

    lut
}

struct HuffReader<R> {
    pub r: R,

    bits: u32,
    num_bits: u8,
}

impl<R: Read> HuffReader<R> {
    pub fn new(r: R) -> HuffReader<R> {
        HuffReader {r: r, bits: 0, num_bits: 0}
    }

    pub fn guarantee(&mut self, n: u8) -> io::Result<()> {
        while self.num_bits < n {
            let byte = try!(self.r.read_u8());

            self.bits |= (byte as u32) << self.num_bits as usize;
            self.num_bits += 8;
        }

        Ok(())
    }

    pub fn byte_align(&mut self) {
        let n = self.bits & 0b111;

        self.bits >>= n as usize;
        self.num_bits -= n as u8;
    }

    pub fn consume(&mut self, n: u8) {
        self.bits >>= n as usize;
        self.num_bits -= n;
    }

    pub fn receive(&mut self, n: u8) -> io::Result<u16> {
        let _ = try!(self.guarantee(n));

        let val = self.bits & ((1 << n as usize) - 1);
        self.consume(n);

        Ok(val as u16)
    }

    pub fn decode_symbol(&mut self, table: &[TableElement]) -> io::Result<u16> {
        let _ = try!(self.guarantee(1));

        loop {
            let index = self.bits & ((1 << TABLESIZE as usize) - 1);

            let (val, size) = match table[index as usize] {
                TableElement::Symbol(val, size) => (val, size),

                TableElement::Table(mask, ref a) => {
                    let index = (self.bits >> TABLESIZE as usize) & mask as u32;

                    match a[index as usize] {
                        TableElement::Symbol(val, size) => (val, size + TABLESIZE),
                        _ 		  => panic!("bad huffman code")
                    }
                }

                TableElement::Nothing => panic!("bad huffman code")
            };

            if size <= self.num_bits {
                self.consume(size);
                return	Ok(val)
            }

            let _ = try!(self.guarantee(size));
        }
    }
}
<|endoftext|>


<commit_before>extern crate num_cpus;

mod render_worker;
mod traced_image;

use std::path::Path;
use std::sync::{Arc, Mutex,
                mpsc::{self, Sender, Receiver},
                atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering},
};
use std::thread::{self, JoinHandle};

use cgmath::{Vector3, Point3};

use glium;
use glium::{VertexBuffer, IndexBuffer, Surface, DrawParameters, Rect, uniform};
use glium::texture::{Texture2d, RawImage2d};
use glium::backend::Facade;

use crate::camera::Camera;
use crate::scene::Scene;
use crate::vertex::Vertex;

use self::render_worker::RenderWorker;
use self::traced_image::TracedImage;

pub trait Intersect<'a, H> {
    fn intersect(&'a self, ray: &Ray) -> Option<H>;
}

#[derive(Clone, Copy, Debug)]
pub struct Ray {
    pub orig: Point3<f32>,
    pub dir: Vector3<f32>,
    // For more efficient ray plane intersections
    pub reciprocal_dir: Vector3<f32>,
    pub length: f32,
}

impl Ray {
    fn new(orig: Point3<f32>, dir: Vector3<f32>, length: f32) -> Ray {
        let reciprocal_dir = 1.0 / dir;
        Ray { orig, dir, reciprocal_dir, length }
    }
}

#[allow(dead_code)]
pub struct RenderCoordinator {
    width: u32,
    height: u32,
    max_iterations: Option<u32>,
    end_x: u32,
    start_y: u32,
    end_y: u32,
    iteration: u32,
}

impl RenderCoordinator {
    fn new(width: u32, height: u32, max_iterations: Option<u32>) -> RenderCoordinator {
        RenderCoordinator {
            width, height, max_iterations,
            end_x: 0, start_y: 0, end_y: 0, iteration: 1
        }
    }

    fn next_block(&mut self) -> Option<Rect> {
        let block_height = 50;
        let block_width = 50;
        if let Some(max) = self.max_iterations {
            if self.iteration > max {
                return None;
            }
        }
        if self.end_y >= self.height && self.end_x >= self.width {
            self.iteration += 1;
            self.start_y = 0;
            self.end_y = block_width.min(self.height);
            self.end_x = 0;
        } else if self.end_x >= self.width {
            self.start_y = self.end_y;
            self.end_y = (self.start_y + block_height).min(self.height);
            self.end_x = 0;
        }
        let start_x = self.end_x;
        self.end_x = (start_x + block_width).min(self.width);
        Some (
            Rect {
                left: start_x,
                bottom: self.start_y,
                width: self.end_x - start_x,
                height: self.end_y - self.start_y,
            }
        )
    }
}

struct PTVisualizer {
    shader: glium::Program,
    vertex_buffer: VertexBuffer<Vertex>,
    index_buffer: IndexBuffer<u32>,
    texture: Texture2d,
}

impl PTVisualizer {
    fn new<F: Facade>(facade: &F, width: u32, height: u32) -> PTVisualizer {
        let vertices = vec!(
            Vertex { pos: [-1.0, -1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [0.0, 0.0] },
            Vertex { pos: [1.0, -1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [1.0, 0.0] },
            Vertex { pos: [1.0, 1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [1.0, 1.0] },
            Vertex { pos: [-1.0, 1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [0.0, 1.0] },
        );
        let vertex_buffer = VertexBuffer::new(facade, &vertices)
            .expect("Failed to create vertex buffer!");
        let indices = vec!(0, 1, 2, 0, 2, 3);
        let index_buffer = IndexBuffer::new(facade,
                                            glium::index::PrimitiveType::TrianglesList,
                                            &indices)
            .expect("Failed to create index buffer!");

        // Image shader
        let vertex_shader_src = include_str!("shaders/image.vert");
        let fragment_shader_src = include_str!("shaders/image.frag");
        let shader = glium::Program::from_source(facade, vertex_shader_src, fragment_shader_src, None)
            .expect("Failed to create program!");

        let texture = Texture2d::empty(facade, width, height).unwrap();

        PTVisualizer {
            shader, vertex_buffer, index_buffer, texture,
        }
    }

    fn render<S: Surface>(&self, target: &mut S) {
        let uniforms = uniform! {
            image: &self.texture,
        };
        let draw_parameters = DrawParameters {
            ..Default::default()
        };
        target.draw(&self.vertex_buffer, &self.index_buffer, &self.shader,
                    &uniforms, &draw_parameters).unwrap();
    }

    fn new_texture<F: Facade>(&mut self, facade: &F, texture_source: RawImage2d<f32>) {
        self.texture = Texture2d::new(facade, texture_source).unwrap();
    }

    fn update_texture(&mut self, rect: Rect, texture_block: RawImage2d<f32>) {
        self.texture.write(rect, texture_block);
    }
}

pub struct PTRenderer {
    visualizer: Option<PTVisualizer>,
    image: TracedImage,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    result_rx: Option<Receiver<(Rect, Vec<f32>)>>,
    message_txs: Vec<Sender<()>>,
    thread_handles: Vec<JoinHandle<()>>,
    ray_count: Arc<AtomicUsize>,
}

impl PTRenderer {
    pub fn new() -> PTRenderer {
        let image = TracedImage::empty(0, 0);
        PTRenderer { visualizer: None,
                     image,
                     coordinator: Arc::new(Mutex::new(RenderCoordinator::new(0, 0, None))),
                     result_rx: None,
                     message_txs: Vec::new(),
                     thread_handles: Vec::new(),
                     ray_count: Arc::new(ATOMIC_USIZE_INIT),
        }
    }

    fn start_render(&mut self, scene: &Arc<Scene>, camera: &Camera, iterations: Option<u32>) {
        let width = camera.width;
        let height = camera.height;
        self.image = TracedImage::empty(width, height);
        self.coordinator = Arc::new(Mutex::new(
            RenderCoordinator::new(width, height, iterations)));
        self.ray_count.store(0, Ordering::SeqCst);

        let (result_tx, result_rx) = mpsc::channel();
        self.result_rx = Some(result_rx);
        for _ in 0..num_cpus::get_physical() {
            let result_tx = result_tx.clone();
            let (message_tx, message_rx) = mpsc::channel();
            self.message_txs.push(message_tx);
            let coordinator = self.coordinator.clone();
            let ray_count = self.ray_count.clone();
            let camera = camera.clone();
            let scene = scene.clone();
            let handle = thread::spawn(move|| {
                let worker = RenderWorker::new(scene.clone(), camera.clone(), coordinator.clone(),
                                               message_rx, result_tx, ray_count);
                worker.run();
            });
            self.thread_handles.push(handle);
        }
    }

    pub fn online_render<F: Facade>(&mut self, facade: &F, scene: &Arc<Scene>, camera: &Camera) {
        self.start_render(scene, camera, None);
        // Create visualizer after staring render so we have the new image source
        if let Some(visualizer) = &mut self.visualizer {
            visualizer.new_texture(facade, self.image.get_texture_source());
        } else {
            self.visualizer = Some(PTVisualizer::new(facade, camera.width, camera.height));
        }
    }

    pub fn offline_render(&mut self, scene: &Arc<Scene>, camera: &Camera, iterations: u32) {
        self.start_render(scene, camera, Some(iterations));

        // Wait for all the threads to finish
        for handle in self.thread_handles.drain(..) {
            handle.join().unwrap();
        }

        for (rect, block) in self.result_rx.as_ref().unwrap().try_iter() {
            self.image.update_block(rect, &block);
        }

    }

    pub fn render<S: Surface>(&mut self, target: &mut S) {
        if let Some(ref rx) = self.result_rx {
            let visualizer = self.visualizer.as_mut().unwrap();
            for (rect, block) in rx.try_iter().take(10) {
                let (rect, texture_block) = self.image.update_block(rect, &block);
                visualizer.update_texture(rect, texture_block);
            }
            visualizer.render(target);
        }
    }

    pub fn stop_threads(&mut self) {
        for sender in &self.message_txs {
            sender.send(()).unwrap();
        }
        for handle in self.thread_handles.drain(..) {
            handle.join().unwrap();
        }
        // Drop channels after join so stop messages are properly reveived
        self.message_txs.clear();
    }

    pub fn save_image(&self, path: &Path) {
        self.image.save_image(path);
    }

    pub fn get_ray_count(&self) -> usize {
        self.ray_count.load(Ordering::Relaxed)
    }
}
<commit_msg>Switch to float textures<commit_after>extern crate num_cpus;

mod render_worker;
mod traced_image;

use std::path::Path;
use std::sync::{Arc, Mutex,
                mpsc::{self, Sender, Receiver},
                atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering},
};
use std::thread::{self, JoinHandle};

use cgmath::{Vector3, Point3};

use glium;
use glium::{VertexBuffer, IndexBuffer, Surface, DrawParameters, Rect, uniform};
use glium::texture::{Texture2d, RawImage2d, MipmapsOption, UncompressedFloatFormat};
use glium::backend::Facade;

use crate::camera::Camera;
use crate::scene::Scene;
use crate::vertex::Vertex;

use self::render_worker::RenderWorker;
use self::traced_image::TracedImage;

pub trait Intersect<'a, H> {
    fn intersect(&'a self, ray: &Ray) -> Option<H>;
}

#[derive(Clone, Copy, Debug)]
pub struct Ray {
    pub orig: Point3<f32>,
    pub dir: Vector3<f32>,
    // For more efficient ray plane intersections
    pub reciprocal_dir: Vector3<f32>,
    pub length: f32,
}

impl Ray {
    fn new(orig: Point3<f32>, dir: Vector3<f32>, length: f32) -> Ray {
        let reciprocal_dir = 1.0 / dir;
        Ray { orig, dir, reciprocal_dir, length }
    }
}

#[allow(dead_code)]
pub struct RenderCoordinator {
    width: u32,
    height: u32,
    max_iterations: Option<u32>,
    end_x: u32,
    start_y: u32,
    end_y: u32,
    iteration: u32,
}

impl RenderCoordinator {
    fn new(width: u32, height: u32, max_iterations: Option<u32>) -> RenderCoordinator {
        RenderCoordinator {
            width, height, max_iterations,
            end_x: 0, start_y: 0, end_y: 0, iteration: 1
        }
    }

    fn next_block(&mut self) -> Option<Rect> {
        let block_height = 50;
        let block_width = 50;
        if let Some(max) = self.max_iterations {
            if self.iteration > max {
                return None;
            }
        }
        if self.end_y >= self.height && self.end_x >= self.width {
            self.iteration += 1;
            self.start_y = 0;
            self.end_y = block_width.min(self.height);
            self.end_x = 0;
        } else if self.end_x >= self.width {
            self.start_y = self.end_y;
            self.end_y = (self.start_y + block_height).min(self.height);
            self.end_x = 0;
        }
        let start_x = self.end_x;
        self.end_x = (start_x + block_width).min(self.width);
        Some (
            Rect {
                left: start_x,
                bottom: self.start_y,
                width: self.end_x - start_x,
                height: self.end_y - self.start_y,
            }
        )
    }
}

struct PTVisualizer {
    shader: glium::Program,
    vertex_buffer: VertexBuffer<Vertex>,
    index_buffer: IndexBuffer<u32>,
    texture: Texture2d,
}

fn create_texture<F: Facade>(facade: &F, texture_source: RawImage2d<f32>) -> Texture2d {
    Texture2d::with_format(facade, texture_source,
                           UncompressedFloatFormat::F32F32F32,
                           MipmapsOption::NoMipmap).unwrap()
}

impl PTVisualizer {
    fn new<F: Facade>(facade: &F, texture_source: RawImage2d<f32>) -> PTVisualizer {
        let vertices = vec!(
            Vertex { pos: [-1.0, -1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [0.0, 0.0] },
            Vertex { pos: [1.0, -1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [1.0, 0.0] },
            Vertex { pos: [1.0, 1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [1.0, 1.0] },
            Vertex { pos: [-1.0, 1.0, 0.0],
                     normal: [0.0, 0.0, 0.0],
                     tex_coords: [0.0, 1.0] },
        );
        let vertex_buffer = VertexBuffer::new(facade, &vertices)
            .expect("Failed to create vertex buffer!");
        let indices = vec!(0, 1, 2, 0, 2, 3);
        let index_buffer = IndexBuffer::new(facade,
                                            glium::index::PrimitiveType::TrianglesList,
                                            &indices)
            .expect("Failed to create index buffer!");

        // Image shader
        let vertex_shader_src = include_str!("shaders/image.vert");
        let fragment_shader_src = include_str!("shaders/image.frag");
        let shader = glium::Program::from_source(facade, vertex_shader_src, fragment_shader_src, None)
            .expect("Failed to create program!");

        let texture = create_texture(facade, texture_source);

        PTVisualizer {
            shader, vertex_buffer, index_buffer, texture,
        }
    }

    fn render<S: Surface>(&self, target: &mut S) {
        let uniforms = uniform! {
            image: &self.texture,
        };
        let draw_parameters = DrawParameters {
            ..Default::default()
        };
        target.draw(&self.vertex_buffer, &self.index_buffer, &self.shader,
                    &uniforms, &draw_parameters).unwrap();
    }

    fn new_texture<F: Facade>(&mut self, facade: &F, texture_source: RawImage2d<f32>) {
        self.texture = create_texture(facade, texture_source);
    }

    fn update_texture(&mut self, rect: Rect, texture_block: RawImage2d<f32>) {
        self.texture.write(rect, texture_block);
    }
}

pub struct PTRenderer {
    visualizer: Option<PTVisualizer>,
    image: TracedImage,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    result_rx: Option<Receiver<(Rect, Vec<f32>)>>,
    message_txs: Vec<Sender<()>>,
    thread_handles: Vec<JoinHandle<()>>,
    ray_count: Arc<AtomicUsize>,
}

impl PTRenderer {
    pub fn new() -> PTRenderer {
        let image = TracedImage::empty(0, 0);
        PTRenderer { visualizer: None,
                     image,
                     coordinator: Arc::new(Mutex::new(RenderCoordinator::new(0, 0, None))),
                     result_rx: None,
                     message_txs: Vec::new(),
                     thread_handles: Vec::new(),
                     ray_count: Arc::new(ATOMIC_USIZE_INIT),
        }
    }

    fn start_render(&mut self, scene: &Arc<Scene>, camera: &Camera, iterations: Option<u32>) {
        let width = camera.width;
        let height = camera.height;
        self.image = TracedImage::empty(width, height);
        self.coordinator = Arc::new(Mutex::new(
            RenderCoordinator::new(width, height, iterations)));
        self.ray_count.store(0, Ordering::SeqCst);

        let (result_tx, result_rx) = mpsc::channel();
        self.result_rx = Some(result_rx);
        for _ in 0..num_cpus::get_physical() {
            let result_tx = result_tx.clone();
            let (message_tx, message_rx) = mpsc::channel();
            self.message_txs.push(message_tx);
            let coordinator = self.coordinator.clone();
            let ray_count = self.ray_count.clone();
            let camera = camera.clone();
            let scene = scene.clone();
            let handle = thread::spawn(move|| {
                let worker = RenderWorker::new(scene.clone(), camera.clone(), coordinator.clone(),
                                               message_rx, result_tx, ray_count);
                worker.run();
            });
            self.thread_handles.push(handle);
        }
    }

    pub fn online_render<F: Facade>(&mut self, facade: &F, scene: &Arc<Scene>, camera: &Camera) {
        self.start_render(scene, camera, None);
        if let Some(visualizer) = &mut self.visualizer {
            visualizer.new_texture(facade, self.image.get_texture_source());
        } else {
            self.visualizer = Some(PTVisualizer::new(facade, self.image.get_texture_source()));
        }
    }

    pub fn offline_render(&mut self, scene: &Arc<Scene>, camera: &Camera, iterations: u32) {
        self.start_render(scene, camera, Some(iterations));

        // Wait for all the threads to finish
        for handle in self.thread_handles.drain(..) {
            handle.join().unwrap();
        }

        for (rect, block) in self.result_rx.as_ref().unwrap().try_iter() {
            self.image.update_block(rect, &block);
        }

    }

    pub fn render<S: Surface>(&mut self, target: &mut S) {
        if let Some(ref rx) = self.result_rx {
            let visualizer = self.visualizer.as_mut().unwrap();
            for (rect, block) in rx.try_iter().take(10) {
                let (rect, texture_block) = self.image.update_block(rect, &block);
                visualizer.update_texture(rect, texture_block);
            }
            visualizer.render(target);
        }
    }

    pub fn stop_threads(&mut self) {
        for sender in &self.message_txs {
            sender.send(()).unwrap();
        }
        for handle in self.thread_handles.drain(..) {
            handle.join().unwrap();
        }
        // Drop channels after join so stop messages are properly reveived
        self.message_txs.clear();
    }

    pub fn save_image(&self, path: &Path) {
        self.image.save_image(path);
    }

    pub fn get_ray_count(&self) -> usize {
        self.ray_count.load(Ordering::Relaxed)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! The purpose of KeyManager is to rotate consensus key (and eventually the network key). It is not
//! responsible for generating the first key and fails if the stores have not been properly setup.
//! During rotation, it first updates the local store, then submits a transaction to rotate to the
//! new key. After some period of time and upon restarts of the process, it will evaluate the
//! current status of the system including:
//! * last rotation time, and rotate if it is too long ago
//! * if the latest key in the store matches the latest key in the ValidatorConfig, upon mismatch
//! it will try to submit a transaction to update the ValidatorConfig to the current key in the
//! store.
//! * if the current key in the ValidatorConfig matches the ValidatorSet, if it does not it
//! evaluates the current time from the last reconfiguration and logs that delta with greater
//! levels of severity depending on the delta.
//!
//! KeyManager talks to Libra via the LibraInterface that may either be a direct link into
//! `LibraDB`/`Executor`, JSON-RPC, or some other concoction.
//! KeyManager talks to its own storage through the `LibraSecureStorage::Storage trait.
#![forbid(unsafe_code)]

use crate::{
    counters::{
        KEYS_STILL_FRESH, LIVENESS_ERROR_ENCOUNTERED, NO_ACTION, ROTATED_IN_STORAGE,
        SUBMITTED_ROTATION_TRANSACTION, UNEXPECTED_ERROR_ENCOUNTERED, WAITING_ON_RECONFIGURATION,
        WAITING_ON_TRANSACTION_EXECUTION,
    },
    libra_interface::LibraInterface,
    logging::{LogEntry, LogEvent, LogSchema},
};
use libra_crypto::ed25519::Ed25519PublicKey;
use libra_global_constants::{CONSENSUS_KEY, OPERATOR_ACCOUNT, OPERATOR_KEY, OWNER_ACCOUNT};
use libra_logger::prelude::*;
use libra_secure_storage::{CryptoStorage, KVStorage};
use libra_secure_time::TimeService;
use libra_types::{
    account_address::AccountAddress,
    account_config::LBR_NAME,
    chain_id::ChainId,
    transaction::{RawTransaction, SignedTransaction, Transaction},
};
use thiserror::Error;

pub mod counters;
pub mod libra_interface;
pub mod logging;

#[cfg(test)]
mod tests;

const GAS_UNIT_PRICE: u64 = 0;
const MAX_GAS_AMOUNT: u64 = 400_000;

/// Defines actions that KeyManager should perform after a check of all associated state.
#[derive(Debug, PartialEq)]
pub enum Action {
    /// There is no need to perform a rotation (keys are still fresh).
    NoAction,
    /// Sufficient time has passed for another key rotation (keys are stale).
    FullKeyRotation,
    /// Storage and the blockchain are inconsistent, submit a new rotation transaction.
    SubmitKeyRotationTransaction,
    /// The validator config and the validator set are inconsistent, wait for reconfiguration.
    WaitForReconfiguration,
    /// Storage and the blockchain are inconsistent, wait for rotation transaction execution.
    WaitForTransactionExecution,
}

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Error, PartialEq)]
pub enum Error {
    #[error("Key mismatch, config: {0}, info: {1}")]
    ConfigInfoKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Key mismatch, config: {0}, storage: {1}")]
    ConfigStorageKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Data does not exist: {0}")]
    DataDoesNotExist(String),
    #[error(
        "The libra_timestamp value on-chain isn't increasing. Last value: {0}, Current value: {1}"
    )]
    LivenessError(u64, u64),
    #[error("Unable to retrieve the account address: {0}, storage error: {1}")]
    MissingAccountAddress(String, String),
    #[error("Storage error: {0}")]
    StorageError(String),
    #[error("ValidatorInfo not found in ValidatorConfig: {0}")]
    ValidatorInfoNotFound(AccountAddress),
    #[error("Unknown error: {0}")]
    UnknownError(String),
}

#[cfg(test)]
impl From<anyhow::Error> for Error {
    fn from(error: anyhow::Error) -> Self {
        Error::UnknownError(format!("{}", error))
    }
}

impl From<libra_secure_storage::Error> for Error {
    fn from(error: libra_secure_storage::Error) -> Self {
        Error::StorageError(error.to_string())
    }
}

pub struct KeyManager<LI, S, T> {
    libra: LI,
    storage: S,
    time_service: T,
    last_checked_libra_timestamp: u64,
    rotation_period_secs: u64, // The frequency by which to rotate all keys
    sleep_period_secs: u64,    // The amount of time to sleep between key management checks
    txn_expiration_secs: u64,  // The time after which a rotation transaction expires
    chain_id: ChainId,
}

impl<LI, S, T> KeyManager<LI, S, T>
where
    LI: LibraInterface,
    S: KVStorage + CryptoStorage,
    T: TimeService,
{
    pub fn new(
        libra: LI,
        storage: S,
        time_service: T,
        rotation_period_secs: u64,
        sleep_period_secs: u64,
        txn_expiration_secs: u64,
        chain_id: ChainId,
    ) -> Self {
        Self {
            libra,
            storage,
            time_service,
            last_checked_libra_timestamp: 0,
            rotation_period_secs,
            sleep_period_secs,
            txn_expiration_secs,
            chain_id,
        }
    }

    /// Begins execution of the key manager by running an infinite loop where the key manager will
    /// periodically wake up, verify the state of the validator keys (e.g., the consensus key), and
    /// initiate a key rotation when required. If something goes wrong that we can't handle, an
    /// error will be returned by this method, upon which the key manager will flag the error and
    /// stop execution.
    pub fn execute(&mut self) -> Result<(), Error> {
        loop {
            info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Pending));

            match self.execute_once() {
                Ok(_) => {
                    info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Success));
                }
                Err(Error::LivenessError(last_value, current_value)) => {
                    // Log the liveness error and continue to execute.
                    let error = Error::LivenessError(last_value, current_value);
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .liveness_error(&error));
                    counters::increment_metric_counter(LIVENESS_ERROR_ENCOUNTERED);
                }
                Err(e) => {
                    // Log the unexpected error and continue to execute.
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .unexpected_error(&e));
                    counters::increment_metric_counter(UNEXPECTED_ERROR_ENCOUNTERED);
                }
            };

            self.sleep();
        }
    }

    fn sleep(&self) {
        info!(LogSchema::new(LogEntry::Sleep)
            .event(LogEvent::Pending)
            .sleep_duration(self.sleep_period_secs));
        self.time_service.sleep(self.sleep_period_secs);

        info!(LogSchema::new(LogEntry::Sleep).event(LogEvent::Success));
    }

    /// Checks the current state of the validator keys and performs any actions that might be
    /// required (e.g., performing a key rotation).
    pub fn execute_once(&mut self) -> Result<(), Error> {
        let action = self.evaluate_status()?;
        self.perform_action(action)
    }

    pub fn compare_storage_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let storage_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        let config_key = validator_config.consensus_public_key;
        if storage_key != config_key {
            return Err(Error::ConfigStorageKeyMismatch(config_key, storage_key));
        }

        Ok(())
    }

    pub fn compare_info_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;
        let validator_info = self.libra.retrieve_validator_info(owner_account)?;

        let info_key = validator_info.consensus_public_key();
        let config_key = validator_config.consensus_public_key;
        if &config_key != info_key {
            return Err(Error::ConfigInfoKeyMismatch(config_key, info_key.clone()));
        }

        Ok(())
    }

    pub fn last_reconfiguration(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.last_reconfiguration()? / 1_000_000)
    }

    pub fn last_rotation(&self) -> Result<u64, Error> {
        Ok(self.storage.get_public_key(CONSENSUS_KEY)?.last_update)
    }

    pub fn libra_timestamp(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.libra_timestamp()? / 1_000_000)
    }

    pub fn resubmit_consensus_key_transaction(&mut self) -> Result<(), Error> {
        let consensus_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        counters::increment_metric_counter(SUBMITTED_ROTATION_TRANSACTION);
        self.submit_key_rotation_transaction(consensus_key)
            .map(|_| ())
    }

    pub fn rotate_consensus_key(&mut self) -> Result<Ed25519PublicKey, Error> {
        let consensus_key = self.storage.rotate_key(CONSENSUS_KEY)?;
        info!(LogSchema::new(LogEntry::KeyRotatedInStorage)
            .event(LogEvent::Success)
            .consensus_key(&consensus_key));
        counters::increment_metric_counter(ROTATED_IN_STORAGE);
        self.submit_key_rotation_transaction(consensus_key)
    }

    pub fn submit_key_rotation_transaction(
        &mut self,
        consensus_key: Ed25519PublicKey,
    ) -> Result<Ed25519PublicKey, Error> {
        let operator_account = self.get_account_from_storage(OPERATOR_ACCOUNT)?;
        let seq_id = self.libra.retrieve_sequence_number(operator_account)?;
        let expiration = self.time_service.now() + self.txn_expiration_secs;

        // Retrieve existing network information as registered on-chain
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let txn = build_rotation_transaction(
            owner_account,
            operator_account,
            seq_id,
            &consensus_key,
            validator_config.validator_network_addresses,
            validator_config.fullnode_network_addresses,
            expiration,
            self.chain_id,
        );

        let operator_pubkey = self.storage.get_public_key(OPERATOR_KEY)?.public_key;
        let txn_signature = self.storage.sign(OPERATOR_KEY, &txn)?;
        let signed_txn = SignedTransaction::new(txn, operator_pubkey, txn_signature);

        self.libra
            .submit_transaction(Transaction::UserTransaction(signed_txn))?;
        info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Success));

        Ok(consensus_key)
    }

    /// Ensures that the libra_timestamp() value registered on-chain is strictly monotonically
    /// increasing.
    fn ensure_timestamp_progress(&mut self) -> Result<(), Error> {
        let current_libra_timestamp = self.libra.libra_timestamp()?;
        if current_libra_timestamp <= self.last_checked_libra_timestamp {
            return Err(Error::LivenessError(
                self.last_checked_libra_timestamp,
                current_libra_timestamp,
            ));
        }

        self.last_checked_libra_timestamp = current_libra_timestamp;
        Ok(())
    }

    /// Evaluates the current status of the key manager by performing various state checks between
    /// secure storage and the blockchain.
    ///
    /// Note: every time this function is called, the libra_timestamp registered on-chain must be
    /// strictly monotonically increasing. This helps to ensure that the blockchain is making
    /// progress. Otherwise, if no progress is being made on-chain, a reconfiguration event is
    /// unlikely, and the key manager will be unable to rotate keys.
    pub fn evaluate_status(&mut self) -> Result<Action, Error> {
        self.ensure_timestamp_progress()?;

        // If this is inconsistent, then we are waiting on a reconfiguration...
        if let Err(Error::ConfigInfoKeyMismatch(..)) = self.compare_info_to_config() {
            warn!(LogSchema::new(LogEntry::WaitForReconfiguration));
            counters::increment_metric_counter(WAITING_ON_RECONFIGURATION);
            return Ok(Action::WaitForReconfiguration);
        }

        let last_rotation = self.last_rotation()?;

        // If this is inconsistent, then the transaction either failed or was never submitted.
        if let Err(Error::ConfigStorageKeyMismatch(..)) = self.compare_storage_to_config() {
            return if last_rotation + self.txn_expiration_secs <= self.time_service.now() {
                Ok(Action::SubmitKeyRotationTransaction)
            } else {
                warn!(LogSchema::new(LogEntry::WaitForTransactionExecution));
                counters::increment_metric_counter(WAITING_ON_TRANSACTION_EXECUTION);
                Ok(Action::WaitForTransactionExecution)
            };
        }

        if last_rotation + self.rotation_period_secs <= self.time_service.now() {
            Ok(Action::FullKeyRotation)
        } else {
            info!(LogSchema::new(LogEntry::KeyStillFresh));
            counters::increment_metric_counter(KEYS_STILL_FRESH);
            Ok(Action::NoAction)
        }
    }

    pub fn perform_action(&mut self, action: Action) -> Result<(), Error> {
        match action {
            Action::FullKeyRotation => {
                info!(LogSchema::new(LogEntry::FullKeyRotation).event(LogEvent::Pending));
                self.rotate_consensus_key().map(|_| ())
            }
            Action::SubmitKeyRotationTransaction => {
                info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Pending));
                self.resubmit_consensus_key_transaction()
            }
            _ => {
                info!(LogSchema::new(LogEntry::NoAction));
                counters::increment_metric_counter(NO_ACTION);
                Ok(())
            }
        }
    }

    fn get_account_from_storage(&self, account_name: &str) -> Result<AccountAddress, Error> {
        self.storage
            .get::<AccountAddress>(account_name)
            .map(|v| v.value)
            .map_err(|e| Error::MissingAccountAddress(account_name.into(), e.to_string()))
    }
}

pub fn build_rotation_transaction(
    owner_address: AccountAddress,
    operator_address: AccountAddress,
    seq_id: u64,
    consensus_key: &Ed25519PublicKey,
    network_addresses: Vec<u8>,
    fullnode_network_addresses: Vec<u8>,
    expiration_timestamp_secs: u64,
    chain_id: ChainId,
) -> RawTransaction {
    let script =
        transaction_builder_generated::stdlib::encode_set_validator_config_and_reconfigure_script(
            owner_address,
            consensus_key.to_bytes().to_vec(),
            network_addresses,
            fullnode_network_addresses,
        );
    RawTransaction::new_script(
        operator_address,
        seq_id,
        script,
        MAX_GAS_AMOUNT,
        GAS_UNIT_PRICE,
        LBR_NAME.to_owned(),
        expiration_timestamp_secs,
        chain_id,
    )
}
<commit_msg>[Key Manager] Don't drop unexpected errors on evaluate_status().<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! The purpose of KeyManager is to rotate consensus key (and eventually the network key). It is not
//! responsible for generating the first key and fails if the stores have not been properly setup.
//! During rotation, it first updates the local store, then submits a transaction to rotate to the
//! new key. After some period of time and upon restarts of the process, it will evaluate the
//! current status of the system including:
//! * last rotation time, and rotate if it is too long ago
//! * if the latest key in the store matches the latest key in the ValidatorConfig, upon mismatch
//! it will try to submit a transaction to update the ValidatorConfig to the current key in the
//! store.
//! * if the current key in the ValidatorConfig matches the ValidatorSet, if it does not it
//! evaluates the current time from the last reconfiguration and logs that delta with greater
//! levels of severity depending on the delta.
//!
//! KeyManager talks to Libra via the LibraInterface that may either be a direct link into
//! `LibraDB`/`Executor`, JSON-RPC, or some other concoction.
//! KeyManager talks to its own storage through the `LibraSecureStorage::Storage trait.
#![forbid(unsafe_code)]

use crate::{
    counters::{
        KEYS_STILL_FRESH, LIVENESS_ERROR_ENCOUNTERED, NO_ACTION, ROTATED_IN_STORAGE,
        SUBMITTED_ROTATION_TRANSACTION, UNEXPECTED_ERROR_ENCOUNTERED, WAITING_ON_RECONFIGURATION,
        WAITING_ON_TRANSACTION_EXECUTION,
    },
    libra_interface::LibraInterface,
    logging::{LogEntry, LogEvent, LogSchema},
};
use libra_crypto::ed25519::Ed25519PublicKey;
use libra_global_constants::{CONSENSUS_KEY, OPERATOR_ACCOUNT, OPERATOR_KEY, OWNER_ACCOUNT};
use libra_logger::prelude::*;
use libra_secure_storage::{CryptoStorage, KVStorage};
use libra_secure_time::TimeService;
use libra_types::{
    account_address::AccountAddress,
    account_config::LBR_NAME,
    chain_id::ChainId,
    transaction::{RawTransaction, SignedTransaction, Transaction},
};
use thiserror::Error;

pub mod counters;
pub mod libra_interface;
pub mod logging;

#[cfg(test)]
mod tests;

const GAS_UNIT_PRICE: u64 = 0;
const MAX_GAS_AMOUNT: u64 = 400_000;

/// Defines actions that KeyManager should perform after a check of all associated state.
#[derive(Debug, PartialEq)]
pub enum Action {
    /// There is no need to perform a rotation (keys are still fresh).
    NoAction,
    /// Sufficient time has passed for another key rotation (keys are stale).
    FullKeyRotation,
    /// Storage and the blockchain are inconsistent, submit a new rotation transaction.
    SubmitKeyRotationTransaction,
    /// The validator config and the validator set are inconsistent, wait for reconfiguration.
    WaitForReconfiguration,
    /// Storage and the blockchain are inconsistent, wait for rotation transaction execution.
    WaitForTransactionExecution,
}

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Error, PartialEq)]
pub enum Error {
    #[error("Key mismatch, config: {0}, info: {1}")]
    ConfigInfoKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Key mismatch, config: {0}, storage: {1}")]
    ConfigStorageKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Data does not exist: {0}")]
    DataDoesNotExist(String),
    #[error(
        "The libra_timestamp value on-chain isn't increasing. Last value: {0}, Current value: {1}"
    )]
    LivenessError(u64, u64),
    #[error("Unable to retrieve the account address: {0}, storage error: {1}")]
    MissingAccountAddress(String, String),
    #[error("Storage error: {0}")]
    StorageError(String),
    #[error("ValidatorInfo not found in ValidatorConfig: {0}")]
    ValidatorInfoNotFound(AccountAddress),
    #[error("Unknown error: {0}")]
    UnknownError(String),
}

#[cfg(test)]
impl From<anyhow::Error> for Error {
    fn from(error: anyhow::Error) -> Self {
        Error::UnknownError(format!("{}", error))
    }
}

impl From<libra_secure_storage::Error> for Error {
    fn from(error: libra_secure_storage::Error) -> Self {
        Error::StorageError(error.to_string())
    }
}

pub struct KeyManager<LI, S, T> {
    libra: LI,
    storage: S,
    time_service: T,
    last_checked_libra_timestamp: u64,
    rotation_period_secs: u64, // The frequency by which to rotate all keys
    sleep_period_secs: u64,    // The amount of time to sleep between key management checks
    txn_expiration_secs: u64,  // The time after which a rotation transaction expires
    chain_id: ChainId,
}

impl<LI, S, T> KeyManager<LI, S, T>
where
    LI: LibraInterface,
    S: KVStorage + CryptoStorage,
    T: TimeService,
{
    pub fn new(
        libra: LI,
        storage: S,
        time_service: T,
        rotation_period_secs: u64,
        sleep_period_secs: u64,
        txn_expiration_secs: u64,
        chain_id: ChainId,
    ) -> Self {
        Self {
            libra,
            storage,
            time_service,
            last_checked_libra_timestamp: 0,
            rotation_period_secs,
            sleep_period_secs,
            txn_expiration_secs,
            chain_id,
        }
    }

    /// Begins execution of the key manager by running an infinite loop where the key manager will
    /// periodically wake up, verify the state of the validator keys (e.g., the consensus key), and
    /// initiate a key rotation when required. If something goes wrong that we can't handle, an
    /// error will be returned by this method, upon which the key manager will flag the error and
    /// stop execution.
    pub fn execute(&mut self) -> Result<(), Error> {
        loop {
            info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Pending));

            match self.execute_once() {
                Ok(_) => {
                    info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Success));
                }
                Err(Error::LivenessError(last_value, current_value)) => {
                    // Log the liveness error and continue to execute.
                    let error = Error::LivenessError(last_value, current_value);
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .liveness_error(&error));
                    counters::increment_metric_counter(LIVENESS_ERROR_ENCOUNTERED);
                }
                Err(e) => {
                    // Log the unexpected error and continue to execute.
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .unexpected_error(&e));
                    counters::increment_metric_counter(UNEXPECTED_ERROR_ENCOUNTERED);
                }
            };

            self.sleep();
        }
    }

    fn sleep(&self) {
        info!(LogSchema::new(LogEntry::Sleep)
            .event(LogEvent::Pending)
            .sleep_duration(self.sleep_period_secs));
        self.time_service.sleep(self.sleep_period_secs);

        info!(LogSchema::new(LogEntry::Sleep).event(LogEvent::Success));
    }

    /// Checks the current state of the validator keys and performs any actions that might be
    /// required (e.g., performing a key rotation).
    pub fn execute_once(&mut self) -> Result<(), Error> {
        let action = self.evaluate_status()?;
        self.perform_action(action)
    }

    pub fn compare_storage_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let storage_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        let config_key = validator_config.consensus_public_key;
        if storage_key != config_key {
            return Err(Error::ConfigStorageKeyMismatch(config_key, storage_key));
        }

        Ok(())
    }

    pub fn compare_info_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;
        let validator_info = self.libra.retrieve_validator_info(owner_account)?;

        let info_key = validator_info.consensus_public_key();
        let config_key = validator_config.consensus_public_key;
        if &config_key != info_key {
            return Err(Error::ConfigInfoKeyMismatch(config_key, info_key.clone()));
        }

        Ok(())
    }

    pub fn last_reconfiguration(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.last_reconfiguration()? / 1_000_000)
    }

    pub fn last_rotation(&self) -> Result<u64, Error> {
        Ok(self.storage.get_public_key(CONSENSUS_KEY)?.last_update)
    }

    pub fn libra_timestamp(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.libra_timestamp()? / 1_000_000)
    }

    pub fn resubmit_consensus_key_transaction(&mut self) -> Result<(), Error> {
        let consensus_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        counters::increment_metric_counter(SUBMITTED_ROTATION_TRANSACTION);
        self.submit_key_rotation_transaction(consensus_key)
            .map(|_| ())
    }

    pub fn rotate_consensus_key(&mut self) -> Result<Ed25519PublicKey, Error> {
        let consensus_key = self.storage.rotate_key(CONSENSUS_KEY)?;
        info!(LogSchema::new(LogEntry::KeyRotatedInStorage)
            .event(LogEvent::Success)
            .consensus_key(&consensus_key));
        counters::increment_metric_counter(ROTATED_IN_STORAGE);
        self.submit_key_rotation_transaction(consensus_key)
    }

    pub fn submit_key_rotation_transaction(
        &mut self,
        consensus_key: Ed25519PublicKey,
    ) -> Result<Ed25519PublicKey, Error> {
        let operator_account = self.get_account_from_storage(OPERATOR_ACCOUNT)?;
        let seq_id = self.libra.retrieve_sequence_number(operator_account)?;
        let expiration = self.time_service.now() + self.txn_expiration_secs;

        // Retrieve existing network information as registered on-chain
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let txn = build_rotation_transaction(
            owner_account,
            operator_account,
            seq_id,
            &consensus_key,
            validator_config.validator_network_addresses,
            validator_config.fullnode_network_addresses,
            expiration,
            self.chain_id,
        );

        let operator_pubkey = self.storage.get_public_key(OPERATOR_KEY)?.public_key;
        let txn_signature = self.storage.sign(OPERATOR_KEY, &txn)?;
        let signed_txn = SignedTransaction::new(txn, operator_pubkey, txn_signature);

        self.libra
            .submit_transaction(Transaction::UserTransaction(signed_txn))?;
        info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Success));

        Ok(consensus_key)
    }

    /// Ensures that the libra_timestamp() value registered on-chain is strictly monotonically
    /// increasing.
    fn ensure_timestamp_progress(&mut self) -> Result<(), Error> {
        let current_libra_timestamp = self.libra.libra_timestamp()?;
        if current_libra_timestamp <= self.last_checked_libra_timestamp {
            return Err(Error::LivenessError(
                self.last_checked_libra_timestamp,
                current_libra_timestamp,
            ));
        }

        self.last_checked_libra_timestamp = current_libra_timestamp;
        Ok(())
    }

    /// Evaluates the current status of the key manager by performing various state checks between
    /// secure storage and the blockchain.
    ///
    /// Note: every time this function is called, the libra_timestamp registered on-chain must be
    /// strictly monotonically increasing. This helps to ensure that the blockchain is making
    /// progress. Otherwise, if no progress is being made on-chain, a reconfiguration event is
    /// unlikely, and the key manager will be unable to rotate keys.
    pub fn evaluate_status(&mut self) -> Result<Action, Error> {
        self.ensure_timestamp_progress()?;

        // Compare the validator config to the validator set
        match self.compare_info_to_config() {
            Ok(()) => { /* Expected */ }
            Err(Error::ConfigInfoKeyMismatch(..)) => {
                warn!(LogSchema::new(LogEntry::WaitForReconfiguration));
                counters::increment_metric_counter(WAITING_ON_RECONFIGURATION);
                return Ok(Action::WaitForReconfiguration);
            }
            Err(e) => return Err(e),
        }

        let last_rotation = self.last_rotation()?;

        // Compare the validator config to secure storage
        match self.compare_storage_to_config() {
            Ok(()) => { /* Expected */ }
            Err(Error::ConfigStorageKeyMismatch(..)) => {
                return if last_rotation + self.txn_expiration_secs <= self.time_service.now() {
                    Ok(Action::SubmitKeyRotationTransaction)
                } else {
                    warn!(LogSchema::new(LogEntry::WaitForTransactionExecution));
                    counters::increment_metric_counter(WAITING_ON_TRANSACTION_EXECUTION);
                    Ok(Action::WaitForTransactionExecution)
                };
            }
            Err(e) => return Err(e),
        };

        if last_rotation + self.rotation_period_secs <= self.time_service.now() {
            Ok(Action::FullKeyRotation)
        } else {
            info!(LogSchema::new(LogEntry::KeyStillFresh));
            counters::increment_metric_counter(KEYS_STILL_FRESH);
            Ok(Action::NoAction)
        }
    }

    pub fn perform_action(&mut self, action: Action) -> Result<(), Error> {
        match action {
            Action::FullKeyRotation => {
                info!(LogSchema::new(LogEntry::FullKeyRotation).event(LogEvent::Pending));
                self.rotate_consensus_key().map(|_| ())
            }
            Action::SubmitKeyRotationTransaction => {
                info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Pending));
                self.resubmit_consensus_key_transaction()
            }
            _ => {
                info!(LogSchema::new(LogEntry::NoAction));
                counters::increment_metric_counter(NO_ACTION);
                Ok(())
            }
        }
    }

    fn get_account_from_storage(&self, account_name: &str) -> Result<AccountAddress, Error> {
        self.storage
            .get::<AccountAddress>(account_name)
            .map(|v| v.value)
            .map_err(|e| Error::MissingAccountAddress(account_name.into(), e.to_string()))
    }
}

pub fn build_rotation_transaction(
    owner_address: AccountAddress,
    operator_address: AccountAddress,
    seq_id: u64,
    consensus_key: &Ed25519PublicKey,
    network_addresses: Vec<u8>,
    fullnode_network_addresses: Vec<u8>,
    expiration_timestamp_secs: u64,
    chain_id: ChainId,
) -> RawTransaction {
    let script =
        transaction_builder_generated::stdlib::encode_set_validator_config_and_reconfigure_script(
            owner_address,
            consensus_key.to_bytes().to_vec(),
            network_addresses,
            fullnode_network_addresses,
        );
    RawTransaction::new_script(
        operator_address,
        seq_id,
        script,
        MAX_GAS_AMOUNT,
        GAS_UNIT_PRICE,
        LBR_NAME.to_owned(),
        expiration_timestamp_secs,
        chain_id,
    )
}
<|endoftext|>


<commit_before>use redox::Vec;

pub struct AvlNode<T> {
    value: T,
    left: Option<AvlNodeId>, // ID for left node
    right: Option<AvlNodeId>, // ID for right node
}

#[derive(Copy, Clone)]
pub struct AvlNodeId {
    index: usize,
    time_stamp: u64,
}

impl AvlNodeId {
    pub fn get<'a, T>(&self, avl: &'a Avl<T>) -> &'a AvlNode<T> {
        let ref slot = avl.nodes[self.index];
        if slot.time_stamp == self.time_stamp {
            slot.node.as_ref().unwrap()
        } else {
            panic!("AvlNodeId had invalid time_stamp");
        }
    }

    pub fn try_get<'a, T>(&self, avl: &'a Avl<T>) -> Option<&'a AvlNode<T>> {
        avl.nodes
           .get(self.index)
           .and_then(|slot| {
               if slot.time_stamp == self.time_stamp {
                   slot.node.as_ref()
               } else {
                   None
               }
           })
    }

    pub fn get_mut<'a, T>(&self, avl: &'a mut Avl<T>) -> &'a mut AvlNode<T> {
        let ref mut slot = avl.nodes[self.index];
        if slot.time_stamp == self.time_stamp {
            slot.node.as_mut().unwrap()
        } else {
            panic!("AvlNodeId had invalid time_stamp");
        }
    }

    pub fn try_get_mut<'a, T>(&self, avl: &'a mut Avl<T>) -> Option<&'a mut AvlNode<T>> {
        avl.nodes
           .get_mut(self.index)
           .and_then(|slot| {
               if slot.time_stamp == self.time_stamp {
                   slot.node.as_mut()
               } else {
                   None
               }
           })
    }
}

pub struct Avl<T> {
    root: usize, // Index of the root node
    nodes: Vec<AvlSlot<T>>,
    free_list: Vec<usize>,
}

impl<T> Avl<T> {
    pub fn new() -> Self {
        Avl {
            root: 0,
            nodes: Vec::new(),
            free_list: Vec::new(),
        }
    }

    pub fn insert(&mut self, value: T) -> AvlNodeId {
        // TODO this is just a placeholder, we need to deal with all the fancy rotation stuff that
        // AVL trees do
        self.allocate_node(value)
    }

    // Performs a left rotation on a tree/subtree.
    // Returns the replace the specified node with
    fn rotate_left(&mut self, node: AvlNodeId) -> AvlNodeId {
        // Keep track of the original node positions
        // For a rotate left, the right child node must exist
        let r = node.get(self).right.unwrap();
        let rl = r.get(self).left;

        let ret = r;
        node.get_mut(self).right = rl;
        ret.get_mut(self).left = Some(node);

        ret
    }

    // Performs a right rotation on a tree/subtree.
    // Returns the replace the specified node with
    fn rotate_right(&mut self, node: AvlNodeId) -> AvlNodeId {
        // Keep track of the original node positions
        // For a rotate right, the left child node must exist
        let l = node.get(self).left.unwrap();
        let lr = l.get(self).right;

        let ret = l;
        node.get_mut(self).left = lr;
        ret.get_mut(self).right = Some(node);

        ret
    }

    // performs a left-right double rotation on a tree/subtree.
    fn rotate_leftright(&mut self, node: AvlNodeId) -> AvlNodeId {
        let l = node.get(self).left.unwrap();
        let new_l = self.rotate_left(l); // Left node needs to exist
        node.get_mut(self).left = Some(new_l);
        self.rotate_right(node)
    }

    // performs a right-left double rotation on a tree/subtree.
    fn rotate_rightleft(&mut self, node: AvlNodeId) -> AvlNodeId {
        let r = node.get(self).right.unwrap();
        let new_r = self.rotate_right(r); // Right node needs to exist
        node.get_mut(self).right = Some(new_r);
        self.rotate_left(node)
    }

    // _ins is the implementation of the binary tree insert function. Lesser values will be stored on
    // the left, while greater values will be stored on the right. No duplicates are allowed.
    /*fn _ins(&mut self, node_index: Option<AvlNodeId>, value: T) -> AvlNodeId {
        let node =
            match node_index {
                Some(node) => {
                    // Node exists, check which way to branch.
                    if n == node->val {
                        return node;
                    else if (n < node->val)
                        _ins(n, node->left);
                    else if (n > node->val)
                        _ins(n, node->right);
                },
                None => {
                    // The node doesn't exist, create it here.
                    self.allocate_node(value)
                },
            };

        rebalance(node)
    }*/

    // _rebalance rebalances the provided node
    fn rebalance(&mut self, node: AvlNodeId) -> AvlNodeId {
        let balance = self.height(node.get(self).left) - self.height(node.get(self).right);
        if balance == 2 { // left
            let lbalance = self.height(node.get(self).left.unwrap().get(self).left) -
                           self.height(node.get(self).left.unwrap().get(self).right);
            if lbalance == 0 || lbalance == 1 { // left left - need to rotate right
                return self.rotate_right(node);
            } else if lbalance == -1 { // left right
                return self.rotate_leftright(node); // function name is just a coincidence
            }
        } else if balance == -2 { // right
            let rbalance = self.height(node.get(self).right.unwrap().get(self).left) -
                           self.height(node.get(self).right.unwrap().get(self).right);
            if rbalance == 1 { // right left
                return self.rotate_rightleft(node); // function name is just a coincidence
            } else if rbalance == 0 || rbalance == -1 { // right right - need to rotate left
                return self.rotate_left(node);
            }
        }

        node
    }

    // height gets the height of a tree or subtree
    fn height(&self, node: Option<AvlNodeId>) -> i64 {
        match node {
            Some(node) => {
                let left_height = self.height(node.get(self).left);
                let right_height = self.height(node.get(self).right);

                if left_height > right_height {
                    left_height+1
                } else {
                    right_height+1
                }
            },
            None => { -1 },
        }
    }

    fn allocate_node(&mut self, value: T) -> AvlNodeId {
        match self.free_list.pop() {
            Some(index) => {
                AvlNodeId { time_stamp: self.nodes[index].time_stamp+1, index: index }
            },
            None => {
                // No free slots, create a new one
                let id = AvlNodeId { index: self.nodes.len(), time_stamp: 0 };
                self.nodes.push(AvlSlot { time_stamp: 0,
                                          node: Some(AvlNode { value: value, left: None, right: None }) });
                id
            },
        }
    }

    fn free_node(&mut self, id: AvlNodeId) -> AvlNode<T> {
        self.free_list.push(id.index);

        // NOTE: We unwrap here, because we trust that `id` points to a valid node, because
        // only we can create and free AvlNodes and their AvlNodeIds
        self.nodes[id.index].node.take().unwrap()
    }
}

struct AvlSlot<T> {
    time_stamp: u64,
    node: Option<AvlNode<T>>,
}
<commit_msg>AvlTree::insert compiles...<commit_after>use redox::Vec;

pub struct AvlNode<T> {
    value: T,
    left: Option<AvlNodeId>, // ID for left node
    right: Option<AvlNodeId>, // ID for right node
}

#[derive(Copy, Clone)]
pub struct AvlNodeId {
    index: usize,
    time_stamp: u64,
}

impl AvlNodeId {
    pub fn get<'a, T: PartialOrd>(&self, avl: &'a Avl<T>) -> &'a AvlNode<T> {
        let ref slot = avl.nodes[self.index];
        if slot.time_stamp == self.time_stamp {
            slot.node.as_ref().unwrap()
        } else {
            panic!("AvlNodeId had invalid time_stamp");
        }
    }

    pub fn try_get<'a, T: PartialOrd>(&self, avl: &'a Avl<T>) -> Option<&'a AvlNode<T>> {
        avl.nodes
           .get(self.index)
           .and_then(|slot| {
               if slot.time_stamp == self.time_stamp {
                   slot.node.as_ref()
               } else {
                   None
               }
           })
    }

    pub fn get_mut<'a, T: PartialOrd>(&self, avl: &'a mut Avl<T>) -> &'a mut AvlNode<T> {
        let ref mut slot = avl.nodes[self.index];
        if slot.time_stamp == self.time_stamp {
            slot.node.as_mut().unwrap()
        } else {
            panic!("AvlNodeId had invalid time_stamp");
        }
    }

    pub fn try_get_mut<'a, T: PartialOrd>(&self, avl: &'a mut Avl<T>) -> Option<&'a mut AvlNode<T>> {
        avl.nodes
           .get_mut(self.index)
           .and_then(|slot| {
               if slot.time_stamp == self.time_stamp {
                   slot.node.as_mut()
               } else {
                   None
               }
           })
    }
}

pub struct Avl<T: PartialOrd> {
    root: usize, // Index of the root node
    nodes: Vec<AvlSlot<T>>,
    free_list: Vec<usize>,
}

impl<T: PartialOrd> Avl<T> {
    pub fn new() -> Self {
        Avl {
            root: 0,
            nodes: Vec::new(),
            free_list: Vec::new(),
        }
    }

    // Inserts a value into the tree, keeping it balanced. Lesser values will be stored on
    // the left, while greater values will be stored on the right. No duplicates are allowed.
    fn insert(&mut self, value: T, node_index: Option<AvlNodeId>) -> AvlNodeId {
        let node =
            match node_index {
                Some(node) => {
                    // Node exists, check which way to branch.
                    if value == node.get(self).value {
                        return node;
                    } else if value < node.get(self).value {
                        let l = node.get(self).left;
                        node.get_mut(self).left = Some(self.insert(value, l));
                    } else if value > node.get(self).value {
                        let r = node.get(self).right;
                        node.get_mut(self).right = Some(self.insert(value, r));
                    }

                    node
                },
                None => {
                    // The node doesn't exist, create it here.
                    self.allocate_node(value)
                },
            };

        self.rebalance(node)
    }

    // Performs a left rotation on a tree/subtree.
    // Returns the replace the specified node with
    fn rotate_left(&mut self, node: AvlNodeId) -> AvlNodeId {
        // Keep track of the original node positions
        // For a rotate left, the right child node must exist
        let r = node.get(self).right.unwrap();
        let rl = r.get(self).left;

        let ret = r;
        node.get_mut(self).right = rl;
        ret.get_mut(self).left = Some(node);

        ret
    }

    // Performs a right rotation on a tree/subtree.
    // Returns the replace the specified node with
    fn rotate_right(&mut self, node: AvlNodeId) -> AvlNodeId {
        // Keep track of the original node positions
        // For a rotate right, the left child node must exist
        let l = node.get(self).left.unwrap();
        let lr = l.get(self).right;

        let ret = l;
        node.get_mut(self).left = lr;
        ret.get_mut(self).right = Some(node);

        ret
    }

    // performs a left-right double rotation on a tree/subtree.
    fn rotate_leftright(&mut self, node: AvlNodeId) -> AvlNodeId {
        let l = node.get(self).left.unwrap();
        let new_l = self.rotate_left(l); // Left node needs to exist
        node.get_mut(self).left = Some(new_l);
        self.rotate_right(node)
    }

    // performs a right-left double rotation on a tree/subtree.
    fn rotate_rightleft(&mut self, node: AvlNodeId) -> AvlNodeId {
        let r = node.get(self).right.unwrap();
        let new_r = self.rotate_right(r); // Right node needs to exist
        node.get_mut(self).right = Some(new_r);
        self.rotate_left(node)
    }

    // _rebalance rebalances the provided node
    fn rebalance(&mut self, node: AvlNodeId) -> AvlNodeId {
        let balance = self.height(node.get(self).left) - self.height(node.get(self).right);
        if balance == 2 { // left
            let lbalance = self.height(node.get(self).left.unwrap().get(self).left) -
                           self.height(node.get(self).left.unwrap().get(self).right);
            if lbalance == 0 || lbalance == 1 { // left left - need to rotate right
                return self.rotate_right(node);
            } else if lbalance == -1 { // left right
                return self.rotate_leftright(node); // function name is just a coincidence
            }
        } else if balance == -2 { // right
            let rbalance = self.height(node.get(self).right.unwrap().get(self).left) -
                           self.height(node.get(self).right.unwrap().get(self).right);
            if rbalance == 1 { // right left
                return self.rotate_rightleft(node); // function name is just a coincidence
            } else if rbalance == 0 || rbalance == -1 { // right right - need to rotate left
                return self.rotate_left(node);
            }
        }

        node
    }

    // height gets the height of a tree or subtree
    fn height(&self, node: Option<AvlNodeId>) -> i64 {
        match node {
            Some(node) => {
                let left_height = self.height(node.get(self).left);
                let right_height = self.height(node.get(self).right);

                if left_height > right_height {
                    left_height+1
                } else {
                    right_height+1
                }
            },
            None => { -1 },
        }
    }

    fn allocate_node(&mut self, value: T) -> AvlNodeId {
        match self.free_list.pop() {
            Some(index) => {
                AvlNodeId { time_stamp: self.nodes[index].time_stamp+1, index: index }
            },
            None => {
                // No free slots, create a new one
                let id = AvlNodeId { index: self.nodes.len(), time_stamp: 0 };
                self.nodes.push(AvlSlot { time_stamp: 0,
                                          node: Some(AvlNode { value: value, left: None, right: None }) });
                id
            },
        }
    }

    fn free_node(&mut self, id: AvlNodeId) -> AvlNode<T> {
        self.free_list.push(id.index);

        // NOTE: We unwrap here, because we trust that `id` points to a valid node, because
        // only we can create and free AvlNodes and their AvlNodeIds
        self.nodes[id.index].node.take().unwrap()
    }
}

struct AvlSlot<T> {
    time_stamp: u64,
    node: Option<AvlNode<T>>,
}
<|endoftext|>