Perceptual-Image-Hashing/src/cache.rs

// Copyright 2016 Drew Short <drew@sothr.com>.
//
// Licensed under the MIT license<LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

extern crate flate2;
extern crate image;
extern crate num;
extern crate sha1;

use std::default::Default;
use std::fs::{create_dir_all, File, remove_dir_all};
use std::io::{Error, ErrorKind, Read, Write};
use std::option::Option;
use std::path::Path;
use std::result::Result;
use std::str::FromStr;

use super::rustc_serialize::json;

use self::flate2::Compression;
use self::flate2::read::ZlibDecoder;
use self::flate2::write::ZlibEncoder;
use self::image::DynamicImage;
use self::sha1::Sha1;

pub const DEFAULT_CACHE_DIR: &'static str = "./.hash_cache";
const CACHED_IMAGE_EXT: &'static str = "png";
const CACHED_MATRIX_EXT: &'static str = "dft";
// Caching version information
const CACHE_VERSION: u32 = 1;
const CACHE_METADATA_FILE: &'static str = "cache.meta";

#[derive(RustcDecodable, RustcEncodable)]
struct CacheMetadata {
    cache_version: u32,
}

impl Default for CacheMetadata {
    fn default() -> CacheMetadata {
        CacheMetadata {
            cache_version: CACHE_VERSION,
        }
    }
}

impl PartialEq<CacheMetadata> for CacheMetadata {
    fn eq(&self, other: &CacheMetadata) -> bool {
        self.cache_version == other.cache_version
    }

    fn ne(&self, other: &CacheMetadata) -> bool {
        !self.eq(&other)
    }
}

/**
* Structure to hold implementation of the cache
*/
#[repr(C)]
pub struct Cache {
    pub cache_dir: String,
    pub use_cache: bool,
}

impl Default for Cache {
    fn default() -> Cache {
        Cache {
            cache_dir: String::from(DEFAULT_CACHE_DIR),
            use_cache: true,
        }
    }
}

impl Cache {
    /**
     * Create the required directories for the cache
     */
    pub fn init(&self) -> Result<(), Error> {
        match create_dir_all(&self.cache_dir) {
            Ok(_) => {
                let metadata_path_str = format!("{}/{}", self.cache_dir, CACHE_METADATA_FILE);
                let metadata_path = Path::new(&metadata_path_str);
                let current_metadata: CacheMetadata = Default::default();
                match File::open(&metadata_path) {
                    Ok(mut file) => {
                        // Metadata file exists, compare them
                        let mut loaded_metadata_string = String::new();
                        match file.read_to_string(&mut loaded_metadata_string) {
                            Ok(_) => {
                                let loaded_metadata: CacheMetadata =
                                    match json::decode(&loaded_metadata_string) {
                                        Ok(data) => data,
                                        Err(_) => CacheMetadata { cache_version: 0 },
                                    };

                                // If they match, continue
                                if current_metadata != loaded_metadata {
                                    // If they don't wipe the cache to start new
                                    match remove_dir_all(&self.cache_dir) {
                                        Ok(_) => match create_dir_all(&self.cache_dir) {
                                            Ok(_) => (),
                                            Err(e) => println!("Error: {}", e),
                                        },
                                        Err(e) => println!("Error: {}", e),
                                    };
                                };
                            }
                            Err(e) => println!("Error: {}", e),
                        };
                    }
                    // Metadata file doesn't exist, do nothing assume all is well,
                    // create new metadata file
                    Err(_) => {}
                };
                let encoded_cache_metadata = json::encode(&current_metadata).unwrap();
                match File::create(&metadata_path) {
                    Ok(mut file) => {
                        let _ = file.write(&encoded_cache_metadata.as_bytes());
                        Ok(())
                    }
                    Err(e) => Err(e),
                }
            }
            Err(e) => Err(e),
        }
    }

    /**
     * Clean the cache directory completely
     */
    pub fn clean(&self) -> Result<(), Error> {
        remove_dir_all(&self.cache_dir)
    }

    /**
     * Get the hash of the desired file and return it as a hex string
     */
    pub fn get_file_hash(&self, path: &Path) -> Result<String, Error> {
        let mut source = File::open(&path)?;
        let mut buf: Vec<u8> = Vec::new();
        source.read_to_end(&mut buf)?;
        let mut sha1 = Sha1::new();
        sha1.update(&buf);
        let digest = sha1.digest();
        // Return the hex result of the hash
        Ok(format!("{}", digest))
    }

    /**
     * Put an image buffer in the cache
     */
    pub fn put_image_in_cache(
        &self,
        path: &Path,
        size: u32,
        image: &DynamicImage,
    ) -> Result<bool, Error> {
        let hash = self.get_file_hash(&path);
        match hash {
            Ok(sha1) => {
                let cache_path_str = format!(
                    "{}/image/{}x{}/{}/{}.{}",
                    self.cache_dir,
                    size,
                    size,
                    &sha1[..10],
                    sha1,
                    CACHED_IMAGE_EXT
                );
                let cache_dir_str =
                    format!("{}/image/{}x{}/{}", self.cache_dir, size, size, &sha1[..10]);
                match create_dir_all(&cache_dir_str) {
                    Ok(_) => {
                        let file_path = Path::new(&cache_path_str);
                        match File::create(file_path) {
                            Ok(_) => {
                                // Save the file into the cache
                                match image.save(file_path) {
                                    Ok(_) => {}
                                    Err(e) => {
                                        println!("Error: {}", e);
                                        return Err(Error::new(ErrorKind::Other, e));
                                    }
                                }
                            }
                            Err(e) => {
                                println!("Unable to create file {:?}", file_path);
                                return Err(e);
                            }
                        }
                    }
                    Err(e) => {
                        println!("Unable to create directory {:?}", &cache_dir_str);
                        return Err(e);
                    }
                }
            }
            Err(e) => {
                println!("Error: {}", e);
                return Err(e);
            }
        }
        Ok(true)
    }

    /**
     * Get an image buffer out of the cache
     */
    pub fn get_image_from_cache(&self, path: &Path, size: u32) -> Option<DynamicImage> {
        if self.use_cache {
            let hash = self.get_file_hash(&path);
            match hash {
                Ok(sha1) => {
                    // Check if the file exists in the cache
                    let cache_path_str = format!(
                        "{}/image/{}x{}/{}/{}.{}",
                        self.cache_dir,
                        size,
                        size,
                        &sha1[..10],
                        sha1,
                        CACHED_IMAGE_EXT
                    );
                    let cached_path = Path::new(&cache_path_str);
                    // Try to open, if it does, then we can read the image in
                    match File::open(&cached_path) {
                        Ok(_) => {
                            let image = image::open(&cached_path).unwrap();
                            Some(image)
                        }
                        // Don't really care here, it just means an existing cached
                        // file doesn't exist, or can't be read.
                        Err(_) => None,
                    }
                }
                Err(e) => {
                    println!("Error: {}", e);
                    None
                }
            }
        } else {
            None
        }
    }

    /**
     * Expects a slice of slices that represents lines in the file
     */
    pub fn put_matrix_in_cache(
        &self,
        path: &Path,
        size: u32,
        file_contents: &Vec<Vec<f64>>,
    ) -> Result<bool, Error> {
        let hash = self.get_file_hash(&path);
        match hash {
            Ok(sha1) => {
                let cache_path_str = format!(
                    "{}/matrix/{}x{}/{}/{}.{}",
                    self.cache_dir,
                    size,
                    size,
                    &sha1[..10],
                    sha1,
                    CACHED_MATRIX_EXT
                );
                let cache_dir_str = format!(
                    "{}/matrix/{}x{}/{}",
                    self.cache_dir,
                    size,
                    size,
                    &sha1[..10]
                );
                match create_dir_all(cache_dir_str) {
                    Ok(_) => {
                        let cached_path = Path::new(&cache_path_str);
                        // Save the file into the cache
                        match File::create(&cached_path) {
                            Ok(mut file) => {
                                let mut compressor =
                                    ZlibEncoder::new(Vec::new(), Compression::default());
                                for row in file_contents {
                                    let mut row_str =
                                        row.iter().fold(String::new(), |acc, &item| {
                                            acc + &format!("{},", item)
                                        });
                                    // remove the last comma
                                    let desire_len = row_str.len() - 1;
                                    row_str.truncate(desire_len);
                                    row_str.push_str("\n");
                                    compressor.write(&row_str.into_bytes())?;
                                }
                                let compressed_matrix = match compressor.finish() {
                                    Ok(data) => data,
                                    Err(e) => {
                                        println!("Unable to compress matrix data: {}", e);
                                        return Err(e);
                                    }
                                };
                                file.write(&compressed_matrix)?;
                                file.flush()?;
                            }
                            Err(e) => {
                                return Err(e);
                            }
                        }
                    }
                    Err(e) => println!("Error: {}", e),
                }
            }
            Err(e) => {
                println!("Error: {}", e);
                return Err(e);
            }
        }
        Ok(true)
    }

    /**
     * Get a matrix out of the cache
     */
    pub fn get_matrix_from_cache(&self, path: &Path, size: u32) -> Option<Vec<Vec<f64>>> {
        if self.use_cache {
            let hash = self.get_file_hash(&path);
            match hash {
                Ok(sha1) => {
                    // Check if the file exists in the cache
                    let cache_path_str = format!(
                        "{}/matrix/{}x{}/{}/{}.{}",
                        self.cache_dir,
                        size,
                        size,
                        &sha1[..10],
                        sha1,
                        CACHED_MATRIX_EXT
                    );
                    let cached_path = Path::new(&cache_path_str);
                    // Try to open, if it does, then we can read the image in
                    match File::open(&cached_path) {
                        Ok(file) => {
                            let mut decoder = ZlibDecoder::new(&file);
                            let mut matrix_data_str = String::new();
                            match decoder.read_to_string(&mut matrix_data_str) {
                                Ok(_) => {}
                                Err(e) => {
                                    println!("Unable to decompress matrix: {}", e);
                                    return None;
                                }
                            };
                            // convert the matrix
                            let matrix: Vec<Vec<f64>> = matrix_data_str
                                .trim()
                                .split("\n")
                                .map(|line| {
                                    line.split(",").map(|f| f64::from_str(f).unwrap()).collect()
                                })
                                .collect();

                            Some(matrix)
                        }
                        // Don't really care here, it just means an existing cached
                        // file doesn't exist, or can't be read.
                        Err(_) => None,
                    }
                }
                Err(e) => {
                    println!("Error: {}", e);
                    None
                }
            }
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use cache::Cache;

    #[test]
    fn test_get_file_hash() {
        let target = "test_images/sample_01_large.jpg";
        let target_path = Path::new(target);
        let cache: Cache = Default::default();
        let hash = cache.get_file_hash(&target_path);
        match hash {
            Ok(v) => {
                println!("Hash: {}", v);
                assert_eq!(v, String::from("4beb6f2d852b75a313863916a1803ebad13a3196"));
            }
            Err(e) => {
                println!("Error: {:?}", e);
                assert!(false);
            }
        }
    }
}