// Copyright 2020 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use {
    crate::{
        error::Error,
        extent::Extent,
        format::ExtentClusterHeader,
        options::ExtractorOptions,
        properties::{DataKind, ExtentKind},
        utils::{RangeOps, ReadAndSeek},
    },
    std::io::{Read, Write},
};

/// Returns `true` if the extent's data needs to be dumped.
fn should_dump_data(extent: &Extent, dump_pii: bool) -> bool {
    let properties = extent.properties();
    // For ExtentKind::Data, we do not need to dump data when DataKind is
    // either Skipped or it is Zeroes.
    // We dump Pii only when we forced to dump Pii or it was `Modified` by the
    // storage software to obfuscate the Pii data.
    match properties.extent_kind {
        ExtentKind::Data => match properties.data_kind {
            DataKind::Modified => true,
            DataKind::Unmodified => true,
            _ => false,
        },
        ExtentKind::Pii => match properties.data_kind {
            DataKind::Modified => true,
            DataKind::Unmodified => dump_pii,
            _ => false,
        },
        _ => false,
    }
}

/// This is in-memory representation of a collection of Extents.
///
/// We hold extents in memory, (coalescing, splitting, merging if needed), till
/// write() is issued.
/// The cluster of extents lives in image file in contiguous location.
#[derive(Debug)]
pub struct ExtentCluster {
    pub(crate) extent_tree: std::collections::BTreeMap<u64, Extent>,
    options: ExtractorOptions,
}

impl ExtentCluster {
    pub fn new(options: &ExtractorOptions) -> Self {
        Self { extent_tree: Default::default(), options: options.clone() }
    }

    // insert_extent's logic is as follows.
    //  * create a list of all affected extents by this insertion.
    //  * for each extent in affected extents, split/merge with current extents.
    //  * replace affected and current extent with the new split/merged extents list.
    fn insert_extent(&mut self, current_extent: Extent) -> Result<(), Error> {
        let mut affected_extents = vec![];

        // Get all extents that may be affected by this extent insertion.
        for ext in self.extent_tree.iter_mut() {
            if !(current_extent.overlaps(&ext.1) || current_extent.is_adjacent(&ext.1)) {
                continue;
            }
            affected_extents.push(ext.1.clone());
        }

        // Remove all the affected extents.
        for ext in &affected_extents {
            self.extent_tree.remove(&ext.start());
        }

        // Perform split/merge of current extent with one affected extent at a time.
        // Note:
        //   1. This is performed on extents in the ascending order of extent.start.
        //   2. The existing extents are assumed to be non-overlapping with all other
        //      existing extents.
        //   3. Order of insertion changes in intermediate state of the cluster.
        let mut new_extents = vec![];
        let mut remaining = Some(current_extent.clone());
        for (i, ext) in affected_extents.iter().enumerate() {
            assert!(current_extent.overlaps(&ext) || current_extent.is_adjacent(&ext));
            // If remaining extent is completely consumed then iterate over the rest of affected
            // extents and just add them.
            // Ex. say we have already inserted extents [2, 5), [5, 8), and [8, 12) all having
            // different priorities. Now the current extent is [5, 8) with lower
            // priority than all the the extents then we may end up consuming current
            // entirely before we parse all of affected_extents.
            match remaining {
                None => new_extents.push(ext.clone()),
                Some(remaining_extent) => {
                    remaining = remaining_extent.split_or_merge(&ext, &mut new_extents);
                    assert!(remaining.is_some() || (i >= affected_extents.len() - 2));
                }
            }
        }

        // Insert all the split/merged extent into the extent tree.
        for ext in new_extents {
            self.extent_tree.insert(ext.start(), ext.clone());
        }

        // It may happen that the current extent maybe unaffected or partially affected.
        // If so insert it into the tree.
        match remaining {
            Some(e) => {
                self.extent_tree.insert(e.start(), e);
            }
            _ => {}
        }

        Ok(())
    }

    /// Adds an extent to extent cluster.
    pub fn add_extent(&mut self, extent: &Extent) -> Result<(), Error> {
        if !extent.storage_range().is_valid() {
            return Err(Error::InvalidRange);
        }
        self.insert_extent(extent.clone())
    }

    /// Returns number of extent in extent cluster.
    pub fn extent_count(&self) -> u64 {
        self.extent_tree.len() as u64
    }

    /// Returns number of data bytes that will be written
    /// in this cluster.
    fn data_size(&self) -> u64 {
        let mut size: u64 = 0;
        for (_, extent) in &self.extent_tree {
            if !should_dump_data(&extent, self.options.force_dump_pii) {
                continue;
            }
            size = size + extent.storage_range().length();
        }
        size
    }

    /// Writes ExtentCluster metadata to out_stream. crc32 to is crc of metadata and any
    /// padding required.
    fn write_metadata(&self, out_stream: &mut dyn Write) -> Result<u64, Error> {
        let mut size =
            ExtentClusterHeader::serialize_to(self.extent_count() as u64, 0, out_stream)?;

        for (_, extent) in &self.extent_tree {
            size = size + extent.write(out_stream)?;
        }
        let zero_fill_len = (((size + self.options.alignment - 1) / self.options.alignment)
            * self.options.alignment)
            - size;
        let zeroes: Vec<u8> = vec![0; zero_fill_len as usize];
        out_stream.write_all(&zeroes).map_err(move |_| Error::WriteFailed)?;
        Ok(size + zero_fill_len)
    }

    /// Write all the data in extents to the out_stream.
    ///
    /// The extent's data is read from the in_stream.
    fn write_data(
        &self,
        in_stream: &mut dyn ReadAndSeek,
        out_stream: &mut dyn Write,
    ) -> Result<u64, Error> {
        let mut size = 0;
        for (_, extent) in &self.extent_tree {
            // No need to dump the data. Only ExtentInfo will be written to the image file.
            if !should_dump_data(&extent, self.options.force_dump_pii) {
                continue;
            }
            let bytes_to_read = extent.storage_range().length();
            let mut read_buffer: Vec<u8> = vec![0; bytes_to_read as usize];
            let offset = extent.storage_range().start;
            // Seek to the storage location.
            in_stream.seek(std::io::SeekFrom::Start(offset)).map_err(move |_| Error::SeekFailed)?;
            let mut r = in_stream.take(bytes_to_read);

            // Read from the storage.
            r.read(&mut read_buffer[..bytes_to_read as usize]).map_err(move |_| {
                eprintln!(
                    "Failed to read {} bytes from {} offset for {:?}",
                    bytes_to_read, offset, extent
                );
                Error::ReadFailed
            })?;

            // Write to the image file
            out_stream.write_all(&read_buffer).map_err(move |_| Error::WriteFailed)?;
            size = size + bytes_to_read;
        }
        Ok(size)
    }

    /// Iterate over the extent_tree and check that they are in ascending order
    /// of start offset and they do not overlap.
    fn check_extent_tree(&self) {
        let mut prev_or: Option<Extent> = None;

        for (_, extent) in &self.extent_tree {
            if prev_or.is_some() {
                let prev = prev_or.unwrap();
                assert!(extent.start() >= prev.end());
            }

            assert!(extent.storage_range().is_valid());
            prev_or = Some(extent.clone());
        }
    }

    /// Writes extent cluster to the image file.
    ///
    /// # Arguments
    ///
    /// `out_stream`    : Points to the image file stream.
    /// `in_stream`     : Stream from where extent data will be read from.
    /// `current_offset`: Starting position of the cluster within the image file.
    /// `last_cluster`  : True if this is the last cluster within the extracted image file.
    ///                   Multi-cluster image files are not yet supported.
    ///
    /// This includes computing checksum of the cluster header and all the extent infos,
    /// then writing cluster header, extent info and extent data to the image file.
    pub fn write(
        &mut self,
        mut in_stream: &mut dyn ReadAndSeek,
        _current_offset: u64,
        last_cluster: bool,
        mut out_stream: &mut dyn Write,
    ) -> Result<u64, Error> {
        if !last_cluster {
            todo!("Support for more than one cluster is not implemented.")
        }

        self.check_extent_tree();
        // Update all the extent location w.r.t. starting of the extent cluster.
        let _data_size = self.data_size();

        let mut size = self.write_metadata(out_stream)?;
        size = size + self.write_data(&mut in_stream, &mut out_stream)?;
        Ok(size)
    }
}

#[cfg(test)]
mod test {
    use {
        super::*,
        crate::{
            extent_cluster::ExtentCluster,
            properties::{DataKind, ExtentKind, ExtentProperties},
        },
        std::io::Cursor,
        std::ops::Range,
    };

    static INSERTED_ADDRESS_START: u64 = 20;
    static INSERTED_ADDRESS_END: u64 = 30;

    static LOW_PRIORITY_PROPERTIES: ExtentProperties =
        ExtentProperties { extent_kind: ExtentKind::Data, data_kind: DataKind::Zeroes };

    static HIGH_PRIORITY_PROPERTIES: ExtentProperties =
        ExtentProperties { extent_kind: ExtentKind::Pii, data_kind: DataKind::Zeroes };

    static INSERTED_PROPERTIES: ExtentProperties =
        ExtentProperties { extent_kind: ExtentKind::Data, data_kind: DataKind::Unmodified };

    static OVERLAPPING_RIGHT_ADDRESS: Range<u64> =
        INSERTED_ADDRESS_END - 5..INSERTED_ADDRESS_END + 6;

    fn inserted_range() -> Range<u64> {
        INSERTED_ADDRESS_START..INSERTED_ADDRESS_END
    }

    fn inserted_extent() -> Extent {
        Extent::new(
            inserted_range(),
            INSERTED_PROPERTIES,
            // data: Some(vec![1; inserted_range().length() as usize]),
            None,
        )
        .unwrap()
    }

    // Non-overlapping extent to the right of inserted_extent().
    fn right_extent() -> Extent {
        Extent::new(
            INSERTED_ADDRESS_END + 5..INSERTED_ADDRESS_END + 15,
            INSERTED_PROPERTIES,
            // data: Some(vec![1; inserted_range().length() as usize]),
            None,
        )
        .unwrap()
    }
    fn setup_extent_cluster() -> (ExtentCluster, Vec<Extent>) {
        let mut cluster = ExtentCluster::new(&Default::default());
        match cluster.add_extent(&inserted_extent()) {
            Err(why) => println!("why: {:?}", why),
            Ok(_) => {}
        }
        return (cluster, vec![inserted_extent().clone()]);
    }

    fn verify(file: &str, line: u32, cluster: &ExtentCluster, extents: &Vec<Extent>) {
        if cluster.extent_tree.len() != extents.len() {
            println!("{}:{} Expected: {:?}\nFound: {:?}", file, line, extents, cluster.extent_tree);
        }
        assert_eq!(cluster.extent_tree.len(), extents.len());

        for (_, ext) in cluster.extent_tree.iter() {
            let mut found = false;
            for inserted in extents.iter() {
                if inserted == ext {
                    found = true;
                    break;
                }
            }
            if !found {
                println!("Could not find {:?}", *ext);
            }
            assert!(found);
        }
    }

    #[test]
    fn test_setup() {
        let (cluster, expected_extents) = setup_extent_cluster();
        assert!(INSERTED_PROPERTIES > LOW_PRIORITY_PROPERTIES);
        assert!(INSERTED_PROPERTIES < HIGH_PRIORITY_PROPERTIES);
        assert!(LOW_PRIORITY_PROPERTIES < HIGH_PRIORITY_PROPERTIES);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_non_overlapping() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let e = Extent::new(40..50, INSERTED_PROPERTIES, None).unwrap();
        cluster.add_extent(&e).unwrap();
        expected_extents.push(e);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_is_adjacent_not_mergable() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let e = Extent::new(
            INSERTED_ADDRESS_END..INSERTED_ADDRESS_END + 5,
            LOW_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&e).unwrap();
        expected_extents.push(e);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_is_adjacent() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let e =
            Extent::new(INSERTED_ADDRESS_END..INSERTED_ADDRESS_END + 5, INSERTED_PROPERTIES, None)
                .unwrap();
        cluster.add_extent(&e).unwrap();
        expected_extents.clear();
        expected_extents.push(
            Extent::new(
                INSERTED_ADDRESS_START..INSERTED_ADDRESS_END + 5,
                INSERTED_PROPERTIES,
                None,
            )
            .unwrap(),
        );
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_adjacent_in_the_middle() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        verify(file!(), line!(), &cluster, &expected_extents);
        let e =
            Extent::new(INSERTED_ADDRESS_END..INSERTED_ADDRESS_END + 5, INSERTED_PROPERTIES, None)
                .unwrap();
        cluster.add_extent(&e).unwrap();
        expected_extents.clear();
        expected_extents.push(
            Extent::new(
                INSERTED_ADDRESS_START..INSERTED_ADDRESS_END + 15,
                INSERTED_PROPERTIES,
                None,
            )
            .unwrap(),
        );
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_adjacent_in_the_middle_not_mergable() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        verify(file!(), line!(), &cluster, &expected_extents);
        let m = Extent::new(
            INSERTED_ADDRESS_END..INSERTED_ADDRESS_END + 5,
            HIGH_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&m).unwrap();
        expected_extents.push(m);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_adjacent_on_both_sides_overlapping_in_middle() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        // Adjacent to left
        let mut m = Extent::new(
            INSERTED_ADDRESS_START - 10..INSERTED_ADDRESS_START,
            LOW_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&m).unwrap();
        expected_extents.push(m.clone());
        m.set_start(INSERTED_ADDRESS_END);
        m.set_end(INSERTED_ADDRESS_END + 10);
        cluster.add_extent(&m).unwrap();
        expected_extents.push(m.clone());

        verify(file!(), line!(), &cluster, &expected_extents);

        let m = Extent::new(
            INSERTED_ADDRESS_START..INSERTED_ADDRESS_END,
            LOW_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&m).unwrap();
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster
            .add_extent(
                &Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), INSERTED_PROPERTIES, None).unwrap(),
            )
            .unwrap();
        expected_extents[0].set_end(OVERLAPPING_RIGHT_ADDRESS.end);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_low_priority() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let mut e =
            Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), LOW_PRIORITY_PROPERTIES, None).unwrap();
        cluster.add_extent(&e).unwrap();
        e.set_start(INSERTED_ADDRESS_END);
        expected_extents.push(e);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_high_priority() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster
            .add_extent(
                &Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), HIGH_PRIORITY_PROPERTIES, None)
                    .unwrap(),
            )
            .unwrap();
        expected_extents[0].set_end(OVERLAPPING_RIGHT_ADDRESS.start);
        expected_extents.push(
            Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), HIGH_PRIORITY_PROPERTIES, None).unwrap(),
        );
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_in_middlex() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        verify(file!(), line!(), &cluster, &expected_extents);

        let middle_extent =
            Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), INSERTED_PROPERTIES, None).unwrap();
        cluster.add_extent(&middle_extent).unwrap();
        expected_extents.clear();
        expected_extents.push(
            Extent::new(INSERTED_ADDRESS_START..right_extent().end(), INSERTED_PROPERTIES, None)
                .unwrap(),
        );
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_in_middle_low_priority() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        verify(file!(), line!(), &cluster, &expected_extents);

        cluster
            .add_extent(
                &Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), LOW_PRIORITY_PROPERTIES, None)
                    .unwrap(),
            )
            .unwrap();
        expected_extents.push(
            Extent::new(
                INSERTED_ADDRESS_END..INSERTED_ADDRESS_END + 5,
                LOW_PRIORITY_PROPERTIES,
                None,
            )
            .unwrap(),
        );
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_in_middle_high_priority() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        verify(file!(), line!(), &cluster, &expected_extents);

        let middle_extent =
            Extent::new(OVERLAPPING_RIGHT_ADDRESS.clone(), HIGH_PRIORITY_PROPERTIES, None).unwrap();
        cluster.add_extent(&middle_extent).unwrap();
        expected_extents[0].set_end(OVERLAPPING_RIGHT_ADDRESS.start);
        expected_extents[1].set_start(OVERLAPPING_RIGHT_ADDRESS.end);
        expected_extents.push(middle_extent);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_multiple() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        let extreme_right_extent = Extent::new(
            right_extent().end() + 10..right_extent().end() + 20,
            INSERTED_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&extreme_right_extent).unwrap();
        expected_extents.push(extreme_right_extent.clone());
        verify(file!(), line!(), &cluster, &expected_extents);

        let overlapping_extent = Extent::new(
            INSERTED_ADDRESS_START - 5..extreme_right_extent.end() + 10,
            INSERTED_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&overlapping_extent).unwrap();
        expected_extents.clear();
        expected_extents.push(overlapping_extent);
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_overlapping_multiple_not_mergeable() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        cluster.add_extent(&right_extent()).unwrap();
        expected_extents.push(right_extent().clone());
        let extreme_right_extent = Extent::new(
            right_extent().end() + 10..right_extent().end() + 20,
            INSERTED_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&extreme_right_extent).unwrap();
        expected_extents.push(extreme_right_extent.clone());
        verify(file!(), line!(), &cluster, &expected_extents);

        let overlapping_extent = Extent::new(
            INSERTED_ADDRESS_START - 5..extreme_right_extent.end() + 10,
            LOW_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&overlapping_extent).unwrap();
        // The overlapping_extent gets divided into multiple extents.
        let mut split_extent = overlapping_extent.clone();
        split_extent.set_end(inserted_extent().start());
        expected_extents.push(split_extent.clone());
        split_extent.set_start(inserted_extent().end());
        split_extent.set_end(right_extent().start());
        expected_extents.push(split_extent.clone());
        split_extent.set_start(right_extent().end());
        split_extent.set_end(extreme_right_extent.start());
        expected_extents.push(split_extent.clone());
        split_extent.set_start(extreme_right_extent.end());
        split_extent.set_end(overlapping_extent.end());
        expected_extents.push(split_extent.clone());
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    #[test]
    fn test_add_splits_existing_extent() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let small_extent = Extent::new(
            inserted_extent().start() + 3..inserted_extent().end() - 3,
            HIGH_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&small_extent).unwrap();
        expected_extents.push(expected_extents[0].clone());
        expected_extents[0].set_end(small_extent.start());
        expected_extents[1].set_start(small_extent.end());
        expected_extents.push(small_extent.clone());
        verify(file!(), line!(), &cluster, &expected_extents);
    }
    #[test]
    fn test_add_splits_existing_extent_at_start() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let small_extent = Extent::new(
            inserted_extent().start()..inserted_extent().end() - 3,
            HIGH_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&small_extent).unwrap();
        expected_extents[0].set_start(small_extent.end());
        expected_extents.push(small_extent.clone());
        verify(file!(), line!(), &cluster, &expected_extents);
    }
    #[test]
    fn test_add_splits_existing_extent_at_end() {
        let (mut cluster, mut expected_extents) = setup_extent_cluster();
        let small_extent = Extent::new(
            inserted_extent().start() + 3..inserted_extent().end(),
            HIGH_PRIORITY_PROPERTIES,
            None,
        )
        .unwrap();
        cluster.add_extent(&small_extent).unwrap();
        expected_extents[0].set_end(small_extent.start());
        expected_extents.push(small_extent.clone());
        verify(file!(), line!(), &cluster, &expected_extents);
    }

    fn dumpable_data_properties() -> ExtentProperties {
        ExtentProperties { extent_kind: ExtentKind::Data, data_kind: DataKind::Modified }
    }

    fn skippable_data_properties() -> ExtentProperties {
        ExtentProperties { extent_kind: ExtentKind::Unmmapped, data_kind: DataKind::Skipped }
    }

    fn pii_data_properties() -> ExtentProperties {
        ExtentProperties { extent_kind: ExtentKind::Pii, data_kind: DataKind::Unmodified }
    }

    fn setup_cluster_write_test(
        dump_pii: bool,
    ) -> (ExtentCluster, ExtractorOptions, Vec<u8>, Cursor<Vec<u8>>) {
        let mut options: ExtractorOptions = Default::default();
        options.force_dump_pii = dump_pii;
        let cluster = ExtentCluster::new(&options);
        let out_buffer: Vec<u8> = vec![];
        let in_buffer = Cursor::new(vec![0; 2 * 1024 * 1024]);

        (cluster, options, out_buffer, in_buffer)
    }

    #[test]
    fn test_extent_cluster_write() {
        let (mut cluster, options, mut out_buffer, mut in_buffer) = setup_cluster_write_test(false);

        let size = cluster.write(&mut in_buffer, 0, true, &mut out_buffer).unwrap();
        assert!(size > 0);
        assert_eq!(size % options.alignment, 0);
        assert_eq!(out_buffer.len() as u64, size);
    }

    #[test]
    fn test_extent_cluster_write_no_data() {
        let (mut cluster, options, mut out_buffer, mut in_buffer) = setup_cluster_write_test(false);
        let size = cluster.write(&mut in_buffer, 0, true, &mut out_buffer).unwrap();
        assert!(size > 0);
        assert_eq!(size % options.alignment, 0);
        assert_eq!(out_buffer.len() as u64, size);

        let properties = skippable_data_properties();
        let mut extent = Extent::new(0..1, properties, None).unwrap();

        // Write as many extents as it takes to fill one alignment unit.
        let extent_count = options.alignment / extent.serialized_size() as u64;
        for i in 0..extent_count {
            let start = i * 2 * options.alignment;
            extent.set_start(start);
            extent.set_end(start + options.alignment);
            cluster.add_extent(&extent).unwrap();
        }

        let mut new_buffer: Vec<u8> = vec![];
        let new_size = cluster.write(&mut in_buffer, 0, true, &mut new_buffer).unwrap();
        assert!(new_size > size);
        assert_eq!(new_size % options.alignment, 0);
        assert_eq!(new_buffer.len() as u64, new_size);
    }

    #[test]
    fn test_extent_cluster_write_with_data() {
        let (mut cluster, options, mut out_buffer, mut in_buffer) = setup_cluster_write_test(false);

        let no_data_extent =
            Extent::new(8192..(2 * 8192), skippable_data_properties(), None).unwrap();
        let data_extent1 =
            Extent::new((2 * 8192)..(3 * 8192), dumpable_data_properties(), None).unwrap();
        let data_extent2 =
            Extent::new((5 * 8192)..(6 * 8192), dumpable_data_properties(), None).unwrap();
        let pii_extent = Extent::new(8 * 8192..(9 * 8192), pii_data_properties(), None).unwrap();

        cluster.add_extent(&no_data_extent).unwrap();
        cluster.add_extent(&data_extent1).unwrap();
        cluster.add_extent(&data_extent2).unwrap();
        cluster.add_extent(&pii_extent).unwrap();

        let new_size = cluster.write(&mut in_buffer, 0, true, &mut out_buffer).unwrap();
        assert_eq!(new_size % options.alignment, 0);
        assert_eq!(out_buffer.len() as u64, new_size);

        // We end up writing three aligned segments.
        // - one cluster header and extents
        // - two data blocks
        assert_eq!(new_size, options.alignment * 3)
    }

    #[test]
    fn test_extent_cluster_write_with_pii() {
        let (mut cluster, options, mut out_buffer, mut in_buffer) = setup_cluster_write_test(true);

        // Add skippable data followed by pii at same range.
        let no_data_extent =
            Extent::new(8192..(15 * 8192), skippable_data_properties(), None).unwrap();
        let pii_extent = Extent::new(8192..(5 * 8192), pii_data_properties(), None).unwrap();

        cluster.add_extent(&no_data_extent).unwrap();
        cluster.add_extent(&pii_extent).unwrap();

        let new_size = cluster.write(&mut in_buffer, 0, true, &mut out_buffer).unwrap();
        assert_eq!(new_size % options.alignment, 0);
        assert_eq!(out_buffer.len() as u64, new_size);

        // We end up writing three aligned segments.
        // - one cluster header and extents
        // - four for pii blocks
        assert_eq!(new_size, options.alignment * 5)
    }

    fn verify_should_dump_data(ekind: ExtentKind, dkind: DataKind, dump_pii: bool, dump: bool) {
        let properties = ExtentProperties { extent_kind: ekind, data_kind: dkind };
        let extent = Extent::new(4..10, properties, None).unwrap();
        assert_eq!(
            should_dump_data(&extent, dump_pii),
            dump,
            "{:?} {} {}",
            properties,
            dump_pii,
            dump
        );
    }

    #[test]
    fn test_should_dump_data() {
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Skipped, true, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Zeroes, true, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Unmodified, true, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Modified, true, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Skipped, false, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Zeroes, false, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Unmodified, false, false);
        verify_should_dump_data(ExtentKind::Unmmapped, DataKind::Modified, false, false);

        verify_should_dump_data(ExtentKind::Unused, DataKind::Skipped, true, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Zeroes, true, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Unmodified, true, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Modified, true, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Skipped, false, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Zeroes, false, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Unmodified, false, false);
        verify_should_dump_data(ExtentKind::Unused, DataKind::Modified, false, false);

        verify_should_dump_data(ExtentKind::Data, DataKind::Skipped, true, false);
        verify_should_dump_data(ExtentKind::Data, DataKind::Zeroes, true, false);
        verify_should_dump_data(ExtentKind::Data, DataKind::Unmodified, true, true);
        verify_should_dump_data(ExtentKind::Data, DataKind::Modified, true, true);
        verify_should_dump_data(ExtentKind::Data, DataKind::Skipped, false, false);
        verify_should_dump_data(ExtentKind::Data, DataKind::Zeroes, false, false);
        verify_should_dump_data(ExtentKind::Data, DataKind::Unmodified, false, true);
        verify_should_dump_data(ExtentKind::Data, DataKind::Modified, false, true);

        verify_should_dump_data(ExtentKind::Pii, DataKind::Skipped, true, false);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Zeroes, true, false);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Unmodified, true, true);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Modified, true, true);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Skipped, false, false);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Zeroes, false, false);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Unmodified, false, false);
        verify_should_dump_data(ExtentKind::Pii, DataKind::Modified, false, true);
    }
}