Initial Commit

.gitignore

@@ -0,0 +1,3 @@
+/target
+/Cargo.lock
+/.idea/

Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "schemsearch"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+hematite-nbt = "0.5.2"
+serde = "1.0.152"

src/lib.rs

@@ -0,0 +1,144 @@
+#![allow(unused_variables)]
+
+use crate::pattern_mapper::match_palette;
+use crate::schematic::Schematic;
+
+mod schematic;
+mod pattern_mapper;
+
+#[derive(Debug, Clone, Copy)]
+pub struct SearchBehavior {
+    ignore_block_data: bool,
+    ignore_block_entities: bool,
+    ignore_entities: bool,
+}
+
+pub fn search(
+    data: &Vec<u8>,
+    pattern: &Vec<u8>,
+    search_behavior: SearchBehavior,
+) -> Vec<(u16, u16, u16)> {
+    let schem: Schematic = parse_schematic(data);
+    let pattern_schem: Schematic = parse_schematic(pattern);
+
+    if schem.width < pattern_schem.width || schem.height < pattern_schem.height || schem.length < pattern_schem.length {
+        return vec![];
+    }
+
+    if pattern_schem.palette.len() > schem.palette.len() {
+        return vec![];
+    }
+
+    let (schem, pattern_schem) = match_palette(&schem, &pattern_schem, search_behavior.ignore_block_data);
+
+    let mut matches: Vec<(u16, u16, u16)> = Vec::new();
+
+    for i in 0..schem.width - pattern_schem.width {
+        for j in 0..schem.height - pattern_schem.height {
+            for k in 0..schem.length - pattern_schem.length {
+                let mut match_found = true;
+                for x in 0..pattern_schem.width {
+                    for y in 0..pattern_schem.height {
+                        for z in 0..pattern_schem.length {
+                            let index = (i + x) as usize + (j + y) as usize * schem.width as usize + (k + z) as usize * schem.width as usize * schem.height as usize;
+                            let pattern_index = x as usize + y as usize * pattern_schem.width as usize + z as usize * pattern_schem.width as usize * pattern_schem.height as usize;
+                            if schem.block_data[index] != pattern_schem.block_data[pattern_index] {
+                                match_found = false;
+                                break;
+                            }
+                        }
+                        if !match_found {
+                            break;
+                        }
+                    }
+                    if !match_found {
+                        break;
+                    }
+                }
+                if match_found {
+                    matches.push((i, j, k));
+                }
+            }
+        }
+    }
+
+    return matches;
+
+}
+
+pub(crate) fn normalize_data(data: &String, ignore_data: bool) -> String {
+    if ignore_data {
+        data.split('[').next().unwrap().to_string()
+    } else {
+        data.clone()
+    }
+}
+
+fn parse_schematic(data: &Vec<u8>) -> Schematic {
+    if data[0] == 0x1f && data[1] == 0x8b {
+        // gzip
+        nbt::from_gzip_reader(data.as_slice()).unwrap()
+    } else {
+        // uncompressed
+        nbt::from_reader(data.as_slice()).unwrap()
+    }
+}
+
+mod tests {
+    use std::path::Path;
+    use crate::pattern_mapper::{match_palette, strip_data};
+    use super::*;
+
+    #[test]
+    fn read_schematic() {
+        let schematic = Schematic::load(Path::new("tests/simple.schem"));
+        assert_eq!(schematic.width as usize * schematic.height as usize * schematic.length as usize, schematic.block_data.len());
+        assert_eq!(schematic.palette_max, schematic.palette.len() as i32);
+        println!("{:?}", schematic);
+    }
+
+    #[test]
+    fn test_parse_function() {
+        let file = std::fs::File::open("tests/simple.schem").expect("Failed to open file");
+        let schematic: Schematic = parse_schematic(&std::io::Read::bytes(file).map(|b| b.unwrap()).collect());
+        assert_eq!(schematic.width as usize * schematic.height as usize * schematic.length as usize, schematic.block_data.len());
+        assert_eq!(schematic.palette_max, schematic.palette.len() as i32);
+        println!("{:?}", schematic);
+    }
+
+    #[test]
+    fn test_strip_schem() {
+        let schematic = Schematic::load(Path::new("tests/simple.schem"));
+        let stripped = strip_data(&schematic);
+
+        assert_eq!(stripped.palette.keys().any(|k| k.contains('[')), false);
+        println!("{:?}", stripped);
+    }
+
+    #[test]
+    fn test_match_palette() {
+        let schematic = Schematic::load(Path::new("tests/simple.schem"));
+        let endstone = Schematic::load(Path::new("tests/endstone.schem"));
+
+        let (matched_schematic, matched_endstone) = match_palette(&schematic, &endstone, true);
+
+        println!("{:?}", matched_schematic);
+        println!("{:?}", matched_endstone);
+    }
+
+    #[test]
+    fn test_search() {
+        let file = std::fs::File::open("tests/simple.schem").expect("Failed to open file");
+        let schematic = &std::io::Read::bytes(file).map(|b| b.unwrap()).collect();
+        let file = std::fs::File::open("tests/endstone.schem").expect("Failed to open file");
+        let pattern = &std::io::Read::bytes(file).map(|b| b.unwrap()).collect();
+
+        let matches = search(schematic, pattern, SearchBehavior {
+            ignore_block_data: true,
+            ignore_block_entities: true,
+            ignore_entities: true,
+        });
+
+        println!("{:?}", matches);
+    }
+}

src/pattern_mapper.rs

@@ -0,0 +1,118 @@
+use nbt::Map;
+use crate::normalize_data;
+use crate::schematic::Schematic;
+
+pub(crate) fn strip_data(schem: &Schematic) -> Schematic {
+    let mut data: Vec<i32> = Vec::new();
+
+    let mut palette: Map<String, i32> = Map::new();
+    let mut palette_max: i32 = 0;
+
+    for block in schem.block_data.iter() {
+        let block_name = schem.palette.iter().find(|(_, &v)| v == *block).expect("Invalid Schematic").0;
+        let block_name = block_name.split('[').next().unwrap().to_string();
+
+        let entry = palette.entry(block_name).or_insert_with(|| {
+            let value = palette_max;
+            palette_max += 1;
+            value
+        });
+        data.push(*entry);
+    }
+
+    Schematic {
+        version: schem.version,
+        data_version: schem.data_version,
+        palette,
+        palette_max,
+        block_data: data,
+        block_entities: schem.block_entities.clone(),
+        height: schem.height,
+        length: schem.length,
+        width: schem.width,
+        metadata: schem.metadata.clone(),
+        offset: schem.offset.clone(),
+        entities: None,
+    }
+}
+
+fn match_palette_adapt(schem: &Schematic, matching_palette: Map<String, i32>, ignore_data: bool) -> Vec<i32> {
+    let mut data: Vec<i32> = Vec::new();
+
+    for x in schem.block_data.iter() {
+        let blockname = schem.palette.iter().find(|(_, &v)| v == *x).expect("Invalid Schematic").0;
+        let blockname = if ignore_data { normalize_data(&blockname, ignore_data) } else { blockname.clone() };
+        let block_id = matching_palette.get(&blockname).unwrap_or(&-1);
+        data.push(*block_id);
+    }
+
+    data
+}
+
+pub(crate) fn match_palette(
+    schem: &Schematic,
+    pattern: &Schematic,
+    ignore_data: bool,
+) -> (Schematic, Schematic) {
+    if ignore_data {
+        match_palette_internal(&strip_data(schem), &strip_data(pattern), ignore_data)
+    } else {
+        match_palette_internal(schem, pattern, ignore_data)
+    }
+}
+
+fn match_palette_internal(
+    schem: &Schematic,
+    pattern: &Schematic,
+    ignore_data: bool,
+) -> (Schematic, Schematic) {
+
+    if schem.palette.len() < pattern.palette.len() {
+        panic!("Schematic palette is larger than pattern palette");
+    }
+
+    let mut matching_palette: Map<String, i32> = Map::new();
+    let mut matching_palette_max: i32 = 0;
+
+    for (block_name, block_id) in pattern.palette.iter() {
+        let block_name = normalize_data(block_name, true);
+        let schem_block_id = pattern.palette.get(&block_name).expect("Pattern block not found in schematic palette");
+        matching_palette.insert(block_name, *schem_block_id);
+        matching_palette_max += 1;
+    }
+
+    let mut data_schem: Vec<i32> = match_palette_adapt(&schem, matching_palette.clone(), true);
+
+    let mut data_pattern: Vec<i32> = match_palette_adapt(&pattern, matching_palette.clone(), true);
+
+    let schem = Schematic {
+        version: schem.version.clone(),
+        data_version: schem.data_version.clone(),
+        palette: matching_palette.clone(),
+        palette_max: matching_palette_max.clone(),
+        block_data: data_schem,
+        block_entities: schem.block_entities.clone(),
+        height: schem.height.clone(),
+        length: schem.length.clone(),
+        width: schem.width.clone(),
+        metadata: schem.metadata.clone(),
+        offset: schem.offset.clone(),
+        entities: None,
+    };
+    let pattern = Schematic {
+        version: pattern.version.clone(),
+        data_version: pattern.data_version.clone(),
+        palette: matching_palette.clone(),
+        palette_max: matching_palette_max.clone(),
+        block_data: data_pattern,
+        block_entities: pattern.block_entities.clone(),
+        height: pattern.height.clone(),
+        length: pattern.length.clone(),
+        width: pattern.width.clone(),
+        metadata: pattern.metadata.clone(),
+        offset: pattern.offset.clone(),
+        entities: None,
+    };
+
+    (schem, pattern)
+}

src/schematic.rs

@@ -0,0 +1,58 @@
+use std::path::Path;
+use nbt::{Map, Value};
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub(crate) struct Schematic {
+    #[serde(rename = "Version")]
+    pub(crate) version: i32,
+    #[serde(rename = "DataVersion")]
+    pub(crate) data_version: i32,
+    #[serde(rename = "Metadata")]
+    pub(crate) metadata: Map<String, Value>,
+    #[serde(rename = "Width")]
+    pub(crate) width: u16,
+    #[serde(rename = "Height")]
+    pub(crate) height: u16,
+    #[serde(rename = "Length")]
+    pub(crate) length: u16,
+    #[serde(rename = "Offset")]
+    pub(crate) offset: [i32; 3],
+    #[serde(rename = "PaletteMax")]
+    pub(crate) palette_max: i32,
+    #[serde(rename = "Palette")]
+    pub(crate) palette: Map<String, i32>,
+    #[serde(rename = "BlockData")]
+    pub(crate) block_data: Vec<i32>,
+    #[serde(rename = "BlockEntities")]
+    pub(crate) block_entities: Vec<BlockEntity>,
+    #[serde(rename = "Entities")]
+    pub(crate) entities: Option<Vec<Entity>>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub(crate) struct BlockEntity {
+    #[serde(rename = "Id")]
+    pub(crate) id: String,
+    #[serde(rename = "Pos")]
+    pub(crate) pos: [i32; 3],
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub(crate) struct Entity {
+    #[serde(rename = "Id")]
+    pub(crate) id: String,
+    #[serde(rename = "Pos")]
+    pub(crate) pos: [i32; 3],
+}
+
+impl Schematic {
+    pub(crate) fn load(path: &Path) -> Schematic {
+        let file = std::fs::File::open(path).expect("Failed to open file");
+        let schematic: Schematic = match nbt::from_gzip_reader(file) {
+            Ok(schem) => schem,
+            Err(e) => panic!("Failed to parse schematic: {}", e),
+        };
+        schematic
+    }
+}