<commit_before>use std::{
    collections::HashMap,
    ffi::OsString,
    fs, io,
    path::{Path, PathBuf},
};

use oci_spec::image::{Descriptor, DescriptorBuilder, MediaType};

use archive_utils::{flate2, tar, zstd};
use common::{
    chrono::Utc,
    eyre::{bail, eyre, Result},
    itertools::Itertools,
    tracing,
};
use hash_utils::{sha2::Digest, sha2::Sha256};

use crate::{blob_writer::BlobWriter, storage::digest_to_parts};

/// A change in a path between two snapshots
///
/// This enum represents the [Change Types](https://github.com/opencontainers/image-spec/blob/main/layer.md#change-types)
/// described in the OCI spec.
#[derive(Debug, PartialEq, Eq, Ord, PartialOrd)]
pub enum Change {
    Added(String),
    Modified(String),
    Removed(String),
}

/// The set of changes between two snapshots
///
/// Represents the set of changes between two filesystem snapshots as described in
/// [OCI Image Layer Filesystem Changeset](https://github.com/opencontainers/image-spec/blob/main/layer.md)
pub struct ChangeSet {
    /// The source directory, on the local filesystem, for the changes
    source_dir: PathBuf,

    /// The destination directory, within the image's root filesystem, for the changes
    dest_dir: PathBuf,

    /// The change items
    pub(crate) items: Vec<Change>,
}

impl ChangeSet {
    /// Create a new set of snapshot changes
    pub fn new<P: AsRef<Path>>(source_dir: P, dest_dir: P, items: Vec<Change>) -> Self {
        let source_dir = source_dir.as_ref().to_path_buf();

        // Parths in tar archive must be relative so stri any leading slash
        let dest_dir = dest_dir.as_ref().to_path_buf();
        let dest_dir = match dest_dir.strip_prefix("/") {
            Ok(dir) => dir.to_owned(),
            Err(_) => dest_dir,
        };

        Self {
            source_dir,
            dest_dir,
            items,
        }
    }

    /// Creates an OCI layer for the set of changes
    ///
    /// This implements the [Representing Changes](https://github.com/opencontainers/image-spec/blob/main/layer.md#representing-changes)
    /// section of the OCI image spec:
    ///
    /// - `Added` and `Modified` paths are added to the archive.
    /// - `Removed` paths are represented as "whiteout" files.
    ///
    /// Note that two SHA256 hashes are calculated, one for the `DiffID` of a changeset (calculated in this function
    /// and used in the image config file) and one for the digest which (calculated by the [`BlobWriter`] and used in the image manifest).
    /// A useful diagram showing how these are calculated and used is available
    /// [here](https://github.com/google/go-containerregistry/blob/main/pkg/v1/remote/README.md#anatomy-of-an-image-upload).
    ///
    /// # Arguments
    ///
    /// - `layout_dir`: the image directory to write the layer to (to the `blob/sha256` subdirectory)
    pub fn write_layer<P: AsRef<Path>>(
        self,
        media_type: &MediaType,
        layout_dir: P,
    ) -> Result<(String, Descriptor)> {
        if self.items.is_empty() {
            return Ok((
                "<empty>".to_string(),
                DescriptorBuilder::default()
                    .media_type(media_type.clone())
                    .digest("<none>")
                    .size(0)
                    .build()?,
            ));
        }

        tracing::info!(
            "Writing image layer from changeset for `{}`",
            self.source_dir.display()
        );

        let mut diffid_hash = Sha256::new();
        let mut blob_writer = BlobWriter::new()?;

        let changes = self.items.len();
        let mut additions: Vec<String> = Vec::new();
        let mut modifications: Vec<String> = Vec::new();
        let mut deletions: Vec<String> = Vec::new();

        {
            enum LayerEncoder<'lt> {
                Plain(&'lt mut BlobWriter),
                Gzip(flate2::write::GzEncoder<&'lt mut BlobWriter>),
                Zstd(zstd::stream::AutoFinishEncoder<'lt, &'lt mut BlobWriter>),
            }

            struct LayerWriter<'lt> {
                diffid_hash: &'lt mut Sha256,
                layer_encoder: LayerEncoder<'lt>,
            }

            impl<'lt> io::Write for LayerWriter<'lt> {
                fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
                    self.diffid_hash.update(buf);
                    match &mut self.layer_encoder {
                        LayerEncoder::Plain(encoder) => encoder.write_all(buf)?,
                        LayerEncoder::Gzip(encoder) => encoder.write_all(buf)?,
                        LayerEncoder::Zstd(encoder) => encoder.write_all(buf)?,
                    }
                    Ok(buf.len())
                }

                fn flush(&mut self) -> io::Result<()> {
                    Ok(())
                }
            }

            let layer_encoder = match media_type {
                MediaType::ImageLayer => LayerEncoder::Plain(&mut blob_writer),
                MediaType::ImageLayerGzip => LayerEncoder::Gzip(flate2::write::GzEncoder::new(
                    &mut blob_writer,
                    flate2::Compression::new(4),
                )),
                MediaType::ImageLayerZstd => LayerEncoder::Zstd(
                    zstd::stream::Encoder::new(&mut blob_writer, 4)?.auto_finish(),
                ),
                _ => bail!("Unhandled media type for layer: {}", media_type.to_string()),
            };

            let mut layer_writer = LayerWriter {
                diffid_hash: &mut diffid_hash,
                layer_encoder,
            };

            let mut archive = tar::Builder::new(&mut layer_writer);

            // Add an entry for the `dest_dir` (and any of its parent) so that ownership (and other
            // metadata) of `source_dir` is maintained. If not done then there are issues with non-root
            // users writing to the `workspace` and `layers` directories and  their subdirectories.
            let mut path = PathBuf::new();
            for part in self.dest_dir.components() {
                path = path.join(part);
                archive.append_path_with_name(&self.source_dir, &path)?;
            }

            // Add each change in sorted order so that the digest of the generated archive
            // does not change for a given change set
            for change in self.items.into_iter().sorted() {
                match change {
                    Change::Added(ref path) | Change::Modified(ref path) => {
                        let source_path = self.source_dir.join(path);
                        let dest_path = self.dest_dir.join(path);

                        let result = if source_path.is_symlink() {
                            match fs::read_link(&source_path).and_then(|target| {<commit_msg>If change_set includes a new folder, the generated sha256 digest will be different, even though the folder name is the same(but the origin super path is different)<commit_after>use std::{
    collections::HashMap,
    ffi::OsString,
    fs, io,
    path::{Path, PathBuf},
};

use oci_spec::image::{Descriptor, DescriptorBuilder, MediaType};

use archive_utils::{
    flate2,
    tar::{self, HeaderMode},
    zstd,
};
use common::{
    chrono::Utc,
    eyre::{bail, eyre, Result},
    itertools::Itertools,
    tracing,
};
use hash_utils::{sha2::Digest, sha2::Sha256};

use crate::{blob_writer::BlobWriter, storage::digest_to_parts};

/// A change in a path between two snapshots
///
/// This enum represents the [Change Types](https://github.com/opencontainers/image-spec/blob/main/layer.md#change-types)
/// described in the OCI spec.
#[derive(Debug, PartialEq, Eq, Ord, PartialOrd)]
pub enum Change {
    Added(String),
    Modified(String),
    Removed(String),
}

/// The set of changes between two snapshots
///
/// Represents the set of changes between two filesystem snapshots as described in
/// [OCI Image Layer Filesystem Changeset](https://github.com/opencontainers/image-spec/blob/main/layer.md)
pub struct ChangeSet {
    /// The source directory, on the local filesystem, for the changes
    source_dir: PathBuf,

    /// The destination directory, within the image's root filesystem, for the changes
    dest_dir: PathBuf,

    /// The change items
    pub(crate) items: Vec<Change>,
}

impl ChangeSet {
    /// Create a new set of snapshot changes
    pub fn new<P: AsRef<Path>>(source_dir: P, dest_dir: P, items: Vec<Change>) -> Self {
        let source_dir = source_dir.as_ref().to_path_buf();

        // Parths in tar archive must be relative so stri any leading slash
        let dest_dir = dest_dir.as_ref().to_path_buf();
        let dest_dir = match dest_dir.strip_prefix("/") {
            Ok(dir) => dir.to_owned(),
            Err(_) => dest_dir,
        };

        Self {
            source_dir,
            dest_dir,
            items,
        }
    }

    /// Creates an OCI layer for the set of changes
    ///
    /// This implements the [Representing Changes](https://github.com/opencontainers/image-spec/blob/main/layer.md#representing-changes)
    /// section of the OCI image spec:
    ///
    /// - `Added` and `Modified` paths are added to the archive.
    /// - `Removed` paths are represented as "whiteout" files.
    ///
    /// Note that two SHA256 hashes are calculated, one for the `DiffID` of a changeset (calculated in this function
    /// and used in the image config file) and one for the digest which (calculated by the [`BlobWriter`] and used in the image manifest).
    /// A useful diagram showing how these are calculated and used is available
    /// [here](https://github.com/google/go-containerregistry/blob/main/pkg/v1/remote/README.md#anatomy-of-an-image-upload).
    ///
    /// # Arguments
    ///
    /// - `layout_dir`: the image directory to write the layer to (to the `blob/sha256` subdirectory)
    pub fn write_layer<P: AsRef<Path>>(
        self,
        media_type: &MediaType,
        layout_dir: P,
    ) -> Result<(String, Descriptor)> {
        if self.items.is_empty() {
            return Ok((
                "<empty>".to_string(),
                DescriptorBuilder::default()
                    .media_type(media_type.clone())
                    .digest("<none>")
                    .size(0)
                    .build()?,
            ));
        }

        tracing::info!(
            "Writing image layer from changeset for `{}`",
            self.source_dir.display()
        );

        let mut diffid_hash = Sha256::new();
        let mut blob_writer = BlobWriter::new()?;

        let changes = self.items.len();
        let mut additions: Vec<String> = Vec::new();
        let mut modifications: Vec<String> = Vec::new();
        let mut deletions: Vec<String> = Vec::new();

        {
            enum LayerEncoder<'lt> {
                Plain(&'lt mut BlobWriter),
                Gzip(flate2::write::GzEncoder<&'lt mut BlobWriter>),
                Zstd(zstd::stream::AutoFinishEncoder<'lt, &'lt mut BlobWriter>),
            }

            struct LayerWriter<'lt> {
                diffid_hash: &'lt mut Sha256,
                layer_encoder: LayerEncoder<'lt>,
            }

            impl<'lt> io::Write for LayerWriter<'lt> {
                fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
                    self.diffid_hash.update(buf);
                    match &mut self.layer_encoder {
                        LayerEncoder::Plain(encoder) => encoder.write_all(buf)?,
                        LayerEncoder::Gzip(encoder) => encoder.write_all(buf)?,
                        LayerEncoder::Zstd(encoder) => encoder.write_all(buf)?,
                    }
                    Ok(buf.len())
                }

                fn flush(&mut self) -> io::Result<()> {
                    Ok(())
                }
            }

            let layer_encoder = match media_type {
                MediaType::ImageLayer => LayerEncoder::Plain(&mut blob_writer),
                MediaType::ImageLayerGzip => LayerEncoder::Gzip(flate2::write::GzEncoder::new(
                    &mut blob_writer,
                    flate2::Compression::new(4),
                )),
                MediaType::ImageLayerZstd => LayerEncoder::Zstd(
                    zstd::stream::Encoder::new(&mut blob_writer, 4)?.auto_finish(),
                ),
                _ => bail!("Unhandled media type for layer: {}", media_type.to_string()),
            };

            let mut layer_writer = LayerWriter {
                diffid_hash: &mut diffid_hash,
                layer_encoder,
            };

            let mut archive = tar::Builder::new(&mut layer_writer);
            archive.mode(HeaderMode::Deterministic);
            // Add an entry for the `dest_dir` (and any of its parent) so that ownership (and other
            // metadata) of `source_dir` is maintained. If not done then there are issues with non-root
            // users writing to the `workspace` and `layers` directories and  their subdirectories.
            let mut path = PathBuf::new();
            for part in self.dest_dir.components() {
                path = path.join(part);
                archive.append_path_with_name(&self.source_dir, &path)?;
            }

            // Add each change in sorted order so that the digest of the generated archive
            // does not change for a given change set
            for change in self.items.into_iter().sorted() {
                match change {
                    Change::Added(ref path) | Change::Modified(ref path) => {
                        let source_path = self.source_dir.join(path);
                        let dest_path = self.dest_dir.join(path);

                        let result = if source_path.is_symlink() {
                            match fs::read_link(&source_path).and_then(|target| {
{'fraction_non_alphanumeric': 0.12186120558651133, 'fraction_numerical': 0.003980027498371807, 'mean_word_length': 1.9663017814981756, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 73, 'https://': 8, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '944525', 'n_tokens_mistral': 3961, 'n_tokens_neox': 3800, 'n_words': 1039}