ED_LRR/rust/src/galaxy.rs

114 lines
2.9 KiB
Rust

use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter};
use std::str;
#[derive(Debug, Clone, Serialize)]
pub struct SystemSerde {
pub id: u32,
pub star_type: String,
pub system: String,
pub body: String,
pub mult: f32,
pub distance: f32,
pub x: f32,
pub y: f32,
pub z: f32,
}
fn get_mult(star_type: &str) -> f32 {
if star_type.contains("White Dwarf") {
return 1.5;
}
if star_type.contains("Neutron") {
return 4.0;
}
1.0
}
#[derive(Debug, Deserialize)]
struct Coords {
x: f32,
y: f32,
z: f32,
}
#[derive(Debug, Deserialize)]
struct Body {
name: String,
#[serde(rename = "type")]
body_type: String,
#[serde(rename = "subType")]
sub_type: String,
#[serde(rename = "distanceToArrival")]
distance: f32,
}
#[derive(Debug, Deserialize)]
struct System {
coords: Coords,
name: String,
bodies: Vec<Body>,
}
pub fn process_galaxy_dump(path: &str) -> std::io::Result<()> {
let fh = File::create("stars.csv")?;
let mut wtr = csv::Writer::from_writer(BufWriter::new(fh));
let mut buffer = String::new();
let mut bz2_reader = std::process::Command::new("7z")
.args(&["x", "-so", path])
.stdout(std::process::Stdio::piped())
.spawn()
.unwrap_or_else(|err| {
eprintln!("Failed to run 7z: {}", err);
eprintln!("Falling back to bzip2");
std::process::Command::new("bzip2")
.args(&["-d", "-c", path])
.stdout(std::process::Stdio::piped())
.spawn()
.expect("Failed to execute bzip2!")
});
let mut reader = BufReader::new(
bz2_reader
.stdout
.as_mut()
.expect("Failed to open stdout of child process"),
);
let mut count = 0;
while let Ok(n) = reader.read_line(&mut buffer) {
if n == 0 {
break;
}
buffer = buffer
.trim()
.trim_end_matches(|c| c == ',')
.trim()
.to_string();
if !buffer.contains("Star") {
continue;
};
if let Ok(sys) = serde_json::from_str::<System>(&buffer) {
for b in &sys.bodies {
if b.body_type == "Star" {
let s = SystemSerde {
id: count,
star_type: b.sub_type.clone(),
distance: b.distance,
mult: get_mult(&b.sub_type),
body: b.name.clone(),
system: sys.name.clone(),
x: sys.coords.x,
y: sys.coords.y,
z: sys.coords.z,
};
wtr.serialize(s)?;
count += 1;
}
}
}
buffer.clear();
}
println!("Total: {}", count);
Ok(())
}