<commit_before>use Result;
use DocId;
use schema::Schema;
use schema::Document;
use schema::Term;
use schema::FieldEntry;
use core::segment_serializer::SegmentSerializer;
use core::index::SegmentInfo;
use core::index::Segment;
use analyzer::SimpleTokenizer;
use core::index::SerializableSegment;
use analyzer::StreamingIterator;
use postings::PostingsWriter;
use fastfield::U32FastFieldsWriter;
use std::clone::Clone;
use schema::Field;
use schema::FieldValue;
use schema::TextIndexingOptions;
use postings::SpecializedPostingsWriter;
use postings::{NothingRecorder, TermFrequencyRecorder, TFAndPositionRecorder};

pub struct SegmentWriter {
    max_doc: DocId,
	tokenizer: SimpleTokenizer,
	per_field_postings_writers: Vec<Box<PostingsWriter>>,
	segment_serializer: SegmentSerializer,
	fast_field_writers: U32FastFieldsWriter,
	fieldnorms_writer: U32FastFieldsWriter,
}

fn create_fieldnorms_writer(schema: &Schema) -> U32FastFieldsWriter {
	let u32_fields: Vec<Field> = schema.fields()
		.iter()
		.enumerate()
		.filter(|&(_, field_entry)| field_entry.is_indexed())
		.map(|(field_id, _)| Field(field_id as u8))
		.collect();
	U32FastFieldsWriter::new(u32_fields)
}

fn posting_from_field_entry(field_entry: &FieldEntry) -> Box<PostingsWriter> {
	match field_entry {
		&FieldEntry::Text(_, ref text_options) => {
			match text_options.get_indexing_options() {
				TextIndexingOptions::TokenizedWithFreq => {
					SpecializedPostingsWriter::<TermFrequencyRecorder>::new_boxed()
				}
				TextIndexingOptions::TokenizedWithFreqAndPosition => {
					SpecializedPostingsWriter::<TFAndPositionRecorder>::new_boxed()
				}
				_ => {
					SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
				}
			}
		}
		&FieldEntry::U32(_, _) => {
			SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
		}
	}
}

impl SegmentWriter {
	pub fn for_segment(segment: Segment, schema: &Schema) -> Result<SegmentWriter> {
		let segment_serializer = try!(SegmentSerializer::for_segment(&segment));
		let per_field_postings_writers = schema.fields()
			  .iter()
			  .map(|field_entry| {
				  posting_from_field_entry(field_entry)
			  })
			  .collect();

		Ok(SegmentWriter {
			max_doc: 0,
			per_field_postings_writers: per_field_postings_writers,
			fieldnorms_writer: create_fieldnorms_writer(schema),
			segment_serializer: segment_serializer,
			tokenizer: SimpleTokenizer::new(),
			fast_field_writers: U32FastFieldsWriter::from_schema(schema),
		})
	}

	// Write on disk all of the stuff that
	// is still on RAM :
	// - the dictionary in an fst
	// - the postings
	// - the segment info
	// The segment writer cannot be used after this, which is
	// enforced by the fact that "self" is moved.
	pub fn finalize(mut self,) -> Result<()> {
		let segment_info = self.segment_info();
		for per_field_postings_writer in self.per_field_postings_writers.iter_mut() {
			per_field_postings_writer.close();
		}
		write(&self.per_field_postings_writers,
			  &self.fast_field_writers,
			  &self.fieldnorms_writer,
			  segment_info,
			  self.segment_serializer)
	}

    pub fn add_document(&mut self, doc: &Document, schema: &Schema) -> Result<()> {
        let doc_id = self.max_doc;
        for (field, field_values) in doc.get_sorted_fields() {
			let mut num_tokens: usize = 0;
			let field_posting_writers: &mut Box<PostingsWriter> = &mut self.per_field_postings_writers[field.0 as usize];
			for field_value in field_values {
				let field_options = schema.get_field_entry(field);
				match *field_options {
					FieldEntry::Text(_, ref text_options) => {
						if text_options.get_indexing_options().is_tokenized() {
							let mut tokens = self.tokenizer.tokenize(field_value.text());
							// right now num_tokens and pos are redundant, but it should
							// change when we get proper analyzers

							let mut pos = 0u32;
							let field = field_value.field();
							loop {
								match tokens.next() {
									Some(token) => {
										let term = Term::from_field_text(field, token);
										field_posting_writers.suscribe(doc_id, pos, term);
										pos += 1;
										num_tokens += 1;
									},
									None => { break; }
								}
							}
						}
						// TODO untokenized yet indexed
					}
					FieldEntry::U32(_, ref u32_options) => {
						if u32_options.is_indexed() {
							let term = Term::from_field_u32(field_value.field(), field_value.u32_value());
							field_posting_writers.suscribe(doc_id, 0.clone(), term);
						}
					}
				}
			}
			self.fieldnorms_writer
				.get_field_writer(field)
				.map(|field_norms_writer| {
					field_norms_writer.add_val(num_tokens as u32)
				});
		}

		self.fieldnorms_writer.fill_val_up_to(doc_id);

		self.fast_field_writers.add_document(&doc);
		let stored_fieldvalues: Vec<&FieldValue> = doc
			.get_fields()
			.iter()
			.filter(|field_value| schema.get_field_entry(field_value.field()).is_stored())
			.collect();
		let doc_writer = self.segment_serializer.get_store_writer();
		try!(doc_writer.store(&stored_fieldvalues));
        self.max_doc += 1;
		Ok(())
    }

	fn segment_info(&self,) -> SegmentInfo {
		SegmentInfo {
			max_doc: self.max_doc
		}
	}
}

fn write(per_field_postings_writers: &Vec<Box<PostingsWriter>>,
		 fast_field_writers: &U32FastFieldsWriter,
		 fieldnorms_writer: &U32FastFieldsWriter,
		 segment_info: SegmentInfo,
	  	mut serializer: SegmentSerializer) -> Result<()> {
		for per_field_postings_writer in per_field_postings_writers.iter() {
			try!(per_field_postings_writer.serialize(serializer.get_postings_serializer()));
		}
		try!(fast_field_writers.serialize(serializer.get_fast_field_serializer()));
		try!(fieldnorms_writer.serialize(serializer.get_fieldnorms_serializer()));
		try!(serializer.write_segment_info(&segment_info));
		try!(serializer.close());
		Ok(())
}

impl SerializableSegment for SegmentWriter {
	fn write(&self, serializer: SegmentSerializer) -> Result<()> {
		write(&self.per_field_postings_writers,
		      &self.fast_field_writers,
			  &self.fieldnorms_writer,
			  self.segment_info(),
		      serializer)
	}
}
<commit_msg>indexing untokenized fields<commit_after>use Result;
use DocId;
use schema::Schema;
use schema::Document;
use schema::Term;
use schema::FieldEntry;
use core::segment_serializer::SegmentSerializer;
use core::index::SegmentInfo;
use core::index::Segment;
use analyzer::SimpleTokenizer;
use core::index::SerializableSegment;
use analyzer::StreamingIterator;
use postings::PostingsWriter;
use fastfield::U32FastFieldsWriter;
use std::clone::Clone;
use schema::Field;
use schema::FieldValue;
use schema::TextIndexingOptions;
use postings::SpecializedPostingsWriter;
use postings::{NothingRecorder, TermFrequencyRecorder, TFAndPositionRecorder};

pub struct SegmentWriter {
    max_doc: DocId,
	tokenizer: SimpleTokenizer,
	per_field_postings_writers: Vec<Box<PostingsWriter>>,
	segment_serializer: SegmentSerializer,
	fast_field_writers: U32FastFieldsWriter,
	fieldnorms_writer: U32FastFieldsWriter,
}

fn create_fieldnorms_writer(schema: &Schema) -> U32FastFieldsWriter {
	let u32_fields: Vec<Field> = schema.fields()
		.iter()
		.enumerate()
		.filter(|&(_, field_entry)| field_entry.is_indexed())
		.map(|(field_id, _)| Field(field_id as u8))
		.collect();
	U32FastFieldsWriter::new(u32_fields)
}

fn posting_from_field_entry(field_entry: &FieldEntry) -> Box<PostingsWriter> {
	match field_entry {
		&FieldEntry::Text(_, ref text_options) => {
			match text_options.get_indexing_options() {
				TextIndexingOptions::TokenizedWithFreq => {
					SpecializedPostingsWriter::<TermFrequencyRecorder>::new_boxed()
				}
				TextIndexingOptions::TokenizedWithFreqAndPosition => {
					SpecializedPostingsWriter::<TFAndPositionRecorder>::new_boxed()
				}
				_ => {
					SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
				}
			}
		}
		&FieldEntry::U32(_, _) => {
			SpecializedPostingsWriter::<NothingRecorder>::new_boxed()
		}
	}
}

impl SegmentWriter {
	pub fn for_segment(segment: Segment, schema: &Schema) -> Result<SegmentWriter> {
		let segment_serializer = try!(SegmentSerializer::for_segment(&segment));
		let per_field_postings_writers = schema.fields()
			  .iter()
			  .map(|field_entry| {
				  posting_from_field_entry(field_entry)
			  })
			  .collect();

		Ok(SegmentWriter {
			max_doc: 0,
			per_field_postings_writers: per_field_postings_writers,
			fieldnorms_writer: create_fieldnorms_writer(schema),
			segment_serializer: segment_serializer,
			tokenizer: SimpleTokenizer::new(),
			fast_field_writers: U32FastFieldsWriter::from_schema(schema),
		})
	}

	// Write on disk all of the stuff that
	// is still on RAM :
	// - the dictionary in an fst
	// - the postings
	// - the segment info
	// The segment writer cannot be used after this, which is
	// enforced by the fact that "self" is moved.
	pub fn finalize(mut self,) -> Result<()> {
		let segment_info = self.segment_info();
		for per_field_postings_writer in self.per_field_postings_writers.iter_mut() {
			per_field_postings_writer.close();
		}
		write(&self.per_field_postings_writers,
			  &self.fast_field_writers,
			  &self.fieldnorms_writer,
			  segment_info,
			  self.segment_serializer)
	}

    pub fn add_document(&mut self, doc: &Document, schema: &Schema) -> Result<()> {
        let doc_id = self.max_doc;
        for (field, field_values) in doc.get_sorted_fields() {
			let mut num_tokens: usize = 0;
			let field_posting_writers: &mut Box<PostingsWriter> = &mut self.per_field_postings_writers[field.0 as usize];
			for field_value in field_values {
				let field_options = schema.get_field_entry(field);
				match *field_options {
					FieldEntry::Text(_, ref text_options) => {
						if text_options.get_indexing_options().is_tokenized() {
							let mut tokens = self.tokenizer.tokenize(field_value.text());
							// right now num_tokens and pos are redundant, but it should
							// change when we get proper analyzers

							let mut pos = 0u32;
							let field = field_value.field();
							loop {
								match tokens.next() {
									Some(token) => {
										let term = Term::from_field_text(field, token);
										field_posting_writers.suscribe(doc_id, pos, term);
										pos += 1;
										num_tokens += 1;
									},
									None => { break; }
								}
							}
						}
						else {
							let term = Term::from_field_text(field, field_value.text());
							field_posting_writers.suscribe(doc_id, 0, term);
						}
					}
					FieldEntry::U32(_, ref u32_options) => {
						if u32_options.is_indexed() {
							let term = Term::from_field_u32(field_value.field(), field_value.u32_value());
							field_posting_writers.suscribe(doc_id, 0, term);
						}
					}
				}
			}
			self.fieldnorms_writer
				.get_field_writer(field)
				.map(|field_norms_writer| {
					field_norms_writer.add_val(num_tokens as u32)
				});
		}

		self.fieldnorms_writer.fill_val_up_to(doc_id);

		self.fast_field_writers.add_document(&doc);
		let stored_fieldvalues: Vec<&FieldValue> = doc
			.get_fields()
			.iter()
			.filter(|field_value| schema.get_field_entry(field_value.field()).is_stored())
			.collect();
		let doc_writer = self.segment_serializer.get_store_writer();
		try!(doc_writer.store(&stored_fieldvalues));
        self.max_doc += 1;
		Ok(())
    }

	fn segment_info(&self,) -> SegmentInfo {
		SegmentInfo {
			max_doc: self.max_doc
		}
	}
}

fn write(per_field_postings_writers: &Vec<Box<PostingsWriter>>,
		 fast_field_writers: &U32FastFieldsWriter,
		 fieldnorms_writer: &U32FastFieldsWriter,
		 segment_info: SegmentInfo,
	  	mut serializer: SegmentSerializer) -> Result<()> {
		for per_field_postings_writer in per_field_postings_writers.iter() {
			try!(per_field_postings_writer.serialize(serializer.get_postings_serializer()));
		}
		try!(fast_field_writers.serialize(serializer.get_fast_field_serializer()));
		try!(fieldnorms_writer.serialize(serializer.get_fieldnorms_serializer()));
		try!(serializer.write_segment_info(&segment_info));
		try!(serializer.close());
		Ok(())
}

impl SerializableSegment for SegmentWriter {
	fn write(&self, serializer: SegmentSerializer) -> Result<()> {
		write(&self.per_field_postings_writers,
		      &self.fast_field_writers,
			  &self.fieldnorms_writer,
			  self.segment_info(),
		      serializer)
	}
}
<|endoftext|>


<commit_before>// Take a look at the license at the top of the repository in the LICENSE file.

use crate::sys::utils::get_all_data;
use crate::{utils, DiskExt, DiskType};

use libc::statvfs;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::mem;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

macro_rules! cast {
    ($x:expr) => {
        u64::from($x)
    };
}

#[doc = include_str!("../../md_doc/disk.md")]
#[derive(PartialEq, Eq)]
pub struct Disk {
    type_: DiskType,
    device_name: OsString,
    file_system: Vec<u8>,
    mount_point: PathBuf,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.device_name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        unsafe {
            let mut stat: statvfs = mem::zeroed();
            let mount_point_cpath = utils::to_cpath(&self.mount_point);
            if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
                let tmp = cast!(stat.f_bsize).saturating_mul(cast!(stat.f_bavail));
                self.available_space = cast!(tmp);
                true
            } else {
                false
            }
        }
    }
}

fn new_disk(
    device_name: &OsStr,
    mount_point: &Path,
    file_system: &[u8],
    removable_entries: &[PathBuf],
) -> Option<Disk> {
    let mount_point_cpath = utils::to_cpath(mount_point);
    let type_ = find_type_for_device_name(device_name);
    let mut total = 0;
    let mut available = 0;
    unsafe {
        let mut stat: statvfs = mem::zeroed();
        if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
            let bsize = cast!(stat.f_bsize);
            let blocks = cast!(stat.f_blocks);
            let bavail = cast!(stat.f_bavail);
            total = bsize.saturating_mul(blocks);
            available = bsize.saturating_mul(bavail);
        }
        if total == 0 {
            return None;
        }
        let mount_point = mount_point.to_owned();
        let is_removable = removable_entries
            .iter()
            .any(|e| e.as_os_str() == device_name);
        Some(Disk {
            type_,
            device_name: device_name.to_owned(),
            file_system: file_system.to_owned(),
            mount_point,
            total_space: cast!(total),
            available_space: cast!(available),
            is_removable,
        })
    }
}

#[allow(clippy::manual_range_contains)]
fn find_type_for_device_name(device_name: &OsStr) -> DiskType {
    // The format of devices are as follows:
    //  - device_name is symbolic link in the case of /dev/mapper/
    //     and /dev/root, and the target is corresponding device under
    //     /sys/block/
    //  - In the case of /dev/sd, the format is /dev/sd[a-z][1-9],
    //     corresponding to /sys/block/sd[a-z]
    //  - In the case of /dev/nvme, the format is /dev/nvme[0-9]n[0-9]p[0-9],
    //     corresponding to /sys/block/nvme[0-9]n[0-9]
    //  - In the case of /dev/mmcblk, the format is /dev/mmcblk[0-9]p[0-9],
    //     corresponding to /sys/block/mmcblk[0-9]
    let device_name_path = device_name.to_str().unwrap_or_default();
    let real_path = fs::canonicalize(device_name).unwrap_or_else(|_| PathBuf::from(device_name));
    let mut real_path = real_path.to_str().unwrap_or_default();
    if device_name_path.starts_with("/dev/mapper/") {
        // Recursively solve, for example /dev/dm-0
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/sd") || device_name_path.starts_with("/dev/vd") {
        // Turn "sda1" into "sda" or "vda1" into "vda"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
    } else if device_name_path.starts_with("/dev/nvme") {
        // Turn "nvme0n1p1" into "nvme0n1"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
        real_path = real_path.trim_end_matches(|c| c == 'p');
    } else if device_name_path.starts_with("/dev/root") {
        // Recursively solve, for example /dev/mmcblk0p1
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/mmcblk") {
        // Turn "mmcblk0p1" into "mmcblk0"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
        real_path = real_path.trim_end_matches(|c| c == 'p');
    } else {
        // Default case: remove /dev/ and expects the name presents under /sys/block/
        // For example, /dev/dm-0 to dm-0
        real_path = real_path.trim_start_matches("/dev/");
    }

    let trimmed: &OsStr = OsStrExt::from_bytes(real_path.as_bytes());

    let path = Path::new("/sys/block/")
        .to_owned()
        .join(trimmed)
        .join("queue/rotational");
    // Normally, this file only contains '0' or '1' but just in case, we get 8 bytes...
    match get_all_data(path, 8)
        .unwrap_or_default()
        .trim()
        .parse()
        .ok()
    {
        // The disk is marked as rotational so it's a HDD.
        Some(1) => DiskType::HDD,
        // The disk is marked as non-rotational so it's very likely a SSD.
        Some(0) => DiskType::SSD,
        // Normally it shouldn't happen but welcome to the wonderful world of IT! :D
        Some(x) => DiskType::Unknown(x),
        // The information isn't available...
        None => DiskType::Unknown(-1),
    }
}

fn get_all_disks_inner(content: &str) -> Vec<Disk> {
    // The goal of this array is to list all removable devices (the ones whose name starts with
    // "usb-"). Then we check if
    let removable_entries = match fs::read_dir("/dev/disk/by-id/") {
        Ok(r) => r
            .filter_map(|res| Some(res.ok()?.path()))
            .filter_map(|e| {
                if e.file_name()
                    .and_then(|x| Some(x.to_str()?.starts_with("usb-")))
                    .unwrap_or_default()
                {
                    e.canonicalize().ok()
                } else {
                    None
                }
            })
            .collect::<Vec<PathBuf>>(),
        _ => Vec::new(),
    };

    content
        .lines()
        .map(|line| {
            let line = line.trim_start();
            // mounts format
            // http://man7.org/linux/man-pages/man5/fstab.5.html
            // fs_spec<tab>fs_file<tab>fs_vfstype<tab>other fields
            let mut fields = line.split_whitespace();
            let fs_spec = fields.next().unwrap_or("");
            let fs_file = fields
                .next()
                .unwrap_or("")
                .replace("\\134", "\\")
                .replace("\\040", " ")
                .replace("\\011", "\t")
                .replace("\\012", "\n");
            let fs_vfstype = fields.next().unwrap_or("");
            (fs_spec, fs_file, fs_vfstype)
        })
        .filter(|(fs_spec, fs_file, fs_vfstype)| {
            // Check if fs_vfstype is one of our 'ignored' file systems.
            let filtered = matches!(
                *fs_vfstype,
                "rootfs" | // https://www.kernel.org/doc/Documentation/filesystems/ramfs-rootfs-initramfs.txt
                "sysfs" | // pseudo file system for kernel objects
                "proc" |  // another pseudo file system
                "tmpfs" |
                "devtmpfs" |
                "cgroup" |
                "cgroup2" |
                "pstore" | // https://www.kernel.org/doc/Documentation/ABI/testing/pstore
                "squashfs" | // squashfs is a compressed read-only file system (for snaps)
                "rpc_pipefs" | // The pipefs pseudo file system service
                "iso9660" // optical media
            );

            !(filtered ||
               fs_file.starts_with("/sys") || // check if fs_file is an 'ignored' mount point
               fs_file.starts_with("/proc") ||
               (fs_file.starts_with("/run") && !fs_file.starts_with("/run/media")) ||
               fs_spec.starts_with("sunrpc"))
        })
        .filter_map(|(fs_spec, fs_file, fs_vfstype)| {
            new_disk(
                fs_spec.as_ref(),
                Path::new(&fs_file),
                fs_vfstype.as_bytes(),
                &removable_entries,
            )
        })
        .collect()
}

pub(crate) fn get_all_disks() -> Vec<Disk> {
    get_all_disks_inner(&get_all_data("/proc/mounts", 16_385).unwrap_or_default())
}

// #[test]
// fn check_all_disks() {
//     let disks = get_all_disks_inner(
//         r#"tmpfs /proc tmpfs rw,seclabel,relatime 0 0
// proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0
// systemd-1 /proc/sys/fs/binfmt_misc autofs rw,relatime,fd=29,pgrp=1,timeout=0,minproto=5,maxproto=5,direct,pipe_ino=17771 0 0
// tmpfs /sys tmpfs rw,seclabel,relatime 0 0
// sysfs /sys sysfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// securityfs /sys/kernel/security securityfs rw,nosuid,nodev,noexec,relatime 0 0
// cgroup2 /sys/fs/cgroup cgroup2 rw,seclabel,nosuid,nodev,noexec,relatime,nsdelegate 0 0
// pstore /sys/fs/pstore pstore rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// none /sys/fs/bpf bpf rw,nosuid,nodev,noexec,relatime,mode=700 0 0
// configfs /sys/kernel/config configfs rw,nosuid,nodev,noexec,relatime 0 0
// selinuxfs /sys/fs/selinux selinuxfs rw,relatime 0 0
// debugfs /sys/kernel/debug debugfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// tmpfs /dev/shm tmpfs rw,seclabel,relatime 0 0
// devpts /dev/pts devpts rw,seclabel,relatime,gid=5,mode=620,ptmxmode=666 0 0
// tmpfs /sys/fs/selinux tmpfs rw,seclabel,relatime 0 0
// /dev/vda2 /proc/filesystems xfs rw,seclabel,relatime,attr2,inode64,logbufs=8,logbsize=32k,noquota 0 0
// "#,
//     );
//     assert_eq!(disks.len(), 1);
//     assert_eq!(
//         disks[0],
//         Disk {
//             type_: DiskType::Unknown(-1),
//             name: OsString::from("devpts"),
//             file_system: vec![100, 101, 118, 112, 116, 115],
//             mount_point: PathBuf::from("/dev/pts"),
//             total_space: 0,
//             available_space: 0,
//         }
//     );
// }
<commit_msg>fix: got nvme0n when passing device name 'nvme0n1'<commit_after>// Take a look at the license at the top of the repository in the LICENSE file.

use crate::sys::utils::get_all_data;
use crate::{utils, DiskExt, DiskType};

use libc::statvfs;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::mem;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

macro_rules! cast {
    ($x:expr) => {
        u64::from($x)
    };
}

#[doc = include_str!("../../md_doc/disk.md")]
#[derive(PartialEq, Eq)]
pub struct Disk {
    type_: DiskType,
    device_name: OsString,
    file_system: Vec<u8>,
    mount_point: PathBuf,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.device_name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        unsafe {
            let mut stat: statvfs = mem::zeroed();
            let mount_point_cpath = utils::to_cpath(&self.mount_point);
            if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
                let tmp = cast!(stat.f_bsize).saturating_mul(cast!(stat.f_bavail));
                self.available_space = cast!(tmp);
                true
            } else {
                false
            }
        }
    }
}

fn new_disk(
    device_name: &OsStr,
    mount_point: &Path,
    file_system: &[u8],
    removable_entries: &[PathBuf],
) -> Option<Disk> {
    let mount_point_cpath = utils::to_cpath(mount_point);
    let type_ = find_type_for_device_name(device_name);
    let mut total = 0;
    let mut available = 0;
    unsafe {
        let mut stat: statvfs = mem::zeroed();
        if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
            let bsize = cast!(stat.f_bsize);
            let blocks = cast!(stat.f_blocks);
            let bavail = cast!(stat.f_bavail);
            total = bsize.saturating_mul(blocks);
            available = bsize.saturating_mul(bavail);
        }
        if total == 0 {
            return None;
        }
        let mount_point = mount_point.to_owned();
        let is_removable = removable_entries
            .iter()
            .any(|e| e.as_os_str() == device_name);
        Some(Disk {
            type_,
            device_name: device_name.to_owned(),
            file_system: file_system.to_owned(),
            mount_point,
            total_space: cast!(total),
            available_space: cast!(available),
            is_removable,
        })
    }
}

#[allow(clippy::manual_range_contains)]
fn find_type_for_device_name(device_name: &OsStr) -> DiskType {
    // The format of devices are as follows:
    //  - device_name is symbolic link in the case of /dev/mapper/
    //     and /dev/root, and the target is corresponding device under
    //     /sys/block/
    //  - In the case of /dev/sd, the format is /dev/sd[a-z][1-9],
    //     corresponding to /sys/block/sd[a-z]
    //  - In the case of /dev/nvme, the format is /dev/nvme[0-9]n[0-9]p[0-9],
    //     corresponding to /sys/block/nvme[0-9]n[0-9]
    //  - In the case of /dev/mmcblk, the format is /dev/mmcblk[0-9]p[0-9],
    //     corresponding to /sys/block/mmcblk[0-9]
    let device_name_path = device_name.to_str().unwrap_or_default();
    let real_path = fs::canonicalize(device_name).unwrap_or_else(|_| PathBuf::from(device_name));
    let mut real_path = real_path.to_str().unwrap_or_default();
    if device_name_path.starts_with("/dev/mapper/") {
        // Recursively solve, for example /dev/dm-0
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/sd") || device_name_path.starts_with("/dev/vd") {
        // Turn "sda1" into "sda" or "vda1" into "vda"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
    } else if device_name_path.starts_with("/dev/nvme") {
        // Turn "nvme0n1p1" into "nvme0n1"
        real_path = match real_path.find('p') {
            Some(idx) => &real_path["/dev/".len()..idx],
            None => &real_path["/dev/".len()..],
        };
    } else if device_name_path.starts_with("/dev/root") {
        // Recursively solve, for example /dev/mmcblk0p1
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/mmcblk") {
        // Turn "mmcblk0p1" into "mmcblk0"
        real_path = match real_path.find('p') {
            Some(idx) => &real_path["/dev/".len()..idx],
            None => &real_path["/dev/".len()..],
        };
    } else {
        // Default case: remove /dev/ and expects the name presents under /sys/block/
        // For example, /dev/dm-0 to dm-0
        real_path = real_path.trim_start_matches("/dev/");
    }

    let trimmed: &OsStr = OsStrExt::from_bytes(real_path.as_bytes());

    let path = Path::new("/sys/block/")
        .to_owned()
        .join(trimmed)
        .join("queue/rotational");
    // Normally, this file only contains '0' or '1' but just in case, we get 8 bytes...
    match get_all_data(path, 8)
        .unwrap_or_default()
        .trim()
        .parse()
        .ok()
    {
        // The disk is marked as rotational so it's a HDD.
        Some(1) => DiskType::HDD,
        // The disk is marked as non-rotational so it's very likely a SSD.
        Some(0) => DiskType::SSD,
        // Normally it shouldn't happen but welcome to the wonderful world of IT! :D
        Some(x) => DiskType::Unknown(x),
        // The information isn't available...
        None => DiskType::Unknown(-1),
    }
}

fn get_all_disks_inner(content: &str) -> Vec<Disk> {
    // The goal of this array is to list all removable devices (the ones whose name starts with
    // "usb-"). Then we check if
    let removable_entries = match fs::read_dir("/dev/disk/by-id/") {
        Ok(r) => r
            .filter_map(|res| Some(res.ok()?.path()))
            .filter_map(|e| {
                if e.file_name()
                    .and_then(|x| Some(x.to_str()?.starts_with("usb-")))
                    .unwrap_or_default()
                {
                    e.canonicalize().ok()
                } else {
                    None
                }
            })
            .collect::<Vec<PathBuf>>(),
        _ => Vec::new(),
    };

    content
        .lines()
        .map(|line| {
            let line = line.trim_start();
            // mounts format
            // http://man7.org/linux/man-pages/man5/fstab.5.html
            // fs_spec<tab>fs_file<tab>fs_vfstype<tab>other fields
            let mut fields = line.split_whitespace();
            let fs_spec = fields.next().unwrap_or("");
            let fs_file = fields
                .next()
                .unwrap_or("")
                .replace("\\134", "\\")
                .replace("\\040", " ")
                .replace("\\011", "\t")
                .replace("\\012", "\n");
            let fs_vfstype = fields.next().unwrap_or("");
            (fs_spec, fs_file, fs_vfstype)
        })
        .filter(|(fs_spec, fs_file, fs_vfstype)| {
            // Check if fs_vfstype is one of our 'ignored' file systems.
            let filtered = matches!(
                *fs_vfstype,
                "rootfs" | // https://www.kernel.org/doc/Documentation/filesystems/ramfs-rootfs-initramfs.txt
                "sysfs" | // pseudo file system for kernel objects
                "proc" |  // another pseudo file system
                "tmpfs" |
                "devtmpfs" |
                "cgroup" |
                "cgroup2" |
                "pstore" | // https://www.kernel.org/doc/Documentation/ABI/testing/pstore
                "squashfs" | // squashfs is a compressed read-only file system (for snaps)
                "rpc_pipefs" | // The pipefs pseudo file system service
                "iso9660" // optical media
            );

            !(filtered ||
               fs_file.starts_with("/sys") || // check if fs_file is an 'ignored' mount point
               fs_file.starts_with("/proc") ||
               (fs_file.starts_with("/run") && !fs_file.starts_with("/run/media")) ||
               fs_spec.starts_with("sunrpc"))
        })
        .filter_map(|(fs_spec, fs_file, fs_vfstype)| {
            new_disk(
                fs_spec.as_ref(),
                Path::new(&fs_file),
                fs_vfstype.as_bytes(),
                &removable_entries,
            )
        })
        .collect()
}

pub(crate) fn get_all_disks() -> Vec<Disk> {
    get_all_disks_inner(&get_all_data("/proc/mounts", 16_385).unwrap_or_default())
}

// #[test]
// fn check_all_disks() {
//     let disks = get_all_disks_inner(
//         r#"tmpfs /proc tmpfs rw,seclabel,relatime 0 0
// proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0
// systemd-1 /proc/sys/fs/binfmt_misc autofs rw,relatime,fd=29,pgrp=1,timeout=0,minproto=5,maxproto=5,direct,pipe_ino=17771 0 0
// tmpfs /sys tmpfs rw,seclabel,relatime 0 0
// sysfs /sys sysfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// securityfs /sys/kernel/security securityfs rw,nosuid,nodev,noexec,relatime 0 0
// cgroup2 /sys/fs/cgroup cgroup2 rw,seclabel,nosuid,nodev,noexec,relatime,nsdelegate 0 0
// pstore /sys/fs/pstore pstore rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// none /sys/fs/bpf bpf rw,nosuid,nodev,noexec,relatime,mode=700 0 0
// configfs /sys/kernel/config configfs rw,nosuid,nodev,noexec,relatime 0 0
// selinuxfs /sys/fs/selinux selinuxfs rw,relatime 0 0
// debugfs /sys/kernel/debug debugfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// tmpfs /dev/shm tmpfs rw,seclabel,relatime 0 0
// devpts /dev/pts devpts rw,seclabel,relatime,gid=5,mode=620,ptmxmode=666 0 0
// tmpfs /sys/fs/selinux tmpfs rw,seclabel,relatime 0 0
// /dev/vda2 /proc/filesystems xfs rw,seclabel,relatime,attr2,inode64,logbufs=8,logbsize=32k,noquota 0 0
// "#,
//     );
//     assert_eq!(disks.len(), 1);
//     assert_eq!(
//         disks[0],
//         Disk {
//             type_: DiskType::Unknown(-1),
//             name: OsString::from("devpts"),
//             file_system: vec![100, 101, 118, 112, 116, 115],
//             mount_point: PathBuf::from("/dev/pts"),
//             total_space: 0,
//             available_space: 0,
//         }
//     );
// }
<|endoftext|>


<commit_before>//! The internals of `VecArena`.
//!
//! `Bitmap` tracks occupied and vacant slots in `VecArena` and provides useful methods for e.g.
//! iterating over all occupied slots, or reserving additional slots.

use std::cmp;
use std::mem;
use std::ptr;

/// Returns the number of bits in one `usize` integer.
#[inline]
fn bits() -> usize {
    mem::size_of::<usize>() * 8
}

/// Returns the number of `usize` integers required to store `len` bits.
#[inline]
fn blocks_for(len: usize) -> usize {
    // Divide `len` by `bits()` and round up.
    len / bits() + ((len % bits() > 0) as usize)
}

/// Given a valid `block` index and `offset` within it, returns the index of that slot as within
/// the whole bitmap.
#[inline]
fn slot_index(block: usize, offset: usize) -> usize {
    block * bits() + offset
}

/// Given a slot `index` in the bitmap, returns it's block index and offset within the block.
#[inline]
fn block_and_offset(index: usize) -> (usize, usize) {
    (index / bits(), index % bits())
}

/// Keeps track of occupied and vacant slots in `VecArena`.
///
/// It's implemented as an array of blocks, where every block tracks only a small contiguous chunk
/// of slots. More precisely: as many slots as there are bits in one `usize`.
///
/// All blocks which are not fully occupied, except the last one, are linked together to form a
/// doubly-linked list. This list allows finding a vacant slot to acquire in `O(1)` and releasing
/// an occupied slot in `O(1)`, assuming the bitmap doesn't grow to accommodate more slots.
///
/// A block consists of:
///
/// * a bit mask (one `usize`), in which zeros are for vacant slots and ones for occupied slots
/// * index of the successor block in the linked list
/// * index of the predecessor block in the linked list
///
/// The last block is tricky to handle because it might not have the same number of slots as other
/// blocks, so it gets special treatment in the implementation.
pub struct Bitmap {
    /// Storage for the following sequences, in this order:
    ///
    /// * bit masks
    /// * indices to successor nodes
    /// * indices to predecessor nodes
    ///
    /// All three sequences are stored in this single contiguous array.
    ///
    /// The most common and costly operation during the lifetime of a bitmap is testing
    /// whether a slot is occupied. Storing bit masks close together improves cache
    /// performance.
    data: *mut usize,

    /// Number of reserved slots.
    len: usize,

    /// Number of occupied slots.
    occupied: usize,

    /// Index of the first block in the linked list, or `!0` if the list is empty.
    head: usize,
}

impl Bitmap {
    /// Constructs a new `Bitmap` with zero slots.
    #[inline]
    pub fn new() -> Self {
        let data = {
            let mut vec = Vec::with_capacity(0);
            let ptr = vec.as_mut_ptr();
            mem::forget(vec);
            ptr
        };
        Bitmap {
            data: data,
            len: 0,
            occupied: 0,
            head: !0,
        }
    }

    /// Returns the number of reserved slots in the bitmap.
    #[inline]
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns the number of occupied slots in the bitmap.
    #[inline]
    pub fn occupied(&self) -> usize {
        self.occupied
    }

    /// Finds a vacant slot, marks as occupied, and returns its index.
    ///
    /// # Panics
    ///
    /// Panics if all slots are occupied.
    pub fn acquire(&mut self) -> usize {
        assert!(self.occupied < self.len, "no vacant slots to acquire, len = {}", self.len);

        let block = if self.head == !0 {
            // The list is empty, so try the last block.
            blocks_for(self.len) - 1
        } else {
            // The list has a head. Take a vacant slot from the head block.
            self.head
        };

        let index = unsafe {
            // Find the rightmost zero bit in the mask. Taking the rightmost zero is always ok,
            // even if this is the last block.
            let offset = (!*self.mask(block)).trailing_zeros() as usize;
            debug_assert!(offset < bits());

            slot_index(block, offset)
        };
        self.acquire_at(index);
        index
    }

    /// Marks the vacant slot at `index` as occupied.
    ///
    /// # Panics
    ///
    /// Panics if the slot is already occupied.
    pub fn acquire_at(&mut self, index: usize) {
        assert!(index < self.len, "`index` out of bounds");

        let (block, offset) = block_and_offset(index);
        unsafe {
            // Mark the slot as occupied in the block's bit mask.
            assert!(*self.mask(block) >> offset & 1 == 0, "occupied slot at `index`");
            *self.mask(block) ^= 1 << offset;

            // If the block has just become fully occupied, remove it from the list.
            if block == self.head && *self.mask(block) == !0 {
                let succ = *self.successor(block);
                self.link_blocks(!0, succ);
                self.head = succ;
            }

            self.occupied += 1;
        }
    }

    /// Releases the occupied slot at `index`.
    ///
    /// # Panics
    ///
    /// Panics if the `index` is out of bounds or the slot is vacant.
    pub fn release(&mut self, index: usize) {
        assert!(index < self.len, "`index` out of bounds");

        let (block, offset) = block_and_offset(index);
        unsafe {
            // Make sure we're not releasing a vacant slot.
            assert!(*self.mask(block) >> offset & 1 == 1, "releasing a vacant slot");

            self.occupied -= 1;

            // If the block is fully occupied, insert it back into the list.
            if *self.mask(block) == !0 {
                let head = self.head;
                self.link_blocks(!0, block);
                self.link_blocks(block, head);
                self.head = block;
            }

            // Mark the slot as vacant in the block's bit mask.
            *self.mask(block) ^= 1 << offset;
        }
    }

    /// Given the required minimal additional number of slots to reserve, returns the number of
    /// slots to reserve in order to keep amortized time complexity.
    fn amortized_reserve(&self, additional: usize) -> usize {
        let len = self.len();
        let required_len = len.checked_add(additional).expect("len overflow");
        let double_len = len.checked_mul(2).expect("len overflow");
        cmp::max(required_len, double_len) - len
    }

    /// Reserves at least `additional` more slots for new elements to be inserted. The arena may
    /// reserve more space to avoid frequent reallocations.
    ///
    /// # Panics
    ///
    /// Panics if the new number of slots overflows `usize`.
    pub fn reserve(&mut self, additional: usize) {
        let amortized = self.amortized_reserve(additional);
        self.reserve_exact(amortized);
    }

    /// Reserves exactly `additional` more slots for new elements to be inserted.
    ///
    /// Note that the allocator may give the bitmap more space than it requests.
    ///
    /// # Panics
    ///
    /// Panics if the new number of slots overflows `usize`.
    pub fn reserve_exact(&mut self, additional: usize) {
        let old_blocks = blocks_for(self.len);
        self.len = self.len.checked_add(additional).expect("len overflow");
        let new_blocks = blocks_for(self.len);

        if new_blocks == old_blocks {
            // We can simply return.
            // The higher unused bits of the last block are always zero anyway.
            return;
        }

        unsafe {
            // If the last block had some unused or vacant slots, push it into the linked list.
            if old_blocks > 0 && *self.mask(old_blocks - 1) != !0 {
                let head = self.head;
                let last = old_blocks - 1;
                self.link_blocks(last, head);
                self.link_blocks(!0, last);
                self.head = last;
            }

            // Allocate a new array.
            let new_data = {
                // Every block contains three `usize` integers, so we need `3 * new_blocks` of
                // space.
                let mut vec = Vec::with_capacity(3 * new_blocks);
                let ptr = vec.as_mut_ptr();
                mem::forget(vec);
                ptr
            };

            // Copy the three old subarrays (bit masks, successor indices, predecessor indices)
            // into the new array.
            for i in 0..3 {
                ptr::copy_nonoverlapping(
                    self.data.offset((old_blocks * i) as isize),
                    new_data.offset((new_blocks * i) as isize),
                    old_blocks);
            }

            // Deallocate the old array.
            Vec::from_raw_parts(self.data, 0, 3 * old_blocks);

            // Set the new bit masks to zero.
            ptr::write_bytes(new_data.offset(old_blocks as isize), 0, new_blocks - old_blocks);

            // Set the new data now because we're about to push blocks into the linked list.
            self.data = new_data;

            // If there are at least two new blocks, that means there is at least one new block
            // that must be pushed into the linked list.
            if old_blocks + 2 <= new_blocks {
                let head = self.head;
                self.link_blocks(new_blocks - 2, head);
                self.link_blocks(!0, old_blocks);
                self.head = old_blocks;

                for block in old_blocks .. new_blocks - 2 {
                    self.link_blocks(block, block + 1);
                }
            }
        }
    }

    /// Returns `true` if the slot at `index` is occupied.
    #[inline]
    pub fn is_occupied(&self, index: usize) -> bool {
        assert!(index < self.len, "`index` out of bounds");

        let (block, offset) = block_and_offset(index);
        unsafe {
            *self.mask(block) >> offset & 1 != 0
        }
    }

    /// Returns an iterator over occupied slots.
    #[inline]
    pub fn iter(&self) -> Iter {
        Iter {
            bitmap: self,
            // block: 0,
            // offset: 0,
            index: 0,
        }
    }

    #[cfg(debug_assertions)]
    pub fn check_invariants(&self) {
        // TODO
    }

    /// Links together blocks `a` and `b` so that `a` comes before `b` in the linked list.
    #[inline]
    unsafe fn link_blocks(&mut self, a: usize, b: usize) {
        if a != !0 { *self.successor(a) = b; }
        if b != !0 { *self.predecessor(b) = a; }
    }

    /// Returns the pointer to the bit mask of `block`.
    #[inline]
    unsafe fn mask(&self, block: usize) -> *mut usize {
        self.data.offset(block as isize)
    }

    /// Returns the pointer to the index of the successor of `block`.
    #[inline]
    unsafe fn successor(&self, block: usize) -> *mut usize {
        self.data.offset((blocks_for(self.len) + block) as isize)
    }

    /// Returns the pointer to the index of the predecessor of `block`.
    #[inline]
    unsafe fn predecessor(&self, block: usize) -> *mut usize {
        self.data.offset((2 * blocks_for(self.len) + block) as isize)
    }

    /// Returns the next occupied slot starting from (and including) the specified `index`.
    pub fn next_occupied(&self, index: usize) -> Option<usize> {
        let num_blocks = blocks_for(self.len);
        let (mut block, mut offset) = block_and_offset(index);

        while block < num_blocks {
            let mask = unsafe { *self.mask(block) };

            if offset == bits() || mask == 0 {
                // Nothing left in the current block or it's empty. Go to the next block.
                block += 1;
                offset = 0;
            } else {
                // Find the next occupied slot in this block.
                while offset < bits() {
                    if mask >> offset & 1 == 1 {
                        return Some(slot_index(block, offset));
                    }
                    offset += 1;
                }
            }
        }
        None
    }
}

impl Drop for Bitmap {
    fn drop(&mut self) {
        unsafe {
            Vec::from_raw_parts(self.data, 0, 2 * blocks_for(self.len));
        }
    }
}

/// Iterates over all occupied slots.
pub struct Iter<'a> {
    /// The bitmap to iterate over.
    bitmap: &'a Bitmap,

    /// Index of the current slot.
    index: usize,
}

impl<'a> Iterator for Iter<'a> {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            index
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let low = 0;
        let high = self.bitmap.len() - self.index;
        (low, Some(high))
    }

    #[inline]
    fn count(self) -> usize {
        self.bitmap.occupied
    }
}
<commit_msg>Implement previous_occupied<commit_after>//! The internals of `VecArena`.
//!
//! `Bitmap` tracks occupied and vacant slots in `VecArena` and provides useful methods for e.g.
//! iterating over all occupied slots, or reserving additional slots.

use std::cmp;
use std::mem;
use std::ptr;

/// Returns the number of bits in one `usize` integer.
#[inline]
fn bits() -> usize {
    mem::size_of::<usize>() * 8
}

/// Returns the number of `usize` integers required to store `len` bits.
#[inline]
fn blocks_for(len: usize) -> usize {
    // Divide `len` by `bits()` and round up.
    len / bits() + ((len % bits() > 0) as usize)
}

/// Given a valid `block` index and `offset` within it, returns the index of that slot as within
/// the whole bitmap.
#[inline]
fn slot_index(block: usize, offset: usize) -> usize {
    block * bits() + offset
}

/// Given a slot `index` in the bitmap, returns it's block index and offset within the block.
#[inline]
fn block_and_offset(index: usize) -> (usize, usize) {
    (index / bits(), index % bits())
}

/// Keeps track of occupied and vacant slots in `VecArena`.
///
/// It's implemented as an array of blocks, where every block tracks only a small contiguous chunk
/// of slots. More precisely: as many slots as there are bits in one `usize`.
///
/// All blocks which are not fully occupied, except the last one, are linked together to form a
/// doubly-linked list. This list allows finding a vacant slot to acquire in `O(1)` and releasing
/// an occupied slot in `O(1)`, assuming the bitmap doesn't grow to accommodate more slots.
///
/// A block consists of:
///
/// * a bit mask (one `usize`), in which zeros are for vacant slots and ones for occupied slots
/// * index of the successor block in the linked list
/// * index of the predecessor block in the linked list
///
/// The last block is tricky to handle because it might not have the same number of slots as other
/// blocks, so it gets special treatment in the implementation.
pub struct Bitmap {
    /// Storage for the following sequences, in this order:
    ///
    /// * bit masks
    /// * indices to successor nodes
    /// * indices to predecessor nodes
    ///
    /// All three sequences are stored in this single contiguous array.
    ///
    /// The most common and costly operation during the lifetime of a bitmap is testing
    /// whether a slot is occupied. Storing bit masks close together improves cache
    /// performance.
    data: *mut usize,

    /// Number of reserved slots.
    len: usize,

    /// Number of occupied slots.
    occupied: usize,

    /// Index of the first block in the linked list, or `!0` if the list is empty.
    head: usize,
}

impl Bitmap {
    /// Constructs a new `Bitmap` with zero slots.
    #[inline]
    pub fn new() -> Self {
        let data = {
            let mut vec = Vec::with_capacity(0);
            let ptr = vec.as_mut_ptr();
            mem::forget(vec);
            ptr
        };
        Bitmap {
            data: data,
            len: 0,
            occupied: 0,
            head: !0,
        }
    }

    /// Returns the number of reserved slots in the bitmap.
    #[inline]
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns the number of occupied slots in the bitmap.
    #[inline]
    pub fn occupied(&self) -> usize {
        self.occupied
    }

    /// Finds a vacant slot, marks as occupied, and returns its index.
    ///
    /// # Panics
    ///
    /// Panics if all slots are occupied.
    pub fn acquire(&mut self) -> usize {
        assert!(self.occupied < self.len, "no vacant slots to acquire, len = {}", self.len);

        let block = if self.head == !0 {
            // The list is empty, so try the last block.
            blocks_for(self.len) - 1
        } else {
            // The list has a head. Take a vacant slot from the head block.
            self.head
        };

        let index = unsafe {
            // Find the rightmost zero bit in the mask. Taking the rightmost zero is always ok,
            // even if this is the last block.
            let offset = (!*self.mask(block)).trailing_zeros() as usize;
            debug_assert!(offset < bits());

            slot_index(block, offset)
        };
        self.acquire_at(index);
        index
    }

    /// Marks the vacant slot at `index` as occupied.
    ///
    /// # Panics
    ///
    /// Panics if the slot is already occupied.
    pub fn acquire_at(&mut self, index: usize) {
        assert!(index < self.len, "`index` out of bounds");

        let (block, offset) = block_and_offset(index);
        unsafe {
            // Mark the slot as occupied in the block's bit mask.
            assert!(*self.mask(block) >> offset & 1 == 0, "occupied slot at `index`");
            *self.mask(block) ^= 1 << offset;

            // If the block has just become fully occupied, remove it from the list.
            if block == self.head && *self.mask(block) == !0 {
                let succ = *self.successor(block);
                self.link_blocks(!0, succ);
                self.head = succ;
            }

            self.occupied += 1;
        }
    }

    /// Releases the occupied slot at `index`.
    ///
    /// # Panics
    ///
    /// Panics if the `index` is out of bounds or the slot is vacant.
    pub fn release(&mut self, index: usize) {
        assert!(index < self.len, "`index` out of bounds");

        let (block, offset) = block_and_offset(index);
        unsafe {
            // Make sure we're not releasing a vacant slot.
            assert!(*self.mask(block) >> offset & 1 == 1, "releasing a vacant slot");

            self.occupied -= 1;

            // If the block is fully occupied, insert it back into the list.
            if *self.mask(block) == !0 {
                let head = self.head;
                self.link_blocks(!0, block);
                self.link_blocks(block, head);
                self.head = block;
            }

            // Mark the slot as vacant in the block's bit mask.
            *self.mask(block) ^= 1 << offset;
        }
    }

    /// Given the required minimal additional number of slots to reserve, returns the number of
    /// slots to reserve in order to keep amortized time complexity.
    fn amortized_reserve(&self, additional: usize) -> usize {
        let len = self.len();
        let required_len = len.checked_add(additional).expect("len overflow");
        let double_len = len.checked_mul(2).expect("len overflow");
        cmp::max(required_len, double_len) - len
    }

    /// Reserves at least `additional` more slots for new elements to be inserted. The arena may
    /// reserve more space to avoid frequent reallocations.
    ///
    /// # Panics
    ///
    /// Panics if the new number of slots overflows `usize`.
    pub fn reserve(&mut self, additional: usize) {
        let amortized = self.amortized_reserve(additional);
        self.reserve_exact(amortized);
    }

    /// Reserves exactly `additional` more slots for new elements to be inserted.
    ///
    /// Note that the allocator may give the bitmap more space than it requests.
    ///
    /// # Panics
    ///
    /// Panics if the new number of slots overflows `usize`.
    pub fn reserve_exact(&mut self, additional: usize) {
        let old_blocks = blocks_for(self.len);
        self.len = self.len.checked_add(additional).expect("len overflow");
        let new_blocks = blocks_for(self.len);

        if new_blocks == old_blocks {
            // We can simply return.
            // The higher unused bits of the last block are always zero anyway.
            return;
        }

        unsafe {
            // If the last block had some unused or vacant slots, push it into the linked list.
            if old_blocks > 0 && *self.mask(old_blocks - 1) != !0 {
                let head = self.head;
                let last = old_blocks - 1;
                self.link_blocks(last, head);
                self.link_blocks(!0, last);
                self.head = last;
            }

            // Allocate a new array.
            let new_data = {
                // Every block contains three `usize` integers, so we need `3 * new_blocks` of
                // space.
                let mut vec = Vec::with_capacity(3 * new_blocks);
                let ptr = vec.as_mut_ptr();
                mem::forget(vec);
                ptr
            };

            // Copy the three old subarrays (bit masks, successor indices, predecessor indices)
            // into the new array.
            for i in 0..3 {
                ptr::copy_nonoverlapping(
                    self.data.offset((old_blocks * i) as isize),
                    new_data.offset((new_blocks * i) as isize),
                    old_blocks);
            }

            // Deallocate the old array.
            Vec::from_raw_parts(self.data, 0, 3 * old_blocks);

            // Set the new bit masks to zero.
            ptr::write_bytes(new_data.offset(old_blocks as isize), 0, new_blocks - old_blocks);

            // Set the new data now because we're about to push blocks into the linked list.
            self.data = new_data;

            // If there are at least two new blocks, that means there is at least one new block
            // that must be pushed into the linked list.
            if old_blocks + 2 <= new_blocks {
                let head = self.head;
                self.link_blocks(new_blocks - 2, head);
                self.link_blocks(!0, old_blocks);
                self.head = old_blocks;

                for block in old_blocks .. new_blocks - 2 {
                    self.link_blocks(block, block + 1);
                }
            }
        }
    }

    /// Returns `true` if the slot at `index` is occupied.
    #[inline]
    pub fn is_occupied(&self, index: usize) -> bool {
        assert!(index < self.len, "`index` out of bounds");

        let (block, offset) = block_and_offset(index);
        unsafe {
            *self.mask(block) >> offset & 1 != 0
        }
    }

    /// Returns an iterator over occupied slots.
    #[inline]
    pub fn iter(&self) -> Iter {
        Iter {
            bitmap: self,
            index: 0,
        }
    }

    #[cfg(debug_assertions)]
    pub fn check_invariants(&self) {
        // TODO
    }

    /// Links together blocks `a` and `b` so that `a` comes before `b` in the linked list.
    #[inline]
    unsafe fn link_blocks(&mut self, a: usize, b: usize) {
        if a != !0 { *self.successor(a) = b; }
        if b != !0 { *self.predecessor(b) = a; }
    }

    /// Returns the pointer to the bit mask of `block`.
    #[inline]
    unsafe fn mask(&self, block: usize) -> *mut usize {
        self.data.offset(block as isize)
    }

    /// Returns the pointer to the index of the successor of `block`.
    #[inline]
    unsafe fn successor(&self, block: usize) -> *mut usize {
        self.data.offset((blocks_for(self.len) + block) as isize)
    }

    /// Returns the pointer to the index of the predecessor of `block`.
    #[inline]
    unsafe fn predecessor(&self, block: usize) -> *mut usize {
        self.data.offset((2 * blocks_for(self.len) + block) as isize)
    }

    /// Returns the next occupied slot starting from (and including) the specified `index`.
    pub fn next_occupied(&self, mut index: usize) -> Option<usize> {
        while index < self.len() {
            let (block, offset) = block_and_offset(index);
            let mask = unsafe { *self.mask(block) };

            if mask != 0 {
                for off in offset .. bits() {
                    if mask >> off & 1 == 1 {
                        return Some(slot_index(block, off));
                    }
                }
            }

            index = slot_index(block + 1, 0);
        }
        None
    }

    /// Returns the previous occupied slot before the specified `index`.
    pub fn previous_occupied(&self, mut index: usize) -> Option<usize> {
        index = cmp::min(index, self.len());
        loop {
            let (block, offset) = block_and_offset(index);
            let mask = unsafe { *self.mask(block) };

            if mask != 0 {
                for off in (0 .. offset).rev() {
                    if mask >> off & 1 == 1 {
                        return Some(slot_index(block, off));
                    }
                }
            }

            if block == 0 {
                break;
            }
            index = slot_index(block - 1, 0);
        }
        None
    }
}

impl Drop for Bitmap {
    fn drop(&mut self) {
        unsafe {
            Vec::from_raw_parts(self.data, 0, 2 * blocks_for(self.len));
        }
    }
}

/// Iterates over all occupied slots.
pub struct Iter<'a> {
    /// The bitmap to iterate over.
    bitmap: &'a Bitmap,

    /// Index of the current slot.
    index: usize,
}

impl<'a> Iterator for Iter<'a> {
    type Item = usize;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            index
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let low = 0;
        let high = self.bitmap.len() - self.index;
        (low, Some(high))
    }

    #[inline]
    fn count(self) -> usize {
        self.bitmap.occupied
    }
}
<|endoftext|>


<commit_before>use std::rc::{Rc};
use std::cell::{RefCell};
use std::collections::HashMap;

pub struct Broker<T: Subscriber> {
    subscriptions: Subscriptions<T>,
}

impl<T: Subscriber> Broker<T> {
    pub fn new() -> Self {
        Broker{ subscriptions: Subscriptions::<T>::new() }
    }

    pub fn publish(&mut self, topic: &str, payload: &[u8]) {
        self.subscriptions.publish(topic, payload);
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        self.subscriptions.subscribe(subscriber, &topics)
    }
}

pub struct Subscriptions<T: Subscriber> {
    subscribers: Vec<Rc<RefCell<T>>>,
}

impl<T: Subscriber> Subscriptions<T> {
    fn new() -> Self {
        return Subscriptions { subscribers: vec![], }
    }

    fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        let subscriber = self.add_or_find(subscriber);
        subscriber.borrow_mut().append_topics(topics);
    }

    fn add_or_find(&mut self, subscriber: Rc<RefCell<T>>) -> Rc<RefCell<T>> {
        for sub in &self.subscribers {
            let sub1 = sub.clone();
            let sub2 = subscriber.clone();
            //this is horrible
            //borrow doesn't return T, it returns Ref<T>
            //since Ref<T> itself doesn't satisfy the template contraint, we need to deref
            //it to get the to the T inside. And since is_same takes borrows, we reapply
            //the ampersand
            if is_same(&*sub1.borrow(), &*sub2.borrow()) {
                return sub2;
            }
        }

        self.subscribers.push(subscriber);
        self.subscribers[self.subscribers.len() - 1].clone()
    }

    fn publish(&mut self, topic: &str, payload: &[u8]) {
        for subscriber in &mut self.subscribers {
            let subscriber = subscriber.clone();
            let mut subscriber = subscriber.borrow_mut();

            for _ in subscriber.topics().iter().filter(|&t| topic_matches(&topic, &t)) {
                subscriber.new_message(payload);
            }
        }
    }
}

//is same identity
fn is_same<T>(lhs: &T, rhs: &T) -> bool {
    return lhs as *const T as i64 == rhs as *const T as i64
}

fn topic_matches(pub_topic: &str, sub_topic: &str) -> bool {
    pub_topic == sub_topic
}

pub trait Subscriber {
    fn new_message(&mut self, bytes: &[u8]);
    fn append_topics(&mut self, topics: &[&str]);
    fn topics(&self) -> Vec<String>;
}

pub trait SubscriberT {
    fn new_message(&mut self, bytes: &[u8]);
}

struct Subscription<T: SubscriberT> {
    subscriber: Rc<RefCell<T>>,
    topics: Vec<String>,
}

impl<T: SubscriberT> Subscription<T> {
    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

struct Node<T: SubscriberT> {
    children: HashMap<String, Box<Node<T>>>,
    leaves: Vec<Subscription<T>>,
}

pub struct SubscriptionsT<T: SubscriberT> {
    tree: Node<T>,
}

impl<T: SubscriberT> SubscriptionsT<T> {
    fn new() -> Self {
        SubscriptionsT { tree: Node { children: HashMap::new(), leaves: vec![] }}
    }

    fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
    }

    fn publish(&self, topic: &str, payload: &[u8]) {

    }
}

#[cfg(test)]
struct TestSubscriberT {
    msgs: Vec<Vec<u8>>,
}

#[cfg(test)]
impl TestSubscriberT {
    fn new() -> Self {
        return TestSubscriberT{msgs: vec![]}
    }
}

#[cfg(test)]
impl SubscriberT for TestSubscriberT {
    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }
}

#[test]
fn test_subt() {
    let mut subs = SubscriptionsT::<TestSubscriberT>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriberT::new()));
    let mut subscriber = sub_ref.clone();
    subs.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    subs.subscribe(subscriber.clone(), &["topics/foo"]);
    subs.publish("topics/foo", &[0, 1, 9]);
    subs.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

}

#[cfg(test)]
struct TestSubscriber {
    msgs: Vec<Vec<u8>>,
    topics: Vec<String>,
}
#[cfg(test)]
impl TestSubscriber {
    fn new() -> Self {
        return TestSubscriber{msgs: vec![], topics: vec![]}
    }
}

#[cfg(test)]
impl Subscriber for TestSubscriber {

    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }

    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[test]
fn test_subscribe() {
    let mut broker = Broker::<TestSubscriber>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriber::new()));
    let subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}

#[test]
fn test_tree() {

}
<commit_msg>UT now identical to 'real' one<commit_after>use std::rc::{Rc};
use std::cell::{RefCell};
use std::collections::HashMap;

pub struct Broker<T: Subscriber> {
    subscriptions: Subscriptions<T>,
}

impl<T: Subscriber> Broker<T> {
    pub fn new() -> Self {
        Broker{ subscriptions: Subscriptions::<T>::new() }
    }

    pub fn publish(&mut self, topic: &str, payload: &[u8]) {
        self.subscriptions.publish(topic, payload);
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        self.subscriptions.subscribe(subscriber, &topics)
    }
}

pub struct Subscriptions<T: Subscriber> {
    subscribers: Vec<Rc<RefCell<T>>>,
}

impl<T: Subscriber> Subscriptions<T> {
    fn new() -> Self {
        return Subscriptions { subscribers: vec![], }
    }

    fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        let subscriber = self.add_or_find(subscriber);
        subscriber.borrow_mut().append_topics(topics);
    }

    fn add_or_find(&mut self, subscriber: Rc<RefCell<T>>) -> Rc<RefCell<T>> {
        for sub in &self.subscribers {
            let sub1 = sub.clone();
            let sub2 = subscriber.clone();
            //this is horrible
            //borrow doesn't return T, it returns Ref<T>
            //since Ref<T> itself doesn't satisfy the template contraint, we need to deref
            //it to get the to the T inside. And since is_same takes borrows, we reapply
            //the ampersand
            if is_same(&*sub1.borrow(), &*sub2.borrow()) {
                return sub2;
            }
        }

        self.subscribers.push(subscriber);
        self.subscribers[self.subscribers.len() - 1].clone()
    }

    fn publish(&mut self, topic: &str, payload: &[u8]) {
        for subscriber in &mut self.subscribers {
            let subscriber = subscriber.clone();
            let mut subscriber = subscriber.borrow_mut();

            for _ in subscriber.topics().iter().filter(|&t| topic_matches(&topic, &t)) {
                subscriber.new_message(payload);
            }
        }
    }
}

//is same identity
fn is_same<T>(lhs: &T, rhs: &T) -> bool {
    return lhs as *const T as i64 == rhs as *const T as i64
}

fn topic_matches(pub_topic: &str, sub_topic: &str) -> bool {
    pub_topic == sub_topic
}

pub trait Subscriber {
    fn new_message(&mut self, bytes: &[u8]);
    fn append_topics(&mut self, topics: &[&str]);
    fn topics(&self) -> Vec<String>;
}

pub trait SubscriberT {
    fn new_message(&mut self, bytes: &[u8]);
}

struct Subscription<T: SubscriberT> {
    subscriber: Rc<RefCell<T>>,
    topics: Vec<String>,
}

impl<T: SubscriberT> Subscription<T> {
    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

struct Node<T: SubscriberT> {
    children: HashMap<String, Box<Node<T>>>,
    leaves: Vec<Subscription<T>>,
}

pub struct SubscriptionsT<T: SubscriberT> {
    tree: Node<T>,
}

impl<T: SubscriberT> SubscriptionsT<T> {
    pub fn new() -> Self {
        SubscriptionsT { tree: Node { children: HashMap::new(), leaves: vec![] }}
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
    }

    pub fn publish(&self, topic: &str, payload: &[u8]) {

    }

    // fn add_or_find(&mut self, subscriber: Rc<RefCell<T>>) -> Rc<RefCell<T>> {
    // }

}

#[cfg(test)]
struct TestSubscriberT {
    msgs: Vec<Vec<u8>>,
}

#[cfg(test)]
impl TestSubscriberT {
    fn new() -> Self {
        return TestSubscriberT{msgs: vec![]}
    }
}

#[cfg(test)]
impl SubscriberT for TestSubscriberT {
    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }
}

#[test]
fn test_subt() {
    let mut broker = SubscriptionsT::<TestSubscriberT>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriberT::new()));
    let mut subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}

#[cfg(test)]
struct TestSubscriber {
    msgs: Vec<Vec<u8>>,
    topics: Vec<String>,
}
#[cfg(test)]
impl TestSubscriber {
    fn new() -> Self {
        return TestSubscriber{msgs: vec![], topics: vec![]}
    }
}

#[cfg(test)]
impl Subscriber for TestSubscriber {

    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }

    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[test]
fn test_subscribe() {
    let mut broker = Broker::<TestSubscriber>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriber::new()));
    let subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}

#[test]
fn test_tree() {

}
<|endoftext|>


<commit_before>use decoder::{Decoder, MAX_COMPONENTS};
use error::{Error, Result, UnsupportedFeature};
use huffman::{fill_default_mjpeg_tables, HuffmanDecoder, HuffmanTable};
use marker::Marker;
use parser::Predictor;
use parser::{
    parse_app, parse_com, parse_dht, parse_dqt, parse_dri, parse_sof, parse_sos,
    AdobeColorTransform, AppData, CodingProcess, Component, Dimensions, EntropyCoding, FrameInfo,
    IccChunk, ScanInfo,
};
use std::io::Read;

impl<R: Read> Decoder<R> {
    /// decode_scan_lossless
    pub fn decode_scan_lossless(
        &mut self,
        frame: &FrameInfo,
        scan: &ScanInfo,
    ) -> Result<(Option<Marker>, Vec<Vec<u16>>)> {
        let ncomp = scan.component_indices.len();
        let npixel = frame.image_size.height as usize * frame.image_size.width as usize;
        assert!(ncomp <= MAX_COMPONENTS);
        let mut results = vec![vec!(0;npixel); ncomp];

        let components: Vec<Component> = scan
            .component_indices
            .iter()
            .map(|&i| frame.components[i].clone())
            .collect();

        // Verify that all required huffman tables has been set.
        if scan
            .dc_table_indices
            .iter()
            .any(|&i| self.dc_huffman_tables[i].is_none())
        {
            return Err(Error::Format(
                "scan makes use of unset dc huffman table".to_owned(),
            ));
        }

        let mut huffman = HuffmanDecoder::new();
        let reader = &mut self.reader;
        let mut mcus_left_until_restart = self.restart_interval;
        let mut expected_rst_num = 0;
        let mut ra = [0u16; MAX_COMPONENTS];
        let mut rb = [0u16; MAX_COMPONENTS];
        let mut rc = [0u16; MAX_COMPONENTS];

        let width = frame.image_size.width as usize;
        let height = frame.image_size.height as usize;

        let mut differences = vec![Vec::with_capacity(npixel); ncomp];
        for mcu_y in 0..height {
            for mcu_x in 0..width {
                if self.restart_interval > 0 {
                    if mcus_left_until_restart == 0 {
                        match huffman.take_marker(reader)? {
                            Some(Marker::RST(n)) => {
                                if n != expected_rst_num {
                                    return Err(Error::Format(format!(
                                        "found RST{} where RST{} was expected",
                                        n, expected_rst_num
                                    )));
                                }

                                huffman.reset();

                                expected_rst_num = (expected_rst_num + 1) % 8;
                                mcus_left_until_restart = self.restart_interval;
                            }
                            Some(marker) => {
                                return Err(Error::Format(format!(
                                    "found marker {:?} inside scan where RST{} was expected",
                                    marker, expected_rst_num
                                )))
                            }
                            None => {
                                return Err(Error::Format(format!(
                                    "no marker found where RST{} was expected",
                                    expected_rst_num
                                )))
                            }
                        }
                    }

                    mcus_left_until_restart -= 1;
                }

                for (i, _component) in components.iter().enumerate() {
                    let dc_table = self.dc_huffman_tables[scan.dc_table_indices[i]]
                        .as_ref()
                        .unwrap();
                    let value = huffman.decode(reader, dc_table)?;
                    let diff = match value {
                        0 => 0,
                        1..=15 => huffman.receive_extend(reader, value)? as i32,
                        16 => 32768,
                        _ => {
                            // Section F.1.2.1.1
                            // Table F.1
                            return Err(Error::Format(
                                "invalid DC difference magnitude category".to_owned(),
                            ));
                        }
                    };
                    differences[i].push(diff);
                }
            }
        }

        if scan.predictor_selection == Predictor::Ra {
            for (i, _component) in components.iter().enumerate() {
                // calculate the top left pixel
                let diff = differences[i][0];
                let prediction = 1 << (frame.precision - scan.point_transform - 1) as i32;
                let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                let result = result << scan.point_transform;
                results[i][0] = result;

                // calculate leftmost column, using top pixel as predictor
                let mut previous = result;
                for mcu_y in 1..height {
                    let diff = differences[i][mcu_y * width];
                    let prediction = previous as i32;
                    let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                    let result = result << scan.point_transform;
                    results[i][mcu_y * width] = result;
                    previous = result;
                }

                // calculate rows, using left pixel as predictor
                for mcu_y in 0..height {
                    for mcu_x in 1..width {
                        let diff = differences[i][mcu_y * width + mcu_x];
                        let prediction = results[i][mcu_y * width + mcu_x - 1] as i32;
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        let result = result << scan.point_transform;
                        results[i][mcu_y * width + mcu_x] = result;
                    }
                }

            }
        } else {
            for mcu_y in 0..height {
                for mcu_x in 0..width {
                    for (i, _component) in components.iter().enumerate() {
                        let diff = differences[i][mcu_y * width + mcu_x];

                        // The following lines could be further optimized, e.g. moving the checks
                        // and updates of the previous values into the prediction function or
                        // iterating such that diagonals with mcu_x + mcu_y = const are computed at
                        // the same time to exploit independent predictions in this case
                        if mcu_x > 0 {
                            ra[i] = results[i][mcu_y * frame.image_size.width as usize + mcu_x - 1];
                        }
                        if mcu_y > 0 {
                            rb[i] = results[i][(mcu_y - 1) * frame.image_size.width as usize + mcu_x];
                            if mcu_x > 0 {
                                rc[i] = results[i]
                                    [(mcu_y - 1) * frame.image_size.width as usize + (mcu_x - 1)];
                            }
                        }
                        let prediction = predict(
                            ra[i] as i32,
                            rb[i] as i32,
                            rc[i] as i32,
                            scan.predictor_selection,
                            scan.point_transform,
                            frame.precision,
                            mcu_x,
                            mcu_y,
                            self.restart_interval > 0
                                && mcus_left_until_restart == self.restart_interval - 1,
                        );
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        results[i][mcu_y * width + mcu_x] = result << scan.point_transform;
                    }
                }
            }
        }

        let mut marker = huffman.take_marker(&mut self.reader)?;
        while let Some(Marker::RST(_)) = marker {
            marker = self.read_marker().ok();
        }
        Ok((marker, results))
    }
}

/// H.1.2.1
fn predict(
    ra: i32,
    rb: i32,
    rc: i32,
    predictor: Predictor,
    point_transform: u8,
    input_precision: u8,
    ix: usize,
    iy: usize,
    restart: bool,
) -> i32 {
    let result = if (ix == 0 && iy == 0) || restart {
        // start of first line or restart
        1 << (input_precision - point_transform - 1)
    } else if iy == 0 {
        // rest of first line
        ra
    } else if ix == 0 {
        // start of other line
        rb
    } else {
        // use predictor Table H.1
        match predictor {
            Predictor::NoPrediction => 0,
            Predictor::Ra => ra,
            Predictor::Rb => rb,
            Predictor::Rc => rc,
            Predictor::RaRbRc1 => ra + rb - rc,
            Predictor::RaRbRc2 => ra + ((rb - rc) >> 1),
            Predictor::RaRbRc3 => rb + ((ra - rc) >> 1),
            Predictor::RaRb => (ra + rb) / 2,
        }
    };
    result
}

pub fn compute_image_lossless(
    frame: &FrameInfo,
    mut data: Vec<Vec<u16>>
) -> Result<Vec<u8>> {
    if data.is_empty() || data.iter().any(Vec::is_empty) {
        return Err(Error::Format("not all components have data".to_owned()));
    }
    let output_size = frame.output_size;
    let components = &frame.components;
    let ncomp = components.len();

    if ncomp == 1 {
        let decoded = convert_to_u8(frame, data.remove(0));
        Ok(decoded)
    } else {
        let mut decoded: Vec<u16> = vec![
            0u16;
            ncomp * output_size.width as usize * output_size.height as usize
        ];
        for (x, chunk) in decoded.chunks_mut(ncomp).enumerate() {
            for (i, (component_data, _)) in data.iter().zip(components.iter()).enumerate() {
                chunk[i] = component_data[x];
            }
        }
        let decoded = convert_to_u8(frame, decoded);
        Ok(decoded)
    }
}

fn convert_to_u8(frame: &FrameInfo, data: Vec<u16>) -> Vec<u8> {
    if frame.precision == 8 {
        data.iter().map(|x| *x as u8).collect()
    } else {
        // we output native endian, which is the standard for image-rs
        let out: Vec<[u8; 2]> = data.iter().map(|x| x.to_ne_bytes()).collect();
        out.iter().flatten().map(|x| *x).collect()
    }
}
<commit_msg>remove unused imports<commit_after>use decoder::{Decoder, MAX_COMPONENTS};
use error::{Error, Result};
use huffman::HuffmanDecoder;
use marker::Marker;
use parser::Predictor;
use parser::{
    Component,  FrameInfo,
    ScanInfo,
};
use std::io::Read;

impl<R: Read> Decoder<R> {
    /// decode_scan_lossless
    pub fn decode_scan_lossless(
        &mut self,
        frame: &FrameInfo,
        scan: &ScanInfo,
    ) -> Result<(Option<Marker>, Vec<Vec<u16>>)> {
        let ncomp = scan.component_indices.len();
        let npixel = frame.image_size.height as usize * frame.image_size.width as usize;
        assert!(ncomp <= MAX_COMPONENTS);
        let mut results = vec![vec!(0, npixel); ncomp];

        let components: Vec<Component> = scan
            .component_indices
            .iter()
            .map(|&i| frame.components[i].clone())
            .collect();

        // Verify that all required huffman tables has been set.
        if scan
            .dc_table_indices
            .iter()
            .any(|&i| self.dc_huffman_tables[i].is_none())
        {
            return Err(Error::Format(
                "scan makes use of unset dc huffman table".to_owned(),
            ));
        }

        let mut huffman = HuffmanDecoder::new();
        let reader = &mut self.reader;
        let mut mcus_left_until_restart = self.restart_interval;
        let mut expected_rst_num = 0;
        let mut ra = [0u16; MAX_COMPONENTS];
        let mut rb = [0u16; MAX_COMPONENTS];
        let mut rc = [0u16; MAX_COMPONENTS];

        let width = frame.image_size.width as usize;
        let height = frame.image_size.height as usize;

        let mut differences = vec![Vec::with_capacity(npixel); ncomp];
        for _mcu_y in 0..height {
            for _mcu_x in 0..width {
                if self.restart_interval > 0 {
                    if mcus_left_until_restart == 0 {
                        match huffman.take_marker(reader)? {
                            Some(Marker::RST(n)) => {
                                if n != expected_rst_num {
                                    return Err(Error::Format(format!(
                                        "found RST{} where RST{} was expected",
                                        n, expected_rst_num
                                    )));
                                }

                                huffman.reset();

                                expected_rst_num = (expected_rst_num + 1) % 8;
                                mcus_left_until_restart = self.restart_interval;
                            }
                            Some(marker) => {
                                return Err(Error::Format(format!(
                                    "found marker {:?} inside scan where RST{} was expected",
                                    marker, expected_rst_num
                                )))
                            }
                            None => {
                                return Err(Error::Format(format!(
                                    "no marker found where RST{} was expected",
                                    expected_rst_num
                                )))
                            }
                        }
                    }

                    mcus_left_until_restart -= 1;
                }

                for (i, _component) in components.iter().enumerate() {
                    let dc_table = self.dc_huffman_tables[scan.dc_table_indices[i]]
                        .as_ref()
                        .unwrap();
                    let value = huffman.decode(reader, dc_table)?;
                    let diff = match value {
                        0 => 0,
                        1..=15 => huffman.receive_extend(reader, value)? as i32,
                        16 => 32768,
                        _ => {
                            // Section F.1.2.1.1
                            // Table F.1
                            return Err(Error::Format(
                                "invalid DC difference magnitude category".to_owned(),
                            ));
                        }
                    };
                    differences[i].push(diff);
                }
            }
        }

        if scan.predictor_selection == Predictor::Ra {
            for (i, _component) in components.iter().enumerate() {
                // calculate the top left pixel
                let diff = differences[i][0];
                let prediction = 1 << (frame.precision - scan.point_transform - 1) as i32;
                let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                let result = result << scan.point_transform;
                results[i][0] = result;

                // calculate leftmost column, using top pixel as predictor
                let mut previous = result;
                for mcu_y in 1..height {
                    let diff = differences[i][mcu_y * width];
                    let prediction = previous as i32;
                    let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                    let result = result << scan.point_transform;
                    results[i][mcu_y * width] = result;
                    previous = result;
                }

                // calculate rows, using left pixel as predictor
                for mcu_y in 0..height {
                    for mcu_x in 1..width {
                        let diff = differences[i][mcu_y * width + mcu_x];
                        let prediction = results[i][mcu_y * width + mcu_x - 1] as i32;
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        let result = result << scan.point_transform;
                        results[i][mcu_y * width + mcu_x] = result;
                    }
                }
            }
        } else {
            for mcu_y in 0..height {
                for mcu_x in 0..width {
                    for (i, _component) in components.iter().enumerate() {
                        let diff = differences[i][mcu_y * width + mcu_x];

                        // The following lines could be further optimized, e.g. moving the checks
                        // and updates of the previous values into the prediction function or
                        // iterating such that diagonals with mcu_x + mcu_y = const are computed at
                        // the same time to exploit independent predictions in this case
                        if mcu_x > 0 {
                            ra[i] = results[i][mcu_y * frame.image_size.width as usize + mcu_x - 1];
                        }
                        if mcu_y > 0 {
                            rb[i] =
                                results[i][(mcu_y - 1) * frame.image_size.width as usize + mcu_x];
                            if mcu_x > 0 {
                                rc[i] = results[i]
                                    [(mcu_y - 1) * frame.image_size.width as usize + (mcu_x - 1)];
                            }
                        }
                        let prediction = predict(
                            ra[i] as i32,
                            rb[i] as i32,
                            rc[i] as i32,
                            scan.predictor_selection,
                            scan.point_transform,
                            frame.precision,
                            mcu_x,
                            mcu_y,
                            self.restart_interval > 0
                                && mcus_left_until_restart == self.restart_interval - 1,
                        );
                        let result = ((prediction + diff) & 0xFFFF) as u16; // modulo 2^16
                        results[i][mcu_y * width + mcu_x] = result << scan.point_transform;
                    }
                }
            }
        }

        let mut marker = huffman.take_marker(&mut self.reader)?;
        while let Some(Marker::RST(_)) = marker {
            marker = self.read_marker().ok();
        }
        Ok((marker, results))
    }
}

/// H.1.2.1
fn predict(
    ra: i32,
    rb: i32,
    rc: i32,
    predictor: Predictor,
    point_transform: u8,
    input_precision: u8,
    ix: usize,
    iy: usize,
    restart: bool,
) -> i32 {
    let result = if (ix == 0 && iy == 0) || restart {
        // start of first line or restart
        1 << (input_precision - point_transform - 1)
    } else if iy == 0 {
        // rest of first line
        ra
    } else if ix == 0 {
        // start of other line
        rb
    } else {
        // use predictor Table H.1
        match predictor {
            Predictor::NoPrediction => 0,
            Predictor::Ra => ra,
            Predictor::Rb => rb,
            Predictor::Rc => rc,
            Predictor::RaRbRc1 => ra + rb - rc,
            Predictor::RaRbRc2 => ra + ((rb - rc) >> 1),
            Predictor::RaRbRc3 => rb + ((ra - rc) >> 1),
            Predictor::RaRb => (ra + rb) / 2,
        }
    };
    result
}

pub fn compute_image_lossless(frame: &FrameInfo, mut data: Vec<Vec<u16>>) -> Result<Vec<u8>> {
    if data.is_empty() || data.iter().any(Vec::is_empty) {
        return Err(Error::Format("not all components have data".to_owned()));
    }
    let output_size = frame.output_size;
    let components = &frame.components;
    let ncomp = components.len();

    if ncomp == 1 {
        let decoded = convert_to_u8(frame, data.remove(0));
        Ok(decoded)
    } else {
        let mut decoded: Vec<u16> =
            vec![0u16; ncomp * output_size.width as usize * output_size.height as usize];
        for (x, chunk) in decoded.chunks_mut(ncomp).enumerate() {
            for (i, (component_data, _)) in data.iter().zip(components.iter()).enumerate() {
                chunk[i] = component_data[x];
            }
        }
        let decoded = convert_to_u8(frame, decoded);
        Ok(decoded)
    }
}

fn convert_to_u8(frame: &FrameInfo, data: Vec<u16>) -> Vec<u8> {
    if frame.precision == 8 {
        data.iter().map(|x| *x as u8).collect()
    } else {
        // we output native endian, which is the standard for image-rs
        let out: Vec<[u8; 2]> = data.iter().map(|x| x.to_ne_bytes()).collect();
        out.iter().flatten().map(|x| *x).collect()
    }
}
<|endoftext|>