<commit_before>use circle::Circle;
use fx32::Fx32;
use fx32::Fx32Small;
use ray2::Ray2;
use vec2::Vec2;
use wide_line::WideLine;

use std::cmp;

use std::io::Error;
use std::io::prelude::*;
use std::io::BufWriter;
use std::fs::File;

#[derive (Clone, Copy)]
pub enum Ray2TraceResult {
	Hit (Fx32, Vec2 <Fx32>, Vec2 <Fx32Small>),
	Pop (Vec2 <Fx32>, Vec2 <Fx32Small>),
	Miss,
}

pub fn fold_closer_result (a: Ray2TraceResult, b: Ray2TraceResult) -> Ray2TraceResult {
	match a {
		Ray2TraceResult::Miss => {
			return b;
		},
		Ray2TraceResult::Hit (a_t, ..) => {
			match b {
				Ray2TraceResult::Miss => {
					return a;
				},
				Ray2TraceResult::Hit (b_t, ..) => {
					if a_t < b_t {
						return a;
					}
					else {
						return b;
					}
				},
				Ray2TraceResult::Pop (..) => {
					return b;
				}
			}
		},
		Ray2TraceResult::Pop (..) => {
			return a;
		},
	}
}

pub fn write_vec2 <T> (writer: &mut T, v: &Vec2 <Fx32>) where T: Write {
	write! (writer, "v {} 0 {}\n", v.x.to_f64 () / 256.0, v.y.to_f64 () / 256.0).unwrap ();
}

pub fn apply_dt (ray: &Ray2, dt: Fx32Small) -> Ray2 {
	Ray2 {
		start: ray.start,
		dir: ray.dir * dt,
	}
}

pub fn test_ray_trace (filename: &str, offset: Fx32) -> Result <(), Error> {
	let scale = 1;
	//let scale_fx = Fx32::from_int (scale);

	let obstacle = vec! [
	Circle {
		center: Vec2 {x: Fx32::from_q (40, scale), y: Fx32::from_q (470 - 34, scale)},
		radius: Fx32::from_q (20, scale),
	},
	Circle {
		center: Vec2 {x: Fx32::from_q (256, scale), y: Fx32::from_q (512, scale)},
		radius: Fx32::from_q (20, scale),
	},
	Circle {
		center: Vec2 {x: Fx32::from_q (480, scale), y: Fx32::from_q (460, scale)},
		radius: Fx32::from_q (20, scale),
	},
	];

	let lines = vec! [
	WideLine {
		start: obstacle [0].center,
		end: obstacle [1].center,
		radius: Fx32::from_q (20, scale),
	},
	WideLine {
		start: obstacle [1].center,
		end: obstacle [2].center,
		radius: Fx32::from_q (20, scale),
	},
	];

	let mut num_bounces = 0;
	let mut num_ticks = 0;

	let inv_dt = 1;
	let gravity = Vec2::<Fx32> {
		x: Fx32::from_q (0, 1),
		y: Fx32::from_q (2, 1),
	};

	let obj_file = try!(File::create(filename));
	let mut writer = BufWriter::new (obj_file);

	let mut vertex_i = 1;
	let mut polyline_start = vertex_i;

	for x in 0..128 {
		let x = x * 4;
		let mut particle = Ray2 {
			start: Vec2 {
				x: Fx32::from_q (x * 2, scale * 2) + offset,
				y: Fx32::from_q (0, scale)
			},
			dir: Vec2 {
				x: Fx32::from_q (0, scale),
				y: Fx32::from_q (1, scale),
			},
		};

		write_vec2 (&mut writer, &particle.start);
		vertex_i += 1;

		let dt = Fx32::from_q (1, inv_dt).to_small ();

		for tick in 0..500 {
			let trace_result = {
				let point_results = obstacle.iter ().map (|obstacle| ray_trace_circle_2 (&apply_dt (&particle, dt), obstacle));

				let line_results = lines.iter ().map (|line| ray_trace_line (&apply_dt (&particle, dt), line));

				point_results.chain (line_results).fold ( Ray2TraceResult::Miss, fold_closer_result)
			};

			match trace_result {
				Ray2TraceResult::Miss => {
					//let air_drag = Fx32 { x: particle.dir.length_sq ().x / -16384 };
					//let air_force = particle.dir * air_drag;
					let new_dir = particle.dir + gravity * dt;// + Vec2::<Fx32>::from (air_force * dt);
					let sum_dir = particle.dir + new_dir;
					let average_dir = Vec2::<Fx32> {
						x: Fx32 { x: sum_dir.x.x / 2 },
						y: Fx32 { x: sum_dir.y.x / 2 },
					};

					particle.start = particle.start + (average_dir * dt);
					particle.dir = new_dir;
				},
				Ray2TraceResult::Pop (ccd_pos, normal) => {
					println! ("{}: Pop from {:?} to {:?}", tick, particle.start, ccd_pos);

					let reflected_dir = particle.dir.reflect (normal);

					let new_dir = reflected_dir + gravity * dt;
					let sum_dir = reflected_dir + new_dir;
					let average_dir = Vec2::<Fx32> {
						x: Fx32 { x: sum_dir.x.x / 2 },
						y: Fx32 { x: sum_dir.y.x / 2 },
					};

					particle.start = ccd_pos + (average_dir * dt);
					particle.dir = new_dir;
				},
				Ray2TraceResult::Hit (_, ccd_pos, normal) => {
					println! ("{}: Hit from {:?} to {:?}", tick, particle.start, ccd_pos);

					particle.start = ccd_pos;
					particle.dir = particle.dir.reflect (normal) * Fx32::from_q (768, 1024);
					//particle.dir = normal;
					num_bounces += 1;
				},
			};

			write_vec2 (&mut writer, &particle.start);
			vertex_i += 1;

			num_ticks += 1;

			if particle.start.y > 768 {
				break;
			}
		}

		for i in polyline_start..vertex_i - 1 {
			try! (write! (writer, "f {} {}\n", i, i + 1));
		}
		polyline_start = vertex_i
	}

	println! ("num_bounces: {}", num_bounces);
	println! ("num_ticks: {}", num_ticks);

	Ok (())
}

pub struct Basis2 {
	pub x: Vec2 <Fx32Small>,
	pub y: Vec2 <Fx32Small>,
}

// Constructs a basis to ray space
// Such that the ray is the X axis
pub fn get_ray_basis (ray: &Ray2, ray_length: Fx32) -> Basis2 {
	let basis_x_big = if ray_length == 0 {
		ray.dir
	}
	else {
		ray.dir / ray_length
	};

	Basis2 {
		x: basis_x_big.to_small (),
		y: basis_x_big.cross ().to_small (),
	}
}

impl Basis2 {
	pub fn to_space (&self, v: &Vec2 <Fx32>) -> Vec2 <Fx32> {
		Vec2::<Fx32> {
			x: *v * self.x,
			y: *v * self.y,
		}
	}

	pub fn from_space (&self, v: &Vec2 <Fx32>) -> Vec2 <Fx32> {
		Vec2::<Fx32> {
			x: (v.x * self.x.x) + (v.y * self.y.x),
			y: (v.x * self.x.y) + (v.y * self.y.y),
		}
	}
}

pub fn ray_trace_line (ray: &Ray2, line: &WideLine) -> Ray2TraceResult {
	// Line has zero length and so zero collision area
	// If we run through the math we may get a 0 / 0 result
	// so just drop it now
	assert! (line.start != line.end);

	let margin = Fx32::from_q (256, 256);

	let line_tangent: Vec2 <Fx32> = line.end - line.start;
	let line_tangent = Vec2::<Fx32> {
		x: line_tangent.x * Fx32::from_q (1, 256),
		y: line_tangent.y * Fx32::from_q (1, 256),
	};
	let line_normal = line_tangent.cross ().normalized ();

	let sdf: Fx32 = (ray.start - line.start) * line_normal;

	// Flip the normal so it's towards the ray
	// This will allow us to extrude the correct side
	let line_normal = if sdf > 0 {
		line_normal
	}
	else {
		-line_normal
	};

	let big_normal: Vec2 <Fx32> = line_normal.into ();

	let along = (ray.start - line.start) * line_tangent;

	if sdf.abs () <= line.radius && along > 0 && along < (line.end - line.start) * line_tangent {
		// Ray has already started inside and we should push it out
		let towards = -Fx32::from (ray.dir * line_normal);

		let safe_point = ray.start + big_normal * (line.radius - sdf.abs () + margin);

		if towards > 0 {
			// Ray is moving deeper inside - Pop it out and reflect it
			return Ray2TraceResult::Pop (safe_point, line_normal);
		}
		else {
			// Ray is moving out - Pop it out but consume its timestep
			return Ray2TraceResult::Pop (safe_point, line_normal);
		}
	}

	let extrude_vector = (line.radius) * Vec2::<Fx32>::from (line_normal);
	// Extrude the line
	let line = WideLine {
		start: line.start + extrude_vector,
		end: line.end + extrude_vector,
		radius: Fx32::from_int (0),
	};

	let ray_length = ray.dir.length ();

	if ray_length == 0 {
		return Ray2TraceResult::Miss;
	}

	let basis = get_ray_basis (ray, ray_length);

	// Transform the line to ray space
	let line = WideLine {
		start: basis.to_space (&(line.start - ray.start)),
		end: basis.to_space (&(line.end - ray.start)),
		radius: line.radius,
	};

	// Line will have zero y difference resulting in a divide by zero
	if line.start.y == line.end.y {
		//return Ray2TraceResult::Miss;
	}

	// Apply SAT
	if line.start.y > 0 && line.end.y > 0 {
		return Ray2TraceResult::Miss;
	}
	else if line.start.y < 0 && line.end.y < 0 {
		return Ray2TraceResult::Miss;
	}

	// If we get here, we know that the obstacle line crosses our ray somewhere

	// Note: Lines must not have 0 length
	let obstacle_t = (-line.start.y) / (line.end.y - line.start.y);

	assert! (obstacle_t >= 0);
	assert! (obstacle_t <= 1, format!("{:?}", ray_length));

	// lerp
	let crossing_x = line.start.x * (Fx32::from_int (1) - obstacle_t) + line.end.x * obstacle_t;

	if crossing_x < 0 || crossing_x > ray_length + margin {
		return Ray2TraceResult::Miss;
	}

	// The line segments intersect
	let t = crossing_x / ray_length;

	let safe_point = ray.start + ray.dir * t + big_normal * margin;

	return Ray2TraceResult::Hit (
		t,
		safe_point,
		line_normal,
	);
}

pub fn ray_trace_circle_2 (ray: &Ray2, circle: &Circle) -> Ray2TraceResult {
	let ray_length = ray.dir.length ();
	let basis = get_ray_basis (ray, ray_length);

	let center_in_ray_space = basis.to_space (&(circle.center - ray.start));

	if center_in_ray_space.x < 0 {
		return Ray2TraceResult::Miss;
	}

	if center_in_ray_space.x > ray_length + circle.radius {
		return Ray2TraceResult::Miss;
	}

	if center_in_ray_space.y.abs () >= circle.radius {
		// Prevents a negative root later on
		return Ray2TraceResult::Miss;
	}

	let ray_space_x = center_in_ray_space.x - (circle.radius.square () - center_in_ray_space.y.square ()).sqrt ();

	if ray_space_x > ray_length || ray_length == 0 {
		return Ray2TraceResult::Miss;
	}

	let t = ray_space_x / ray_length;
	let ccd_pos = ray.start + ray.dir * t;

	if t <= 1 {
		return Ray2TraceResult::Hit (
			Fx32 { x: cmp::max (t.x, Fx32::from_int (0).x) },
			ccd_pos,
			((ccd_pos - circle.center) / circle.radius).to_small (),
		);
	}
	else {
		return Ray2TraceResult::Miss;
	}
}
<commit_msg>I had been integrating with the wrong velocity<commit_after>use circle::Circle;
use fx32::Fx32;
use fx32::Fx32Small;
use ray2::Ray2;
use vec2::Vec2;
use wide_line::WideLine;

use std::cmp;

use std::io::Error;
use std::io::prelude::*;
use std::io::BufWriter;
use std::fs::File;

#[derive (Clone, Copy)]
pub enum Ray2TraceResult {
	Hit (Fx32, Vec2 <Fx32>, Vec2 <Fx32Small>),
	Pop (Vec2 <Fx32>, Vec2 <Fx32Small>),
	Miss,
}

pub fn fold_closer_result (a: Ray2TraceResult, b: Ray2TraceResult) -> Ray2TraceResult {
	match a {
		Ray2TraceResult::Miss => {
			return b;
		},
		Ray2TraceResult::Hit (a_t, ..) => {
			match b {
				Ray2TraceResult::Miss => {
					return a;
				},
				Ray2TraceResult::Hit (b_t, ..) => {
					if a_t < b_t {
						return a;
					}
					else {
						return b;
					}
				},
				Ray2TraceResult::Pop (..) => {
					return b;
				}
			}
		},
		Ray2TraceResult::Pop (..) => {
			return a;
		},
	}
}

pub fn write_vec2 <T> (writer: &mut T, v: &Vec2 <Fx32>) where T: Write {
	write! (writer, "v {} 0 {}\n", v.x.to_f64 () / 256.0, v.y.to_f64 () / 256.0).unwrap ();
}

pub fn apply_dt (ray: &Ray2, dt: Fx32Small) -> Ray2 {
	Ray2 {
		start: ray.start,
		dir: ray.dir * dt,
	}
}

pub fn test_ray_trace (filename: &str, offset: Fx32) -> Result <(), Error> {
	let scale = 1;
	//let scale_fx = Fx32::from_int (scale);

	let obstacle = vec! [
	Circle {
		center: Vec2 {x: Fx32::from_q (40, scale), y: Fx32::from_q (470 - 34, scale)},
		radius: Fx32::from_q (20, scale),
	},
	Circle {
		center: Vec2 {x: Fx32::from_q (256, scale), y: Fx32::from_q (512, scale)},
		radius: Fx32::from_q (20, scale),
	},
	Circle {
		center: Vec2 {x: Fx32::from_q (480, scale), y: Fx32::from_q (460, scale)},
		radius: Fx32::from_q (20, scale),
	},
	];

	let lines = vec! [
	WideLine {
		start: obstacle [0].center,
		end: obstacle [1].center,
		radius: Fx32::from_q (20, scale),
	},
	WideLine {
		start: obstacle [1].center,
		end: obstacle [2].center,
		radius: Fx32::from_q (20, scale),
	},
	];

	let mut num_bounces = 0;
	let mut num_ticks = 0;

	let inv_dt = 1;
	let gravity = Vec2::<Fx32> {
		x: Fx32::from_q (0, 1),
		y: Fx32::from_q (2, 1),
	};

	let obj_file = try!(File::create(filename));
	let mut writer = BufWriter::new (obj_file);

	let mut vertex_i = 1;
	let mut polyline_start = vertex_i;

	for x in 0..32 {
		let x = x * 16;
		let mut particle = Ray2 {
			start: Vec2 {
				x: Fx32::from_q (x * 2, scale * 2) + offset,
				y: Fx32::from_q (0, scale)
			},
			dir: Vec2 {
				x: Fx32::from_q (0, scale),
				y: Fx32::from_q (1, scale),
			},
		};

		write_vec2 (&mut writer, &particle.start);
		vertex_i += 1;

		let dt = Fx32::from_q (1, inv_dt).to_small ();

		for tick in 0..500 {
			let trace_result = {
				let point_results = obstacle.iter ().map (|obstacle| ray_trace_circle_2 (&apply_dt (&particle, dt), obstacle));

				let line_results = lines.iter ().map (|line| ray_trace_line (&apply_dt (&particle, dt), line));

				point_results.chain (line_results).fold ( Ray2TraceResult::Miss, fold_closer_result)
			};

			match trace_result {
				Ray2TraceResult::Miss => {
					//let air_drag = Fx32 { x: particle.dir.length_sq ().x / -16384 };
					//let air_force = particle.dir * air_drag;
					let new_dir = particle.dir + gravity * dt;// + Vec2::<Fx32>::from (air_force * dt);
					let sum_dir = particle.dir + new_dir;
					let average_dir = Vec2::<Fx32> {
						x: Fx32 { x: sum_dir.x.x / 2 },
						y: Fx32 { x: sum_dir.y.x / 2 },
					};

					particle.start = particle.start + (particle.dir * dt);
					particle.dir = new_dir;
				},
				Ray2TraceResult::Pop (ccd_pos, normal) => {
					println! ("{}: Pop from {:?} to {:?}", tick, particle.start, ccd_pos);

					let reflected_dir = particle.dir.reflect (normal);

					let new_dir = reflected_dir + gravity * dt;
					let sum_dir = reflected_dir + new_dir;
					let average_dir = Vec2::<Fx32> {
						x: Fx32 { x: sum_dir.x.x / 2 },
						y: Fx32 { x: sum_dir.y.x / 2 },
					};

					particle.start = ccd_pos;// + (average_dir * dt);
					particle.dir = new_dir;
				},
				Ray2TraceResult::Hit (_, ccd_pos, normal) => {
					println! ("{}: Hit from {:?} to {:?}", tick, particle.start, ccd_pos);

					particle.start = ccd_pos;
					particle.dir = particle.dir.reflect (normal) * Fx32::from_q (768, 1024);
					//particle.dir = normal;
					num_bounces += 1;
				},
			};

			write_vec2 (&mut writer, &particle.start);
			vertex_i += 1;

			num_ticks += 1;

			if particle.start.y > 768 {
				break;
			}
		}

		for i in polyline_start..vertex_i - 1 {
			try! (write! (writer, "f {} {}\n", i, i + 1));
		}
		polyline_start = vertex_i
	}

	println! ("num_bounces: {}", num_bounces);
	println! ("num_ticks: {}", num_ticks);

	Ok (())
}

pub struct Basis2 {
	pub x: Vec2 <Fx32Small>,
	pub y: Vec2 <Fx32Small>,
}

// Constructs a basis to ray space
// Such that the ray is the X axis
pub fn get_ray_basis (ray: &Ray2, ray_length: Fx32) -> Basis2 {
	let basis_x_big = if ray_length == 0 {
		ray.dir
	}
	else {
		ray.dir / ray_length
	};

	Basis2 {
		x: basis_x_big.to_small (),
		y: basis_x_big.cross ().to_small (),
	}
}

impl Basis2 {
	pub fn to_space (&self, v: &Vec2 <Fx32>) -> Vec2 <Fx32> {
		Vec2::<Fx32> {
			x: *v * self.x,
			y: *v * self.y,
		}
	}

	pub fn from_space (&self, v: &Vec2 <Fx32>) -> Vec2 <Fx32> {
		Vec2::<Fx32> {
			x: (v.x * self.x.x) + (v.y * self.y.x),
			y: (v.x * self.x.y) + (v.y * self.y.y),
		}
	}
}

pub fn ray_trace_line (ray: &Ray2, line: &WideLine) -> Ray2TraceResult {
	// Line has zero length and so zero collision area
	// If we run through the math we may get a 0 / 0 result
	// so just drop it now
	assert! (line.start != line.end);

	let margin = Fx32::from_q (0, 256);

	let line_tangent: Vec2 <Fx32> = line.end - line.start;
	let line_tangent = Vec2::<Fx32> {
		x: line_tangent.x * Fx32::from_q (1, 256),
		y: line_tangent.y * Fx32::from_q (1, 256),
	};
	let line_normal = line_tangent.cross ().normalized ();

	let sdf: Fx32 = (ray.start - line.start) * line_normal;

	// Flip the normal so it's towards the ray
	// This will allow us to extrude the correct side
	let line_normal = if sdf > 0 {
		line_normal
	}
	else {
		-line_normal
	};

	let big_normal: Vec2 <Fx32> = line_normal.into ();

	let along = (ray.start - line.start) * line_tangent;

	if sdf.abs () <= line.radius && along > 0 && along < (line.end - line.start) * line_tangent {
		// Ray has already started inside and we should push it out
		let towards = -Fx32::from (ray.dir * line_normal);

		let safe_point = ray.start + big_normal * (line.radius - sdf.abs () + margin);

		if towards > 0 {
			// Ray is moving deeper inside - Pop it out and reflect it
			return Ray2TraceResult::Pop (safe_point, line_normal);
		}
		else {
			// Ray is moving out - Pop it out but consume its timestep
			return Ray2TraceResult::Pop (safe_point, line_normal);
		}
	}

	let extrude_vector = (line.radius) * Vec2::<Fx32>::from (line_normal);
	// Extrude the line
	let line = WideLine {
		start: line.start + extrude_vector,
		end: line.end + extrude_vector,
		radius: Fx32::from_int (0),
	};

	let ray_length = ray.dir.length ();

	if ray_length == 0 {
		return Ray2TraceResult::Miss;
	}

	let basis = get_ray_basis (ray, ray_length);

	// Transform the line to ray space
	let line = WideLine {
		start: basis.to_space (&(line.start - ray.start)),
		end: basis.to_space (&(line.end - ray.start)),
		radius: line.radius,
	};

	// Line will have zero y difference resulting in a divide by zero
	if line.start.y == line.end.y {
		//return Ray2TraceResult::Miss;
	}

	// Apply SAT
	if line.start.y > 0 && line.end.y > 0 {
		return Ray2TraceResult::Miss;
	}
	else if line.start.y < 0 && line.end.y < 0 {
		return Ray2TraceResult::Miss;
	}

	// If we get here, we know that the obstacle line crosses our ray somewhere

	// Note: Lines must not have 0 length
	let obstacle_t = (-line.start.y) / (line.end.y - line.start.y);

	assert! (obstacle_t >= 0);
	assert! (obstacle_t <= 1, format!("{:?}", ray_length));

	// lerp
	let crossing_x = line.start.x * (Fx32::from_int (1) - obstacle_t) + line.end.x * obstacle_t;

	if crossing_x < 0 || crossing_x > ray_length + margin {
		return Ray2TraceResult::Miss;
	}

	// The line segments intersect
	let t = crossing_x / ray_length;

	let safe_point = ray.start + ray.dir * t + big_normal * margin;

	return Ray2TraceResult::Hit (
		t,
		safe_point,
		line_normal,
	);
}

pub fn ray_trace_circle_2 (ray: &Ray2, circle: &Circle) -> Ray2TraceResult {
	let ray_length = ray.dir.length ();
	let basis = get_ray_basis (ray, ray_length);

	let center_in_ray_space = basis.to_space (&(circle.center - ray.start));

	if center_in_ray_space.x < 0 {
		return Ray2TraceResult::Miss;
	}

	if center_in_ray_space.x > ray_length + circle.radius {
		return Ray2TraceResult::Miss;
	}

	if center_in_ray_space.y.abs () >= circle.radius {
		// Prevents a negative root later on
		return Ray2TraceResult::Miss;
	}

	let ray_space_x = center_in_ray_space.x - (circle.radius.square () - center_in_ray_space.y.square ()).sqrt ();

	if ray_space_x > ray_length || ray_length == 0 {
		return Ray2TraceResult::Miss;
	}

	let t = ray_space_x / ray_length;
	let ccd_pos = ray.start + ray.dir * t;

	if t <= 1 {
		return Ray2TraceResult::Hit (
			Fx32 { x: cmp::max (t.x, Fx32::from_int (0).x) },
			ccd_pos,
			((ccd_pos - circle.center) / circle.radius).to_small (),
		);
	}
	else {
		return Ray2TraceResult::Miss;
	}
}
<|endoftext|>


<commit_before>//! An example of using syntect for testing syntax definitions.
//! Basically exactly the same as what Sublime Text can do,
//! but without needing ST installed
// To run tests only for a particular package, while showing the operations, you could use:
// cargo run --example syntest -- --debug testdata/Packages/Makefile/
// to specify that the syntax definitions should be parsed instead of loaded from the dump file,
// you can tell it where to parse them from - the following will execute only 1 syntax test after
// parsing the sublime-syntax files in the JavaScript folder:
// cargo run --example syntest testdata/Packages/JavaScript/syntax_test_json.json testdata/Packages/JavaScript/
extern crate syntect;
extern crate walkdir;
#[macro_use]
extern crate lazy_static;
extern crate regex;
//extern crate onig;
use syntect::parsing::{SyntaxSet, ParseState, ScopeStack, Scope};
use syntect::highlighting::ScopeSelectors;
use syntect::easy::{ScopeRegionIterator};

use std::path::Path;
use std::io::{BufRead, BufReader};
use std::fs::File;
use std::cmp::{min, max};
use walkdir::{DirEntry, WalkDir};
use std::str::FromStr;
use regex::Regex;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SyntaxTestHeaderError {
    MalformedHeader,
    SyntaxDefinitionNotFound,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SyntaxTestFileResult {
    FailedAssertions(usize, usize),
    Success(usize),
}

lazy_static! {
    pub static ref SYNTAX_TEST_HEADER_PATTERN: Regex = Regex::new(r#"(?xm)
            ^(?P<testtoken_start>\s*\S+)
            \s+SYNTAX\sTEST\s+
            "(?P<syntax_file>[^"]+)"
            \s*(?P<testtoken_end>\S+)?$
        "#).unwrap();
    pub static ref SYNTAX_TEST_ASSERTION_PATTERN: Regex = Regex::new(r#"(?xm)
        \s*(?:
            (?P<begin_of_token><-)|(?P<range>\^+)
        )(.*)$"#).unwrap();
}

#[derive(Debug)]
struct AssertionRange<'a> {
    begin_char: usize,
    end_char: usize,
    scope_selector_text: &'a str,
    is_pure_assertion_line: bool,
}

#[derive(Debug)]
struct ScopedText {
    scope: Vec<Scope>,
    char_start: usize,
    text_len: usize,
}

#[derive(Debug)]
struct RangeTestResult {
    column_begin: usize,
    column_end: usize,
    success: bool,
}

fn get_line_assertion_details<'a>(testtoken_start: &str, testtoken_end: Option<&str>, line: &'a str) -> Option<AssertionRange<'a>> {
    // if the test start token specified in the test file's header is on the line
    if let Some(index) = line.find(testtoken_start) {
        let (before_token_start, token_and_rest_of_line) = line.split_at(index);

        if let Some(captures) = SYNTAX_TEST_ASSERTION_PATTERN.captures(&token_and_rest_of_line[testtoken_start.len()..]) {
            let mut sst = captures.get(3).unwrap().as_str(); // get the scope selector text
            let mut only_whitespace_after_token_end = true;

            if let Some(token) = testtoken_end { // if there is an end token defined in the test file header
                if let Some(end_token_pos) = sst.find(token) { // and there is an end token in the line
                    let (ss, after_token_end) = sst.split_at(end_token_pos); // the scope selector text ends at the end token
                    sst = &ss;
                    only_whitespace_after_token_end = after_token_end.trim_right().is_empty();
                }
            }
            return Some(AssertionRange {
                begin_char: index + if captures.get(2).is_some() { testtoken_start.len() + captures.get(2).unwrap().start() } else { 0 },
                end_char: index + if captures.get(2).is_some() { testtoken_start.len() + captures.get(2).unwrap().end() } else { 1 },
                scope_selector_text: sst,
                is_pure_assertion_line: before_token_start.trim_left().is_empty() && only_whitespace_after_token_end, // if only whitespace surrounds the test tokens on the line, then it is a pure assertion line
            });
        }
    }
    None
}

fn process_assertions(assertion: &AssertionRange, test_against_line_scopes: &Vec<ScopedText>) -> Vec<RangeTestResult> {
    // format the scope selector to include a space at the beginning, because, currently, ScopeSelector expects excludes to begin with " -"
    // and they are sometimes in the syntax test as ^^^-comment, for example
    let selector = ScopeSelectors::from_str(&format!(" {}", &assertion.scope_selector_text)).unwrap();
    // find the scope at the specified start column, and start matching the selector through the rest of the tokens on the line from there until the end column is reached
    let mut results = Vec::new();
    for scoped_text in test_against_line_scopes.iter().skip_while(|s|s.char_start + s.text_len <= assertion.begin_char).take_while(|s|s.char_start < assertion.end_char) {
        let match_value = selector.does_match(scoped_text.scope.as_slice());
        let result = RangeTestResult {
            column_begin: max(scoped_text.char_start, assertion.begin_char),
            column_end: min(scoped_text.char_start + scoped_text.text_len, assertion.end_char),
            success: match_value.is_some()
        };
        results.push(result);
    }
    // don't ignore assertions after the newline, they should be treated as though they are asserting against the newline
    let last = test_against_line_scopes.last().unwrap();
    if last.char_start + last.text_len < assertion.end_char {
        let match_value = selector.does_match(last.scope.as_slice());
        let result = RangeTestResult {
            column_begin: max(last.char_start + last.text_len, assertion.begin_char),
            column_end: assertion.end_char,
            success: match_value.is_some()
        };
        results.push(result);
    }
    results
}

/// If `parse_test_lines` is `false` then lines that only contain assertions are not parsed
fn test_file(ss: &SyntaxSet, path: &Path, parse_test_lines: bool, debug: bool) -> Result<SyntaxTestFileResult, SyntaxTestHeaderError> {
    use syntect::util::debug_print_ops;
    let f = File::open(path).unwrap();
    let mut reader = BufReader::new(f);
    let mut line = String::new();

    // read the first line from the file - if we have reached EOF already, it's an invalid file
    if reader.read_line(&mut line).unwrap() == 0 {
        return Err(SyntaxTestHeaderError::MalformedHeader);
    }

    line = line.replace("\r", &"");

    // parse the syntax test header in the first line of the file
    let header_line = line.clone();
    let search_result = SYNTAX_TEST_HEADER_PATTERN.captures(&header_line);
    let captures = search_result.ok_or(SyntaxTestHeaderError::MalformedHeader)?;

    let testtoken_start = captures.name("testtoken_start").unwrap().as_str();
    let testtoken_end = captures.name("testtoken_end").map_or(None, |c|Some(c.as_str()));
    let syntax_file = captures.name("syntax_file").unwrap().as_str();

    // find the relevant syntax definition to parse the file with - case is important!
    println!("The test file references syntax definition file: {}", syntax_file);
    let syntax = ss.find_syntax_by_path(syntax_file).ok_or(SyntaxTestHeaderError::SyntaxDefinitionNotFound)?;

    // iterate over the lines of the file, testing them
    let mut state = ParseState::new(syntax);
    let mut stack = ScopeStack::new();

    let mut current_line_number = 1;
    let mut test_against_line_number = 1;
    let mut scopes_on_line_being_tested = Vec::new();
    let mut previous_non_assertion_line = line.to_string();

    let mut assertion_failures: usize = 0;
    let mut total_assertions: usize = 0;

    loop { // over lines of file, starting with the header line
        let mut line_only_has_assertion = false;
        let mut line_has_assertion = false;
        if let Some(assertion) = get_line_assertion_details(testtoken_start, testtoken_end, &line) {
            let result = process_assertions(&assertion, &scopes_on_line_being_tested);
            total_assertions += &assertion.end_char - &assertion.begin_char;
            for failure in result.iter().filter(|r|!r.success) {
                let length = failure.column_end - failure.column_begin;
                let text: String = previous_non_assertion_line.chars().skip(failure.column_begin).take(length).collect();
                println!("  Assertion selector {:?} \
                    from line {:?} failed against line {:?}, column range {:?}-{:?} \
                    (with text {:?}) \
                    has scope {:?}",
                    assertion.scope_selector_text.trim(),
                    current_line_number, test_against_line_number, failure.column_begin, failure.column_end,
                    text,
                    scopes_on_line_being_tested.iter().skip_while(|s|s.char_start + s.text_len <= failure.column_begin).next().unwrap_or(scopes_on_line_being_tested.last().unwrap()).scope
                );
                assertion_failures += failure.column_end - failure.column_begin;
            }
            line_only_has_assertion = assertion.is_pure_assertion_line;
            line_has_assertion = true;
        }
        if !line_only_has_assertion || parse_test_lines {
            if !line_has_assertion { // ST seems to ignore lines that have assertions when calculating which line the assertion tests against
                scopes_on_line_being_tested.clear();
                test_against_line_number = current_line_number;
                previous_non_assertion_line = line.to_string();
            }
            if debug && !line_only_has_assertion {
                println!("-- debugging line {} -- scope stack: {:?}", current_line_number, stack);
            }
            let ops = state.parse_line(&line);
            if debug && !line_only_has_assertion {
                if ops.is_empty() && !line.is_empty() {
                    println!("no operations for this line...");
                } else {
                    debug_print_ops(&line, &ops);
                }
            }
            let mut col: usize = 0;
            for (s, op) in ScopeRegionIterator::new(&ops, &line) {
                stack.apply(op);
                if s.is_empty() { // in this case we don't care about blank tokens
                    continue;
                }
                if !line_has_assertion {
                    // if the line has no assertions on it, remember the scopes on the line so we can test against them later
                    let len = s.chars().count();
                    scopes_on_line_being_tested.push(
                        ScopedText {
                            char_start: col,
                            text_len: len,
                            scope: stack.as_slice().to_vec()
                        }
                    );
                    // TODO: warn when there are duplicate adjacent (non-meta?) scopes, as it is almost always undesired
                    col += len;
                }
            }
        }

        line.clear();
        current_line_number += 1;
        if reader.read_line(&mut line).unwrap() == 0 {
            break;
        }
        line = line.replace("\r", &"");
    }
    if assertion_failures > 0 {
        Ok(SyntaxTestFileResult::FailedAssertions(assertion_failures, total_assertions))
    } else {
        Ok(SyntaxTestFileResult::Success(total_assertions))
    }
}

fn main() {
    let mut args: Vec<String> = std::env::args().collect();
    let debug_arg = args.iter().position(|s| s == "--debug");
    if debug_arg.is_some() {
        args.remove(debug_arg.unwrap());
    }

    let tests_path = if args.len() < 2 {
        "."
    } else {
        &args[1]
    };
    let syntaxes_path = if args.len() == 3 {
        &args[2]
    } else {
        ""
    };

    // load the syntaxes from disk if told to
    // (as opposed to from the binary dumps)
    // this helps to ensure that a recompile isn't needed
    // when using this for syntax development
    let mut ss = if syntaxes_path.is_empty() {
        SyntaxSet::load_defaults_newlines() // note we load the version with newlines
    } else {
        SyntaxSet::new()
    };
    if !syntaxes_path.is_empty() {
        println!("loading syntax definitions from {}", syntaxes_path);
        ss.load_syntaxes(&syntaxes_path, true).unwrap(); // note that we load the version with newlines
        ss.link_syntaxes();
    }

    let exit_code = recursive_walk(&ss, &tests_path, debug_arg.is_some());
    println!("exiting with code {}", exit_code);
    std::process::exit(exit_code);

}

fn recursive_walk(ss: &SyntaxSet, path: &str, debug: bool) -> i32 {
    let mut exit_code: i32 = 0; // exit with code 0 by default, if all tests pass
    let walker = WalkDir::new(path).into_iter();
    for entry in walker.filter_entry(|e|e.file_type().is_dir() || is_a_syntax_test_file(e)) {
        let entry = entry.unwrap();
        if entry.file_type().is_file() {
            println!("Testing file {}", entry.path().display());
            let result = test_file(&ss, entry.path(), true, debug);
            println!("{:?}", result);
            if exit_code != 2 { // leave exit code 2 if there was an error
                if let Err(_) = result { // set exit code 2 if there was an error
                    exit_code = 2;
                } else if let Ok(ok) = result {
                    if let SyntaxTestFileResult::FailedAssertions(_, _) = ok {
                        exit_code = 1; // otherwise, if there were failures, exit with code 1
                    }
                }
            }
        }
    }
    exit_code
}

fn is_a_syntax_test_file(entry: &DirEntry) -> bool {
    entry.file_name()
         .to_str()
         .map(|s| s.starts_with("syntax_test_"))
         .unwrap_or(false)
}
<commit_msg>Add --time option to syntest<commit_after>//! An example of using syntect for testing syntax definitions.
//! Basically exactly the same as what Sublime Text can do,
//! but without needing ST installed
// To run tests only for a particular package, while showing the operations, you could use:
// cargo run --example syntest -- --debug testdata/Packages/Makefile/
// to specify that the syntax definitions should be parsed instead of loaded from the dump file,
// you can tell it where to parse them from - the following will execute only 1 syntax test after
// parsing the sublime-syntax files in the JavaScript folder:
// cargo run --example syntest testdata/Packages/JavaScript/syntax_test_json.json testdata/Packages/JavaScript/
extern crate syntect;
extern crate walkdir;
#[macro_use]
extern crate lazy_static;
extern crate regex;
//extern crate onig;
use syntect::parsing::{SyntaxSet, ParseState, ScopeStack, Scope};
use syntect::highlighting::ScopeSelectors;
use syntect::easy::{ScopeRegionIterator};

use std::path::Path;
use std::io::{BufRead, BufReader};
use std::fs::File;
use std::cmp::{min, max};
use std::time::Instant;
use walkdir::{DirEntry, WalkDir};
use std::str::FromStr;
use regex::Regex;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SyntaxTestHeaderError {
    MalformedHeader,
    SyntaxDefinitionNotFound,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SyntaxTestFileResult {
    FailedAssertions(usize, usize),
    Success(usize),
}

lazy_static! {
    pub static ref SYNTAX_TEST_HEADER_PATTERN: Regex = Regex::new(r#"(?xm)
            ^(?P<testtoken_start>\s*\S+)
            \s+SYNTAX\sTEST\s+
            "(?P<syntax_file>[^"]+)"
            \s*(?P<testtoken_end>\S+)?$
        "#).unwrap();
    pub static ref SYNTAX_TEST_ASSERTION_PATTERN: Regex = Regex::new(r#"(?xm)
        \s*(?:
            (?P<begin_of_token><-)|(?P<range>\^+)
        )(.*)$"#).unwrap();
}

#[derive(Debug)]
struct AssertionRange<'a> {
    begin_char: usize,
    end_char: usize,
    scope_selector_text: &'a str,
    is_pure_assertion_line: bool,
}

#[derive(Debug)]
struct ScopedText {
    scope: Vec<Scope>,
    char_start: usize,
    text_len: usize,
}

#[derive(Debug)]
struct RangeTestResult {
    column_begin: usize,
    column_end: usize,
    success: bool,
}

fn get_line_assertion_details<'a>(testtoken_start: &str, testtoken_end: Option<&str>, line: &'a str) -> Option<AssertionRange<'a>> {
    // if the test start token specified in the test file's header is on the line
    if let Some(index) = line.find(testtoken_start) {
        let (before_token_start, token_and_rest_of_line) = line.split_at(index);

        if let Some(captures) = SYNTAX_TEST_ASSERTION_PATTERN.captures(&token_and_rest_of_line[testtoken_start.len()..]) {
            let mut sst = captures.get(3).unwrap().as_str(); // get the scope selector text
            let mut only_whitespace_after_token_end = true;

            if let Some(token) = testtoken_end { // if there is an end token defined in the test file header
                if let Some(end_token_pos) = sst.find(token) { // and there is an end token in the line
                    let (ss, after_token_end) = sst.split_at(end_token_pos); // the scope selector text ends at the end token
                    sst = &ss;
                    only_whitespace_after_token_end = after_token_end.trim_right().is_empty();
                }
            }
            return Some(AssertionRange {
                begin_char: index + if captures.get(2).is_some() { testtoken_start.len() + captures.get(2).unwrap().start() } else { 0 },
                end_char: index + if captures.get(2).is_some() { testtoken_start.len() + captures.get(2).unwrap().end() } else { 1 },
                scope_selector_text: sst,
                is_pure_assertion_line: before_token_start.trim_left().is_empty() && only_whitespace_after_token_end, // if only whitespace surrounds the test tokens on the line, then it is a pure assertion line
            });
        }
    }
    None
}

fn process_assertions(assertion: &AssertionRange, test_against_line_scopes: &Vec<ScopedText>) -> Vec<RangeTestResult> {
    // format the scope selector to include a space at the beginning, because, currently, ScopeSelector expects excludes to begin with " -"
    // and they are sometimes in the syntax test as ^^^-comment, for example
    let selector = ScopeSelectors::from_str(&format!(" {}", &assertion.scope_selector_text)).unwrap();
    // find the scope at the specified start column, and start matching the selector through the rest of the tokens on the line from there until the end column is reached
    let mut results = Vec::new();
    for scoped_text in test_against_line_scopes.iter().skip_while(|s|s.char_start + s.text_len <= assertion.begin_char).take_while(|s|s.char_start < assertion.end_char) {
        let match_value = selector.does_match(scoped_text.scope.as_slice());
        let result = RangeTestResult {
            column_begin: max(scoped_text.char_start, assertion.begin_char),
            column_end: min(scoped_text.char_start + scoped_text.text_len, assertion.end_char),
            success: match_value.is_some()
        };
        results.push(result);
    }
    // don't ignore assertions after the newline, they should be treated as though they are asserting against the newline
    let last = test_against_line_scopes.last().unwrap();
    if last.char_start + last.text_len < assertion.end_char {
        let match_value = selector.does_match(last.scope.as_slice());
        let result = RangeTestResult {
            column_begin: max(last.char_start + last.text_len, assertion.begin_char),
            column_end: assertion.end_char,
            success: match_value.is_some()
        };
        results.push(result);
    }
    results
}

/// If `parse_test_lines` is `false` then lines that only contain assertions are not parsed
fn test_file(ss: &SyntaxSet, path: &Path, parse_test_lines: bool, debug: bool) -> Result<SyntaxTestFileResult, SyntaxTestHeaderError> {
    use syntect::util::debug_print_ops;
    let f = File::open(path).unwrap();
    let mut reader = BufReader::new(f);
    let mut line = String::new();

    // read the first line from the file - if we have reached EOF already, it's an invalid file
    if reader.read_line(&mut line).unwrap() == 0 {
        return Err(SyntaxTestHeaderError::MalformedHeader);
    }

    line = line.replace("\r", &"");

    // parse the syntax test header in the first line of the file
    let header_line = line.clone();
    let search_result = SYNTAX_TEST_HEADER_PATTERN.captures(&header_line);
    let captures = search_result.ok_or(SyntaxTestHeaderError::MalformedHeader)?;

    let testtoken_start = captures.name("testtoken_start").unwrap().as_str();
    let testtoken_end = captures.name("testtoken_end").map_or(None, |c|Some(c.as_str()));
    let syntax_file = captures.name("syntax_file").unwrap().as_str();

    // find the relevant syntax definition to parse the file with - case is important!
    println!("The test file references syntax definition file: {}", syntax_file);
    let syntax = ss.find_syntax_by_path(syntax_file).ok_or(SyntaxTestHeaderError::SyntaxDefinitionNotFound)?;

    // iterate over the lines of the file, testing them
    let mut state = ParseState::new(syntax);
    let mut stack = ScopeStack::new();

    let mut current_line_number = 1;
    let mut test_against_line_number = 1;
    let mut scopes_on_line_being_tested = Vec::new();
    let mut previous_non_assertion_line = line.to_string();

    let mut assertion_failures: usize = 0;
    let mut total_assertions: usize = 0;

    loop { // over lines of file, starting with the header line
        let mut line_only_has_assertion = false;
        let mut line_has_assertion = false;
        if let Some(assertion) = get_line_assertion_details(testtoken_start, testtoken_end, &line) {
            let result = process_assertions(&assertion, &scopes_on_line_being_tested);
            total_assertions += &assertion.end_char - &assertion.begin_char;
            for failure in result.iter().filter(|r|!r.success) {
                let length = failure.column_end - failure.column_begin;
                let text: String = previous_non_assertion_line.chars().skip(failure.column_begin).take(length).collect();
                println!("  Assertion selector {:?} \
                    from line {:?} failed against line {:?}, column range {:?}-{:?} \
                    (with text {:?}) \
                    has scope {:?}",
                    assertion.scope_selector_text.trim(),
                    current_line_number, test_against_line_number, failure.column_begin, failure.column_end,
                    text,
                    scopes_on_line_being_tested.iter().skip_while(|s|s.char_start + s.text_len <= failure.column_begin).next().unwrap_or(scopes_on_line_being_tested.last().unwrap()).scope
                );
                assertion_failures += failure.column_end - failure.column_begin;
            }
            line_only_has_assertion = assertion.is_pure_assertion_line;
            line_has_assertion = true;
        }
        if !line_only_has_assertion || parse_test_lines {
            if !line_has_assertion { // ST seems to ignore lines that have assertions when calculating which line the assertion tests against
                scopes_on_line_being_tested.clear();
                test_against_line_number = current_line_number;
                previous_non_assertion_line = line.to_string();
            }
            if debug && !line_only_has_assertion {
                println!("-- debugging line {} -- scope stack: {:?}", current_line_number, stack);
            }
            let ops = state.parse_line(&line);
            if debug && !line_only_has_assertion {
                if ops.is_empty() && !line.is_empty() {
                    println!("no operations for this line...");
                } else {
                    debug_print_ops(&line, &ops);
                }
            }
            let mut col: usize = 0;
            for (s, op) in ScopeRegionIterator::new(&ops, &line) {
                stack.apply(op);
                if s.is_empty() { // in this case we don't care about blank tokens
                    continue;
                }
                if !line_has_assertion {
                    // if the line has no assertions on it, remember the scopes on the line so we can test against them later
                    let len = s.chars().count();
                    scopes_on_line_being_tested.push(
                        ScopedText {
                            char_start: col,
                            text_len: len,
                            scope: stack.as_slice().to_vec()
                        }
                    );
                    // TODO: warn when there are duplicate adjacent (non-meta?) scopes, as it is almost always undesired
                    col += len;
                }
            }
        }

        line.clear();
        current_line_number += 1;
        if reader.read_line(&mut line).unwrap() == 0 {
            break;
        }
        line = line.replace("\r", &"");
    }
    if assertion_failures > 0 {
        Ok(SyntaxTestFileResult::FailedAssertions(assertion_failures, total_assertions))
    } else {
        Ok(SyntaxTestFileResult::Success(total_assertions))
    }
}

fn main() {
    let mut args: Vec<String> = std::env::args().collect();
    let debug_arg = args.iter().position(|s| s == "--debug");
    if debug_arg.is_some() {
        args.remove(debug_arg.unwrap());
    }

    let time_arg = args.iter().position(|s| s == "--time");
    if time_arg.is_some() {
        args.remove(time_arg.unwrap());
    }

    let tests_path = if args.len() < 2 {
        "."
    } else {
        &args[1]
    };
    let syntaxes_path = if args.len() == 3 {
        &args[2]
    } else {
        ""
    };

    // load the syntaxes from disk if told to
    // (as opposed to from the binary dumps)
    // this helps to ensure that a recompile isn't needed
    // when using this for syntax development
    let mut ss = if syntaxes_path.is_empty() {
        SyntaxSet::load_defaults_newlines() // note we load the version with newlines
    } else {
        SyntaxSet::new()
    };
    if !syntaxes_path.is_empty() {
        println!("loading syntax definitions from {}", syntaxes_path);
        ss.load_syntaxes(&syntaxes_path, true).unwrap(); // note that we load the version with newlines
        ss.link_syntaxes();
    }

    let exit_code = recursive_walk(&ss, &tests_path, debug_arg.is_some(), time_arg.is_some());
    println!("exiting with code {}", exit_code);
    std::process::exit(exit_code);

}

fn recursive_walk(ss: &SyntaxSet, path: &str, debug: bool, time: bool) -> i32 {
    let mut exit_code: i32 = 0; // exit with code 0 by default, if all tests pass
    let walker = WalkDir::new(path).into_iter();
    for entry in walker.filter_entry(|e|e.file_type().is_dir() || is_a_syntax_test_file(e)) {
        let entry = entry.unwrap();
        if entry.file_type().is_file() {
            println!("Testing file {}", entry.path().display());
            let start = Instant::now();
            let result = test_file(&ss, entry.path(), true, debug);
            let elapsed = start.elapsed();
            println!("{:?}", result);
            if time {
                let ms = (elapsed.as_secs() * 1_000) + (elapsed.subsec_nanos() / 1_000_000) as u64;
                println!("{} ms for file {}", ms, entry.path().display());
            }
            if exit_code != 2 { // leave exit code 2 if there was an error
                if let Err(_) = result { // set exit code 2 if there was an error
                    exit_code = 2;
                } else if let Ok(ok) = result {
                    if let SyntaxTestFileResult::FailedAssertions(_, _) = ok {
                        exit_code = 1; // otherwise, if there were failures, exit with code 1
                    }
                }
            }
        }
    }
    exit_code
}

fn is_a_syntax_test_file(entry: &DirEntry) -> bool {
    entry.file_name()
         .to_str()
         .map(|s| s.starts_with("syntax_test_"))
         .unwrap_or(false)
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use windows::processor::{create_processor, CounterValue, Processor, Query};

use sys::disk::{new_disk, Disk, DiskType};

use std::collections::HashMap;
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};

use winapi::ctypes::c_void;

use winapi::shared::minwindef::{BYTE, DWORD, MAX_PATH, TRUE};
use winapi::um::fileapi::{
    CreateFileW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, OPEN_EXISTING,
};
use winapi::um::handleapi::CloseHandle;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::sysinfoapi::{
    GetSystemInfo, GetTickCount64, SYSTEM_INFO,
};
use winapi::um::winioctl::{
    DISK_GEOMETRY, IOCTL_STORAGE_QUERY_PROPERTY, IOCTL_DISK_GET_DRIVE_GEOMETRY,
};
use winapi::um::winnt::{
    BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE,
};
use winapi::um::winbase::DRIVE_FIXED;

pub struct KeyHandler {
    pub unique_id: String,
    pub win_key: Vec<u16>,
}

impl KeyHandler {
    pub fn new(unique_id: String, win_key: Vec<u16>) -> KeyHandler {
        KeyHandler {
            unique_id: unique_id,
            win_key: win_key,
        }
    }
}

/*#[allow(non_snake_case)]
#[allow(unused)]
unsafe fn browser() {
    use winapi::um::pdh::{PdhBrowseCountersA, PDH_BROWSE_DLG_CONFIG_A};
    use winapi::shared::winerror::ERROR_SUCCESS;

    let mut BrowseDlgData: PDH_BROWSE_DLG_CONFIG_A = ::std::mem::zeroed();
    let mut CounterPathBuffer: [i8; 255] = ::std::mem::zeroed();
    const PERF_DETAIL_WIZARD: u32 = 400;
    let text = b"Select a counter to monitor.\0";

    BrowseDlgData.set_IncludeInstanceIndex(FALSE as u32);
    BrowseDlgData.set_SingleCounterPerAdd(TRUE as u32);
    BrowseDlgData.set_SingleCounterPerDialog(TRUE as u32);
    BrowseDlgData.set_LocalCountersOnly(FALSE as u32);
    BrowseDlgData.set_WildCardInstances(TRUE as u32);
    BrowseDlgData.set_HideDetailBox(TRUE as u32);
    BrowseDlgData.set_InitializePath(FALSE as u32);
    BrowseDlgData.set_DisableMachineSelection(FALSE as u32);
    BrowseDlgData.set_IncludeCostlyObjects(FALSE as u32);
    BrowseDlgData.set_ShowObjectBrowser(FALSE as u32);
    BrowseDlgData.hWndOwner = ::std::ptr::null_mut();
    BrowseDlgData.szReturnPathBuffer = CounterPathBuffer.as_mut_ptr();
    BrowseDlgData.cchReturnPathLength = 255;
    BrowseDlgData.pCallBack = None;
    BrowseDlgData.dwCallBackArg = 0;
    BrowseDlgData.CallBackStatus = ERROR_SUCCESS as i32;
    BrowseDlgData.dwDefaultDetailLevel = PERF_DETAIL_WIZARD;
    BrowseDlgData.szDialogBoxCaption = text as *const _ as usize as *mut i8;
    let ret = PdhBrowseCountersA(&mut BrowseDlgData as *mut _);
    println!("browser: {:?}", ret);
    for x in CounterPathBuffer.iter() {
        print!("{:?} ", *x);
    }
    println!("");
    for x in 0..256 {
        print!("{:?} ", *BrowseDlgData.szReturnPathBuffer.offset(x));
    }
    println!("");
}*/

pub fn init_processors() -> Vec<Processor> {
    unsafe {
        let mut sys_info: SYSTEM_INFO = zeroed();
        GetSystemInfo(&mut sys_info);
        let mut ret = Vec::with_capacity(sys_info.dwNumberOfProcessors as usize + 1);
        for nb in 0..sys_info.dwNumberOfProcessors {
            ret.push(create_processor(&format!("CPU {}", nb + 1)));
        }
        ret.insert(0, create_processor("Total CPU"));
        ret
    }
}

pub unsafe fn open_drive(drive_name: &[u16], open_rights: DWORD) -> HANDLE {
    CreateFileW(drive_name.as_ptr(),
                open_rights,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                ::std::ptr::null_mut(),
                OPEN_EXISTING,
                0,
                ::std::ptr::null_mut())
}

pub unsafe fn get_drive_size(handle: HANDLE) -> u64 {
    let mut pdg: DISK_GEOMETRY = ::std::mem::zeroed();
    let mut junk = 0;
    let result = DeviceIoControl(handle,
                                 IOCTL_DISK_GET_DRIVE_GEOMETRY,
                                 ::std::ptr::null_mut(),
                                 0,
                                 &mut pdg as *mut DISK_GEOMETRY as *mut c_void,
                                 size_of::<DISK_GEOMETRY>() as DWORD,
                                 &mut junk,
                                 ::std::ptr::null_mut());
    if result == TRUE {
        *pdg.Cylinders.QuadPart() as u64 * pdg.TracksPerCylinder as u64
        * pdg.SectorsPerTrack as u64 * pdg.BytesPerSector as u64
    } else {
        0
    }
}

pub unsafe fn get_disks() -> Vec<Disk> {
    let mut disks = Vec::new();
    let drives = GetLogicalDrives();
    if drives == 0 {
        return disks;
    }
    for x in 0..size_of::<DWORD>() * 8 {
        if (drives >> x) & 1 == 0 {
            continue
        }
        let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];
        if GetDriveTypeW(mount_point.as_ptr()) != DRIVE_FIXED {
            continue
        }
        let mut name = [0u16; MAX_PATH + 1];
        let mut file_system = [0u16; 32];
        if GetVolumeInformationW(mount_point.as_ptr(),
                                 name.as_mut_ptr(),
                                 name.len() as DWORD,
                                 ::std::ptr::null_mut(),
                                 ::std::ptr::null_mut(),
                                 ::std::ptr::null_mut(),
                                 file_system.as_mut_ptr(),
                                 file_system.len() as DWORD) == 0 {
            continue
        }
        let mut pos = 0;
        for x in name.iter() {
            if *x == 0 {
                break
            }
            pos += 1;
        }
        let name = String::from_utf16_lossy(&name[..pos]);
        let name = OsStr::new(&name);

        pos = 0;
        for x in file_system.iter() {
            if *x == 0 {
                break
            }
            pos += 1;
        }
        let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

        let drive_name = [b'\\' as u16, b'\\' as u16, b'.' as u16, b'\\' as u16, b'A' as u16 + x as u16, b':' as u16, 0];
        let handle = open_drive(&drive_name, 0);
        if handle == INVALID_HANDLE_VALUE {
            disks.push(new_disk(name, &mount_point, &file_system, DiskType::Unknown(-1), 0));
            CloseHandle(handle);
            continue
        }
        let disk_size = get_drive_size(handle);
        /*let mut spq_trim: ffi::STORAGE_PROPERTY_QUERY = ::std::mem::zeroed();
        spq_trim.PropertyId = ffi::StorageDeviceTrimProperty;
        spq_trim.QueryType = ffi::PropertyStandardQuery;
        let mut dtd: ffi::DEVICE_TRIM_DESCRIPTOR = ::std::mem::zeroed();*/
        #[allow(non_snake_case)]
        #[repr(C)]
        struct STORAGE_PROPERTY_QUERY {
            PropertyId: i32,
            QueryType: i32,
            AdditionalParameters: [BYTE; 1]
        }
        #[allow(non_snake_case)]
        #[repr(C)]
        struct DEVICE_TRIM_DESCRIPTOR {
            Version: DWORD,
            Size: DWORD,
            TrimEnabled: BOOLEAN,
        }
        let mut spq_trim = STORAGE_PROPERTY_QUERY {
            PropertyId: 8i32,
            QueryType: 0i32,
            AdditionalParameters: [0],
        };
        let mut dtd: DEVICE_TRIM_DESCRIPTOR = ::std::mem::zeroed();

        let mut dw_size = 0;
        if DeviceIoControl(handle,
                           IOCTL_STORAGE_QUERY_PROPERTY,
                           &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                           size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                           &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                           size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                           &mut dw_size,
                           ::std::ptr::null_mut()) == 0 ||
           dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD {
            disks.push(new_disk(name, &mount_point, &file_system, DiskType::Unknown(-1),
                                disk_size));
            CloseHandle(handle);
            continue
        }
        let is_ssd = dtd.TrimEnabled != 0;
        CloseHandle(handle);
        disks.push(new_disk(name, &mount_point, &file_system,
                            if is_ssd { DiskType::SSD } else { DiskType::HDD },
                            disk_size));
    }
    disks
}

#[allow(non_snake_case)]
pub unsafe fn load_symbols() -> HashMap<String, u32> {
    use winapi::um::winreg::{HKEY_PERFORMANCE_DATA, RegQueryValueExA};

    let mut cbCounters = 0;
    let mut dwType = 0;
    let mut ret = HashMap::new();

    let _dwStatus = RegQueryValueExA(HKEY_PERFORMANCE_DATA,
                                     b"Counter 009\0".as_ptr() as *const _,
                                     ::std::ptr::null_mut(),
                                     &mut dwType as *mut i32 as *mut _,
                                     ::std::ptr::null_mut(),
                                     &mut cbCounters as *mut i32 as *mut _);

    let mut lpmszCounters = Vec::with_capacity(cbCounters as usize);
    lpmszCounters.set_len(cbCounters as usize);
    let _dwStatus = RegQueryValueExA(HKEY_PERFORMANCE_DATA,
                                     b"Counter 009\0".as_ptr() as *const _,
                                     ::std::ptr::null_mut(),
                                     &mut dwType as *mut i32 as *mut _,
                                     lpmszCounters.as_mut_ptr(),
                                     &mut cbCounters as *mut i32 as *mut _);

    for (pos, s) in lpmszCounters.split(|x| *x == 0)
                                 .collect::<Vec<_>>()
                                 .chunks(2)
                                 .filter(|&x| x.len() == 2 && !x[0].is_empty() && !x[1].is_empty())
                                 .map(|x| (u32::from_str_radix(String::from_utf8(x[0].to_vec()).unwrap().as_str(), 10).unwrap(), String::from_utf8(x[1].to_vec()).expect("invalid string"))) {
        ret.insert(s, pos as u32);
    }
    ret
}

pub fn get_translation(s: &String, map: &HashMap<String, u32>) -> Option<String> {
    use winapi::um::pdh::PdhLookupPerfNameByIndexW;

    if let Some(index) = map.get(s) {
        let mut size: usize = 0;
        unsafe {
            let _res = PdhLookupPerfNameByIndexW(::std::ptr::null(),
                                                *index,
                                                ::std::ptr::null_mut(),
                                                &mut size as *mut usize as *mut _);
            let mut v = Vec::with_capacity(size);
            v.set_len(size);
            let _res = PdhLookupPerfNameByIndexW(::std::ptr::null(),
                                                *index,
                                                v.as_mut_ptr() as *mut _,
                                                &mut size as *mut usize as *mut _);
            return Some(String::from_utf16(&v[..size - 1]).expect("invalid utf16"));
        }
    }
    None
}

pub fn add_counter(s: String, query: &mut Query, keys: &mut Option<KeyHandler>, counter_name: String,
               value: CounterValue) {
    let mut full = s.encode_utf16().collect::<Vec<_>>();
    full.push(0);
    if query.add_counter(&counter_name, full.clone(), value) {
        *keys = Some(KeyHandler::new(counter_name, full));
    }
}

pub fn get_uptime() -> u64 {
    unsafe { GetTickCount64() / 1000 }
}
<commit_msg>[issue #137] Fix the System::new() panicking on some PCs with Windows 10 (where system performance counter has empty content)<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use windows::processor::{create_processor, CounterValue, Processor, Query};

use sys::disk::{new_disk, Disk, DiskType};

use std::collections::HashMap;
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};

use winapi::ctypes::c_void;

use winapi::shared::minwindef::{BYTE, DWORD, MAX_PATH, TRUE};
use winapi::um::fileapi::{
    CreateFileW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, OPEN_EXISTING,
};
use winapi::um::handleapi::CloseHandle;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::sysinfoapi::{
    GetSystemInfo, GetTickCount64, SYSTEM_INFO,
};
use winapi::um::winioctl::{
    DISK_GEOMETRY, IOCTL_STORAGE_QUERY_PROPERTY, IOCTL_DISK_GET_DRIVE_GEOMETRY,
};
use winapi::um::winnt::{
    BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE,
};
use winapi::um::winbase::DRIVE_FIXED;

pub struct KeyHandler {
    pub unique_id: String,
    pub win_key: Vec<u16>,
}

impl KeyHandler {
    pub fn new(unique_id: String, win_key: Vec<u16>) -> KeyHandler {
        KeyHandler {
            unique_id: unique_id,
            win_key: win_key,
        }
    }
}

/*#[allow(non_snake_case)]
#[allow(unused)]
unsafe fn browser() {
    use winapi::um::pdh::{PdhBrowseCountersA, PDH_BROWSE_DLG_CONFIG_A};
    use winapi::shared::winerror::ERROR_SUCCESS;

    let mut BrowseDlgData: PDH_BROWSE_DLG_CONFIG_A = ::std::mem::zeroed();
    let mut CounterPathBuffer: [i8; 255] = ::std::mem::zeroed();
    const PERF_DETAIL_WIZARD: u32 = 400;
    let text = b"Select a counter to monitor.\0";

    BrowseDlgData.set_IncludeInstanceIndex(FALSE as u32);
    BrowseDlgData.set_SingleCounterPerAdd(TRUE as u32);
    BrowseDlgData.set_SingleCounterPerDialog(TRUE as u32);
    BrowseDlgData.set_LocalCountersOnly(FALSE as u32);
    BrowseDlgData.set_WildCardInstances(TRUE as u32);
    BrowseDlgData.set_HideDetailBox(TRUE as u32);
    BrowseDlgData.set_InitializePath(FALSE as u32);
    BrowseDlgData.set_DisableMachineSelection(FALSE as u32);
    BrowseDlgData.set_IncludeCostlyObjects(FALSE as u32);
    BrowseDlgData.set_ShowObjectBrowser(FALSE as u32);
    BrowseDlgData.hWndOwner = ::std::ptr::null_mut();
    BrowseDlgData.szReturnPathBuffer = CounterPathBuffer.as_mut_ptr();
    BrowseDlgData.cchReturnPathLength = 255;
    BrowseDlgData.pCallBack = None;
    BrowseDlgData.dwCallBackArg = 0;
    BrowseDlgData.CallBackStatus = ERROR_SUCCESS as i32;
    BrowseDlgData.dwDefaultDetailLevel = PERF_DETAIL_WIZARD;
    BrowseDlgData.szDialogBoxCaption = text as *const _ as usize as *mut i8;
    let ret = PdhBrowseCountersA(&mut BrowseDlgData as *mut _);
    println!("browser: {:?}", ret);
    for x in CounterPathBuffer.iter() {
        print!("{:?} ", *x);
    }
    println!("");
    for x in 0..256 {
        print!("{:?} ", *BrowseDlgData.szReturnPathBuffer.offset(x));
    }
    println!("");
}*/

pub fn init_processors() -> Vec<Processor> {
    unsafe {
        let mut sys_info: SYSTEM_INFO = zeroed();
        GetSystemInfo(&mut sys_info);
        let mut ret = Vec::with_capacity(sys_info.dwNumberOfProcessors as usize + 1);
        for nb in 0..sys_info.dwNumberOfProcessors {
            ret.push(create_processor(&format!("CPU {}", nb + 1)));
        }
        ret.insert(0, create_processor("Total CPU"));
        ret
    }
}

pub unsafe fn open_drive(drive_name: &[u16], open_rights: DWORD) -> HANDLE {
    CreateFileW(drive_name.as_ptr(),
                open_rights,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                ::std::ptr::null_mut(),
                OPEN_EXISTING,
                0,
                ::std::ptr::null_mut())
}

pub unsafe fn get_drive_size(handle: HANDLE) -> u64 {
    let mut pdg: DISK_GEOMETRY = ::std::mem::zeroed();
    let mut junk = 0;
    let result = DeviceIoControl(handle,
                                 IOCTL_DISK_GET_DRIVE_GEOMETRY,
                                 ::std::ptr::null_mut(),
                                 0,
                                 &mut pdg as *mut DISK_GEOMETRY as *mut c_void,
                                 size_of::<DISK_GEOMETRY>() as DWORD,
                                 &mut junk,
                                 ::std::ptr::null_mut());
    if result == TRUE {
        *pdg.Cylinders.QuadPart() as u64 * pdg.TracksPerCylinder as u64
        * pdg.SectorsPerTrack as u64 * pdg.BytesPerSector as u64
    } else {
        0
    }
}

pub unsafe fn get_disks() -> Vec<Disk> {
    let mut disks = Vec::new();
    let drives = GetLogicalDrives();
    if drives == 0 {
        return disks;
    }
    for x in 0..size_of::<DWORD>() * 8 {
        if (drives >> x) & 1 == 0 {
            continue
        }
        let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];
        if GetDriveTypeW(mount_point.as_ptr()) != DRIVE_FIXED {
            continue
        }
        let mut name = [0u16; MAX_PATH + 1];
        let mut file_system = [0u16; 32];
        if GetVolumeInformationW(mount_point.as_ptr(),
                                 name.as_mut_ptr(),
                                 name.len() as DWORD,
                                 ::std::ptr::null_mut(),
                                 ::std::ptr::null_mut(),
                                 ::std::ptr::null_mut(),
                                 file_system.as_mut_ptr(),
                                 file_system.len() as DWORD) == 0 {
            continue
        }
        let mut pos = 0;
        for x in name.iter() {
            if *x == 0 {
                break
            }
            pos += 1;
        }
        let name = String::from_utf16_lossy(&name[..pos]);
        let name = OsStr::new(&name);

        pos = 0;
        for x in file_system.iter() {
            if *x == 0 {
                break
            }
            pos += 1;
        }
        let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

        let drive_name = [b'\\' as u16, b'\\' as u16, b'.' as u16, b'\\' as u16, b'A' as u16 + x as u16, b':' as u16, 0];
        let handle = open_drive(&drive_name, 0);
        if handle == INVALID_HANDLE_VALUE {
            disks.push(new_disk(name, &mount_point, &file_system, DiskType::Unknown(-1), 0));
            CloseHandle(handle);
            continue
        }
        let disk_size = get_drive_size(handle);
        /*let mut spq_trim: ffi::STORAGE_PROPERTY_QUERY = ::std::mem::zeroed();
        spq_trim.PropertyId = ffi::StorageDeviceTrimProperty;
        spq_trim.QueryType = ffi::PropertyStandardQuery;
        let mut dtd: ffi::DEVICE_TRIM_DESCRIPTOR = ::std::mem::zeroed();*/
        #[allow(non_snake_case)]
        #[repr(C)]
        struct STORAGE_PROPERTY_QUERY {
            PropertyId: i32,
            QueryType: i32,
            AdditionalParameters: [BYTE; 1]
        }
        #[allow(non_snake_case)]
        #[repr(C)]
        struct DEVICE_TRIM_DESCRIPTOR {
            Version: DWORD,
            Size: DWORD,
            TrimEnabled: BOOLEAN,
        }
        let mut spq_trim = STORAGE_PROPERTY_QUERY {
            PropertyId: 8i32,
            QueryType: 0i32,
            AdditionalParameters: [0],
        };
        let mut dtd: DEVICE_TRIM_DESCRIPTOR = ::std::mem::zeroed();

        let mut dw_size = 0;
        if DeviceIoControl(handle,
                           IOCTL_STORAGE_QUERY_PROPERTY,
                           &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                           size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                           &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                           size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                           &mut dw_size,
                           ::std::ptr::null_mut()) == 0 ||
           dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD {
            disks.push(new_disk(name, &mount_point, &file_system, DiskType::Unknown(-1),
                                disk_size));
            CloseHandle(handle);
            continue
        }
        let is_ssd = dtd.TrimEnabled != 0;
        CloseHandle(handle);
        disks.push(new_disk(name, &mount_point, &file_system,
                            if is_ssd { DiskType::SSD } else { DiskType::HDD },
                            disk_size));
    }
    disks
}

#[allow(non_snake_case)]
pub unsafe fn load_symbols() -> HashMap<String, u32> {
    use winapi::um::winreg::{HKEY_PERFORMANCE_DATA, RegQueryValueExA};

    let mut cbCounters = 0;
    let mut dwType = 0;
    let mut ret = HashMap::new();

    let _dwStatus = RegQueryValueExA(HKEY_PERFORMANCE_DATA,
                                     b"Counter 009\0".as_ptr() as *const _,
                                     ::std::ptr::null_mut(),
                                     &mut dwType as *mut i32 as *mut _,
                                     ::std::ptr::null_mut(),
                                     &mut cbCounters as *mut i32 as *mut _);

    let mut lpmszCounters = Vec::with_capacity(cbCounters as usize);
    lpmszCounters.set_len(cbCounters as usize);
    let _dwStatus = RegQueryValueExA(HKEY_PERFORMANCE_DATA,
                                     b"Counter 009\0".as_ptr() as *const _,
                                     ::std::ptr::null_mut(),
                                     &mut dwType as *mut i32 as *mut _,
                                     lpmszCounters.as_mut_ptr(),
                                     &mut cbCounters as *mut i32 as *mut _);

    for (pos, s) in lpmszCounters.split(|x| *x == 0)
                                 .collect::<Vec<_>>()
                                 .chunks(2)
                                 .filter(|&x| x.len() == 2 && !x[0].is_empty() && !x[1].is_empty())
                                 .map(|x| (u32::from_str_radix(String::from_utf8(x[0].to_vec()).unwrap().as_str(), 10).unwrap(), String::from_utf8(x[1].to_vec()).expect("invalid string"))) {
        ret.insert(s, pos as u32);
    }
    ret
}

pub fn get_translation(s: &String, map: &HashMap<String, u32>) -> Option<String> {
    use winapi::um::pdh::PdhLookupPerfNameByIndexW;

    if let Some(index) = map.get(s) {
        let mut size: usize = 0;
        unsafe {
            let _res = PdhLookupPerfNameByIndexW(::std::ptr::null(),
                                                *index,
                                                ::std::ptr::null_mut(),
                                                &mut size as *mut usize as *mut _);
            if size == 0 {
                return Some(String::new());
            }
            else {
                let mut v = Vec::with_capacity(size);
                v.set_len(size);
                let _res = PdhLookupPerfNameByIndexW(::std::ptr::null(),
                                                    *index,
                                                    v.as_mut_ptr() as *mut _,
                                                    &mut size as *mut usize as *mut _);
                return Some(String::from_utf16(&v[..size - 1]).expect("invalid utf16"));
            }
        }
    }
    None
}

pub fn add_counter(s: String, query: &mut Query, keys: &mut Option<KeyHandler>, counter_name: String,
               value: CounterValue) {
    let mut full = s.encode_utf16().collect::<Vec<_>>();
    full.push(0);
    if query.add_counter(&counter_name, full.clone(), value) {
        *keys = Some(KeyHandler::new(counter_name, full));
    }
}

pub fn get_uptime() -> u64 {
    unsafe { GetTickCount64() / 1000 }
}
<|endoftext|>


<commit_before>use std::collections::HashSet;
use std::fs::{read_dir, read_link, remove_file, metadata};
use std::io::{self, stdout, stderr};
use std::ffi::OsStr;
use std::path::Path;
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH, Duration};
use std::os::unix::fs::MetadataExt;

use argparse::{ArgumentParser, PushConst, StoreTrue, StoreOption};
use scan_dir::ScanDir;
use humantime;

use super::setup;
use super::Wrapper;
use container::util::clean_dir;
use file_util::{read_visible_entries, Lock};
use wrapper::build::get_version_hash;

#[derive(Clone, Copy)]
enum Action {
    Temporary,
    Old,
    Unused,
    Everything,
    Transient,
}

pub fn clean_cmd(wrapper: &Wrapper, cmdline: Vec<String>)
    -> Result<i32, String>
{
    let mut global = false;
    let mut dry_run = false;
    let mut actions = vec!();
    let mut duration = None::<humantime::Duration>;
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Performs various cleanup tasks
            ");
        ap.refer(&mut actions)
          .add_option(&["--tmp", "--tmp-folders"],
                PushConst(Action::Temporary),
                "Clean temporary containers (failed builds)")
          .add_option(&["--old", "--old-containers"], PushConst(Action::Old), "
                Clean old versions of containers (those which doesn't have a \
                symlink in .vagga)")
          .add_option(&["--unused"], PushConst(Action::Unused), "
                Clean unused containers, or versions thereof. (This is not \
                `--old` for historical reasons, we will probably merge the \
                commands later on)")
          .add_option(&["--transient"], PushConst(Action::Transient),
                "Clean unneeded transient folders (left from containers with
                 `write-mode` set to transient-something). The pid of process
                 is checked for liveness first.")
          .add_option(&["--everything"], PushConst(Action::Everything),
                "Clean whole `.vagga` folder. Useful when deleting a project.
                 With ``--global`` cleans whole storage-dir and cache-dir")
          .required();
        ap.refer(&mut global)
          .add_option(&["--global"], StoreTrue,
                "Apply cleanup command to all containers. Works only \
                if `storage-dir` is configured in settings");
        ap.refer(&mut dry_run)
          .add_option(&["-n", "--dry-run"], StoreTrue,
                "Dry run. Don't delete everything, just print");
        ap.refer(&mut duration)
          .add_option(&["--at-least"], StoreOption, "
            Only in combination with `--unused`. Treat as unused \
            containers that are unused for specified time, rather than \
            the ones not used by current version of config");
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(1);
            }
        }
    }
    if global && wrapper.ext_settings.storage_dir.is_none() {
        error!("The --global flag is only meaningful if you configure \
            storage-dir in settings");
        return Ok(2);
    }
    let duration = duration.map(|x| x.into());
    for action in actions.iter() {
        let res = match *action {
            Action::Temporary => clean_temporary(wrapper, global, dry_run),
            Action::Old => clean_old(wrapper, global, dry_run),
            Action::Unused => {
                clean_unused(wrapper, global, duration, dry_run)
            }
            Action::Transient => clean_transient(wrapper, global, dry_run),
            Action::Everything => clean_everything(wrapper, global, dry_run),
        };
        match res {
            Ok(()) => {}
            Err(err) => {
                error!("Error cleaning up: {}", err);
                return Ok(3);
            }
        }
    }
    return Ok(0);
}

fn clean_dir_wrapper(path: &Path,
        remove_dir_itself: bool, dry_run: bool) -> Result<(), String> {
    // TODO(tailhook) chroot to dir for removing
    if dry_run {
        println!("Would remove {:?}", path);
    } else {
        let mut n = path.to_path_buf().into_os_string();
        n.push(".lock");
        let lock_name = Path::new(&n);
        let lock_guard = if lock_name.exists() {
            match Lock::exclusive(&lock_name) {
                Ok(x) => Some(x),
                Err(e) => {
                    error!("Failed to lock {:?}: {}, skipping", lock_name, e);
                    return Ok(());
                }
            }
        } else {
            None
        };
        debug!("Removing {:?}", path);
        try!(clean_dir(path, remove_dir_itself));
        if let Some(_lock) = lock_guard {
            try!(remove_file(lock_name)
                .map_err(|e| format!("Error removing lock file {:?}: {}",
                    lock_name, e)));
        }
    }
    Ok(())
}

fn clean_everything(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        if let Some(ref cache_dir) = wrapper.ext_settings.cache_dir {
            try!(clean_dir_wrapper(&cache_dir, false, dry_run));
        }
        if let Some(ref storage_dir) = wrapper.ext_settings.storage_dir {
            try!(clean_dir_wrapper(&storage_dir, false, dry_run));
        }
    } else {
        let base = match try!(setup::get_vagga_base(
            wrapper.project_root, wrapper.ext_settings))
        {
            Some(base) => base,
            None => {
                warn!("No vagga directory exists");
                return Ok(());
            }
        };
        try!(clean_dir_wrapper(&base, true, dry_run));
        let inner = wrapper.project_root.join(".vagga");
        if base != inner {
            try!(clean_dir_wrapper(&inner, true, dry_run));
        }

    }
    return Ok(());
}

fn clean_temporary(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }
    let base = match try!(setup::get_vagga_base(
        wrapper.project_root, wrapper.ext_settings))
    {
        Some(base) => base,
        None => {
            warn!("No vagga directory exists");
            return Ok(());
        }
    };
    let roots = base.join(".roots");
    for entry in try_msg!(read_dir(&roots),
        "Can't read dir {r:?}: {err}", r=roots)
    {
        let entry = try_msg!(entry, "Can't read dir {r:?}: {err}", r=roots);
        let typ = try_msg!(entry.file_type(),
            "Can't stat {p:?}: {err}", p=entry.path());
        if typ.is_dir() &&
           entry.file_name()[..].to_str().map(|n| n.starts_with(".tmp"))
                                         .unwrap_or(false)
        {
            try!(clean_dir_wrapper(&entry.path(), true, dry_run));
        }
    }

    return Ok(());
}

fn clean_old(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }
    let base = match try!(setup::get_vagga_base(
        wrapper.project_root, wrapper.ext_settings))
    {
        Some(base) => base,
        None => {
            warn!("No vagga directory exists");
            return Ok(());
        }
    };
    let useful: HashSet<String> = try_msg!(
        read_visible_entries(&wrapper.project_root.join(".vagga")),
            "Can't read vagga directory: {err}")
        .into_iter()
        .filter_map(|path| read_link(&path)
             .map_err(|e| warn!("Can't readlink {:?}: {}", path, e))
             .ok()
             .and_then(|f| {
                 // The container name is next to the last component
                 f.iter().rev().nth(1)
                 .and_then(|x| x.to_str()).map(ToString::to_string)
             }))
        .collect();

    info!("Useful images {:?}", useful);
    try!(clean_dirs_except(&base.join(".roots"), &useful, dry_run));

    return Ok(());
}

fn clean_transient(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }
    let base = match try!(setup::get_vagga_base(
        wrapper.project_root, wrapper.ext_settings))
    {
        Some(base) => base,
        None => {
            warn!("No vagga directory exists");
            return Ok(());
        }
    };
    let procfs = Path::new("/proc");
    for entry in try_msg!(read_dir(&base.join(".transient")),
                        "Can't read .vagga/.transient dir: {err}")
    {
        let entry = try_msg!(entry, "Error reading .vagga/transient: {err}");
        if let Some(fname) = entry.file_name()[..].to_str() {
            if let Some(idx) = fname.find('.') {
                if u32::from_str(&fname[idx+1..]).is_ok() &&
                    procfs.join(&fname[idx+1..]).exists()
                {
                    continue;
                }
            }
        }
        try!(clean_dir_wrapper(&entry.path(), true, dry_run));
    }

    return Ok(());
}

fn clean_dirs_except<P: AsRef<Path>>(roots: P, useful: &HashSet<String>,
    dry_run: bool)
    -> Result<(), String>
{
    let roots = roots.as_ref();
    for entry in try_msg!(read_dir(&roots),
                         "Can't read dir {dir:?}: {err}", dir=roots)
    {
        let entry = try_msg!(entry,
                             "Can't read dir {dir:?}: {err}", dir=roots);
        let path = entry.path();
        let typ = try_msg!(entry.file_type(),
            "Can't stat {p:?}: {err}", p=path);
        if !typ.is_dir() {
            if path.extension() == Some(OsStr::new("lock")) &&
               path.with_extension("").is_dir()
            {
                debug!("Skipping lock file {:?}", path);
            } else {
                try_msg!(remove_file(&path),
                    "Can't remove file {p:?}: {err}", p=path);
            }
        } else if !typ.is_dir() || entry.file_name()[..].to_str()
            .map(|n| !useful.contains(&n.to_string()))
            .unwrap_or(false)
        {
            try!(clean_dir_wrapper(&entry.path(), true, dry_run));
        }
    }
    Ok(())
}

fn clean_unused(wrapper: &Wrapper, global: bool, duration: Option<Duration>,
    dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }

    try!(setup::setup_base_filesystem(
        wrapper.project_root, wrapper.ext_settings));

    let mut useful: HashSet<String> = HashSet::new();
    if let Some(duration) = duration {
        let unixtime = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
        let cut_off = (unixtime - duration).as_secs() as i64;
        ScanDir::dirs().skip_hidden(false).read("/vagga/base/.roots", |iter| {
            for (entry, name) in iter {
                let luse_path = entry.path().join("last_use");
                match metadata(&luse_path) {
                    Ok(ref meta) if meta.mtime() > cut_off => {
                        useful.insert(name);
                    }
                    Ok(_) => {}
                    Err(ref e) if e.kind() == io::ErrorKind::NotFound => {}
                    Err(e) => {
                        error!("Error trying to stat {:?}: {}", luse_path, e);
                    }
                }
            }
        }).map_err(|e| {
            error!("Error reading `.vagga/.roots`: {}", e);
        }).ok();
    } else {
        for (name, _) in &wrapper.config.containers {
            if let Some(version) = try!(get_version_hash(name, wrapper)) {
                useful.insert(format!("{}.{}", name, &version[..8]));
            }
        }
    }
    info!("Useful images {:?}", useful);
    try!(clean_dirs_except("/vagga/base/.roots", &useful, dry_run));

    return Ok(());
}
<commit_msg>Fix race condition in parallel clean (fixes #302)<commit_after>use std::collections::HashSet;
use std::fs::{read_dir, read_link, remove_file, metadata};
use std::io::{self, stdout, stderr};
use std::ffi::OsStr;
use std::path::Path;
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH, Duration};
use std::os::unix::fs::MetadataExt;

use argparse::{ArgumentParser, PushConst, StoreTrue, StoreOption};
use scan_dir::ScanDir;
use humantime;

use super::setup;
use super::Wrapper;
use container::util::clean_dir;
use file_util::{read_visible_entries, Lock};
use wrapper::build::get_version_hash;

#[derive(Clone, Copy)]
enum Action {
    Temporary,
    Old,
    Unused,
    Everything,
    Transient,
}

pub fn clean_cmd(wrapper: &Wrapper, cmdline: Vec<String>)
    -> Result<i32, String>
{
    let mut global = false;
    let mut dry_run = false;
    let mut actions = vec!();
    let mut duration = None::<humantime::Duration>;
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Performs various cleanup tasks
            ");
        ap.refer(&mut actions)
          .add_option(&["--tmp", "--tmp-folders"],
                PushConst(Action::Temporary),
                "Clean temporary containers (failed builds)")
          .add_option(&["--old", "--old-containers"], PushConst(Action::Old), "
                Clean old versions of containers (those which doesn't have a \
                symlink in .vagga)")
          .add_option(&["--unused"], PushConst(Action::Unused), "
                Clean unused containers, or versions thereof. (This is not \
                `--old` for historical reasons, we will probably merge the \
                commands later on)")
          .add_option(&["--transient"], PushConst(Action::Transient),
                "Clean unneeded transient folders (left from containers with
                 `write-mode` set to transient-something). The pid of process
                 is checked for liveness first.")
          .add_option(&["--everything"], PushConst(Action::Everything),
                "Clean whole `.vagga` folder. Useful when deleting a project.
                 With ``--global`` cleans whole storage-dir and cache-dir")
          .required();
        ap.refer(&mut global)
          .add_option(&["--global"], StoreTrue,
                "Apply cleanup command to all containers. Works only \
                if `storage-dir` is configured in settings");
        ap.refer(&mut dry_run)
          .add_option(&["-n", "--dry-run"], StoreTrue,
                "Dry run. Don't delete everything, just print");
        ap.refer(&mut duration)
          .add_option(&["--at-least"], StoreOption, "
            Only in combination with `--unused`. Treat as unused \
            containers that are unused for specified time, rather than \
            the ones not used by current version of config");
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(1);
            }
        }
    }
    if global && wrapper.ext_settings.storage_dir.is_none() {
        error!("The --global flag is only meaningful if you configure \
            storage-dir in settings");
        return Ok(2);
    }
    let duration = duration.map(|x| x.into());
    for action in actions.iter() {
        let res = match *action {
            Action::Temporary => clean_temporary(wrapper, global, dry_run),
            Action::Old => clean_old(wrapper, global, dry_run),
            Action::Unused => {
                clean_unused(wrapper, global, duration, dry_run)
            }
            Action::Transient => clean_transient(wrapper, global, dry_run),
            Action::Everything => clean_everything(wrapper, global, dry_run),
        };
        match res {
            Ok(()) => {}
            Err(err) => {
                error!("Error cleaning up: {}", err);
                return Ok(3);
            }
        }
    }
    return Ok(0);
}

fn clean_dir_wrapper(path: &Path,
        remove_dir_itself: bool, dry_run: bool) -> Result<(), String> {
    // TODO(tailhook) chroot to dir for removing
    if dry_run {
        println!("Would remove {:?}", path);
    } else {
        let mut n = path.to_path_buf().into_os_string();
        n.push(".lock");
        let lock_name = Path::new(&n);
        let lock_guard = if lock_name.exists() {
            match Lock::exclusive(&lock_name) {
                Ok(x) => Some(x),
                Err(e) => {
                    error!("Failed to lock {:?}: {}, skipping", lock_name, e);
                    return Ok(());
                }
            }
        } else {
            None
        };
        debug!("Removing {:?}", path);
        try!(clean_dir(path, remove_dir_itself));
        if let Some(_lock) = lock_guard {
            try!(remove_file(lock_name)
                .map_err(|e| format!("Error removing lock file {:?}: {}",
                    lock_name, e)));
        }
    }
    Ok(())
}

fn clean_everything(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        if let Some(ref cache_dir) = wrapper.ext_settings.cache_dir {
            try!(clean_dir_wrapper(&cache_dir, false, dry_run));
        }
        if let Some(ref storage_dir) = wrapper.ext_settings.storage_dir {
            try!(clean_dir_wrapper(&storage_dir, false, dry_run));
        }
    } else {
        let base = match try!(setup::get_vagga_base(
            wrapper.project_root, wrapper.ext_settings))
        {
            Some(base) => base,
            None => {
                warn!("No vagga directory exists");
                return Ok(());
            }
        };
        try!(clean_dir_wrapper(&base, true, dry_run));
        let inner = wrapper.project_root.join(".vagga");
        if base != inner {
            try!(clean_dir_wrapper(&inner, true, dry_run));
        }

    }
    return Ok(());
}

fn clean_temporary(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }
    let base = match try!(setup::get_vagga_base(
        wrapper.project_root, wrapper.ext_settings))
    {
        Some(base) => base,
        None => {
            warn!("No vagga directory exists");
            return Ok(());
        }
    };
    let roots = base.join(".roots");
    for entry in try_msg!(read_dir(&roots),
        "Can't read dir {r:?}: {err}", r=roots)
    {
        let entry = try_msg!(entry, "Can't read dir {r:?}: {err}", r=roots);
        let typ = try_msg!(entry.file_type(),
            "Can't stat {p:?}: {err}", p=entry.path());
        if typ.is_dir() &&
           entry.file_name()[..].to_str().map(|n| n.starts_with(".tmp"))
                                         .unwrap_or(false)
        {
            try!(clean_dir_wrapper(&entry.path(), true, dry_run));
        }
    }

    return Ok(());
}

fn clean_old(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }
    let base = match try!(setup::get_vagga_base(
        wrapper.project_root, wrapper.ext_settings))
    {
        Some(base) => base,
        None => {
            warn!("No vagga directory exists");
            return Ok(());
        }
    };
    let useful: HashSet<String> = try_msg!(
        read_visible_entries(&wrapper.project_root.join(".vagga")),
            "Can't read vagga directory: {err}")
        .into_iter()
        .filter_map(|path| read_link(&path)
             .map_err(|e| warn!("Can't readlink {:?}: {}", path, e))
             .ok()
             .and_then(|f| {
                 // The container name is next to the last component
                 f.iter().rev().nth(1)
                 .and_then(|x| x.to_str()).map(ToString::to_string)
             }))
        .collect();

    info!("Useful images {:?}", useful);
    try!(clean_dirs_except(&base.join(".roots"), &useful, dry_run));

    return Ok(());
}

fn clean_transient(wrapper: &Wrapper, global: bool, dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }
    let base = match try!(setup::get_vagga_base(
        wrapper.project_root, wrapper.ext_settings))
    {
        Some(base) => base,
        None => {
            warn!("No vagga directory exists");
            return Ok(());
        }
    };
    let procfs = Path::new("/proc");
    for entry in try_msg!(read_dir(&base.join(".transient")),
                        "Can't read .vagga/.transient dir: {err}")
    {
        let entry = try_msg!(entry, "Error reading .vagga/transient: {err}");
        if let Some(fname) = entry.file_name()[..].to_str() {
            if let Some(idx) = fname.find('.') {
                if u32::from_str(&fname[idx+1..]).is_ok() &&
                    procfs.join(&fname[idx+1..]).exists()
                {
                    continue;
                }
            }
        }
        try!(clean_dir_wrapper(&entry.path(), true, dry_run));
    }

    return Ok(());
}

fn clean_dirs_except<P: AsRef<Path>>(roots: P, useful: &HashSet<String>,
    dry_run: bool)
    -> Result<(), String>
{
    let roots = roots.as_ref();
    for entry in try_msg!(read_dir(&roots),
                         "Can't read dir {dir:?}: {err}", dir=roots)
    {
        let entry = try_msg!(entry,
                             "Can't read dir {dir:?}: {err}", dir=roots);
        let path = entry.path();
        let typ = try_msg!(entry.file_type(),
            "Can't stat {p:?}: {err}", p=path);
        if !typ.is_dir() {
            if path.extension() == Some(OsStr::new("lock")) &&
               path.with_extension("").is_dir()
            {
                debug!("Skipping lock file {:?}", path);
            } else {
                match remove_file(&path) {
                    Ok(()) => {}
                    // File is deleted while we were scanning
                    Err(ref e) if e.kind() == io::ErrorKind::NotFound => {}
                    Err(ref e) => {
                        return Err(
                            format!("Can't remove file {:?}: {}", path, e));
                    }
                }
            }
        } else if !typ.is_dir() || entry.file_name()[..].to_str()
            .map(|n| !useful.contains(&n.to_string()))
            .unwrap_or(false)
        {
            try!(clean_dir_wrapper(&entry.path(), true, dry_run));
        }
    }
    Ok(())
}

fn clean_unused(wrapper: &Wrapper, global: bool, duration: Option<Duration>,
    dry_run: bool)
    -> Result<(), String>
{
    if global {
        panic!("Global cleanup is not implemented yet");
    }

    try!(setup::setup_base_filesystem(
        wrapper.project_root, wrapper.ext_settings));

    let mut useful: HashSet<String> = HashSet::new();
    if let Some(duration) = duration {
        let unixtime = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
        let cut_off = (unixtime - duration).as_secs() as i64;
        ScanDir::dirs().skip_hidden(false).read("/vagga/base/.roots", |iter| {
            for (entry, name) in iter {
                let luse_path = entry.path().join("last_use");
                match metadata(&luse_path) {
                    Ok(ref meta) if meta.mtime() > cut_off => {
                        useful.insert(name);
                    }
                    Ok(_) => {}
                    Err(ref e) if e.kind() == io::ErrorKind::NotFound => {}
                    Err(e) => {
                        error!("Error trying to stat {:?}: {}", luse_path, e);
                    }
                }
            }
        }).map_err(|e| {
            error!("Error reading `.vagga/.roots`: {}", e);
        }).ok();
    } else {
        for (name, _) in &wrapper.config.containers {
            if let Some(version) = try!(get_version_hash(name, wrapper)) {
                useful.insert(format!("{}.{}", name, &version[..8]));
            }
        }
    }
    info!("Useful images {:?}", useful);
    try!(clean_dirs_except("/vagga/base/.roots", &useful, dry_run));

    return Ok(());
}
<|endoftext|>


<commit_before>use rustc_index::bit_set::BitSet;
use rustc_middle::mir::visit::Visitor;
use rustc_middle::mir::{self, Local, Location};

/// The set of locals in a MIR body that do not have `StorageLive`/`StorageDead` annotations.
///
/// These locals have fixed storage for the duration of the body.
//
// FIXME: Currently, we need to traverse the entire MIR to compute this. We should instead store it
// as a field in the `LocalDecl` for each `Local`.
#[derive(Debug, Clone)]
pub struct AlwaysLiveLocals(BitSet<Local>);

impl AlwaysLiveLocals {
    pub fn new(body: &mir::Body<'tcx>) -> Self {
        let mut ret = AlwaysLiveLocals(BitSet::new_filled(body.local_decls.len()));

        let mut vis = StorageAnnotationVisitor(&mut ret);
        vis.visit_body(body);

        ret
    }

    pub fn into_inner(self) -> BitSet<Local> {
        self.0
    }
}

impl std::ops::Deref for AlwaysLiveLocals {
    type Target = BitSet<Local>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Removes locals that have `Storage*` annotations from `AlwaysLiveLocals`.
struct StorageAnnotationVisitor<'a>(&'a mut AlwaysLiveLocals);

impl Visitor<'tcx> for StorageAnnotationVisitor<'_> {
    fn visit_statement(&mut self, statement: &mir::Statement<'tcx>, _location: Location) {
        use mir::StatementKind::{StorageDead, StorageLive};
        if let StorageLive(l) | StorageDead(l) = statement.kind {
            (self.0).0.remove(l);
        }
    }
}
<commit_msg>Visit only statements in always live locals<commit_after>use rustc_index::bit_set::BitSet;
use rustc_middle::mir::{self, Local};

/// The set of locals in a MIR body that do not have `StorageLive`/`StorageDead` annotations.
///
/// These locals have fixed storage for the duration of the body.
//
// FIXME: Currently, we need to traverse the entire MIR to compute this. We should instead store it
// as a field in the `LocalDecl` for each `Local`.
#[derive(Debug, Clone)]
pub struct AlwaysLiveLocals(BitSet<Local>);

impl AlwaysLiveLocals {
    pub fn new(body: &mir::Body<'tcx>) -> Self {
        let mut always_live_locals = AlwaysLiveLocals(BitSet::new_filled(body.local_decls.len()));

        for block in body.basic_blocks() {
            for statement in &block.statements {
                use mir::StatementKind::{StorageDead, StorageLive};
                if let StorageLive(l) | StorageDead(l) = statement.kind {
                    always_live_locals.0.remove(l);
                }
            }
        }

        always_live_locals
    }

    pub fn into_inner(self) -> BitSet<Local> {
        self.0
    }
}

impl std::ops::Deref for AlwaysLiveLocals {
    type Target = BitSet<Local>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
<|endoftext|>