/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt;
use std::fmt::Debug;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use bytes::Bytes;
use fbthrift::compact_protocol;
use quickcheck::single_shrinker;
use quickcheck::Arbitrary;
use quickcheck::Gen;

use crate::blob::Blob;
use crate::blob::BlobstoreValue;
use crate::blob::RawBundle2Blob;
use crate::errors::ErrorKind;
use crate::thrift;
use crate::typed_hash::RawBundle2Id;
use crate::typed_hash::RawBundle2IdContext;

/// An enum representing contents of a raw bundle2 blob
#[derive(Clone, Eq, PartialEq)]
pub enum RawBundle2 {
    Bytes(Bytes),
}

impl RawBundle2 {
    pub fn new_bytes<B: Into<Bytes>>(b: B) -> Self {
        RawBundle2::Bytes(b.into())
    }

    pub(crate) fn from_thrift(fc: thrift::RawBundle2) -> Result<Self> {
        match fc {
            thrift::RawBundle2::Bytes(bytes) => Ok(RawBundle2::Bytes(bytes.into())),
            thrift::RawBundle2::UnknownField(x) => bail!(ErrorKind::InvalidThrift(
                "RawBundle2".into(),
                format!("unknown rawbundle2 field: {}", x)
            )),
        }
    }

    pub fn size(&self) -> usize {
        match *self {
            RawBundle2::Bytes(ref bytes) => bytes.len(),
        }
    }

    /// Whether this starts with a particular string.
    #[inline]
    pub fn starts_with(&self, needle: &[u8]) -> bool {
        match self {
            RawBundle2::Bytes(b) => b.starts_with(needle),
        }
    }

    pub fn into_bytes(self) -> Bytes {
        match self {
            RawBundle2::Bytes(bytes) => bytes,
        }
    }

    pub fn as_bytes(&self) -> &Bytes {
        match self {
            RawBundle2::Bytes(bytes) => bytes,
        }
    }

    pub(crate) fn into_thrift(self) -> thrift::RawBundle2 {
        match self {
            // TODO (T26959816) -- allow Thrift to represent binary as Bytes
            RawBundle2::Bytes(bytes) => thrift::RawBundle2::Bytes(bytes.to_vec()),
        }
    }

    pub fn from_encoded_bytes(encoded_bytes: Bytes) -> Result<Self> {
        let thrift_tc = compact_protocol::deserialize(encoded_bytes.as_ref())
            .with_context(|| ErrorKind::BlobDeserializeError("RawBundle2".into()))?;
        Self::from_thrift(thrift_tc)
    }
}

impl BlobstoreValue for RawBundle2 {
    type Key = RawBundle2Id;

    fn into_blob(self) -> RawBundle2Blob {
        let mut context = RawBundle2IdContext::new();
        context.update(self.as_bytes());
        let id = context.finish();
        let thrift = self.into_thrift();
        let data = compact_protocol::serialize(thrift);
        Blob::new(id, data)
    }

    fn from_blob(blob: RawBundle2Blob) -> Result<Self> {
        let thrift_tc = compact_protocol::deserialize(blob.data().as_ref())
            .with_context(|| ErrorKind::BlobDeserializeError("RawBundle2".into()))?;
        Self::from_thrift(thrift_tc)
    }
}

impl Debug for RawBundle2 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            RawBundle2::Bytes(ref bytes) => write!(f, "RawBundle2::Bytes(length {})", bytes.len()),
        }
    }
}

impl Arbitrary for RawBundle2 {
    fn arbitrary(g: &mut Gen) -> Self {
        RawBundle2::new_bytes(Vec::arbitrary(g))
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        single_shrinker(RawBundle2::new_bytes(vec![]))
    }
}

#[cfg(test)]
mod test {
    use quickcheck::quickcheck;

    use super::*;

    quickcheck! {
        fn thrift_roundtrip(fc: RawBundle2) -> bool {
            let thrift_fc = fc.clone().into_thrift();
            let fc2 = RawBundle2::from_thrift(thrift_fc)
                .expect("thrift roundtrips should always be valid");
            fc == fc2
        }

        fn blob_roundtrip(cs: RawBundle2) -> bool {
            let blob = cs.clone().into_blob();
            let cs2 = RawBundle2::from_blob(blob)
                .expect("blob roundtrips should always be valid");
            cs == cs2
        }
    }

    #[test]
    fn bad_thrift() {
        let thrift_fc = thrift::RawBundle2::UnknownField(-1);
        RawBundle2::from_thrift(thrift_fc).expect_err("unexpected OK - unknown field");
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use fbthrift::compact_protocol;

use crate::blob::Blob;
use crate::blob::BlobstoreValue;
use crate::blob::RedactionKeyListBlob;
use crate::thrift;
use crate::typed_hash::RedactionKeyListId;
use crate::typed_hash::RedactionKeyListIdContext;

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct RedactionKeyList {
    pub keys: Vec<String>,
}

impl RedactionKeyList {
    fn into_thrift(self) -> thrift::RedactionKeyList {
        thrift::RedactionKeyList { keys: self.keys }
    }

    fn from_thrift(t: thrift::RedactionKeyList) -> Result<Self> {
        Ok(Self { keys: t.keys })
    }

    pub fn from_bytes(serialized: &[u8]) -> Result<Self> {
        Self::from_thrift(
            compact_protocol::deserialize(serialized)
                .with_context(|| anyhow!("While deserializing RedactionKeyList"))?,
        )
    }
}

impl BlobstoreValue for RedactionKeyList {
    type Key = RedactionKeyListId;

    fn into_blob(self) -> RedactionKeyListBlob {
        let thrift = self.into_thrift();
        let data = compact_protocol::serialize(thrift);
        let mut context = RedactionKeyListIdContext::new();
        context.update(&data);
        let id = context.finish();
        Blob::new(id, data)
    }

    fn from_blob(blob: Blob<Self::Key>) -> Result<Self> {
        Self::from_bytes(blob.data().as_ref())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt;
use std::str::FromStr;

use abomonation_derive::Abomonation;
use anyhow::Error;
use lazy_static::lazy_static;
use regex::Regex;
use serde_derive::Serialize;
use sql::mysql;

use crate::errors::ErrorKind;

pub const REPO_ID_PREFIX: &str = "repo";
pub const REPO_ID_SUFFIX: &str = ".";
pub const REPO_ID_SUFFIX_PATTERN: &str = r"\.";

pub const EPH_ID_PREFIX: &str = "eph";
pub const EPH_ID_SUFFIX: &str = ".";
pub const EPH_ID_SUFFIX_PATTERN: &str = r"\.";

lazy_static! {
    /// Matches the repo prefix for repo-specific keys.
    pub static ref REPO_PREFIX_REGEX: Regex = Regex::new(
        format!(r"^{}(\d{{3}}\d+){}", REPO_ID_PREFIX, REPO_ID_SUFFIX_PATTERN
    ).as_str()).unwrap();

    /// Matches the ephemeral and repo prefix for repo-specific ephemeral
    /// keys.
    pub static ref EPH_REPO_PREFIX_REGEX: Regex = Regex::new(
        format!(
            r"^{}(\d+){}{}(\d{{3}}\d+){}",
            EPH_ID_PREFIX,
            EPH_ID_SUFFIX_PATTERN,
            REPO_ID_PREFIX,
            REPO_ID_SUFFIX_PATTERN,
    ).as_str()).unwrap();
}

/// Represents a repository. This ID is used throughout storage.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug, Hash, Abomonation)]
#[derive(Serialize, mysql::OptTryFromRowField)]
pub struct RepositoryId(i32);

impl RepositoryId {
    // TODO: (rain1) T30368905 Instead of using this struct directly, have a wrapper around it that
    // only accepts a u32.
    #[inline]
    pub const fn new(id: i32) -> Self {
        Self(id)
    }

    #[inline]
    pub fn id(&self) -> i32 {
        self.0
    }

    /// Generate a prefix suitable for a blobstore.
    #[inline]
    pub fn prefix(&self) -> String {
        // Generate repo0001, repo0002, etc.
        format!("{}{:04}{}", REPO_ID_PREFIX, self.0, REPO_ID_SUFFIX)
    }
}

impl fmt::Display for RepositoryId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl Default for RepositoryId {
    fn default() -> Self {
        Self::new(0)
    }
}

impl FromStr for RepositoryId {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        s.parse::<u32>()
            .map_err(|_| ErrorKind::FailedToParseRepositoryId(s.to_owned()).into())
            .map(|repoid| Self::new(repoid as i32))
    }
}

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

    #[test]
    fn prefix_generation() {
        assert_eq!(RepositoryId::new(0).prefix().as_str(), "repo0000.");
        assert_eq!(RepositoryId::new(1).prefix().as_str(), "repo0001.");
        assert_eq!(RepositoryId::new(99).prefix().as_str(), "repo0099.");
        assert_eq!(RepositoryId::new(456).prefix().as_str(), "repo0456.");
        assert_eq!(RepositoryId::new(9999).prefix().as_str(), "repo9999.");
        assert_eq!(RepositoryId::new(12000).prefix().as_str(), "repo12000.");
    }

    #[test]
    fn prefix_match() {
        // Check generated prefixes match the expected form
        assert!(REPO_PREFIX_REGEX.is_match(RepositoryId::new(0).prefix().as_str()));
        assert!(REPO_PREFIX_REGEX.is_match(RepositoryId::new(1).prefix().as_str()));
        assert!(REPO_PREFIX_REGEX.is_match(RepositoryId::new(99).prefix().as_str()));
        assert!(REPO_PREFIX_REGEX.is_match(RepositoryId::new(456).prefix().as_str()));
        assert!(REPO_PREFIX_REGEX.is_match(RepositoryId::new(9999).prefix().as_str()));
        assert!(REPO_PREFIX_REGEX.is_match(RepositoryId::new(12000).prefix().as_str()));
    }

    #[test]
    fn prefix_not_match() {
        // Check we dont match unexpected forms
        assert!(!REPO_PREFIX_REGEX.is_match(""));
        assert!(!REPO_PREFIX_REGEX.is_match("."));
        assert!(!REPO_PREFIX_REGEX.is_match(".repo0000."));
        assert!(!REPO_PREFIX_REGEX.is_match("repo0000"));
        assert!(!REPO_PREFIX_REGEX.is_match("repo(0000)"));
        assert!(!REPO_PREFIX_REGEX.is_match("repo0000a"));
        assert!(!REPO_PREFIX_REGEX.is_match("repo00a0."));
        assert!(!REPO_PREFIX_REGEX.is_match("flat/repo0000."));
        assert!(!REPO_PREFIX_REGEX.is_match("repo."));
        assert!(!REPO_PREFIX_REGEX.is_match("repo0."));
        assert!(!REPO_PREFIX_REGEX.is_match("repo00."));
        assert!(!REPO_PREFIX_REGEX.is_match("repo000."));
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::str::FromStr;

use abomonation_derive::Abomonation;
use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use ascii::AsciiStr;
use ascii::AsciiString;
use faster_hex::hex_decode;
use faster_hex::hex_encode;
use quickcheck::single_shrinker;
use quickcheck::Arbitrary;
use quickcheck::Gen;
use serde_derive::Deserialize;
use serde_derive::Serialize;
use sha1::Digest;

use crate::errors::ErrorKind;
use crate::thrift;

pub const SHA1_HASH_LENGTH_BYTES: usize = 20;
pub const SHA1_HASH_LENGTH_HEX: usize = SHA1_HASH_LENGTH_BYTES * 2;

pub const NULL: Sha1 = Sha1([0; SHA1_HASH_LENGTH_BYTES]);

/// Raw SHA-1 hash
///
/// Mercurial bases all its hashing on SHA-1, but this type is only used to build
/// more specific typed hashes.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[derive(Serialize, Deserialize, Abomonation)]
pub struct Sha1([u8; SHA1_HASH_LENGTH_BYTES]);

impl Sha1 {
    /// Construct a `Sha1` from an array of SHA1_HASH_LENGTH_BYTES (20) bytes containing a
    /// SHA-1 (ie, *not* a hash of the bytes).
    pub fn from_bytes<B: AsRef<[u8]>>(bytes: B) -> Result<Sha1> {
        let bytes = bytes.as_ref();
        if bytes.len() != SHA1_HASH_LENGTH_BYTES {
            bail!(ErrorKind::InvalidSha1Input(format!(
                "need exactly {} bytes",
                SHA1_HASH_LENGTH_BYTES
            )));
        } else {
            let mut ret = [0; SHA1_HASH_LENGTH_BYTES];
            ret.copy_from_slice(bytes);
            Ok(Sha1(ret))
        }
    }

    pub fn from_thrift(h: thrift::Sha1) -> Result<Self> {
        if h.0.len() != SHA1_HASH_LENGTH_BYTES {
            bail!(ErrorKind::InvalidThrift(
                "Sha1".into(),
                format!(
                    "wrong length: expected {}, got {}",
                    SHA1_HASH_LENGTH_BYTES,
                    h.0.len()
                )
            ));
        }
        let mut arr = [0u8; SHA1_HASH_LENGTH_BYTES];
        arr.copy_from_slice(&h.0[..]);
        Ok(Sha1(arr))
    }

    /// Construct a `Sha1` from an array of SHA1_HASH_LENGTH_BYTES bytes.
    #[inline]
    pub const fn from_byte_array(arr: [u8; SHA1_HASH_LENGTH_BYTES]) -> Sha1 {
        Sha1(arr)
    }

    /// Extract this hash's underlying byte array.
    #[inline]
    pub fn into_byte_array(self) -> [u8; SHA1_HASH_LENGTH_BYTES] {
        self.0
    }

    /// Construct a `Sha1` from a hex-encoded `AsciiStr`.
    #[inline]
    pub fn from_ascii_str(s: &AsciiStr) -> Result<Sha1> {
        Self::from_str(s.as_str())
    }

    pub fn to_hex(&self) -> AsciiString {
        let mut v = [0; SHA1_HASH_LENGTH_HEX];

        // This can only panic if buffer size of Vec isn't correct, which would be
        // a programming error.
        hex_encode(self.as_ref(), &mut v).expect("failed to hex encode");

        unsafe {
            // A hex string is always a pure ASCII string.
            AsciiString::from_ascii_unchecked(v)
        }
    }

    pub fn into_thrift(self) -> thrift::Sha1 {
        thrift::Sha1(self.0.into())
    }
}

/// Context for incrementally computing a `Sha1` hash.
#[derive(Clone)]
pub struct Context(sha1::Sha1);

/// Compute the `Sha1` for a slice of bytes.
impl<'a> From<&'a [u8]> for Sha1 {
    fn from(data: &[u8]) -> Sha1 {
        let mut sha1 = sha1::Sha1::new();
        sha1.update(data);

        Sha1(sha1.finalize().into())
    }
}

/// Get a reference to the underlying bytes of a `Sha1`
impl AsRef<[u8]> for Sha1 {
    fn as_ref(&self) -> &[u8] {
        &self.0[..]
    }
}

impl FromStr for Sha1 {
    type Err = Error;

    fn from_str(s: &str) -> Result<Sha1> {
        if s.len() != SHA1_HASH_LENGTH_HEX {
            bail!(ErrorKind::InvalidSha1Input(format!(
                "need exactly {} hex digits",
                SHA1_HASH_LENGTH_HEX
            )));
        }

        let mut ret = Sha1([0; SHA1_HASH_LENGTH_BYTES]);
        match hex_decode(s.as_bytes(), &mut ret.0) {
            Ok(_) => Ok(ret),
            Err(_) => bail!(ErrorKind::InvalidSha1Input("bad hex character".into())),
        }
    }
}

impl Display for Sha1 {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt(&self.to_hex(), fmt)
    }
}

/// Custom `Debug` output for `Sha1` so it prints in hex.
impl Debug for Sha1 {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Sha1({})", self)
    }
}

impl Arbitrary for Sha1 {
    fn arbitrary(g: &mut Gen) -> Self {
        let mut bytes = [0; SHA1_HASH_LENGTH_BYTES];
        // The null hash is special, so give it a 5% chance of happening
        if usize::arbitrary(g) % SHA1_HASH_LENGTH_BYTES >= 1 {
            for b in bytes.iter_mut() {
                *b = u8::arbitrary(g);
            }
        }
        Sha1::from_bytes(bytes).unwrap()
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        single_shrinker(NULL)
    }
}

#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[derive(Serialize, Deserialize, Abomonation)]
/// Raw SHA-1 hash prefix.
/// Internal implementation is the inclusive range of Sha1 objects.
/// If can be build from a from a hex-encoded string (len <= SHA1_HASH_LENGTH_HEX (40))
/// or from an array of bytes (len <= SHA1_HASH_LENGTH_BYTES (20)).
pub struct Sha1Prefix(Sha1, Sha1);

impl Sha1Prefix {
    /// Construct a `Sha1Prefix` from an array of bytes.
    pub fn from_bytes<B: AsRef<[u8]> + ?Sized>(bytes: &B) -> Result<Self> {
        let bytes = bytes.as_ref();
        if bytes.len() > SHA1_HASH_LENGTH_BYTES {
            bail!(ErrorKind::InvalidSha1Input(format!(
                "prefix needs to be less or equal to {} bytes",
                SHA1_HASH_LENGTH_BYTES
            )))
        } else {
            static SHA1_MIN: [u8; SHA1_HASH_LENGTH_BYTES] = [0x00; SHA1_HASH_LENGTH_BYTES];
            static SHA1_MAX: [u8; SHA1_HASH_LENGTH_BYTES] = [0xff; SHA1_HASH_LENGTH_BYTES];

            let min_tail = &SHA1_MIN[bytes.len()..];
            let max_tail = &SHA1_MAX[bytes.len()..];

            Ok(Sha1Prefix(
                Sha1::from_bytes(bytes.iter().chain(min_tail).copied().collect::<Vec<_>>())?,
                Sha1::from_bytes(bytes.iter().chain(max_tail).copied().collect::<Vec<_>>())?,
            ))
        }
    }

    #[inline]
    /// Get a reference to the underlying bytes of the `Sha1` lower bound object.
    pub fn min_as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }

    #[inline]
    /// Get a reference to the underlying bytes of the `Sha1` inclusive upper bound object.
    pub fn max_as_ref(&self) -> &[u8] {
        self.1.as_ref()
    }

    /// Convert into sha1 if it is the full prefix of the hash.
    #[inline]
    pub fn into_sha1(self) -> Option<Sha1> {
        if self.0 == self.1 { Some(self.0) } else { None }
    }

    pub fn to_hex(&self) -> AsciiString {
        let mut v_min_hex = &mut [0; SHA1_HASH_LENGTH_HEX][..];
        hex_encode(self.0.as_ref(), v_min_hex).expect("failed to hex encode");
        let v_max_hex = &mut [0; SHA1_HASH_LENGTH_HEX][..];
        hex_encode(self.1.as_ref(), v_max_hex).expect("failed to hex encode");
        for i in 0..SHA1_HASH_LENGTH_HEX {
            if v_min_hex[i] != v_max_hex[i] {
                v_min_hex = &mut v_min_hex[..i];
                break;
            }
        }
        unsafe {
            // A hex string is always a pure ASCII string.
            AsciiString::from_ascii_unchecked(v_min_hex)
        }
    }
}

/// Construct a `Sha1Prefix` from a hex-encoded string.
impl FromStr for Sha1Prefix {
    type Err = Error;
    fn from_str(s: &str) -> Result<Sha1Prefix> {
        if s.len() > SHA1_HASH_LENGTH_HEX {
            bail!(ErrorKind::InvalidSha1Input(format!(
                "prefix needs to be less or equal {} hex digits",
                SHA1_HASH_LENGTH_HEX
            )));
        }
        let min_tail: String = String::from_utf8(vec![b'0'; SHA1_HASH_LENGTH_HEX - s.len()])?;
        let max_tail: String = String::from_utf8(vec![b'f'; SHA1_HASH_LENGTH_HEX - s.len()])?;
        Ok(Sha1Prefix(
            Sha1::from_str(&(s.to_owned() + &min_tail))?,
            Sha1::from_str(&(s.to_owned() + &max_tail))?,
        ))
    }
}

/// Custom `Display` output for `Sha1Prefix` so it prints in hex.
impl Display for Sha1Prefix {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt(&self.to_hex(), fmt)
    }
}

/// Custom `Debug` output for `Sha1Prefix` so it prints in hex.
impl Debug for Sha1Prefix {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Sha1Prefix({})", self)
    }
}

impl Context {
    /// Construct a `Context`
    pub fn new() -> Context {
        Context(sha1::Sha1::new())
    }

    /// Update a context from something that can be turned into a `&[u8]`
    pub fn update<T>(&mut self, data: T)
    where
        T: AsRef<[u8]>,
    {
        self.0.update(data.as_ref())
    }

    pub fn finish(self) -> Sha1 {
        Sha1(self.0.finalize().into())
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use quickcheck::quickcheck;
    use quickcheck::TestResult;

    use super::*;

    #[rustfmt::skip]
    const NILHASH: Sha1 = Sha1([0xda, 0x39, 0xa3, 0xee,
                                0x5e, 0x6b, 0x4b, 0x0d,
                                0x32, 0x55, 0xbf, 0xef,
                                0x95, 0x60, 0x18, 0x90,
                                0xaf, 0xd8, 0x07, 0x09]);

    #[test]
    fn test_null() {
        assert_eq!(NULL, Sha1([0_u8; SHA1_HASH_LENGTH_BYTES]));
    }

    #[test]
    fn test_nil() {
        let nil = Sha1::from(&[][..]);
        assert_eq!(nil, NILHASH);
    }

    #[test]
    fn parse_ok() {
        assert_eq!(
            NULL,
            Sha1::from_str("0000000000000000000000000000000000000000").unwrap()
        );
        assert_eq!(
            NILHASH,
            Sha1::from_str("da39a3ee5e6b4b0d3255bfef95601890afd80709").unwrap()
        );
        assert_eq!(
            NILHASH,
            Sha1::from_str("DA39A3EE5E6B4B0D3255BFEF95601890AFD80709").unwrap()
        );
    }

    #[test]
    fn test_display() {
        assert_eq!(
            format!("{}", NULL),
            "0000000000000000000000000000000000000000"
        );
        assert_eq!(
            format!("{}", NILHASH),
            "da39a3ee5e6b4b0d3255bfef95601890afd80709"
        );
    }

    #[test]
    fn test_parse_and_display_prefix() {
        // even length
        assert_eq!(
            format!("{}", Sha1Prefix::from_str("da39a3").unwrap()),
            "da39a3"
        );
        // odd length
        assert_eq!(
            format!("{}", Sha1Prefix::from_str("da39a").unwrap()),
            "da39a"
        );
        // max length
        assert_eq!(
            format!(
                "{}",
                Sha1Prefix::from_str("da39a3ee5e6b4b0d3255bfef95601890afd80709").unwrap()
            ),
            "da39a3ee5e6b4b0d3255bfef95601890afd80709"
        );
        // capital case
        assert_eq!(
            format!(
                "{}",
                Sha1Prefix::from_str("DA39A3EE5E6B4B0D3255BFEF95601890AFD80709").unwrap()
            ),
            "da39a3ee5e6b4b0d3255bfef95601890afd80709"
        );
        // zero length
        assert_eq!(format!("{}", Sha1Prefix::from_str("").unwrap()), "");
    }

    #[test]
    fn parse_bad() {
        match Sha1::from_str("") {
            Ok(_) => panic!("unexpected OK - zero len"),
            Err(_) => {}
        };
        match Sha1::from_str("da39a3ee5e6b4b0d3255bfef95601890afd8070") {
            // one char missing
            Ok(_) => panic!("unexpected OK - trunc"),
            Err(_) => {}
        };
        match Sha1::from_str("xda39a3ee5e6b4b0d3255bfef95601890afd8070") {
            // one char bad
            Ok(_) => panic!("unexpected OK - badchar end"),
            Err(_) => {}
        };
        match Sha1::from_str("da39a3ee5e6b4b0d3255bfef95601890afd8070x") {
            // one char bad
            Ok(_) => panic!("unexpected OK - badchar end"),
            Err(_) => {}
        };
        match Sha1::from_str("da39a3ee5e6b4b0d325Xbfef95601890afd80709") {
            // one char missing
            Ok(_) => panic!("unexpected OK - trunc"),
            Err(_) => {}
        };
    }

    #[test]
    fn parse_thrift() {
        let null_thrift = thrift::Sha1(vec![0; SHA1_HASH_LENGTH_BYTES].into());
        assert_eq!(NULL, Sha1::from_thrift(null_thrift.clone()).unwrap());
        assert_eq!(NULL.into_thrift(), null_thrift);

        let nil_thrift = thrift::Sha1(NILHASH.0.into());
        assert_eq!(NILHASH, Sha1::from_thrift(nil_thrift.clone()).unwrap());
        assert_eq!(NILHASH.into_thrift(), nil_thrift);
    }

    #[test]
    fn parse_thrift_bad() {
        Sha1::from_thrift(thrift::Sha1(vec![].into())).expect_err("unexpected OK - zero len");
        Sha1::from_thrift(thrift::Sha1(vec![0; 19].into())).expect_err("unexpected OK - too short");
        Sha1::from_thrift(thrift::Sha1(vec![0; 21].into())).expect_err("unexpected OK - too long");
    }

    quickcheck! {
        fn parse_roundtrip(v: Vec<u8>) -> TestResult {
            if v.len() != SHA1_HASH_LENGTH_BYTES {
                return TestResult::discard()
            }
            let h = Sha1::from_bytes(v).unwrap();
            let s = format!("{}", h);
            let sh = s.parse().unwrap();

            TestResult::from_bool(h == sh)
        }

        fn thrift_roundtrip(h: Sha1) -> bool {
            let v = h.into_thrift();
            let sh = Sha1::from_thrift(v)
                .expect("converting a valid Thrift structure should always work");
            h == sh
        }

        fn to_hex_roundtrip(h: Sha1) -> bool {
            let v = h.to_hex();
            let sh = Sha1::from_ascii_str(&v).unwrap();
            h == sh
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(dead_code)]
#![allow(clippy::mutable_key_type)] // false positive: Bytes is not inner mutable

use std::collections::BTreeMap;
use std::fmt::Debug;

use anyhow::bail;
use anyhow::Context;
use anyhow::Ok;
use anyhow::Result;
use async_recursion::async_recursion;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::Storable;
use bounded_traversal::bounded_traversal_ordered_stream;
use bounded_traversal::OrderedTraversal;
use bytes::Bytes;
use context::CoreContext;
use derivative::Derivative;
use futures::stream;
use futures::FutureExt;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use itertools::Itertools;
use nonzero_ext::nonzero;
use once_cell::sync::OnceCell;
use smallvec::SmallVec;
use sorted_vector_map::sorted_vector_map;
use sorted_vector_map::SortedVectorMap;

use crate::blob::Blob;
use crate::blob::BlobstoreValue;
use crate::thrift;
use crate::typed_hash::IdContext;
use crate::typed_hash::MononokeId;
use crate::ThriftConvert;

pub trait MapValue: ThriftConvert + Debug + Clone + Send + Sync + 'static {
    type Id: MononokeId<Thrift = thrift::ShardedMapNodeId, Value = ShardedMapNode<Self>>;
    type Context: IdContext<Id = Self::Id>;
}

type SmallBinary = SmallVec<[u8; 24]>;

#[derive(Derivative)]
#[derivative(PartialEq, Debug)]
#[derive(Clone, Eq)]
pub enum ShardedMapNode<Value: MapValue> {
    Intermediate {
        prefix: SmallBinary,
        value: Option<Value>,
        edges: SortedVectorMap<u8, ShardedMapEdge<Value>>,
        #[derivative(PartialEq = "ignore", Debug = "ignore")]
        size: OnceCell<usize>,
    },
    Terminal {
        // The key is the original map key minus the prefixes and edges from all
        // intermediate nodes in the path to this node.
        values: SortedVectorMap<SmallBinary, Value>,
    },
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ShardedMapEdge<Value: MapValue> {
    size: usize,
    child: ShardedMapChild<Value>,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum ShardedMapChild<Value: MapValue> {
    Id(Value::Id),
    Inlined(ShardedMapNode<Value>),
}

impl<Value: MapValue> ShardedMapEdge<Value> {
    async fn load_child(
        self,
        ctx: &CoreContext,
        blobstore: &impl Blobstore,
    ) -> Result<ShardedMapNode<Value>> {
        self.child.load(ctx, blobstore).await
    }

    fn from_thrift(t: thrift::ShardedMapEdge) -> Result<Self> {
        Ok(Self {
            size: t.size.try_into().context("Failed to parse size to usize")?,
            child: ShardedMapChild::from_thrift(t.child)?,
        })
    }

    fn into_thrift(self) -> thrift::ShardedMapEdge {
        thrift::ShardedMapEdge {
            size: self.size as i64,
            child: self.child.into_thrift(),
        }
    }

    /// Return this MapNode's id if it's not inlined
    pub fn id(&self) -> Option<Value::Id> {
        self.child.id()
    }
}

impl<Value: MapValue> ShardedMapChild<Value> {
    pub async fn load(
        self,
        ctx: &CoreContext,
        blobstore: &impl Blobstore,
    ) -> Result<ShardedMapNode<Value>> {
        match self {
            Self::Inlined(inlined) => Ok(inlined),
            Self::Id(id) => id.load(ctx, blobstore).await.map_err(Into::into),
        }
    }

    fn from_thrift(t: thrift::ShardedMapChild) -> Result<Self> {
        Ok(match t {
            thrift::ShardedMapChild::inlined(inlined) => {
                Self::Inlined(ShardedMapNode::from_thrift(inlined)?)
            }
            thrift::ShardedMapChild::id(id) => Self::Id(Value::Id::from_thrift(id)?),
            thrift::ShardedMapChild::UnknownField(_) => bail!("Unknown variant"),
        })
    }

    fn into_thrift(self) -> thrift::ShardedMapChild {
        match self {
            Self::Inlined(inlined) => thrift::ShardedMapChild::inlined(inlined.into_thrift()),
            Self::Id(id) => thrift::ShardedMapChild::id(id.into_thrift()),
        }
    }

    /// Return this MapNode's id if it's not inlined
    pub fn id(&self) -> Option<Value::Id> {
        match self {
            Self::Id(v) => Some(*v),
            Self::Inlined(_) => None,
        }
    }
}

impl<Value: MapValue> Default for ShardedMapChild<Value> {
    fn default() -> Self {
        Self::Inlined(Default::default())
    }
}

impl<Value: MapValue> Default for ShardedMapNode<Value> {
    fn default() -> Self {
        Self::Terminal {
            values: SortedVectorMap::new(),
        }
    }
}

impl<Value: MapValue> Default for ShardedMapEdge<Value> {
    fn default() -> Self {
        Self {
            size: 0,
            child: Default::default(),
        }
    }
}

/// Returns longest common prefix of a and b.
fn common_prefix<'a>(a: &'a [u8], b: &'a [u8]) -> &'a [u8] {
    let lcp = a.iter().zip(b.iter()).take_while(|(a, b)| a == b).count();
    // Panic safety: lcp is at most a.len()
    &a[..lcp]
}

impl<Value: MapValue> ShardedMapNode<Value> {
    fn intermediate(
        prefix: SmallBinary,
        value: Option<Value>,
        edges: SortedVectorMap<u8, ShardedMapEdge<Value>>,
    ) -> Self {
        Self::Intermediate {
            prefix,
            value,
            edges,
            size: Default::default(),
        }
    }

    /// Given a key, what's the value for that key, if any?
    // See the detailed description of the logic in docs/sharded_map.md
    #[async_recursion]
    pub async fn lookup(
        &self,
        ctx: &CoreContext,
        blobstore: &impl Blobstore,
        key: &[u8],
    ) -> Result<Option<Value>> {
        Ok(match self {
            // Case 1: Do lookup directly on the inlined map
            Self::Terminal { values } => values.get(key).cloned(),
            Self::Intermediate {
                prefix,
                value,
                edges,
                ..
            } => {
                if let Some(key) = key.strip_prefix(prefix.as_slice()) {
                    if let Some((first, rest)) = key.split_first() {
                        if let Some(edge) = edges.get(first) {
                            // Case 2: Recurse, either inlined or first fetching from the blobstore
                            match &edge.child {
                                ShardedMapChild::Inlined(node) => {
                                    node.lookup(ctx, blobstore, rest).await?
                                }
                                ShardedMapChild::Id(id) => {
                                    id.load(ctx, blobstore)
                                        .await?
                                        .lookup(ctx, blobstore, rest)
                                        .await?
                                }
                            }
                        } else {
                            // Case 3: No edge from this node to the next byte of the key
                            None
                        }
                    } else {
                        // Case 4: The node for this key is this intermediate node, not a terminal node
                        value.clone()
                    }
                } else {
                    // Case 5: Key doesn't match prefix
                    None
                }
            }
        })
    }

    /// Given a map and replacements, return the map with the replacements.
    fn update_map(
        mut map: BTreeMap<SmallBinary, Value>,
        replacements: impl IntoIterator<Item = (Bytes, Option<Value>)>,
        deleter: impl Fn(Value),
    ) -> Result<SortedVectorMap<SmallBinary, Value>> {
        for (key, value) in replacements {
            let key = SmallVec::from_iter(key);
            match value {
                Some(value) => map.insert(key, value),
                None => map.remove(&key),
            }
            .map(&deleter);
        }
        Ok(map.into())
    }

    /// Prepend all keys in this node with the given prefix.
    fn prepend(&mut self, prefix: SmallBinary) {
        match self {
            Self::Terminal { values } => {
                *values = std::mem::take(values)
                    .into_iter()
                    .update(|(k, _)| {
                        k.insert_from_slice(0, &prefix);
                    })
                    .collect()
            }
            Self::Intermediate {
                prefix: cur_prefix, ..
            } => {
                cur_prefix.insert_from_slice(0, &prefix);
            }
        }
    }

    fn shard_size() -> Result<usize> {
        if cfg!(test) {
            Ok(5)
        } else {
            thrift::MAP_SHARD_SIZE
                .try_into()
                .context("Failed to parse shard size")
        }
    }

    /// Create a new map from this map with given replacements. It is a generalization of
    /// adding and removing, and should be faster than doing all operations separately.
    /// It does not rely on the added keys not existing or the removed keys existing.
    // See the detailed description of the logic in docs/sharded_map.md
    #[async_recursion]
    pub async fn update(
        self,
        ctx: &CoreContext,
        blobstore: &impl Blobstore,
        replacements: BTreeMap<Bytes, Option<Value>>,
        // Called for all deletions
        deleter: impl Fn(Value) + Send + Copy + 'async_recursion,
    ) -> Result<Self> {
        let shard_size = Self::shard_size()?;
        match self {
            Self::Terminal { values } => {
                let values = Self::update_map(values.into_iter().collect(), replacements, deleter)?;
                if values.len() <= shard_size {
                    // Case 1: values is small enough, return a terminal node
                    Ok(Self::Terminal { values })
                } else {
                    // Case 2: This will become a intermediate node
                    // Let's reuse the logic to add values to a intermediate node by creating
                    // an empty one.
                    let lcp = values
                        .keys()
                        .map(|k| k.as_slice())
                        .reduce(common_prefix)
                        .unwrap_or(b"");
                    // Setting the correct prefix is not necessary for correctness, but it avoids
                    // having Case 3 + Case 4.3.2 + compression unnecessarily.
                    Self::intermediate(SmallBinary::from_slice(lcp), None, Default::default())
                        .update(
                            ctx,
                            blobstore,
                            values
                                .into_iter()
                                .map(|(k, v)| (Bytes::copy_from_slice(k.as_ref()), Some(v)))
                                .collect(),
                            deleter,
                        )
                        .await
                }
            }
            Self::Intermediate {
                mut prefix,
                mut value,
                mut edges,
                ..
            } => {
                // LCP only considered added keys
                let lcp = replacements
                    .iter()
                    .filter_map(|(k, v)| v.as_ref().map(|_| k))
                    .fold(prefix.as_slice(), |lcp, key| common_prefix(lcp, key))
                    .len();
                if lcp < prefix.len() {
                    // Case 3: The prefix of all keys is smaller than `prefix`
                    // Let's create two new nodes and recursively update them.
                    // Right: Bytes lcp + 1 .. size
                    let prefix_right = prefix.drain(lcp + 1..).collect();
                    // Middle: Byte lcp
                    // unwrap safety: lcp + 1 > 0, so split_off leaves at least one element at prefix
                    let mid_byte = prefix.pop().unwrap();
                    // Left: Bytes 0 .. lcp
                    let prefix_left = prefix;
                    let right_node = Self::intermediate(prefix_right, value, edges);
                    let left_node = Self::intermediate(
                        prefix_left,
                        None,
                        // Design decision: all intermediate nodes are inlined
                        sorted_vector_map! {
                            mid_byte =>
                            ShardedMapEdge {
                                size: right_node.size(),
                                child: ShardedMapChild::Inlined(right_node),
                            },
                        },
                    );
                    left_node
                        .update(ctx, blobstore, replacements, deleter)
                        .await
                } else {
                    // Case 4: All added keys traverse the long edge (have the prefix `prefix`)
                    let mut partitioned = BTreeMap::<u8, BTreeMap<Bytes, Option<Value>>>::new();
                    // Step 4.1: Strip prefixes, and partition replacements.
                    replacements.into_iter().for_each(|(k, v)| {
                        match k.strip_prefix(prefix.as_slice()) {
                            None => {
                                // Only deletions might not have the correct prefix
                                debug_assert!(v.is_none());
                            }
                            Some(rest) => {
                                if let Some((first, rest)) = rest.split_first() {
                                    partitioned
                                        .entry(*first)
                                        .or_default()
                                        // Panic safety: rest was produced from k
                                        .insert(k.slice_ref(rest), v);
                                } else {
                                    std::mem::replace(&mut value, v).map(deleter);
                                }
                            }
                        }
                    });

                    // Step 4.2: Recursively update partitioned children
                    let replaced_futures = partitioned
                        .into_iter()
                        .map(|(next_byte, replacements)| {
                            let edge = edges.remove(&next_byte).unwrap_or_default();
                            async move {
                                let node = edge.load_child(ctx, blobstore).await?;
                                let replaced_node =
                                    node.update(ctx, blobstore, replacements, deleter).await?;
                                Ok((next_byte, replaced_node))
                            }
                        })
                        .collect::<Vec<_>>();
                    let replaced = stream::iter(replaced_futures)
                        .buffer_unordered(100)
                        .try_collect::<Vec<_>>()
                        .await?;
                    let mut new_children = BTreeMap::new();
                    for (next_byte, replaced_node) in replaced {
                        if !replaced_node.is_empty() {
                            let previous = new_children.insert(next_byte, replaced_node);
                            debug_assert!(previous.is_none());
                        }
                    }

                    let new_size: usize = edges.values().map(|edge| edge.size).sum::<usize>()
                        + new_children.values().map(|v| v.size()).sum::<usize>()
                        + value.iter().len();

                    if new_size <= shard_size {
                        // Case 4.3.1: Compress node into terminal node.
                        // For simplicity, reuse into_entries.
                        // In practice, all children will be terminal nodes, so nothing extra
                        // will be unecessarily persisted into the blobstore.
                        for (byte, node) in new_children {
                            debug_assert!(matches!(node, Self::Terminal { .. }));
                            let previous = edges.insert(
                                byte,
                                ShardedMapEdge {
                                    size: node.size(),
                                    child: ShardedMapChild::Inlined(node),
                                },
                            );
                            debug_assert!(previous.is_none());
                        }
                        let values = Self::intermediate(prefix, value, edges)
                            .into_entries(ctx, blobstore)
                            // Extending SortedVectorMap 1 by 1 will be fast because into_entries
                            // returns elements in order
                            .try_collect()
                            .await?;

                        Ok(Self::Terminal { values })
                    } else {
                        // Case 4.3.2: This will continue being a intermediate node, let's
                        // inline what's necessary and store everything
                        let new_edges = stream::iter(new_children)
                            .map(|(byte, node)| async move {
                                let size = node.size();
                                let child = match &node {
                                    // Design decision: Inline all intermediate nodes and store
                                    // terminal nodes separated
                                    Self::Intermediate { .. } => ShardedMapChild::Inlined(node),
                                    Self::Terminal { .. } => ShardedMapChild::Id(
                                        node.into_blob().store(ctx, blobstore).await?,
                                    ),
                                };
                                Ok((byte, ShardedMapEdge { size, child }))
                            })
                            .buffer_unordered(100)
                            .try_collect::<Vec<_>>()
                            .await?;
                        for (byte, edge) in new_edges {
                            let previous = edges.insert(byte, edge);
                            debug_assert!(previous.is_none());
                        }
                        debug_assert!(!edges.is_empty());
                        if edges.len() == 1 && value.is_none() {
                            // Unwrap safety: edges.len() == 1 above
                            let (byte, edge) = edges.into_iter().next().unwrap();
                            let mut child = edge.load_child(ctx, blobstore).await?;
                            prefix.push(byte);
                            child.prepend(prefix);
                            Ok(child)
                        } else {
                            Ok(Self::intermediate(prefix, value, edges))
                        }
                    }
                }
            }
        }
    }

    /// Iterates through all values in the map, asynchronously and only loading
    /// blobs as needed.
    // See the detailed description of the logic in docs/sharded_map.md
    pub fn into_entries<'a>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a impl Blobstore,
    ) -> impl Stream<Item = Result<(SmallBinary, Value)>> + 'a {
        self.into_prefix_entries(ctx, blobstore, &[])
    }

    /// Iterates through all values like `into_entries`, but also leaks information about
    /// the shard ids for each encountered ShardMapNode which is stored in the blobstore
    pub fn into_sharded_entries<'a>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a impl Blobstore,
    ) -> impl Stream<Item = Result<ShardedTraversalOutput<'a, Value>>> + 'a {
        self.into_sharded_prefix_entries(ctx, blobstore, &[])
    }

    pub fn into_prefix_entries<'a>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a impl Blobstore,
        prefix: &'a [u8],
    ) -> impl Stream<Item = Result<(SmallBinary, Value)>> + 'a {
        bounded_traversal_ordered_stream(
            nonzero!(256usize),
            nonzero!(256usize),
            vec![(
                self.size(),
                (SmallBinary::new(), prefix, ShardedMapChild::Inlined(self)),
            )],
            move |(mut cur_prefix, remaining_prefix, child): (
                SmallBinary,
                &[u8],
                ShardedMapChild<Value>,
            )| {
                async move {
                    Ok(match child.load(ctx, blobstore).await? {
                        // Case 1. Prepend all keys with cur_prefix and output elements
                        Self::Terminal { values } => values
                            .into_iter()
                            .filter(|(k, _)| k.starts_with(remaining_prefix))
                            .map(|(key, value)| {
                                let mut full_key = cur_prefix.clone();
                                full_key.extend(key);
                                OrderedTraversal::Output((full_key, value))
                            })
                            .collect::<Vec<_>>(),
                        // Case 2. Recurse
                        Self::Intermediate {
                            prefix: new_prefix,
                            value,
                            edges,
                            ..
                        } => {
                            let remaining_prefix = if remaining_prefix.len() >= new_prefix.len() {
                                if let Some(new_remaining) =
                                    remaining_prefix.strip_prefix(new_prefix.as_slice())
                                {
                                    new_remaining
                                } else {
                                    // prefix doesn't match
                                    return Ok(vec![]);
                                }
                            } else if new_prefix.starts_with(remaining_prefix) {
                                &[]
                            } else {
                                // prefix doesn't match
                                return Ok(vec![]);
                            };
                            // Step 2-a. Extend cur_prefix
                            cur_prefix.extend(new_prefix);
                            let cur_prefix = &cur_prefix;
                            remaining_prefix
                                .is_empty()
                                .then_some(value)
                                .flatten()
                                // Step 2-b. If value is present (and prefix empty), output (cur_prefix, value)
                                .map(|value| OrderedTraversal::Output((cur_prefix.clone(), value)))
                                .into_iter()
                                // Step 2-c. Copy prefix, append byte, and recurse.
                                .chain(edges.into_iter().filter_map(|(byte, edge)| {
                                    let (first, rest) =
                                        remaining_prefix.split_first().unwrap_or((&byte, &[]));
                                    if *first == byte {
                                        let mut new_prefix = cur_prefix.clone();
                                        new_prefix.push(byte);
                                        let size_prediction = edge.size;
                                        Some(OrderedTraversal::Recurse(
                                            size_prediction,
                                            (new_prefix, rest, edge.child),
                                        ))
                                    } else {
                                        // Byte didn't match prefix
                                        None
                                    }
                                }))
                                .collect::<Vec<_>>()
                        }
                    })
                }
                .boxed()
            },
        )
    }

    // Traverse the `ShardedMap` down to the next blob, be it a non-inlined `ShardedMapChild` or a `Value`
    // and return the current state of the traversal at this point.
    // This allows us to traverse large sharded entries without loading unnecessary blobs if we
    // don't need to. For instance, it can be used to diff sharded manifests and decide not to
    // load non-inlined shards which are the same between both manifests.
    fn into_sharded_prefix_entries<'a>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a impl Blobstore,
        prefix: &'a [u8],
    ) -> impl Stream<Item = Result<ShardedTraversalOutput<'a, Value>>> + 'a {
        bounded_traversal_ordered_stream(
            nonzero!(256usize),
            nonzero!(256usize),
            vec![(
                self.size(),
                (SmallBinary::new(), prefix, ShardedMapChild::Inlined(self)),
            )],
            move |(cur_prefix, remaining_prefix, child): (
                SmallBinary,
                &[u8],
                ShardedMapChild<Value>,
            )| {
                async move {
                    Self::sharded_traversal_step(
                        ctx,
                        blobstore,
                        cur_prefix,
                        remaining_prefix,
                        child,
                    )
                    .await
                }
                .boxed()
            },
        )
    }

    async fn sharded_traversal_step<'a>(
        ctx: &'a CoreContext,
        blobstore: &'a impl Blobstore,
        mut cur_prefix: SmallBinary,
        remaining_prefix: &'a [u8],
        child: ShardedMapChild<Value>,
    ) -> Result<
        Vec<
            OrderedTraversal<
                ShardedTraversalOutput<'a, Value>,
                (SmallBinary, &'a [u8], ShardedMapChild<Value>),
            >,
        >,
    > {
        // When this `ShardedMapChild` is inlined, we're not really loading from the blobstore but we fetch from memory instead
        Ok(match child.load(ctx, blobstore).await? {
            // Case 1. Prepend all keys with cur_prefix and output elements
            Self::Terminal { values } => values
                .into_iter()
                .filter(|(k, _)| k.starts_with(remaining_prefix))
                .map(|(key, value)| {
                    let mut full_key = cur_prefix.clone();
                    full_key.extend(key);
                    OrderedTraversal::Output(ShardedTraversalOutput::Value(full_key, value))
                })
                .collect::<Vec<_>>(),
            // Case 2. Recurse
            Self::Intermediate {
                prefix: new_prefix,
                value,
                edges,
                ..
            } => {
                let remaining_prefix = if remaining_prefix.len() >= new_prefix.len() {
                    if let Some(new_remaining) =
                        remaining_prefix.strip_prefix(new_prefix.as_slice())
                    {
                        new_remaining
                    } else {
                        // prefix doesn't match
                        return Ok(vec![]);
                    }
                } else if new_prefix.starts_with(remaining_prefix) {
                    &[]
                } else {
                    // prefix doesn't match
                    return Ok(vec![]);
                };
                // Step 2-a. Extend cur_prefix
                cur_prefix.extend(new_prefix);
                let cur_prefix = &cur_prefix;
                remaining_prefix
                    .is_empty()
                    .then_some(value)
                    .flatten()
                    // Step 2-b. If value is present (and prefix empty), output (cur_prefix, value)
                    .map(|value| {
                        OrderedTraversal::Output(ShardedTraversalOutput::Value(
                            cur_prefix.clone(),
                            value,
                        ))
                    })
                    .into_iter()
                    // Step 2-c. Copy prefix, append byte, and recurse.
                    .chain(edges.into_iter().filter_map(|(byte, edge)| {
                        let (first, rest) = remaining_prefix.split_first().unwrap_or((&byte, &[]));
                        if *first == byte {
                            let mut new_prefix = cur_prefix.clone();
                            new_prefix.push(byte);
                            let size_prediction = edge.size;
                            // TODO (Pierre): don't recurse here, but this is where we
                            // return the shard and maybe its index, maybe its
                            // prefix, maybe both
                            match edge.child {
                                ShardedMapChild::Inlined(_) => Some(OrderedTraversal::Recurse(
                                    size_prediction,
                                    (new_prefix, rest, edge.child),
                                )),
                                // If this shard lives in the blobstore, return its id
                                // alongside the current prefix and the rest so we
                                // can resume the traversal later
                                ShardedMapChild::Id(id) => Some(OrderedTraversal::Output(
                                    ShardedTraversalOutput::Shard(new_prefix, rest, id),
                                )),
                            }
                        } else {
                            // Byte didn't match prefix
                            None
                        }
                    }))
                    .collect::<Vec<_>>()
            }
        })
    }

    pub fn is_empty(&self) -> bool {
        match self {
            Self::Terminal { values } => values.is_empty(),
            Self::Intermediate { .. } => self.size() == 0,
        }
    }

    fn size(&self) -> usize {
        match self {
            Self::Terminal { values } => values.len(),
            Self::Intermediate {
                value, edges, size, ..
            } => *size.get_or_init(|| {
                value.iter().len() + edges.values().map(|edge| edge.size).sum::<usize>()
            }),
        }
    }
}

pub enum ShardedTraversalOutput<'a, Value: MapValue> {
    /// We found a value. Return the key/value pair
    Value(SmallBinary, Value),
    /// We found a non-inlined shard. Return the current prefix, the remainder and the shard id.
    Shard(SmallBinary, &'a [u8], Value::Id),
}

impl<Value: MapValue> ThriftConvert for ShardedMapNode<Value> {
    const NAME: &'static str = "ShardedMapNode";
    type Thrift = thrift::ShardedMapNode;

    fn from_thrift(t: thrift::ShardedMapNode) -> Result<Self> {
        Ok(match t {
            thrift::ShardedMapNode::intermediate(intermediate) => Self::Intermediate {
                prefix: intermediate.prefix.0,
                value: intermediate
                    .value
                    .as_ref()
                    .map(ThriftConvert::from_bytes)
                    .transpose()?,
                edges: intermediate
                    .edges
                    .into_iter()
                    .map(|(k, e)| Ok((k as u8, ShardedMapEdge::from_thrift(e)?)))
                    .collect::<Result<_>>()?,
                size: Default::default(),
            },
            thrift::ShardedMapNode::terminal(terminal) => Self::Terminal {
                values: terminal
                    .values
                    .into_iter()
                    .map(|(k, v)| Ok((k.0, Value::from_bytes(&v)?)))
                    .collect::<Result<_>>()?,
            },
            thrift::ShardedMapNode::UnknownField(_) => bail!("Unknown map node variant"),
        })
    }

    fn into_thrift(self) -> thrift::ShardedMapNode {
        match self {
            Self::Intermediate {
                prefix,
                value,
                edges,
                ..
            } => thrift::ShardedMapNode::intermediate(thrift::ShardedMapIntermediateNode {
                prefix: thrift::small_binary(prefix),
                value: value.map(ThriftConvert::into_bytes),
                edges: edges
                    .into_iter()
                    .map(|(k, e)| (k as i8, e.into_thrift()))
                    .collect(),
            }),
            Self::Terminal { values } => {
                thrift::ShardedMapNode::terminal(thrift::ShardedMapTerminalNode {
                    values: values
                        .into_iter()
                        .map(|(k, v)| (thrift::small_binary(k), v.into_bytes()))
                        .collect(),
                })
            }
        }
    }
}

impl<Value: MapValue> BlobstoreValue for ShardedMapNode<Value> {
    type Key = Value::Id;

    fn into_blob(self) -> Blob<Self::Key> {
        let data = self.into_bytes();
        let id = Value::Context::id_from_data(&data);
        Blob::new(id, data)
    }

    fn from_blob(blob: Blob<Self::Key>) -> Result<Self> {
        Self::from_bytes(blob.data())
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use async_trait::async_trait;
    use blobstore::BlobstoreKeyParam;
    use blobstore::BlobstoreKeyRange;
    use blobstore::BlobstoreKeySource;
    use blobstore::LoadableError;
    use context::CoreContext;
    use fbinit::FacebookInit;
    use futures::TryStreamExt;
    use memblob::Memblob;
    use pretty_assertions::assert_eq;
    use quickcheck::Arbitrary;
    use quickcheck::Gen;
    use quickcheck::QuickCheck;
    use quickcheck::TestResult;
    use quickcheck::Testable;
    use ShardedMapNode::*;

    use super::*;
    use crate::impl_typed_hash;
    use crate::private::Blake2;
    use crate::typed_hash::BlobstoreKey;

    #[derive(Debug, Clone, Copy, Eq, PartialEq)]
    pub struct MyType(i32);

    #[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug, Hash)]
    pub struct ShardedMapNodeMyId(Blake2);

    impl_typed_hash! {
        hash_type => ShardedMapNodeMyId,
        thrift_hash_type => thrift::ShardedMapNodeId,
        value_type => ShardedMapNode<MyType>,
        context_type => ShardedMapNodeMyContext,
        context_key => "mytype.mapnode",
    }

    impl MapValue for MyType {
        type Id = ShardedMapNodeMyId;
        type Context = ShardedMapNodeMyContext;
    }

    impl ThriftConvert for MyType {
        const NAME: &'static str = "MyType";
        type Thrift = i32;
        fn into_thrift(self) -> Self::Thrift {
            self.0
        }
        fn from_thrift(t: Self::Thrift) -> Result<Self> {
            Ok(MyType(t))
        }
    }

    type TestShardedMap = ShardedMapNode<MyType>;

    fn terminal(values: Vec<(&str, i32)>) -> TestShardedMap {
        ShardedMapNode::Terminal {
            values: values
                .into_iter()
                .map(|(k, v)| (SmallVec::from_slice(k.as_bytes()), MyType(v)))
                .collect(),
        }
    }

    fn intermediate(
        prefix: &str,
        value: Option<i32>,
        children: Vec<(char, TestShardedMap)>,
    ) -> TestShardedMap {
        ShardedMapNode::Intermediate {
            prefix: SmallVec::from_slice(prefix.as_bytes()),
            value: value.map(MyType),
            edges: children
                .into_iter()
                .map(|(c, v)| {
                    (
                        c as u32 as u8,
                        ShardedMapEdge {
                            size: v.size(),
                            child: ShardedMapChild::Inlined(v),
                        },
                    )
                })
                .collect(),
            size: Default::default(),
        }
    }

    /// Returns an example map based on the picture on https://fburl.com/2fqtp2rk
    fn example_map() -> TestShardedMap {
        let abac = terminal(vec![
            ("ab", 7),
            ("aba", 8),
            ("akkk", 9),
            ("ate", 10),
            ("axi", 11),
        ]);
        let abal = terminal(vec![("aba", 5), ("ada", 6)]);
        let a = intermediate("ba", Some(12), vec![('c', abac), ('l', abal)]);
        let o = terminal(vec![("miojo", 1), ("miux", 2), ("mundo", 3), ("mungal", 4)]);
        // root
        intermediate("", None, vec![('a', a), ('o', o)])
    }

    const EXAMPLE_ENTRIES: &[(&str, i32)] = &[
        ("aba", 12),
        ("abacab", 7),
        ("abacaba", 8),
        ("abacakkk", 9),
        ("abacate", 10),
        ("abacaxi", 11),
        ("abalaba", 5),
        ("abalada", 6),
        ("omiojo", 1),
        ("omiux", 2),
        ("omundo", 3),
        ("omungal", 4),
    ];

    fn assert_round_trip(map: TestShardedMap) {
        let map_t = map.clone().into_thrift();
        // This is not deep equality through blobstore
        assert_eq!(ShardedMapNode::from_thrift(map_t).unwrap(), map);
    }

    #[derive(Clone)]
    struct MapHelper(TestShardedMap, CoreContext, Memblob);
    impl MapHelper {
        fn size(&self) -> usize {
            self.0.size()
        }

        async fn lookup(&self, key: &str) -> Result<Option<i32>> {
            let v = self.0.lookup(&self.1, &self.2, key.as_bytes()).await?;
            Ok(v.map(|my_type| my_type.0))
        }

        fn entries(&self) -> impl Stream<Item = Result<(String, i32)>> + '_ {
            self.0
                .clone()
                .into_entries(&self.1, &self.2)
                .and_then(|(k, v)| async move { Ok((String::from_utf8(k.to_vec())?, v.0)) })
        }

        fn prefix_entries<'a>(
            &'a self,
            prefix: &'a str,
        ) -> impl Stream<Item = Result<(String, i32)>> + 'a {
            self.0
                .clone()
                .into_prefix_entries(&self.1, &self.2, prefix.as_bytes())
                .and_then(|(k, v)| async move { Ok((String::from_utf8(k.to_vec())?, v.0)) })
        }

        async fn assert_entries(&self, entries: &[(&str, i32)]) -> Result<()> {
            assert_eq!(
                self.entries().try_collect::<Vec<_>>().await?,
                entries
                    .iter()
                    .map(|(k, v)| (String::from(*k), *v))
                    .collect::<Vec<_>>()
            );
            Ok(())
        }

        async fn assert_prefix_entries(&self, prefix: &str, entries: &[(&str, i32)]) -> Result<()> {
            assert_eq!(
                self.prefix_entries(prefix).try_collect::<Vec<_>>().await?,
                entries
                    .iter()
                    .map(|(k, v)| (String::from(*k), *v))
                    .collect::<Vec<_>>()
            );
            Ok(())
        }

        async fn add_remove(
            &mut self,
            to_add: &[(&str, i32)],
            to_remove: &[&str],
        ) -> Result<Vec<i32>> {
            let map = std::mem::take(&mut self.0);
            let (send, recv) = crossbeam::channel::unbounded();
            self.0 = map
                .update(
                    &self.1,
                    &self.2,
                    to_add
                        .iter()
                        .map(|(k, v)| (Bytes::copy_from_slice(k.as_bytes()), Some(MyType(*v))))
                        .chain(
                            to_remove
                                .iter()
                                .map(|k| (Bytes::copy_from_slice(k.as_bytes()), None)),
                        )
                        .collect(),
                    |x| send.send(x.0).unwrap(),
                )
                .await?;
            self.validate().await?;
            Ok(recv.try_iter().collect())
        }

        #[async_recursion]
        async fn validate(&self) -> Result<()> {
            let size = self.size();
            assert_eq!(self.0.is_empty(), size == 0);
            match &self.0 {
                Terminal { values } => assert!(values.len() <= 5),
                Intermediate {
                    prefix: _,
                    value,
                    edges,
                    size: _,
                } => {
                    let children_size: usize = stream::iter(
                        edges
                            .into_iter()
                            .map(|(_, e)| async move {
                                anyhow::Ok((e.size, e.child.clone().load(&self.1, &self.2).await?))
                            })
                            // prevent compiler bug
                            .collect::<Vec<_>>(),
                    )
                    .buffer_unordered(100)
                    .and_then(|(size, child)| async move {
                        let child = Self(child, self.1.clone(), self.2.clone());
                        child.validate().await?;
                        assert_eq!(child.size(), size);
                        Ok(size)
                    })
                    .try_collect::<Vec<_>>()
                    .await?
                    .into_iter()
                    .sum();
                    assert_eq!(children_size + value.iter().len(), size);
                }
            }
            Ok(())
        }

        #[async_recursion]
        async fn inner_inline_all(
            &self,
            mut map: ShardedMapNode<MyType>,
        ) -> Result<ShardedMapNode<MyType>> {
            match &mut map {
                Terminal { .. } => {}
                Intermediate { edges, .. } => {
                    for (_, edge) in edges {
                        let node = std::mem::take(&mut edge.child)
                            .load(&self.1, &self.2)
                            .await?;
                        edge.child = ShardedMapChild::Inlined(self.inner_inline_all(node).await?);
                    }
                }
            }
            Ok(map)
        }

        async fn inline_all(&mut self) -> Result<()> {
            let map = std::mem::take(&mut self.0);
            self.0 = self.inner_inline_all(map).await?;
            Ok(())
        }

        async fn child(&self, key: char) -> Result<Self> {
            let child = match &self.0 {
                Terminal { .. } => bail!("terminal"),
                Intermediate { edges, .. } => {
                    edges
                        .get(&(key as u8))
                        .unwrap()
                        .child
                        .clone()
                        .load(&self.1, &self.2)
                        .await?
                }
            };
            Ok(Self(child, self.1.clone(), self.2.clone()))
        }

        fn assert_terminal(&self, values_len: usize) {
            match &self.0 {
                Intermediate { .. } => panic!("not terminal"),
                Terminal { values } => assert_eq!(values.len(), values_len),
            }
        }
        fn assert_intermediate(&self, child_count: usize) {
            match &self.0 {
                Terminal { .. } => panic!("not intermediate"),
                Intermediate { edges, .. } => assert_eq!(edges.len(), child_count),
            }
        }
        fn assert_prefix(&self, prefix: &str) {
            match &self.0 {
                Terminal { .. } => panic!("not intermediate"),
                Intermediate {
                    prefix: my_prefix, ..
                } => assert_eq!(my_prefix.as_slice(), prefix.as_bytes()),
            }
        }
    }

    #[test]
    fn basic_test() {
        let empty = ShardedMapNode::<MyType>::Terminal {
            values: Default::default(),
        };
        assert!(empty.is_empty());
        assert_eq!(empty.size(), 0);
        let empty = ShardedMapNode::<MyType>::Intermediate {
            value: None,
            edges: Default::default(),
            prefix: Default::default(),
            size: Default::default(),
        };
        assert!(empty.is_empty());
        assert_eq!(empty.size(), 0);
        assert!(ShardedMapNode::<MyType>::default().is_empty());

        let map = terminal(vec![("ab", 3), ("cd", 5)]);
        assert!(!map.is_empty());
        assert_round_trip(map);

        let map = example_map();
        assert!(!map.is_empty());
        assert_eq!(map.size(), 12);
        assert_round_trip(map);
    }

    #[fbinit::test]
    async fn lookup_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();

        let map = MapHelper(example_map(), ctx, blobstore);
        // Case 2 > Case 1
        assert_eq!(map.lookup("omiux").await?, Some(2));
        // Case 3
        assert_eq!(map.lookup("inexistent").await?, None);
        // Case 2 > Case 5
        assert_eq!(map.lookup("abxio").await?, None);
        // Case 2 > Case 4
        assert_eq!(map.lookup("aba").await?, Some(12));
        // Case 2 > Case 2 > Case 1
        assert_eq!(map.lookup("abacakkk").await?, Some(9));
        assert_eq!(map.lookup("abacakk").await?, None);
        // Case 4
        assert_eq!(map.lookup("").await?, None);
        Ok(())
    }

    #[fbinit::test]
    async fn into_entries_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();

        let map = MapHelper(example_map(), ctx, blobstore);
        map.assert_entries(EXAMPLE_ENTRIES).await?;
        map.assert_prefix_entries("", EXAMPLE_ENTRIES).await?;
        map.assert_prefix_entries("aba", &EXAMPLE_ENTRIES[0..8])
            .await?;
        map.assert_prefix_entries("abaca", &EXAMPLE_ENTRIES[1..6])
            .await?;
        map.assert_prefix_entries("omi", &EXAMPLE_ENTRIES[8..10])
            .await?;
        map.assert_prefix_entries("om", &EXAMPLE_ENTRIES[8..])
            .await?;
        map.assert_prefix_entries("o", &EXAMPLE_ENTRIES[8..])
            .await?;
        map.assert_prefix_entries("ban", &[]).await?;
        Ok(())
    }

    async fn get_all_keys(
        ctx: &CoreContext,
        blobstore: &impl BlobstoreKeySource,
    ) -> Result<impl Iterator<Item = String>> {
        let data = blobstore
            .enumerate(
                ctx,
                &BlobstoreKeyParam::Start(BlobstoreKeyRange {
                    begin_key: String::new(),
                    end_key: String::new(),
                }),
            )
            .await?;
        if data.next_token.is_some() {
            unimplemented!();
        }
        let mut data: Vec<_> = data.keys.into_iter().collect();
        data.sort();
        Ok(data.into_iter())
    }

    async fn assert_all_keys(
        ctx: &CoreContext,
        blobstore: &impl BlobstoreKeySource,
        keys: Vec<&str>,
    ) -> Result<()> {
        assert_eq!(
            get_all_keys(ctx, blobstore).await?.collect::<Vec<_>>(),
            keys.into_iter().map(String::from).collect::<Vec<_>>()
        );
        Ok(())
    }

    async fn assert_key_count(
        ctx: &CoreContext,
        blobstore: &impl BlobstoreKeySource,
        count: usize,
    ) -> Result<()> {
        assert_eq!(get_all_keys(ctx, blobstore).await?.count(), count);
        Ok(())
    }

    #[fbinit::test]
    async fn store_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        let map = example_map();
        map.into_blob().store(&ctx, &blobstore).await?;
        assert_all_keys(&ctx, &blobstore, vec!["mytype.mapnode.blake2.9239707907ceb346d7146c601f131ab7c598a8e98441c2934817c964f0a2c270"]).await?;
        Ok(())
    }

    #[fbinit::test]
    async fn update_basic_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        let mut map = MapHelper(Default::default(), ctx.clone(), blobstore.clone());
        map.assert_entries(&[]).await?;
        map.add_remove(EXAMPLE_ENTRIES, &[]).await?;
        map.assert_entries(EXAMPLE_ENTRIES).await?;
        assert_all_keys(
            &ctx,
            &blobstore,
            vec!["mytype.mapnode.blake2.8c8f56e9612f7cc94187729e4eff067bf56bb239019e25c8421243a60e4d1fb9",
                "mytype.mapnode.blake2.e808435da65aa2e0f61db333fba1904e57b1b46dff2d3a5c263d0016750f1f0d",
                "mytype.mapnode.blake2.f8728f3aabc9b8083db7150e9c2636c2971dac0784fac59cee3b7b6908165476"],
        )
        .await?;
        {
            // Let's compare it to our hand-written map
            let mut map = map.clone();
            map.inline_all().await?;
            map.assert_entries(EXAMPLE_ENTRIES).await?;
            assert_eq!(map.0, example_map());
        }
        let deleted = map.add_remove(&[], &["abalaba", "non_existing"]).await?;
        assert_eq!(deleted, vec![5]);
        assert_eq!(map.0.size(), EXAMPLE_ENTRIES.len() - 1);
        assert_key_count(&ctx, &blobstore, 4).await?;
        let deleted = map.add_remove(&[], &["abalada"]).await?;
        assert_eq!(deleted, vec![6]);
        // Intermeditate node should now have 1 child, but also a value
        assert_key_count(&ctx, &blobstore, 4).await?;
        let child = map.child('a').await?;
        match child.0 {
            Terminal { .. } => bail!("not intermediate"),
            Intermediate { value, edges, .. } => {
                assert!(value.is_some());
                assert_eq!(edges.len(), 1);
            }
        }
        let deleted = map.add_remove(&[], &["aba"]).await?;
        assert_eq!(deleted, vec![12]);
        // Intermediate node without a value should be merged
        assert_key_count(&ctx, &blobstore, 5).await?;
        map.assert_entries(&[
            ("abacab", 7),
            ("abacaba", 8),
            ("abacakkk", 9),
            ("abacate", 10),
            ("abacaxi", 11),
            ("omiojo", 1),
            ("omiux", 2),
            ("omundo", 3),
            ("omungal", 4),
        ])
        .await?;
        map.child('a').await?.assert_terminal(5);
        assert!(
            map.add_remove(&[], &["abalada", "abalaba", "aba"])
                .await?
                .is_empty()
        );
        assert_eq!(
            map.add_remove(&[("potato", 1000), ("abacaxi", 1001)], &[])
                .await?,
            vec![11]
        );
        Ok(())
    }

    #[fbinit::test]
    async fn update_tricky_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        let mut map = MapHelper(Default::default(), ctx.clone(), blobstore.clone());
        map.add_remove(
            &[
                ("A11", 1),
                ("A12", 2),
                ("A13", 3),
                ("A21", 1),
                ("A22", 2),
                ("A23", 3),
            ],
            &[],
        )
        .await?;
        map.assert_intermediate(2);
        map.child('1').await?.assert_terminal(3);
        map.child('2').await?.assert_terminal(3);
        // LCP of keys is smaller than prefix only due to removals.
        map.add_remove(&[("A14", 4)], &["cz", "A", "A31"]).await?;
        map.assert_intermediate(2);
        map.child('1').await?.assert_terminal(4);
        map.child('2').await?.assert_terminal(3);
        map.add_remove(
            &[
                ("B11", 1),
                ("B21", 1),
                ("B22", 2),
                ("B23", 3),
                ("B24", 4),
                ("B31", 1),
            ],
            &["A11", "A12", "A13", "A14", "A21", "A22", "A23"],
        )
        .await?;
        map.assert_intermediate(3);
        map.child('1').await?.assert_terminal(1);
        map.child('2').await?.assert_terminal(4);
        map.child('3').await?.assert_terminal(1);
        map.add_remove(&[], &[""]).await?;
        Ok(())
    }

    #[fbinit::test]
    async fn update_tricky_deletes_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        let mut map = MapHelper(Default::default(), ctx.clone(), blobstore.clone());
        map.add_remove(EXAMPLE_ENTRIES, &[]).await?;
        // Removing something that is not a prefix of the intermediate node
        // can panic if not done correctly
        map.add_remove(&[], &[""]).await?;
        map.add_remove(&[], &["a"]).await?;
        map.add_remove(&[], &["ab"]).await?;
        // Bug where we might mismatch prefix of deleted keys
        map.add_remove(&[], &["abx"]).await?;
        assert_eq!(map.size(), 12);
        map.add_remove(&[], &["abxlada"]).await?;
        assert_eq!(map.size(), 12);
        // Let's play with the value of an intermediate node and assert all is still good:
        map.add_remove(&[], &["aba"]).await?;
        assert_eq!(map.size(), 11);
        let child = map.child('a').await?;
        assert_eq!(child.size(), 7);
        assert_eq!(map.lookup("aba").await?, None);
        map.add_remove(&[("aba", 0)], &[]).await?;
        assert_eq!(map.size(), 12);
        assert_eq!(map.lookup("aba").await?, Some(0));
        map.add_remove(&[("aba", -1)], &[]).await?;
        assert_eq!(map.size(), 12);
        assert_eq!(map.lookup("aba").await?, Some(-1));
        Ok(())
    }

    #[fbinit::test]
    async fn update_cases_test(fb: FacebookInit) -> Result<()> {
        // Let's try to do updates that cause different cases and assert it all works out
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        let mut map = MapHelper(Default::default(), ctx.clone(), blobstore.clone());
        // Case 1
        map.add_remove(&[("_a", 1), ("_b", 2), ("_c", 3)], &[])
            .await?;
        map.assert_terminal(3);
        // Case 2 > Case 4 > (Recursive Case 1's) > Case 4.3.2
        map.add_remove(&[("_d", 4), ("_e", 5), ("_f", 6)], &[])
            .await?;
        map.assert_intermediate(6);
        map.child('d').await?.assert_terminal(1);
        // Case 4 > (Recursive Case 1) > case 4.3.1
        map.add_remove(&[], &["_b"]).await?;
        map.assert_terminal(5);
        // Case 2 > Case 3 > Case 4 > (Recursive Case 1, Case 2>...) > Case 4.3.2
        map.add_remove(&[("_b", 2), ("z", -1)], &[]).await?;
        map.assert_intermediate(2);
        map.assert_prefix("");
        map.child('_').await?.assert_intermediate(6);
        map.child('z').await?.assert_terminal(1);
        // Case 4 > (Rec Case 1) > Case 4.3.2 + merge
        map.add_remove(&[], &["z"]).await?;
        map.assert_intermediate(6);
        map.assert_prefix("_");
        // Case 3 > Case 4 > Rec Case 1 > Case 4.3.2
        map.add_remove(&[("y", -2)], &[]).await?;
        map.assert_intermediate(2);
        map.assert_prefix("");
        // Case 4 > (Rec Case 1, Case 4>Rec Case1>Case 4.3.1) > Case 4.3.1
        map.add_remove(&[], &["_b", "y"]).await?;
        map.assert_terminal(5);
        map.assert_entries(&[("_a", 1), ("_c", 3), ("_d", 4), ("_e", 5), ("_f", 6)])
            .await?;
        Ok(())
    }

    #[fbinit::test]
    async fn update_case_3_test(fb: FacebookInit) -> Result<()> {
        // Let's try an update that causes case 3 and do detailed asserting
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        let mut map = MapHelper(Default::default(), ctx.clone(), blobstore.clone());
        map.add_remove(
            &[
                ("abc1", 1),
                ("abc2", 2),
                ("abc3", 3),
                ("abc4", 4),
                ("abc5", 5),
                ("abc6", 6),
            ],
            &[],
        )
        .await?;
        map.assert_intermediate(6);
        map.assert_prefix("abc");
        map.add_remove(&[("a1", 1)], &[]).await?;
        map.assert_intermediate(2);
        map.assert_prefix("a");
        let childb = map.child('b').await?;
        childb.assert_prefix("c");
        childb.assert_intermediate(6);
        let child1 = map.child('1').await?;
        child1.assert_terminal(1);
        map.assert_entries(&[
            ("a1", 1),
            ("abc1", 1),
            ("abc2", 2),
            ("abc3", 3),
            ("abc4", 4),
            ("abc5", 5),
            ("abc6", 6),
        ])
        .await?;
        Ok(())
    }

    #[fbinit::test]
    fn round_trip_quickcheck(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let blobstore = Memblob::default();
        use tokio::runtime::Runtime;

        struct Roundtrip(Runtime, CoreContext, Memblob);
        impl Testable for Roundtrip {
            fn result(&self, gen: &mut Gen) -> TestResult {
                let res = self.0.block_on(async {
                    let values: BTreeMap<String, i32> = Arbitrary::arbitrary(gen);
                    let mut queries: Vec<String> = Arbitrary::arbitrary(gen);
                    let keys: Vec<&String> = values.keys().collect();
                    for _ in 0..values.len() / 2 {
                        queries.push(gen.choose(&keys).unwrap().to_string());
                    }
                    let mut map = MapHelper(Default::default(), self.1.clone(), self.2.clone());
                    let adds = values
                        .iter()
                        .map(|(k, v)| (k.as_str(), *v))
                        .collect::<Vec<_>>();
                    map.add_remove(&adds, &[]).await?;
                    if map.size() != values.len() {
                        return Ok(false);
                    }
                    for k in queries {
                        let correct_v = values.get(&k);
                        let test_v = map.lookup(&k).await?;
                        if correct_v.cloned() != test_v {
                            return Ok(false);
                        }
                    }
                    let test_roundtrip = map.entries().try_collect::<BTreeMap<_, _>>().await?;
                    Ok(values == test_roundtrip)
                });
                TestResult::from_bool(matches!(res, Result::Ok(true)))
            }
        }

        QuickCheck::new().quickcheck(Roundtrip(Runtime::new().unwrap(), ctx, blobstore));
    }
}