use std::cmp::min;
use crypto::{Hash, PublicKey, HASH_SIZE};
use super::super::StorageKey;
pub const BRANCH_KEY_PREFIX: u8 = 0;
pub const LEAF_KEY_PREFIX: u8 = 1;
pub const KEY_SIZE: usize = HASH_SIZE;
pub const PROOF_PATH_SIZE: usize = KEY_SIZE + 2;
pub const PROOF_PATH_KIND_POS: usize = 0;
pub const PROOF_PATH_LEN_POS: usize = KEY_SIZE + 1;
pub trait ProofMapKey: StorageKey {}
impl ProofMapKey for Hash {}
impl ProofMapKey for PublicKey {}
impl ProofMapKey for [u8; KEY_SIZE] {}
impl StorageKey for [u8; KEY_SIZE] {
fn size(&self) -> usize {
KEY_SIZE
}
fn write(&self, buffer: &mut [u8]) {
buffer.copy_from_slice(self.as_ref())
}
fn read(buffer: &[u8]) -> Self {
let mut value = [0; KEY_SIZE];
value.copy_from_slice(buffer);
value
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ChildKind {
Left,
Right,
}
impl ::std::ops::Not for ChildKind {
type Output = ChildKind;
fn not(self) -> ChildKind {
match self {
ChildKind::Left => ChildKind::Right,
ChildKind::Right => ChildKind::Left,
}
}
}
#[derive(Copy, Clone)]
pub struct ProofPath {
bytes: [u8; PROOF_PATH_SIZE],
start: u16,
}
impl ProofPath {
pub fn new<K: ProofMapKey>(key: &K) -> ProofPath {
debug_assert_eq!(key.size(), KEY_SIZE);
let mut data = [0; PROOF_PATH_SIZE];
data[0] = LEAF_KEY_PREFIX;
key.write(&mut data[1..KEY_SIZE + 1]);
data[PROOF_PATH_LEN_POS] = 0;
ProofPath::from_raw(data)
}
pub fn is_leaf(&self) -> bool {
self.bytes[0] == LEAF_KEY_PREFIX
}
pub fn as_bytes(&self) -> &[u8] {
&self.bytes
}
fn from_raw(raw: [u8; PROOF_PATH_SIZE]) -> ProofPath {
debug_assert!(
(raw[PROOF_PATH_KIND_POS] != LEAF_KEY_PREFIX) || (raw[PROOF_PATH_LEN_POS] == 0),
"ProofPath is inconsistent"
);
ProofPath {
bytes: raw,
start: 0,
}
}
fn set_end(&mut self, end: Option<u8>) {
if let Some(pos) = end {
self.bytes[0] = BRANCH_KEY_PREFIX;
self.bytes[PROOF_PATH_LEN_POS] = pos as u8;
} else {
self.bytes[0] = LEAF_KEY_PREFIX;
self.bytes[PROOF_PATH_LEN_POS] = 0;
};
}
}
pub(crate) trait BitsRange {
fn start(&self) -> u16;
fn end(&self) -> u16;
fn len(&self) -> u16 {
self.end() - self.start()
}
fn is_empty(&self) -> bool {
self.end() == self.start()
}
fn bit(&self, idx: u16) -> ChildKind {
debug_assert!(self.start() + idx < self.end());
let pos = self.start() + idx;
let chunk = self.raw_key()[(pos / 8) as usize];
let bit = pos % 8;
let value = (1 << bit) & chunk;
if value != 0 {
ChildKind::Right
} else {
ChildKind::Left
}
}
fn start_from(&self, pos: u16) -> Self;
fn prefix(&self, len: u16) -> Self;
fn suffix(&self, len: u16) -> Self;
fn starts_with(&self, other: &Self) -> bool {
self.common_prefix_len(other) == other.len()
}
fn raw_key(&self) -> &[u8];
fn common_prefix_len(&self, other: &Self) -> u16 {
if self.start() != other.start() {
0
} else {
let from = self.start() / 8;
let to = min((self.end() + 7) / 8, (other.end() + 7) / 8);
let max_len = min(self.len(), other.len());
for i in from..to {
let x = self.raw_key()[i as usize] ^ other.raw_key()[i as usize];
if x != 0 {
let tail = x.trailing_zeros() as u16;
return min(i * 8 + tail - self.start(), max_len);
}
}
max_len
}
}
}
impl BitsRange for ProofPath {
fn start(&self) -> u16 {
self.start
}
fn end(&self) -> u16 {
if self.is_leaf() {
KEY_SIZE as u16 * 8
} else {
u16::from(self.bytes[PROOF_PATH_LEN_POS])
}
}
fn start_from(&self, pos: u16) -> Self {
debug_assert!(pos <= self.end());
let mut key = ProofPath::from_raw(self.bytes);
key.start = pos;
key
}
fn prefix(&self, len: u16) -> Self {
let end = self.start + len;
let key_len = self.raw_key().len() as u16 * 8;
debug_assert!(end < key_len);
let mut key = ProofPath::from_raw(self.bytes);
key.start = self.start;
key.set_end(Some(end as u8));
key
}
fn suffix(&self, len: u16) -> Self {
self.start_from(self.start() + len)
}
fn raw_key(&self) -> &[u8] {
&self.bytes[1..KEY_SIZE + 1]
}
}
impl PartialEq for ProofPath {
fn eq(&self, other: &Self) -> bool {
self.len() == other.len() && self.starts_with(other)
}
}
impl ::std::fmt::Debug for ProofPath {
fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
let mut bits = String::with_capacity(KEY_SIZE * 9);
for byte in 0..self.raw_key().len() {
let chunk = self.raw_key()[byte];
for bit in (0..8).rev() {
let i = (byte * 8 + bit) as u16;
if i < self.start() || i >= self.end() {
bits.push('_');
} else {
bits.push(if (1 << bit) & chunk == 0 { '0' } else { '1' });
}
}
bits.push('|');
}
f.debug_struct("ProofPath")
.field("start", &self.start())
.field("end", &self.end())
.field("bits", &bits)
.finish()
}
}
impl StorageKey for ProofPath {
fn size(&self) -> usize {
PROOF_PATH_SIZE
}
fn write(&self, buffer: &mut [u8]) {
buffer.copy_from_slice(&self.bytes);
if !self.is_leaf() {
let right = (self.end() as usize + 7) / 8;
if self.end() % 8 != 0 {
buffer[right] &= !(255u8 << (self.end() % 8));
}
for i in buffer.iter_mut().take(KEY_SIZE + 1).skip(right + 1) {
*i = 0
}
}
}
fn read(buffer: &[u8]) -> Self::Owned {
debug_assert_eq!(buffer.len(), PROOF_PATH_SIZE);
let mut data = [0; PROOF_PATH_SIZE];
data.copy_from_slice(buffer);
ProofPath::from_raw(data)
}
}
#[test]
fn test_proof_path_storage_key_leaf() {
let key = ProofPath::new(&[250; 32]);
let mut buf = vec![0; PROOF_PATH_SIZE];
key.write(&mut buf);
let key2 = ProofPath::read(&buf);
assert_eq!(buf[0], LEAF_KEY_PREFIX);
assert_eq!(buf[33], 0);
assert_eq!(&buf[1..33], &[250; 32]);
assert_eq!(key2, key);
}
#[test]
fn test_proof_path_storage_key_branch() {
let mut key = ProofPath::new(&[255u8; 32]);
key = key.prefix(11);
key = key.suffix(5);
let mut buf = vec![0; PROOF_PATH_SIZE];
key.write(&mut buf);
let mut key2 = ProofPath::read(&buf);
key2.start = 5;
assert_eq!(buf[0], BRANCH_KEY_PREFIX);
assert_eq!(buf[33], 11);
assert_eq!(&buf[1..3], &[255, 7]);
assert_eq!(&buf[3..33], &[0; 30]);
assert_eq!(key2, key);
}
#[test]
fn test_proof_path_suffix() {
let b = ProofPath::from_raw(*b"\x00\x01\x02\xFF\x0C0000000000000000000000000000\x20");
assert_eq!(b.len(), 32);
assert_eq!(b.bit(0), ChildKind::Right);
assert_eq!(b.bit(7), ChildKind::Left);
assert_eq!(b.bit(8), ChildKind::Left);
assert_eq!(b.bit(9), ChildKind::Right);
assert_eq!(b.bit(15), ChildKind::Left);
assert_eq!(b.bit(16), ChildKind::Right);
assert_eq!(b.bit(20), ChildKind::Right);
assert_eq!(b.bit(23), ChildKind::Right);
assert_eq!(b.bit(26), ChildKind::Right);
assert_eq!(b.bit(27), ChildKind::Right);
assert_eq!(b.bit(31), ChildKind::Left);
let b2 = b.suffix(8);
assert_eq!(b2.len(), 24);
assert_eq!(b2.bit(0), ChildKind::Left);
assert_eq!(b2.bit(1), ChildKind::Right);
assert_eq!(b2.bit(7), ChildKind::Left);
assert_eq!(b2.bit(12), ChildKind::Right);
assert_eq!(b2.bit(15), ChildKind::Right);
let b3 = b2.suffix(24);
assert_eq!(b3.len(), 0);
let b4 = b.suffix(1);
assert_eq!(b4.bit(6), ChildKind::Left);
assert_eq!(b4.bit(7), ChildKind::Left);
assert_eq!(b4.bit(8), ChildKind::Right);
}
#[test]
fn test_proof_path_prefix() {
let b = ProofPath::from_raw(*b"\x00\x83wertyuiopasdfghjklzxcvbnm123456\x08");
assert_eq!(b.len(), 8);
assert_eq!(b.prefix(1).bit(0), ChildKind::Right);
assert_eq!(b.prefix(1).len(), 1);
}
#[test]
fn test_proof_path_len() {
let b = ProofPath::from_raw(*b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00");
assert_eq!(b.len(), 256);
}
#[test]
#[should_panic(expected = "self.start() + idx < self.end()")]
fn test_proof_path_at_overflow() {
let b = ProofPath::from_raw(*b"\x00qwertyuiopasdfghjklzxcvbnm123456\x0F");
b.bit(32);
}
#[test]
#[should_panic(expected = "pos <= self.end()")]
fn test_proof_path_suffix_overflow() {
let b = ProofPath::from_raw(*b"\x00qwertyuiopasdfghjklzxcvbnm123456\xFF");
assert_eq!(b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00".len(), 34);
b.suffix(255).suffix(2);
}
#[test]
#[should_panic(expected = "self.start() + idx < self.end()")]
fn test_proof_path_suffix_bit_overflow() {
let b = ProofPath::from_raw(*b"\x00qwertyuiopasdfghjklzxcvbnm123456\xFF");
b.suffix(1).bit(255);
}
#[test]
fn test_proof_path_common_prefix_len() {
let b1 = ProofPath::from_raw(*b"\x01abcd0000000000000000000000000000\x00");
let b2 = ProofPath::from_raw(*b"\x01abef0000000000000000000000000000\x00");
assert_eq!(b1.common_prefix_len(&b1), 256);
let c = b1.common_prefix_len(&b2);
assert_eq!(c, 17);
let c = b2.common_prefix_len(&b1);
assert_eq!(c, 17);
let b1 = b1.suffix(9);
let b2 = b2.suffix(9);
let c = b1.common_prefix_len(&b2);
assert_eq!(c, 8);
let b3 = ProofPath::from_raw(*b"\x01\xFF0000000000000000000000000000000\x00");
let b4 = ProofPath::from_raw(*b"\x01\xF70000000000000000000000000000000\x00");
assert_eq!(b3.common_prefix_len(&b4), 3);
assert_eq!(b4.common_prefix_len(&b3), 3);
assert_eq!(b3.common_prefix_len(&b3), 256);
let b3 = b3.suffix(30);
assert_eq!(b3.common_prefix_len(&b3), 226);
let b3 = b3.prefix(200);
assert_eq!(b3.common_prefix_len(&b3), 200);
let b5 = ProofPath::from_raw(*b"\x01\xF00000000000000000000000000000000\x00");
assert_eq!(b5.prefix(0).common_prefix_len(&b3), 0);
}
#[test]
fn test_proof_path_is_leaf() {
let b = ProofPath::from_raw(*b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00");
assert_eq!(b.len(), 256);
assert_eq!(b.suffix(4).is_leaf(), true);
assert_eq!(b.suffix(8).is_leaf(), true);
assert_eq!(b.suffix(250).is_leaf(), true);
assert_eq!(b.prefix(16).is_leaf(), false);
}
#[test]
fn test_proof_path_is_branch() {
let b = ProofPath::from_raw(*b"\x00qwertyuiopasdfghjklzxcvbnm123456\xFF");
assert_eq!(b.len(), 255);
assert_eq!(b.is_leaf(), false);
}
#[test]
fn test_proof_path_debug_leaf() {
use std::fmt::Write;
let b = ProofPath::from_raw(*b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00");
let mut buf = String::new();
write!(&mut buf, "{:?}", b).unwrap();
assert_eq!(
buf,
"ProofPath { start: 0, end: 256, bits: \"01110001|01110111|01100101|01110010|01110100|0111\
1001|01110101|01101001|01101111|01110000|01100001|01110011|01100100|01100110|01100111|0110\
1000|01101010|01101011|01101100|01111010|01111000|01100011|01110110|01100010|01101110|0110\
1101|00110001|00110010|00110011|00110100|00110101|00110110|\" }"
);
}
#[test]
fn test_proof_path_debug_branch() {
use std::fmt::Write;
let b = ProofPath::from_raw(*b"\x00qwertyuiopasdfghjklzxcvbnm123456\xF0").suffix(12);
let mut buf = String::new();
write!(&mut buf, "{:?}", b).unwrap();
assert_eq!(
buf,
"ProofPath { start: 12, end: 240, bits: \"________|0111____|01100101|01110010|01110100|011\
11001|01110101|01101001|01101111|01110000|01100001|01110011|01100100|01100110|01100111|011\
01000|01101010|01101011|01101100|01111010|01111000|01100011|01110110|01100010|01101110|011\
01101|00110001|00110010|00110011|00110100|________|________|\" }"
);
}