/*
 * 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.
 */

//! Blob store on local disk.
//!
//! See [Zstore] for the main structure.

use std::borrow::Cow;
use std::fmt;
use std::path::Path;
use std::path::PathBuf;

use indexedlog::log as ilog;
use indexedlog::DefaultOpenOptions;
use indexedlog::OpenWithRepair;
use lazy_static::lazy_static;
use lru_cache::LruCache;
use minibytes::Bytes;
use parking_lot::Mutex;
use serde::Deserialize;
use serde::Serialize;
use sha1::Digest;
use sha1::Sha1;
use tracing::debug_span;
use tracing::info_span;
use tracing::trace_span;
pub use types::Id20;

/// An append-only local-disk blob storage.
///
/// Blobs are addressed by their SHA1 digest, compressed using zstd algorithm.
/// Blobs with similar contents can be stored as a chain of zstd dictionary
/// compression results to save space. That is, a newer version is compressed
/// using an existing version as a zstd dictionary.
///
/// The name `Zstore` was chosen because the prefix `zst` is the name of the
/// compression algorithm.
pub struct Zstore {
    dir: PathBuf,
    log: ilog::Log,
    pub delta_opts: DeltaOptions,
    cache: Mutex<LruCache<Id20, Bytes>>,
}

impl DefaultOpenOptions<ilog::OpenOptions> for Zstore {
    fn default_open_options() -> ilog::OpenOptions {
        ilog::OpenOptions::new()
            .index("id", |_| -> Vec<_> {
                // The offset of `start` should match mincode serialization layout
                // of `Delta`.
                let start = 0;
                let end = start + Id20::len();
                vec![ilog::IndexOutput::Reference(start as u64..end as u64)]
            })
            .create(true)
            .flush_filter(Some(|context, data| {
                // At flush time, there might be data written by other processes.
                // Drop duplicated entries.
                let id = &data[0..Id20::len()];
                if let Ok(mut iter) = context.log.lookup(Self::ID20_INDEX, id) {
                    if iter.nth(0).is_some() {
                        return Ok(ilog::FlushFilterOutput::Drop);
                    }
                }
                Ok(ilog::FlushFilterOutput::Keep)
            }))
    }
}

/// Options for opening a [`Zstore`].
pub struct OpenOptions {
    cache_size: usize,
}

impl Default for OpenOptions {
    fn default() -> Self {
        Self { cache_size: 1 }
    }
}

impl OpenOptions {
    /// Constructs [`OpenOptions`] with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the LRU cache size.
    pub fn cache_size(mut self, size: usize) -> Self {
        self.cache_size = size;
        self
    }

    /// Load or create [`Zstore`] at the given directory.
    pub fn open(&self, dir: &Path) -> crate::Result<Zstore> {
        let log = Zstore::default_open_options().open_with_repair(dir)?;
        // The cache_size should be greater than len(metalog.keys()))
        Ok(Zstore {
            dir: dir.to_path_buf(),
            log,
            delta_opts: Default::default(),
            cache: Mutex::new(LruCache::new(self.cache_size)),
        })
    }
}

impl Zstore {
    const ID20_INDEX: usize = 0;

    /// Load or create [`Zstore`] at the given directory.
    /// Use the default [`OpenOptions`].
    pub fn open(dir: impl AsRef<Path>) -> crate::Result<Zstore> {
        OpenOptions::default().open(dir.as_ref())
    }

    /// Insert a new blob to the store. Return its identity (SHA1).
    ///
    /// `candidate_base_ids` can be used to provide a list of "similar" blob
    /// identities that can be used as zstd dictionaries to achieve better
    /// compression. Note: delta base outside the candidate list can also be
    /// chosen to optimize the delta chain length.
    ///
    /// Writes are buffered. Call [Zstore::flush] to actually write them to
    /// disk. See [indexedlog::Log::append] and [indexedlog::Log::sync] for
    /// details.
    pub fn insert(&mut self, data: &[u8], candidate_base_ids: &[Id20]) -> crate::Result<Id20> {
        let id = sha1(data);

        if self.contains(id)? {
            return Ok(id);
        }

        debug_span!(
            "Zstore::insert",
            data_len = data.len(),
            id = &AsRef::<str>::as_ref(&id.to_hex())
        )
        .in_scope(|| {
            // Base line: delta against b"".
            let compressed = zstdelta::diff(b"", data)?;
            let chain_bytes = compressed.len();
            let mut best_delta = Delta {
                id,
                base_id: *EMPTY_ID20,
                depth: 1, // "EMPTY" has depth 0
                subchain_len: 0,
                chain_bytes,
                data: Cow::Owned(compressed),
            };

            // Attempt to use delta bases.
            for &base_id in candidate_base_ids {
                if let Some(delta) = self.create_delta(id, base_id, data, false)? {
                    if delta.data.len() < best_delta.data.len() {
                        best_delta = delta;
                    }
                }
            }

            // Insert the delta to the blob store.
            let bytes = mincode::serialize(&best_delta)?;
            self.log.append(bytes)?;
            Ok(id)
        })
    }

    /// Write blobs to disk.
    ///
    /// See [indexedlog::Log::sync] for details.
    /// Return the size of the main log in bytes.
    pub fn flush(&mut self) -> crate::Result<u64> {
        info_span!("Zstore::flush").in_scope(|| Ok(self.log.flush()?))
    }

    /// Get the content of the specified blob.
    pub fn get(&self, id: Id20) -> crate::Result<Option<Bytes>> {
        debug_span!("Zstore::get", id = &AsRef::<str>::as_ref(&id.to_hex())).in_scope(|| match self
            .get_delta(id)?
        {
            None => Ok(None),
            Some(delta) => Ok(Some(self.resolve(delta)?)),
        })
    }

    /// Check if the store contains the given blob identity.
    pub fn contains(&self, id: Id20) -> crate::Result<bool> {
        debug_span!("Zstore::contains", id = &AsRef::<str>::as_ref(&id.to_hex())).in_scope(|| {
            let mut results = self.log.lookup(Self::ID20_INDEX, id)?;
            Ok(results.next().transpose().map(|v| v.is_some())?)
        })
    }

    /// Create a new [`Delta`] using the specified delta base candidate.
    /// Satisfy limitations specified by [`DeltaOptions`].
    /// Return `None` if a suitable delta cannot be created.
    fn create_delta<'a>(
        &'a self,
        id: Id20,
        base_id: Id20,
        data: &'a [u8],
        mut preserve_depth: bool,
    ) -> crate::Result<Option<Delta<'a>>> {
        if base_id == *EMPTY_ID20 {
            // Pointless to create a delta against EMPTY_ID20.
            return Ok(None);
        }

        let base_delta = match self.get_delta(base_id)? {
            Some(v) => v,
            None => {
                // base_id does not exist.
                return Ok(None);
            }
        };

        // See docstring above Delta.depth for how this works.
        if base_delta.depth >= self.delta_opts.max_depth {
            // Cannot go deeper.
            preserve_depth = true;
        }

        if preserve_depth && base_delta.subchain_len + 1 >= self.delta_opts.max_subchain_len {
            // Creating a delta based on base_id will exceed the max_subchain_len
            // limit. So attempt to use a "parent" delta instead.
            let current_depth = base_delta.depth;
            let mut current_parent = base_delta;
            while let Some(parent) = self.get_delta(current_parent.base_id)? {
                if parent.id == *EMPTY_ID20 {
                    // Avoid infinite loop.
                    return Ok(None);
                }
                if parent.depth < current_depth {
                    // Find a parent. But it can also exceeds the
                    // max_subchain_len limit. So call create_delta
                    // recursively.
                    return self.create_delta(id, parent.id, data, true);
                } else {
                    current_parent = parent
                }
            }
            Ok(None)
        } else {
            let (depth, subchain_len) = if preserve_depth {
                // Create a new delta at the current depth.
                (base_delta.depth, base_delta.subchain_len + 1)
            } else {
                // Haven't reached the maximum delta depth. So attempt to
                // increase the delta depth here.
                (base_delta.depth + 1, 0)
            };
            let base_bytes = base_delta.chain_bytes;
            let bytes = zstdelta::diff(&self.resolve(base_delta)?, data)?;
            let chain_bytes = base_bytes + bytes.len();
            if chain_bytes
                > self
                    .delta_opts
                    .max_chain_bytes
                    .min(data.len() << self.delta_opts.max_chain_factor_log)
            {
                // The delta length is not suitable for a delta.
                Ok(None)
            } else {
                // Create a new delta at the next depth.
                Ok(Some(Delta {
                    id,
                    base_id,
                    depth,
                    subchain_len,
                    chain_bytes,
                    data: Cow::Owned(bytes),
                }))
            }
        }
    }

    /// Decode a [`Delta`]. Do not apply delta chain to get full text.
    fn get_delta<'a>(&'a self, id: Id20) -> crate::Result<Option<Delta<'a>>> {
        if id == *EMPTY_ID20 {
            return Ok(Some(Delta {
                id,
                base_id: id,
                depth: 0,
                subchain_len: 0,
                chain_bytes: 0,
                data: Cow::Borrowed(b""),
            }));
        }
        let mut results = self.log.lookup(0, id)?;
        match results.next() {
            None => Ok(None),
            Some(Ok(bytes)) => {
                let result = mincode::deserialize(bytes)?;
                Ok(Some(result))
            }
            Some(Err(err)) => Err(err.into()),
        }
    }

    /// Apply delta chains recursively to reconstruct full text.
    fn resolve<'a>(&'a self, delta: Delta<'a>) -> crate::Result<Bytes> {
        if delta.id == *EMPTY_ID20 {
            return Ok(Bytes::from_static(b""));
        }
        {
            let mut cache = self.cache.lock();
            if let Some(content) = cache.get_mut(&delta.id) {
                return Ok(content.clone());
            }
        }

        trace_span!(
            "Zstore::resolve",
            id = &AsRef::<str>::as_ref(&delta.id.to_hex()),
            base_id = &AsRef::<str>::as_ref(&delta.base_id.to_hex()),
            depth = delta.depth,
            subchain_len = delta.subchain_len,
            chain_bytes = delta.chain_bytes,
            data_len = delta.data.len(),
        )
        .in_scope(|| {
            match self.get_delta(delta.base_id)? {
                Some(base_delta) => {
                    // PERF: some caching would avoid N^2 chain application.
                    let base_bytes = self.resolve(base_delta)?;
                    let bytes = Bytes::from(zstdelta::apply(&base_bytes, &delta.data)?);
                    {
                        let mut cache = self.cache.lock();
                        cache.insert(delta.id, bytes.clone());
                    }
                    Ok(bytes)
                }
                None => Err(self.error(format!(
                    "incomplete delta chain: {} -> {} chain is broken",
                    delta.id.to_hex(),
                    delta.base_id.to_hex()
                ))),
            }
        })
    }

    fn error(&self, message: impl fmt::Display) -> crate::Error {
        crate::Error(format!("{:?}: {}", &self.dir, message))
    }
}

// -------- Debug APIs --------

/// Represent delta relationships.
pub struct DebugDeltaTree {
    id: Id20,
    chain_len: usize,
    depth: usize,
    subchain_len: usize,
    children: Vec<DebugDeltaTree>,
}

#[allow(dead_code)]
impl Zstore {
    /// Reconstruct a tree of deltas for debugging purpose.
    fn debug_delta_tree(&self) -> crate::Result<DebugDeltaTree> {
        let mut root = DebugDeltaTree {
            id: *EMPTY_ID20,
            chain_len: 0,
            depth: 0,
            subchain_len: 0,
            children: Vec::new(),
        };

        fn insert<'a>(tree: &'a mut DebugDeltaTree, delta: Delta) -> &'a mut DebugDeltaTree {
            let id = tree
                .children
                .iter()
                .enumerate()
                .find(|(_, c)| c.id == delta.id)
                .map(|(id, _)| id);
            match id {
                Some(id) => &mut tree.children[id],
                None => {
                    tree.children.push(DebugDeltaTree {
                        id: delta.id,
                        chain_len: tree.chain_len + 1,
                        depth: delta.depth,
                        subchain_len: delta.subchain_len,
                        children: Vec::new(),
                    });
                    tree.children.last_mut().unwrap()
                }
            }
        }

        for entry in self.log.iter() {
            let id = &self.log.index_func(Self::ID20_INDEX, entry?)?[0];
            let mut id = Id20::from_slice(id).unwrap();
            let mut chain: Vec<Delta> = Vec::new();
            while id != *EMPTY_ID20 {
                if let Some(delta) = self.get_delta(id)? {
                    id = delta.base_id;
                    chain.push(delta);
                } else {
                    return Err(crate::Error(format!(
                        "unexpected broken chain around {}",
                        id.to_hex()
                    )));
                }
            }

            let mut tree = &mut root;
            for delta in chain.into_iter().rev() {
                tree = insert(tree, delta);
            }
        }

        Ok(root)
    }
}

impl fmt::Debug for DebugDeltaTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.id == *EMPTY_ID20 {
            // Write header.
            write!(f, "Chain Len| Depth |Subchain Len| Chain ID\n")?;
        }
        write!(
            f,
            "{:8} | {:5} | {:10} | {}{}\n",
            self.chain_len,
            self.depth,
            self.subchain_len,
            " ".repeat(self.chain_len),
            &self.id.to_hex()[..6],
        )?;
        for child in self.children.iter() {
            write!(f, "{:?}", child)?;
        }
        Ok(())
    }
}

// -------- Utilities --------

pub fn sha1(data: &[u8]) -> Id20 {
    trace_span!("sha1", data_len = data.len()).in_scope(|| {
        let mut hasher = Sha1::new();
        hasher.update(data);
        let id: [u8; 20] = hasher.finalize().into();
        Id20::from_byte_array(id)
    })
}

// -------- Options --------

/// Options for deltas.
///
/// `max_depth` and `max_subchain_len` reshapes linear deltas to a tree,
/// reducing delta-chain length.
///
/// For example, the following graph shows 1 full text and 12 deltas.
/// The maximum chain length is 6. If those 12 deltas are chained
/// linearly, the maximum chain length would be 12.
///
/// ```plain,ignore
/// Tree            | Base  | Parent | Depth | Chain        | Sub Chain
/// --------------------------------------------------------------------
/// Rev 0           | EMPTY | -      | 0     |              |
///  +-- Rev 1      | Rev 0 | Rev 0  | 1     | 1            | 1
///  |    +-- Rev 2 | Rev 1 | Rev 1  | 2     | 1-2          | 2
///  |    +-- Rev 3 | Rev 2 | Rev 1  | 2     | 1-2-3        | 2-3
///  |    +-- Rev 4 | Rev 3 | Rev 1  | 2     | 1-2-3-4      | 2-3-4
///  +-- Rev 5      | Rev 1 | Rev 0  | 1     | 1-5          | 1-5
///  |    +-- Rev 6 | Rev 5 | Rev 5  | 2     | 1-5-6        | 6
///  |    ...       |       |        |       |              |
///  +-- Rev 9      | Rev 5 | Rev 0  | 1     | 1-5-9        | 9
///       ...       |       |        |       |              |
///       +-- Rev 12| Rev 9 | Rev 11 | 2     |1-5-9-10-11-12| 10-11-12
///                   ^       ^        ^       ^              ^
///                   |       |        |       |         revs in same depth
///                   |       |        |       see "Tree", path from root to
///                   |       |        |       rev, including siblings
///                   |       |        see "Tree"
///                   |       rev with one less depth
///                   delta base
///
/// In this example, `max_depth` is 2, `max_subchain_len` is 3.
/// ```
///
/// Let `d` be `max_depth`, `n` be `max_subchain_len`, this would rewrite a
/// linear chain of `Sum[n**i, {i, 1, d}]` (evaluates to
/// `(n * (n ** d - 1)) / (n - 1)`) deltas into a tree with a maximum chain
/// length `n * d`.
///
/// Table of linear chain length (i.e. a linear chain of such length can
/// be reshaped without using full texts):
///
/// ```plain,ignore
///     |    d=2      3      4      5      6      7
/// -----------------------------------------------
/// n=2 |      6     14     30     62    126    254
///   3 |     12     39    120    363   1092   3279
///   4 |     20     84    340   1364   5460  21844
///   5 |     30    155    780   3905  19530  97655
///   6 |     42    258   1554   9330  55986 335922
///   7 |     56    399   2800  19607 137256 960799
/// ```
///
/// Table of average reshaped chain length:
///
/// ```plain,ignore
///     |    d=2      3      4      5      6      7
/// -----------------------------------------------
/// n=2 |    2.5    3.6    4.9    6.2    7.6    9.1
///   3 |    3.5    5.2    7.1    9.0   11.0   13.0
///   4 |    4.5    6.8    9.2   11.7   14.2   16.7
///   5 |    5.5    8.3   11.3   14.3   17.3   20.3
///   6 |    6.5    9.8   13.3   16.8   20.3   23.8
///   7 |    7.5   11.4   15.3   19.3   23.3   27.3
/// ```
///
/// In general, a larger `n` helps space usage for shorter chains,
/// a larger `d` helps handling longer chains.
pub struct DeltaOptions {
    /// Maximum depth of a delta.
    ///
    /// See [`DeltaOptions`] for explanation.
    pub max_depth: usize,

    /// Maximum length of a subchain.
    ///
    /// See [`DeltaOptions`] for explanation.
    pub max_subchain_len: usize,

    /// Maximum bytes for the total compressed bytes in the delta chain.
    pub max_chain_bytes: usize,

    /// Do not use delta, if the total compressed bytes in the delta chain
    /// exceeds `uncompressed_source_data.len() << max_chain_factor_log`.
    pub max_chain_factor_log: u8,

    /// Prevent constructing this struct.
    _private: (),
}

impl Default for DeltaOptions {
    fn default() -> Self {
        DeltaOptions {
            max_depth: 5,
            max_subchain_len: 4,
            max_chain_bytes: 500_000_000,
            // Stop using delta if the compressed chain exceeds
            // 2x uncompressed full text size. This matches
            // mercurial revlog behavior.
            max_chain_factor_log: 1,
            _private: (),
        }
    }
}

lazy_static! {
    static ref EMPTY_ID20: Id20 = sha1(b"");
}

// -------- Serde Structures --------

#[derive(Serialize, Deserialize)]
struct Delta<'a> {
    // ATTENTION: The 0..20 byte slice of mincode serialized Delta is expected to be the id.
    // This also requires the hgid::tuple serialization strategy.
    /// Pre-calculated checksum of the full data.
    #[serde(with = "types::serde_with::hgid::tuple")]
    id: Id20,

    /// Delta base Id. Use `EMPTY_ID20` if there is no real delta base.
    #[serde(with = "types::serde_with::hgid::tuple")]
    base_id: Id20,

    /// Depth of the delta. See [`DeltaOptions`] for how this works.
    depth: usize,

    /// Count of deltas with the same depth in the current chain, excluding
    /// self.
    ///
    /// ```plain,ignore
    /// Depth:         0  1  1  1  2  2  2
    /// Chain:         _                   Depth 0
    ///                 \-o--o--o          Depth 1 (subchain)
    ///                          \-o--o--o Depth 2 (subchain)
    /// Chain Len:     0  1  2  3  4  5  6
    /// Subchain Len:  0  0  1  2  0  1  2
    /// ```
    ///
    /// See also `depth`.
    ///
    /// If `subchain_len` reaches the `max_subchain_len` limit. This delta will
    /// no longer be a delta base candidate, its parent (the nearest delta in
    /// the chain with a smaller depth) will be considered instead.
    subchain_len: usize,

    /// Delta chain size in total. Help cutting down a delta chain.
    chain_bytes: usize,

    /// Delta content.
    ///
    /// Output of [`zstdelta::diff`] (owned) or [`mincode::deserialize`]
    /// (borrowed).
    #[serde(borrow)]
    data: Cow<'a, [u8]>,
}

// -------- Tests --------

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

    use super::*;

    #[test]
    fn test_simple_case() {
        let contents = [
            &b"11111111111111111111111111111111"[..],
            &b"1111111111111111111111111111112"[..],
            &b"111111111111111111111111111113"[..],
            &b""[..],
        ];
        test_round_trips(&contents);
    }

    #[test]
    fn test_dedup_on_flush() {
        let dir = TempDir::new().unwrap();
        let mut zstore = Zstore::open(&dir).unwrap();
        let id = zstore.insert(b"123456", &[]).unwrap();
        zstore.flush().unwrap();

        let mut zstore1 = Zstore::open(&dir).unwrap();
        let mut zstore2 = Zstore::open(&dir).unwrap();
        let id1 = zstore1.insert(b"1234567", &[id]).unwrap();
        let id2 = zstore2.insert(b"1234567", &[]).unwrap();
        assert_eq!(id1, id2);

        // Because id1 == id2, they will be de-duplicated by the flush_filter
        // on ilog::Log::flush().
        let size1 = zstore1.flush().unwrap();
        let size2 = zstore2.flush().unwrap();
        assert_eq!(size1, size2);
    }

    quickcheck! {
        fn test_random_round_trips(contents: Vec<Vec<u8>>) -> bool {
            test_round_trips(&contents);
            true
        }
    }

    fn test_round_trips(contents: &[impl AsRef<[u8]>]) {
        let dir = TempDir::new().unwrap();
        let mut zstore = Zstore::open(&dir).unwrap();

        let mut ids: Vec<Id20> = Vec::new();
        for content in contents.iter() {
            let mut base_ids: Vec<Id20> = Vec::new();
            if !ids.is_empty() {
                let pseudo_rand = AsRef::<[u8]>::as_ref(&ids[ids.len() - 1])[0] as usize + 1;
                base_ids.push(ids[(ids.len() - 1usize) % pseudo_rand]);
            }
            let id = zstore.insert(content.as_ref(), &base_ids).unwrap();
            ids.push(id);
        }

        for (i, id) in ids.iter().enumerate() {
            assert_eq!(zstore.get(*id).unwrap().unwrap(), contents[i].as_ref());
        }

        zstore.flush().unwrap();
        let zstore = Zstore::open(&dir).unwrap();
        for (i, id) in ids.iter().enumerate() {
            assert_eq!(zstore.get(*id).unwrap().unwrap(), contents[i].as_ref());
        }
    }

    /// Generate noise that is hard to compress.
    fn generate_noise(approximated_len: usize) -> String {
        (0..(approximated_len / 41))
            .map(|i| sha1(&[(i % 100) as u8, (i % 101) as u8][..]).to_hex())
            .collect::<Vec<_>>()
            .join("|")
    }

    /// Adjust `DeltaOptions`. Insert `contents` as linear delta chain.
    /// Print the delta tree to a `String`.
    fn show_tree(contents: &[impl AsRef<[u8]>], set_opts: fn(&mut DeltaOptions)) -> String {
        let dir = TempDir::new().unwrap();
        let mut zstore = Zstore::open(&dir).unwrap();
        set_opts(&mut zstore.delta_opts);

        let mut id = *EMPTY_ID20;
        for content in contents {
            let next_id = zstore.insert(content.as_ref(), &[id]).unwrap();
            id = next_id;
        }

        let tree = zstore.debug_delta_tree().unwrap();
        format!("\n{:?}", tree)
    }

    #[test]
    fn test_delta_options() {
        let noise = generate_noise(4000);
        // Similar contents. This ensures that delta will be used.
        let contents: Vec<Vec<u8>> = (0..50)
            .map(|i| format!("{}{}{}", noise, i, noise).as_bytes().to_vec())
            .collect();

        // Test that the delta trees effectively limit the max
        // delta chain length.
        //
        // Set n (max_subchain_len) and d (max_depth) to 3. Check:
        // - The chain length is bounded to 9 (n * d).
        // - 39 (see dostring of DeltaOptions) deltas are used before full text
        //   ("Depth" is 1).
        assert_eq!(
            show_tree(&contents[..42], |opts| {
                opts.max_depth = 3;
                opts.max_subchain_len = 3;
            }),
            // Hint: In case this test is broken by zstd version change, run
            // `fbcode/experimental/quark/grep-rs/cargo-test-i.py --lib` from
            // the `src` directory to auto-update the test.
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  12a9a7
       2 |     2 |          0 |   d176e3
       3 |     3 |          0 |    480cf2
       4 |     3 |          1 |     7b2274
       5 |     3 |          2 |      9195fa
       3 |     2 |          1 |    68a44f
       4 |     3 |          0 |     049af1
       5 |     3 |          1 |      6be3a2
       6 |     3 |          2 |       0565d7
       4 |     2 |          2 |     fcf8f6
       5 |     3 |          0 |      0a012d
       6 |     3 |          1 |       316785
       7 |     3 |          2 |        7c562e
       2 |     1 |          1 |   5b8328
       3 |     2 |          0 |    8802f4
       4 |     3 |          0 |     35b10a
       5 |     3 |          1 |      565d3f
       6 |     3 |          2 |       c2a84a
       4 |     2 |          1 |     aad27f
       5 |     3 |          0 |      4ad6aa
       6 |     3 |          1 |       dd14d5
       7 |     3 |          2 |        aea69c
       5 |     2 |          2 |      0c4075
       6 |     3 |          0 |       ffa7b4
       7 |     3 |          1 |        3cfff1
       8 |     3 |          2 |         0c98c3
       3 |     1 |          2 |    091e0d
       4 |     2 |          0 |     98715e
       5 |     3 |          0 |      47258d
       6 |     3 |          1 |       7802ae
       7 |     3 |          2 |        cec4d0
       5 |     2 |          1 |      fa80d4
       6 |     3 |          0 |       1ca3c5
       7 |     3 |          1 |        5292f3
       8 |     3 |          2 |         a2ff59
       6 |     2 |          2 |       758c14
       7 |     3 |          0 |        465a4a
       8 |     3 |          1 |         3af16a
       9 |     3 |          2 |          002504
       1 |     1 |          0 |  409bca
       2 |     2 |          0 |   cca7b8
       3 |     3 |          0 |    6972af
"#
        );

        // Test that if n = d = 0, full text will always be used.
        assert_eq!(
            show_tree(&contents[..5], |opts| {
                opts.max_depth = 0;
                opts.max_subchain_len = 0;
            }),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  12a9a7
       1 |     1 |          0 |  d176e3
       1 |     1 |          0 |  480cf2
       1 |     1 |          0 |  7b2274
       1 |     1 |          0 |  9195fa
"#
        );

        // Test that if either n or d is a large number, it's similar to a
        // linear chain.
        assert_eq!(
            show_tree(&contents[..5], |opts| {
                opts.max_depth = 1000;
                opts.max_subchain_len = 0;
            }),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  12a9a7
       2 |     2 |          0 |   d176e3
       3 |     3 |          0 |    480cf2
       4 |     4 |          0 |     7b2274
       5 |     5 |          0 |      9195fa
"#
        );
        assert_eq!(
            show_tree(&contents[..5], |opts| {
                opts.max_depth = 0;
                opts.max_subchain_len = 1000;
            }),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  12a9a7
       2 |     1 |          1 |   d176e3
       3 |     1 |          2 |    480cf2
       4 |     1 |          3 |     7b2274
       5 |     1 |          4 |      9195fa
"#
        );
        assert_eq!(
            show_tree(&contents[..5], |opts| {
                opts.max_depth = 1000;
                opts.max_subchain_len = 1000;
            }),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  12a9a7
       2 |     2 |          0 |   d176e3
       3 |     3 |          0 |    480cf2
       4 |     4 |          0 |     7b2274
       5 |     5 |          0 |      9195fa
"#
        );

        // Test that `max_chain_bytes` is effective.
        assert_eq!(
            show_tree(&contents[..10], |opts| {
                opts.max_chain_bytes = 2120;
            }),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  12a9a7
       2 |     2 |          0 |   d176e3
       3 |     3 |          0 |    480cf2
       1 |     1 |          0 |  7b2274
       2 |     2 |          0 |   9195fa
       3 |     3 |          0 |    68a44f
       1 |     1 |          0 |  049af1
       2 |     2 |          0 |   6be3a2
       3 |     3 |          0 |    0565d7
       1 |     1 |          0 |  fcf8f6
"#
        );

        // Test that `max_chain_factor_log` is effective.
        assert_eq!(
            show_tree(
                &[&noise[0..100], &noise[50..150], &noise[100..200]],
                |opts| {
                    // chain bytes should < full text length * 2 (100 * 2).
                    opts.max_chain_factor_log = 1;
                }
            ),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  ed1309
       2 |     2 |          0 |   ab2eea
       1 |     1 |          0 |  f4d55f
"#
        );
    }

    #[test]
    fn test_delta_candidates() {
        let noise = generate_noise(5000);
        let noise = noise.as_bytes();

        // Delta candidates are not used if it cannot beat full text.
        let dir = TempDir::new().unwrap();
        let mut zstore = Zstore::open(&dir).unwrap();
        let id1 = zstore.insert(&noise[0..2000], &[]).unwrap();
        let id2 = zstore.insert(&noise[2000..], &[id1]).unwrap();
        assert_eq!(
            format!("\n{:?}", zstore.debug_delta_tree().unwrap()),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  6e06bf
       1 |     1 |          0 |  4ccc52
"#
        );

        // For multiple delta candidates, the one that compresses better is used.
        zstore.insert(&noise[1500..3500], &[id1, id2]).unwrap();
        assert_eq!(
            format!("\n{:?}", zstore.debug_delta_tree().unwrap()),
            r#"
Chain Len| Depth |Subchain Len| Chain ID
       0 |     0 |          0 | da39a3
       1 |     1 |          0 |  6e06bf
       1 |     1 |          0 |  4ccc52
       2 |     2 |          0 |   a078da
"#
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Convert between foreign types and containers of foreign types.
//! - dag::Id <-> i32 (Javascript / WASM has issues with u64)
//! - dag::Vertex <-> String
//! - Containers (Vec, BTreeMap, BTreeSet, HashMap) of these types.

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::hash::Hash;

use dag::Group;
use dag::Id;
use dag::Vertex;

/// Similar to Into but bypass Rust's orphan rule.
pub trait Convert<T> {
    fn convert(self) -> T;
}

impl Convert<Vertex> for String {
    fn convert(self) -> Vertex {
        Vertex::from(self.into_bytes())
    }
}

impl Convert<String> for Vertex {
    fn convert(self) -> String {
        std::str::from_utf8(self.as_ref()).unwrap().to_string()
    }
}

impl Convert<i32> for Id {
    fn convert(self) -> i32 {
        let mut v: i32 = (self.0 - self.group().min_id().0) as i32;
        if self.group() == Group::NON_MASTER {
            // Use negative number for non-master group.
            v = -v - 1;
        }
        v
    }
}

impl Convert<Id> for i32 {
    fn convert(self) -> Id {
        if self < 0 {
            Group::NON_MASTER.min_id() + (1 - self) as u64
        } else {
            Id(self as u64)
        }
    }
}

impl Convert<String> for String {
    fn convert(self) -> String {
        self
    }
}

impl<T1, T2> Convert<Vec<T2>> for Vec<T1>
where
    T1: Convert<T2>,
{
    fn convert(self) -> Vec<T2> {
        self.into_iter().map(|v| v.convert()).collect()
    }
}

impl<T1, T2> Convert<Vec<T2>> for BTreeSet<T1>
where
    T1: Convert<T2>,
{
    fn convert(self) -> Vec<T2> {
        self.into_iter().map(|v| v.convert()).collect()
    }
}

impl<T1, T2> Convert<BTreeSet<T2>> for Vec<T1>
where
    T1: Convert<T2>,
    T2: Ord,
{
    fn convert(self) -> BTreeSet<T2> {
        self.into_iter().map(|v| v.convert()).collect()
    }
}

impl<T1, T2> Convert<BTreeSet<T2>> for BTreeSet<T1>
where
    T1: Convert<T2>,
    T2: Ord,
{
    fn convert(self) -> BTreeSet<T2> {
        self.into_iter().map(|v| v.convert()).collect()
    }
}

impl<T1, T2> Convert<Option<T2>> for Option<T1>
where
    T1: Convert<T2>,
{
    fn convert(self) -> Option<T2> {
        self.map(|v| v.convert())
    }
}

impl<K1, K2, V1, V2> Convert<HashMap<K2, V2>> for HashMap<K1, V1>
where
    K1: Convert<K2>,
    V1: Convert<V2>,
    K2: Eq + Hash,
{
    fn convert(self) -> HashMap<K2, V2> {
        self.into_iter()
            .map(|(k, v)| (k.convert(), v.convert()))
            .collect()
    }
}

impl<K1, K2, V1, V2> Convert<BTreeMap<K2, V2>> for BTreeMap<K1, V1>
where
    K1: Convert<K2>,
    V1: Convert<V2>,
    K2: Eq + Ord,
{
    fn convert(self) -> BTreeMap<K2, V2> {
        self.into_iter()
            .map(|(k, v)| (k.convert(), v.convert()))
            .collect()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::BTreeSet;
use std::collections::HashMap;
use std::sync::Arc;

use dag::nameset::SyncNameSetQuery;
use dag::ops::DagAddHeads;
use dag::ops::DagAlgorithm;
use dag::ops::IdConvert;
use dag::ops::IdDagAlgorithm;
use dag::Group;
use dag::Id;
use dag::IdSet;
use dag::MemDag;
use dag::Set;
use dag::Vertex;
use dag::VertexListWithOptions;
use nonblocking::non_blocking_result as r;
use serde::Serialize;
use wasm_bindgen::prelude::*;

use crate::Convert;
use crate::JsResult;

#[wasm_bindgen]
pub struct JsDag(MemDag);

#[wasm_bindgen]
pub struct JsIdDag(Arc<dyn IdDagAlgorithm + Send + Sync>);

#[wasm_bindgen]
pub struct JsSet(Set);

#[wasm_bindgen]
pub struct JsIdSet(IdSet);

// Like dag::IdSegment but use i32 instead of Id/u64.
// Also see Convert between Id and i32 below.
//
// Id converts to u64 natively. But i64/u64 can cause WASM errors like:
//
//   ERROR in ./src/wasm-bindings/pkg/wasm_bindings_bg.wasm
//   Import "__wbindgen_bigint_from_u64" from "./wasm_bindings_bg.js" with Non-JS-compatible Func
//   Signature (i64 as parameter) can only be used for direct wasm to wasm dependencies
//
// i32 is used instead of u64 to make WASM happy.
#[derive(Serialize)]
struct IdSegment {
    low: i32,
    high: i32,
    parents: Vec<i32>,
    hasRoot: bool,
    level: u8,
}

#[wasm_bindgen]
impl JsDag {
    #[wasm_bindgen(constructor)]
    pub fn new() -> Self {
        Self(MemDag::new())
    }

    /// addHeads(parents: {name: [name]}, main_heads: [name])
    #[wasm_bindgen]
    pub fn addHeads(&mut self, parents: JsValue, main_heads: JsValue) -> JsResult<()> {
        let parents: HashMap<String, Vec<String>> = serde_wasm_bindgen::from_value(parents)?;
        let parents: HashMap<Vertex, Vec<Vertex>> = parents.convert();
        let heads: VertexListWithOptions = {
            let main_heads: Vec<String> = serde_wasm_bindgen::from_value(main_heads)?;
            let main_heads: Vec<Vertex> = main_heads.convert();
            let all_parents: BTreeSet<Vertex> = parents.values().flatten().cloned().collect();
            let all_children: BTreeSet<Vertex> =
                parents.keys().cloned().collect::<BTreeSet<Vertex>>();
            let all_heads: Vec<Vertex> = all_children.difference(&all_parents).cloned().collect();
            VertexListWithOptions::from(main_heads)
                .with_highest_group(Group::MASTER)
                .chain(all_heads)
        };
        r(self.0.add_heads(&parents, &heads))?;
        Ok(())
    }

    #[wasm_bindgen]
    pub fn render(&self) -> JsResult<JsValue> {
        let result = dag::render::render_namedag_structured(&self.0, None).unwrap();
        Ok(serde_wasm_bindgen::to_value(&result)?)
    }

    #[wasm_bindgen]
    pub fn renderSubset(&self, subset: &JsSet) -> JsResult<JsValue> {
        let result =
            dag::render::render_namedag_structured(&self.0, Some(subset.0.clone())).unwrap();
        Ok(serde_wasm_bindgen::to_value(&result)?)
    }

    // trait DagAlgorithm {{

    #[wasm_bindgen]
    pub fn parentNames(&self, name: String) -> JsResult<JsValue> {
        let name = name.convert();
        let names = r(self.0.parent_names(name))?;
        let names: Vec<String> = names.convert();
        Ok(serde_wasm_bindgen::to_value(&names)?)
    }

    #[wasm_bindgen]
    pub fn all(&self) -> JsResult<JsSet> {
        let set = r(self.0.all())?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn mainGroup(&self) -> JsResult<JsSet> {
        let set = r(self.0.master_group())?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn dirty(&self) -> JsResult<JsSet> {
        let set = r(self.0.dirty())?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn sort(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.sort(&set.0))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn ancestors(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.ancestors(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn parents(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.parents(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn firstAncestors(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.first_ancestors(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn heads(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.heads(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn children(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.children(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn roots(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.roots(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn merges(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.merges(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn gca(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.gca_all(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn commonAncestors(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.common_ancestors(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn headsAncestors(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.heads_ancestors(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn range(&self, roots: &JsSet, heads: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.range(roots.0.clone(), heads.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn only(&self, reachable: &JsSet, unreachable: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.only(reachable.0.clone(), unreachable.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn descendants(&self, set: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.descendants(set.0.clone()))?;
        Ok(JsSet(set))
    }

    #[wasm_bindgen]
    pub fn reachableRoots(&self, roots: &JsSet, heads: &JsSet) -> JsResult<JsSet> {
        let set = r(self.0.reachable_roots(roots.0.clone(), heads.0.clone()))?;
        Ok(JsSet(set))
    }

    // Gateway to low-level Ids.
    #[wasm_bindgen]
    pub fn idDag(&self) -> JsResult<JsIdDag> {
        let id_dag = self.0.id_dag_snapshot()?;
        Ok(JsIdDag(id_dag))
    }

    // }} // trait DagAlgorithm

    // trait IdConvert {{

    #[wasm_bindgen]
    pub fn vertexId(&self, name: String) -> JsResult<JsValue> {
        let id: Option<Id> = r(self.0.vertex_id_optional(&name.convert()))?;
        let id: Option<i32> = id.convert();
        Ok(serde_wasm_bindgen::to_value(&id)?)
    }

    #[wasm_bindgen]
    pub fn vertexName(&self, id: i32) -> JsResult<String> {
        let id: Id = id.convert();
        let name = r(self.0.vertex_name(id))?;
        let name: String = name.convert();
        Ok(name)
    }

    #[wasm_bindgen]
    pub fn convtainsVertexIdLocally(&self, ids: JsValue) -> JsResult<JsValue> {
        let ids: Vec<i32> = serde_wasm_bindgen::from_value(ids)?;
        let ids: Vec<Id> = ids.convert();
        let result: Vec<bool> = r(self.0.contains_vertex_id_locally(&ids))?;
        Ok(serde_wasm_bindgen::to_value(&result)?)
    }

    #[wasm_bindgen]
    pub fn convtainsVertexNameLocally(&self, names: JsValue) -> JsResult<JsValue> {
        let names: Vec<String> = serde_wasm_bindgen::from_value(names)?;
        let names: Vec<Vertex> = names.convert();
        let result: Vec<bool> = r(self.0.contains_vertex_name_locally(&names))?;
        Ok(serde_wasm_bindgen::to_value(&result)?)
    }

    // }} // trait IdConvert
}

#[wasm_bindgen]
impl JsIdDag {
    // trait IdDagAlgorithm {{

    #[wasm_bindgen]
    pub fn parentIds(&self, id: i32) -> JsResult<JsValue> {
        let id: Id = id.convert();
        let ids: Vec<Id> = self.0.parent_ids(id)?;
        let ids: Vec<i32> = ids.convert();
        Ok(serde_wasm_bindgen::to_value(&ids)?)
    }

    #[wasm_bindgen]
    pub fn all(&self) -> JsResult<JsIdSet> {
        let set = self.0.all()?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn mainGroup(&self) -> JsResult<JsIdSet> {
        let set = self.0.master_group()?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn ancestors(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.ancestors(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn parents(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.parents(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn firstAncestors(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.first_ancestors(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn heads(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.heads(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn children(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.children(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn roots(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.roots(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn merges(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.merges(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn gca(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.gca_all(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn commonAncestors(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.common_ancestors(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn headsAncestors(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.heads_ancestors(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn range(&self, roots: &JsIdSet, heads: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.range(roots.0.clone(), heads.0.clone())?;
        Ok(JsIdSet(set))
    }

    #[wasm_bindgen]
    pub fn descendants(&self, set: &JsIdSet) -> JsResult<JsIdSet> {
        let set = self.0.descendants(set.0.clone())?;
        Ok(JsIdSet(set))
    }

    // Gateway to get segments information.
    #[wasm_bindgen]
    pub fn segments(&self, set: &JsIdSet, max_level: u8) -> JsResult<JsValue> {
        let id_segments = self
            .0
            .id_set_to_id_segments_with_max_level(&set.0, max_level)?;
        let id_segments: Vec<IdSegment> = id_segments
            .into_iter()
            .map(|s| IdSegment {
                low: s.low.convert(),
                high: s.high.convert(),
                parents: s.parents.convert(),
                hasRoot: s.has_root,
                level: s.level,
            })
            .collect();
        Ok(serde_wasm_bindgen::to_value(&id_segments)?)
    }

    // }} // trait IdDagAlgorithm
}

#[wasm_bindgen]
impl JsSet {
    #[wasm_bindgen(constructor)]
    pub fn new(value: JsValue) -> JsResult<JsSet> {
        if value.is_null() || value.is_undefined() {
            Ok(Self(Set::empty()))
        } else {
            let names: Vec<String> = serde_wasm_bindgen::from_value(value)?;
            let names: Vec<Vertex> = names.convert();
            let set = Set::from_static_names(names);
            Ok(Self(set))
        }
    }

    #[wasm_bindgen]
    pub fn toString(&self) -> String {
        format!("{:?}", &self.0)
    }

    #[wasm_bindgen]
    pub fn toJSON(&self) -> JsResult<JsValue> {
        let names: Vec<Vertex> = SyncNameSetQuery::iter(&self.0)?.collect::<Result<Vec<_>, _>>()?;
        let names: Vec<String> = names.convert();
        Ok(serde_wasm_bindgen::to_value(&names)?)
    }

    #[wasm_bindgen]
    pub fn count(&self) -> JsResult<usize> {
        Ok(self.0.count()?)
    }

    #[wasm_bindgen]
    pub fn skip(&self, n: u32) -> JsResult<JsSet> {
        Ok(Self(self.0.skip(n as _)))
    }

    #[wasm_bindgen]
    pub fn take(&self, n: u32) -> JsResult<JsSet> {
        Ok(Self(self.0.take(n as _)))
    }

    #[wasm_bindgen]
    pub fn union(&self, rhs: &JsSet) -> JsResult<JsSet> {
        Ok(Self(self.0.union(&rhs.0)))
    }

    #[wasm_bindgen]
    pub fn intersection(&self, rhs: &JsSet) -> JsResult<JsSet> {
        Ok(Self(self.0.intersection(&rhs.0)))
    }

    #[wasm_bindgen]
    pub fn difference(&self, rhs: &JsSet) -> JsResult<JsSet> {
        Ok(Self(self.0.difference(&rhs.0)))
    }
}

#[wasm_bindgen]
impl JsIdSet {
    #[wasm_bindgen(constructor)]
    pub fn new(value: JsValue) -> JsResult<JsIdSet> {
        if value.is_null() || value.is_undefined() {
            Ok(Self(Default::default()))
        } else {
            let ids: Vec<i32> = serde_wasm_bindgen::from_value(value)?;
            let ids: Vec<Id> = ids.convert();
            let set = IdSet::from_spans(ids);
            Ok(Self(set))
        }
    }

    #[wasm_bindgen]
    pub fn toString(&self) -> String {
        format!("{:?}", &self.0)
    }

    #[wasm_bindgen]
    pub fn toJSON(&self) -> JsResult<JsValue> {
        let ids: Vec<Id> = self.0.iter_desc().collect();
        let ids: Vec<i32> = ids.convert();
        Ok(serde_wasm_bindgen::to_value(&ids)?)
    }

    #[wasm_bindgen]
    pub fn count(&self) -> JsResult<u32> {
        Ok(self.0.count() as _)
    }

    #[wasm_bindgen]
    pub fn skip(&self, n: u32) -> JsResult<JsIdSet> {
        Ok(Self(self.0.skip(n as _)))
    }

    #[wasm_bindgen]
    pub fn take(&self, n: u32) -> JsResult<JsIdSet> {
        Ok(Self(self.0.take(n as _)))
    }

    #[wasm_bindgen]
    pub fn union(&self, rhs: &JsIdSet) -> JsResult<JsIdSet> {
        Ok(Self(self.0.union(&rhs.0)))
    }

    #[wasm_bindgen]
    pub fn intersection(&self, rhs: &JsIdSet) -> JsResult<JsIdSet> {
        Ok(Self(self.0.intersection(&rhs.0)))
    }

    #[wasm_bindgen]
    pub fn difference(&self, rhs: &JsIdSet) -> JsResult<JsIdSet> {
        Ok(Self(self.0.difference(&rhs.0)))
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::BTreeMap;
use std::collections::BTreeSet;

use dag::Vertex;
use wasm_bindgen::prelude::*;

use crate::Convert;

/// ASCII graph -> {name: List[name]}
#[wasm_bindgen]
pub fn drawdag(text: &str) -> JsValue {
    let parent_map: BTreeMap<String, BTreeSet<String>> = drawdag::parse(text);
    let parent_map: BTreeMap<Vertex, Vec<Vertex>> = parent_map.convert();
    // drawdag does not prevent cycles. Break cycles here to avoid infinite loops.
    let parent_func = dag::utils::break_parent_func_cycle(|name: Vertex| {
        Ok(parent_map.get(&name).cloned().unwrap_or_default())
    });
    let result: BTreeMap<Vertex, Vec<Vertex>> = parent_map
        .keys()
        .map(|k| (k.clone(), parent_func(k.clone()).unwrap()))
        .collect();
    let result: BTreeMap<String, Vec<String>> = result.convert();
    serde_wasm_bindgen::to_value(&result).unwrap()
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use serde::Serialize;
use wasm_bindgen::prelude::*;

// Matches std::error::Error.
#[derive(Serialize)]
pub struct JsError(String);

// See https://rustwasm.github.io/wasm-bindgen/reference/types/result.html
impl Into<JsValue> for JsError {
    fn into(self) -> JsValue {
        JsValue::from_str(&self.0)
    }
}

impl<T: std::error::Error> From<T> for JsError {
    fn from(error: T) -> Self {
        to_js_error(&error)
    }
}

fn to_js_error(error: &dyn std::error::Error) -> JsError {
    let mut message = error.to_string();
    if let Some(inner_error) = error.source() {
        message += &format!(" ({})", to_js_error(inner_error).0);
    }
    JsError(message)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// WASM exports must be "pub". But they are not used in this crate.
#![allow(dead_code)]
// Use fooBar naming for exported to match Javascript style.
#![allow(non_snake_case)]

use wasm_bindgen::prelude::*;

#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

mod convert;
mod dag;
mod drawdag;
mod error;
mod tracing;

pub use convert::Convert;
pub use error::JsError;
pub type JsResult<T> = std::result::Result<T, JsError>;

#[wasm_bindgen(start)]
pub fn setup() {
    ::console_error_panic_hook::set_once();
    ::dag::iddag::set_default_seg_size(3);
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::borrow::Cow;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::Mutex;

use serde::Serialize;
use tracing::span::Attributes;
use tracing::span::Id;
use tracing::span::Record;
use tracing::Event;
use tracing::Level;
use tracing::Metadata;
use tracing::Subscriber;
use wasm_bindgen::prelude::*;

#[derive(Serialize, Clone)]
struct JsTracingSubscriber {
    #[serde(skip)]
    level: Level,
    spans: Arc<Mutex<Vec<MetadataMap>>>,
    log: Arc<Mutex<Vec<(Action, u32)>>>,
}

#[derive(Serialize)]
enum Action {
    Enter,
    Exit,
    Event,
}

#[derive(Serialize)]
#[serde(transparent)]
struct MetadataMap(BTreeMap<&'static str, Cow<'static, str>>);

struct FieldVisitor<'a> {
    out: &'a mut MetadataMap,
}

impl MetadataMap {
    fn from_metadata(metadata: &Metadata<'static>) -> Self {
        let mut map = BTreeMap::new();
        map.insert("name", metadata.name().into());
        if let Some(module_path) = metadata.module_path() {
            map.insert("module_path", module_path.into());
        }
        if let Some(line) = metadata.line() {
            map.insert("line", line.to_string().into());
        }
        Self(map)
    }
}

impl Subscriber for JsTracingSubscriber {
    fn enabled(&self, metadata: &Metadata<'_>) -> bool {
        metadata.level() >= &self.level
    }
    fn new_span(&self, span: &Attributes<'_>) -> Id {
        let id = self.push_metadata(span.metadata(), None);
        Id::from_u64(id as _)
    }
    fn record(&self, span_id: &Id, span_values: &Record<'_>) {
        let id = (span_id.into_u64() - 1) as usize;
        let mut spans = self.spans.lock().unwrap();
        let mut visitor = FieldVisitor {
            out: &mut spans[id],
        };
        span_values.record(&mut visitor);
    }
    fn record_follows_from(&self, span_id: &Id, follows_id: &Id) {
        let _ = (span_id, follows_id);
    }
    fn event(&self, event: &Event<'_>) {
        let id = self.push_metadata(event.metadata(), Some(event));
        self.push_log(Action::Event, id);
    }
    fn enter(&self, span_id: &Id) {
        self.push_log(Action::Enter, span_id.into_u64() as _);
    }
    fn exit(&self, span_id: &Id) {
        self.push_log(Action::Exit, span_id.into_u64() as _);
    }
}

impl JsTracingSubscriber {
    fn new(level: Level) -> Self {
        Self {
            level,
            spans: Default::default(),
            log: Default::default(),
        }
    }

    fn push_metadata(&self, metadata: &Metadata<'static>, event: Option<&Event>) -> usize {
        let mut spans = self.spans.lock().unwrap();
        let mut map = MetadataMap::from_metadata(metadata);
        let mut visitor = FieldVisitor { out: &mut map };
        if let Some(event) = event {
            event.record(&mut visitor);
        }
        spans.push(map);
        spans.len()
    }

    fn push_log(&self, action: Action, id: usize) {
        let mut log = self.log.lock().unwrap();
        log.push((action, id as _))
    }

    fn to_jsvalue(&self) -> JsValue {
        serde_wasm_bindgen::to_value(self).unwrap()
    }
}

impl<'a> FieldVisitor<'a> {
    fn record(&mut self, field: &tracing::field::Field, value: impl ToString) {
        let key = field.name();
        let value = value.to_string();
        self.out.0.insert(key, value.into());
    }
}

impl<'a> tracing::field::Visit for FieldVisitor<'a> {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        self.record(field, format!("{:?}", value));
    }
    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
        self.record(field, value)
    }
    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
        self.record(field, value)
    }
    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
        self.record(field, value)
    }
    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        self.record(field, value)
    }
    fn record_error(
        &mut self,
        field: &tracing::field::Field,
        value: &(dyn std::error::Error + 'static),
    ) {
        self.record(field, value)
    }
}

/// Calls f() with tracing captured. Returns trace.
#[wasm_bindgen]
pub fn withTracing(f: &js_sys::Function, level: Option<u32>) -> Result<JsValue, JsValue> {
    let this = JsValue::NULL;
    let level: Level = level.map(u32_to_level).unwrap_or(Level::TRACE);
    let subscriber = JsTracingSubscriber::new(level);
    // Ideally we can return f() as a second return value. But it seems tricky to do so
    // with the current version of wasm-bindgen (0.2.83).
    let _ = tracing::subscriber::with_default(subscriber.clone(), || f.call0(&this))?;
    let trace = subscriber.to_jsvalue();
    Ok(trace)
}

fn u32_to_level(v: u32) -> Level {
    match v {
        0 => Level::ERROR,
        1 => Level::WARN,
        2 => Level::INFO,
        3 => Level::DEBUG,
        _ => Level::TRACE,
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;

use crate::datatype::AbstractDomain;
use crate::graph::Graph;
use crate::graph::SuccessorNodes;
use crate::wpo::WeakPartialOrdering;
use crate::wpo::WpoIdx;

// Unlike the C++ version, we don't treat this as a base interface
// for all iterators, instead, it's only a trait for concrete type
// to analyze nodes and edges. And this can be composited in the
// fixpoint iterator type. We will revisit this part in the following
// diff for analyzer interface.
pub trait FixpointIteratorTransformer<G: Graph, D: AbstractDomain> {
    /// The *current_state* could be updated in place.
    fn analyze_node(&mut self, n: G::NodeId, current_state: &mut D);

    fn analyze_edge(&mut self, e: G::EdgeId, exit_state_at_src: &D) -> D;
}

pub struct MonotonicFixpointIteratorContext<G: Graph, D: AbstractDomain> {
    init_value: D,
    local_iterations: HashMap<G::NodeId, u32>,
    global_iterations: HashMap<G::NodeId, u32>,
}

impl<G, D> MonotonicFixpointIteratorContext<G, D>
where
    G: Graph,
    D: AbstractDomain,
{
    pub fn get_local_iterations_for(&self, n: G::NodeId) -> u32 {
        *self.local_iterations.get(&n).unwrap_or(&0)
    }

    pub fn get_global_iterations_for(&self, n: G::NodeId) -> u32 {
        *self.global_iterations.get(&n).unwrap_or(&0)
    }

    pub fn get_init_value(&self) -> &D {
        &self.init_value
    }

    fn increase_iteration_count(n: G::NodeId, table: &mut HashMap<G::NodeId, u32>) {
        *table.entry(n).or_default() += 1;
    }

    pub fn increase_iteration_count_for(&mut self, n: G::NodeId) {
        Self::increase_iteration_count(n, &mut self.local_iterations);
        Self::increase_iteration_count(n, &mut self.global_iterations);
    }

    pub fn reset_local_iteration_count_for(&mut self, n: G::NodeId) {
        *self.local_iterations.entry(n).or_default() = 0;
    }

    pub fn new(init_value: D) -> Self {
        Self {
            init_value,
            local_iterations: Default::default(),
            global_iterations: Default::default(),
        }
    }

    pub fn with_nodes(mut self, nodes: &HashSet<G::NodeId>) -> Self {
        for &node in nodes {
            *self.global_iterations.entry(node).or_default() = 0;
            *self.local_iterations.entry(node).or_default() = 0;
        }
        self
    }
}

pub struct MonotonicFixpointIterator<
    'g,
    G: Graph,
    D: AbstractDomain,
    T: FixpointIteratorTransformer<G, D>,
> {
    graph: &'g G,
    entry_states: HashMap<G::NodeId, D>,
    exit_states: HashMap<G::NodeId, D>,
    transformer: T,
    wpo: WeakPartialOrdering<G::NodeId>,
}

impl<'g, G, D, T> MonotonicFixpointIterator<'g, G, D, T>
where
    G: Graph,
    D: AbstractDomain,
    T: FixpointIteratorTransformer<G, D>,
{
    pub fn new<SN>(g: &'g G, cfg_size_hint: usize, transformer: T, successors_nodes: &SN) -> Self
    where
        SN: SuccessorNodes<NodeId = G::NodeId>,
    {
        let wpo = WeakPartialOrdering::new(g.entry(), g.size(), successors_nodes);
        Self {
            graph: g,
            entry_states: HashMap::with_capacity(cfg_size_hint),
            exit_states: HashMap::with_capacity(cfg_size_hint),
            transformer,
            wpo,
        }
    }

    pub fn run(&mut self, init_value: D) {
        self.clear();

        let mut context = MonotonicFixpointIteratorContext::new(init_value);
        let wpo_counter: Vec<AtomicU32> =
            (0..self.wpo.size()).map(|_| Default::default()).collect();

        let mut worklist = VecDeque::new();
        let entry_idx = self.wpo.get_entry();
        worklist.push_front(entry_idx);
        assert_eq!(self.wpo.get_num_preds(entry_idx), 0);

        let mut process_node = |wpo_idx: WpoIdx, worklist: &mut VecDeque<WpoIdx>| {
            assert_eq!(
                wpo_counter[wpo_idx as usize].load(Ordering::Relaxed),
                self.wpo.get_num_preds(wpo_idx)
            );

            wpo_counter[wpo_idx as usize].store(0, Ordering::Relaxed);

            if !self.wpo.is_exit(wpo_idx) {
                self.analyze_vertex(&context, self.wpo.get_node(wpo_idx));

                for &succ_idx in self.wpo.get_successors(wpo_idx) {
                    let old_counter =
                        wpo_counter[succ_idx as usize].fetch_add(1, Ordering::Relaxed);
                    if old_counter + 1 == self.wpo.get_num_preds(succ_idx) {
                        worklist.push_back(succ_idx);
                    }
                }

                return;
            }

            let head_idx = self.wpo.get_head_of_exit(wpo_idx);
            let head = self.wpo.get_node(head_idx);
            let current_state = self.entry_states.entry(head).or_insert_with(D::bottom);
            let mut new_state = D::bottom();
            Self::compute_entry_state(
                self.graph,
                &self.exit_states,
                &mut self.transformer,
                &context,
                head,
                &mut new_state,
            );

            if new_state.leq(current_state) {
                context.reset_local_iteration_count_for(head);
                *current_state = new_state;

                for &succ_idx in self.wpo.get_successors(wpo_idx) {
                    let old_counter =
                        wpo_counter[succ_idx as usize].fetch_add(1, Ordering::Relaxed);
                    if old_counter + 1 == self.wpo.get_num_preds(succ_idx) {
                        worklist.push_back(succ_idx);
                    }
                }
            } else {
                Self::extrapolate(&context, head, current_state, new_state);
                context.increase_iteration_count_for(head);
                for (&component_idx, &num) in self.wpo.get_num_outer_preds(wpo_idx) {
                    assert!(component_idx != entry_idx);
                    let old_counter =
                        wpo_counter[component_idx as usize].fetch_add(num, Ordering::Relaxed);
                    if old_counter + num == self.wpo.get_num_preds(component_idx) {
                        worklist.push_back(component_idx);
                    }
                }

                if head_idx == entry_idx {
                    worklist.push_back(head_idx);
                }
            }
        };

        while let Some(idx) = worklist.pop_front() {
            process_node(idx, &mut worklist);
        }

        for counter in wpo_counter {
            assert_eq!(counter.load(Ordering::Relaxed), 0);
        }
    }

    /// Default strategy for applying widening operator (apply
    /// join at the first iteration and then widening in all
    /// rest iterations).
    pub fn extrapolate(
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
        current_state: &mut D,
        new_state: D,
    ) {
        if 0 == context.get_global_iterations_for(n) {
            // TODO: we need to revisit this design, should we use
            // move or clone of domain?
            current_state.join_with(new_state);
        } else {
            current_state.widen_with(new_state);
        }
    }

    fn get_state_at_or_bottom(states: &HashMap<G::NodeId, D>, n: G::NodeId) -> Cow<'_, D> {
        if let Some(state) = states.get(&n) {
            Cow::Borrowed(state)
        } else {
            Cow::Owned(D::bottom())
        }
    }

    pub fn get_entry_state_at(&self, n: G::NodeId) -> Cow<'_, D> {
        Self::get_state_at_or_bottom(&self.entry_states, n)
    }

    pub fn get_exit_state_at(&self, n: G::NodeId) -> Cow<'_, D> {
        Self::get_state_at_or_bottom(&self.exit_states, n)
    }

    pub fn clear(&mut self) {
        self.entry_states.clear();
        self.entry_states.shrink_to_fit();
        self.exit_states.clear();
        self.exit_states.shrink_to_fit();
    }

    pub fn set_all_to_bottom(&mut self, all_nodes: &HashSet<G::NodeId>) {
        for &node in all_nodes {
            self.entry_states
                .entry(node)
                .and_modify(|s| *s = D::bottom())
                .or_insert_with(D::bottom);
            self.exit_states
                .entry(node)
                .and_modify(|s| *s = D::bottom())
                .or_insert_with(D::bottom);
        }
    }

    pub fn compute_entry_state(
        graph: &'g G,
        exit_states: &HashMap<G::NodeId, D>,
        transformer: &mut T,
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
        entry_state: &mut D,
    ) {
        if n == graph.entry() {
            entry_state.join_with(context.get_init_value().clone());
        }

        for e in graph.predecessors(n) {
            let d = Self::get_state_at_or_bottom(exit_states, graph.source(e));
            entry_state.join_with(transformer.analyze_edge(e, &d));
        }
    }

    pub fn analyze_vertex(
        &mut self,
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
    ) {
        let entry_state = self.entry_states.entry(n).or_insert_with(D::bottom);
        Self::compute_entry_state(
            self.graph,
            &self.exit_states,
            &mut self.transformer,
            context,
            n,
            entry_state,
        );
        let exit_state = self
            .exit_states
            .entry(n)
            .and_modify(|s| *s = entry_state.clone())
            .or_insert_with(|| entry_state.clone());
        self.transformer.analyze_node(n, exit_state);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::borrow::Borrow;
use std::fmt::Debug;
use std::hash::Hash;
use std::marker::PhantomData;

use smallvec::SmallVec;

pub const DEFAULT_GRAPH_SUCCS_NUM: usize = 4;

/// Graph trait used in Sparta fixpoint iterator.
///
/// Clients can either maintain their own structure and implement
/// this trait, or use third-party graph crates and implement
/// this trait as a wrapper.
/// The constant generic parameter S is an estimated number of successor
/// nodes that is just enough for most nodes. A larger/smaller
/// value will not cause incorrect result, it only affects
/// performance.
// NOTE: due to the status of Rust const generics, we can not define
// S as an associated constant since it is part of the return type for
// method predecessors and successors.
pub trait Graph<const S: usize = DEFAULT_GRAPH_SUCCS_NUM> {
    type NodeId: Copy + Hash + Eq + Ord + Debug;
    type EdgeId: Copy;

    /// Entry node.
    fn entry(&self) -> Self::NodeId;

    /// Exit node.
    fn exit(&self) -> Self::NodeId;

    /// Predecessors of n.
    fn predecessors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; S]>;

    /// Sucessors of n.
    fn successors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; S]>;

    /// The source node of e.
    fn source(&self, e: Self::EdgeId) -> Self::NodeId;

    /// The target node of e.
    fn target(&self, e: Self::EdgeId) -> Self::NodeId;

    /// The number of all nodes.
    fn size(&self) -> usize;
}

/// A limited view of the graph that allows retrieving successor nodes of a given node.
///
/// Clients can create a new type to implement it, or they can use the default
/// implementation provided on types that implemented `Graph`.
pub trait SuccessorNodes<const S: usize = DEFAULT_GRAPH_SUCCS_NUM> {
    type NodeId: Copy + Hash + Eq;

    fn get_succ_nodes(&self, n: Self::NodeId) -> SmallVec<[Self::NodeId; S]>;
}

/// Implement `SuccessorNodes` for any type implemented `Graph`.
///
/// This is a convenient facility for clients who don't want to
/// create a new type to implement `SuccessorNodes`. Note that
/// this may be less efficient than custom implementation of
/// SuccessorNodes since it can only use available trait methods
/// such as `successors`.
impl<const S: usize, T: Graph<S>> SuccessorNodes<S> for T {
    type NodeId = T::NodeId;

    fn get_succ_nodes(&self, n: Self::NodeId) -> SmallVec<[Self::NodeId; S]> {
        self.successors(n)
            .into_iter()
            .map(|edge_idx| self.target(edge_idx))
            .collect()
    }
}

/// A reversed graph used for backward analysis.
pub struct ReversedGraph<G: Graph<S>, B: Borrow<G>, const S: usize = DEFAULT_GRAPH_SUCCS_NUM> {
    graph: B,
    phantom: PhantomData<*const G>,
}

// Type alias to ease specifying generic parameters.
pub type ReversedRefGraph<'a, G, const S: usize = DEFAULT_GRAPH_SUCCS_NUM> =
    ReversedGraph<G, &'a G, S>;
pub type ReversedIntoGraph<G, const S: usize = DEFAULT_GRAPH_SUCCS_NUM> = ReversedGraph<G, G, S>;

/// Graph type that implements this trait can create reversed graph.
pub trait ReverseGraph<const S: usize = DEFAULT_GRAPH_SUCCS_NUM> {
    type G: Graph<S>;

    /// Reversed graph.
    fn rev(&self) -> ReversedRefGraph<Self::G, S>;

    /// Reversed graph that consumes the original graph.
    fn into_rev(self) -> ReversedIntoGraph<Self::G, S>;
}

/// Default `ReverseGraph` implementation for any type that implements `Graph`
///
/// This can be chained multiple times, e.g., `graph.rev().rev()` becomes
/// the same as `graph`. This behavior is the same as `std::iter::Iterator::rev()`.
impl<const S: usize, T: Graph<S>> ReverseGraph<S> for T {
    type G = T;

    fn rev(&self) -> ReversedRefGraph<T, S> {
        ReversedGraph {
            graph: self,
            phantom: PhantomData,
        }
    }

    fn into_rev(self) -> ReversedIntoGraph<T, S> {
        ReversedGraph {
            graph: self,
            phantom: PhantomData,
        }
    }
}

impl<G: Graph<S>, B: Borrow<G>, const S: usize> Graph<S> for ReversedGraph<G, B, S> {
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;

    fn entry(&self) -> Self::NodeId {
        self.graph.borrow().exit()
    }

    fn exit(&self) -> Self::NodeId {
        self.graph.borrow().entry()
    }

    fn predecessors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; S]> {
        self.graph.borrow().successors(n)
    }

    fn successors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; S]> {
        self.graph.borrow().predecessors(n)
    }

    fn source(&self, e: Self::EdgeId) -> Self::NodeId {
        self.graph.borrow().target(e)
    }

    fn target(&self, e: Self::EdgeId) -> Self::NodeId {
        self.graph.borrow().source(e)
    }

    fn size(&self) -> usize {
        self.graph.borrow().size()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#![allow(dead_code)]

pub mod datatype;
pub mod fixpoint_iter;
pub mod graph;
pub mod wpo;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Debug;
use std::hash::Hash;
use std::iter::FromIterator;

use petgraph::unionfind::UnionFind;

use crate::graph::SuccessorNodes;

pub type WpoIdx = u32;

pub struct WpoNodeData<NodeId: Copy + Hash + Ord> {
    node: NodeId,
    size: usize,
    successors: BTreeSet<WpoIdx>,
    predessors: BTreeSet<WpoIdx>,
    num_outer_preds: HashMap<WpoIdx, u32>,
}

impl<NodeId> WpoNodeData<NodeId>
where
    NodeId: Copy + Hash + Ord,
{
    pub fn new(node: NodeId, size: usize) -> Self {
        Self {
            node,
            size,
            successors: BTreeSet::new(),
            predessors: BTreeSet::new(),
            num_outer_preds: HashMap::new(),
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
pub enum WpoNodeType {
    Head,
    Plain,
    Exit,
}

pub struct WpoNode<NodeId: Copy + Ord + Hash> {
    ty: WpoNodeType,
    data: WpoNodeData<NodeId>,
}

impl<NodeId> WpoNode<NodeId>
where
    NodeId: Copy + Ord + Hash + Debug,
{
    pub fn plain(node: NodeId, size: usize) -> Self {
        Self {
            ty: WpoNodeType::Plain,
            data: WpoNodeData::new(node, size),
        }
    }

    pub fn head(node: NodeId, size: usize) -> Self {
        Self {
            ty: WpoNodeType::Head,
            data: WpoNodeData::new(node, size),
        }
    }

    pub fn exit(node: NodeId, size: usize) -> Self {
        Self {
            ty: WpoNodeType::Exit,
            data: WpoNodeData::new(node, size),
        }
    }

    pub fn new(ty: WpoNodeType, node: NodeId, size: usize) -> Self {
        Self {
            ty,
            data: WpoNodeData::new(node, size),
        }
    }

    pub fn is_plain(&self) -> bool {
        self.ty == WpoNodeType::Plain
    }

    pub fn is_head(&self) -> bool {
        self.ty == WpoNodeType::Head
    }

    pub fn is_exit(&self) -> bool {
        self.ty == WpoNodeType::Exit
    }

    pub fn get_node(&self) -> NodeId {
        self.data.node
    }

    pub fn get_successors(&self) -> &BTreeSet<WpoIdx> {
        &self.data.successors
    }

    pub fn get_predecessors(&self) -> &BTreeSet<WpoIdx> {
        &self.data.predessors
    }

    pub fn get_num_preds(&self) -> u32 {
        self.get_predecessors().len() as u32
    }

    pub fn get_num_outer_preds(&self) -> &HashMap<WpoIdx, u32> {
        assert_eq!(
            self.ty,
            WpoNodeType::Exit,
            "Node {:#?} is not Exit",
            self.data.node
        );
        &self.data.num_outer_preds
    }

    pub fn get_size(&self) -> usize {
        self.data.size
    }

    fn add_successor(&mut self, idx: WpoIdx) {
        self.data.successors.insert(idx);
    }

    fn add_predecessor(&mut self, idx: WpoIdx) {
        self.data.predessors.insert(idx);
    }

    fn is_successor(&self, idx: WpoIdx) -> bool {
        self.get_successors().contains(&idx)
    }

    pub fn inc_num_outer_preds(&mut self, idx: WpoIdx) {
        assert_eq!(
            self.ty,
            WpoNodeType::Exit,
            "Node {:#?} is not Exit",
            self.data.node
        );
        *self.data.num_outer_preds.entry(idx).or_default() += 1;
    }
}

pub struct WeakPartialOrdering<NodeId: Copy + Hash + Ord> {
    /// All nodes under WPO.
    nodes: Vec<WpoNode<NodeId>>,
    /// All top level nodes.
    toplevel: Vec<WpoIdx>,
    /// Post depth first numbering for each node.
    post_dfn: HashMap<NodeId, u32>,
}

impl<NodeId> WeakPartialOrdering<NodeId>
where
    NodeId: Copy + Hash + Ord + Debug,
{
    pub fn new<SN>(root: NodeId, size: usize, successors_nodes: &SN) -> Self
    where
        SN: SuccessorNodes<NodeId = NodeId>,
    {
        if successors_nodes.get_succ_nodes(root).is_empty() {
            let mut wpo = Self {
                nodes: vec![],
                toplevel: vec![],
                post_dfn: HashMap::new(),
            };
            wpo.nodes.push(WpoNode::plain(root, 1));
            wpo.toplevel.push(0);
            wpo.post_dfn.insert(root, 1);
            wpo
        } else {
            WeakPartialOrderingImpl::new().build(size, root, successors_nodes)
        }
    }

    pub fn size(&self) -> usize {
        self.nodes.len()
    }

    pub fn get_entry(&self) -> WpoIdx {
        (self.nodes.len() - 1) as WpoIdx
    }

    pub fn get_successors(&self, idx: WpoIdx) -> &BTreeSet<WpoIdx> {
        self.nodes[idx as usize].get_successors()
    }

    pub fn get_predecessors(&self, idx: WpoIdx) -> &BTreeSet<WpoIdx> {
        self.nodes[idx as usize].get_predecessors()
    }

    pub fn get_num_preds(&self, idx: WpoIdx) -> u32 {
        self.nodes[idx as usize].get_num_preds()
    }

    pub fn get_num_outer_preds(&self, exit: WpoIdx) -> &HashMap<WpoIdx, u32> {
        self.nodes[exit as usize].get_num_outer_preds()
    }

    pub fn get_head_of_exit(&self, exit: WpoIdx) -> WpoIdx {
        exit + 1
    }

    pub fn get_exit_of_head(&self, head: WpoIdx) -> WpoIdx {
        head - 1
    }

    pub fn get_node(&self, idx: WpoIdx) -> NodeId {
        self.nodes[idx as usize].get_node()
    }

    pub fn is_plain(&self, idx: WpoIdx) -> bool {
        self.nodes[idx as usize].is_plain()
    }

    pub fn is_head(&self, idx: WpoIdx) -> bool {
        self.nodes[idx as usize].is_head()
    }

    pub fn is_exit(&self, idx: WpoIdx) -> bool {
        self.nodes[idx as usize].is_exit()
    }

    pub fn is_from_outside(&self, head: NodeId, pred: NodeId) -> bool {
        self.get_post_dfn(head) < self.get_post_dfn(pred)
    }

    fn get_post_dfn(&self, n: NodeId) -> u32 {
        // If the key does not exist, meaning that node is not
        // finished yet, return default value 0.
        self.post_dfn.get(&n).copied().unwrap_or_default()
    }
}

// This private type is only used to build the actual WPO.
struct WeakPartialOrderingImpl<NodeId: Copy + Hash + Ord> {
    nodes: Vec<WpoNode<NodeId>>,
    toplevel: Vec<WpoIdx>,
    post_dfn: HashMap<NodeId, u32>,
    // A map from NodeId to post DFN.
    dfn: HashMap<NodeId, u32>,
    dfn_to_node: Vec<NodeId>,
    cross_fwd_edges: HashMap<u32, Vec<(u32, u32)>>,
    back_preds: HashMap<u32, Vec<u32>>,
    // Tree edges (map from node to its predecessors).
    non_back_preds: HashMap<u32, Vec<u32>>,
    next_dfn: u32,
    // Map from dfn to WpoIdx
    dfn_to_wpo_idx: Vec<WpoIdx>,
    // Next WpoIdx to assign
    next_idx: WpoIdx,
}

impl<NodeId> WeakPartialOrderingImpl<NodeId>
where
    NodeId: Copy + Hash + Ord + Debug,
{
    pub fn new() -> Self {
        // I really don't want to add `Default` bound to `NodeId`, so let's
        // have a bit tedious code here to give user side more flexibility.
        Self {
            next_dfn: 1u32,
            nodes: vec![],
            toplevel: vec![],
            post_dfn: HashMap::new(),
            dfn: HashMap::new(),
            dfn_to_node: vec![],
            cross_fwd_edges: HashMap::new(),
            back_preds: HashMap::new(),
            non_back_preds: HashMap::new(),
            dfn_to_wpo_idx: vec![],
            next_idx: 0,
        }
    }

    fn add_node(&mut self, dfn_i: u32, vertex: u32, sz: u32, ty: WpoNodeType) {
        self.dfn_to_wpo_idx[dfn_i as usize] = self.next_idx;
        self.next_idx += 1;
        self.nodes.push(WpoNode::new(
            ty,
            // dfn reserves 0, so should subtract 1 here.
            self.dfn_to_node[vertex as usize - 1],
            sz as usize,
        ));
    }

    fn node_of(&mut self, dfn_i: u32) -> &mut WpoNode<NodeId> {
        let idx = self.index_of(dfn_i) as usize;
        &mut self.nodes[idx]
    }

    fn index_of(&self, dfn_i: u32) -> u32 {
        self.dfn_to_wpo_idx[dfn_i as usize]
    }

    fn add_successor(
        &mut self,
        from: u32,
        to: u32,
        exit: u32,
        outer_pred: bool,
        for_outer_preds: &mut Vec<(WpoIdx, WpoIdx)>,
    ) {
        let from_idx = self.index_of(from);
        let to_idx = self.index_of(to);
        if !self.nodes[from_idx as usize].is_successor(to_idx) {
            if outer_pred {
                for_outer_preds.push((to_idx, self.index_of(exit)));
            }
            self.nodes[from_idx as usize].add_successor(to_idx);
            self.nodes[to_idx as usize].add_predecessor(from_idx);
        }
    }

    fn build_auxilary<SN>(&mut self, size: usize, root: NodeId, successors_nodes: &SN)
    where
        SN: SuccessorNodes<NodeId = NodeId>,
    {
        // Since 0 is reserved for undiscovered nodes, the total number of nodes
        // would be size + 1.
        let mut dft_dsets = UnionFind::<u32>::new(size + 1);
        let mut stack = Vec::new();
        let mut next_post_dfn = 1u32;
        let mut visited = HashMap::new();
        let mut ancestor = HashMap::new();

        let get_dfn = |n: NodeId, dfn: &HashMap<NodeId, u32>| {
            // If the key does not exist, meaning that node is not
            // discovered yet, return default value 0.
            dfn.get(&n).copied().unwrap_or_default()
        };
        let set_dfn = |n: NodeId, num: u32, dfn: &mut HashMap<NodeId, u32>| {
            dfn.insert(n, num);
        };

        stack.push((root, false, 0u32));

        while let Some((node, finished, pred)) = stack.pop() {
            if finished {
                self.post_dfn.insert(node, next_post_dfn);
                next_post_dfn += 1;

                let vertex = get_dfn(node, &self.dfn);
                visited.insert(vertex, true);

                dft_dsets.union(vertex, pred);
                ancestor.insert(dft_dsets.find_mut(pred), pred);
            } else {
                if get_dfn(node, &self.dfn) != 0 {
                    // Skip forward edges.
                    continue;
                }

                let vertex = self.next_dfn;
                self.next_dfn += 1;
                self.dfn_to_node.push(node);
                set_dfn(node, vertex, &mut self.dfn);
                ancestor.insert(vertex, vertex);

                stack.push((node, true, pred));

                let successors = successors_nodes.get_succ_nodes(node);
                for &succ_node in successors.iter().rev() {
                    let succ = get_dfn(succ_node, &self.dfn);
                    if 0 == succ {
                        stack.push((succ_node, false, vertex));
                    } else if visited.get(&succ).copied().unwrap_or_default() {
                        let lca = ancestor.get(&dft_dsets.find_mut(succ)).copied().unwrap();
                        self.cross_fwd_edges
                            .entry(lca)
                            .or_default()
                            .push((vertex, succ));
                    } else {
                        self.back_preds.entry(succ).or_default().push(vertex);
                    }
                }

                if pred != 0 {
                    self.non_back_preds.entry(vertex).or_default().push(pred);
                }
            }
        }

        // Number of dfn should be equal or smaller (if there is unreachable node)
        // than grpah size + 1 (number 0 for undiscovered).
        assert!(self.next_dfn as usize <= size + 1);
    }

    fn build<SN>(
        mut self,
        size: usize,
        root: NodeId,
        successors_nodes: &SN,
    ) -> WeakPartialOrdering<NodeId>
    where
        SN: SuccessorNodes<NodeId = NodeId>,
    {
        // Step 1: construct auxilary data structures, including
        // classifying edges, finding lowest common ancestors
        // of cross/forward edges.
        self.build_auxilary(size, root, successors_nodes);

        // Step 2: start constructing WPO.
        let mut dsets = UnionFind::<u32>::new(self.next_dfn as usize);
        // Union find does not guarantee that the root of a subset has
        // always the minimum DFN, so we need to maintain this information.
        // Used for creating exit nodes.
        let mut exit_next_dfn = self.next_dfn;
        // Initializaiton.
        let mut rep: Vec<u32> = (0..self.next_dfn).collect();
        let mut exit: Vec<u32> = (0..self.next_dfn).collect();
        let mut origin: Vec<Vec<(u32, u32)>> = (0..self.next_dfn)
            .map(|v| {
                self.non_back_preds
                    .get(&v)
                    .map_or_else(std::vec::Vec::new, |non_back_preds_v| {
                        non_back_preds_v.iter().map(|&p| (p, v)).collect()
                    })
            })
            .collect();

        self.dfn_to_wpo_idx.resize(2 * self.next_dfn as usize, 0);
        let mut for_outer_preds = Vec::<(WpoIdx, WpoIdx)>::new();
        let mut components_sizes = vec![0u32; self.next_dfn as usize];
        let mut parent = HashMap::<WpoIdx, WpoIdx>::new();

        // In descending order, excluding 0 which is for undiscovered.
        for h in (1..self.next_dfn).rev() {
            // Restore cross/forward edges
            if let Some(edges) = self.cross_fwd_edges.get(&h) {
                for &(u, v) in edges {
                    let rep_v = rep[dsets.find(v) as usize];
                    self.non_back_preds.entry(rep_v).or_default().push(u);
                    origin[rep_v as usize].push((u, v));
                }
            }

            // Find nested SCCs.
            let mut is_scc = false;
            let mut backpreds_h = HashSet::<u32>::new();
            if let Some(preds) = self.back_preds.get(&h) {
                for &v in preds {
                    if v != h {
                        backpreds_h.insert(rep[dsets.find(v) as usize]);
                    } else {
                        is_scc = true;
                    }
                }
            }

            if !backpreds_h.is_empty() {
                is_scc = true;
            }

            let mut nested_sccs_h = backpreds_h.clone();
            let mut worklist_h = Vec::from_iter(backpreds_h.iter().copied());
            while let Some(v) = worklist_h.pop() {
                if let Some(preds) = self.non_back_preds.get(&v) {
                    for &p in preds {
                        let rep_p = rep[dsets.find(p) as usize];
                        if !nested_sccs_h.contains(&rep_p) && rep_p != h {
                            worklist_h.push(rep_p);
                            nested_sccs_h.insert(rep_p);
                        }
                    }
                }
            }

            // h represents a trivial SCC.
            if !is_scc {
                components_sizes[h as usize] = 1;
                self.add_node(h, h, 1, WpoNodeType::Plain);
                continue;
            }

            // Initialize size to 2 for head and exit.
            let mut sz_h = 2;
            for &v in nested_sccs_h.iter() {
                sz_h += components_sizes[v as usize];
            }
            components_sizes[h as usize] = sz_h;

            // Add new exit.
            let x_h = exit_next_dfn;
            exit_next_dfn += 1;
            self.add_node(x_h, h, sz_h, WpoNodeType::Exit);
            // Wpo index of head is then exit + 1 for the same component.
            self.add_node(h, h, sz_h, WpoNodeType::Head);

            if backpreds_h.is_empty() {
                // Scheduling constraints from h to x_h.
                self.add_successor(h, x_h, x_h, false, &mut for_outer_preds);
            } else {
                for p in backpreds_h {
                    self.add_successor(exit[p as usize], x_h, x_h, false, &mut for_outer_preds);
                }
            }

            // Scheduling constraints between WPOs for nested SCCs.
            for &v in nested_sccs_h.iter() {
                for &(u, vv) in origin[v as usize].iter() {
                    let x_u = exit[rep[dsets.find(u) as usize] as usize];
                    let x_v = exit[v as usize];
                    self.add_successor(x_u, vv, x_v, x_v != v, &mut for_outer_preds);
                }
            }

            // Merging all reps in nested SCCs to h
            for &v in nested_sccs_h.iter() {
                dsets.union(v, h);
                rep[dsets.find(v) as usize] = h;
                parent.insert(self.index_of(v), self.index_of(h));
            }

            exit[h as usize] = x_h;
        }

        // Scheduling constraints between WPOs for maximal SCCs.
        self.toplevel.reserve(self.next_dfn as usize);
        for v in 1..self.next_dfn {
            if rep[dsets.find(v) as usize] == v {
                let v_idx = self.index_of(v);
                self.toplevel.push(v_idx);
                parent.insert(v_idx, v_idx);

                for &(u, vv) in origin[v as usize].iter() {
                    let x_u = exit[rep[dsets.find(u) as usize] as usize];
                    let x_v = exit[v as usize];
                    self.add_successor(x_u, vv, x_v, x_v != v, &mut for_outer_preds);
                }
            }
        }

        // Compute num_outer_preds.
        for &(v, x_max) in for_outer_preds.iter() {
            let mut h = if self.nodes[v as usize].is_head() {
                v
            } else {
                *parent.get(&v).unwrap()
            };
            let mut x = h - 1;
            while x != x_max {
                self.nodes[x as usize].inc_num_outer_preds(v);
                h = *parent.get(&h).unwrap();
                x = h - 1;
            }
            self.nodes[x as usize].inc_num_outer_preds(v);
        }

        WeakPartialOrdering {
            nodes: self.nodes,
            toplevel: self.toplevel,
            post_dfn: self.post_dfn,
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

pub trait AbstractDomain: Clone + Eq {
    fn bottom() -> Self;
    fn top() -> Self;
    fn is_bottom(&self) -> bool;
    fn is_top(&self) -> bool;
    fn leq(&self, rhs: &Self) -> bool;

    fn join(mut self, rhs: Self) -> Self {
        self.join_with(rhs);
        self
    }

    fn meet(mut self, rhs: Self) -> Self {
        self.meet_with(rhs);
        self
    }

    fn widen(mut self, rhs: Self) -> Self {
        self.widen_with(rhs);
        self
    }

    fn narrow(mut self, rhs: Self) -> Self {
        self.narrow_with(rhs);
        self
    }

    fn join_with(&mut self, rhs: Self);
    fn meet_with(&mut self, rhs: Self);
    fn widen_with(&mut self, rhs: Self);
    fn narrow_with(&mut self, rhs: Self);
}